2025-11-24T16:43:16.687108

In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning

Wakayama, Suzuki

This paper develops a finite-sample statistical theory for in-context learning (ICL), analyzed within a meta-learning framework that accommodates mixtures of diverse task types. We introduce a principled risk decomposition that separates the total ICL risk into two orthogonal components: Bayes Gap and Posterior Variance. The Bayes Gap quantifies how well the trained model approximates the Bayes-optimal in-context predictor. For a uniform-attention Transformer, we derive a non-asymptotic upper bound on this gap, which explicitly clarifies the dependence on the number of pretraining prompts and their context length. The Posterior Variance is a model-independent risk representing the intrinsic task uncertainty. Our key finding is that this term is determined solely by the difficulty of the true underlying task, while the uncertainty arising from the task mixture vanishes exponentially fast with only a few in-context examples. Together, these results provide a unified view of ICL: the Transformer selects the optimal meta-algorithm during pretraining and rapidly converges to the optimal algorithm for the true task at test time.

academic

文脈内学習は証明可能なベイズ推論である：メタ学習の一般化理論

基本情報

論文ID: 2510.10981
タイトル: In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning
著者: Tomoya Wakayama (理化学研究所 AIP), Taiji Suzuki (東京大学、理化学研究所 AIP)
分類: stat.ML cs.LG
発表日: 2025年10月13日 (arXiv プレプリント)
論文リンク: https://arxiv.org/abs/2510.10981v1

要約

本論文は、複数の異なるタスク型の混合に対応するメタ学習フレームワーク内で、文脈内学習(ICL)の有限サンプル統計理論を確立している。論文は原則的なリスク分解を導入し、総ICLリスクを2つの直交成分に分解する：ベイズギャップ(Bayes Gap)と事後分散(Posterior Variance)。ベイズギャップは、訓練済みモデルがベイズ最適文脈予測器にどの程度近いかを定量化する。均一注意Transformerに対して、論文はこのギャップの非漸近上界を導出し、事前訓練プロンプト数と文脈長への依存関係を明確に示している。事後分散は、固有のタスク不確実性を表すモデル非依存的なリスクである。重要な発見は、この項が真の潜在タスクの難度によってのみ決定され、タスク混合からの不確実性は少数の文脈サンプルで指数関数的に急速に消失することである。

研究背景と動機

問題背景

GPT-3以来、大規模言語モデルは顕著な文脈内学習能力を示している。すなわち、少数の入出力例のみから新しいタスクに適応でき、パラメータ更新を必要としない。この現象は様々なデータセットとタスク形式で普遍的に見られ、現代的なLLMワークフローの中核である。

研究動機

理論の欠落：ICLが暗黙的なベイズ推論の形式として広く認識されているにもかかわらず、既存の理論はICLとベイズ推論の理論的関係を十分に活用していない
実践的必要性：現代的なLLM展開は共通の制約に直面している——推論時のプロンプトが短く、上流の事前訓練は異質なタスク型をカバーし、有限サンプルの予測誤差の具体的な分析が必要である
理論的空白：既存の理論は(i)事前訓練規模Nとプロンプト長pを共同に結合し、(ii)異質なタスク型の混合に対応できる統計理論が不足している

既存手法の限界

初期の理論は特定のアーキテクチャと設定における情報論的分析または非パラメトリック率に主に焦点を当てていた
pとNの共同効果を完全に捉えることができていない
混合タスク設定におけるICL動作の理論的説明が不足している

核心的貢献

原則的なリスク分解：ICLリスクの直交分解を提案：ICL risk = Bayes Gap + Posterior Variance
非漸近上界：均一注意Transformerに対してベイズギャップの非漸近上界を提供し、事前訓練プロンプト数Nと文脈長pの結合依存関係を明確にする： $E[R_{BG}(M_{\hat{\theta}})] \lesssim m^{-2\alpha/d_{eff}} + \frac{m}{pN} + \frac{1}{N}$
タスク識別理論：タスク混合において、事後分布がタスクインデックス上で指数関数的に急速に真のタスクに集中し、ICLが真のタスクの最適アルゴリズムに急速に収束することを証明
分布シフト安定性：入力分布シフト下での安定性を特性化し、ベイズギャップが分布間のWasserstein距離に比例して増加することを証明

方法の詳細

タスク定義

論文はT個の異なるタスク型の有限混合に対応するメタ学習フレームワークを考察する：

プロンプト生成プロセス：

タスク型をサンプリング： $I \sim \text{Categorical}(\alpha)$
$I=i$ が与えられた場合、タスク関数をサンプリング： $f \sim P_{F_i}$
$k=1,\ldots,p+1$ $k = 1, \dots, p + 1$ に対して：
- 入力をサンプリング： $x_k \overset{i.i.d.}{\sim} P_X$
- 出力を生成： $y_k = f(x_k) + \varepsilon_k$
長さpのプロンプトを形成： $P = (x_1,y_1,\ldots,x_p,y_p,x_{p+1})$

モデルアーキテクチャ

均一注意Transformer： $M_\theta(P^k) := \rho_\theta\left(\frac{1}{k}\sum_{i=1}^k \phi_\theta(x_i,y_i), x_{k+1}\right)$

ここで：

特徴エンコーダ $\phi_\theta: U \to \Delta_{m-1}$ ：深さ $D_\phi$ のフィードフォワードReLUネットワーク、その後に再正規化層
デコーダ $\rho_\theta: \Delta_{m-1} \times C \to \mathbb{R}$ ：深さ $D_\rho$ のフィードフォワードReLUネットワーク

ベイズ最適予測器

ICLリスク最小化はベイズリスク最小化と等価であり、最適予測器は事後平均である： $M_{\text{Bayes}}(P^k) := E_{I\sim P_{I|D^k}} E_{f\sim P_{F_I|D^k}}[f(x_{k+1})]$

技術的革新点

置換不変性理論基礎：ベイズ予測器の置換不変性を証明し、均一注意アーキテクチャに理論的支援を提供
逐次学習理論の応用：プロンプト内のp個の文脈サンプルを処理するために逐次学習理論を利用し、従来の学習理論とN個のメタ訓練プロンプトを結合
最適輸送近似理論：ソフトヒストグラムに基づく分割単位を構築してプロンプトをエンコードし、離散1-Wasserstein距離上のMcShane拡張を通じてベイズ予測器を近似

実験設定

理論分析フレームワーク

論文は主に理論分析を提供し、以下の設定を採用している：

仮定条件：

仮定1：有界タスク関数 $|f(x)| \leq B_f$
仮定2：有界入力と条件独立性 $\|x\|_2 \leq B_X$

ネットワーク規模：

特徴エンコーダ： $S(\phi_\theta) \leq C_\phi m^{1/d_{eff}}$
デコーダ： $S(\rho_\theta) \leq C_\rho m^{1/2}$

評価指標

ICLリスクは以下のように定義される： $R(M) = \frac{1}{p}\sum_{k=1}^p E_{I,f,D^k,x_{k+1}}\left[(f(x_{k+1}) - M(P^k))^2\right]$

実験結果

主要な理論的結果

定理1（リスク分解）： $R(M) = R_{BG}(M) + R_{PV}$ ここで：

ベイズギャップ： $R_{BG}(M) := \frac{1}{p}\sum_{k=1}^p E[(M(P^k) - M_{\text{Bayes}}(P^k))^2]$
事後分散： $R_{PV} := \frac{1}{p}\sum_{k=1}^p E[\text{Var}_{f\sim P(f|D^k)}(f(x_{k+1}))]$

定理2（ベイズギャップ上界）： Hölder条件下で、均一注意Transformerに対して： $E[R_{BG}(M_{\hat{\theta}})] \lesssim m^{-2\alpha/d_{eff}} + \frac{m}{pN}\text{polylog}(pN) + \frac{1}{N}\text{polylog}(pN)$

$m^* \asymp (pN)^{d_{eff}/(d_{eff}+2\alpha)}$ を選択すると： $E[R_{BG}(M_{\hat{\theta}})] \lesssim (pN)^{-2\alpha/(d_{eff}+2\alpha)} + N^{-1}$

定理3（事後分散分析）：対数尤度比条件下で： $E_{D^k,x|I=i^*}[\text{Var}_{f|D^k}\{f(x)\}] \leq \inf_M \sup_{f\in F_{i^*}} E[(f(x_{k+1}) - M(P^k))^2|f] + 5B_f^2\left(\frac{1-\alpha_{i^*}}{\alpha_{i^*}}e^{-D_{\min}k/2} + (T-1)e^{-Ck}\right)$

主要な発見

最適メタアルゴリズム選択：Transformerは事前訓練期間中に最適メタアルゴリズムを選択し、速度 $\propto m/(pN)$ はpとNの共同効果を明確にする
指数関数的タスク識別：混合タスク設定において、タスク事後は真のタスクインデックスに指数関数的に急速に集中し、不可約誤差は真のタスクのミニマックスリスクに収束
分布シフト安定性：入力分布シフト下で、ベイズギャップはWasserstein距離に比例して増加し、事後分散はターゲット領域内在の特性を保持