Performer——将注意力复杂度从序列长度平方降至线性的近似方法
引言
今天给大家介绍的外文精品博客是来自Google Brain研究员关于Performer模型 (opens new window)的讲解。该博客并未涉及太多的数学推导，而是从整体设计思想来讲解Performer，包括广义注意力的定义以及如何节省显存与计算时间。在介绍完Performer的设计思路之后，作者还用实验证明了该方法的在长序列任务上的有效性，最后还介绍了Performer能够应用的场景以及研究结论。
作者介绍
Krzysztof Choromanski Google Brain首席研究员，同时也是Performer的一作。
Lucy Colwell Google Brain研究员。
译者说
当今Transformer模型可以说是在nlp领域处于统治者地位的模型，其超强的性能使得在其出现以后的大多数研究都以Transformer为骨架进行设计。但由于Transformer的核心结构是注意力模块，而原始计算注意力的复杂度是序列长度的平方（自注意力中Q×K的结果是一个L×L的注意力矩阵，L为序列长度），因此想要将Transformer模型应用到长序列任务上不是一件容易的事。同时，本文后续提到的复杂度的概念均是针对序列长度而言的。
在博客的前半部分，作者讲述了前人在使用注意力建模长序列这一问题上的一些尝试，其中主流方法就是稀疏化注意力矩阵，同时也有研究人员表示能够将稀疏后的矩阵看作一张无向图，即把图上的节点看作token，把边看做attention，两个节点有边则表明这两个token的注意力操作没有被稀疏。但作者也提到这些方法都有一些不足：
稀疏算子的加速往往只有部分设备支持，这就导致该加速手段并不能轻松地应用到所有设备上。
前人的稀疏方法除了Big Bird对稀疏后的矩阵的表示能力有一个理论性的探讨与证明，其余工作都没有保证稀疏后的注意力矩阵的表示能力不受很大影响。
这些工作大多只考虑了以Transformer为基础的模型和生成式预训练模型。
稀疏注意力矩阵往往因为不具备与原始注意力矩阵相当的表示能力，所以需要更多的层数来弥补这一缺陷，从而无法直接使用预训练好的模型参数。这就付出了更多的训练代价，包括重新预训练时间与能量消耗（电费）。
存在着一些稀疏注意力矩阵无法代替原始注意力矩阵的领域，如指针网络（Pointer Network (opens new window)）。
因此Performer的作者提出一种由线性复杂度计算近似原始注意力矩阵的方法，这种方法能够将注意力计算降低到线性复杂度的同时依然保留原始注意力矩阵的各种特性，从而在很大程度上避免了上述不足。具体来说，作者首先使用特殊的指数变换和高斯投射变换得到Q和K矩阵的近似Q'和K'（先把Q*K从softmax中提取出来），然后重新安排了QKV矩阵乘法的顺序，先用K'乘V，再用Q'乘上步计算结果（即K'乘V的结果），调整后的计算复杂度降低为线性复杂度（可以参考本文FAVOR+方法的第一张图）。实验方面，作者通过几组实验证明该方法是具备线性时间与空间复杂度的。同时还证明使用了Performer的模型速度与不使用attention机制的Transformer相当。
原博客精华内容概括
1.广义注意力矩阵
首先，为了分析注意力矩阵是如何具备表示能力的，作者对注意力矩阵进行了拆分。但由于原始注意力矩阵是在Q矩阵与K矩阵相乘后使用了一个softmax非线性变换，因此无法直接将原注意力矩阵直接拆分为Q矩阵和K矩阵。但值得注意的是，我们依然可以将原注意力矩阵拆分为经过非线性变换的Q矩阵与K矩阵的乘积。如下图所示，等式的左侧是代表注意力矩阵，等式的右侧则是Q‘矩阵与K’矩阵。其中Q'矩阵与K'矩阵是原始Q矩阵与K矩阵的非线性变换。
如果将注意力矩阵看作是Q矩阵和K矩阵的非线性变换矩阵的乘积，那么原始的注意力矩阵就可以看作是将Q矩阵与K矩阵进行特殊的指数变换与高斯投射变换。由于Attention矩阵的计算是先计算Q矩阵与K矩阵的乘积，再进行非线性变换，而作者希望先对Q矩阵与K矩阵做非线性变换再做矩阵乘，因此这里Q矩阵与K矩阵的变换是特殊的指数变换（对标原始注意力计算中Softmax操作包含的指数变换）与高斯投射变换（对标Q矩阵与K矩阵在原始注意力计算中包含的线性变换）。而这种变换是非常稠密的，因为他并没有缩小矩阵的大小，这就导致变换后的矩阵乘法依然具有序列长度平方的复杂度。从这个角度出发，作者认为可以将这种非线性变换用一些其他近似方法，如核方法（大致思想就是将低维的向量映射到高维后，使用一些核函数直接使用低维向量来计算高维向量的内积而并不直接计算两个高维向量的内积，这样有助于加速计算。在这里则是将QK矩阵的维度降维后使用核函数来计算降维后的QK矩阵的乘积来近似未降维的QK矩阵的乘积。关于核函数的详细解析在这里 (opens new window)），替代得到广义上的注意力矩阵计算方法。虽然对于大多数核方法来说解析解并不存在，但是作者认为，由于在使用近似替代的过程中并不依赖于解析解，因此使用该方法依然是有效的。
在Performer中，作者使用正值随机特征（即该变换得到的矩阵元素都是正值，且该变换为非线性变换）来对Q、K矩阵近似变换。这种方法能够在带来更稳定的训练过程的同时还能更好地近似原始注意力矩阵（因为经过Softmax之后矩阵的元素都是正的）。关于这种近似方法的细节与数学推导可以移步原论文或其他论文解读 (opens new window)了解。
2.FAVOR
在得到Q、K矩阵的近似矩阵后，我们可以使用线性复杂度计算与存储原始注意力矩阵的近似矩阵。但即便是如此，显示的计算注意力矩阵依然会带来平方复杂度的显存消耗，那么另一个优化方法则是重新安排矩阵乘的顺序，从而达到以线性计算复杂度和线性空间复杂度的代价完成Q、K、V的计算。而这就是FAVOR+。
上图我们可以看到如果先计算与
的乘积则需要序列长度平方复杂度的显存消耗，而如果先计算
与V的乘积就只会有线性复杂度的显存消耗（这里我们只关心序列长度的复杂度，因为本文假设序列长度非常长），从而达到节省显存的目的。
但同时我们注意到，由于没有显示地构造注意力矩阵，这种计算方法无法直接使用原始注意力计算中的future mask矩阵，而mask矩阵却是非常重要。例如在自回归机器翻译系统的训练中，我们不能让单词在注意力机制中看到它后面的单词，而只能让它看到它自己与它前面的单词（因为自回归机器翻译是逐个地完成译文单词的生成，因此在模型实际生成译文的时候，正在生成的单词 之后的单词还没有被模型生成出来，所以在训练阶段也不能让模型见到未来的单词）。为了解决这一问题，作者采用了前缀和计算来达到future mask的效果，并且还能节省常规注意力矩阵计算时future mask所需的显存。（详细计算过程见论文附录B.1）
3.实验结果
实验方面，作者通过几组实验证明该方法是具备线性时间与空间复杂度的。同时还证明使用了Performer的模型速度基本上是最优的（即与完全不使用注意力机制的模型有相似的性能与复杂度）。其中下图中画X的线代表去除注意力计算的Transformer模型，可以看到Performer与这条线具有相当的复杂度量级（即随着序列长度增长，模型计算时间的增长趋势与无注意力计算的模型一致），而仅在常数项有些区别。
之后作者还设置了一组实验表示Performer与原有的注意力计算方法是兼容的。这里作者将训练好的网络的参数迁移到Performer上，实验发现稍稍经过微调的Performer就能够达到Transformer模型收敛后的效果了。这一实验表明，Performer可以减少重新训练网络的训练代价。
4.潜在应用方面：蛋白质分析
在生物医学领域，由于蛋白质本身可以看作是由20个氨基酸为单词组成的句子，因此蛋白质分析任务与NLP任务具有强相关性。基于这个特性，也有许多人尝试使用基于Transformer为骨架的模型来完成这一任务。但是他们往往会遇到蛋白质序列较长导致模型输入序列过长的问题。这时，Performer的线性时空复杂度特性就完美地契合这一任务。作者通过实验发现，尽管理论上分析使用Softmax作为非线性变换的Performer能够更好地拟合原始注意力矩阵，但是使用ReLU作为非线性变换的Performer具有最好的实验效果。
同时，作者还分析了Performer的注意力矩阵的可视图，发现Performer训练得到的注意力权重与Transformer的注意力权重非常类似，同时Performer也能在完全没有除了数据集以外的其他先验知识的情况下注意到相似的氨基酸，从而赋予更高的注意力权重（下图的D-E与F-Y）。而Performer所具备的线性时空复杂度允许研究人员可以在相同设备上训练更大规模的模型从而带来更好的模型性能。
下图则展示了在同等设备下，使用Performer方法，从而能够训练更大的模型带来的比普通Transformer更好的性能。
结论
根据实验与理论分析，基于核方法设计的Performer具有非稀疏特性，且能够与现有模型兼容（能够直接将原有参数移植到Performer上）。同时该方法也能与其它基于注意力的方法兼容，如已经开源的FAVOR与Reformer结合的工作 (opens new window)。同时作者希望Performer能够引发人们关于注意力矩阵压缩、Transformer模型本身以及核方法的其他思考。
来源 | Google AI Blog
译者 | 张译 穆永誉
单位| 东北大学自然语言处理实验室
