通过 Performer 架构再探注意力机制

2020 年 11 月 30 日 谷歌开发者

文 / Krzysztof Choromanski 和 Lucy Colwell，Google Research 研究员

Transformer 模型已在多个领域中取得世界前沿成果 (SOTA)，包括自然语言、对话、图像甚至是音乐。Transformer 架构的核心组成部分是注意力模块，通过它计算输入序列中所有位置对的相似性得分。然而，随着输入序列的增长，注意力机制的扩展能力却无法与之匹配。这需要计算时间呈平方增长来生成所有相似性得分，以及存储空间的平方增长来构建一个矩阵存储这些得分。

图像
https://openai.com/blog/image-gpt/
音乐
https://magenta.tensorflow.org/music-transformer

对于需要长距离注意力的应用，我们提出了几种快速且节省空间的替代方法（例如内存缓存技术），但更普遍的方法是依靠稀疏注意力 (Sparse attention)。稀疏注意力从序列而不是所有可能的位置对来计算有限的相似性得分，从而减少注意力机制的计算时间和内存需求，最终生成稀疏矩阵而不是完整矩阵。如 Sparse Transformer、Longformer、Routing Transformer、Reformer 和 Big Bird 等方法所展示的那样，这些稀疏条目可以手动提出、通过优化方法找到、通过学习获得，甚至可以进行随机化。由于稀疏矩阵也可以用图和边的概念来表示，因此稀疏化方法也受到了图神经网络文献的启发，并与注意力建立了图注意力网络中所概述的特定关系。此类基于稀疏性的架构通常需要额外的注意力层来隐式生成完整的注意力机制。

稀疏注意力
https://openai.com/blog/sparse-transformer/

标准稀疏化技术，左侧：稀疏模式示例，其中 token 仅注意其他邻近的 token；右侧：在图注意力网络中，token 仅注意图中的邻近 token，该 token 的相关性应高于其他节点。请参阅 Efficient Transformers: A Survey，全面了解各种不同类别的方法

遗憾的是，稀疏注意力方法仍有一些局限：

它们需要高效的稀疏矩阵乘法运算，但这并不是所有加速器都能做到的；
它们通常无法为自己的表示能力提供严格的理论保证；
它们主要针对 Transformer 模型和生成式预训练进行优化；
它们通常会堆叠更多的注意力层来弥补稀疏表示，因此难以与其他预训练模型一起使用，需要重新训练并消耗大量能源。

除上述缺点以外，稀疏注意力机制在通常情况下仍不足以解决需要应用常规注意力方法的所有问题（例如指针网络）。此外，还存在一些无法稀疏化的运算，例如常用的 softmax 运算。该运算可标准化注意力机制中的相似性分数，且在业界被广泛应用的推荐系统中有着大量应用。

业界被广泛应用的推荐系统
https://research.google/pubs/pub45530/

为解决这些问题，我们提出了 Performer，这是一种具有可线性扩展注意力机制的 Transformer 架构，因此可以加快训练速度，同时允许模型处理更长的长度，从而满足某些图像数据集（如 ImageNet64）和文本数据集（如 PG-19）的训练需求。Performer 使用了高效（线性）的广义注意力框架，可支持使用基于不同相似性度量（内核）的多种注意力机制。我们使用了新的通过正正交随机特征获得快速注意力 (FAVOR+) 算法来实现这一框架。该算法为注意力机制提供了可扩展的低方差和无偏估计 ，并可通过随机特征图分解来表示（具体来说，就是常规的 softmax-attention）。在保持了线性的空间和时间复杂度的情况下，该方法还确保了良好的准确率，并且可以应用到独立的 softmax 运算。

Performer
https://arxiv.org/abs/2009.14794
通过正正交随机特征获得快速注意力 (FAVOR+) 算法
https://github.com/google-research/google-research/tree/master/performer/fast_self_attention

广义注意力机制

在最初的注意力机制中，分别对应矩阵行与列的 query 和 key 输入会进行相乘，然后通过 softmax 运算形成一个注意力矩阵，用来存储相似性得分。值得注意的是，虽然在这种方法中，将 query-key 的乘积传递给非线性 softmax 运算后，便无法再将其分解为最初的 query 和 key。但是它可以将注意力矩阵分解为一个将原始 query 和 key 作为输入的随机非线性函数的乘积，也被称作随机特征，这样就可以更加高效地对相似性信息进行编码。

左侧：标准注意力矩阵，包含每一对条目的所有相似性得分（通过对 query 和 key 进行 softmax 运算得到），表示为 q 和 k；右侧：标准注意力矩阵可通过低秩的随机矩阵 Q′ 和 K′ 来近似，其中的行对原始 query/key 的潜在随机非线性函数进行编码。对于常规的 softmax-attention，该转换非常紧凑，并且涉及指数函数以及随机高斯投影

常规的 softmax-attention 可以看作是由指数函数和高斯投影定义的一个非线性函数特例。请注意，我们也可以进行反向推理，方法是先实现更广义的非线性函数，再隐式定义 query-key 乘积中的其他类型的相似性度量（内核）。我们基于早期核方法方面的研究，确立了广义注意力机制。尽管对于大多核方法来说，闭式公式并不存在，但由于我们的机制并不依赖于闭式公式，因此仍然可以应用。

核方法
https://papers.nips.cc/paper/3182-random-features-for-large-scale-kernel-machines

据我们所知，我们的研究首次证明了任意注意力矩阵都可以使用随机特征在下游 Transformer 应用中实现有效的近似计算。我们实现这一点所使用的新机制是使用正随机特征，即原始 query 和 key 的正值非线性函数，这对于避免训练过程中的不稳定性至关重要，并实现了对常规 softmax 注意力的更准确近似。

迈向 FAVOR+：通过矩阵关联性实现快速注意力

上文所述的分解使我们能够以线性而非呈平方增长的内存复杂度的方式存储隐式注意力矩阵。而我们还可以通过分解获得一个线性时间注意力机制。虽然原始注意力机制可以通过将已存储的注意力矩阵与值输入相乘得到最终结果，但在分解注意力矩阵之后，我们还可以重新排列矩阵乘法，以对常规注意力机制的结果进行近似，而不需要显示构建大小呈平方增长的注意力矩阵。基于这个方法，我们最终得到了新的算法 FAVOR+。

左侧：标准注意力模块计算，通过使用注意力矩阵 A 和值张量 V 进行矩阵乘法来计算最终的预期结果；右侧：通过解耦低秩分解 A 中使用的矩阵 Q′ 和 K′，并按照虚线框中指示的顺序执行矩阵乘法，我们便获得了一个线性注意力矩阵，不再需要显式构建 A 或其近似

上述分析与所谓的双向注意力机制（即没有过去和未来概念的非因果注意力机制）相关。针对输入序列中的 token 不会注意后方其他 token 的单向（即因果）注意力机制，我们对方法进行了微调，以使用前缀和计算 (Prefix-Sum Computation)，此计算方法仅存储矩阵计算的运行总数，而不存储显式的下三角常规注意力矩阵。

左侧：标准单向注意力机制需要屏蔽注意力矩阵以获得其下三角部分；右侧：左侧的无偏近似可通过前缀和机制获得，其中用于 key 和值向量的随机特征图的外积前缀和通过动态构建实现，并通过查询随机特征向量进行左乘计算，以在最终矩阵中获得新行

性能

我们首先对 Performer 的空间和时间复杂度进行了基准测试，其结果表明，注意力计算的加速和内存占用的减少在实践中近乎达到最优的效果，也就是说，非常接近在模型中完全不使用注意力机制的性能。

在以时间 (T) 和长度 (L) 都以对数坐标表示的坐标图中，常规 Transformer 模型的双向计时。达到 GPU 内存上限后，线条中断。黑线（带有 X 标记）表示使用“伪”注意力模块时，可实现的最大内存压缩和注意力加速，其基本上绕过了注意力计算，并代表模型的最优效率。Performer 模型几乎能够在注意力组件中达到这一最佳性能

我们进一步证明了，通过使用我们提出的无偏置近似 softmax 函数， Performer 模型在稍微进行微调之后可以向后兼容预训练的 Transformer 模型，从而有可能在提升推理速度的同时降低能耗，并且不需要对现有的预训练模型重新训练。

在使用 10 亿单词基准 (LM1B) 数据集时，我们将原始的预训练 Transformer 的权重迁移至 Performer 模型，获得了 0.07 的初始非零准确率（橙色虚线）。但是，在微调之后，Performer 的准确率只需用原始梯度步数的一小部分即可迅速恢复

10 亿单词基准
https://arxiv.org/abs/1312.3005

应用示例：蛋白质建模

蛋白质是具有复杂 3D 结构的大分子，拥有生命必不可少的特定功能。和单词一样，我们也可以将蛋白质指定为线性序列，其中每个字符代表 20 个氨基酸基本模块之一。将 Transformer 应用于大型未标记的蛋白质序列语料库（例如 UniRef）后，其生成的模型可用于精确预测已折叠的功能大分子。Performer-ReLU（使用基于 ReLU 的注意力机制，这是与 softmax 不同的广义注意力机制实例）在蛋白质序列数据建模方面表现良好，而 Performer-Softmax 的性能可以与 Transformer 相媲美，这与我们的研究理论结果所预测的一致。

蛋白质序列建模的性能表现。训练 = 虚线，验证 = 实线，单向 = (U)，双向 = (B)。在运行的所有训练中，我们均使用 ProGen (2019) 中的 36 层模型参数，每个参数均使用 16x16 TPU-v2。在给定相应的计算约束后，每次运行的批次大小都可以达到最大

下方展示了一个蛋白质 Performer 模型的可视化效果图，该模型使用基于 ReLU 的近似注意力机制进行训练。使用 Performer 来估计氨基酸之间的相似性，可以恢复与知名的替换矩阵 (Substitution Matrice) 相似的结构，而该矩阵本来需要通过分析仔细策划的序列比对上的进化替代模式 (Evolutionary Substitution Pattern) 获得。从更广泛的范围上看，我们发现局部和全局注意力模式与基于蛋白质数据训练的 Transformer 模型相一致。Performer 的密集注意力近似有可能捕获跨多个蛋白质序列的全局交互作用。为执行概念验证，我们在串联蛋白长序列上进行了模型训练。在这种情况下，常规的 Transformer 模型会出现内存溢出，而 Performer 由于具有良好的空间利用效率，则不会出现这一问题。

左侧：使用注意力权重估计的氨基酸相似性矩阵。在仅访问了蛋白质序列且无需事先了解生物化学信息的情况下，该模型可以识别高度相似的氨基酸对，例如 (D,E) 和 (F,Y)。中间：通过 BPT1_BOVIN 蛋白质的 4 层（行）和 3 个选定的头（列）建立的注意力矩阵，显示了局部和全局注意力模式

在最大长度达 8192（通过连接单个蛋白质序列获得）的序列上的性能表现。为适应 TPU 内存，减小了 Transformer 的大小（层数和嵌入尺寸）

结论

我们的工作为基于非稀疏性方法和基于核方法的 Transformer 解释的最新进展做出了贡献。我们的方法可与其他技术（如可逆层）实现互操作，甚至还将 FAVOR 与 Reformer 的代码进行了整合。我们在此处提供了本论文、Performer 代码以及蛋白质语言建模代码的相关链接。我们相信，我们的研究为注意力机制、Transformer 架构甚至核方法开辟了全新的研究方向。

FAVOR 与 Reformer 的代码
https://github.com/google/trax/blob/master/trax/layers/research/sparsity.py#L215
本论文
https://arxiv.org/abs/2009.14794
Performer 代码
http://github.com/google-research/google-research/tree/master/performer
蛋白质语言建模代码
http://github.com/google-research/google-research/tree/master/protein_lm

致谢

此项研究由 Performer 核心设计师 Krzysztof Choromanski（Google Brain 团队，技术和研究主管）、Valerii Likhosherstov（剑桥大学）和 Xingyou Song（Google Brain 团队）合作完成，David Dohan、Andreea Gane、Tamas Sarlos、Peter Hawkins、Jared Davis、Afroz Mohiuddin、Lukasz Kaiser、David Belanger、Lucy Colwell 和 Adrian Weller 亦有贡献。同时，特别感谢应用科学团队共同领导了将高效 Transformer 架构应用于蛋白质序列数据的研究工作。

还要感谢 Joshua Meier、John Platt 和 Tom Weingarten 就生物学数据与我们展开了许多富有成果的讨论，并对本文稿给予了建设性的评论，以及 Yi Tay 和 Mostafa Dehghani 对比较基线的讨论。最后，感谢 Nikita Kitaev 和 Wojciech Gajewski 对 Reformer 展开的多次讨论，以及 Aurko Roy 与 Ashish Vaswani 对 Routing Transformer 展开的多次讨论。

如果您想详细了解 本文讨论 的相关内容，请参阅以下文档。这些文档深入探讨了这篇文章中提及的许多主题：