DeepSeek开源新一代注意力机制NSA，梁文锋亲自挂名，长上下文建模更高效！
DeepSeek开源新一代注意力机制NSA，梁文锋亲自挂名，长上下文建模更高效！
DeepSeek联合创始人梁文锋亲自挂名的研究团队在arXiv上发表了一篇重要论文，提出了一种全新的注意力机制架构NSA（Native Sparse Attention，原生稀疏注意力）。该架构在保持模型性能的同时显著提升了计算效率，为开源AI社区提供了宝贵的参考。
由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个百分点。这一结果表明，稀疏注意力的预训练不仅没有损害模型能力，反而通过过滤无关注意力路径增强了模型的推理性能。
图丨全注意力基线和NSA在一般基准测试上的预训练性能比较（来源：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社区提供了宝贵的参考。