Transformer在下一個token預(yù)測任務(wù)上的SGD訓(xùn)練動態(tài)

【導(dǎo)讀】AI理論再進(jìn)一步，破解ChatGPT指日可待？

Transformer架構(gòu)已經(jīng)橫掃了包括自然語言處理、計算機(jī)視覺、語音、多模態(tài)等多個領(lǐng)域，不過目前只是實驗效果非常驚艷，對Transformer工作原理的相關(guān)研究仍然十分有限。

其中最大謎團(tuán)在于，Transformer為什么僅依靠一個「簡單的預(yù)測損失」就能從梯度訓(xùn)練動態(tài)（gradient training dynamics）中涌現(xiàn)出高效的表征？

最近田淵棟博士公布了團(tuán)隊的最新研究成果，以數(shù)學(xué)嚴(yán)格方式，分析了1層Transformer（一個自注意力層加一個解碼器層）在下一個token預(yù)測任務(wù)上的SGD訓(xùn)練動態(tài)。

論文鏈接：https://arxiv.org/abs/2305.16380

這篇論文打開了自注意力層如何組合輸入token動態(tài)過程的黑盒子，并揭示了潛在的歸納偏見的性質(zhì)。

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

從均勻分布的注意力（uniform attention）開始，對于要預(yù)測的特定下一個token，模型逐漸關(guān)注不同的key token，而較少關(guān)注那些出現(xiàn)在多個next token窗口中的常見token

對于不同的token，模型會逐漸降低注意力權(quán)重，遵循訓(xùn)練集中的key token和query token之間從低到高共現(xiàn)的順序。

有趣的是，這個過程不會導(dǎo)致贏家通吃，而是由兩層學(xué)習(xí)率控制的相變而減速，最后變成（幾乎）固定的token組合，在合成和真實世界的數(shù)據(jù)上也驗證了這種動態(tài)。

田淵棟博士是Meta人工智能研究院研究員、研究經(jīng)理，圍棋AI項目負(fù)責(zé)人，其研究方向為深度增強(qiáng)學(xué)習(xí)及其在游戲中的應(yīng)用，以及深度學(xué)習(xí)模型的理論分析。先后于2005年及2008年獲得上海交通大學(xué)本碩學(xué)位，2013年獲得美國卡耐基梅隆大學(xué)機(jī)器人研究所博士學(xué)位。

曾獲得2013年國際計算機(jī)視覺大會（ICCV）馬爾獎提名（Marr Prize Honorable Mentions），ICML2021杰出論文榮譽提名獎。

曾在博士畢業(yè)后發(fā)布《博士五年總結(jié)》系列，從研究方向選擇、閱讀積累、時間管理、工作態(tài)度、收入和可持續(xù)的職業(yè)發(fā)展等方面對博士生涯總結(jié)心得和體會。

揭秘1層Transformer

基于Transformer架構(gòu)的預(yù)訓(xùn)練模型通常只包括非常簡單的監(jiān)督任務(wù)，比如預(yù)測下一個單詞、填空等，但卻可以為下游任務(wù)提供非常豐富的表征，實在是令人費解。

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

研究人員認(rèn)為，了解Transformer的訓(xùn)練動態(tài)（training dynamics）是很重要的，也就是說，在訓(xùn)練過程中，可學(xué)習(xí)參數(shù)是如何隨時間變化的。

文章首先以嚴(yán)謹(jǐn)數(shù)學(xué)定義的方式，形式化描述了1層無位置編碼Transformer的SGD在下一個token預(yù)測（GPT系列模型常用的訓(xùn)練范式）上的訓(xùn)練動態(tài)。

1層的Transformer包含一個softmax自注意力層和預(yù)測下一個token的解碼器層。

在假設(shè)序列很長，而且解碼器的學(xué)習(xí)速度比自注意力層快的情況下，證明了訓(xùn)練期間自注意力的動態(tài)行為：

1. 頻率偏差Frequency Bias

模型會逐漸關(guān)注那些與query token大量共現(xiàn)的key token，而對那些共現(xiàn)較少的token降低注意力。

2. 判別偏差Discrimitive Bias

模型更關(guān)注那些在下一個要預(yù)測的token中唯一出現(xiàn)的獨特token，而對那些在多個下一個token中出現(xiàn)的通用token失去興趣。

這兩個特性表明，自注意力隱式地運行著一種判別式掃描（discriminative scanning）的算法，并存在歸納偏差（inductive bias），即偏向于經(jīng)常與query token共同出現(xiàn)的獨特的key token

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

學(xué)習(xí)的最后階段并沒有收斂到任何梯度為零的鞍點，而是進(jìn)入了一個注意力變化緩慢的區(qū)域（即隨時間變化的對數(shù)），并出現(xiàn)參數(shù)凍結(jié)和學(xué)會（learned）。

研究結(jié)果進(jìn)一步表明，相變的開始是由學(xué)習(xí)率控制的：大的學(xué)習(xí)率會產(chǎn)生稀疏的注意力模式，而在固定的自注意力學(xué)習(xí)率下，大的解碼器學(xué)習(xí)率會導(dǎo)致更快的相變和密集的注意力模式。

研究人員將工作中發(fā)現(xiàn)的SGD動態(tài)命名為掃描（scan）和snap：

掃描階段：自注意力集中在key tokens上，即不同的、經(jīng)常與下一個預(yù)測token同時出現(xiàn)的token；其他所有token的注意力都下降。

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

這一現(xiàn)象在簡單的真實世界數(shù)據(jù)實驗中也得到驗證，使用SGD在WikiText上訓(xùn)練的1層和3層Transformer的最低自注意力層進(jìn)行觀察，可以發(fā)現(xiàn)即使在整個訓(xùn)練過程中學(xué)習(xí)率保持不變，注意力也會在訓(xùn)練過程中的某一時刻凍結(jié)，并變得稀疏。
責(zé)任編輯：彭菁