本文作者李宏康，美國倫斯勒理工大學電氣、計算機與系統工程系在讀博士生，本科畢業于中國科學技術大學。研究方向包括深度學習理論，大語言模型理論，統計機器學習等等。目前已在 ICLR/ICML/Neurips 等 AI 頂會發表多篇論文。

上下文學習 (in-context learning, 簡寫為 ICL) 已經在很多 LLM 有關的應用中展現了強大的能力，但是對其理論的分析仍然比較有限。人們依然試圖理解為什么基于 Transformer 架構的 LLM 可以展現出 ICL 的能力。

近期，一個來自美國倫斯勒理工大學和 IBM 研究院的團隊從優化和泛化理論的角度分析了帶有非線性注意力模塊 (attention) 和多層感知機 (MLP) 的 Transformer 的 ICL 能力。他們特別從理論端證明了單層 Transformer 首先在 attention 層根據 query 選擇一些上下文示例，然后在 MLP 層根據標簽嵌入進行預測的 ICL 機制。該文章已收錄在 ICML 2024。

• 論文題目：How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?
• 論文地址：https://arxiv.org/pdf/2402.15607

背景介紹

上下文學習 in context learning (ICL)

上下文學習 (ICL) 是一種新的學習范式，在大語言模型 (LLM) 中非常流行。它具體是指在測試查詢 (testing query)

前添加 N 個測試樣本 testing examples (上下文)，即測試輸入

和測試輸出

的組合，從而構成一個 testing prompt:

，作為模型的輸入以引導模型作出正確的推斷。這種方式不同于經典的對預訓練模型進行微調的方式，它不需要改變模型的權重，從而更加的高效。

ICL 理論工作的進展

近期的很多理論工作都是基于 [1] 所提出的研究框架，即人們可以直接使用 prompt 的格式來對 Transformer 進行訓練 (這一步也可以理解為在模擬一種簡化的 LLM 預訓練模式)，從而使得模型具有 ICL 能力。已有的理論工作聚焦于模型的表達能力 (expressive power) 的角度 [2]。他們發現，人們能夠找到一個有著 “完美” 的參數的 Transformer 可以通過前向運算執行 ICL，甚至隱含地執行梯度下降等經典機器學習算法。但是這些工作無法回答為什么 Transformer 可以被訓練成這樣 “完美” 的，具有 ICL 能力的參數。因此，還有一些工作試圖從 Transformer 的訓練或泛化的角度理解 ICL 機制 [3,4]。不過，受制于分析 Transformer 結構的復雜性，這些工作目前止步于研究線性回歸任務，而所考慮的模型通常會略去 Transformer 中的非線形部分。

本文從優化和泛化理論的角度分析了帶有非線性 attention 和 MLP 的 Transformer 的 ICL 能力和機制：

• 基于一個簡化的分類模型，本文具體量化了數據的特征如何影響了一層單頭 Transformer 的域內 (in-domain) 和域外 (out-of-domain, OOD) 的 ICL 泛化能力。
• 本文進一步闡釋了 ICL 是如何通過被訓練的 Transformer 來實現了。
• 基于被訓練的 Transformer 的特點，本文還分析了在 ICL 推斷的時候使用基于幅值的模型剪枝 (magnitude-based pruning) 的可行性。

理論部分

問題描述

本文考慮一個二分類問題，即將

通過一個任務

映射到

。為了解決這樣的一個問題，本文構建了 prompt 來進行學習。這里的 prompt 被表示為：

訓練網絡為一個單層單頭 Transformer：

預訓練過程是求解一個對所有訓練任務的經驗風險最小化 (empirical risk minimization)。損失函數使用的是適合二分類問題的 Hinge loss，訓練算法是隨機梯度下降。

本文定義了兩種 ICL 泛化的情況。一個是 in-domain 的，即泛化的時候測試數據的分布和訓練數據一樣，注意這個情況里面測試任務不必和訓練任務一樣，即這里已經考慮了對未見任務 (unseen task) 的泛化。另一個是 out-of-domain 的，即測試、訓練數據分布不一樣。

本文還涉及了在 ICL 推斷的時候進行 magnitude-based pruning 的分析，這里的剪枝方式是指對于訓練得到的中的各個神經元，根據其幅值大小，進行從小到大的刪除。

對數據和任務的構建

這一部分請參考原文的 Section 3.2，這里只做一個概述。本文的理論分析是基于最近比較火熱的 feature learning 路線，即通常將數據假設為可分（通常是正交）的 pattern，從而推導出基于不同 pattern 的梯度變化。本文首先定義了一組 in-domain-relevant (IDR) pattern 用于決定 in-domain 任務的分類，和一組與任務無關的 in-domain-irrelevant (IDI) pattern，這些 pattern 之間互相正交。IDR pattern 有

個，IDI pattern 有

個。一個

被表示為一個 IDR pattern 和一個 IDI pattern 的和。一個 in-domain 任務就被定義為基于某兩個 IDR pattern 的分類問題。

類似地，本文通過定義 out-of-domain-relevant (ODR) pattern 和 out-of-domain-irrelevant (ODI) pattern，可以刻畫 OOD 泛化時候的數據和任務。

本文對 prompt 的表示可以用下圖的例子來闡述，其中

是 IDR pattern，

是 IDI pattern。這里在做的任務是基于 x 中的

做分類，如果是

那么其標簽為 + 1，對應于 +q，如果是

那么其標簽為 - 1，對應于 -q。α，α' 分別被定義為訓練和測試 prompt 中跟 query 的 IDR/ODR pattern 一樣的上下文示例。下圖中的例子里面，

。

理論結果

首先，對于 in-domain 的情況，本文先給了一個 condition 3.2 來規定訓練任務需要滿足的條件，即訓練任務需要覆蓋所有的 IDR pattern 和標簽。然后 in-domain 的結果如下：

這里表明：1，訓練任務的數量只需要在全部任務中占比達到滿足 condition 3.2 的小比例，我們就可以對 unseen task 實現很好的泛化；2，跟當前任務相關的 IDR pattern 在 prompt 中的比例越高，就可以以更少的訓練數據，訓練迭代次數，以及更短的 training/testing prompt 實現理想的泛化。

接下來是 out-of-domain 泛化的結果。

這里說明，如果 ODR pattern 是 IDR pattern 的線性組合且系數和大于 1，那么此時 OOD ICL 泛化可以達到理想的效果。這個結果給出了在 ICL 的框架下，好的 OOD 泛化所需要的訓練和測試數據之間的內在聯系。該定理也通過 GPT-2 的實驗得到了驗證。如下圖所示，當 (12) 中的系數和

大于 1 的時候，OOD 分類可以達到理想的結果。與此同時，當

，即 prompt 中和分類任務相關的 ODR/IDR pattern 比例越高的時候，所需要的 context 長度越小。

然后，本文給出了帶有 magnitude-based pruning 的 ICL 泛化結果。

這個結果表明，首先，訓練得到的

中有一部分（常數比例）神經元的幅值很小，而剩下的相對比較大（公式 14）。當我們只枝剪小神經元的時候，對泛化結果基本沒有影響，而當枝剪比例增加到要剪大神經元的時候，泛化誤差會隨之顯著變大（公式 15，16）。以下實驗驗證了定理 3.7。下圖 A 中淺藍色的豎線表示訓練得到的

呈現出了公式 14 的結果。而對小神經元進行枝剪不會使泛化變差，這個結果符合理論。圖 B 反映出當 prompt 中和任務相關的上下文越多的時候，我們可以允許更大的枝剪比例以達到相同的泛化性能。

ICL 機制

通過對預訓練過程的刻畫，本文得到了單層單頭非線性 Transformer 做 ICL 的內在機制，這一部分在原文的 Section 4。該過程可以用下圖表示。

簡而言之，attention 層會選擇和 query 的 ODR/IDR pattern 一樣的上下文，賦予它們幾乎全部 attention 權重，然后 MLP 層會重點根據 attention 層輸出中的標簽嵌入來作出最后的分類。

總結

本文講解了在 ICL 當中，非線性 Transformer 的訓練機制，以及對于新任務和分布偏移數據的泛化能力。理論結果對于設計 prompt 選擇算法和 LLM 剪枝算法有一定實際意義。

本文轉自 機器之心 ，作者：機器之心

原文鏈接:??https://mp.weixin.qq.com/s/SJQiIp1W5kwWSVJaOXA9yA??