田淵棟新作：打開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的最低自注意力層進行觀察，可以發現即使在整個訓練過程中學習率保持不變，注意力也會在訓練過程中的某一時刻凍結，并變得稀疏。
