稀疏注意力 Sparse Attention：局部与全局结合的注意力模式

FreeGuideOnline 最新 2026-06-21

什么是稀疏注意力？

在标准自注意力机制中，每个词元（token）会与序列中所有其他词元计算注意力权重，这种全局计算带来了平方级的内存和计算复杂度 (O(n^2))。当处理长文本、高分辨率图像或长序列数据时，标准注意力会迅速成为瓶颈。

稀疏注意力（Sparse Attention） 是一类高效注意力变体，它不再让每个词元关注所有词元，而是有选择地只关注一部分其他词元。通过设计特定的稀疏连接模式，使复杂度从 (O(n^2)) 降至接近线性，同时尽量保持模型的表达能力。

核心思想是：用局部+全局的混合模式取代全连接，让信息依然能够流动，但成本大幅降低。

标准注意力的计算瓶颈

先回顾标准缩放点积注意力：

[ \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V ]

其中 (Q, K, V \in \mathbb{R}^{n \times d})，(n) 为序列长度，(d) 为维度。计算流程：

  • 计算相似度矩阵 (QK^T)：复杂度 (O(n^2 d))
  • softmax 操作：(O(n^2))
  • 加权求和：(O(n^2 d))

当 (n) 到达数千甚至上万时，内存和速度都难以承受。稀疏注意力的目标就是减少参与计算的位置对数量。

稀疏注意力如何工作：模式与原理

稀疏注意力通过定义一个注意力掩码矩阵 (M) 来限制每个查询可以关注的位置：

[ \text{SparseAttn}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \odot M \right) V ]

其中 (M) 中的元素为 0 或 (-\infty)（softmax 前），控制哪些位置被屏蔽。

常见的稀疏模式

稀疏模式决定了信息流动的路径。一个好的模式需要在全局视野和局部细节之间取得平衡。

局部与全局结合的稀疏模式详解

局部注意力（Local Attention）

只关注每个词元周围一个固定窗口大小的邻居。

  • 优点：完美捕捉局部上下文，计算量与 (n \times w) 成正比（(w) 为窗口大小，通常远小于 (n)）。
  • 缺点：长距离依赖需要多层才能逐渐传播，信息流动慢。

掩码示例（窗口大小为 2）

位置  0 1 2 3 4
0     1 1 0 0 0
1     1 1 1 0 0
2     0 1 1 1 0
3     0 0 1 1 1
4     0 0 0 1 1

全局注意力（Global Attention）

指定少数几个特殊词元（如句子开头符号 [CLS] 或特殊全局标记）关注所有位置，同时所有位置也可以关注这些全局词元。

  • 优点：为模型提供“摘要路线”，长距离信息可以快速通过这些全局节点传播。
  • 缺点：如果全局词元太少，信息带宽受限；过多则计算量回升。

掩码示例（位置0和位置2为全局标记）

位置  0 1 2 3 4
0     1 1 1 1 1    ← 全局标记0关注全部
1     1 0 1 0 0    ← 普通词元关注全局标记
2     1 1 1 1 1    ← 全局标记2关注全部
3     1 0 1 0 0
4     1 0 1 0 0

带全局节点的局部注意力（Local + Global）

这是最实用的一类稀疏模式，将局部窗口与少量的全局节点结合：

  • 每个词元既在其局部窗口内执行注意力，也关注预先定义的全局词元。
  • 全局词元可以关注所有位置，扮演信息枢纽的角色。
  • 复杂度：(O(n \times (w + g)))，其中 (w) 是窗口大小，(g) 是全局节点数量。
graph LR
    A[全局节点] -->|关注所有| B[序列]
    B -->|关注全局节点+窗口| A

其他经典稀疏模式

  • 带状稀疏（Band）：纯局部窗口的泛化，可以使用滑动窗口，窗口大小固定。
  • 扩张/空洞窗口（Dilated Window）：像膨胀卷积一样，以固定间隔跳过一些位置，在窗口内关注，用更少的计算捕捉更远的上下文。
  • 随机稀疏（Random Sparse）：每个查询随机采样少量键，防止模型过拟合固定模式，增强鲁棒性。
  • 块状稀疏（Block Sparse）：将序列分块，查询只与特定块内的键交互，便于硬件加速。

代表模型中的稀疏注意力实现

Longformer：滑动窗口+全局注意力

Longformer 提出了三种稀疏注意力模式，并可混合使用：

  1. 滑动窗口注意力：每个词元关注其左右各 (w/2) 个词元（典型 (w=512)）。计算量随序列长度线性增长。
  2. 膨胀滑动窗口：在窗口基础上引入间隙，以指数级增加感受野。
  3. 全局注意力：对任务特定词元（如分类符号、问题符号）使用全局注意力，其余词元使用滑动窗口。 Longformer 通过自定义 CUDA 内核高效实现了这种模式，使得处理超过 4000 长度的文档成为可能。

BigBird：随机+局部+全局

BigBird 证明，稀疏注意力也能保持 Transformer 的通用近似能力。它同时使用三类边：

  • 随机注意力：每个查询随机关注 (r) 个其他位置。
  • 局部窗口注意力：固定大小的滑动窗口。
  • 全局注意力：少量全局词元（如 CLS 或可学习的全局令牌）。

这种组合被证明在理论上可以模拟全注意力，且存储和计算开销降至 (O(n))（适中的 (w, r, g) 下）。

ETC（Extended Transformer Construction）：全局-局部-随机

ETC 在编码器中将输入分为全局–长格式结构，同样结合了全局注意力和局部注意力，并加入随机连接以覆盖长距依赖。

稀疏注意力的优点与权衡

优点

  • 线性复杂度：将序列长度依赖从二次降低到线性（或接近线性），支持数万甚至更长的序列。
  • 内存占用小：显存中只需存储稀疏的注意力矩阵，不易 OOM。
  • 长文本能力：能直接处理长文档、长对话、基因组序列等，无需截断或分段处理。

权衡与挑战

  • 信息损失：某些细粒度、跨段落的长距离关系可能因稀疏而丢失，需要全局节点或随机连接补偿。
  • 实现复杂性：高效实现需要自定义稀疏算子或块计算，尤其在 GPU 上，因为稀疏矩阵乘法的硬件利用率不如密集计算。
  • 任务相关性：对于只需要局部上下文的任务（如语法纠错），稀疏注意力足够；而对于需要复杂全局推理的任务（如文档级问答、共指消解），需要仔细调整全局节点的位置和数量。

实践：设计自己的稀疏注意力

如果你尝试为模型添加稀疏注意力，可以从以下步骤开始：

  1. 定义窗口长度 (w)：从 128 或 256 开始，根据任务局部依赖范围调整。
  2. 插入全局词元：至少放置 2 个全局节点（例如 BERT 的 [CLS] 和 [SEP]），或为文档各段落插入专用全局标记。
  3. 补充随机边（可选）：增加 10% ~ 20% 的随机注意力连接，提升鲁棒性，尤其适合理论完备性要求高的场景。
  4. 验证替换效果：先用小型数据集消融，比较与全注意力的精度差异。
  5. 高效实现：如果训练规模大，考虑使用块稀疏库（如 torch_blocksparse）或基于滑动窗口的优化操作（如 flash_attention 的局部版本）。

小结与延伸阅读

稀疏注意力通过主动放弃全连接模式，以较小的精度代价换取了处理长序列的能力。局部和全局相结合的方案是其中应用最广、效果最稳定的设计。掌握稀疏注意力，你可以：

  • 让普通 Transformer 处理整本小说级别的输入。
  • 设计更高效的视觉 Transformer（如 Swin Transformer 中的窗口注意力）。
  • 在生产环境中大幅降低推理延迟和成本。

如果你想深入了解，建议阅读以下论文：

  • Longformer: The Long-Document Transformer (Beltagy et al., 2020)
  • Big Bird: Transformers for Longer Sequences (Zaheer et al., 2020)
  • ETC: Encoding Long and Structured Inputs in Transformers (Ainslie et al., 2020)

稀疏注意力的世界远不止这些基础模式，还有可学习稀疏模式、基于内容的路由注意力等方向，等待你继续探索。