稀疏注意力论文解读
稀疏注意力通过将 O(n²) 的全局注意力替换为稀疏的注意力模式，在保持模型性能的同时大幅降低计算和内存开销，是处理长序列的关键技术方向。
动机
标准自注意力的每个 Token 需要关注所有其他 Token，形成 O(n²) 的全连接注意力图。但实际中，大部分 Token 对之间的注意力权重接近零——注意力天然是稀疏的。
稀疏注意力模式
局部窗口注意力
每个 Token 仅关注邻近的 w 个 Token：
python
class LocalAttention(nn.Module):
"""局部窗口注意力"""
def __init__(self, d_model, window_size=256):
super().__init__()
self.window_size = window_size
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)
def forward(self, x):
B, T, D = x.shape
w = self.window_size
# 填充使长度整除窗口
pad = (w - T % w) % w
x_pad = F.pad(x, (0, 0, 0, pad))
T_pad = T + pad
# 重塑为窗口
x_windows = x_pad.reshape(B, T_pad // w, w, D)
qkv = self.qkv(x_windows)
q, k, v = qkv.chunk(3, dim=-1)
# 窗口内注意力
attn = (q @ k.transpose(-2, -1)) / (D ** 0.5)
attn = F.softmax(attn, dim=-1)
out = (attn @ v).reshape(B, T_pad, D)[:, :T, :]
return self.proj(out)组合稀疏模式
Longformer 提出了组合稀疏模式：
- 滑动窗口：局部上下文建模
- 扩张滑动窗口：增大感受野
- 全局注意力：少数关键 Token（如 [CLS]）使用全局注意力
BigBird 的块稀疏注意力
BigBird 将注意力组织为块级别：
- 随机块：随机选择块间连接，保证图连通性
- 窗口块：局部滑动窗口
- 全局块：关键 Token 全局连接
稀疏注意力的理论保证
Beltagy et al. (2020) 证明，包含随机 + 窗口 + 全局三种模式的稀疏注意力图是 Turing 完备的，理论上可以逼近全注意力。这为稀疏注意力的有效性提供了理论基础。
主要论文对比
| 模型 | 注意力模式 | 复杂度 | 最大长度 |
|---|---|---|---|
| Transformer | 全连接 | O(n²) | ~512 |
| Sparse Transformer | 固定稀疏模式 | O(n√n) | ~8K |
| Longformer | 窗口+全局 | O(n×w) | ~16K |
| BigBird | 块稀疏 | O(n) | ~32K |
| Linformer | 低秩近似 | O(n) | ~64K |
实际应用
稀疏注意力在以下场景特别有价值：
- 长文档处理：法律合同、学术论文全文理解
- 代码生成：需要理解整个代码库上下文
- 对话系统：维护完整对话历史
- 基因组学：处理极长 DNA 序列