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

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

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

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

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

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

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

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

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

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

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

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

揭秘1層Transformer

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

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

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

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

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

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

1. 頻率偏差Frequency Bias

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

2. 判別偏差Discrimitive Bias

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

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

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

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

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

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

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

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

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