為什么 transformer 性能這么好？它給眾多大語言模型帶來的上下文學(xué)習(xí) (In-Context Learning) 能力是從何而來？在人工智能領(lǐng)域里，transformer 已成為深度學(xué)習(xí)中的主導(dǎo)模型，但人們對(duì)于它卓越性能的理論基礎(chǔ)卻一直研究不足。

最近，來自 Google AI、蘇黎世聯(lián)邦理工學(xué)院、Google DeepMind 研究人員的新研究嘗試為我們揭開謎底。在新研究中，他們對(duì) transformer 進(jìn)行了逆向工程，尋找到了一些優(yōu)化方法。論文《Uncovering mesa-optimization algorithms in Transformers》：

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

作者證明，最小化通用自回歸損失會(huì)產(chǎn)生在 Transformer 的前向傳遞中運(yùn)行的基于輔助梯度的優(yōu)化算法。這種現(xiàn)象最近被稱為「mesa 優(yōu)化（mesa-optimization）」。此外，研究人員發(fā)現(xiàn)所得的 mesa 優(yōu)化算法表現(xiàn)出上下文中的小樣本學(xué)習(xí)能力，與模型規(guī)模無關(guān)。因此，新的結(jié)果對(duì)此前大語言模型中出現(xiàn)的小樣本學(xué)習(xí)的原理進(jìn)行了補(bǔ)充。

研究人員認(rèn)為：Transformers 的成功基于其在前向傳遞中實(shí)現(xiàn) mesa 優(yōu)化算法的架構(gòu)偏差：(i) 定義內(nèi)部學(xué)習(xí)目標(biāo)，以及 (ii) 對(duì)其進(jìn)行優(yōu)化。

圖 1：新假設(shè)的說明：優(yōu)化自回歸 Transformer fθ 的權(quán)重 θ 會(huì)產(chǎn)生在模型前向傳播中實(shí)現(xiàn)的 mesa 優(yōu)化算法。作為輸入序列 s_1, . 。。， s_t 被處理到時(shí)間步 t，Transformer (i) 創(chuàng)建一個(gè)由輸入 - 目標(biāo)關(guān)聯(lián)對(duì)組成的內(nèi)部訓(xùn)練集，(ii) 通過結(jié)果數(shù)據(jù)集定義內(nèi)部目標(biāo)函數(shù)，用于衡量內(nèi)部模型的性能 使用權(quán)重 W，(iii) 優(yōu)化該目標(biāo)并使用學(xué)習(xí)的模型生成未來的預(yù)測(cè)。

該研究的貢獻(xiàn)包括：

• 概括了 von Oswald 等人的理論，并展示了從理論上，Transformers 是如何通過使用基于梯度的方法優(yōu)化內(nèi)部構(gòu)建的目標(biāo)來自回歸預(yù)測(cè)序列下一個(gè)元素的。
• 通過實(shí)驗(yàn)對(duì)在簡單序列建模任務(wù)上訓(xùn)練的 Transformer 進(jìn)行了逆向工程，并發(fā)現(xiàn)強(qiáng)有力的證據(jù)表明它們的前向傳遞實(shí)現(xiàn)了兩步算法：(i) 早期自注意力層通過分組和復(fù)制標(biāo)記構(gòu)建內(nèi)部訓(xùn)練數(shù)據(jù)集，因此隱式地構(gòu)建內(nèi)部訓(xùn)練數(shù)據(jù)集。定義內(nèi)部目標(biāo)函數(shù)，(ii) 更深層次優(yōu)化這些目標(biāo)以生成預(yù)測(cè)。
• 與 LLM 類似，實(shí)驗(yàn)表明簡單的自回歸訓(xùn)練模型也可以成為上下文學(xué)習(xí)者，而即時(shí)調(diào)整對(duì)于改善 LLM 的上下文學(xué)習(xí)至關(guān)重要，也可以提高特定環(huán)境中的表現(xiàn)。
• 受發(fā)現(xiàn)注意力層試圖隱式優(yōu)化內(nèi)部目標(biāo)函數(shù)的啟發(fā)，作者引入了 mesa 層，這是一種新型注意力層，可以有效地解決最小二乘優(yōu)化問題，而不是僅采取單個(gè)梯度步驟來實(shí)現(xiàn)最優(yōu)。實(shí)驗(yàn)證明單個(gè) mesa 層在簡單的順序任務(wù)上優(yōu)于深度線性和 softmax 自注意力 Transformer，同時(shí)提供更多的可解釋性。

• 在初步的語言建模實(shí)驗(yàn)后發(fā)現(xiàn)，用 mesa 層替換標(biāo)準(zhǔn)的自注意力層獲得了有希望的結(jié)果，證明了該層具有強(qiáng)大的上下文學(xué)習(xí)能力。

基于最近人們的工作表明，經(jīng)過明確訓(xùn)練來解決上下文中的小樣本任務(wù)的 transformer 可以實(shí)現(xiàn)梯度下降（GD）算法。在這里，作者展示了這些結(jié)果可以推廣到自回歸序列建模 —— 這是訓(xùn)練 LLM 的典型方法。

首先分析在簡單線性動(dòng)力學(xué)上訓(xùn)練的 transformer，其中每個(gè)序列由不同的 W* 生成 - 以防止跨序列記憶。在這個(gè)簡單的設(shè)置中，作者展示了 transformer 創(chuàng)建 mesa 數(shù)據(jù)集，然后使用預(yù)處理的 GD 優(yōu)化 mesa 目標(biāo)。

該研究在聚合相鄰序列元素的 token 結(jié)構(gòu)上訓(xùn)練深度 transformer。有趣的是，這種簡單的預(yù)處理會(huì)產(chǎn)生極其稀疏的權(quán)重矩陣（只有不到 1% 的權(quán)重非零），從而產(chǎn)生逆向工程算法。

對(duì)于單層線性自注意力，權(quán)重對(duì)應(yīng)一個(gè) GD 步驟。對(duì)于深度 transformer，可解釋性就變得困難。該研究依靠線性探測(cè)并檢查隱藏激活是否可以預(yù)測(cè)自回歸目標(biāo)或預(yù)處理輸入。

有趣的是，兩種探測(cè)方法的可預(yù)測(cè)性都會(huì)隨著網(wǎng)絡(luò)深度的增加而逐漸提高。這一發(fā)現(xiàn)表明模型中隱藏著預(yù)處理的 GD。

圖 2：對(duì)經(jīng)過訓(xùn)練的線性自注意力層進(jìn)行逆向工程。

該研究發(fā)現(xiàn)，在構(gòu)建中使用所有自由度時(shí)，可以完美地?cái)M合訓(xùn)練層，不僅包括學(xué)習(xí)的學(xué)習(xí)率 η，還包括一組學(xué)習(xí)的初始權(quán)重 W_0。重要的是，如圖 2 所示，學(xué)得的 one-step 算法的性能仍然遠(yuǎn)遠(yuǎn)優(yōu)于單個(gè) mesa 層。

我們可以注意到，在簡單的權(quán)重設(shè)置下，很容易通過基礎(chǔ)優(yōu)化發(fā)現(xiàn)，該層可以最優(yōu)地解決此處研究的任務(wù)。該結(jié)果證明了硬編碼歸納偏差有利于 mesa 優(yōu)化的優(yōu)勢(shì)。

憑借對(duì)多層案例的理論見解，先分析深度線性和 softmax 僅注意 Transformer。作者根據(jù) 4 通道結(jié)構(gòu)設(shè)置輸入格式，，這對(duì)應(yīng)于選擇 W_0 = 0。

與單層模型一樣，作者在訓(xùn)練模型的權(quán)重中看到了清晰的結(jié)構(gòu)。作為第一個(gè)逆向工程分析，該研究利用這個(gè)結(jié)構(gòu)并構(gòu)建一個(gè)算法（RevAlg-d，其中 d 表示層數(shù)），每個(gè)層頭包含 16 個(gè)參數(shù)（而不是 3200 個(gè)）。作者發(fā)現(xiàn)這種壓縮但復(fù)雜的表達(dá)式可以描述經(jīng)過訓(xùn)練的模型。特別是，它允許以幾乎無損的方式在實(shí)際 Transformer 和 RevAlg-d 權(quán)重之間進(jìn)行插值。

雖然 RevAlg-d 表達(dá)式解釋了具有少量自由參數(shù)的經(jīng)過訓(xùn)練的多層 Transformer，但很難將其解釋為 mesa 優(yōu)化算法。因此，作者采用線性回歸探測(cè)分析（Alain & Bengio，2017；Akyürek et al.，2023）來尋找假設(shè)的 mesa 優(yōu)化算法的特征。

在圖 3 所示的深度線性自注意力 Transformer 上，我們可以看到兩個(gè)探針都可以線性解碼，解碼性能隨著序列長度和網(wǎng)絡(luò)深度的增加而增加。因此，基礎(chǔ)優(yōu)化發(fā)現(xiàn)了一種混合算法，該算法在原始 mesa-objective Lt (W) 的基礎(chǔ)上逐層下降，同時(shí)改進(jìn) mesa 優(yōu)化問題的條件數(shù)。這導(dǎo)致  mesa-objective Lt (W) 快速下降。此外可以看到性能隨著深度的增加而顯著提高。

因此可以認(rèn)為自回歸 mesa-objective Lt (W) 的快速下降是通過對(duì)更好的預(yù)處理數(shù)據(jù)進(jìn)行逐步（跨層）mesa 優(yōu)化來實(shí)現(xiàn)的。

圖 3：對(duì)構(gòu)建的 token 輸入進(jìn)行逆向工程的多層 Transformer 訓(xùn)練。

這表明，如果 transformer 在構(gòu)建的 token 上進(jìn)行訓(xùn)練，它就會(huì)通過 mesa 優(yōu)化進(jìn)行預(yù)測(cè)。有趣的是，當(dāng)直接給出序列元素時(shí)，transformer 會(huì)自行通過對(duì)元素進(jìn)行分組來構(gòu)造 token，研究團(tuán)隊(duì)將其稱為「創(chuàng)建 mesa 數(shù)據(jù)集」。

結(jié)論

該研究表明，當(dāng)在標(biāo)準(zhǔn)自回歸目標(biāo)下針對(duì)序列預(yù)測(cè)任務(wù)進(jìn)行訓(xùn)練時(shí)，Transformer 模型能夠開發(fā)基于梯度的推理算法。因此，在多任務(wù)、元學(xué)習(xí)設(shè)置下獲得的最新結(jié)果也可以轉(zhuǎn)化到傳統(tǒng)的自監(jiān)督 LLM 訓(xùn)練設(shè)置中。

此外，該研究還發(fā)現(xiàn)學(xué)得的自回歸推理算法可以在無需重新訓(xùn)練的情況下重新調(diào)整用途，以解決有監(jiān)督的上下文學(xué)習(xí)任務(wù)，從而在單個(gè)統(tǒng)一框架內(nèi)解釋結(jié)果。

那么，這些與上下文學(xué)習(xí)（in-context learning）有什么關(guān)系呢？該研究認(rèn)為：在自回歸序列任務(wù)上訓(xùn)練 transformer 后，它實(shí)現(xiàn)了適當(dāng)?shù)?mesa 優(yōu)化，因此可以進(jìn)行少樣本（few-shot）上下文學(xué)習(xí)，而無需任何微調(diào)。

該研究假設(shè) LLM 也存在 mesa 優(yōu)化，從而提高了其上下文學(xué)習(xí)能力。有趣的是，該研究還觀察到，為 LLM 有效調(diào)整 prompt 也可以帶來上下文學(xué)習(xí)能力的實(shí)質(zhì)性改進(jìn)。

感興趣的讀者可以閱讀論文原文，了解更多研究內(nèi)容。