Native Sparse Attention（NSA）：重新定义长上下文建模的效率与性能

编者注：DeepSeek团队的最新论文：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

论文链接：https://arxiv.org/abs/2502.11089

值得一提的是，幻方科技、DeepSeek 创始人梁文锋也是论文的作者之一。以下是论文解读：

引言

随着大语言模型（LLM）在复杂推理、代码生成和多轮对话等任务中的广泛应用，长上下文建模能力成为新一代模型的核心需求。然而，传统的全注意力机制（Full Attention）在处理长序列时面临高昂的计算成本，尤其是序列长度超过万级时，注意力计算占用了70%以上的延迟。尽管现有稀疏注意力方法试图通过减少计算量提升效率，但它们往往在硬件对齐性、训练支持或性能保持上存在显著缺陷。DeepSeek团队提出的Native Sparse Attention（NSA），通过动态分层稀疏策略和硬件优化设计，为这一难题提供了全新的解决方案。

NSA的核心设计

NSA的核心理念是分层稀疏，通过三路并行注意力机制平衡全局与局部信息的捕捉：

1. 压缩注意力（Compressed Attention）
   将序列划分为连续块，通过可学习的MLP生成压缩的块级表示，大幅减少需处理的键值对数量。例如，块长度l=32l=32时，压缩后的键值数量仅为原始序列的1/32，显著降低计算开销。
2. 选择注意力（Selected Attention）
   基于压缩块的注意力得分，动态筛选出对当前查询最重要的细粒度块（如Top-16块）。这种块级选择策略不仅符合注意力得分的空间连续性规律，还通过硬件友好的连续内存访问优化了GPU计算效率。
3. 滑动窗口注意力（Sliding Window）
   保留局部窗口内的原始键值（如最近512个令牌），防止模型因过度依赖压缩和选择路径而忽略局部模式。三路注意力的输出通过门控机制动态加权融合，确保信息的完整性和灵活性。

创新亮点

1. 硬件对齐的算法设计

NSA针对现代GPU的Tensor Core特性优化内存访问模式：

• 组中心数据加载（Group-Centric Data Loading）：以GQA（Grouped-Query Attention）组为单位加载查询和键值块，减少内存访问碎片化。
• 连续块计算：选择注意力仅加载连续的键值块，避免随机索引导致的低效内存访问，充分利用GPU的并行计算能力。

实验显示，NSA在64k序列长度下实现了11.6倍解码加速、9.0倍前向传播加速和6.0倍反向传播加速，且加速比随序列长度增长持续提升（图1）。

2. 端到端的可训练性

NSA首次实现了原生可训练的稀疏注意力架构：

• 梯度反向传播支持：通过可微的块选择机制（基于压缩注意力得分派生），避免传统方法中离散操作（如聚类、哈希）导致的梯度断裂问题。
• 训练稳定性：独立的注意力分支设计（压缩、选择、滑动窗口）减少了梯度干扰，确保模型在长序列预训练中稳定收敛（图4）。

性能表现

通用任务与长上下文任务

在27B参数的模型上，NSA在多项基准测试中全面超越全注意力基线（表1-2）：

• 通用任务：在MMLU（知识）、GSM8K（数学推理）、HumanEval（代码生成）等任务中，NSA平均得分提升1.3%，尤其在推理任务（DROP、GSM8K）上优势显著。
• 长上下文任务：在64k长度的“大海捞针”测试中，NSA实现100%的检索准确率（图5）。在LongBench多文档QA和代码任务中，NSA以0.469的平均分领先全注意力模型（0.437）和其他稀疏方法（如H2O、InfLLM）。

推理与训练效率

• 推理加速：64k序列的解码速度提升11.6倍，内存访问量仅为全注意力的1/11（表4）。
• 训练成本：通过Triton优化的稀疏核，NSA在预训练中减少30%的计算资源消耗，同时保持更优的收敛曲线（图6）。

讨论与启示

NSA的成功揭示了稀疏注意力设计的两个关键原则：

1. 硬件感知的算法设计：计算效率的提升需紧密结合硬件特性（如内存带宽、Tensor Core利用率），而非单纯追求理论计算量的减少。
2. 训练与推理的协同优化：仅优化推理阶段的稀疏性会导致模型偏离预训练轨迹，而原生可训练的稀疏架构能更充分地挖掘长上下文潜力。

未来，NSA的模块化设计有望扩展至多模态、实时交互等场景，例如处理超长视频序列或实时代码仓库分析。此外，其分层稀疏思想也可启发更高效的MoE（Mixture-of-Experts）模型优化。

结语

NSA通过创新的分层稀疏架构和硬件对齐优化，在长上下文建模中实现了效率与性能的兼得。它不仅为当前LLM的规模化应用提供了实用工具，更重新定义了稀疏注意力机制的设计范式——从“计算量削减”转向“硬件-算法-训练”的全栈协同优化。这一突破或将加速下一代千亿级长上下文模型的落地，推动AI在复杂任务中的边界进一步拓展。

以下为中文翻译版：