2月18日，DeepSeek公布了一项由梁文锋亲自参与的研究论文成果——原生稀疏注意力（Native Sparse Attention, NSA）机制。这是DeepSeek团队在稀疏注意力领域的创新性工作，结合了算法创新和硬件优化，旨在解决长上下文建模中的计算瓶颈。
近年来，我们见证了长文本建模在 AI 领域的重要性日益凸显。无论是深度推理、代码库生成、还是多轮对话，都离不开模型对长序列信息的有效处理能力。像 OpenAI 的 o-series 模型、DeepSeek-R1、以及 Google Gemini 1.5 Pro 等，都展现了处理超长文本的强大潜力。
传统Transformer模型中的全注意力（Full Attention）机制面临着高昂的计算成本和内存需求。这不仅限制了模型在处理长文本时的效率，也影响了其在实际应用中的可行性。因此，如何在保证模型性能的同时，降低注意力机制的计算成本和内存需求，成为当前NLP领域亟待解决的关键问题。
DeepSeek论文主要解决长文本处理中注意力机制的高计算成本和内存需求问题，提出了一种结合硬件优化和训练感知设计的稀疏注意力方法。该方法在保证高性能的同时，显著提高处理速度，并支持高效部署和端到端训练。论文显示，NSA不仅将大语言模型处理64k长文本的速度最高提升11.6倍，更在通用基准测试中实现性能反超传统全注意力模型。（文末附论文原文下载链接！）
Native Sparse Attention（NSA）
NSA 采用动态分层稀疏策略，结合粗粒度令牌压缩和细粒度令牌选择，保留全局上下文感知和局部精度。该方法有两个关键创新：
一是硬件适配系统：通过优化块级稀疏注意力，实现对 Tensor Core 的高效利用和内存访问的优化，确保算术强度的平衡，大幅提升速度；
二是训练感知设计：借助高效的算法和反向传播操作符，实现稳定的端到端训练，在不牺牲模型性能的前提下降低预训练计算量。
方法
从具体实现来看，NSA 在整体框架上对传统注意力机制进行了优化。它将原始的键值对替换为更紧凑、信息更密集的表示。NSA 通过三种映射策略来构建这些新的键值对，分别是压缩（compression）、选择（selection）和滑动窗口（sliding window），设计硬件对齐的稀疏注意力内核，提高计算效率：
-
令牌压缩将连续的键值块聚合为块级表示，减少计算负担；
-
令牌选择基于块级重要性分数，选择关键令牌块，保留细粒度信息；
-
滑动窗口单独处理局部上下文，防止局部模式干扰其他分支学习。
为了进一步提升效率，NSA 在硬件优化的内核设计上也下足功夫。针对当前主流的具有共享 KV 缓存的架构（如 GQA 和 MQA），NSA 实现了硬件对齐的稀疏注意力内核。在处理稀疏选择注意力时，引入了独特的查询分组策略，以解决内存访问效率低下的问题。这种设计通过消除冗余的 KV 传输和平衡计算负载，实现了接近最优的算术强度。
测试与性能
在实验环节，研究人员对 NSA 进行了全面评估。在预训练设置上，采用了结合分组查询注意力（GQA）和专家混合（MoE）的骨干模型，参数规模达 270 亿，其中有 30 亿个活跃参数。模型包含 30 层，隐藏层维度为 2560。针对 GQA 和 MoE，分别进行了细致的参数配置。同时，为确保训练稳定性，对首层的 MoE 进行了特殊处理。实验对比了 NSA 与全注意力基线模型以及其他先进的稀疏注意力方法。
在通用基准测试中，在 9 项指标中的 7 项上超越了包括全注意力模型在内的所有基线模型，展现出强大的性能；在长上下文评估中，NSA 在 64k 上下文检索任务中准确率完美，在 LongBench 基准测试中平均得分最高；在思维链推理评估中，NSA-R 在 8k 和 16k 上下文设置下，准确率显著高于全注意力的对比模型。
训练速度上，NSA 随着上下文长度增加，加速比逐渐增大，在 64k 上下文长度时，前向和反向传播速度分别提升 9.0 倍和 6.0 倍；解码速度上，NSA 内存访问量大幅减少，在 64k 上下文长度时，解码速度提升 11.6 倍。
综合来看，NSA 作为一种创新的稀疏注意力机制，在长上下文建模中展现出卓越的性能和效率。它不仅在模型性能上与全注意力模型相当甚至超越，还在计算效率上实现了显著提升，为下一代大语言模型的发展提供了重要的技术支持，具有广阔的应用前景。
SDNLAB后台回复
0220 NSA
获取论文原文！
