Skip to content

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, ReformerFlashAttention-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)

AI Knowledge Base — 持续积累