＂Linformer＂ 拍了拍 ＂被吊打 Transformers 的後浪們＂

論文標題：《Linformer: Self-Attention with Linear Complexity》

來源：ACL 2020

連結：https://arxiv.org/abs/2006.04768

1 引言

近年來，大型的 Transformer 模型刷遍了各大 NLP 任務榜單，取得了非凡的成功。但對於長序列、訓練和部署這些，模型的成本卻高得嚇人。因為 Transformer 中用到的自注意力與長度n呈現出

Linformer 與其它 Transformer 變體的算法複雜度一覽

本研究基於自注意力是低秩的觀察，在理論和實踐中都證實了注意力矩陣可以由一個低秩矩陣來近似。我們將原本的尺度化的點積注意力拆解成了多個更小的線性投射的注意力。這剛好是對原注意力去做低秩因式分解。我們在 BookCorpus 和 英文的 Wiki 上用 MLM 的訓練目標預訓練了一個模型，並在 GLUE 的各個任務上，以及 情感分析任務 IMDB reviews 上進行微調。結果顯示，在顯著的速度提升基礎上，我們的模型與原版的 Transformer 相當，甚至還好於原版。

我們先來看一下原版的 Transformer 是怎樣計算的，以及它的問題在哪裡。接著說一說近年來對該問題的解決思路有哪些，各自的局限是什麼。最後提出本論文方法，用理論證明為什麼它比之前的方法都好，並用實驗驗證理論的可靠性。

2 原版 Transformer 的問題

多頭自注意力在不同位置捕獲到的信息

Transformer [1] 基於多頭自注意力 (MHA)，能讓模型的不同的位置互相注意到不同子空間的表徵信息。它的計算方式如下：

多頭注意力的架構

其中，

其中，

3 相關工作

混合精度訓練迭代

混合精度 [2]: 針對 MHA 部分的計算複雜度高，一種妥協的方法是用半精度 ( half-precision) 或混合精度 (mixed-precision) 訓練[3]。這種方法又被訓練時對權重量化，用直通估計器 (Straight-Through Estimator) 去近似梯度的偽量化技術進一步提升改進[4,5]。本論文採用了混合精度訓練方法。

知識蒸餾的流程

知識蒸餾 [6]: 知識蒸餾力圖把大模型的老師模型學到的知識教授給輕量的學生模型。最終用學生模型來推斷。但這種方法有些缺點。它並未考慮老師模型訓練過程中也需要加速。而且，學生模型通常要承受表現變差的風險。比如我們要對 12 層的 BERT 對 6 層的 BERT 進行蒸餾，學生模型在各個指標任務上會有大約 2.5% 的表現下降 [7]。

稀疏自注意力的思路

稀疏注意 [8]: 一種流行的做法是讓每個位置的 token 去注意重要的局部(對角線附近)，而不是整個序列，從而把稀疏性引入注意力層。或者我們也可以把

局部敏感性哈希

局部敏感性哈希 [10]: 最近提出的 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 的秩是更低的。

以下，我們會為以上頻譜結果提供一個理論分析。

定理一：自注意力是低秩的。對於矩陣

證明：

基於上面的式子，我們的

可以用 JL 定律得出，對於矩陣

通俗地說，一個高維空間中的點集，可以被線性地鑲嵌到低維空間中，且其空間結構只遭受較小的形變。JL 定理的證明說明了，如何用隨機投影法來明確求出這個變換，且該算法只需要多項式時間。降維是有代價的。如果要求儘可能地減少形變，被嵌入的低維空間則不能很低。反過來，如果要儘可能地壓縮，則形變會不可避免地增加。最終JL定理給出的結論是，將維數下降到樣本數的對數級，更兼容的變換是線性的，顯式的，且可以被快速計算的。這部分的詳細證明可以見論文附錄。

當矩陣

其中，

5 模型

我們提出了一種新的自注意力機制能讓我們以線性空間和時間複雜度計算矩陣

這樣，我們就用

定理2：線性的自注意力。對於任意

證明：其主要思想是基於 JL 定理。我們先證明矩陣

當

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.