文 | 煉丹學徒
編 | 小軼

我們忽略掉引言和介紹，直接把工作的效果丟上來，相信就足夠令自然語言生成的相關同學心動——對于任何一個已有的Transformer生成模型，只需根據本文算法更改attention的計算順序，就可以實現

• 成倍速度提升！

• 顯存使用量降低到原來百分之個位數！

• 不需要重新訓練！

• 保證輸出結果與原來完全一致!

以BART為例，本文方法可以把顯存使用率降低為原來的96分之一！是的，不需要在效率和質量中做權衡！無腦地將本文策略應用到你的Transformer里，龐大的自回歸預訓練的生成模型速度也會變得可以接受！你甚至可以大膽地去和蒸餾模型、剪枝模型、（半）非自回歸模型比較速度。

仔細想想，我們自然語言生成的過程中，其實只有編碼和解碼是必須要計算的開銷，而作者們發現，顯卡計算的時間遠小于CPU操作和顯存IO的時間，并最終進行理論分析降低顯存的耗時、優化代碼降低CPU操作的耗時，顯著降低顯存占用和提升生成速度。本文正是聚焦在了顯存優化的部分。

對于該方法的理論描述先是發表在了ICML 2021上。之后文章作者又將EL-Attention等相關技術封裝成了一個工具包供大家一鍵調用，項目名稱FastSeq，在2021 ACL Demo paper里獲得了5 5 4的高分，并被兩位審稿人推薦為best demo paper。

感興趣的讀者可以直接安裝FastSeq工具包，僅需要一行代碼引入該庫函數，只要你用的是常見的Facebook Fairseq或者Huggingface Transformers中的模型，import 完 FastSeq，甚至不需要改代碼，就可以獲得如下的加速效果：

論文題目：
EL-Attention: Memory Efficient Lossless Attention for Generation

FastSeq項目鏈接：
https://github.com/microsoft/fastseq

論文鏈接：
https://arxiv.org/pdf/2105.04779.pdf

Arxiv訪問慢的小伙伴也可以在 【夕小瑤的賣萌屋】訂閱號后臺回復關鍵詞 【0609】 下載論文PDF~

簡單回顧Transformer的注意力計算

注意力層中，輸入是Q，K，V即query、key、value矩陣，輸出是Q、K、V隱狀態維度相同，與Q的批大小、序列長度相同的隱狀態矩陣。訓練過程中，自注意力層Q=K=V=隱狀態H；編碼器-解碼器注意力層中，Q=解碼器的隱狀態H，K=V=編碼器的隱狀態H。推斷過程中，自注意力層K=V=H是已經輸出的前文隱狀態，Q是預測的下一個詞；編碼器-解碼器注意力層中，K=V=H是編碼器的隱狀態，Q是解碼器里預測的下一個詞。

計算時，我們先把輸入的Q，K，V線性變換，得到多頭的隱狀態變小一些的（i代表第幾個頭），然后對于每一個頭，點乘歸一獲得注意力分布，用這個加權把的值取過來，再把這個個頭的低緯度信息線性方便換到之前Q，K，V的維度上作為這個頭的隱狀態計算結果，把每個頭i的隱狀態結果相加獲得最終結果。在自回歸推斷時，無論在自注意力層，還是編碼器-解碼器注意力層中，Q都是一個單詞，而K和V為輸入編碼后的隱狀態或者已經解碼的前文，都是比較長的內容。為了表達方便，后續描述中，為經過線性變換前的隱狀態，經過線性變換后的記為，表示多頭注意力中第i頭的內容。（詳細的公式描述可以看推導章節）

推斷過程中，由于需要進行beam search，所以往往把編碼器的輸出重復beam size份并cache起來。其次，因為每次只能預測下一個詞，所以自注意力層、編碼器解碼器注意力層里的會被cache起來避免重復計算，而是要預測的下一個詞的隱狀態，因此不可能被cache起來，是我們想要計算的東西。

Transformer Beam Search為啥這么慢

我們回憶一下Transformer生成訓練的時候，forward一次的速度是非常快的，但是 為什么真正去 beam search 然后推斷 inference 的時候卻很慢。

首先我們知道，訓練和推斷的時候，編碼器端的運行是相同的，所以變慢的原因都在解碼器端。即使我們設置beam search增加了一些計算量，但是實際上我們等待的時間遠遠大于理論上增加的計算量，把常見的生成任務的測試集完整的生成一遍結果，動輒幾個小時的等待時間，到底花在哪里了呢？通過每個調用函數的時間消耗分析，作者得出了結論是：推斷的代碼中，把完整的矩陣運算打散成了每次只能預測后續一個詞，零散的運算（從訓練時 teacher forcing 的完整矩陣的Q,K,V計算，變成了推斷時每次Q都只有一個詞，去和K,V自回歸地計算若干遍）從而顯存的帶寬成為了推斷速度的瓶頸。

由于有cache技術的存在，beam search 時我們往往把計算過的隱狀態存起來反復使用以避免重復計算（如果不cache，會更慢，因為要反復計算重復內容。后續分析會告訴我們，cache的速度瓶頸在顯存IO，不cache的速度瓶頸在計算速度），頻繁的顯存內容搬運和粗放的顯存使用，導致GPU memory IO的時間超過了計算時間，顯卡一直在等顯存內容的搬運。如果再加上去除重復的輸出等等CPU的操作，速度就更慢了。

我們可以再看一下本篇推送引言部分的推斷時長分布圖。左側優化前的推斷時間里，CPU相關的后處理占用了最多的時間，消耗了6.8秒；其次是庫函數中往往支持去除相鄰的連續的多少個詞的連續出現的問題，也就是圖中的ngram block函數，去處理反復生成相同單詞短語的問題，消耗了4.5秒。顯存的搬運也是時間的大頭，3.5秒，比真正解碼計算的時間3s要多。編碼只用了最少的時間，因為只有一次簡單的forward。所有的這些時間里，只有編碼和解碼是必須消耗的，EL-Attention解決掉了cache的問題，FaseSeq項目的其他部分解決了CPU相關計算的問題，最終把不必要的計算去除，優化達到耗時最少。本篇后續只介紹EL-Attention部分的提升。

Transformer 推斷過程顯存IO瓶頸

根據論文作者的分析，Transformer自然語言生成時的顯存IO瓶頸主要由以下三個問題組成：

1）在解碼器中的編碼器-解碼器注意力子層，把編碼器的輸出經過每個子層不同的線性變換得到每一層都不一樣的多頭矩陣矩陣存儲。這就導致，層的解碼器，需要把encoded hidden states存遍。甚至由于開了beam search，當前Transformer的各個庫函數中，解碼器中的每一層都還把自己層計算出來的編碼器K,V又要再重復beam size遍，占用了大量的顯存空間。解碼器中的 自注意力子層 也有相同的問題，存儲的同樣是經過線性變換后的多頭矩陣。

2）在beam search過程中，因為每一步的寬度搜索，都會導致beam candidates的得分發生變化從而導致重新排序，以及生成結束符時從candidates隊列向finished隊列搬運的過程，從而導致大量的memory IO消耗。

3）在顯卡中，如果兩個三維矩陣運算時，他們的第一維大小相同，則運算通過并行運算其中的各個二維矩陣運算完成。推斷過程中，Q只是下一個詞的隱狀態，而K，V則顯存占用比Q大得多，描述整個上文/輸入信息。Q對K和V的運算，反復加載大量顯存占用的K和V，增大IO吞吐量負擔。（EL-Attention后面則減小query的第一維，增大query第二維，從而通過一次矩陣運算得到完整的各個頭的計算結果，避免了反復加載key的值）

優化方案

后續的一切優化和計算的更改都是保證計算結果與原始Transformer完全一致的情況下展開和推導的

為了推導出更適合推斷過程的計算順序，降低存儲量，讓矩陣的運算更高效，還能保持輸出結果一致，本章節介紹EL-Attention如何進行MultiHead Attention（多頭注意力計算）的等效替換。本章節里，仍然是使用，代表線性變換前的隱狀態,代表線性變換后的低多頭注意力里第i頭的結果，。

相比于存儲解碼器段每層計算過的，EL-Attention只緩存經過線性變換之前的隱狀態H，由當前要預測詞的Q和線性變換前的H，直接計算得到注意力層的結果，從而將原始的注意力計算

變成：

很明顯的我們看到，原始計算里使用的是經過線性變換后的多頭進行計算，而EL-Attention中，則直接使用輸入的進行計算。這就是本文的核心做法，只cache隱狀態H而非多個低維度的多頭，從而進行更高效的矩陣運算，顯著減少顯存占用。

其中，和是兩個線性變換。為了表達方便，我們略去了部分計算比如矩陣運算中的bias。完整的計算方法可以看下面的推導章節。此時，我們可以拋棄所有的計算過的緩存，從而只緩存一份隱狀態 即可。其中，無需把Q計算到隱狀態變小的多頭狀態進行零散矩陣運算，而是直接在原本的hidden size進行更加完整的矩陣運算，詳細內容見推導章節如下：

推導

本章節我們一起看一下推導，確保EL-Attention的計算結果是和MultiHead Attention完全一致的。回顧傳統的注意力計算方法，是將輸入的 Q,K,V 線性變換得到維度更小，但是多份的多頭隱狀態，對于每一個頭i，進行注意力計算，然后再用 線性變換到之前隱狀態的大維度，把每個頭的隱狀態加起來。

我們假設原本 的隱狀態是維的（ 均為 維），多頭數為 ，每一頭的隱狀態是 維的，則 ，，，。

其中，，。我們記 ，則：

其中，

最終我們得到，

在推斷過程中，。

降低了多少

直觀的減少顯存使用

顯存占用：假設編碼器-解碼器注意力層 beam search 的 size 大小為b，解碼器層，則原始的 beam search會緩存 倍的encoded隱狀態。其中的倍是因為當前的庫函數實現不佳，重復beam size份造成的，可以簡單的優化掉，剩下的倍通過EL-Attention優化掉。即，編碼器-解碼器注意力子層中，把encoded hidden states的顯存占用降為。類似的，解碼器的自注意力子層中，可以把顯存占用降低1/2。

詳細的計算復雜度和顯存優化

進一步分析，EL-Attention分析注意力計算中三個步驟的計算復雜度和顯存占用復雜度。它把注意力的計算分解成三部分進行分析，第一部分是Build Key and Value(即原本計算中的把H線性變換到多頭的)，第二部分是Build Query（即原本計算中的把Q線性變換到多頭的），第三部分是進行注意力的計算。

首先看Build Key and Value，傳統的做法中，如果不cache，則需要每次進行的計算，然后把計算結果存起來（存儲復雜度）。他的計算復雜度高，需要反復重新計算，cache則相反。而EL-Attention中，由于直接使用原始的輸入K，V進行計算，無需計算出多頭的那些，因此計算和顯存都為0。

其次是Build Query，對于要預測的下一個詞的計算是繞不開的，所以無論傳統做法中是否cache，Q都要被計算到多頭的，因此計算復雜度和顯存使用相同。EL-Attetnion的這一步是函數，由于多乘了將多頭的低hidden size隱狀態變成原本的高hidden size計算，因此此處顯存多使用了（多頭數）倍。然而這個其實很小，因為畢竟Q只有后續要預測的那一個單詞的隱狀態。最終是注意力計算部分，可以看到，因為沒有緩存那些計算過的，EL-Attention的計算復雜度增大為倍，與此同時，顯存消耗降低了。

為了比較上述三個步驟，用計算換取減少顯存的操作是否收益大于付出，EL-Attention使用下面的圖來表示這種權衡的收益。下圖中，橫軸是顯存的使用量，縱軸是計算量，面積代表時間消耗。傳統做法的時間消耗由三部分組成，圖中為無邊框的藍色的大圈，灰色的大圈和橙色的小圈。EL-Attention的時間消耗由兩部分組成，虛線邊框的灰色小圈和橙色圈，可以看到，由于重新平衡了指令密度，顯存消耗和計算消耗，總時間消耗（兩個虛線邊框圓的總面積）明顯小于傳統做法（三個無邊框圓的總面積）。

實驗結果

首先，因為EL-Attention優化后的輸出結果與優化前的Transformer模型完全一致，不需要重新訓練，只需要優化推斷的計算順序，因此performance和輸出結果，原始論文中沒有展示。為了分析速度，首先，EL-Attention使用固定的假輸入去分析速度影響，他固定了編碼器端輸入1024長，然后嘗試不同的解碼器段長度、不同的beam size去比較EL-Attention和原始attention的速度。我們可以看到，cache機制雖然增加了顯存使用，但因為避免了重復計算，明顯比不cahe的速度快，而EL-Attention則又明顯的優于帶cache的beam search生成。

此外，EL-Attention在真實模型和數據集上開展試驗。它使用Transformer，BART，GPT-2作為實驗模型，其中Transformer和BART為編碼器-解碼器結構，GPT-2為只有解碼器的結構，在SQuAD 1.1問題生成、XSum摘要任務、CNN/DM摘要任務上開展試驗。beam size越大，EL-Attention的加速效果越明顯，EL-Attention作者很保守的把所有模型的beam size都開的比較小，只有4，就有了若干倍的速度提升。

EL-Attention由于顯著地減少了顯存的占用，所以可以在有限的顯存里，把batch size開大很多倍。通過增大batch size的方法，繼續提高GPU的使用率和推斷吞吐量（下表的顯存占用對比令人吃驚）：

總結

EL-Attention通過分析自然語言生成中的速度瓶頸，精確定位到了顯存IO的問題，然后通過理論分析顯存的計算方案，找到了若干致命問題，對于已經訓練好的模型，通過提出新的計算順序和算法來在對原輸出無損的情況下，優化計算量和顯存使用，從而達到了降顯存、加速生成的效果。

萌屋作者：煉丹學徒

在微軟搬磚的聯培博士在讀生，擅長烹飪和摸魚，被迫掌握豐富的增肥和減肥經驗。祝大家吃好喝好，減肥成功。

作品推薦

把數據集刷穿是什么體驗？MetaQA已100%準確率

Transformer太大了，我要把它微調成RNN

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎

總結

以上是生活随笔為你收集整理的超硬核 ICML’21 | 如何使自然语言生成提速五倍，且显存占用减低99%的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。