田淵棟新作：打開1層Transformer黑盒，注意力機制沒那麼神秘

Transformer架構已經橫掃了包括自然語言處理、電腦視覺、語音、多模態等多個領域，不過目前只是實驗效果非常驚艷，對Transformer工作原理的相關研究仍然十分有限。

其中最大謎團在於，Transformer為什麼只依靠一個「簡單的預測損失」就能從梯度訓練動態（gradient training dynamics）中湧現出高效的表徵？

最近田徑棟博士公佈了團隊的最新研究成果，以數學嚴格方式，分析了1層Transformer（一個自註意力層加一個解碼器層）在下一個token預測任務上的SGD訓練動態。

論文連結：https://arxiv.org/abs/2305.16380

#這篇論文打開了自註意力層如何組合輸入token動態過程的黑盒子，並揭示了潛在的歸納偏見的性質。

具體來說，在沒有位置編碼、長輸入序列、以及解碼器層比自註意力層學習更快的假設下，研究人員證明了自註意力就是一個判別式掃描演算法（discriminative scanning algorithm）：

從均勻分佈的注意力（uniform attention）開始，對於要預測的特定下一個token，模型逐漸關注不同的key token，而較少關注那些出現在多個next token窗口中的常見token

#對於不同的token，模型會逐漸降低注意力權重，遵循訓練集中的key token和query token之間從低到高共現的順序。

有趣的是，這個過程不會導致贏家通吃，而是由兩層學習率控制的相變而減速，最後變成（幾乎）固定的token組合，在合成和真實世界的數據上也驗證了這種動態。

田淵棟博士是Meta人工智慧研究院研究員、研究經理，圍棋AI專案負責人，其研究方向為深度增強學習及其在遊戲中的應用，以及深度學習模型的理論分析。先後於2005年及2008年取得上海交通大學本碩學位，2013年獲得美國卡內基美隆大學機器人研究所博士學位。

曾獲2013年國際電腦視覺大會（ICCV）馬爾獎提名（Marr Prize Honorable Mentions），ICML2021傑出論文榮譽提名獎。

曾在博士畢業後發布《博士五年總結》系列，從研究方向選擇、閱讀積累、時間管理、工作態度、收入和可持續的職業發展等方面對博士生涯總結心得和體會。

揭秘1層Transformer

基於Transformer架構的預訓練模型通常只包含非常簡單的監督任務，例如預測下一個單字、填空等，但卻可以為下游任務提供非常豐富的表徵，實在令人費解。

之前的工作雖然已經證明了Transformer本質上就是一個通用近似器（universal approximator），但先前常用的機器學習模型，例如kNN、核SVM、多層感知機等其實也是通用近似器，這種理論無法解釋這兩類模型在效能上的巨大差距。

研究人員認為，了解Transformer的訓練動態（training dynamics）是很重要的，也就是說，在訓練過程中，可學習參數是如何隨時間變化的。

文章首先以嚴謹數學定義的方式，形式化描述了1層無位置編碼Transformer的SGD在下一個token預測（GPT系列模型常用的訓練範式）上的訓練動態。

1層的Transformer包含一個softmax自註意力層和預測下一個token的解碼器層。

在假設序列很長，而且解碼器的學習速度比自註意力層快的情況下，證明了訓練期間自註意力的動態行為：

1. 頻率偏差Frequency Bias

##模型會逐漸關注那些與query token大量共現的key token，而對那些共現較少的token降低注意力。

2. 判別偏差Discrimitive Bias

##模型更關注那些在下一個要預測的token中唯一出現的獨特token，而對那些在多個下一個token中出現的通用token失去興趣。

這兩個特性表明，自註意力隱式地運行著一種判別式掃描（discriminative scanning）的演算法，並存在歸納偏差（inductive bias），即偏向於經常與query token共同出現的獨特的key token

此外，雖然自註意力層在訓練過程中趨向於變得更加稀疏，但正如頻率偏差所暗示的，模型因為訓練動態中的相變（phase transition），所以不會崩潰為獨熱（one hot）。

學習的最後階段並沒有收斂到任何梯度為零的鞍點，而是進入了一個注意力變化緩慢的區域（即隨時間變化的對數），並出現參數凍結和學會（learned）。

研究結果進一步表明，相變的開始是由學習率控制的：大的學習率會產生稀疏的注意力模式，而在固定的自註意力學習率下，大的解碼器學習率會導致更快的相變和密集的注意力模式。

研究人員將工作中發現的SGD動態命名為掃描（scan）和snap：

掃描階段：自註意力集中在key tokens上，即不同的、經常與下一個預測token同時出現的token；其他所有token的注意力都下降。

snap階段：注意力全中幾乎凍結，token組合固定。

這現像在簡單的真實世界資料實驗中也得到驗證，使用SGD在WikiText上訓練的1層和3層Transformer的最低自註意力層進行觀察，可以發現即使在整個訓練過程中學習率保持不變，注意力也會在訓練過程中的某一時刻凍結，並變得稀疏。

以上是田淵棟新作：打開1層Transformer黑盒，注意力機制沒那麼神秘的詳細內容。更多資訊請關注PHP中文網其他相關文章！