线性Transformer应该不是你要等的那个模型
?PaperWeekly 原創 ·?作者?|?蘇劍林
單位?|?追一科技
研究方向?|?NLP、神經網絡
筆者之前已經多次討論過線性 Attention 的相關內容。介紹線性 Attention 的邏輯大體上都是：標準 Attention 具有 的平方復雜度，是其主要的“硬傷”之一，于是我們 復雜度的改進模型，也就是線性 Attention。有些讀者看到線性 Attention 的介紹后，就一直很期待我們發布基于線性 Attention 的預訓練模型，以緩解他們被 BERT 的算力消耗所折騰的“死去活來”之苦。
然而，本文要說的是：抱有這種念頭的讀者可能要失望了，標準 Attention 到線性 Attention 的轉換應該遠遠達不到你的預期，而 BERT 那么慢的原因也并不是因為標準 Attention 的平方復雜度。
BERT之反思
按照直觀理解，平方復雜度換成線性復雜度不應該要“突飛猛進”才對嘛？怎么反而“遠遠達不到預期”？出現這個疑惑的主要原因，是我們一直以來都沒有仔細評估一下常規的 Transformer 模型（如BERT）的整體計算量。
很多讀者都已經知道，Transformer 的結構大體上是 Embedding 層加若干個 Transformer 層，Embedding 層的計算量很少，我們主要關心 Transformer 層。忽略殘差、Layer Normalization 等計算量比較小的層不計，每個 Transformer 層主要組成就是兩個子層：Self Attention（簡稱 SA）和 FeedForward Network（簡稱 FFN）。
雖然 Transformer 的開山之作聲稱“Attention is all you need” [1]，但是也有不少工作論證了殘差、FFN 等模塊的必要性了，比如《Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth》[2]。
現在問大家一個問題：
你覺得是SA計算量大還是FFN計算量大？
評估計算量
毋庸置疑，SA 的復雜度是 ，而FFN的復雜度則是 ，如果你直接憑此就想當然地說 SA 計算量比 FFN 大，那就錯了！
我們知道加法比乘法快很多，所以在估計計算量的時候我們主要計算要做多少次乘法，神經網絡里邊，主要的運算是矩陣相乘，不難估計按照定義一個 的矩陣乘以一個 的矩陣要做 abc 次乘法，所以 abc 就是兩個矩陣相乘的復雜度了，這是我們估算 Transformer 復雜度的依據。
設 n 為序列長度，d 為 head_size（base 版是 64），h 為 head 的數目（base 版是 12），那么 hd 就是我們通常說的“hidden_size”（base 版是 768）。對于 SA 來說，一開始是 的投影變換，即 的矩陣乘以 的矩陣做3次，因此計算量是 ；然后是 h 個 Attention 頭的運算，每個頭先是 的 與 的 相乘得到 的 Attention 矩陣（softmax 和歸一化的計算量暫且忽略），然后 的矩陣與 的 相乘得到 的矩陣，這兩步的計算量都是 ，所以總計算量是 ；最后的輸出還有一個投影變換，也是 的矩陣乘以 的矩陣，計算量是 。所以，SA 的總計算量是
至于 FFN 就比較簡單了，它就是兩個全連接層，也就是兩個矩陣變換（激活函數的計算量也忽略不計），一般的參數設置是：第一層是 的矩陣乘以 的矩陣，第二層就是 的矩陣乘以 的矩陣。所以總計算量是
這樣一來，如果 SA 的計算量比 FFN 大，就意味著
對于 base 版來說，這意味著 ！也就是說，只有當序列長度超過 1536 時，SA 的計算量才大于 FFN，在這之前，都是線性復雜度的 FFN 占主導！
這還不止，由上面的結果我們可以得到 Transformer 層總的計算量為
它是關于 n 的一次項和二次項的求和，當 n 足夠大時，復雜度自然是 ，然而二次項占主導的條件是
對于 base 版來說，這意味著 ！也就是說，當序列長度接近 5000 時，Transformer 的復雜度才真正體現出二次性！
綜合的結論
綜合上述結果，我們可以得到結論：對于 base 版來說，當序列長度不超過 1536 時，Transformer 的復雜度都是近乎線性的；當序列長度超過 1536 時，Transformer 的計算量逐漸以 Attention 為主，復雜度慢慢趨于二次方，直到長度超過 4608，才真正以二次項為主。當然這個邊界只是一個估計，實際情況可能有所偏差，大家就此感知一下范圍和數量級就好。
筆者以前也建議過很多讀者，對于不超過 2000 長度的“長文本”任務，直接用 NEZHA [3] 或者 RoFormer 這種不限長度的模型試試，不要想太多的技巧，原因也是如此。你想再多技巧，也頂多是降到線性復雜度，而在這個長度范圍內模型本身就是近乎線性的，各種技巧也省不了多少。
對于老老實實用 BERT base 的讀者來說，maxlen 一般不超過 512，遠低于上述界限，因此就不要再說 Attention 的平方復雜度費硬件之類的吐槽了，因為事實是：
BERT 之所以慢，主要是因為它真的大，而不是因為 Attention 的平方復雜度。
“線性”含義
至于對線性 Attention “遠遠達不到預期”而感到疑惑的另一個原因，則是沒有從實際情況分析線性 Attention 的計算量，以至于對線性 Attention 期待過高。
線性 Attention 的介紹可以參考《線性 Attention 的探索：Attention 必須有個 Softmax 嗎？》，這里不做重復。簡單來說，線性 Attention 就是按照 的順序算注意力。所以按照前面的估算方法，線性 Attention 每個頭運算的計算量就是 ，而標準 Attention 則是 ，因此如果 ，那么線性 Attention 是比標準 Attention 要省計算量的。（注：實現線性效率的 Attention 也不止這一種思路，但總的而言復雜度是相似的，因此下面的結論也有代表性。）
對于 base 版來說，那就是 ，這個界還是很容易達到的，所以有些讀者可能會想“能省一點是一點”、“不用白不用”。然而，這是假設了標準 Attention 與線性 Attention 都用同一個 d 的前提下得出的結果。
而認真琢磨過《Performer：用隨機投影將 Attention 的復雜度線性化》、《Transformer 升級之路：3、從? Performer 到線性 Attention》的讀者都知道，線性 Attention 有著比標準 Attention 更嚴重的“低秩瓶頸”，所以如果切換為線性 Attention 后還用同一個? d，那么線性 Attention 的效果將會明顯下降，而如果要保留大致相同的效果，那么線性 Attention 要用更大的 d（一般是原來的 4 倍左右）。
這樣一來，線性 Attention 的計算量應該是 ，如果線性 Attention 要比標準 Attention 快，那么就要 ，對于 base 版來說，就是 ，這也超出了一般讀者所能用到的范圍了。況且換成線性 Attention 后，前面關于 SA 和 FFN 的計算量結論依然存在，即大部分序列長度下占主導計算量的還是FFN等線性運算，換了線性 Attention 后也無法感覺到明顯的速度提升。所以，總的來說
你要不是成千上萬的序列長度，就不要想著換線性 Attention 了。
再翻翻論文
事實上，就算不進行上述分析，只要認真讀過關于 Attention 效率改進相關工作的讀者，從論文中的某些圖片就可以得到類似的結論：所謂更“高效”的 Attention，一般都只適用于成千上萬的序列長度，只有在這個場景下性能才有明顯提升。
比如較早的工作 Sparse Transformers [4]，里邊有一張圖顯示出處理的序列長度都是 3000+ 的：
比如大名鼎鼎的 Reformer [5]，演示性能的序列長度都是以 K 為單位的：
大家頗多好評的 Longformer [6] 也是如此：
還有 Google 關于線性 Attention 的經典之作 Performer [7]，顯示出哪怕序列長度是 ，Performer 與 Transformer 的差距也不能說特別顯著：
最后是比較新的工作 Luna，提供了一個比較綜合的對比表格，同樣支持我們的結論：
從已有的各個高效 Attention 的工作中，我們可以得出結論：這些改進工作所關心的序列長度主要都是以千為單位的，有明顯計算效率提升的序列長度基本上都要好幾千；當然，我們前面的討論主要針對的還是時間復雜度，對于空間復雜度，也就是顯存占用量，降低的幅度一般要比時間復雜度提升的幅度的要大，但總體而言都是長序列才有價值。
換個期待吧
所以，如果你的序列長度還只是一兩百，那么就完全不要期望 Attention 本身的改進了，老老實實換個小點的模型就好。你可以期望未來會有更小的模型能達到同樣好的效果，但是不要期望同樣大的模型通過修改 Attention 來提升效率，因為說白了，就算把 Attention 完全去掉，也提升不了多少性能。
參考文獻
[1] https://arxiv.org/abs/1706.03762
[2] https://arxiv.org/abs/2103.03404
[3]https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/NEZHA-TensorFlow
[4] https://arxiv.org/abs/1904.10509
[5] https://arxiv.org/abs/2001.04451
[6] https://arxiv.org/abs/2004.05150
[7] https://arxiv.org/abs/2009.14794
特別鳴謝
感謝 TCCI 天橋腦科學研究院對于 PaperWeekly 的支持。TCCI 關注大腦探知、大腦功能和大腦健康。
更多閱讀
#投 稿?通 道#
?讓你的文字被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。
?????稿件基本要求：
? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?
? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題
? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算
?????投稿通道：
? 投稿郵箱：hr@paperweekly.site?
? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者
? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿
△長按添加PaperWeekly小編
????
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
·
總結
以上是生活随笔為你收集整理的线性Transformer应该不是你要等的那个模型的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 深国投是什么公司
- 下一篇： 今日arXiv精选 | 12篇EMNLP