一只小狐貍帶你解鎖煉丹術&NLP秘籍

作者：蘇劍林 (來自追一科技，人稱“蘇神”)

前言

《Attention is All You Need》一文發布后，基于Multi-Head Attention的Transformer模型開始流行起來，而去年發布的BERT模型更是將Transformer模型的熱度推上了又一個高峰。當然，技術的探索是無止境的，改進的工作也相繼涌現：有改進預訓練任務的，比如XLNET的PLM、ALBERT的SOP等；有改進歸一化的，比如Post-Norm向Pre-Norm的改變，以及T5中去掉了Layer Norm里邊的beta參數等；也有改進模型結構的，比如Transformer-XL等；有改進訓練方式的，比如ALBERT的參數共享等...

以上的這些改動，都是在Attention外部進行改動的，也就是說它們都默認了Attention的合理性，沒有對Attention本身進行改動。而本文我們則介紹關于兩個新結果：它們針對Multi-Head Attention中可能存在建模瓶頸，提出了不同的方案來改進Multi-Head Attention。兩篇論文都來自Google，并且做了相當充分的實驗，因此結果應該是相當有說服力的了 。

arxiv訪問慢的小伙伴也可以在訂閱號后臺回復關鍵詞【0415】下載論文PDF。

再小也不能小key_size

第一個結果來自文章《Low-Rank Bottleneck in Multi-head Attention Models》，它明確地指出了Multi-Head Attention里邊的表達能力瓶頸，并提出通過增大key_size的方法來緩解這個瓶頸。

Multi-Head Attention

首先簡單回顧一下Multi-Head Attention，Multi-Head Attention的基礎是自然是Single-Head Attention，也叫Scaled-Dot Attention，定義如下：

（公式可以左右滑動哦）

其中。而Multi-Head Attention，就是將分別用個不同的投影矩陣投影次，然后分別做次Single-Head Attention，最后把結果拼接起來，即

Attention里有個瓶頸

在實際使用中，一般具有相同的特征維度（即hidden_size），比如BERT Base里邊是768；一般選擇12、16、24等，比如BERT base里邊是12；確定了之后，通常的選擇是讓投影矩陣，也就是說，每個Attention Head里邊，是將原始的維投影到維，然后在進行Attention運算，輸出也是維，最后把個維的結果拼接起來，得到一個維的輸出。這里的我們通常稱為head_size。

在Attention中，關鍵的一步是

這一步是描述了與的兩兩向量之間的聯系，我們可以將看成一個二元聯合分布（實際上是個一元分布，不過這個細節并不重要），如果序列長度都為，也就是每個元有個可能的取值，那么這個分布共有個值。

但是，我們將分別投影到低維后，各自的參數量只有，總的參數量是，所以式就相當于用的參數量去逼近一個本身有個值的量，而我們通常有，尤其是比較大時更是如此，因此這種建模有點“強模型所難”，這就是原論文中的“低秩瓶頸（Low-Rank Bottleneck）”的含義.

不妨試試增大key_size？

那么，解決辦法是什么呢？直接的想法是讓增大，所以要不就是減少head的數目，要不就是增加hidden_size大小。但是更多的Attention Head本身也能增強模型的表達能力，所以為了緩解低秩瓶頸而減少的做法可能得不償失；如果增加的話，那自然是能夠增強模型整體表達能力的，但整個模型的規模與計算量也會劇增，似乎也不是一個好選擇。

那沒有其他辦法了嗎？有！當我們用投影矩陣將都投影到低維時，前面都是將它們投影到維，但其實它們的維度不一定要相等，而是只需要保證的維度相等就行了（因為要做內積），為了區別，我們通常稱的維度為key_size，的維度才叫head_size，改變key_size的大小而不改變head_size的話，也不影響模型的hidden_size。

所以，這篇論文提出來的解決方法就是增大模型的key_size，它能增加Attention的表達能力，并且不改變模型整體的hidden_size，計算量上也只是稍微增加了一點。

補充說明：事實上原論文考慮的是同時增大key_size和head_size、然后Multi-Head Attention的輸出拼接之后再用一個變換矩陣降維，但筆者認為由于拼接降維這一步只是一個線性變換，所以本質上的提升還是來源于增大key_size，所以本文只強調了增大key_size這一步。此外，如果同時增大key_size和head_size，那么會導致計算量和顯存消耗都明顯增加，而只增大key_size的話，增加的資源消耗就小很多了。

來看看實驗結果～

增大key_size這個想法很簡單，也容易實現，但是否真的有效呢？我們來看看原論文的實驗結果，其實驗都是以BERT為baseline的，實驗結果圖表很多，大家直接看原論文為好，這里只分享比較有代表性的一個：

這個結果顯示，如果固定一個比較大的key_size（比如128），那么我們可以調整模型的hidden_size和head數，使得參數量可以跟原始的BERT設計一致，但是效果更優！所以，增加key_size確實是有意義的，哪怕將總體參數量重新調整到原來的一樣大，也能一定程度上提升模型的效果。這無疑對我們設計新的Transformer模型（尤其是小規模的模型）有重要的指導作用。

最后，附上我們預訓練的兩個增大了key_size的RoBERTa小模型，歡迎大家使用（我們稱之為RoBERTa⁺）：

https://github.com/ZhuiyiTechnology/pretrained-models

再缺也不能缺Talking

對Multi-Head Attention改進的第二個結果來自論文《Talking-Heads Attention》，這篇論文雖然沒有顯式地指出它跟前一篇論文的聯系，但筆者認為它們事實上在解決同一個問題，只不過思路不一樣：它指出當前的Multi-Head Attention每個head的運算是相互孤立的，而通過將它們聯系（Talking）起來，則可以得到更強的Attention設計，即標題的“Talking-Heads Attention”。

從單一分布到混合分布

在前一篇論文里邊，我們提到了低秩瓶頸，也就是由于key_size太小所以表達能力不足，因此softmax之后無法很好地建議完整的二元分布。為了緩解這個問題，除了增大key_size之外，還有沒有其他方法呢？有，比如這篇論文使用的混合分布思路。

所謂混合分布，就是多個簡單分布的疊加（比如加權平均），它能極大地增強原分布的表達能力。典型的例子是高斯混合模型：我們知道高斯分布只是一個常見的簡單分布，但多個高斯分布疊加而成的高斯混合分布（也叫高斯混合模型，GMM）就是一個更強的分布，理論上來說，只要疊加的高斯分布足夠多，高斯混合分布能逼近任意概率分布。這個例子告訴我們，想要增加Attention中分布的表達能力，又不想增加key_size，那么可以考慮疊加多個低秩分布。

那么“多個”低秩分布哪里來呢？不是有Multi-Head嘛，每個head都帶有一個低秩分布，就直接用它們疊加就行了，這就是Talking-Heads Attention了。具體來說，它的形式是：

寫起來很復雜，事實上很簡單，就是在“之后、softmax之前”用一個參數矩陣將各個的結果疊加一下而已。這樣就把原本是孤立的各個Attention Head聯系了起來，即做了一個簡單的Talking。

對上述公式，做兩點補充說明：

1、簡單起見，上述公式中筆者省去了縮放因子，如果有需要，讀者自行補充上去即可；

2、更一般的Talking-Heads Attention允許可以在這一步進行升維，即疊加出多于個混合分布，然后再用另一個參數矩陣降維，但這并不是特別重要的改進，所以不在主要篇幅介紹。

再來看看實驗結果

是不是真的有效，當然還是得靠實驗結果來說話。這篇論文的實驗陣容可謂空前強大，它同時包含了BERT、ALBERT、T5為baseline的實驗結果！眾所周知，BERT、ALBERT、T5均是某個時間段的NLP最優模型，尤其是T5還是處在superglue的榜首，并且遠超出第二名很多，而這個Talking-Heads Attention則幾乎是把它們的輝煌戰績又刷到了一個新高度！

還是那句話，具體的實驗結果大家自己看論文去，這里展示一個比較典型的結果：

這個結果顯示，使用Talking-Head Attention情況下，保持hidden_size不變，head數目越大（相應地key_size和head_size都越小），效果越優。這看起來跟前一篇增大key_size的結論矛盾，但事實上這正說明了混合分布對分布擬合能力明顯提升作用，能夠將key_size縮小時本身變弱的單一分布，疊加成擬合能力更強大的分布。當然，這不能說明就直接設key_size=1就好了，因為key_size=1時計算量會遠遠大于原始的BERT base，應用時需要根據實際情況平衡效果和計算量。

上述表格只是原論文實驗結果的冰山一角，這里再放出一個實驗表格，讓大家感受感受它的實驗陣容：

幾乎每個任務、每個超參組合都做了實驗，并給出實驗結果。如此強大的實驗陣容，基本上也就只有Google能搞出來了，而且整篇論文明顯是濃濃的“T5 Style”（還沒看過T5論文的讀者，可以去Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》感受一下），果不其然，作者之一Noam Shazeer也正是T5的作者之一。

筆者只想說，這種龐大的實驗轟炸，仿佛在向我們宣告著：

不用質疑，該調的參數我們都調了，就我們的Talking-Heads Attention最好～

插曲：神奇的論文畫風

話說回來，筆者在Arxiv上首次刷到《Talking-Heads Attention》這篇論文時，第一感覺是一篇垃圾論文。為啥？因為它的畫風是這樣的：

誰能想象到，一篇如此強大的論文，里邊居然沒有一條數學公式，取而代之的全是偽代碼！！其實偽代碼都算不上，感覺更像是直接把實驗中的Python代碼復制到了論文中，還是復制到論文主體上！筆者印象里，只有那些不入流的水論文才會這樣做，所以筆者看到的第一想法就是水文一篇。也就Google的大佬們才能這么任性，要不是耐著心多掃了幾眼，要不是不小心掃到了T5等字眼，要不是回去看作者居然清一色是Google的，這篇強大的論文就被筆者當作垃圾論文放到回收站了 不過，任性還是有任性的代價的，這篇實驗陣容這么強大又這么有效的論文，發布至今也有一個多月了，但似乎也沒什么反響，估計也跟這個任性的風格有關系～

來自文末的小結

本文介紹了兩個關于Multi-Head Attention的后續改進工作，雖然改進細節不一致，但可以說它們都是針對“低秩瓶頸”這個問題而提出的，有種殊途同歸之感。兩個工作都來自Google，實驗內容都很豐富，所以結果都比較有說服力，正在做模型結構改進工作的讀者可以參考參考。

arxiv訪問慢的小伙伴也可以在訂閱號后臺回復關鍵詞【0415】下載論文PDF。

可

能

喜

歡

• ACL2020 | FastBERT：放飛BERT的推理速度

• ACL2020 | 對話數據集Mutual：論對話邏輯，BERT還差的很遠

• LayerNorm是Transformer的最優解嗎？

• 如何優雅地編碼文本中的位置信息？三種positioanl encoding方法簡述

• 知乎搜索框背后的Query理解和語義召回技術

• 詳解ERNIE-Baidu進化史及應用場景

• 深入解析GBDT二分類算法（附代碼實現）
  

夕小瑤的賣萌屋

_

關注&星標小夕，帶你解鎖AI秘籍

訂閱號主頁下方「撩一下」有驚喜哦

總結

以上是生活随笔為你收集整理的Google | 突破瓶颈，打造更强大的Transformer的全部內容，希望文章能夠幫你解決所遇到的問題。

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