DeepSeek推出NSA：加速AI模型训练与推理
2月18日，人工智能公司DeepSeek发布了一项突破性研究成果——新型稀疏注意力机制NSA（Native Sparse Attention），相关论文详细介绍了这一技术如何实现超快速的长上下文训练与推理。NSA最引人注目的特点是在训练阶段应用稀疏性，这在业内实属罕见。通过这种方式，NSA能够在训练和推理过程中显著提升速度，尤其是在解码阶段，速度提升高达11.6倍。
值得关注的是，DeepSeek创始人兼CEO梁文锋亲自参与了该项目的一线研究工作，并在作者排名中位列倒数第二。这意味着他不仅是管理者，也是技术贡献者。此外，该论文的第一作者Jingyang Yuan是在实习期间完成的这项创新研究。
据DeepSeek介绍，NSA的核心在于其三大组件：动态分层稀疏策略、粗粒度token压缩和精粒度token选择。这三个组件协同工作，既提高了效率，又保证了模型对全局长上下文的感知能力和局部精确性。
NSA专为现代硬件优化设计，原生支持模型训练，在加速推理的同时降低预训练成本，并且对模型性能没有明显影响。实验表明，采用NSA机制的模型在通用基准、长上下文任务以及基于指令的推理上，表现与全注意力模型相当，甚至更优。
在配备8块A100 GPU的计算集群上，NSA的前向传播和反向传播速度分别比全注意力快9倍和6倍。由于减少了内存访问量，NSA在长序列解码时，速度提升更加显著。
现有稀疏注意力机制的局限性与NSA的创新
长文本建模是下一代语言模型的关键能力，但传统注意力机制的计算复杂度限制了其在长序列上的应用。在处理长达64k的上下文时，注意力计算可能占据总延迟的70%至80%。因此，稀疏注意力机制应运而生，旨在通过选择性计算关键的查询键对来减少计算开销。
然而，尽管许多稀疏注意力方法在理论上降低了计算复杂度，但它们在实际推理中并未能显著降低延迟。一些方法仅在自回归解码阶段应用稀疏性，而预填充阶段仍需进行密集计算。另一些方法则只关注预填充阶段的稀疏性，导致在某些情况下无法实现全阶段加速。此外，部分稀疏方法无法与现代高效的解码架构兼容，导致KV缓存访问量仍然较高，无法充分发挥稀疏性优势。更重要的是，现有的稀疏注意力方法大多仅应用于推理阶段，缺乏对训练阶段的支持。
DeepSeek推出NSA，旨在填补这些空白，解决以下两大问题：一是事后稀疏化导致的性能退化；二是现有稀疏方法难以满足长序列训练的效率需求。
NSA：软硬件协同优化的解决方案
NSA的核心思想是通过动态分层稀疏策略，结合粗粒度的token压缩和细粒度的token选择，来保持全局上下文感知能力和局部精确性。
NSA的架构包含三个并行的注意力分支：压缩注意力（compressed attention）、选择性注意力（selected attention）和滑动窗口注意力（sliding attention）。压缩注意力通过将键（key）和值（value）聚合成块级表示来捕捉粗粒度的语义信息，减轻计算负担。选择性注意力通过块选择机制保留重要的细粒度信息，给每个块分配重要性分数，并选择排名靠前的块进行注意力计算。滑动窗口注意力则专注于局部上下文信息，防止模型过度依赖局部模式。
为了实现高效的稀疏注意力计算，NSA还针对现代硬件进行了优化。DeepSeek在Triton上实现了硬件对齐的稀疏注意力内核，并专注于共享KV缓存的架构，如GQA和MQA。关键的优化策略包括以组为中心的数据加载、共享KV加载和网格循环调度，从而实现了近乎最优的计算强度平衡。
实验结果：性能超越基线模型，训练推理速度显著提升
为了验证NSA在实际训练和推理场景中的性能，DeepSeek使用了一个结合GQA和MoE的骨干架构作为样本模型，总参数量为27B（其中3B为活跃参数）。
实验结果表明，采用NSA的模型在多个通用基准测试中，总体性能优于所有基线模型，包括全注意力模型，在9项指标中有7项表现最佳。尤其值得一提的是，NSA在推理相关基准测试中取得了显著提升。
在长上下文任务中，NSA在64k上下文的“大海捞针”测试中实现了超高的检索精度。在LongBench上，NSA在多跳QA任务和代码理解任务中表现优于所有基线，显示出在复杂长文本推理任务上的优势。此外，NSA还能与推理模型结合，通过知识蒸馏和监督微调，在32k长度的数学推理任务上获得链式数学推理能力。在AIME 24基准测试中，NSA-R（稀疏注意力变体）在8k和16k上下文设置下均显著优于全注意力-R。
在计算效率方面，在8-GPU A100系统上，随着上下文长度的增加，NSA的加速效果愈发显著。在64k上下文长度时，NSA的前向传播速度提升了9倍，反向传播速度提升了6倍。在解码速度方面，NSA在64k上下文长度时实现了高达11.6倍的速度提升。
总结与展望
DeepSeek推出的NSA为开源AI研究贡献了一项重要成果，有望加速长文本建模和相关应用的发展。尽管NSA已经取得了显著的成果，DeepSeek研究团队也指出了未来的改进方向，例如进一步优化稀疏注意力模式的学习过程，以及探索更高效的硬件实现方式。
目前由于访问人数较多导致DeepSeek服务器超负荷，大家可以从另外2个渠道去使用，不会卡：
渠道一：硅基流动（SiliconFlow）：AI人工智能云服务平台
DeepSeek官网下载： 【点击登录】
