注意力机制真的可以“集中注意力”吗？

原创苏剑林 PaperWeekly

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 月之暗面

研究方向 | NLP、神经网络

之前在《Transformer升级之路：从Performer到线性Attention》、《为什么现在的LLM都是Decoder-only的架构？》等文章中，我们从 Attention 矩阵的“秩”的角度探讨了 Attention 机制，并曾经判断线性 Attention 不如标准 Attention 的关键原因正是“低秩瓶颈”。

然而，这一解释对于双向的 Encoder 模型或许成立，但却难以适用于单向的 Decoder 模型，因为 Decoder 的 Attention 矩阵的上三角部分是被 mask 掉的，留下的下三角矩阵必然是满秩的，而既然都是满秩了，那么低秩瓶颈问题似乎就不复存在了。

所以，“低秩瓶颈”并不能完全解释线性 Attention 的能力缺陷。在这篇文章中，笔者试图寻求另一个角度的解释。简单来说，与标准 Attention 相比，线性 Attention 更难“集中注意力”，从而难以准确地定位到关键 token，这大概是它效果稍逊一筹的主要原因。

稀疏程度

在文章《从熵不变性看Attention的Scale操作》中，我们就从“集中注意力”的角度考察过 Attention 机制，当时我们以信息熵作为“集中程度”的度量，熵越低，表明 Attention 越有可能集中在某个 token 上。

但是，对于一般的 Attention 机制来说，Attention 矩阵可能是非归一化的，比如《FLASH：可能是近来最有意思的高效Transformer设计》介绍的 GAU 模块，以及《相对位置编码Transformer的一个理论缺陷与对策》所引入的归一化 Attention，甚至从更一般的 Non-Local Neural Networks [1] 角度来看，Attention 矩阵还未必是非负的。这些非归一化的乃至非负的 Attention 矩阵自然就不适用于信息熵了，因为信息熵是针对概率分布的。

为此，我们考虑在《如何度量数据的稀疏程度？》介绍的形式的稀疏程度指标：

该指标跟信息熵相似，越小意味着对应的随机向量越稀疏，越稀疏意味着越有可能“一家独大”，这对应于概率中的 one hot 分布，跟信息熵不同的是，它适用于一般的随机变量或者向量。

简化形式

对于注意力机制，我们记，其中，那么

接下来都考虑的极限。假设，那么可以设，其中，于是

注意所服从的分布是一个各向同性的分布，与《n 维空间下两个随机向量的夹角分布》[2] 推导的化简思路一样，由于各向同性的原因，相关的数学期望只与的模长有关，跟它的方向无关，于是我们可以将简化为，那么对的数学期望就可以简化为

其中是一个随机标量。

两个例子

现在可以对常见的一些进行计算对比了。目前最常用的 Attention 机制是，此时求期望只是常规的一维高斯积分，容易算得

当或时，都有，也就是理论上标准Attention确实可以任意稀疏地“集中注意力”，同时这也告诉了我们让注意力更集中的方法：增大的模长，或者增大各个之间的方差，换言之拉开的差距。

另一个例子是笔者喜欢的 GAU（Gated Attention Unit），它在开始提出的时候是（不过笔者后来自己用的时候复原为 Softmax 了，参考《FLASH：可能是近来最有意思的高效Transformer设计》和《听说Attention与Softmax更配》），此时积分没有那么简单，不过也可以直接用 Mathematica 硬算，结果是

其中。式子很恐怖，但是无所谓，画图即可：

▲ relu2注意力的稀疏程度曲线图

可以看到，只有时，原版 GAU 的稀疏度才有机会趋于 0。这也很直观，当偏置项小于 0 时，才有更多的机会让的结果为 0，从而实现稀疏。这个结果也说明了跟的标准注意力不同，的 bias 项可能会对的 GAU 有正面帮助。

极简线性

下面我们再来看一个最简单的例子：不加，或者等价地说。这种情况下对应的就是最简单的一种线性 Attention，同样可以用 Mathematica 硬算得：

下面是几个不同的函数图像：

▲ 极简线性注意力的稀疏程度曲线图

注意，此时的是关于偶函数（读者不妨尝试证明一下），所以时图像跟它相反数的图像是一样的，因此上图只画了的结果。从图中可以看出，不加任何激活函数的线性 Attention 的稀疏程度并不能接近 0，而是存在一个较高的下限，这意味着当输入序列足够长时，这种线性 Attention 并没有办法“集中注意力”到关键位置上。

一般线性

从《线性Attention的探索：Attention必须有个Softmax吗？》我们知道，线性 Attention 的一般形式为，其中是值域非负的激活函数。我们记，那么，并且可以写出

这是关于非负型线性 Attention 的一般结果，现在还没做任何近似，其中分别是序列的均值向量和协方差矩阵。

从这个结果可以看出，非负型线性 Attention 也可能任意稀疏（即），只需要均值趋于0，或者协方差趋于，也就是说序列的信噪比尽可能小。

然而序列是一个非负向量序列，信噪比很小的非负序列意味着序列中大部分元素都是相近的，于是这样的序列能表达的信息有限，也意味着线性 Attention 通常只能表示绝对位置的重要性（比如 Attention 矩阵即某一列都是 1），而无法很好地表达相对位置的重要性，这本质上也是线性 Attention 的低秩瓶颈的体现。

为了更形象地感知的变化规律，我们不妨假设一种最简单的情况：的每一个分量是独立同分布的，这时候均值向量可以简化为，协方差矩阵则可以简化为，那么的公式可以进一步简化为

从这个结果可以看出，要想线性注意力变得稀疏，一个方向是增大，即降低序列的信噪比，另一个方向则是增大，该因子最大值是，其中是的维数，所以增大它意味着要增大，而增大了意味着提高了注意力矩阵的秩的上限，这跟低秩瓶颈视角的结论一样，只有增大才能从根本上缓解线性 Attention 的不足。

特别地，我们在《Transformer升级之路：作为无限维的线性Attention》[3] 也分析过，标准 Attention 也可以理解为一种无限维的线性 Attention，也就是说理论上只有将 d 增加到无穷大，才能彻底弥合两者的差距。

线性衰减

最后，我们来看一下在《Google新作试图“复活”RNN：RNN能否再次辉煌？》介绍过的线性 RNN 模型系列，它们的特点是带有一个显式的递归，这可以看成一个简单的 Attention：

其中。可以算出

当时，只要，总有，所以对于带有显式 Decay 的线性 RNN 模型来说，稀疏性是不成问题的，它的问题是只能表达随着相对位置增大而衰减的、固定不变的注意力，从而无法自适应地关注到距离足够长的 Context。

文章小结

本文提出了通过 Attention 矩阵的稀疏程度来考察不同 Attention 机制潜力的思路，得出二次型 Attention 机制有可能实现任意稀疏的 Attention 矩阵，线性 Attention 则并不容易实现这种稀疏，或者只能实现绝对位置相关的稀疏，这可能是线性 Attention 能力有所限制的原因之一。