DeepSeek开源注意力机制NSA，梁文锋挂名，长上下文建模更高效！

在马斯克的 Grok 3 发布会反复与 DeepSeek 进行比较之后，DeepSeek 不语，只是又在 X 上公布了一项新的技术成果。

图丨相关推文（来源：X）
由 DeepSeek 联合创始人梁文锋亲自挂名的研究团队，在arXiv上发表了一篇题为“Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention”的论文，提出了一种全新的注意力机制架构 NSA（Native Sparse Attention，原生稀疏注意力）。
熟悉 DeepSeek-R1 的用户都知道，这款模型虽然在许多方面表现出色，但比较遗憾的一点在于，其输入上下文能力方面相对不足。而这次发布的 NSA 架构就主要致力于解决当前大模型长文本处理中的关键瓶颈问题。
在传统注意力机制中，当序列长度达到 64K 时，注意力计算可能占用总延迟的 70-80%，这种计算开销已经成为制约模型性能的重要因素。
NSA 的核心技术创新主要体现在两个方面。首先是其独特的分层稀疏注意力设计。该架构将输入序列按时间维度划分为连续的 block，并通过三条并行的注意力分支进行处理：压缩注意力（Compressed Attention）通过可学习的 MLP 将每个 block 压缩成单一表示，用于捕获粗粒度的全局信息；选择性注意力（Selected Attention）则保留最重要的 fine-grained token 信息；滑动窗口注意力（Sliding Attention）用于处理近期的局部上下文。这种分层设计使得模型能够在保持表达能力的同时大幅降低计算复杂度。

图丨 NSA 架构概览（来源：arXiv）
第二个创新点在于其硬件友好的实现优化。研究团队基于 Triton 开发了专门的 kernel，通过“Group-Centric Data Loading”策略，将同一GQA（Grouped-Query Attention）组内的所有 query head 同时加载到 SRAM 中处理。这种设计不仅最大化了 Tensor Core 的利用率，还通过优化的循环调度消除了冗余的 KV 数据传输。特别是在处理 block 化的稀疏注意力时，NSA 采用了连续的内存访问模式，这与现代 GPU 架构高度契合。
值得一提的是，NSA 突破性地实现了端到端的可训练稀疏注意力。与现有方法往往在推理阶段才引入稀疏化不同，NSA 从预训练阶段就开始使用稀疏注意力机制。这种“原生”设计使得注意力模块能够与模型其他组件协同优化，形成更优的稀疏模式。为了支持稳定训练，研究团队还为三条注意力分支设计了独立的 key 和 value 参数，这种设计虽然带来了轻微的参数开销，但有效防止了局部模式对其他分支学习的干扰。

图丨 NSA 的内核设计（来源：arXiv）
在具体实现上，NSA 在处理 64K 长度序列时，每个解码步骤只需要加载 ⌊(s-l)/d⌋ 个压缩 token、nl' 个选择性 token 和 w 个近邻 token，其中 s 是缓存序列长度，l 是 block 长度，d 是滑动步长，n 是选择的 block 数量，w 是滑动窗口大小。这种设计使得内存访问量随序列长度的增长维持在一个较低水平，从而实现了接近理论极限的加速效果。
研究团队对 NSA 进行了全方位的性能验证。实验采用了一个基于 GQA 和 MoE（Mixture-of-Experts）的 27B 参数 backbone，包含 30 层网络结构，隐藏维度为 2560。为确保实验的可比性，研究团队采用了与全量注意力模型完全相同的训练流程，包括在 270B token 的 8K 长度文本上进行预训练，随后使用 YaRN 方法在 32K 长度文本上进行延续训练和监督微调。
在通用能力评测中，NSA 展现出了超出预期的表现。在涵盖知识、推理和编程能力的九项基准测试中，包括 MMLU、MMLU-PRO、CMMLU、BBH、GSM8K、MATH、DROP、MBPP 和 HumanEval，NSA 在七项上超越了全量注意力基线。尤其值得关注的是在推理相关任务上的显著提升，如在 DROP 任务上提升了 4.2 个百分点，在 GSM8K 上提升了 3.4 个百分点。这一结果表明，稀疏注意力的预训练不仅没有损害模型能力，反而通过过滤无关注意力路径增强了模型的推理性能。

（来源：arXiv）
在长文本处理能力的专项测试上，NSA 的优势更为突出。在 64K 长度的“大海捞针”测试中，NSA 实现了全位置的完美检索准确率。在 LongBench 评测集上，NSA 的平均得分达到 0.469，显著超过了包括全量注意力在内的所有基线方法。具体来看，在多跳问答任务 HPQ 和 2Wiki 上分别提升了 8.7 和 5.1 个百分点，在代码理解任务 LCC 上提升了 6.9 个百分点，在段落检索任务 PassR-en 上提升了 7.5 个百分点。

图丨解码期间每次注意力操作的内存访问量（来源：arXiv）
在推理能力的深入测试中，研究团队还探索了 NSA 在进阶数学推理方面的表现。通过从 DeepSeek-R1 模型蒸馏数学推理能力，在 10B 个 32K 长度的数学推理轨迹上进行监督微调后，NSA 在美国邀请数学竞赛基准测试上取得了显著进展。在 8K 上下文限制下，NSA 比基线模型提升了 7.5 个百分点；在扩展到 16K 上下文时，仍保持了 5.4 个百分点的优势。这一结果验证了 NSA 在保持长程逻辑依赖方面的独特优势。
在计算效率方面，NSA 的表现同样令人瞩目。在 64K 序列长度的场景下，在解码、前向传播和反向传播三个阶段分别实现了 11.6 倍、9.0 倍和 6.0 倍的加速比。更重要的是，这种加速优势会随着序列长度的增加而进一步扩大，这对于未来更长上下文的处理提供了可行方案。
尽管 NSA 取得了显著的成果，但也还存在几个值得深入探索的方向。比如，稀疏注意力模式的学习过程还有优化空间。目前的方案虽然实现了端到端训练，但如何让模型学习到更优的稀疏模式，特别是在更大规模模型上的表现，还需要进一步研究。此外，NSA 提供的 Triton 实现为业界提供了很好的参考，但在实际部署中，还需要考虑不同硬件平台的适配、推理服务的稳定性等问题。
不过，NSA 的表现已经证明：通过精心的算法设计和硬件协同优化，我们可以在保持模型性能的同时显著提升计算效率，其出现无疑又为整个开源 AI 社区提供了宝贵的参考。
参考资料：
1.https://arxiv.org/abs/2502.11089
运营/排版：何晨龙