
自注意力機制使Transformer成為可解釋性較強的模型之一,基於query-key的點積自注意力作為自注意力機制的核心元件看起來是必不可少的:點積自注意力使成對的token被完全連線,並且能夠對長距離依賴的資訊進行建模。但是點積自注意力的計算真的是必須的嗎?越來越多的工作深入探索自注意力,並給出了否定的回答。本文從點積自注意力的原理出發,並以兩篇參考文獻以例,重新思考自注意力機制,進一步對點積自注意力的可替代性展開討論。
Transformer模型的核心為自注意力機制,同時點積自注意力是最必不可少的核心元件。我們透過下面的公式[1]理解點積自注意力的含義:
具體說來,在點積自注意力中,透過對同一個輸入 經過不同的投影矩陣 得到 , , 。為了方便後文理解,這裡將統一點積自注意力及權重的計算表示式:
點積自注意力是如何提供了強大的建模能力呢?考察其幾何含義上:基於query-key的點積計算確定了序列中每一個token相對於所有其他token的重要程度,來學習自對齊矩陣。點積自注意力機制可以理解為一個模擬基於內容的檢索過程,過程的核心是成對token間的互動。
然而,特定於例項的互動容易使注意力權重產生自由波動,並且相對於整個序列來說依舊錶徵區域性的資訊變化,缺乏一致的全域性語境。越來越多的工作深入探索自注意力並發現:自注意力旨在生成自對齊矩陣;透過query-key點積計算自注意力可能是不必要的,甚至可以說完全不需要基於內容的自注意力。下面本文將以兩篇參考文獻以例,重新思考自注意力機制,分析點積自注意力的可替代性並解釋文獻中自對齊矩陣的生成方式。
Google Research提出了一種合成器模型[3]。該模型直接合成自對齊矩陣,無需標準自注意力中成對token的點積互動。論文中介紹了四種合成模式:全連接合成器、隨機合成器、分解全連接合成器和分解隨機合成器,下面將對這四種合成模式展開分析。
仔細思考公式(1)和(2)可以發現,注意力權重矩陣的計算過程本質上是 從(表示N個token,每個token的維度為d)的輸入得到的過程。若不採用成對token的點積互動,一種簡單的方法是透過全連線層來進行維度間的變換。論文中提出了全連接合成器,即兩層全連線層加上線性整流啟用函式(ReLU)生成B矩陣。如圖1(b)所示,這一步驟中,序列中的每個token都是獨立計算的,token之間並不產生任何互動。最後,注意力權重矩陣經過Softmax操作與value向量相乘得到最終的注意力輸出。上述過程與點積自注意力模型相似,雖然全連接合成器也利用區域性樣本資訊,但作者用全連線層替換掉了標準自注意力中的query-key點積計算。

如圖1(b)所示,為了得到 大小的注意力權重矩陣B,全連接合成器將序列中的每個token進行獨立對映。在此基礎上,作者進一步提出了隨機合成器(圖1(c)),隨機合成器透過隨機初始化得到注意力權重B,B可以選擇隨訓練更新或不更新。該合成器模型既不依賴於成對的token也不依賴於內容,且過程中利用全域性樣本資訊,學習一種任務式的對齊方式。
全連接合成器在生成注意力權重矩陣的過程為網路添加了 的引數,而隨機合成器在此過程為網路添加了 的引數。如果在序列很長的情況下,上述兩種模型均具有很大的引數量。因此,作者針對兩種合成器分別提出了對應的分解模型來減少引數量及防止過擬合。
分解全連接合成器模型針對可能過大的序列長度 ,先經過全連接合成器生成兩個矩陣記為 和 , 。然後將 中表示token的向量重複 次, 中表示token的向量重複 次後,做元素積。上述過程可將兩個分解矩陣恢復到B,減小模型引數量。類似地,分解隨機合成器模型先經過隨機合成器生成兩個矩陣記為 和 。再將 與 的轉置做外積恢復到B,同樣減小了模型引數量。表1將上述合成器的特性做了彙總,更多細節見表1所示:

由表1可知,上述四種合成器模型都代替了query-key的點積互動,且模型引數量都相對減少了。
此文提出可將上述合成器模型和點積方法以加和的方式組合使用生成自對齊矩陣。作者分別在機器翻譯、語言模型、文字生成、多工自然語言處理等任務上進行了實驗。在機器翻譯和多工自然語言處理等任務上,混合合成器(Random + Dot)模型[2]與其他模型相比,取得最好的效果。兩個分解方法雖不如混合模型取得更多的提升,但在一定程度上減少了模型引數量。
然而,對於文字生成任務而言,合成器的各個模型表現不一。與一些預訓練模型相比,合成器模型在單一性的任務上表現良好,但是訓練出的模型泛化性差,遷移能力不足。
點積自注意力本質上關注區域性領域內token資訊。基於此,You等人[4]提出用固定編碼自注意力來代替可學習的點積自注意力。該方法無需點積注意力或基於內容的注意力。實驗發現,用高斯分佈代替編碼器和解碼器中的自注意力對機器翻譯的結果幾乎沒有影響,並且使用高斯分佈可以還略微提高推理速度、節省記憶體。
由上文點積自注意力原理可知:第 個token的自注意力權重是透過查詢詞向量 和所有的鍵向量 點乘得到的。You等人提出使用以查詢詞位置 、 或 為中心的高斯分佈代替點積自注意力權重。高斯分佈從幾何含義上刻畫序列中第 個token相對於序列中所有其他token的重要程度(區域性視窗內,與查詢詞相鄰位置的重要程度較高)。公式(3)表示計算固定編碼高斯注意力公式:
固定編碼自注意力本質上是求值向量的加權平均。公式(3)中高斯分佈代替了成對token間的點積互動,且 的均值和標準差均為輸入 完全無關超參量。以查詢詞“to”為例,圖2表示每個頭中的注意力權重都是在區域性視窗內以不同token為中心的高斯分佈。

此文提出了固定編碼自注意力,它雖然不需要學習任何引數,其效能仍可以與點積自注意力相媲美。在IWSLT16 En-De[5]及其他機器翻譯資料集上,編碼器和解碼器中具有固定編碼自注意力的Transformer(HC-SA)與具有點積自注意力的Transformer(BASE)相比,獲得幾乎相等的 BLEU 分數[6,7]。其最佳驗證結果是將編碼器自注意力權重替換為以 和 為中心的高斯分佈,解碼器自注意力權重替換為以 和 為中心的高斯分佈。實驗發現,使用高斯分佈可以還略微提高推理速度、節省記憶體。

Google提出了一個無需計算點積注意力的合成器模型。在機器翻譯、語言建模和對話生成等多個任務上,合成器注意力與基於點積自注意力的Transformer(BASE)相比,效能表現優秀。在對話生成任務上,成對token間的點積互動反而降低效能。如此看來,自注意力機制中點積自注意力並非不可替代的。
You等人提出的HC-SA之所以能獲得與BASE相似的效能,我們認為前饋神經網路層(FFN)發揮了重要的作用,FFN可以學習到固定編碼自注意力的損失並補償。此外,文中實驗還證明了,當採用固定編碼交叉注意力時,機器翻譯的效能會大幅度下降。我們推測固定編碼交叉注意力的失敗可能是因為解碼器的前饋層不夠強大,無法補償固定編碼器交叉注意力的資訊損失。
本文以點積自注意力的原理出發,並以上述兩篇參考文獻以例,重新思考自注意力機制。我們發現雖然點積自注意力提供了強大的建模能力和可解釋性,但在一些任務中成對token的點積互動並不是必須的。不論是Google Research 提出的可直接合成自對齊矩陣的合成器,還是You等人提出的固定編碼自注意力,在多種任務上都表現良好。並且於模型引數量和推理速度而言,非點積自注意力也具有明顯的優勢。不得不說,無需點積或內容的自注意力是對自注意力機制的深刻認識和突破。
但合成器模型、固定編碼自注意力、點積自注意力本身沒有絕對的優劣,三者表現與具體的任務相關。為何三者在不同任務中表現存在明顯差異?它們分別適合於哪些文字任務?在計算機視覺領域是否同樣適用?其背後的原理值得進一步挖掘和探索。
[1] Vaswani, A., Shazeer, N.M., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., & Polosukhin, I. (2017). Attention is All you Need. ArXiv, abs/1706.03762.
[2] Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683, 2019.
[3] Tay, Y., Bahri, D., Metzler, D., Juan, D., Zhao, Z., & Zheng, C. (2021). Synthesizer: Rethinking Self-Attention in Transformer Models. ArXiv, abs/2005.00743.
[4] You, W., Sun, S., & Iyyer, M. (2020). Hard-Coded Gaussian Attention for Neural Machine Translation. ACL.
[5] Cettolo, M., Jan, N., Sebastian, S., Bentivogli, L., Cattoni, R., & Federico, M. (2016). The IWSLT 2016 Evaluation Campaign.
[6] Voita, E., Talbot, D., Moiseev, F., Sennrich, R., & Titov, I. (2019). Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. ACL.
[7] Michel, P., Levy, O., & Neubig, G. (2019). Are Sixteen Heads Really Better than One? NeurIPS.


壁仞科技研究院作為壁仞科技的前沿研究部門,旨在研究新型智慧計算系統的關鍵技術,重點關注新型架構,先進編譯技術和設計方法學,並將逐漸拓展研究方向,探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則,將積極投入各類產學研合作並參與開源社群的建設,為相關領域的技術進步做出自己的貢獻。


關鍵詞
矩陣
資訊
效能
自注意力機制
區域性