DeepSeek原生稀疏注意力（Native Sparse Attention, NSA）算法介绍

李升伟 整理

DeepSeek 提出的原生稀疏注意力（Native Sparse Attention, NSA）算法是一种创新的注意力机制，旨在解决大语言模型（LLM）在处理长序列数据时的计算瓶颈问题。NSA 通过结合算法优化和硬件对齐设计，显著提升了长文本建模的效率，同时保持了模型的性能。以下是 NSA 的核心特点和技术细节：

1. 核心创新

NSA 的核心创新在于其动态分层稀疏策略和硬件友好的设计：

• 动态分层稀疏策略：

NSA 将输入序列按时间维度划分为连续的块，并通过三条并行的注意力分支进行处理：

• 压缩注意力（Compressed Attention）：

通过可学习的 MLP 将每个块压缩成单一表示，用于捕获粗粒度的全局信息。

• 选择性注意力（Selected Attention）：
  保留最重要的细粒度 token 信息，确保局部精度。

• 滑动窗口注意力（Sliding Attention）：

处理近期的局部上下文信息。这种分层设计使得 NSA 能够在保持表达能力的同时，大幅降低计算复杂度。

• 硬件对齐优化：

NSA 通过专门设计的 Triton 内核，最大化 GPU 的 Tensor Core 利用率，并通过优化的内存访问模式减少冗余数据传输。这种设计使得 NSA 在现代 GPU 架构上能够高效运行45。

2. 端到端可训练性

NSA 的一个显著特点是其端到端的可训练性。与现有方法通常在推理阶段引入稀疏化不同，NSA 从预训练阶段就开始使用稀疏注意力机制。这种“原生”设计使得注意力模块能够与模型其他组件协同优化，形成更优的稀疏模式46。

3. 性能表现

NSA 在多项基准测试中表现出色：

• 长文本处理：

在 64k 长度的“大海捞针”测试中，NSA 实现了全位置的完美检索准确率。在 LongBench 评测集上，NSA 的平均得分显著超过了全注意力基线45。

• 推理能力：

在推理相关任务（如 DROP 和 GSM8K）上，NSA 表现出了显著的提升，表明其通过过滤无关注意力路径增强了模型的推理性能45。

• 计算效率：

在 64k 序列长度的场景下，NSA 在解码、前向传播和反向传播三个阶段分别实现了 11.6 倍、9.0 倍和 6.0 倍的加速比14。

4. 应用前景

NSA 的提出为大语言模型的长序列处理提供了新的解决方案，尤其是在需要处理超长上下文的任务（如代码生成、多轮对话和复杂推理）中表现出色。其高效的计算能力和硬件优化设计也为未来更长上下文的处理提供了可行方案136。

5. 开源与社区贡献

NSA 的相关论文和实现已在 arXiv 上公开，并提供了 Triton 内核的实现代码，为开源 AI 社区提供了宝贵的参考46。

综上所述，NSA 通过创新的分层稀疏设计和硬件优化，显著提升了长序列处理的效率和性能，为大语言模型的进一步发展提供了重要支持。

（来自deepseek问答。）