作者： 邱震宇（ 華泰證券股份有限公司 算法工程師）

知乎專欄： 我的ai之路

今天我想給大家介紹這樣一篇論文：Multi-Head Attention: Collaborate Instead of Concatenate。作者均來自

洛桑聯邦理工學院_百度百科 baike.baidu.com

看過我文章的同學肯定知道，我一直在關注bert模型的性能優化相關研究，而這篇論文正好是與transformer的性能優化相關，並且我認爲它的方法不需要做太多的適配就能應用在預訓練模型上面，實用性較高，因此推薦給大家。

衆所周知，經典的transformer架構中採用了multi-head attention機制來引導模型從不同角度學習不同的語義信息，從各種實驗對比中也能發現多頭機制確實能夠提升模型在NLP任務上的精度。然而，隨着目前大規模預訓練模型的普及，多頭注意力機制在帶來精度提升的同時，也增加了計算的成本，帶來了性能上的限制。

因此最近兩年，有些研究人員嘗試從不同的維度去探討是否能從多頭機制上去優化transformer的性能。有些工作重點關注了多頭中每個頭的注意力到底捕捉了哪些語義信息，頭與頭之間捕捉的信息是否有冗餘，例如這篇論文：Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned，提出了一種量化注意力頭重要程度的方法。還有一些工作更加激進，提出了多頭注意力機制是否有必要的疑問，例如這篇論文：Are sixteen heads really better than one。它對transformer中的每個頭都做了消融實驗，探討了每個頭在不同下游NLP任務上的作用，最後提出了一種迭代式地剪枝注意力頭的方法。

與上述工作不同，本篇論文並非直接對注意力頭進行結構性剪枝，而是關注所有注意力頭捕捉的通用信息，試圖將這些信息提取出來作爲sharing weights，每個頭各自關注自己獨有的工作，從而減少多頭注意力計算時的成本。下面我就詳細得爲大家解讀這篇論文的工作。

單個注意力頭的減負

在那篇經典的Attention is all you need論文中，對於注意力分數的計算是這樣的：

其中，當X和Y是同一個序列時，就是自注意力模型，此時 。

然而，在各種版本的transformer實現中，上述各種線性映射計算是附加bias的，即，其中  。因爲在各種深度學習框架中，默認支持broadcasting，所以這裏公式上引入了 ,相當於實現了broadcasting。

在引入了bias後，我們重新對進行展開，可得：

備註一下：論文這裏的公式貌似有點問題，最後一項應該是我推導出的項。

最後兩項在做softmax的時候可以捨棄掉，爲什麼呢？其實很簡單，我們得到的Attention分數是一個T*T的矩陣，而 和  得到的都是一個T*1的向量，最後通過 重複了T列擴充成了矩陣，因此每一行上，它的每一列的值都是相同的，因爲softmax針對的是列維度，因此後兩項對於整體的attention計算來說是一個常量，又因爲：

因此最後兩項計算可以捨棄。又因爲前面兩項中，不存在，因此我們甚至不用去定義這個bias項。

另外，對於上述推導式的第一項，由於其計算了Query和key的相互關係，因此相當於捕捉了上下文的相關信息，而第二項只包含了key的content信息，相當於捕捉了原文內容上的信息。

多頭注意力的整合

傳統的transformer中，對於不同的注意力採取的整合方式是直接拼接，如下所示：

其中，表示注意力頭的數量。我們再定義  爲總的query/key的列空間維度，而 爲每個注意力頭中query/key的列空間維度。

根據過往的研究，我們可以發現所有注意力頭之間捕捉的信息肯定是存在冗餘的。單單研究不同頭中key或者query映射矩陣的相似度是不夠的，根據我們在單頭注意力機制中的公式推導，可以知道， 同時捕捉了上下文的信息，因此可以對該項進行驗證分析。論文考慮瞭如下一種情況，即假設兩個注意力頭中的key/query映射矩陣通過一個相同的旋轉方陣  轉換爲相同的矩陣，即：

此時，兩個頭的注意力分數是相同的，因爲：

(R是一個旋轉方陣，轉置相當於逆旋轉)。爲了能夠忽視引入旋轉矩陣帶來的人工誤差，下面主要是分析不同頭之間 的相似程度。論文設計了一種對比實驗，對比A側爲對每個頭上  做PCA分解後，計算不同主成分數量的累積variance值。對比B側則是將所有頭的  進行拼接後，對拼接後的矩陣做PCA分解，計算不同主成分數量上的累積variance值。

PCA之前我有文章科普過，相當於是找一個投影空間，讓所有高維的點在這個空間上的投影儘量分開，即方差儘量大。因此variance代表了不同主成分上包含原始矩陣的信息量。variance越高，信息量就越多。

實驗結果如圖所示：

可以看到，右邊圖例上，對於拼接後的矩陣，其variance累積曲線很快就接近了1.0，表明其有一大部分成分包含的信息比較少，而左邊圖例上的曲線相對較爲平滑，表明單個頭上的 主成分包含的信息分佈較爲均勻。

從上述對比圖例可以看出，頭與頭之間的通用信息還是比較多的。論文指出拼接後的 只需要大概1/3的維度就足夠捕捉絕大部分的信息了。

提取通用信息

既然注意力頭之間存在那麼多的通用信息，那麼如何進行實際操作將其單獨提取出來呢？論文設計了一個混合向量， ， 表示整合完所有注意力頭之後的輸出的映射矩陣維度。這個向量可以通過跟模型一起學習得到。然後我們將其代入原始的多頭注意力計算中，得到：

其中 被所有注意力頭共享。上述式子我們可以理解爲 用於捕捉所有注意力頭之間的通用信息，而 則是幫助捕捉每個頭各自獨有的信息。上述整合方式有兩種作用：

1、注意力頭的表示方式更加靈活，注意力頭的維度可以根據實際情況進行改變。的維度可以根據實際情況進行設置，若  ，則相當於傳統的transformer形式，若  ，則相當於對transformer模型進行了壓縮。

2、參數計算更加高效， 是共享於所有的注意力頭的，每一輪訓練只需要計算一次。

明顯可以看出來，原始的拼接方式做多頭注意力機制是上述CollabHead方式的一種特例，此時 ，而 是一個由1和0兩種元素組成的向量，其中1的元素位置爲其對應注意力頭的映射矩陣在拼接後的整體矩陣中的位置。如圖所示：

其中， 。M使得模型在整合注意力頭的時候，讓每個注意力頭之間都互相獨立。

如果我們的目的是進行模型的壓縮，提升模型性能，那麼可以通過設計不同形式的M來實現，比如如下幾種模式：

b模式相當於對不同的head抽取不同維度的矩陣信息；c模式則是讓所有head都共享映射矩陣；d模式則是在共享映射矩陣的基礎上，進一步壓縮最終輸出的整合矩陣的維度，達到壓縮維度的效果。

實驗

論文作者利用上述優化後的transformer架構進行了NMT的實驗，主體網絡爲encoder-decoder，實驗結果如下：

可以看到，collabHead在維度縮減到1/4時，仍然能保持跟原始維度相近的效果，說明經過壓縮之後，模型只損失了較少的信息。

論文說到這裏還沒結束，下面要說的纔是我比較關注的內容，即如何對預訓練後的bert模型應用collabHead，從而提升bert的inference效率。

More collabHead on Bert

通過上述實驗，已經能夠證明CollabHead模式相比原始的簡單拼接模式，在提升性能的同時，只會損失很小的精度。而Bert模型的主體架構也是transformer，因此它也可以利用這個優化達到性能提升的效果。最簡單的方式就是在預訓練的時候就採用這種架構，而論文也比較推薦這種方式。但是對於我們這種硬件條件有限制的企業來說，從頭預訓練一個模型似乎不太現實，那麼有沒有辦法直接在finetune過程實施這種優化，從而達到性能提升的效果呢？

答案當然是可以的。論文提出了一種re-parameterize方式，直接對預訓練模型中的attention權重進行張量分解，使得分解後得到的矩陣能分別對應上述優化中的各個參數。

該方法的主要核心爲Tucker 張量分解，或者更具體得說是CP分解。（CP分解是Tucker分解的一種簡化特例）。對於張量分解，可以參考其他博客中的講解zhuanlan.zhihu.com/p/25 。這裏就簡單介紹一下。

假設當前有一個張量 ,Tucker分解可以將該張量分解爲如下形式：

,Tucker分解可以將該張量分解爲如下形式：

其中，

，更具體的有：

代表矩陣的外積。下圖能夠直觀表示上述過程：圖參考自sandia.gov/~tgkolda/pub

矩陣A，B，C通常稱爲因子矩陣，在一定程度上包含了原始X中各個維度上的主要信息，而矩陣G通常稱爲核張量，用於表徵不同因子矩陣之間互相關聯的程度。

回到我們的多頭注意力優化問題上。我可以嘗試對多頭拼接後的key/query映射權重參數進行張量分解：

然而我們可以對上述張量分解進行簡化。由於我們關注的是不同注意力頭中對齊後的映射矩陣上的值（ 簡單來說就是head_1中的矩陣的第一個元素與head_2中的矩陣的第一個元素進行對比 ），因此可以令G中非對角線上的值均爲0，另外我們期望分解後得到的矩陣列空間維度均爲 ,可以通過該參數調節模型壓縮的程度，此時可以令  。最後我們就得到一個超對角G，同時它也是一個立方，根據以上條件，可以將Tucker分解簡化爲CP分解，具體如下：

,其中 分別表示X不同維度上的因子向量，我們可以分別用因子矩陣A,B,C來表示對應所有因子向量的組合。具體來說，上述還可以表示爲：

下圖爲其矩陣分解圖例，看完應該會對上述分解有個直觀的理解：圖參考自sandia.gov/~tgkolda/pub

備註：張量分解其實也是一個優化問題，一般通過計算分解前和分解後矩陣對應元素的mse loss來進行優化學習。因此分解過程也是需要一定的時間成本。

那麼接下來，我們可以通過分解，使其分解的目標矩陣分別對應  。另外，根據第一小節介紹的單注意力頭的優化結果，結合CollabHead的注意力頭整合方式，可以得到如下式子：

其中，第一項可以在預訓練權重參數中拿取所有注意力頭的 ，並通過stack操作得到，對其進行張量分解得到的三個矩陣，最後完成第一項的計算；第二項其實也很簡單，對每個注意力頭，令  ,我們只需要拿到預訓練權重參數中的對應參數，並計算得到 就可以了。

當然整體的re-parameterize操作除了上述步驟外，還有一些應用的小技巧。在實際應用時，通常按如下步驟進行：

1、使用原始的bert模型對下游任務進行finetune，得到一個finetune-bert-task模型。

2、對finetune-bert-task進行re-parameterize操作，得到壓縮後的re-finetune-bert-task模型。

3、使用re-finetune-bert-task模型對下游任務再進行少量迭代的finetune。

經過實驗驗證，步驟3對於最後模型的精度還是很有幫助的，建議保留。

實驗驗證

我本來是想將該方法在我的NER任務中進行實際驗證，但是發現好多張量分解的工具在tensorflow的靜態圖（尤其是estimator模式）下不太適配。如果有的同學對這方面實現興趣，可以看一下tensorly框架，它支持tensorflow2.0的動態圖模式、pytorch以及MXNET，在github上的一個issue上搜到其貌似也支持靜態圖，配置如下：

import tensorly as tl

tl.set_backend("tensorflow_graph")

另外還有一個框架tensorD,我試了一下也有bug。好在論文作者非常良心得放出了基於pytorch的開源代碼，鏈接如下：github.com/epfml/collab。核心內容主要在兩個文件：collaborative_attention.py 以及swap.py。前者主要定義了CollabHead的優化邏輯。後者則是定義瞭如何將預訓練模型中的權重參數進行re-parameterize。下面主要看一下re-parameterize的內容。代碼如下：

new_layer = CollaborativeAttention(
            dim_input=layer.dim_input,
            dim_value_all=layer.dim_value_all,
            dim_key_query_all=dim_shared_query_key,
            dim_output=layer.dim_output,
            num_attention_heads=layer.num_attention_heads,
            output_attentions=False,
            attention_probs_dropout_prob=layer.attention_probs_dropout_prob,
            use_dense_layer=layer.use_dense_layer,
            use_layer_norm=layer.use_layer_norm,
            mixing_initialization=MixingMatrixInit.CONCATENATE,
        )
_, factors = parafac(
                WQWKT_per_head.detach(), dim_shared_query_key, init="random", tol=tol
            )
WQ_shared, mixing, WK_shared = factors
new_layer.key.weight.data.copy_(WK_shared.transpose(0, 1))
new_layer.query.weight.data.copy_(WQ_shared.transpose(0, 1))
new_layer.mixing.data.copy_(mixing)

首先，定義出帶CollabHead結構的attention模型，然後調用了tensorly的CP分解parafac方法，得到三個因子矩陣。最後分別將三個矩陣權重參數賦給計算圖中對應的變量中。

bq_per_head = layer.bQ.reshape([layer.num_attention_heads, -1])
content_bias = bq_per_head.unsqueeze(1) @ WK_per_head
content_bias = content_bias.squeeze(1)
new_layer.content_bias.weight.data.copy_(content_bias)

上述操作即 的部分。

我在GLUE的MRPC任務上進行了實際的驗證，結果如下：

1、對於精度損失而言，原始的bert模型在驗證集上的f1大概是88。對re-parameterize後的bert模型再進行finetune後，其在驗證集上的f1大概是86。雖然跟論文中的結果有些差距，但是精度損失還是比較可觀的。

2、對於性能提升而言，我是將原始的隱層維度768減去一半。整體來說，inference的性能大概有30%左右的提升。但是我得提一點，就是在做張量分解的時候，論文中只要幾分鐘，可是我對12層的attention做了張量分解，總共花了有18分鐘，雖然對訓練的效率並沒有拖太多後腿，但是跟論文數據確實是有點差距，後續還是要深入tensorly框架研究一下其內部機制。

3、簡單說一下該方法與原始transformer在模型參數上的一些對比。原始的attention head計算機制在忽略bias的情況下，參數量爲 ，本論文中的由於增加了混合向量，因此參數量爲 ，我們可以通過控制 的大小來控制壓縮程度。另外增加的  參數，因爲注意力頭數值通常比隱層維度數值小很多，因此其帶來的參數量增加在一定程度上可以被低消。在計算的FLOPS對比上，原始的計算機制的FLOPS爲  ，而本論文中的方法爲  ,因爲  ,因此兩者FLOPS的比率爲 。

另外，根據論文所說，該方法還可以與其他模型壓縮方法一起使用，例如和albert還有DistilBert等聯合使用，效果也還不錯，感興趣的同學可以進一步驗證。

小結

本次解讀的論文主要通過分析transformer中注意力頭之間的冗餘信息，並設計了一種優化後的多注意力頭整合方法，將通用的信息提取出來共享於所有注意力頭，讓每個注意力頭可以專注於捕捉獨有的信息。此方法經過驗證，可以在提升性能的同時只會損失極小的精度。另外，本論文還提出對預訓練的bert模型中的attention權重進行張量分解，使得本文的多注意力頭整合方法同樣可以適用於預訓練模型，拓展了本方法的適用範圍，個人認爲該方法值得在實際業務中進行嘗試。

由於本人的數學功底不深，因此對於張量分解這塊內容的解讀比較淺，還有很多問題亟待解答，例如是不是可以通過其他分解方式更高效得去分解bert模型權重？這個就需要感興趣的同學繼續深入探索了。

本文由作者授權AINLP原創發佈於公衆號平臺，歡迎投稿，AI、NLP均可。 原文鏈接，點擊"閱讀原文"直達：

https://zhuanlan.zhihu.com/p/156820477

推薦閱讀

這個NLP工具，玩得根本停不下來

如何讓Bert在finetune小數據集時更“穩”一點

模型壓縮實踐系列之——bert-of-theseus，一個非常親民的bert壓縮方法

徵稿啓示| 200元稿費+5000DBC（價值20個小時GPU算力）

文本自動摘要任務的“不完全”心得總結番外篇——submodular函數優化

Node2Vec 論文+代碼筆記

模型壓縮實踐收尾篇——模型蒸餾以及其他一些技巧實踐小結

中文命名實體識別工具（NER）哪家強？

學自然語言處理，其實更應該學好英語

斯坦福大學NLP組Python深度學習自然語言處理工具Stanza試用

關於AINLP

AINLP 是一個有趣有AI的自然語言處理社區，專注於 AI、NLP、機器學習、深度學習、推薦算法等相關技術的分享，主題包括文本摘要、智能問答、聊天機器人、機器翻譯、自動生成、知識圖譜、預訓練模型、推薦系統、計算廣告、招聘信息、求職經驗分享等，歡迎關注！加技術交流羣請添加AINLPer(id：ainlper)，備註工作/研究方向+加羣目的。

閱讀至此了，分享、點贊、在看三選一吧:pray:

相關文章