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

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

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

作者證明，最小化通用自回歸損失會產生在 Transformer 的前向傳遞中運行的基于輔助梯度的優化算法。這種現象最近被稱為「mesa 優化（mesa-optimization）」。此外，研究人員發現所得的 mesa 優化算法表現出上下文中的小樣本學習能力，與模型規模無關。因此，新的結果對此前大語言模型中出現的小樣本學習的原理進行了補充。

研究人員認為：Transformers 的成功基于其在前向傳遞中實現 mesa 優化算法的架構偏差：(i) 定義內部學習目標，以及 (ii) 對其進行優化。

圖 1：新假設的說明：優化自回歸 Transformer fθ 的權重 θ 會產生在模型前向傳播中實現的 mesa 優化算法。作為輸入序列 s_1, . 。。， s_t 被處理到時間步 t，Transformer (i) 創建一個由輸入 - 目標關聯對組成的內部訓練集，(ii) 通過結果數據集定義內部目標函數，用于衡量內部模型的性能 使用權重 W，(iii) 優化該目標并使用學習的模型生成未來的預測。

該研究的貢獻包括：

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

• 在初步的語言建模實驗后發現，用 mesa 層替換標準的自注意力層獲得了有希望的結果，證明了該層具有強大的上下文學習能力。

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

首先分析在簡單線性動力學上訓練的 transformer，其中每個序列由不同的 W* 生成 - 以防止跨序列記憶。在這個簡單的設置中，作者展示了 transformer 創建 mesa 數據集，然后使用預處理的 GD 優化 mesa 目標。

該研究在聚合相鄰序列元素的 token 結構上訓練深度 transformer。有趣的是，這種簡單的預處理會產生極其稀疏的權重矩陣（只有不到 1% 的權重非零），從而產生逆向工程算法。

對于單層線性自注意力，權重對應一個 GD 步驟。對于深度 transformer，可解釋性就變得困難。該研究依靠線性探測并檢查隱藏激活是否可以預測自回歸目標或預處理輸入。

有趣的是，兩種探測方法的可預測性都會隨著網絡深度的增加而逐漸提高。這一發現表明模型中隱藏著預處理的 GD。

圖 2：對經過訓練的線性自注意力層進行逆向工程。

該研究發現，在構建中使用所有自由度時，可以完美地擬合訓練層，不僅包括學習的學習率 η，還包括一組學習的初始權重 W_0。重要的是，如圖 2 所示，學得的 one-step 算法的性能仍然遠遠優于單個 mesa 層。

我們可以注意到，在簡單的權重設置下，很容易通過基礎優化發現，該層可以最優地解決此處研究的任務。該結果證明了硬編碼歸納偏差有利于 mesa 優化的優勢。

憑借對多層案例的理論見解，先分析深度線性和 softmax 僅注意 Transformer。作者根據 4 通道結構設置輸入格式，，這對應于選擇 W_0 = 0。

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

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

在圖 3 所示的深度線性自注意力 Transformer 上，我們可以看到兩個探針都可以線性解碼，解碼性能隨著序列長度和網絡深度的增加而增加。因此，基礎優化發現了一種混合算法，該算法在原始 mesa-objective Lt (W) 的基礎上逐層下降，同時改進 mesa 優化問題的條件數。這導致  mesa-objective Lt (W) 快速下降。此外可以看到性能隨著深度的增加而顯著提高。

因此可以認為自回歸 mesa-objective Lt (W) 的快速下降是通過對更好的預處理數據進行逐步（跨層）mesa 優化來實現的。

圖 3：對構建的 token 輸入進行逆向工程的多層 Transformer 訓練。

這表明，如果 transformer 在構建的 token 上進行訓練，它就會通過 mesa 優化進行預測。有趣的是，當直接給出序列元素時，transformer 會自行通過對元素進行分組來構造 token，研究團隊將其稱為「創建 mesa 數據集」。

結論

該研究表明，當在標準自回歸目標下針對序列預測任務進行訓練時，Transformer 模型能夠開發基于梯度的推理算法。因此，在多任務、元學習設置下獲得的最新結果也可以轉化到傳統的自監督 LLM 訓練設置中。

此外，該研究還發現學得的自回歸推理算法可以在無需重新訓練的情況下重新調整用途，以解決有監督的上下文學習任務，從而在單個統一框架內解釋結果。

那么，這些與上下文學習（in-context learning）有什么關系呢？該研究認為：在自回歸序列任務上訓練 transformer 后，它實現了適當的 mesa 優化，因此可以進行少樣本（few-shot）上下文學習，而無需任何微調。

該研究假設 LLM 也存在 mesa 優化，從而提高了其上下文學習能力。有趣的是，該研究還觀察到，為 LLM 有效調整 prompt 也可以帶來上下文學習能力的實質性改進。

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