2025年7月30日，在奥地利维也纳举行的ACL2025大会上，由DeepSeek与北京大学联合研发、梁文锋作为主要作者的论文《Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention》荣获大会最佳论文奖。
这项突破性研究重新定义了传统注意力机制的实现方式，通过创新的硬件对齐设计和原生可训练稀疏注意力方案，显著提升了计算效率。评审专家一致认为，该技术为长文本处理领域带来了革命性的进步，其性能表现和训练效率均实现了质的飞跃。
这项殊荣不仅是对技术创新的肯定，更凸显了其解决AI领域一个长期痛点——高效长文本处理——的巨大价值。
当我们考虑现代人工智能面临的挑战时，长文本处理无疑是其中最令人头疼的难题之一。设想你正在阅读一本长篇小说，需要同时记住开头的重要线索、中间的转折情节以及眼前的具体细节。对于AI模型来说，这种"全局记忆"能力正是处理长文本的核心需求。然而，传统的注意力机制就像一个需要同时关注所有信息的完美主义者，虽然不会遗漏任何细节，但计算成本却随着文本长度呈平方级增长，就如同一个人试图同时专注于房间里的每一个角落。
当处理包含64,000个词汇的长文档时，传统注意力计算会消耗整个系统70-80%的计算资源，这种巨大的计算负担严重制约了长文本AI应用的发展。正是在这样的背景下，研究团队开发了NSA这一突破性解决方案，它就像一个智慧的阅读者，知道何时需要关注全局概况，何时需要深入细节，何时只需浏览近期内容。
破解长文本处理的核心挑战
更令人印象深刻的是，NSA不仅在算法层面实现了突破，还通过专门的硬件优化确保这种理论上的效率提升能够转化为实际的速度改进。研究团队发现，许多现有的稀疏注意力方法虽然在理论上减少了计算量，但在实际部署中却无法实现相应的速度提升，这主要是因为它们没有充分考虑现代GPU的内存访问模式和计算特性。NSA通过精心设计的内核实现，特别是针对分组查询注意力（GQA）和多查询注意力（MQA）架构的优化，确保了稀疏性能够真正转化为实际的性能提升。
在训练方面，NSA的另一个重要创新是实现了端到端的可训练稀疏性。传统方法通常只在推理阶段应用稀疏性，而在训练时仍然使用完整的注意力机制，这种不一致性会导致性能下降。NSA则从训练开始就采用稀疏注意力模式，让模型能够学会如何在稀疏约束下最优地分配注意力资源，就像训练一个人从小就学会有选择性地关注重要信息。
三重注意力分支的精妙设计
NSA的架构设计体现了对注意力机制本质的深刻理解。在处理每个查询时，系统会将前面的键值对组织成时间块，然后通过三个专门的注意力路径进行处理。
压缩注意力分支的工作原理类似于制作一部电影的预告片。它将连续的文本块聚合成块级表示，每个压缩后的键值对都包含了整个块的精华信息。这种压缩过程通过一个可学习的多层感知器（MLP）实现，该网络会考虑块内的位置编码信息，确保重要的序列关系不会在压缩过程中丢失。通过设置较小的滑动步长，系统能够避免信息碎片化，保持语义的连续性。
选择性注意力分支则更像一个智能的信息筛选器。它首先将键值序列划分为选择块，然后计算每个块的重要性分数。这里的巧妙之处在于，系统利用压缩注意力计算过程中产生的中间注意力分数来评估选择块的重要性，这样就避免了额外的计算开销。对于使用分组查询注意力的模型，系统会确保同一组内的所有查询头选择相同的块，这样可以最小化解码过程中的键值缓存加载量。
滑动窗口注意力分支的设计解决了一个重要的训练问题。在深度学习中，局部模式通常适应更快，可能会主导学习过程，阻止模型有效学习压缩和选择分支的特征。通过将局部上下文处理隔离到专门的分支中，其他分支能够专注于学习各自的特征模式，避免被局部模式干扰。为了进一步防止分支间的梯度干扰，系统为三个分支提供了独立的键值表示，虽然增加了少量计算开销，但显著提高了训练稳定性。
硬件优化的关键突破
NSA在硬件优化方面的创新确保了理论效率能够转化为实际的性能提升。研究团队针对现代GPU的特性，特别是Tensor Core的利用和内存访问模式，设计了专门的内核实现。
传统的FlashAttention策略是将时间上连续的查询块加载到SRAM中，但这种方式在稀疏注意力场景下会导致效率低下，因为块内的查询可能需要访问不连续的键值块。NSA采用了不同的查询分组策略：对于查询序列上的每个位置，系统将分组查询注意力组内的所有查询头加载到SRAM中，因为它们共享相同的稀疏键值块。
这种设计的核心特点包括以组为中心的数据加载模式，在每个内循环中加载组内所有头在位置t的查询以及它们共享的稀疏键值块索引。系统顺序加载由索引指定的连续键值块到SRAM中，以最小化内存加载开销。由于内循环长度（与选择块数量成正比）对于不同查询块保持基本一致，系统将查询和输出循环放在Triton的网格调度器中，简化和优化了内核设计。
这种设计通过消除冗余的键值传输和平衡GPU流式多处理器间的计算工作负载，实现了接近最优的算术强度。在64k长度的序列处理中，NSA在解码、前向传播和后向传播阶段分别实现了11.6倍、9.0倍和6.0倍的加速比。
全面的实验验证
研究团队在一个27B参数的transformer模型上进行了全面的实验验证，该模型结合了分组查询注意力和专家混合架构，包含30层，隐藏维度为2560。模型在270B个8k长度的文本token上进行预训练，随后使用YaRN技术在32k长度的文本上进行长上下文适应训练。
在通用基准测试中，尽管采用了稀疏架构，NSA在多项评估指标上都达到或超过了完整注意力基线的性能。在MMLU、MMLU-PRO、CMMLU等知识性基准上，NSA展现出了与完整注意力相当的能力。特别值得注意的是，在推理相关的基准测试中，NSA表现出了显著的优势，在DROP任务上提升了0.042，在GSM8K上提升了0.034。这种改进可能源于稀疏注意力预训练机制迫使模型专注于最重要的信息，通过过滤无关的注意力路径来增强性能。
在LongBench评估中，NSA获得了0.469的最高平均分，比完整注意力提升了0.032，比Exact-Top方法提升了0.046。特别是在需要复杂长上下文推理的任务上，NSA表现出色，在多跳问答任务HPQ和2Wiki上分别比完整注意力提升了0.087和0.051，在代码理解任务LCC上超出其他方法0.069。
推理能力的重大突破
为了验证NSA与先进下游训练范式的兼容性，研究团队评估了其通过后训练获得链式思维数学推理能力的表现。团队使用来自DeepSeek-R1的知识蒸馏，对10B个32k长度的数学推理轨迹进行监督微调，产生了两个可比较的模型：完整注意力基线模型和NSA稀疏变体模型。
在美国数学邀请赛（AIME 24）基准测试中，研究结果显示NSA在8k上下文设置下比完整注意力基线提升了0.075，在16k上下文下这种优势仍然保持0.054的提升。这些结果验证了原生稀疏注意力的两个关键优势：预训练的稀疏注意力模式能够高效捕获对复杂数学推导至关重要的长距离逻辑依赖关系，以及硬件对齐的架构设计在不发生灾难性遗忘的情况下维持了足够的上下文密度来支持不断增长的推理深度。
效率分析的深入洞察
在计算效率方面，NSA在8-GPU A100系统上的表现证明了其实用价值。基于Triton的NSA实现与Triton版本的FlashAttention-2进行公平的速度比较，结果显示NSA随着上下文长度的增加获得了逐渐更大的加速比，在64k上下文长度下实现了9.0倍的前向加速和6.0倍的后向加速。
在解码速度方面，注意力的解码速度主要由内存访问瓶颈决定，这与键值缓存加载量密切相关。NSA在每个解码步骤中只需要加载压缩token、选择token和邻近token，随着解码长度的增加表现出显著的延迟减少，在64k上下文长度下实现了11.6倍的加速比。这种内存访问效率的优势也随着更长的序列而放大。
研究团队还对不同token选择策略进行了深入的比较分析。基于聚类的策略虽然理论上可行，但面临三个重要挑战：动态聚类机制引入的非平凡计算开销、专家并行组执行时间不平衡导致的算子优化困难，以及强制定期重聚类和块顺序训练协议产生的实现约束。其他块级选择策略则面临选择操作的不可微分性或启发式重要性分数计算的低召回率问题。
注意力模式的可视化发现
通过对预训练27B完整注意力模型的注意力图可视化，研究团队发现了有趣的模式：注意力分数倾向于表现出块状聚类特征，附近的键通常显示出相似的注意力分数。这一观察为NSA的设计提供了重要启发，表明基于空间连续性选择键块可能是一种有效的方法。块状聚类现象表明序列中相邻的token可能与查询token共享某些语义关系，这支撑了在连续token块而非单个token上操作的稀疏注意力机制设计。
NSA代表了稀疏注意力研究的一个重要里程碑。这项工作不仅在算法层面实现了创新，更重要的是通过硬件优化确保了理论效率能够转化为实际的性能提升。对于AI研究者和工程师来说，NSA提供了一个可行的解决方案来处理长文本任务，同时保持了模型的性能和训练的可行性。随着大型语言模型向更长上下文能力发展，NSA这样的硬件感知稀疏注意力架构可能会成为未来AI系统的标准配置。
这项研究的影响不仅限于技术层面，它为整个AI行业指明了一个方向：在追求更强大能力的同时，我们必须同样重视计算效率和实用性。NSA证明了通过精心的算法设计和硬件优化，我们可以在不牺牲性能的前提下显著提高效率，这为构建更加实用和可持续的AI系统开辟了新的可能性。
论文地址：
https://arxiv.org/pdf/2502.11089
END
本文来自至顶AI实验室，一个专注于探索生成式AI前沿技术及其应用的实验室。致力于推动生成式AI在各个领域的创新与突破，挖掘其潜在的应用场景，为企业和个人提供切实可行的解决方案。
Q&A
Q1：NSA稀疏注意力机制相比传统注意力有什么具体优势？
A：NSA主要有三大优势：首先是显著的速度提升，在64k长度序列处理中实现了11.6倍的解码加速、9.0倍的前向加速和6.0倍的后向加速；其次是端到端可训练性，从预训练开始就采用稀疏模式，让模型学会最优的注意力分配；最后是硬件优化设计，确保理论效率能转化为实际性能提升，特别针对现代GPU的内存访问模式进行了优化。
Q2：NSA的三重注意力分支是如何协同工作的？
Q3：NSA能否完全替代传统的完整注意力机制？
A：目前看来，NSA在长文本处理场景下表现出色，在多项基准测试中达到或超过完整注意力的性能，特别是在推理和长上下文任务上有明显优势。不过，NSA主要针对长文本处理进行优化，在短文本场景下可能无法充分发挥其效率优势。未来可能会根据具体应用场景选择合适的注意力机制，NSA很可能成为长文本AI应用的标准选择。
