梁文锋亲自挂名DeepSeek发布新论文

由 DeepSeek 联合创始人梁文锋亲自挂名的研究团队，在 arXiv 上发表了一篇题为“Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention”的论文，提出了一种全新的注意力机制架构 NSA（Native Sparse Attention，原生稀疏注意力）。

背景

长序列建模对于下一代语言模型至关重要，但传统的注意力机制由于计算复杂度高，面临着显著的计算挑战。稀疏注意力机制被认为是提高效率的有前景的方向。NSA通过动态层次化的稀疏策略，结合粗粒度的令牌压缩和细粒度的令牌选择，既保留了全局上下文感知能力，又保持了局部精度。
熟悉 DeepSeek-R1 的用户都知道，这款模型虽然在许多方面表现出色，但比较遗憾的一点在于，其输入上下文能力方面相对不足。而这次发布的 NSA 架构就主要致力于解决当前大模型长文本处理中的关键瓶颈问题。

NSA的核心创新包括两个方面

硬件对齐的算法优化

通过平衡计算强度的算法设计，实现了显著的加速，并针对现代硬件进行了实现优化。为了最大化利用现代硬件资源，研究团队基于Triton开发了专门的kernel。以下是主要优化点：

• Group-Centric Data Loading：将同一GQA（Grouped-Query Attention）组内的所有query head同时加载到SRAM中处理，最大化Tensor Core的利用率。
• 共享KV Fetching：通过连续加载key/value块到SRAM中，减少冗余数据传输。
• Outer Loop on Grid：利用Triton的grid调度器简化和优化内核，确保高效的内存访问模式。

这些优化设计与现代GPU架构高度契合，实现了接近理论极限的加速效果。

分层稀疏注意力设计

端到端训练能通过高效的算法和反向传播操作符，实现了可训练的稀疏注意力机制，减少了预训练计算成本，同时保持了模型性能。
NSA通过独特的分层稀疏注意力设计，大幅降低了计算复杂度，同时保持了模型的表达能力。具体来说，该架构将输入序列按时间维度划分为连续的block，并通过三条并行的注意力分支进行处理：

• 压缩注意力（Compressed Attention）：通过可学习的MLP将每个block压缩成单一表示，用于捕获粗粒度的全局信息。
• 选择性注意力（Selected Attention）：保留最重要的fine-grained token信息，确保模型能够聚焦于关键细节。
• 滑动窗口注意力（Sliding Attention）：专注于处理近期的局部上下文，捕捉短距离依赖关系。
  这种分层设计使得模型能够在保持强大表达能力的同时，显著降低计算复杂度。

算法设计关键部分

（1）令牌压缩（Token Compression）

块级聚合：将连续令牌块压缩为单一表示，例如每32个令牌压缩为一个块，降低计算量。
数学形式：
其中 l为块长，d为滑动步长，MLP含块内位置编码。

（2）令牌选择（Token Selection）

• 块级重要性评分：利用压缩注意力分数生成块级重要性分布，选择Top-n关键块。
• 硬件友好设计：连续块访问适配GPU内存带宽，支持GQA组内共享索引，减少KV缓存加载。

（3）滑动窗口（Sliding Window）

• 局部上下文隔离：显式处理最近512个令牌，防止局部模式主导全局学习。
• 门控融合：通过可学习门控动态加权三路注意力输出。

结果

研究团队对 NSA 进行了全方位的性能验证。实验采用了一个基于 GQA 和 MoE（Mixture-of-Experts）的 27B 参数 backbone，包含 30 层网络结构，隐藏维度为 2560。为确保实验的可比性，研究团队采用了与全量注意力模型完全相同的训练流程，包括在 270B token 的 8K 长度文本上进行预训练，随后使用 YaRN 方法在 32K 长度文本上进行延续训练和监督微调。

在通用能力评测中，NSA 展现出了超出预期的表现。在涵盖知识、推理和编程能力的九项基准测试中，包括 MMLU、MMLU-PRO、CMMLU、BBH、GSM8K、MATH、DROP、MBPP 和 HumanEval，NSA 在七项上超越了全量注意力基线。
尤其值得关注的是在推理相关任务上的显著提升，如在 DROP 任务上提升了 4.2 个百分点，在 GSM8K 上提升了 3.4 个百分点。这一结果表明，稀疏注意力的预训练不仅没有损害模型能力，反而通过过滤无关注意力路径增强了模型的推理性能。

在长文本处理能力的专项测试上，NSA 的优势更为突出。具体来看，在多跳问答任务 HPQ 和 2Wiki 上分别提升了 8.7 和 5.1 个百分点，在代码理解任务 LCC 上提升了 6.9 个百分点，在段落检索任务 PassR-en 上提升了 7.5 个百分点。

结论与展望

NSA通过动态分层稀疏策略和硬件对齐优化，为长上下文建模提供了高效且可扩展的解决方案。其核心贡献包括：

• 端到端可训练性：支持从预训练到微调的完整流程，避免性能损失；

• 硬件级加速：适配Tensor Core和内存带宽，实现理论计算与实测加速一致；

• 任务自适应能力：在通用任务、长上下文推理和数学推导中均表现优异。

未来方向包括：探索更灵活的稀疏模式、优化MoE架构下的负载均衡，以及扩展至多模态长序列处理。NSA的代码与模型已开源，为LLM的长上下文应用提供了新的基础设施。