"Linformer" 拍了拍 "被吊打 Transformers 的后浪们"

原创 · 作者 | 龚俊民

学校 | 新南威尔士大学硕士生

公司 | Vivo AI Lab

研究方向 | 自然语言处理和可解释学习

来自 | AINLP

论文标题：《Linformer: Self-Attention with Linear Complexity》

来源：ACL 2020

链接：https://arxiv.org/abs/2006.04768

1 引言

近年来，大型的 Transformer 模型刷遍了各大 NLP 任务榜单，取得了非凡的成功。但对于长序列、训练和部署这些，模型的成本却高得吓人。因为 Transformer 中用到的自注意力与长度n呈现出 的时间和空间复杂度。本论文提出用低秩矩阵逼近的方法实现一种新的自注意力机制，名为 Linformer。它的空间和时间复杂度从原版的 降到了 。

Linformer 与其它 Transformer 变体的算法复杂度一览

本研究基于自注意力是低秩的观察，在理论和实践中都证实了注意力矩阵可以由一个低秩矩阵来近似。我们将原本的尺度化的点积注意力拆解成了多个更小的线性投射的注意力。这刚好是对原注意力去做低秩因式分解。我们在 BookCorpus 和 英文的 Wiki 上用 MLM 的训练目标预训练了一个模型，并在 GLUE 的各个任务上，以及 情感分析任务 IMDB reviews 上进行微调。结果显示，在显著的速度提升基础上，我们的模型与原版的 Transformer 相当，甚至还好于原版。

我们先来看一下原版的 Transformer 是怎样计算的，以及它的问题在哪里。接着说一说近年来对该问题的解决思路有哪些，各自的局限是什么。最后提出本论文方法，用理论证明为什么它比之前的方法都好，并用实验验证理论的可靠性。

2 原版 Transformer 的问题

多头自注意力在不同位置捕获到的信息

Transformer [1] 基于多头自注意力 (MHA)，能让模型的不同的位置互相注意到不同子空间的表征信息。它的计算方式如下：

多头注意力的架构

其中， , , 都是输入的嵌入矩阵， 是序列长度， 是嵌入维度， 是注意力头的数量，每个注意力头的计算方式如下：

其中， , , , 三个都为待训练学习到的参数。 , 都是隐维度的投射空间。 部分计算的开销是非常昂贵的。它需要把序列中每个位置的 token 都两两组合，即需要将两个 的矩阵相乘，时间空间复杂度都是 。这部分计算成为了 Transformer 的瓶颈。

3 相关工作

混合精度训练迭代

混合精度 [2]: 针对 MHA 部分的计算复杂度高，一种妥协的方法是用半精度 ( half-precision) 或混合精度 (mixed-precision) 训练[3]。这种方法又被训练时对权重量化，用直通估计器 (Straight-Through Estimator) 去近似梯度的伪量化技术进一步提升改进[4,5]。本论文采用了混合精度训练方法。

知识蒸馏的流程

知识蒸馏 [6]: 知识蒸馏力图把大模型的老师模型学到的知识教授给轻量的学生模型。最终用学生模型来推断。但这种方法有些缺点。它并未考虑老师模型训练过程中也需要加速。而且，学生模型通常要承受表现变差的风险。比如我们要对 12 层的 BERT 对 6 层的 BERT 进行蒸馏，学生模型在各个指标任务上会有大约 2.5% 的表现下降 [7]。

稀疏自注意力的思路

稀疏注意 [8]: 一种流行的做法是让每个位置的 token 去注意重要的局部(对角线附近)，而不是整个序列，从而把稀疏性引入注意力层。或者我们也可以把 部分的计算分成不同的区块，只计算选择的区块注意力。这种做法能把整体复杂度降到 [9]。但这类稀疏注意力的时空效率提升的代价是牺牲性能，大约 20% 的加速会导致 2% 的表现下降。

局部敏感性哈希

局部敏感性哈希 [10]: 最近提出的 Reformer 使用了局部敏感性哈希 (LSH) 来把自注意力的复杂度降到 。但是，实践中 Reformer 只有在序列长度大于 2048 的时候，才会有显著效率提升。而且，Reformer 中用到的多轮 LSH 实际上会增加序列的操作，从而进一步降低最终的效率增益。

提升优化器的效率: 微批量训练技术[11]，即把一个批量的数据分成许多更小的微批量来分别计算梯度，再累加，以便避免大批次撑爆 GPU 内存。梯度检查点 [12] 通过缓存部分激活层来节省内存。未缓存的激活层会在最近的检查点反向传播的时候重新计算。两种方法都可以节省内存，但增大了计算量。

综上所述，大部分的方法在同时减少空间复杂度和时间复杂度时存在局限。而我们研究出的 Transformer 变体可以克服这种局限，把时间和空间复杂度都同时降下来。

4 自注意力是低秩的

Figure 1: 左边两幅图分别是自注意力矩阵的奇异值分解 (n=512) 频谱分析。Y轴为归一化后的矩阵P的累积奇异值。右图是其热图

首先，我们来说一说为什么注意力是低秩的。这一点可以从已有的预训练语言模型 RoBERTa [13] 去分析。首先，我们对上下文映射矩阵P (之前定义的) 进行了频谱分析。一个是 12 层的 RoBERTa，在 WiKi103 数据集上用 MLM 预训练，在 IMDB 上分类。我们对模型的不同层、不同注意力头对应的矩阵 P，都进行了奇异值分解 SVD，并把超过 10k 的句子上的归一化的累积奇异值做了平均。结果显示沿着不同层、不同注意力头和不同的任务，都呈现出一个清晰的长尾分布。这表明，矩阵 P 中的大部分信息都可以由少量最大的奇异值来恢复。如 Figure 1 中的热图可见，更高层的 Transformers 中，会比更低层的 Transformers 有更大的偏度。这意味着，在更高层，更多信息集中在少量最大的奇异值中，且 矩阵 P 的秩是更低的。

以下，我们会为以上频谱结果提供一个理论分析。

定理一：自注意力是低秩的。对于矩阵 中任意满足 的列向量， ，以及 都存在一个低秩矩阵 满足以下式子：

证明：

基于上面的式子，我们的 是一个 的对角矩阵。该证明的主要思路是约翰逊-林登斯特劳斯定理(Johnson–Lindenstrauss theorem)[14]。我们构建一个服从独立同分布 的低秩矩阵

可以用 JL 定律得出，对于矩阵 中的任意列向量 ，当 时，会有：

通俗地说，一个高维空间中的点集，可以被线性地镶嵌到低维空间中，且其空间结构只遭受较小的形变。JL 定理的证明说明了，如何用随机投影法来明确求出这个变换，且该算法只需要多项式时间。降维是有代价的。如果要求尽可能地减少形变，被嵌入的低维空间则不能很低。反过来，如果要尽可能地压缩，则形变会不可避免地增加。最终JL定理给出的结论是，将维数下降到样本数的对数级，更兼容的变换是线性的，显式的，且可以被快速计算的。这部分的详细证明可以见论文附录。

当矩阵 有了低秩性质之后，一个直观想法是用奇异值分解 SVD 去用低秩矩阵 来近似P，它满足：

其中， 都为对应奇异向量的最大奇异值。基于定理1 和 低秩矩阵近似定理 (Eckart–Young–Mirsky Theorem) [15]，我们可以用 以 的误差量和 时间复杂度和空间复杂度去近似自注意力矩阵。然而，这一方法需要对每一个自注意力矩阵都计算奇异值分解，带来了额外复杂度。因此，我们提出另一种低秩近似方法来避免这一额外开销。

5 模型

我们提出了一种新的自注意力机制能让我们以线性空间和时间复杂度计算矩阵 。其核心思想是，用两个线性投射矩阵 来计算 矩阵。我们首先将正交的维度为 的 和 投射到 维度的上下文映射矩阵 。然后，我们就可以用尺度化的点积注意力算 维度的 。

这样，我们就用 计算了每个注意力头的矩阵。此计算过程的复杂度为 ，其中， ，且不会随着序列长度 的变化而变化，我们可以显著降低时间和空间的消耗。以下定理我们将证明，当 时，我们可以在 的误差范围内，近似 。

定理2：线性的自注意力。对于任意 ，以及任意 ， 如果 ，那么一定存在矩阵 ，对于矩阵 中任意的行向量 ，会有：

证明：其主要思想是基于 JL 定理。我们先证明矩阵 中的每个行向量 ，和矩阵 中的每个列向量 ，都满足：

当 的时候，如果我们不用矩阵 是低秩的信息，则 的选择依赖于序列长度 。但我们用自注意力矩阵是低秩这一事实后， 的选择就可以独立于序列长度 。其详细证明见附录2。

Linformer 还使用了其它的一些提升表现和效率的技巧。

参数共享 我们实验了三种层级的参数共享。

• Headwise: 所有注意力头共享 E 和 F 参数。
• Key-Value: 键值参数共享。对于每一层，所有的注意力头的键值映射矩阵共享参数同一参数， 。
• Layerwise: 所有层参数共享。对于所有层，都共享投射矩阵 E。

一个12层，12个头的 Transformers 模型，在分别用上述参数共享后，其参数矩阵的个数分别为 24，12，和 1。

适应性的投影维度 因为注意力矩阵是低秩的，我们可以动态调整k的大小，让更高层的 Transformers 选用更小的 k。

其它投影方法 除了用两个矩阵线性近似，还可以用均值/最大池化，或卷积的方式来近似。

6 实验和结果

模型用的 RoBERTa 的架构，语料用了 BookCorpus 和 英文的维基百科，大约 3300M 单词的预训练语料。所有实验模型的训练目标都是 MLM，在 64 张 Tesla V100 GPUs 上训练了 250k 的迭代。

预训练的实验目标要验证三个问题：

• 要如何选择 k？k 增加，表现更好。但 k 越大，其计算复杂度和空间复杂度越接近于原版的 Transformer。
• 要如何选择参数共享策略？参数共享的越多，内存占用越少，但性能也会随之下降。我们要参数共享多少？
• 序列长度变化如何影响？如果复杂度是线性的，在 k 值固定的情况下，序列增加，其消耗的时间和空间也应该是线性增长。

我们用模型的困惑度来作为模型表现指标。困惑度越低，则模型训练得越好。(a) 图展示出，在其它条件相同下，随着 k 值变大，模型的困惑度越低，表示其训练得越好。(b) 图试图说明，即便我们把 k 从 256 降低到它的一半 128，其困惑度也不会显著增加。(c) 图比较的是不同参数共享下的效果。三者的效果相当，我们可以极端一点采用 Layerwise 参数共享。(d) 图比较了在固定k的情况下，不同的序列长度对其困惑度的影响。结果显示随着长度增加，初始的困惑度长序列会比较大。但最终收敛之后，不同长度序列的困惑度都差不多。这验证了 Linformer 是线性复杂度

在下游任务中。我们也可以看到，使用 Linformer 架构训练的 RoBERTa 与原版 Transformer 架构训练的 RoBERTa 效果相当。在 k = 128 时，略逊于 RoBERTa，平均的差异小于 0.01%。当增大 k 后，二者的差别几乎可以忽略不计。我们还可以在增大 k 后，再用参数共享的技巧来减少空间消耗。其性能的损耗几乎不变。而且实验结果也表明，Linformer 的表现主要取决于投影维度 k ，而不是 n/k。

左图为节省的时间，右图为节省的内存

上表呈现出不同序列长度 和投影维度 的加速比。可以看出，当时间、空间复杂度都降到线性后，模型可以支持更长的序列，且享有的加速增益越大。

7 结论

综上所述，线性时空复杂度最大的用途就是保护环境从我做起。其次是让 transformer 在超长文本和图像上做预训练变得更兼容。最后是能大大节省自己训练预训练语言模型的开销，将费用变得平民化。

Reference

[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998–6008, 2017.

[2] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1710.03740, 2017.

[3] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations), pp. 48–53, 2019.

[4] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2704–2713, 2018.

[5] Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave, Remi Gribonval, Herve Jegou, and Armand Joulin. Training with quantization noise for extreme fixed-point compression. arXiv preprint arXiv:2004.07320, 2020.

[6] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015.

[7] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.

[8] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.

[9] Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint arXiv:1911.02972, 2019.

[10] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.

[11] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in Neural Information Processing Systems, pp. 103–112, 2019.

[12] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.

[13] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 2019.

[14] W Johnson J Lindenstrauss. Extensions of lipschitz maps into a hilbert space. Contemp. Math, 26: 189–206, 1984.

[15] Carl Eckart and Gale Young. The approximation of one matrix by another of lower rank. Psychome-trika, 1(3):211–218, 1936.

添加个人微信，备注： 昵称-学校（公司）-方向 ， 即可获得

1. 快速学习深度学习五件套资料

2. 进入高手如云DL&NLP交流群

记得备注呦