原創 · 作者 | 龔俊民
學校 | 新南威爾士大學碩士生
公司 | Vivo AI Lab
研究方向 | 自然語言處理和可解釋學習
來自 | AINLP
論文標題：《Linformer: Self-Attention with Linear Complexity》
來源：ACL 2020
鏈接：https://arxiv.org/abs/2006.04768
1 引言
近年來，大型的 Transformer 模型刷遍了各大 NLP 任務榜單，取得了非凡的成功。但對於長序列、訓練和部署這些，模型的成本卻高得嚇人。因爲 Transformer 中用到的自注意力與長度n呈現出 的時間和空間複雜度。本論文提出用低秩矩陣逼近的方法實現一種新的自注意力機制，名爲 Linformer。它的空間和時間複雜度從原版的 降到了 。
本研究基於自注意力是低秩的觀察，在理論和實踐中都證實了注意力矩陣可以由一個低秩矩陣來近似。我們將原本的尺度化的點積注意力拆解成了多個更小的線性投射的注意力。這剛好是對原注意力去做低秩因式分解。我們在 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 自注意力是低秩的
首先，我們來說一說爲什麼注意力是低秩的。這一點可以從已有的預訓練語言模型 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交流羣
記得備註呦
