摘要：长上下文建模对于下一代语言模型至关重要，然而标准注意力机制的高计算成本带来了显著的计算挑战。稀疏注意力为提高效率同时保持模型能力提供了一个有前景的方向。我们提出了NSA（原生可训练的稀疏注意力机制），该机制将算法创新与硬件对齐的优化相结合，以实现高效的长上下文建模。NSA采用动态分层稀疏策略，将粗粒度的标记压缩与细粒度的标记选择相结合，既保留了全局上下文感知能力，又保证了局部精确性。我们的方法在稀疏注意力设计方面取得了两项关键创新：（1）我们通过算术强度平衡的算法设计实现了显著的速度提升，并针对现代硬件进行了实现优化。（2）我们实现了端到端的训练，减少了预训练计算量，同时没有牺牲模型性能。如论文图1所示，实验表明，使用NSA预训练的模型在通用基准测试、长上下文任务和指令推理方面均保持或超过了全注意力模型的表现。同时，在64k长度的序列上，NSA在解码、前向传播和后向传播方面均实现了相对于全注意力的显著速度提升，验证了其在模型整个生命周期中的高效性。Huggingface链接：Paper page，论文链接：2502.11089

1. 引言

背景与动机：

• 长上下文建模的重要性：长上下文建模是下一代大型语言模型的关键能力，它对于深入推理、代码生成、多轮自主代理系统等应用至关重要。然而，标准注意力机制的高计算成本成为处理长序列时的瓶颈。
• 稀疏注意力的潜力：稀疏注意力机制通过选择性地计算关键查询-键对，可以在显著降低计算开销的同时保持模型性能，为解决长上下文建模的计算挑战提供了有前景的方向。

研究目标：

• 本文提出了一种名为NSA（Native Sparse Attention）的原生可训练稀疏注意力机制，旨在通过结合算法创新与硬件对齐的优化，实现高效的长上下文建模。

2. 稀疏注意力方法的重新思考

高效推理的幻觉：

• 许多现有的稀疏注意力方法虽然实现了注意力计算的稀疏性，但未能在实际推理延迟上取得相应减少。这主要是由于以下挑战：
  • 阶段受限的稀疏性：一些方法仅在自回归解码阶段应用稀疏性，而在预填充阶段需要高计算成本的预处理（如注意力图计算和索引构建）。
  • 与现代注意力架构的不兼容性：某些稀疏注意力方法难以适应像多查询注意力（MQA）和分组查询注意力（GQA）这样的现代解码高效架构，这些架构通过跨多个查询头共享键值对显著减少了内存访问瓶颈。

可训练稀疏性的神话：

• 现有的稀疏注意力方法主要侧重于推理阶段，而忽略了训练阶段的计算挑战。这导致了以下问题：
  • 性能退化：推理阶段的后置稀疏性迫使模型偏离其预训练优化轨迹，可能导致模型性能下降。
  • 训练效率低下：许多理论上可训练的稀疏注意力方法在实际训练中效率低下，这主要是由于非连续内存访问模式和不可训练的组件。

原生稀疏性的必要性：

• 本文提出的NSA机制旨在解决推理效率和训练可行性问题，通过设计一种原生稀疏注意力框架，既提高计算效率又满足训练要求。

3. 方法论

背景：

• 注意力机制：在语言建模中广泛使用，其中每个查询标记qt​通过计算与所有先前键标记k:t​的相关性分数，来生成值标记v:t​的加权和。
• 算术强度：定义为计算操作与内存访问的比率，它内在地决定了硬件上的算法优化目标。

总体框架：

• NSA通过用更紧凑、信息密集的关键值对集合K~t​,V~t​替换原始键值对k:t​,v:t​，来优化注意力输出。这些紧凑的关键值对是基于当前查询qt​和上下文内存k:t​,v:t​动态构建的。

算法设计：

• 标记压缩：通过将序列块中的键或值聚合为块级表示，获得压缩的键和值，这些表示捕获了整个块的信息，减少了注意力计算的计算负担。
• 标记选择：仅使用压缩键和值可能会丢失重要的细粒度信息，因此NSA还通过选择性地保留个别键和值来保留这些信息。选择策略基于块的重要性分数，这些分数可以通过压缩标记的注意力计算中间结果来推导。
• 滑动窗口：为了处理局部模式，NSA引入了一个专用的滑动窗口分支，该分支显式地处理局部上下文，允许其他分支（压缩和选择）专注于学习它们各自的特征，而不受局部模式的干扰。

内核设计：

• 为了实现与FlashAttention相当的速度提升，NSA在Triton上实现了硬件对齐的稀疏注意力内核。内核设计的特点包括以GQA组为中心的数据加载、共享的键值获取以及在外层循环上使用网格调度器。

4. 实验

预训练设置：

• 实验采用了一个结合GQA和混合专家（MoE）的270亿参数模型骨干，模型包含30层，隐藏维度为2560。MoE结构采用DeepSeekMoE，包含72个路由专家和2个共享专家。

基线方法：

• 除了与全注意力基线进行比较外，还评估了几种最先进的推理阶段稀疏注意力方法，包括H2O、infLLM、Quest和Exact-Top。

性能比较：

• 通用基准测试：在MMLU、MMLU-PRO、CMMLU等基准测试上，尽管具有稀疏性，NSA在大多数指标上仍优于或持平于全注意力基线。
• 长上下文基准测试：在LongBench上，NSA在包括单文档QA、多文档QA、合成代码任务在内的多个子集上均优于基线方法，包括全注意力模型。
• 指令推理评估：在AIME 24基准测试上，经过监督微调后的NSA-R模型在8k和16k序列长度下均显著优于全注意力基线模型Full Attention-R。

5. 效率分析

训练速度：

• 在8个A100 GPU的系统上，与全注意力相比，NSA在64k长度序列上实现了高达9.0倍的前向传播和6.0倍的后向传播速度提升。随着序列长度的增加，速度优势变得更加明显。

解码速度：

• NSA的解码速度主要由内存访问瓶颈决定，与全注意力相比，NSA在64k长度序列上实现了高达11.6倍的速度提升。随着解码长度的增加，延迟显著降低。

6. 讨论

替代标记选择策略的挑战：

• 在设计NSA之前，作者探索了将现有稀疏注意力方法适应到训练阶段的可能性，但遇到了各种挑战，如非平凡的计算开销、操作器优化困难和实现约束。

注意力分数块模式可视化：

• 通过可视化全注意力变压器的注意力图，作者观察到注意力分数往往表现出块状聚类特性，这为NSA的设计提供了灵感，即基于空间连续性选择键块可能是一个有前景的方法。

7. 相关工作

固定稀疏模式：

• 这类方法允许查询仅在固定窗口内计算注意力，如SlidingWindow和StreamingLLM。这些方法有效降低了内存和计算成本，但其固定的忽略上下文模式限制了它们在需要全上下文理解的任务上的性能。

动态标记修剪：

• 这类方法通过动态地逐出或选择键值对来减少KV缓存的内存使用，如H2O和SnapKV。这些方法能够根据查询的重要性动态调整稀疏性，但实现高效的训练和推理仍然具有挑战性。

查询感知选择：

• 这类方法通过估计每个块的重要性来选择重要的键值块，如Quest和InfLLM。这些方法能够根据查询内容自适应地调整稀疏性，但在实现高效的端到端训练方面仍面临挑战。

8. 结论

本文提出了NSA，一种硬件对齐的稀疏注意力架构，用于高效的长上下文建模。通过结合分层标记压缩和块级标记选择，NSA在加速训练和推理的同时保持了全注意力模型的性能。实验结果表明，NSA在通用基准测试、长上下文基准测试和指令推理评估上均表现出色，同时在计算延迟方面实现了显著的速度提升。

详细分析

算法创新：

• NSA通过动态分层稀疏策略，将粗粒度的标记压缩与细粒度的标记选择相结合，实现了全局上下文感知和局部精确性的平衡。这种策略既减少了计算开销，又保留了模型的关键信息。

硬件优化：

• NSA针对现代GPU硬件进行了实现优化，通过算术强度平衡的算法设计和硬件对齐的内核设计，实现了显著的速度提升。这些优化包括以GQA组为中心的数据加载、共享的键值获取以及在外层循环上使用网格调度器。

端到端训练：

• NSA支持端到端的训练，能够在不牺牲模型性能的情况下减少预训练计算量。这对于开发更强大的长上下文模型至关重要，因为它允许模型在训练过程中自适应地调整稀疏性模式。

实验验证：

• 通过在多个基准测试上的实验评估，作者证明了NSA在保持或超过全注意力模型性能的同时，显著提高了计算效率。特别是在长上下文任务和指令推理方面，NSA表现出了出色的性能。

未来工作：

• 尽管NSA在稀疏注意力设计方面取得了显著进展，但仍有许多挑战需要解决。例如，如何进一步优化NSA的内核设计以提高在更长序列上的效率？如何将NSA扩展到其他类型的语言模型任务中？这些问题将是未来研究的重点。

综上所述，NSA作为一种原生可训练的稀疏注意力机制，为高效的长上下文建模提供了新的解决方案。通过结合算法创新与硬件对齐的优化，NSA在保持模型性能的同时显著提高了计算效率，为开发更强大的下一代语言模型奠定了坚实的基础。