Appearance
Sparse Attention
稀疏注意力 (Sparse Attention) 通过限制每个 token 只关注部分其他 token,将自注意力的复杂度从 $O(n^2)$ 降低到 $O(n \log n)$ 或 $O(n)$,从而支持处理更长的序列。这是解决标准 Transformer 在长序列上可扩展性问题的关键方向。
问题背景
标准密集注意力的复杂度:
- 计算: $O(n^2)$ 每层注意力的时间复杂度
- 内存: $O(n^2)$ 注意力矩阵的存储复杂度
- 实际限制: 8K 上下文已是许多模型的实际上限,长文档/视频/音频需要更长序列
稀疏注意力的主要类型
1. 固定模式稀疏 (Fixed Pattern Sparsity)
预定义固定的稀疏模式,不随输入变化:
- 滑动窗口 (Sliding Window): 每个 token 只关注周围 $w$ 个邻居 token
- 膨胀 (Dilated): 在滑动窗口内跳跃采样,进一步降低复杂度
- 块稀疏 (Block Sparse): 将序列分块,块内密集注意力,块间稀疏连接
代表: Sparse Transformers (OpenAI, 2019)
2. 全局+局部注意力 (Global + Local)
部分 token 具有全局视野,其他 token 只局部关注:
- 全局 token: CLS token、句子第一个 token 等,可见所有位置
- 局部 token: 只关注周围窗口内的 token
代表: Longformer (Allen AI, 2020)
3. 随机+窗口+全局 (Random + Window + Global)
结合多种连接类型,理论上证明可近似全注意力:
- 随机连接: 每个 token 随机连接少量远处 token
- 窗口连接: 局部邻居连接
- 全局连接: 特殊 token 的全局视野
代表: Big Bird (Google, 2021)
4. 核方法/特征映射 (Kernel Methods)
将 softmax 注意力重写为核函数内积,实现线性复杂度:
- Performer: 使用正交随机特征映射 (FAVOR+) 近似 softmax 内积
- Linear Attention: 将 softmax 改为线性核,利用矩阵乘法结合律实现 $O(n)$ 复杂度
5. 哈希方法 (Hashing-based)
通过局部敏感哈希 (LSH) 将相似 token 映射到同一桶内:
- Reformer: LSH 注意力 + 可逆层,实现线性复杂度
稀疏注意力 vs FlashAttention
| 对比维度 | 稀疏注意力 | FlashAttention |
|---|---|---|
| 核心思想 | 近似/稀疏化注意力矩阵 | 精确注意力,IO 优化 |
| 复杂度 | $O(n)$ 到 $O(n \log n)$ | 仍为 $O(n^2)$,但常数因子大降 |
| 内存访问 | 依赖具体模式 | 分块加载到 SRAM,减少 HBM 访问 |
| 精度 | 近似,可能有信息损失 | 精确,无近似 |
| 适用场景 | 极长序列 (>100K) | 中等长度 (<100K),训练推理加速 |
| 代表方法 | Longformer, BigBird, Reformer | FlashAttention-1/2/3 |
关键区别: 稀疏注意力通过减少计算量来降低复杂度,而 FlashAttention 通过 IO 优化在不改变算法复杂度的情况下大幅加速。两者可以结合使用。
关键论文
| 论文 | 作者 | 年份 | 核心贡献 |
|---|---|---|---|
| "Longformer" | Beltagy et al. | 2020 (Allen AI) | 滑动窗口+膨胀+全局注意力,线性复杂度 |
| "Big Bird" | Zaheer et al. | 2021 (Google) | 随机+窗口+全局,理论上近似全注意力 |
| "Reformer" | Kitaev et al. | 2020 (Google) | LSH 注意力,可逆层,线性复杂度 |
| "Performer" | Choromanski et al. | 2021 (Google) | FAVOR+ 核方法,无偏/有偏近似 softmax |
| "Sparse Transformers" | Child et al. | 2019 (OpenAI) | 步长/固定模式稀疏,生成长序列 |
| "FlashAttention" | Dao et al. | 2022 (Stanford) | IO-aware 精确注意力,分块计算 |
重要框架与工具
- FlashAttention: 已成为 LLM 训练事实标准,支持到 FlashAttention-3
- xFormers: Meta 开源模块化注意力实现(稀疏/内存高效/Flash)
- Hugging Face Transformers: 集成 Longformer、BigBird、Reformer 等
- vLLM / SGLang: 利用 PagedAttention 实现高效稀疏/变长推理
与相关概念的关系
- Sparse vs Linear Attention: Linear Attention 将 softmax 改为核特征图,实现真正线性复杂度
- Sparse vs Recurrence: RNN/State Space Models (Mamba, RWKV) 是序列建模的另一条路径,与稀疏注意力竞争
- Sparse vs Long Context: 稀疏注意力是实现长上下文的关键技术之一,但不是唯一方案
Sources
- Beltagy et al., "Longformer: The Long-Document Transformer" (arXiv 2020)
- Zaheer et al., "Big Bird: Transformers for Longer Sequences" (NeurIPS 2021)
- Kitaev et al., "Reformer: The Efficient Transformer" (ICLR 2020)
- Choromanski et al., "Rethinking Attention with Performers" (ICML 2021)
- Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (NeurIPS 2022)
- Tay et al., "Efficient Transformers: A Survey" (arXiv 2020)