编译 | 陈骏达

编辑 | 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 研究贡献的最新成果。