DeepSeek梁文锋合著论文获ACL最佳论文奖，提出全新NSA稀疏注意力

在刚刚举办的计算语言学和自然语言处理领域的顶级国际会议 ACL 上，由 DeepSeek 创始人梁文锋亲自署名的论文从 8000 多篇论文中脱颖而出，成为本届大会的最佳论文之一。

（来源：资料图）

该论文的其他合作者分别来自北京大学和美国华盛顿大学，担任第一作者的是 Jingyang Yuan。

图 | Jingyang Yuan（来源：https://ieeexplore.ieee.org/autho）

在这篇论文中，他们提出了一种名为 NSA 的本机可训练的稀疏注意力机制，该机制将算法创新与硬件优化相结合，以实现高效的长上下文建模。

NSA 采用了一种动态分层稀疏策略，将粗粒度 token 压缩与细粒度 token 选择相结合，以同时保持全局上下文感知和局部精确性。

这一方法通过以下两项关键创新推动了稀疏注意力设计的发展：

一方面，他们通过设计算术强度平衡的算法，并结合针对现代硬件的实现优化，实现了显著的速度提升。

另一方面，他们实现了端到端的训练，在不牺牲模型性能的前提下减少了预训练的计算量。

如下图所示，实验表明，在通用基准测试、长上下文任务和基于指令的推理任务中，使用 NSA 预训练的模型表现与全注意力模型相当或更优。

（来源：https://arxiv.org/pdf/2502.11089）

同时，在解码、前向传播和后向传播方面，NSA 在 64k 长度序列上相较于全注意力机制实现了显著加速，验证了其在整个模型生命周期中的效率。

图 | 相关论文（来源：https://arxiv.org/pdf/2502.11089）

长上下文建模是下一代大模型的关键能力

研究界日益认识到，长上下文建模是下一代大模型的关键能力，其推动因素是现实世界中的各种应用，包括深度推理、库级代码生成和多轮自主代理系统。然而，标准注意力机制的高计算成本带来了巨大的计算挑战。稀疏注意力机制为在保持模型性能的同时提高效率提供了一个有前景的方向。

最近领域内的一些突破，包括 OpenAI 的 o 系列模型、DeepSeek-R1 和 Gemini 1.5 Pro，使模型能够处理整个代码库、长文档，在数千个 token 上保持连贯的多轮对话，并在较长依赖关系中进行复杂推理。

然而，随着序列长度的增加，基础注意力机制的高度复杂性成为关键的延迟瓶颈。

理论估计表明，在解码 64k 长度上下文时，采用 softmax 架构进行注意力计算占总延迟的 70% 至 80%，这凸显出人们迫切需要更高效的注意力机制。

实现高效长上下文建模的一种自然方法是利用 softmax 注意力机制的固有稀疏性，即选择性地计算关键的查询-键对，这可以在保持性能的同时显著降低计算开销。近期的研究进展通过多种策略展示了这一潜力：键值（KV）缓存淘汰方法、分块 KV 缓存选择方法，以及基于采样、聚类或哈希的选择方法。

尽管这些策略颇具前景，但现有的稀疏注意力方法在实际部署中往往表现不佳。许多方法未能实现与其理论增益相当的速度提升；此外，大多数方法缺乏有效的训练时间支持，无法充分利用注意力机制的稀疏模式。

为了克服这些局限性，部署有效的稀疏注意力机制必须应对两个关键挑战：

（1）解决硬件适配的推理加速挑战：将理论计算减少转化为实际速度提升，需要在预填充和解码阶段都进行硬件友好的算法设计，以便缓解内存访问和硬件调度瓶颈；

（2）解决训练感知算法设计的挑战：通过可训练算子实现端到端的计算，以便降低训练成本，同时还得保持模型性能。这些要求对于实际应用实现快速长上下文推理或训练至关重要。当同时考虑这两个方面时，现有方法仍然存在显著的差距。

因此，为了实现更有效且高效的稀疏注意力机制，他们提出了 NSA 这种原生可训练的稀疏注意力架构，并集成了分层 token 建模。如下图所示，NSA 通过将键和值组织成时间块，并通过三条注意力路径来处理它们，从而减少每次查询的计算量：压缩的粗粒度 token、选择性保留的细粒度 token 以及用于获取局部上下文信息的滑动窗口。

（来源：https://arxiv.org/pdf/2502.11089）

然后可以实现专门的核函数，以最大限度地提高其实用效率。针对上述关键需求，NSA 推出了两项核心创新：

首先是硬件对齐系统：针对 Tensor Core 的利用率和内存访问，优化块级稀疏注意力机制，确保算术强度平衡。

其次是训练感知设计：通过高效算法和反向操作符实现稳定的端到端训练。这一优化使 NSA 能够同时支持高效部署和端到端训练。

研究中，该团队通过在真实世界语言语料库上的综合实验来评估 NSA。在拥有 2600 亿个 token 的 270 亿参数的 Transformer 主干网络上进行预训练后，他们评估了 NSA 在通用语言评估、长上下文评估和思维链推理评估中的表现。并进一步比较了在英伟达 A100 GPU 上采用优化 Triton 实现的核速度。实验结果表明，NSA 的性能与全注意力基线相当或更优，同时优于现有的稀疏注意力方法。此外，与全注意力机制相比，NSA 在解码、前向和后向阶段均实现了显著加速，且随着序列长度的增加，加速比也随之提高。这些结果验证了他们的分层稀疏注意力设计有效地平衡了模型能力和计算效率。

三个角度评估 NSA

实验中，他们从以下三个角度来评估 NSA：通用基准性能、长上下文基准性能以及思维链推理性能，并与全注意力基线以及当前最先进的稀疏注意力方法进行比较。

遵循最先进的大模型的常见做法，他们的实验采用了一种结合分组查询注意力（GQA，Grouped-Query Attention）和混合专家（MoE，Mixture-of-Experts）的骨干网络，总参数为 270 亿，其中活跃参数为 30 亿。该模型由 30 层组成，隐藏层维度为 2560。对于 GQA，他们将组数设置为 4，总共有 64 个注意力头。对于 MoE，他们采用了 DeepSeek MoE 结构，该结构包含 72 个路由专家和 2 个共享专家，并将 top-k 专家数设置为 6。为确保训练稳定性，第一层中的 MoE 被替换为 SwiGLU 形式的 MLP。

该团队在论文中指出，其所提出的架构在计算成本和模型性能之间实现了有效的权衡。其将全注意力模型和稀疏注意力模型均在 2700 亿个 8k 长度文本的 token 上进行预训练，随后使用 YaRN 在 32k 长度文本上进行持续训练和监督微调，以实现长上下文适应。两种模型都经过充分训练以实现完全收敛，以确保公平比较。如下图所示，NSA 和全注意力（Full Attention）基线的预训练损失曲线呈现出稳定且平滑的下降趋势，且 NSA 模型的表现始终优于全注意力模型。

（来源：https://arxiv.org/pdf/2502.11089）

除了与全注意力（Full Attention）方法进行比较外，他们还评估了几种最先进的推理阶段稀疏注意力方法：H2O、infLLM、Quest 和 Exact-Top，这些方法首先计算全注意力得分，并选择与每个查询相对应的前几个得分最高的键，然后计算这些位置上的注意力。与此同时，这些方法涵盖了多种稀疏注意力范式。

在一般评估中，当大多数样本的长度位于稀疏注意力基线的局部上下文窗口内时，这些方法实际上等同于全注意力方法。因此，在这种设置之下，该团队仅展示了 NSA 与全注意力基线之间的比较结果。在长文本评估中，他们对所有基线方法进行了比较，并将所有稀疏注意力方法的稀疏度设置为相同，以便确保比较的公平性。对于需要进行长文本监督微调的思维链推理评估，他们仅将比较范围限定在全注意力模型上，因为稀疏注意力基线模型不支持训练。

在一系列涵盖知识、推理和编码能力的综合基准测试上，该团队对预训练的 NSA 和全注意力基线进行了评估，这些基准测试包括 MMLU、MMLU-PRO、CMMLU、BBH、GSM8K、MATH、DROP、MBPP 和 HumanEval。

图 | 对比结果（来源：https://arxiv.org/pdf/2502.11089）

尽管 NSA 较为简洁，但其整体表现却十分出色，在 9 项指标中有 7 项优于包括全注意力机制在内的所有基线模型。这表明，尽管在较短序列上，NSA 可能无法充分发挥其效率优势，但其表现仍然强劲。值得注意的是，NSA 在推理相关基准测试中表现出了显著的提升（DROP：+0.042，GSM8K：+0.034），这表明该团队所使用的预训练有助于模型发展出专门的注意力机制。这种稀疏注意力预训练机制迫使模型聚焦于最重要的信息，通过过滤掉无关注意力路径中的噪声，从而能够潜在地提升性能。在不同评估中的一致表现也验证了 NSA 作为一种通用架构的稳健性。

同时，他们还进行了长上下文评估。如下图所示，在 64k 上下文的“大海捞针”测试中，NSA 在所有位置均实现了完美的检索准确率。

（来源：https://arxiv.org/pdf/2502.11089）

这种表现源于该团队的分层稀疏注意力设计方案，该设计方案结合了压缩 token 以便实现高效的全局上下文扫描，以及能够通过选择 token 实现精确的局部信息检索。粗粒度压缩则能以较低的计算成本识别出相关的上下文块，而对选定 token 的 token 级注意力能够确保关键细粒度信息的保留。这种设计使得 NSA 既能保持全局感知，又能确保局部精确性。

该团队还在 LongBench 上对 NSA 与最先进的稀疏注意力方法和全注意力基线进行了评估。为确保稀疏性一致，他们将所有稀疏注意力基线中每个查询激活的 token 设置为 2560 个，这对应于在处理 32k 序列长度时，NSA 中激活的 token 的平均数量。继 StreamLLM 之后，该 token 预算包括前 128 个 token 和 512 个本地 token。他们排除了 LongBench 中的某些子集，因为它们在所有模型中的得分都很低，可能无法提供有意义的比较。如下表所示，NSA 取得了最高的平均分 0.469，优于所有基线方法（比全注意力机制高出 0.032，比精确顶部方法高出 0.046）。

（来源：https://arxiv.org/pdf/2502.11089）

这一优异表现源于以下两项关键创新：第一，他们独创的稀疏注意力设计，能够在预训练期间对稀疏模式进行端到端的优化，促进稀疏注意力模块与其他模型组件之间的同步适应；第二，分层稀疏注意力机制实现了局部和全局信息处理之间的平衡。

值得注意的是，NSA 在需要长上下文复杂推理的任务上表现出色，在多跳问答任务（HPQ 和 2Wiki）上相较于全注意力机制分别提升了+0.087 和+0.051，在代码理解任务上超越了基线模型（LCC：+0.069），并且在段落检索任务上优于其他方法（PassR-en：+0.075）。这些结果验证了 NSA 在处理多样化长上下文挑战的能力，其原生预训练的稀疏注意力在学习任务最优模式方面提供了额外优势。

在思维链推理评估中，为了评估 NSA 与先进下游训练范式的兼容性，该团队研究了其通过后训练获取思维链数学推理能力的能力。鉴于强化学习在较小规模模型上的效果有限，他们采用了 DeepSeek-R1 的知识蒸馏方法，利用 100 亿个 32k 长度数学推理轨迹的 token 进行监督微调（SFT，supervised fine-tuning）。这产生了两个可比较的模型：Full Attention-R（全注意力基线模型）和 NSA-R（该团队的稀疏变体模型）。他们在 AIME 24 基准测试中评估了这两个模型。为了验证推理深度的影响，他们在两种上下文限制（8k 和 16k 个 token）下进行了实验，以便衡量扩展推理链是否能提高准确性。

如下表所示，在 8k 上下文设置下，NSA-R 的准确率显著高于全注意力-R（+0.075），且在 16k 上下文设置下，这一优势仍然存在（+0.054）。

（来源：https://arxiv.org/pdf/2502.11089）

这些结果验证了原生稀疏注意力机制的两大关键优势：（1）预训练的稀疏注意力模式能够高效捕捉对复杂数学推导至关重要的长距离逻辑依赖关系；（2）本次架构采用与硬件相匹配的设计，保持了足够的上下文密度，以支持不断增加的推理深度，同时避免灾难性遗忘。

这种在不同上下文长度上的一致性表现证明，当稀疏注意力机制被原生整合到训练流程中时，它能够胜任高级推理任务。

（来源：https://openreview.net/profile?id=~Jingyang_Yuan1）

另据悉，担任本次论文第一作者兼通讯作者的 Jingyang Yuan 于 2022 年获得北京大学计算机科学学士学位，预计将于 2027 年从北京大学博士毕业，他的研究方向包括图表示学习、神经物理模拟、大模型等。

参考资料：

https://ieeexplore.ieee.org/author/37090050318

https://scholar.google.com/citations?user=mDwlqfkAAAAJ&hl=en

https://arxiv.org/pdf/2502.11089

运营/排版：何晨龙