DeepSeek革命性NSA注意力机制问世！梁文锋上阵，长文本推理能力飙升

来源：新智元
本文约6000字，建议阅读5分钟
DeepSeek官方再发布新论文！

在这篇论文中，来自DeepSeek、北大和华盛顿大学的研究人员提出了一种全新的注意力机制NSA。

具体来说，它是一种面向硬件且支持原生训练的稀疏注意力机制（Sparse Attention），专为超快长上下文训练与推理设计。其核心组成包括——

• 动态分层稀疏策略

• 粗粒度的token压缩

• 细粒度的token选择

通过针对现代硬件的优化设计，NSA（原生可训练稀疏注意力机制）在保证性能的同时显著提升了推理速度，并有效降低了预训练成本。

在通用基准测试、长文本处理以及基于指令的推理任务中，它的表现均能达到甚至超越传统全注意力（Full Attention）模型的水平。

值得一提的是，DeepSeek创始人梁文锋这次不仅亲自上阵，并且还亲自提交了论文。

论文地址：https://arxiv.org/abs/2502.11089

目前，这成为了该论文的最大亮点之一。

网友们也不忘艾特整天放小道消息的奥特曼：看，DeepSeek又有新论文了。

部署高效稀疏注意力机制的关键挑战

随着序列长度的增加，注意力计算在整体计算成本中变得越来越占主导地位，这对长上下文处理提出了重大挑战。

尤其在深入推理、代码库级别的代码生成以及多轮自主智能体系统等多种现实应用的推动下，这一趋势尤为明显。

近期OpenAI的o系列模型、DeepSeek-R1和Gemini 1.5 Pro，在数千个token的范围内保持连贯的多轮对话，同时执行复杂的长距离依赖推理了。

然而，随着序列长度的增加，传统注意力机制的高计算复杂度，成为关键的时延瓶颈。

理论估算表明，在处理64k长度上下文时，使用softmax结构的注意力计算占总时延迟的70–80%，更凸显了对更高效注意力机制的迫切需求。

部署高效的稀疏注意力机制必须应对两个关键挑战：硬件对齐的推理加速和训练感知的算法设计。

为此，研究团队提出了NSA，一种原生可训练的稀疏注意力架构，并集成了分层token建模。同时，还设计了专用计算内核，以最大化实际计算效率。

如图2所示，NSA通过将键值对组织为时间块（temporal blocks），并通过三条注意力路径进行计算，从而降低每个查询的计算量。

NSA引入了两个核心创新点：

1. 硬件对齐系统：优化块级稀疏注意力，使其充分利用Tensor Core并优化内存访问，从而实现平衡的算术强度。
2. 训练感知设计：通过高效算法和反向传播运算符实现稳定的端到端训练，使NSA能够同时支持高效推理与完整训练流程。

上图左：该框架通过三个并行的注意力分支处理输入序列。对于给定的查询，前面的键和值被处理成压缩注意力以用于粗粒度模式、选择注意力以用于重要token块，以及滑动注意力以用于局部上下文。

上图右：可视化每个分支产生的不同注意力模式。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。

方法概述

具体而言，在注意力机制中，每个查询token q_t与所有前面的键k_:t计算相关性得分，以生成值v_:t的加权和。形式上，对于长度为t的输入序列，注意力操作定义为：

其中Attn表示注意力函数：

这里，𝛼_𝑡,𝑖表示q_𝑡和k_𝑖之间的注意力权重，𝑑_𝑘是键的特征维度。随着序列长度的增加，注意力计算在整体计算成本中的占比越来越大，给长序列上下文处理带来了显著的挑战。

算术强度是计算操作与内存访问的比率，它本质上决定了硬件上的算法优化。

对于计算任务，当算术强度超过这个临界阈值时，任务变为计算受限（受限于GPU的FLOPS），而低于该阈值时则变为内存受限（受限于内存带宽）。

具体到因果自注意力机制，在训练和预填充阶段，批量矩阵乘法和注意力计算表现出高算术强度，使得这些阶段在现代加速器上成为计算受限阶段。

相比之下，自回归解码阶段则为内存带宽受限，因为每次前向传递只生成一个token，同时需要加载整个键值缓存，导致算术强度较低。

就这导致了不同的优化目标——在训练和预填充阶段减少计算成本，而在解码阶段减少内存访问。

NSA的整体框架

为了利用注意力机制在自然稀疏模式下的潜力，研究人员提出用更紧凑、信息更密集的表示键值对 ，来替换方程1中的原始键值对k_:𝑡，v_:𝑡，并根据每个查询q_𝑡进行优化。

具体而言，将优化后的注意力输出正式定义如下：

其中， 是基于当前查询q_𝑡和上下文记忆k_:𝑡，v_:𝑡动态构建的。

可以设计不同的映射策略，来生成不同类别的 ，并将它们组合如下：

如图2所示，NSA采用了三种映射策略 C = {cmp,slc,win}，分别表示对键值对的压缩（compression）、选择（selection）和滑动窗口（sliding window）。

是与策略c对应的门控得分（gate score），通过MLP和sigmoid激活函数从输入特征中学习得到。

令𝑁_𝑡表示重映射后的键值对总数：

通过确保𝑁_𝑡 ≪ 𝑡，就能维持高稀疏率，从而优化计算效率。

接下来，就是重新映射策略的设计，包括token压缩、token选取和滑动窗口的具体算法。

关键算法组件

核心算法是重映射策略fK和fV的设计：token压缩、token选择和滑动窗口。

1. token压缩：将连续的键或值块聚合为块级表示，得到压缩后的键和值。压缩表示能够捕获较粗粒度的高级语义信息，并减少注意力机制的计算负担。
2. token选择：用于识别并保留最相关的token，同时减少计算开销。其中，按块选择在空间连续的块中处理键和值序列：首先将键值序列划分为选择块；为每个块分配重要性分数，识别出最重要的块以用于注意力计算。而Top-𝑛块选择在获得选择块重要性得分后，按块重要性得分排序，选择top-𝑛稀疏区块中的token，被选中的键和值会参与注意力计算。
3. 滑动窗口：专门的滑动窗口分支，明确处理局部上下文，使得其他分支（压缩和选择）可以专注于学习各自的特征，而不会被局部模式所干扰。

这三种策略都是为了提高计算和内存效率。token压缩减少了数据量，token选择减少了计算量，滑动窗口限制了需要关注的序列长度。

硬件优化的内核设计

新研究使用Triton实现了与硬件对齐的稀疏注意力内核。

遵循当前最先进的LLM，专注于具有共享KV缓存的架构，为稀疏选择注意力引入了专门的内核设计。

为了解决内存访问效率低下问题，关键优化在于不同的查询分组策略：对于查询序列上的每个位置，将GQA组内的所有查询头（它们共享相同的稀疏KV块）加载到SRAM中。

NSA的内核设计采用了以下策略：

1. 查询加载（Grid Loop）：内核按GQA（Grid Query Attention）组加载查询。每个GQA组包含共享相同稀疏键值块的查询头。这种方法减少了冗余的键值传输，提高了内存访问效率。
2. 键值获取（Inner Loop）：在每个查询加载后，内核获取相应的稀疏键值块。这些键值块存储在高带宽内存（HBM）中。内核将这些键值块传输到片上静态随机存取存储器（SRAM）中，以进行注意力计算。
3. 注意力计算：内核在SRAM中执行注意力计算。绿色块表示存储在SRAM中的数据，蓝色块表示存储在HBM中的数据。这种内存层次结构优化了数据传输和计算效率。

所提出的内核架构关键特征， 如下：

1. 以组为中心的数据加载：对于每个内部循环，在位置𝑡加载该组中所有头部的查询𝑄 ∈ R[ℎ, 𝑑_𝑘]，以及它们共享的稀疏键值块索引I。
2. 共享KV获取：在内部循环中，按顺序加载由I𝑡索引的连续键值块到SRAM中，作为𝐾 ∈ R[𝐵_𝑘, 𝑑_𝑘]和𝑉 ∈ R[𝐵_𝑘, 𝑑_𝑣]，以最小化内存加载，其中𝐵𝑘是满足𝐵_𝑘 | 𝑙′的内核块大小。
3. 网格外部循环：由于内部循环的长度（与选择的块数𝑛成比例）对于不同的查询块几乎相同，我们将查询/输出循环放入Triton的网格调度器中，以简化和优化内核。

这种设计通过以下方式实现了接近最优的算术强度：（1）通过组共享消除冗余的K传输；（2）在GPU流式多处理器之间平衡计算负载。

革命性NSA，碾压全注意力

接下来，研究人员从三个方面评估了NSA，将其与全注意力基线，以及SOTA稀疏注意力方法进行比较：

• 通用基准性能
• 长文本基准性能
• 思维链推理性能

如下图4所示，NSA和全注意力基线的预训练损失曲线，呈现出稳定和平滑的下降，NSA始终优于全注意力模型。

通用基准评估

研究人员在一套全面的基准测试中评估了预训练的NSA和全注意力基线，覆盖了知识、推理和编码能力。

结果如表1所示，尽管NSA具有稀疏性，但它在总体性能上表现优越，在9个指标中有7个超过了包括全注意力在内的所有基线。

这表明，尽管NSA在较短序列上可能未能充分利用其效率优势，但它仍然表现出强大的性能。

值得注意的是，NSA在推理相关的基准测试中显示出显著的提升（DROP: +0.042，GSM8K: +0.034）。

这进一步证明了，预训练有助于模型发展专门的注意力机制。

NSA稀疏注意力预训练机制迫使模型集中关注最重要的信息，可能通过过滤掉来自无关注意力路径的噪声来增强性能。

不同评估中的一致表现，也验证了NSA作为通用架构的稳健性。

长文本评估

如图5所示，NSA在64k上下文中的「大海捞针」测试中，在所有位置上都实现了完美的检索准确率。

这一性能得益于层次化的稀疏注意力设计，该设计结合了压缩token以实现高效的全局上下文扫描，以及用于精确局部信息检索的选择token。

粗粒度的压缩在低计算成本下识别相关的上下文块，而对选择token的token级注意力则确保了关键细粒度信息的保留。

这种设计，使NSA能够同时保持全局感知能力和局部精确性。

此外，研究人员还在LongBench基准上，对NSA进行了评估。

如表2所示，NSA达下了最高的平均分0.469，优于所有基准（比全注意力高0.032，比Exact-Top高0.046）。

思维链推理评估

考虑到RL在小规模模型上的效果有限，研究人员采用了从DeepSeek-R1进行知识蒸馏的方法，使用100亿个32k长度的数学推理轨迹进行监督微调（SFT）。

由此，产生了两个模型：Full Attention-R（全注意力基准）和NSA-R（稀疏变体）。

在AIME 24基准测试上，研究人员评估了这两个模型。

结果如下表3所示，在8k上下文设置下，NSA-R比Full Attention-R获得了更高的准确率（高出0.075），这种优势在16k上下文中依然保持（高出0.054）。

这些结果验证了原生稀疏注意力的两个关键优势：

1. 预训练的稀疏注意力模式，能够有效捕获复杂数学推导中至关重要的长程逻辑依赖关系；
2. 架构采用硬件对齐设计，在增加推理深度的同时保持足够的上下文密度，避免了灾难性遗忘。

在不同上下文长度下的持续优异表现，证实了当稀疏注意力原生集成到训练流程中时，它在高级推理任务中的可行性。

64k上下文，前向传播9倍速飙升

研究团队在配备8个A100 GPU的系统上评估了NSA相对于全注意力机制（Full Attention）的计算效率。

训练速度

为了评估训练速度，研究团队将基于Triton的NSA实现与基于Triton的FlashAttention-2进行了比较，以确保在同一后端进行公平的速度比较。

实验结果表明，随着上下文长度的增加，NSA实现了逐渐提高的加速，在64k上下文长度下实现了高达9.0倍的前向加速和6.0倍的反向加速。值得注意的是，随着序列长度的增加，速度优势变得更加明显。

这种加速源于NSA的硬件对齐算法设计，旨在最大限度地提高稀疏注意架构的效率：分块内存访问模式通过合并加载最大限度地利用了Tensor Core。内核中精细的循环调度消除了冗余的KV传输。

下图6比较了Triton实现的NSA和全注意力与Triton实现的FlashAttention-2，以确保在同一后端进行公平的速度比较。

解码速度

注意力机制的解码速度主要受限于内存访问瓶颈，这与KV缓存的加载量密切相关。

随着解码长度的增加，本文中的方法显著降低了延迟，在64k上下文长度下实现了高达11.6倍的加速。

如下表4所示NSA在长序列解码时具有显著的效率优势。通过减少内存访问量，NSA能够显著提高解码速度，尤其是在处理较长的上下文时。

这种内存访问效率的优势也会随着序列的增长而放大。

讨论

研究团队反思了NSA的开发过程，并讨论了在探索不同稀疏注意力策略过程中获得的关键见解。

虽然本文中的方法展示了有希望的结果，但了解替代策略遇到的挑战并分析注意力分布模式，为未来的研究方向提供了有价值的背景。

研究团队首先考察了促使他们做出设计选择的替代token选择策略所面临的挑战，然后通过可视化来深入了解注意力分布模式。

替代token选择策略