编译 | 陈骏达

编辑 | Panken

智东西2月18日报道，今天下午，DeepSeek团队发布一篇新论文，介绍了一种改进的稀疏注意力机制NSA，可用于超快速的长上下文训练与推理。NSA以性价比极高的方式，罕见地在训练阶段应用稀疏性，在训推场景中均实现速度的明显提升，特别是在解码阶段实现了高达11.6倍的提升。

让人眼前一亮的是，DeepSeek创始人兼CEO梁文锋这次出现在了合著名单之中，在作者排名中位列倒数第二。这意味着他作为项目管理者，参与了一线的研究工作。另外，这篇论文的第一作者Jingyang Yuan是在实习期间完成的这项研究。

据DeepSeek介绍，NSA具有三大核心组件：动态分层稀疏策略、粗粒度token压缩、精粒度token选择。通过三大组件的配合，既提升了效率，也保留了模型对全局长上下文的感知能力和局部精确性。

这一机制专门针现代硬件进行优化设计，原生支持模型训练，在加速推理的同时降低预训练成本，对性能也无明显影响。采用NSA机制的模型在通用基准、长上下文任务和基于指令的推理上，与全注意力模型相当或表现更优。

在8卡A100计算集群上，NSA的前向传播和反向传播速度分别比全注意力快9倍和6倍，由于减少了内存访问量，NSA在长序列解码时相较于全注意力模型速度显著提升。

▲NSA在部分测试上的表现（图源：DeepSeek）

论文链接：arxiv.org/abs/2502.11089

一、现有稀疏注意力机制存在明显缺陷，DeepSeek希望填补空白

长文本建模是下一代语言模型的关键能力，但传统注意力机制的高复杂度限制了其在长序列上的应用。

例如，在解码64k长度的上下文时，注意力计算占据了总延迟的70%至80%。因此，稀疏注意力机制应运而生，通过选择性计算关键的查询键对来减少计算开销。

然而，尽管许多稀疏注意力方法在理论上减少了计算复杂度，但这些方法在实际推理中未能显著降低延迟。

一些方法仅在自回归解码阶段应用稀疏性，而预填充阶段仍需进行密集计算（如H2O）；另一些方法仅关注预填充阶段的稀疏性（如MInference），导致在某些工作负载下无法实现全阶段加速。

还有部分稀疏方法无法适应现代高效的解码架构（如MQA和GQA），导致KV缓存访问量仍然较高，无法充分发挥稀疏性优势。

此外，现有的稀疏注意力方法大多仅在推理阶段应用稀疏性，缺乏对训练阶段的支持。

NSA旨在通过针对硬件特点的推理加速和适用于训练的算法设计，填补这一空白。DeepSeek推出NSA主要希望解决两大问题：

一是事后稀疏化导致性能退化，如预训练模型的检索头易被剪枝；

二是现有稀疏方法难以应对长序列训练的效率需求。现有方法存在非可训练组件和低效反向传播等问题，阻碍了高效训练和长上下文模型的发展。

二、软硬件协同深度优化，无限逼近计算强度最优解

NSA的核心思想是通过动态分层稀疏策略，结合粗粒度的token压缩和细粒度的token选择，以保留全局上下文感知能力和局部精确性。

下方是NSA架构的概览，左侧，NSA将输入序列通过三个并行的注意力分支处理：压缩注意力（compressed attention）、选择性注意力（selected attention）和滑动窗口注意力（sliding attention）。右侧是对每个分支产生的不同注意力模式的可视化。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。

▲NSA架构概览（图源：DeepSeek）

其中，压缩注意力通过将键（key）和值（value）聚合成块（block）级表示来捕捉粗粒度的语义信息。这种压缩表示能够捕捉更粗粒度的高层语义信息，并减轻注意力计算的负担。

不过，仅使用压缩后的键（key）和值（value）可能会丢失重要的细粒度信息，DeepSeek引入了选择性注意力，通过块选择机制保留重要的细粒度信息。

他们给每个块分配了重要性分数，根据块的重要性分数选择排名前n的块，并将这些块中的标记用于注意力计算。这种方法在保留关键信息的同时，显著降低了计算负担。

在注意力机制中，局部模式通常会快速适应并主导学习过程，可能会阻碍模型从压缩和选择token中有效学习。滑动窗口注意力可以回应这一问题，通过专注于局部上下文信息，防止模型过度依赖局部模式。

为了实现高效的稀疏注意力计算，NSA还针对现代硬件进行了优化。

具体来看，DeepSeek在Triton上实现了硬件对齐的稀疏注意力内核。鉴于多头自注意力（MHA）内存密集且解码效率低，他们专注于共享KV缓存的架构，如分组查询注意力（GQA）和多查询注意力（MQA），这些架构与当前最先进的LLMs一致。

DeepSeek的关键优化策略是采用不同的查询分组策略，通过以下关键特性实现了近乎最优的计算强度平衡：

1、以组为中心的数据加载：在每个内循环中，加载组内所有头的查询及其共享的稀疏KV块索引。

2、共享KV加载：在内循环中，连续加载KV块以最小化内存加载。

3、网格循环调度：由于内循环长度在不同查询块中几乎相同，将查询/输出循环放在Triton的网格调度器中，以简化和优化内核。

三、超越多款基线模型，训练提速6-9倍、推理最高提速11.6倍

为测试NSA机制在实际训练、推理场景中的表现，DeepSeek用当前最先进的LLM的常见实践，使用了一个结合分组查询注意力（GQA）和混合专家（MoE）的骨干架构作为样本模型。这一模型总参数量为27B，其中3B为活跃参数。

在这一模型的基础上，DeepSeek使用了NSA、全注意力以及其它注意力机制，并进行了评估。

在多个通用基准测试中，采用NSA的模型尽管具有稀疏性，但其总体性能优于所有基线模型，包括全注意力模型，在9项指标中有7项表现最佳。

▲采用不同注意力机制的模型在多个通用基准测试上的表现（图源：DeepSeek）

NSA在较短序列上可能未能充分发挥其效率优势，但其性能依然强劲。值得注意的是，NSA在推理相关基准测试中取得了显著提升，这表明NSA的预训练机制有助于模型开发专门的注意力机制，能迫使模型专注于最重要的信息，通过过滤掉无关的注意力路径中的噪声，潜在地提升了性能。

在长上下文任务中，NSA在64k上下文的“大海捞针（neddle in a haystack）”测试中实现了超强的的检索精度。这得益于其分层稀疏注意力设计，通过粗粒度的压缩token实现高效的全局上下文扫描，并通过细粒度的选择标记保留关键信息，从而在全局感知和局部精确性之间取得平衡。

▲大海捞针测试结果（图源：DeepSeek）

在LongBench上，NSA在多跳QA任务和代码理解任务中表现优于所有基线，还显示出在复杂长文本推理任务上的优势。这些结果表明，NSA的原生稀疏注意力机制不仅提升了模型性能，还为长文本任务提供了更优的解决方案。

NSA机制还能与推理模型进行结合，适配前沿的后训练方式。DeepSeek使用了从DeepSeek-R1蒸馏获得的知识和监督微调（SFT）的方式，使采用NSA的模型在32k长度的数学推理任务上获得链式数学推理能力。

实验中，NSA-R（稀疏注意力变体）和全注意力-R（基线模型）在具有挑战性的AIME 24基准测试上进行了对比。结果显示，NSA-R在8k和16k上下文设置下均显著优于全注意力-R（分别高出0.075和0.054），验证了其在复杂推理任务中的优势。

▲NSA-R（稀疏注意力变体）和全注意力-R（基线模型）在AIME上的表现（图源：DeepSeek）

DeepSeek还在8-GPU A100系统上，对NSA的计算效率与全注意力机制进行了对比。

在训练速度方面，随着上下文长度的增加，NSA的加速效果愈发显著。在64k上下文长度时，NSA的前向传播速度提升了9倍，反向传播速度提升了6倍。

这种加速主要得益于NSA的硬件对齐设计：块状的内存访问模式通过合并加载最大化了Tensor Core的利用率，内核中精细的循环调度消除了冗余的KV传输。

在解码速度方面，注意力机制的解码速度主要受限于KV缓存加载的内存瓶颈。随着解码长度的增加，NSA的延迟显著降低，在64k上下文长度时实现了高达11.6倍的速度提升。这种内存访问效率的优势随着序列长度的增加而更加明显。

结语：DeepSeek持续给开源AI惊喜

尽管NSA取得了显著的成果，但DeepSeek研究团队也指出了一些可能的改进方向。例如，进一步优化稀疏注意力模式的学习过程，以及探索更高效的硬件实现方式。

正如DeepSeek之前发布的所有技术报告那样，这篇详解NSA机制的论文内容详实，对NSA机制中涉及的技术细节阐释清晰，可操作性强，是DeepSeek给开源AI研究贡献的最新成果。