自己挖坑自己填，谷歌大改Transformer注意力( 二 ) 机器之心报道机器之心编辑部

该研究首次证明了，任意注意力矩阵都可以通过随机特征在下游 Transformer 应用中实现有效地近似。实现这一点的的新机制是使用正随机特征，即原始 query 和 key 的正直非线性函数，这对于避免训练过程中的不稳定性至关重要，并实现了对常规 softmax 注意力的更准确近似。
新算法 FAVOR+：通过矩阵相关性实现快速注意力
上文描述的分解允许我们以线性而非二次内存复杂度的方式存储隐式注意力矩阵。我们还可以通过分解获得一个线性时间注意力机制。虽然在分解注意力矩阵之后，原始注意力机制与具有值输入的存储注意力矩阵相乘以获得最终结果，我们可以重新排列矩阵乘法以近似常规注意力机制的结果，并且不需要显式地构建二次方大小的注意力矩阵。最终生成了新算法 FAVOR+ 。
文章插图
左：标准注意力模块计算，其中通过执行带有矩阵 A 和值张量 V 的矩阵乘法来计算最终的预期结果；右：通过解耦低秩分解 A 中使用的矩阵 Q′和 K′以及按照虚线框中指示的顺序执行矩阵乘法，研究者获得了一个线性注意力矩阵，同时不用显式地构建 A 或其近似。
上述分析与双向注意力（即非因果注意力）相关，其中没有 past 和 future 的概念。对于输入序列中没有注意前后 token 的单向（即因果）注意力而言，研究者稍微修改方法以使用前缀和计算（prefix-sum computation），它们只存储矩阵计算的运行总数，而不存储显式的下三角常规注意力矩阵。
文章插图
左：标准单向注意力需要 mask 注意力矩阵以获得其下三角部分；右：LHS 上的无偏近似可以通过前缀和获得，其中用于 key 和值向量的随机特征图的外积（outer-product）前缀和实现动态构建，并通过 query 随机特征向量进行左乘计算，以在最终矩阵中获得新行（new row）。
性能
研究者首先对 Performer 的空间和时间复杂度进行基准测试，结果表明，注意力的加速比和内存减少在实证的角度上近乎最优，也就是说，这非常接近在模型中根本不使用注意力机制的情况。
文章插图
在以时间（T）和长度（L）为度量的双对数坐标轴中，常规 Transformer 模型的双向 timing 。
研究者进一步证明，使用无偏 softmax 近似，该 Performer 模型在稍微进行微调之后可以向后兼容预训练 Transformer 模型，从而在提升推理速度的同时降低能耗，并且不需要从头训练预先存在的模型。
文章插图
在 One Billion Word Benchmark (LM1B) 数据集上，研究者将原始预训练 Transformer 的权重迁移至 Performer 模型，使得初始非零准确度为 0.07（橙色虚线）。但在微调之后， Performer 的准确度在很少的梯度步数之后迅速恢复。
应用示例：蛋白质建模
蛋白质具有复杂的 3D 结构，是生命必不可少的拥有特定功能的大分子。和单词一样，蛋白质可以被看做线性序列，每个字符代表一种氨基酸。将 Transformers 应用于大型未标记的蛋白质序列语料库，生成的模型可用于精确预测折叠功能大分子。正如该研究理论结果所预测的那样， Performer-ReLU 在蛋白质序列数据建模方面表现良好，而 Performer-Softmax 与 Transformer 性能相媲美。
文章插图
Performer 在蛋白质序列建模时的性能。
下面可视化一个蛋白质 Performer 模型，该模型使用基于 ReLU 的近似注意力机制进行训练。研究者发现， Performer 的密集注意力近似有可能捕捉到跨多个蛋白质序列的全局相互作用。作为概念的证明，研究者在串联蛋白长序列上训练模型，这使得常规的 Transformer 模型内存过载。但由于具有良好的空间利用效率， Performer 不会出现这一问题。
文章插图
左：从注意力权重估计氨基酸相似性矩阵。该模型可以识别高度相似的氨基酸对，例如 (D,E) 和 (F,Y) 。
文章插图
Performer 和 Transformer 在长度为 8192 的蛋白质序列上的性能。
随着 Transformer 的频繁跨界，越来越多的研究者开始关注其内存占用和计算效率的问题，比如机器之心前段时间介绍的《抛弃注意力，比 EfficientNet 快 3.5 倍，类 Transformer 新模型跨界视觉任务实现新 SOTA》。在那篇文章中，研究者提出了一种名为「lambda」的层，这些层提供了一种捕获输入和一组结构化上下文元素之间长程交互的通用框架。类似的改进还在不断涌现，我们也将持续关注。