原创 小可怜Gemini 2025-02-19 09:34 广东
省流：价格便宜，性能更好
现在的大语言模型是越来越强大，从32K 到 2 Million，个个都能处理超长篇幅的内容，无论是代码库、长文档还是多轮对话，都不在话下。
但是，能力越大，开销也越大传统的 “全注意力机制” 在处理长文本时，计算量会 平方级 增长，速度越来越慢，成本也蹭蹭上涨，这可咋整？成本上涨后，用户成本也得上涨...
别慌！ 这篇论文就给我们带来了一个非常给力的解决方案 —— 原生稀疏注意力 (Native Sparse Attention, NSA)。
简单来说，NSA 就像给模型的 “眼睛” 配上了 “智能眼镜”，让它只关注重点信息，忽略无关紧要的内容，从而大大提升效率！
上图 Figure 1 就清晰地展示了 NSA 的强大之处：
• 左图： 在各种评测基准上（通用 benchmark、长文本任务、推理评测），使用了 NSA 的模型，性能 不降反升，甚至超越了传统的全注意力模型！
• 右图： 在处理 64k 长度的超长文本时，NSA 在解码、前向传播和反向传播各个阶段，都比全注意力模型实现了 显著的加速！
这简直是鱼和熊掌兼得啊！既要又要还要党的福音！
🤔 稀疏注意力是啥？为啥它能提速？ 🤔
可能有些小伙伴对 “稀疏注意力” 这个概念还不太熟悉。没关系，我来给大家简单科普一下。
传统的“全注意力”机制，就像 “地毯式搜索”，每个词都要和所有词计算关联度，计算量非常大。
而“稀疏注意力”，就像 “重点巡逻”，模型会学习到只关注 重要的词，忽略掉那些不重要的词，这样就大大减少了计算量，速度自然就提升上去了。
这篇论文提出的 NSA 更牛的地方在于，它不仅仅是 “稀疏” 而已，还做到了 “原生可训练” 和 “硬件对齐”！
🌟 NSA 的两大创新：既高效又好用！ 🌟
这篇论文的核心贡献，可以总结为两大创新点：
1. 硬件对齐的系统优化: 为了让稀疏注意力真正在硬件上跑得飞快，NSA 在算法设计上就充分考虑了现代 GPU 的特性，例如 Tensor Core 的高效利用 和 内存访问优化。作者们设计了专门的 kernel (你可以理解为更底层的优化代码)，让 NSA 在各种硬件平台上都能发挥出极致的性能。
2. 训练感知的算法设计: 很多稀疏注意力方法都只关注推理加速，但 NSA 更进一步，它从 训练阶段 就开始引入稀疏性。这意味着模型在训练的时候，就学会了如何更高效地利用稀疏注意力模式，而不是像以前的方法那样，先用全注意力模型训练好，再在推理时 “打补丁” 用稀疏注意力。 这种 “原生可训练” 的设计，不仅降低了训练成本，还保证了模型性能！
为了实现这两个创新点，NSA 采用了 “动态分层稀疏策略” (dynamic hierarchical sparse strategy)。 如图 Figure 2 所示，NSA 把注意力机制分成了 三条 “通路” (three attention paths)：
• 压缩注意力 (Compressed Attention): 把长文本分成一个个 “时间块” (temporal blocks)，然后把每个块 压缩 成一个更简洁的 “粗粒度” (coarse-grained) 表示。 这样模型可以快速浏览全局信息，就像看文章的 “大纲”。
• 选择性注意力 (Selective Attention): 模型会 选择性地保留 一些 重要的 “token 块” (token blocks)，关注文本中的关键信息，就像阅读时文章的重点段落。
• 滑动窗口注意力 (Sliding Window Attention): 对于 局部上下文信息，NSA 还保留了 滑动窗口 (sliding windows) 这种经典的方式，确保模型不会忽略 临近词之间的关系，就像阅读时关注上下文。
这三种 “通路” 各司其职，优势互补，共同构成了 NSA 高效又强大的注意力机制！
🏆 实验结果：全面超越现有方法！ 🏆
为了验证 NSA 的效果，作者们做了大量的实验，结果那是相当的亮眼！
1. 通用 Benchmark 评测：性能不输 Full Attention，推理任务更胜一筹！
如上图所示，在 MMLU, BBH, GSM8K, DROP 等多个通用benchmark 上， NSA 模型都取得了优异的成绩，平均性能甚至超过了 Full Attention 模型！ 尤其是在 DROP 和 GSM8K 这些需要更强推理能力的 benchmark 上，NSA 的优势更加明显。
2. 长文本评测：LongBench 全面领先，Needle-in-a-Haystack 完美检索！
在 Needle-in-a-Haystack 测试中，NSA 在 64k 超长上下文 下，完美找回 “大海捞针” 般的目标信息，检索准确率 高达 100%！
Table 2 展示了 LongBench 长文本 benchmark 上的评测结果。 NSA 模型在各个子任务上都表现出色，平均分也超过了 Full Attention 和其他 Sparse Attention 方法， 尤其在 HPQ 和 2Wiki 多跳问答任务上，优势巨大！
3. Chain-of-Thought 推理评测：NSA-R更擅长复杂数学推理！
1. NSA-R 显著优于 Full Attention-R：
在 8192 的 token 限制下，NSA-R 的得分为 0.121，而 Full Attention-R 的得分仅为 0.046。 NSA-R 的得分是 Full Attention-R 的近 2.6 倍。
在 16384 的 token 限制下，NSA-R 的得分为 0.146，而 Full Attention-R 的得分为 0.092。 NSA-R 的得分仍然明显高于 Full Attention-R，约为其 1.6 倍。
2. 增加生成 Token 限制对两个模型都有正面影响：
Full Attention-R: 从 8192 到 16384 的 token 限制，Full Attention-R 的得分从 0.046 提升到 0.092，几乎翻倍。
NSA-R: 从 8192 到 16384 的 token 限制，NSA-R 的得分也从 0.121 提升到 0.146，虽然提升幅度不如 Full Attention-R 那么显著，但仍然是正向提升。
可以看到，经过 Chain-of-Thought 指令微调 后，NSA-R 模型的准确率显著高于 Full Attention-R 模型，尤其是在 16k 长文本输入下，优势依然明显！
4. 效率评测：训练、推理速度大幅提升！，且内存访问大大减少
Figure 6 展示了在 A100 GPU 上，NSA 的 Triton kernel 与 FlashAttention-2 的 Triton kernel 速度对比。 可以看到， 随着文本长度的增加，NSA 的加速效果越来越明显，64k 长度时，前向传播速度提升 9 倍，反向传播速度提升 6 倍！
上图则展示了解码阶段的内存访问量对比。 NSA 的内存访问量远低于 Full Attention， 64k 长度时，内存访问量减少了 11.6 倍！ 这也解释了 NSA 解码速度大幅提升的原因。
这些实验结果都充分证明了 NSA 在性能和效率上的优越性，是下一代长文本大语言模型的理想选择！
这篇论文提出的 原生稀疏注意力 (NSA) 机制，是一项非常出色的创新工作。它不仅在算法设计上巧妙地融合了 分层稀疏策略，更重要的是，它 从硬件和训练的角度 出发，真正解决了稀疏注意力在实际应用中面临的效率和性能问题。
不愧是成本控制大师 DeepSeek 的作品，我愿意称他为古希腊掌握账本增效的神...
