Transformer升级之路：从Performer到线性Attention

原创苏剑林 PaperWeekly

©PaperWeekly 原创 · 作者｜苏剑林

单位｜追一科技

研究方向｜NLP、神经网络

看过笔者之前的文章线性Attention的探索：Attention 必须有个 Softmax 吗？和 Performer：用随机投影将 Attention 的复杂度线性化的读者，可能会觉得本文的标题有点不自然，因为是先有线性 Attention 然后才有 Performer 的，它们的关系为“Performer 是线性 Attention 的一种实现，在保证线性复杂度的同时保持了对标准 Attention 的近似”，所以正常来说是“从线性 Attention 到 Performer”才对。

然而，本文并不是打算梳理线性 Attention 的发展史，而是打算反过来思考 Performer 给线性 Attention 所带来的启示，所以是“从 Performer 到线性 Attention”。

激活函数

线性 Attention 的常见形式是：

其中是值域非负的激活函数。那么如何选取这个激活函数呢？Performer 告诉我们，应该选择指数函数：

首先，我们来看它跟已有的结果有什么不一样。在《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》[1] 给出的选择是：

我们知道正是在 x=0 处的一阶泰勒展开，因此这个选择其实已经相当接近了。

此外，这个方案还跟《Efficient Attention: Attention with Linear Complexities》[2] 一文中引入的双重 softmax 来构建线性 Attention 的设计很相似，在那种设计中有

，相比直接只不过归一化的位置有所不同。

简单推导

为什么说 Performer 告诉我们激活函数的最佳选择是呢？我们来看 Performer 找到的将标准 Attention 线性化的映射：

简单来说，Performer 找到了一个映射，使得 d 维向量被映射为了 m 维向量，并且满足近似关系，此时：

最后一个等式表明，往里边乘以一个常数（哪怕这个常数跟有关），Performer 的结果完全不改变，这意味着将映射改为：

Performer 的结果不会有任何变化。当然，这里这一项还不能去掉，但是如果我们假设不会波动太大，它并不是 Attention 的主要因素，那么这一项也相当于一个常数，于是最终的映射（近似地）等价为：

这个看上去已经简化很多的映射该怎么理解呢？其实 m 个随机向量拼成了一个的矩阵，它将 d 维的映射为了 m 维的向量，然后加上激活函数得到了。我们知道 Attention 的都有一个全连接层变换，如果我们将这个的映射矩阵整合到全连接层中，那么剩下的就是一个激活函数了！

所以这就是最优激活函数的来源了，只要我们将的输出维度从 d 维改为 m 维，然后配合的激活函数，那么理论上它就有 Performer 的拟合能力，甚至更强，因为 Performer 的矩阵是一个固定的随机矩阵，而这里我们相当于把该矩阵也设为可训练了，还去掉了低秩约束，空间是比 Performer 更大的。

低秩问题

不管是本文的主角 Performer，还是之前在 Nyströmformer：基于矩阵分解的线性化 Attention 方案介绍的 Nyströmformer，它们的思路都是“寻找一个能逼近标准 Attention 的线性 Attention”。那么一个很自然的问题就是：标准 Attention 有什么好的？哪里值得大家向它对齐？

从信息损失的角度来看，标准 Attention 矩阵的“秩”可能更大，即更接近可逆矩阵，这意味着它能保留更多有效信息。具体来说，Attention 矩阵是一个的矩阵，它由通过而来，要注意的是，这里的 d 是 Attention 的key_size，比如对于 BERT base 来说它只是 64，而 n 往往比较大，这说明的秩不超过d，而且，即离满秩还远得很。

不过，softmax 的关键运算是，一个矩阵如果每个元素取指数的话，那么新矩阵的秩是可能增加的！所以标准 Attention 矩阵有升秩的可能性，意味着它蕴含了更有效处理信息的能力。

相比之下，线性 Attention 矩阵是的形式，所以线性 Attention 矩阵的秩一定不超过 m，而为了弥补秩的损失，所以一般要设置 m > d，在 Performer 的实验中选择的是 m = 4d，也就是 key_size 扩大为 4 倍，秩的重要性可见一斑。当然，扩大了 key_size，一个直接的后果是处理短序列的时候，线性 Attention 还比标准 Attention 要慢，这是线性 Attention 的固有瓶颈。

关于 Attention 矩阵的秩的理论分析，也有一些论文可以参考，比如《Low-Rank Bottleneck in Multi-head Attention Models》[3] 就指出哪怕在标准 Attention 中，低秩性也是一个严重的瓶颈，增大 key_size 可以提升性能；

上个月的《Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth》[4] 则指出，如果没有残差和 FFN，那么标准 Attention 有极大的风险退化为秩等于 1 的简单变换。连标准 Attention 这个有“升秩潜力”的模型都有低秩问题，更不用说线性 Attention 这种本身秩就有上限的模型了。

所以，一句话就是：用线性 Attention 需要用更大的 key_size 来维持矩阵的秩。

集中注意

我们还可以从稀疏性角度来理解标准 Attention 的好处。直观来想，既然是“注意力机制”，那么肯定需要“集中注意力”，如果太分散，那么可能就相当于平均池化了，而“集中注意力”，意味着每个 token 应该只能显著地关联到若干个 token，用数学的话说，那就是意味着 Attention 矩阵是稀疏的，或者说至少要具备变得稀疏的可能性。

对于标准 Attention 来说，它通过 softmax 来归一化：

其中指数函数起到了一个放大的作用，只要各个本身能拉开一定差距，那么会进一步放大这种差距，结果就是归一化之后除了最大值的那几个位置之外，剩下的概率都很接近于 0 了，这说明标准 Attention 是有潜力“集中注意力”的。

而对于线性 Attention 来说，它是直接内积的结果，没有得到的进一步放大，所以它的注意力是比较稠密的，在序列长度较大的时候，它往往就很接近平均池化了。要缓解这一点，还是需要增大 key_size，来放大差距，直观来说，就是 n 向量放到一个低维空间太“挤”了，换到更高维的空间就“松”一些了。

怎么样验证稀疏的重要性呢？笔者曾经尝试过，将线性 Attention 的 Attention 矩阵先算出来，然后强行截断 Attention 矩阵（也就是每个token只跟前后几个 token 做 attention，变成局部形式的 Attention）让它变得稀疏，结果发现这种截断后的线性 Attention 效果明显好于全矩阵的线性 Attention。

这就肯定了稀疏的重要性了，当然，这样把 Attention 矩阵先算出来然后前行截断的方式，使得线性 Attention 的复杂度不再是线性的了，因此不具备实用价值，仅用于理论验证。

还有一个实验现象可以辅助证明稀疏的重要性，那就是线性 Attention 做语言模型或者解码器的时候，效果是跟标准 Attention 差不了多少的，这时候线性 Attention 变成了单向的 RNN（参考《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》[1] ），等价于 Attention 矩阵变成了下三角阵，也是更稀疏了。相比之下，如果用不稀疏的双向的线性 Attention 直接做 MLM 模型，则掉点会相当明显。

更重要的是，稀疏性和前一节提到的秩是有密切关联的，甚至可以说它们是“一体两面”：适当的稀疏化方法能提高矩阵的秩！比如做语言模型的下三角 Attention 矩阵，只要对角线元素非零（往往都能达到），那么这时候的矩阵直接就是满秩可逆阵了！还有笔者实验的局部 Attention 截断，也能增加矩阵的秩，比如极端情况下，每个 token 只跟自身做 attention，那么 Attention 矩阵就是满秩的单位阵了！

文章小结

本文从 Performer 出发思考了线性 Attention 的一些问题，包括关于线性 Attention 的激活函数选择，以及线性 Attention 的瓶颈所在（低秩性、稀疏性），总的结论是，线性 Attention 的最佳激活函数应当是指数函数，而有效的 Attention 机制应当具备更高的秩和更大的稀疏性。