Sparse Attention

稀疏注意力是一类注意力机制：每个查询只对选定的一部分键评分，而不是对序列中的所有键评分。对 N 个词元进行完整的自注意力，每个头需要计算 N² 个分数，因此文档长度翻倍，注意力开销就会变成四倍；而每个词元只关注 k 个其他词元的模式，开销为 N·k。注意力模式可以预先固定，例如局部窗口；也可以根据内容选择，例如为每个查询选取得分最高的键。OpenAI 于 2019 年推出的 Sparse Transformer 使用步幅模式，以 O(N√N) 的开销建模长达数万步的序列；当前的大型语言模型（LLM）则使用稀疏层来处理超过 100,000 个词元的上下文窗口。

常见稀疏模式#

固定模式会预先决定要计算哪些位置对的分数。滑动窗口注意力让每个词元只能看到前面的 w 个词元：Mistral 7B使用 4,096 个词元的窗口；OpenAI 的开放权重 gpt-oss 模型则交替使用密集层和 128 个词元的窗口层，这与 GPT-3 使用的密集层和局部带状层交替模式相同。每个窗口层都会将信息向前传递 w 个词元，因此即使堆叠多层，远距离词元之间仍然可以建立联系。步幅模式会加入每隔 s 个位置的较早词元，使信息也能在单层中跨越较远距离。全局词元会关注所有位置，也会被所有位置关注：Longformer将这类词元分配给分类词元等任务词元，BigBird还加入随机连接，而注意力汇聚点则让每个查询都能看到最前面的几个词元，从而保持流式生成的稳定性。

基于内容的模式会根据数据选择位置对。DeepSeek 的 Native Sparse Attention（NSA，2025）将键块的压缩摘要、最相关块的选择结果和滑动窗口结合起来，并采用端到端训练。DeepSeek-V3.2中的 DeepSeek Sparse Attention 会运行一个小型“闪电索引器”，为每个查询给所有更早的词元评分，然后只对得分最高的 2,048 个词元计算完整注意力，因此无论上下文增长多长，每个查询的主要注意力开销都保持不变。生产环境中的模型通常会组合使用这些规则：

让模式真正节省计算#

如果在密集分数矩阵上应用掩码，仍然会计算所有分数；要实现加速，内核必须跳过矩阵中的空白块。Flash attention内核直接支持滑动窗口，而 PyTorch 的 FlexAttention 可将任意掩码函数编译为块稀疏 GPU 内核。它的 create_block_mask 展示了这种内核在 1,024 个词元的窗口加上四个汇聚词元时会跳过哪些块：

from torch.nn.attention.flex_attention import create_block_mask

seq_len, window, sinks = 8192, 1024, 4


def local_plus_sinks(b, h, q_idx, kv_idx):
    """Causal attention to the previous 1,024 tokens plus the first 4 tokens."""
    causal = q_idx >= kv_idx
    return causal & ((q_idx - kv_idx <= window) | (kv_idx < sinks))


# FlexAttention 会将掩码划分为 128 x 128 的块，以便其内核跳过空块
block_mask = create_block_mask(local_plus_sinks, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device="cpu")
print(block_mask)

打印结果给出 sparsity=85.47%，即 8,192 × 8,192 分数矩阵中不包含任何允许位置对的 128 × 128 块所占比例，并将模式绘制为表示窗口的对角带，以及左侧表示汇聚词元的一列。当 flex_attention 在 torch.compile 下运行时，其内核会完全跳过这些块。

稀疏注意力与相关术语的区别#

• Flash attention：一种考虑 I/O 的内核，它通过在快速片上内存中分块处理，对所有位置对计算精确注意力。它的分块也是块稀疏内核跳过的单位。
• 线性注意力：保留所有位置对，但用特征映射取代 softmax，因此所有过去的键和值都能压缩进固定大小的状态。
• 可变形注意力：一种视觉机制，每个查询会关注特征图中学习得到的偏移位置上的少数采样点，而不是固定模式中的位置。
• 混合专家模型（MoE）：参数稀疏性的一种形式，每个词元只激活少数专家网络，而其注意力模式仍然是密集的。

应用场景#

长上下文语言模型是主要应用场景：稀疏层可以限制计算量，窗口层还可以限制 KV 缓存，使模型能够在单个请求中总结长篇报告或搜索大型代码库。BigBird 将其稀疏模式用于 DNA 序列，以预测启动子区域；DNA 输入通常包含数千个碱基对，远长于典型句子。在计算机视觉中，词元数量会随图像分辨率增加，因此模型会将注意力限制在局部区域：Swin Transformer在移位窗口内计算注意力。YOLO12是由布法罗大学和中国科学院大学的研究人员提出、可在 Ultralytics 软件包中运行的注意力型检测器，它使用区域注意力：对于 640 像素图像，其步幅为 16 的主干网络阶段会将 40 × 40 的特征图分成四个水平条带，每个条带包含 400 个词元，并且只在每个条带内计算注意力，评分数量仅为对全部 1,600 个词元进行完整注意力所需 256 万个位置对的四分之一。