【论文阅读】原生稀疏注意力NSA：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

本文链接：https://blog.csdn.net/qq1198768105/article/details/145738281

上周才刚看完DeepSeek的R1，这周它又出新文章了，他们的产出速度确实快。

1. 简介

这篇文章相较于DeepSeek-R1这篇文章，聚焦范围更加小，R1的重点在于提出了一个文本生成的训练策略和蒸馏策略，这篇文章则是提出了一个注意力机制NSA，主要解决的是长序列做注意力时带来的效率问题。通篇文章看下来，它的实际意义可能比较局限，因此本文仅关注其主要内容，对于具体细节和实验结果并不进一步细究。

论文标题：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
论文链接：https://arxiv.org/abs/2502.11089

2. 动机和相关工作

动机：目前大模型的底层核心架构还是Transformer，Transformer的一个特点是在训练时，需要计算每个Token和其它Token之间的注意力，计算效率为 $O(n^2)$ 。当输入的token很长时，就需要花费很长时间进行计算，非常消耗资源。

目前，已有如下几个思路进行解决长序列注意力的效率问题：

固定稀疏模式(Fixed Sparse Pattern)：大致思路是取一个局部窗口，只在这个窗口中进行注意力的计算。
动态token裁剪(Dynamic Token Pruning)：有点类似于模型剪枝，大致思路是仅仅动态保留一部分最重要的token计算注意力。
序列感知选择(Query-Aware Selection)：大致思路是将几个token打包成块，在块之间计算注意力。

上述方法虽然比直接做全部Token注意力(Full Attention)计算效率要高，但模型的性能同样容易受到较大影响。

NSA就说，我既要加速计算，又要让模型性能不受影响，甚至更高。

下面这张图是NSA的直观效果，速度比全做注意力提升了数倍，性能也略有提升。

在这里插入图片描述

3. 主要方法

NSA主要采用以下几个方法：

粗粒度压缩（Token Compression）
具体做法是：将Token按照位置划分为若干连续块（block），然后将每个块中的内容压缩成单个表示。
细粒度选择（Token Selection）
考虑到直接将Token进行压缩，会让一些细节信息，比如某些重要token，会被块直接给平均表示掉。因此，这里再将一些重要的细粒度Token筛选出来，具体做法是：再次将Token执行分块(可以和上一步分块一样，也可以不一样)，利用注意力分数对其进行重要性评估，最后选取Top-K个最重要的块，将里面的Token取出来和其它的块进行注意力计算。
局部滑动窗口（Local Sliding Window）
分块操作带来的另一个弊病时局部信息丢失。比如，两个原本连续的Token，因划分成不同的block，导致这两个连续信息会丢失。补救方式是运用一个滑动窗口，使其处理局部的上下文信息。
门控网络融合
上述三点实际上视作不同的分支，每个分支各有侧重点。最终，总的注意力为三路分支的加权融合。这里权重通过门控网络来学习得到。由于所有的策略参数都是可微，可学习的，作为整个模型的一部分，因此这个稀疏注意力框架也叫原生(Native)。

此外，在主要方法一节，文章还单独提到了一些硬件底层的创新。下图展示了NSA的内核设计，有点类似于FlashAttention，通过算法在静态随机存取存储器(SRAM)和高带宽存储器(HBM)的优化，进一步提升了计算效率。

在这里插入图片描述

总结

这篇文章主要只是展现了一项新的注意力计算方法。这个方法是一些基础策略的组合，虽然在文章中测试的效果不错，但未必值得借鉴。

其一是它额外引入了很多策略性的超参数，无疑增加了调参的成本；
其二是它引入了额外的模型参数，这可能会影响到模型整体的稳定性，虽然在文中实验性能不降反升，但换个场景未必有效；
其三是它的效率提升很大程度上来自于硬件底层的优化，这种方式基本于GPU设备强绑定，如果设备不同，可能需要额外的工程优化。

因此，我觉得这篇文章的影响力不会像R1这么大，目前受到关注的原因主要还是媒体对DeepSeek的炒作。

参考资料

[1] （3万字长文）深入解析 DeepSeek 最新论文提出的 NSA 注意力机制 https://blog.csdn.net/qq_38334677/article/details/145729713
[2] 刚刚，DeepSeek放出重磅论文！梁文锋亲自参与！
https://mp.weixin.qq.com/s/mlC5CFJ06tSjVJdATuM69Q
[3] 刚刚！DeepSeek梁文锋亲自挂名，公开新注意力架构NSA
https://mp.weixin.qq.com/s/SlhNdGNFpGlicIkMUE21og