Native Sparse Attention 论文简介

用更少计算资源，处理更长文本：解读Native Sparse Attention（NSA）的革新设计

近年来，随着大语言模型（LLM）在长文本理解、代码生成和多轮对话等场景中的广泛应用，如何高效处理长上下文序列成为技术发展的核心挑战。传统注意力机制（如Transformer中的全注意力）在序列长度增加时面临计算量爆炸式增长的问题，导致推理和训练成本急剧上升。

近期，一篇题为《Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention》的论文提出了一种名为NSA（原生稀疏注意力）的创新方案，不仅在计算效率上实现了显著提升，还首次实现了端到端的可训练稀疏注意力设计。本文将从问题背景、技术原理和实验结果三个方面解读这一工作。

---

一、长上下文建模的困境与突破口

传统注意力机制的瓶颈

全注意力机制的计算复杂度为(O(n^2))，当处理64k长度的序列时，注意力计算可能占据总延迟的70%-80%。尽管已有多种稀疏注意力方法（如KV缓存淘汰、块状选择或哈希聚类）试图减少计算量，但它们大多面临两大问题：

1. 推理加速效果有限：许多方法仅在预填充（Prefilling）或解码（Decoding）阶段实现加速，无法覆盖全流程；
2. 难以支持端到端训练：离散化操作（如聚类或哈希）导致梯度无法回传，模型无法动态学习稀疏模式。

NSA的核心目标

NSA的提出旨在解决上述问题，其设计目标包括：

• 硬件对齐的高效实现：通过块状内存访问和计算优化，充分利用GPU的Tensor Core性能；
• 原生可训练的稀疏架构：支持从预训练到推理全流程的稀疏注意力模式学习。

---

二、NSA的技术设计：分而治之的注意力策略

NSA的核心思想是通过动态分层稀疏策略，将注意力计算分解为三个并行分支，分别处理不同粒度的信息：

1. 压缩分支（Compression）

• 粗粒度聚合：将连续的令牌块（如每32个令牌）压缩为一个块级表示，通过MLP和位置编码保留语义信息。
• 作用：快速捕捉全局上下文，减少需要处理的总令牌数。

2. 选择分支（Selection）

• 细粒度保留：基于压缩分支的中间注意力分数，动态选择最重要的令牌块（如Top-16块）。
• 硬件优化：块状连续访问适配GPU内存带宽，避免随机读取的性能损失。

3. 滑动窗口分支（Sliding Window）

• 局部上下文维护：保留最近的512个令牌，确保模型不丢失局部细节信息。
• 防止“捷径学习”：独立计算局部注意力，避免全局分支被局部模式主导。

动态门控融合

三个分支的输出通过可学习的门控权重（基于MLP和Sigmoid）动态融合，平衡全局与局部信息的贡献。

---

三、实验结果：效率与性能的双重突破

论文在27B参数的模型上进行了全面实验，结果显示NSA在以下方面表现突出：

1. 性能表现

• 通用任务：在MMLU、GSM8K等9项基准测试中，NSA在7项超越全注意力模型；
• 长上下文任务：在LongBench上平均得分提升3.2%，且在64k长度的“大海捞针”测试中实现100%准确率；
• 推理能力：通过知识蒸馏后，NSA在数学竞赛AIME中的准确率显著高于全注意力模型。

2. 计算效率

• 训练加速：64k序列长度下，前向传播加速9倍，反向传播加速6倍；
• 解码加速：内存访问量减少至全注意力的1/11.6，解码速度提升11.6倍。

3. 硬件利用率

通过Triton实现的定制化内核，NSA充分利用GPU的SRAM和Tensor Core，在块状数据加载和共享KV缓存的设计下，算术强度（计算与内存访问比）接近硬件极限。

---

四、意义与启示

NSA的突破不仅在于效率提升，更在于其**“原生可训练”**的设计理念。传统稀疏方法通常在预训练后引入稀疏化，导致模型偏离优化轨迹；而NSA从预训练阶段开始学习稀疏模式，使模型能够自适应地筛选关键信息。这一思路为未来长上下文模型的开发提供了新方向：

• 更经济的模型训练：减少长文本预训练的计算成本；
• 更灵活的架构扩展：适配MoE、GQA等先进模型结构；
• 更广泛的应用场景：支持代码生成、多轮对话等需长序列处理的任务。

---

结语

NSA通过算法与硬件的深度协同，证明了稀疏注意力机制在大模型时代的潜力。其设计理念——将效率优化与模型能力统一于端到端的训练框架中——或将成为下一代语言模型的核心范式。

论文链接：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention