基于Transformer的大模型是如何運(yùn)行的？Meta從全局和上下文學(xué)習(xí)揭秘

本文旨在更好地理解基于 Transformer 的大型語(yǔ)言模型（LLM）的內(nèi)部機(jī)制，以提高它們的可靠性和可解釋性。

隨著大型語(yǔ)言模型（LLM）在使用和部署方面的不斷增加，打開黑箱并了解它們的內(nèi)部工作原理變得越來(lái)越重要。更好地理解這些模型是如何做出決策的，這對(duì)改進(jìn)模型和減輕其故障（如幻覺(jué)或推理錯(cuò)誤）至關(guān)重要。眾所周知，最近 LLM 成功的一個(gè)重要因素是它們能夠從上下文中學(xué)習(xí)和推理。LLM 對(duì)這些上下文的學(xué)習(xí)能力通常歸功于 Transformer 架構(gòu)，特別是自注意力塊的使用，其能夠仔細(xì)選擇輸入序列，進(jìn)而推理出可信的下一個(gè) token。此外，預(yù)測(cè)可能需要全局知識(shí)，如語(yǔ)法規(guī)則或一般事實(shí)，這些可能不會(huì)出現(xiàn)在上下文中，需要存儲(chǔ)在模型中。我們不禁會(huì)疑問(wèn)，為什么基于 Transformer 的模型非常擅長(zhǎng)使用它們的上下文來(lái)預(yù)測(cè)新的 token，這種能力是如何在訓(xùn)練中產(chǎn)生的？帶著這些問(wèn)題，來(lái)自 Meta AI 的研究者進(jìn)行了深入的研究。他們通過(guò)研究合成設(shè)置下 Transformer 的學(xué)習(xí)機(jī)制，揭示了其全局和上下文學(xué)習(xí)的平衡，并將權(quán)重矩陣解釋為聯(lián)想記憶，為理解和優(yōu)化 Transformer 提供了基礎(chǔ)。

論文地址：https://arxiv.org/pdf/2306.00802.pdf首先要了解的是在訓(xùn)練過(guò)程中 Transformer 是如何發(fā)現(xiàn)這些能力的。為此，該研究引入了一個(gè)合成數(shù)據(jù)集，該數(shù)據(jù)集由二元語(yǔ)言模型生成的序列組成。然后，模型需要依靠上下文學(xué)習(xí)來(lái)對(duì)特定的二元序列進(jìn)行良好的預(yù)測(cè)，而全局二元可以根據(jù)當(dāng)前 token 的全局統(tǒng)計(jì)數(shù)據(jù)進(jìn)行猜測(cè)。雖然單層的 Transformer 無(wú)法可靠地預(yù)測(cè)上下文二元，但該研究發(fā)現(xiàn)通過(guò)開發(fā)感應(yīng)頭（induction head）機(jī)制的雙層 Transformer 取得了成功，即擁有兩個(gè)注意力頭的 circuit，其允許 Transformer 從上下文 [???, a, b,???, a] 中預(yù)測(cè) b，并且在 Transformer 語(yǔ)言模型中似乎無(wú)處不在。這種感應(yīng)頭（induction head）機(jī)制在 Transformer 語(yǔ)言模型中是普遍存在的，并且取得了成功。更進(jìn)一步的，為了更好的了解上下文機(jī)制是怎樣出現(xiàn)在訓(xùn)練過(guò)程中的，該研究在隨機(jī)初始化時(shí)凍結(jié)了一些層（包括嵌入和值矩陣）來(lái)進(jìn)一步簡(jiǎn)化模型架構(gòu)。這樣一來(lái)研究重點(diǎn)轉(zhuǎn)移到注意力和前饋機(jī)制，同時(shí)避免了學(xué)習(xí)表征的困難。與此同時(shí)，這種簡(jiǎn)化還為單個(gè)權(quán)重矩陣引入了一個(gè)自然模型作為聯(lián)想記憶。自然模型可以通過(guò)它們的外積存儲(chǔ)輸入 - 輸出或鍵 - 值對(duì)嵌入。隨機(jī)高維嵌入由于其接近正交性而特別適合這種觀點(diǎn)。總結(jié)而言，該研究的貢獻(xiàn)可概括為：

• 本文引入了一種新的合成設(shè)置來(lái)研究全局和上下文學(xué)習(xí)：序列遵循二元語(yǔ)言模型，其中一些二元在序列中變化，而另一些不會(huì)。

• 本文將 Transformer 的權(quán)重矩陣視為學(xué)習(xí)存儲(chǔ)特定嵌入對(duì)的聯(lián)想記憶，并以此為任務(wù)推導(dǎo)出一個(gè)簡(jiǎn)化但更可解釋的模型。

• 本文對(duì)訓(xùn)練動(dòng)態(tài)進(jìn)行了細(xì)致的實(shí)證研究：首先學(xué)習(xí)全局二元，然后以自上而下的方式學(xué)習(xí)適當(dāng)?shù)挠洃?，形成感?yīng)頭。

• 本文給出了訓(xùn)練動(dòng)力學(xué)的理論見解，展示了如何通過(guò)在噪聲輸入中找到信號(hào)，在種群損失上進(jìn)行一些自上而下的梯度步驟來(lái)恢復(fù)所需的聯(lián)想記憶。

方法介紹接著該研究介紹了合成數(shù)據(jù)設(shè)置，這樣能夠仔細(xì)研究感應(yīng)頭機(jī)制在訓(xùn)練過(guò)程中的發(fā)展以及 Transformer 如何學(xué)習(xí)利用上下文信息的。雙元數(shù)據(jù)模型：模型序列由一個(gè)通用的雙元語(yǔ)言模型（即馬爾可夫鏈）組成，每個(gè)序列的生成方式如下：下圖 2 可視化了測(cè)試序列上的注意力圖，這表明該模型已經(jīng)學(xué)習(xí)了感應(yīng)頭機(jī)制。接著該研究介紹了 Transformer 聯(lián)想記憶觀點(diǎn)：因?yàn)閹缀跽坏那度?，?quán)重矩陣表現(xiàn)為聯(lián)想記憶，將成對(duì)的嵌入存儲(chǔ)為其外積的加權(quán)和。研究引入了一個(gè)具有固定隨機(jī)嵌入的簡(jiǎn)化 Transformer 模型，將用這種想法產(chǎn)生對(duì)學(xué)習(xí)動(dòng)力學(xué)的精確理解。此外，該研究提出了一個(gè)有用的觀點(diǎn)，將 Transformer 中的模型權(quán)重視為高維嵌入向量的聯(lián)想記憶。感應(yīng)頭機(jī)制可以通過(guò)以下外積矩陣作為記憶來(lái)獲得，而其他所有權(quán)重則固定為隨機(jī)初始化狀態(tài)： 實(shí)驗(yàn)圖 3 研究了在迭代 300 次之前凍結(jié)不同層對(duì)訓(xùn)練動(dòng)態(tài)的影響。全局 vs 上下文學(xué)習(xí)。從圖 4（左 / 右）可以看出，當(dāng)聯(lián)合訓(xùn)練所有層時(shí)，全局二元統(tǒng)計(jì)的學(xué)習(xí)速度往往比感應(yīng)頭更快，這可以從早期迭代中的 loss 和 KL 的快速下降中看出。此外，從圖 4（左）中看到，數(shù)據(jù)分布的變化會(huì)對(duì)上下文機(jī)制的學(xué)習(xí)速度產(chǎn)生重大影響。該研究觀察到以下情況可能會(huì)使上下文學(xué)習(xí)減慢：(i) 較少數(shù)量的觸發(fā)器 K， (ii) 僅使用少有的固定觸發(fā)器，以及 (iii) 使用隨機(jī)觸發(fā)器而不是固定觸發(fā)器。該研究還在圖 4（中間）中顯示，在訓(xùn)練時(shí)將輸出 token 分布更改為二元分布會(huì)降低準(zhǔn)確率，這表明，使用更多樣化的訓(xùn)練分布可以產(chǎn)生具有更好泛化準(zhǔn)確率的模型，并且只需少量的額外訓(xùn)練成本。更多研究?jī)?nèi)容，請(qǐng)參考原論文。