ChatGPT 等基于 Transformer 的大語言模型具備極強(qiáng)的在上下文中學(xué)習(xí)（In-Context Learning，ICL）的能力：輸入少量示例樣本，即能夠正確回答同類問題。如何理解這種 ICL 能力？

本文作者實(shí)驗(yàn)發(fā)現(xiàn)并證明一種 ICL 的新機(jī)制：自動算法選擇，可以允許單一 Transformer 模型在不同輸入數(shù)據(jù)上選擇執(zhí)行完全不同的，適合該數(shù)據(jù)的學(xué)習(xí)算法，類似統(tǒng)計與機(jī)器學(xué)習(xí)專家能夠現(xiàn)實(shí)完成的工作。基于量化的 Transformer 構(gòu)造，文章一并給出 Transformer 實(shí)現(xiàn) ICL 的一套全面的統(tǒng)計理論，包含近似精度，預(yù)測表現(xiàn)，以及預(yù)訓(xùn)練的樣本復(fù)雜度。

圖片

論文地址：https://arxiv.org/abs/2306.04637

Transformer 能在 ICL 中完成機(jī)器學(xué)習(xí)任務(wù)

ChatGPT 等基于 Transformer 的大模型可以根據(jù)輸入的文本，自上下文中學(xué)習(xí)。如何系統(tǒng)地理解這種能力？NeurIPS 2022 的一篇論文（Garg et al. 2022）考察了 Transformer 從上下文中進(jìn)行機(jī)器學(xué)習(xí)任務(wù)的能力。

將 N 個訓(xùn)練樣本與 1 個測試樣本 (x_1, y_1, …, x_N, y_N, x_{N+1}) 作為一個序列輸入 Transformer，要求 Transformer 輸出 y_{N+1}。這些樣本來自簡單的統(tǒng)計模型，例如線性模型，但每個序列由不同的模型參數(shù)（w_\star）生成。Transformer 如果想總是正確地預(yù)測 y_{N+1}，那么就需要從訓(xùn)練樣本中學(xué)習(xí)真正的參數(shù) w_\star，并利用其進(jìn)行預(yù)測。

Garg et al. 發(fā)現(xiàn)，訓(xùn)練好的 Transformer 總是能夠精準(zhǔn)地預(yù)測 y_{N+1}，并且預(yù)測表現(xiàn)能夠媲美該數(shù)據(jù)上的最優(yōu)算法。例如線性模型上，Transformer 的預(yù)測效果可以媲美最小二乘法（Least Squares），稀疏線性模型上媲美 Lasso，決策樹上能超過 Gradient Boosting。

圖片

Figure 1: Transformer 能在 ICL 中完成機(jī)器學(xué)習(xí)任務(wù) (Garg et al. 2022)

Transformer 雖然在各個任務(wù)中實(shí)現(xiàn)最優(yōu)算法，然而這些算法都只是標(biāo)準(zhǔn)的機(jī)器學(xué)習(xí)算法，遠(yuǎn)不能完全解釋 ChatGPT 等大模型強(qiáng)大的 ICL 能力。還存在更強(qiáng)的機(jī)制嗎？

自動算法選擇

現(xiàn)實(shí)生活中，統(tǒng)計學(xué)家與機(jī)器學(xué)習(xí)專家會如何分析數(shù)據(jù)？給定一個數(shù)據(jù)集，統(tǒng)計學(xué)家會先確定數(shù)據(jù)的結(jié)構(gòu)，規(guī)模等，然后根據(jù)數(shù)據(jù)的特點(diǎn)選擇最適合的算法。如果不確定哪個算法合適，則會同時嘗試多個算法，然后利用驗(yàn)證集（validation split）或交叉驗(yàn)證（cross-validation）等選擇表現(xiàn)最好的算法。

本文作者發(fā)現(xiàn)，Transformer 也能夠進(jìn)行類似的自動算法選擇。自動算法選擇允許一個單獨(dú)的 Transformer 模型，在不同的 ICL 問題上選擇不同的算法，類似統(tǒng)計學(xué)家可以現(xiàn)實(shí)完成的工作。

作者給出兩種一般的算法選擇機(jī)制，從理論上證明 Transformer 模型可以實(shí)現(xiàn)這兩個機(jī)制，并且實(shí)驗(yàn)上驗(yàn)證了 Transformer 能夠近似實(shí)現(xiàn)這兩種機(jī)制，達(dá)到了比單一機(jī)器學(xué)習(xí)算法更強(qiáng)的效果。

機(jī)制 1：用驗(yàn)證集做算法選擇

在這一機(jī)制中，Transformer 先將輸入數(shù)據(jù)分為訓(xùn)練集和驗(yàn)證集。接下來在訓(xùn)練集上同時執(zhí)行 K 個算法，然后在驗(yàn)證集上測試 K 個算法的表現(xiàn)，最終用表現(xiàn)最好的算法 k_star 給出預(yù)測。

Figure 2: 用驗(yàn)證集做算法選擇（右）及實(shí)例（左）。

應(yīng)用這一機(jī)制，Transformer 可以完成一大類算法選擇。如下圖當(dāng)中，通過恰當(dāng)?shù)念A(yù)訓(xùn)練，Transformer 可同時實(shí)現(xiàn)在兩個帶不同正則化的 ridge regression 算法，并對具體數(shù)據(jù)分布實(shí)現(xiàn)較優(yōu)的那個算法。進(jìn)一步，對這一任務(wù)，我們在理論上也能夠證明 Transformer 能夠近似整個任務(wù)的 Bayes 最優(yōu)表現(xiàn)。

Figure 3: 單一 Transformer 可以在兩個帶不同噪音的線性模型中同時接近最優(yōu)。在每個模型中，Transformer 的預(yù)測都接近該模型上 Bayes 最優(yōu)算法（帶不同正則化的 ridge regression）。

機(jī)制 2：提前對數(shù)據(jù)分布進(jìn)行檢驗(yàn)

在這一機(jī)制中，Transformer 通過提前檢驗(yàn)數(shù)據(jù)分布（如計算一些統(tǒng)計量），來決定恰當(dāng)?shù)乃惴ā＠缭谙聢D當(dāng)中，單一的 Transformer 可以在回歸問題上實(shí)現(xiàn)回歸算法（如線性回歸），在分類問題上實(shí)現(xiàn)分類算法（如 Logistic Regression）。

Figure 4: 提前對數(shù)據(jù)分布進(jìn)行檢驗(yàn)（右）及實(shí)例（左）。

Figure 5: 單一 Transformer 可以同時在回歸問題和分類問題上接近最優(yōu)表現(xiàn)：在回歸問題上表現(xiàn)接近最小二乘法，在分類問題上接近 Logistic Regression。

理論框架

除了自動算法選擇，本文的另一大貢獻(xiàn)是給出了 Transformer 進(jìn)行 ICL 的一套完整的統(tǒng)計學(xué)習(xí)理論，涵蓋

• Transformer 如何實(shí)現(xiàn)各種機(jī)器學(xué)習(xí)算法的具體構(gòu)造，如最小二乘法，Lasso，ridge regression，解廣義線性模型的凸優(yōu)化算法，神經(jīng)網(wǎng)絡(luò)上的梯度下降法等；
• 對上述 Transformer 大小（層數(shù)、attention head 個數(shù)、權(quán)重矩陣的范數(shù)）的精確界；
• 上述 Transformer 在標(biāo)準(zhǔn)統(tǒng)計假設(shè)下，在 ICL 中的預(yù)測表現(xiàn)；
• Transformer 實(shí)現(xiàn)自動算法選擇時類似的保證；
• 通過預(yù)訓(xùn)練 Transformer 達(dá)到上述效果的樣本復(fù)雜度。

這套理論給出了 Transformer 進(jìn)行 ICL 的一整套分析框架。作者相信這一框架可以推廣到一大類相關(guān)問題當(dāng)中，給出類似的理論保證。

結(jié)語

本文從理論和實(shí)驗(yàn)上發(fā)現(xiàn) Transformer 模型在 ICL 中能夠進(jìn)行自動算法選擇，并給出了一整套進(jìn)行 ICL 的理論框架。

基于本文的結(jié)論還有很大的探索空間，例如其它進(jìn)行 ICL 或自動算法選擇的機(jī)制；在 ICL 中逼近 Bayes 最優(yōu)表現(xiàn)的其它機(jī)制；預(yù)訓(xùn)練的 Transformer 如何實(shí)現(xiàn)算法選擇的內(nèi)部機(jī)理；對其它 ICL 任務(wù)的分析。作者相信，對這些問題的進(jìn)一步探索，能對大模型有更多有趣的發(fā)現(xiàn)。

作者簡介

本文作者 Yu Bai 現(xiàn)任 Salesforce Research 資深研究科學(xué)家。Fan Chen 本科畢業(yè)于北京大學(xué)，即將博士入學(xué)麻省理工大學(xué)。Huan Wang、Caiming Xiong 分別現(xiàn)任 Salesforce Research 研究主管及副總裁。Song Mei 現(xiàn)任 加州大學(xué)伯克利統(tǒng)計系助理教授。