In-context learning as implicit Bayesian inference

Context. GPT-3 learns a task from a few input-output examples in its prompt with no gradient update, although nothing in next-token pretraining asks for this and a prompt of concatenated unrelated examples is an unnatural, low-probability document. Xie et al. ask what property of the pretraining distribution makes this happen, assuming the model has learned that distribution exactly; Raventós et al. ask, for transformers trained on regression prompts, when the learned predictor stops being the Bayes rule for its pretraining distribution. If every pretraining document is generated by a hidden concept that persists across the document, then predicting the next token already means inferring that concept, and a prompt works by supplying evidence about it. In-context learning is then the posterior predictive of the pretraining mixture, and it still works when the prompt is out of distribution, because the format mismatch lowers the log-likelihood by a bounded amount per example while the evidence grows with the example length. A trained transformer stays Bayesian for its pretraining distribution only while that distribution has few tasks: past a task-diversity threshold it abandons the Bayes estimator for the tasks it has seen, accepts a higher training loss, and converges to the Bayes estimator for the distribution that generated them, which is what generalizes.

Setting (Xie et al.). A concept $\theta$ with prior $p(\theta)$ fixes the transition matrix of an HMM, and a pretraining document is a draw from the mixture $\int p(o_{1:T}\mid\theta)p(\theta)\,d\theta$ . A prompt is unlike a document: its examples $O_i=[x_i,y_i]$ of length $k$ are independent, each restarted from a prompt start distribution, run under the prompt's concept $\theta^*$ and followed by a delimiter. The in-context predictor is $\arg\max_yp(y\mid S_n,x_{\rm test})$ under the pretraining distribution, compared with the best predictor $\arg\max_yp_{\rm prompt}(y\mid x_{\rm test})$ , whose margin over the runner-up is $\Delta$ . The assumptions bound the log-likelihood lost to the prompt format: delimiter states emit only the delimiter, transitions into them are below $c_2$ under $\theta\ne\theta^*$ and above $c_1$ under $\theta^*$ , their start probabilities lie in $[c_3,c_4]$ , start probabilities under $\theta^*$ exceed $c_8$ , and the prompt start distribution is within $\Delta/4$ in total variation of the post-delimiter distribution. This gives two per-example bounds on that loss, $\epsilon_{\rm delim}=2\log\frac{c_2}{c_1}+\log\frac{c_4}{c_3},\qquad \epsilon_{\rm start}=\log\frac1{c_8}.$

Theorems 1 to 3. If the evidence in one example exceeds the mismatch in one example, $\sum_{j=1}^k\mathrm{KL}_j(\theta^*\|\theta)\gt\epsilon_{\rm start}+\epsilon_{\rm delim}\qquad\text{for every }\theta\ne\theta^*,$ with $\mathrm{KL}_j$ the expected KL between the $j$ -th token conditionals of the prompt distribution and of concept $\theta$ , then the in-context predictor converges to the optimal one as $n\to\infty$ , and its 0-1 risk to the Bayes risk; the condition does not involve $n$ . The proof treats $p(\theta^*)$ as a point mass, so the theorem concerns a family in which $\theta^*$ has positive prior mass, such as a finite one. When some concepts violate the condition, the limiting excess 0-1 risk is at most $g^{-1}\big(O(\gamma_{\theta^*}(\epsilon_{\rm start}+\epsilon_{\rm delim})/(k-1))\big)$ , with $\gamma_{\theta^*}$ a Fisher-information condition number and $g$ a calibration function of the logistic loss that is roughly linear on the useful range, so longer examples help like $1/k$ (Theorem 2, assuming the 0-1 and logistic risks have the same minimizers). Theorem 3 drops $\gamma_{\theta^*}$ and the continuity assumption when the test input has uniformly random length.

Step 1

The predictor is a posterior average whose weights are likelihood ratios to the true concept. Marginalize the concept and the test example's start state, and divide by $p(S_n,x_{\rm test}\mid\theta^*)$ : the weight on $\theta$ is $e^{n\,r_n(\theta)}p(\theta),\qquad r_n(\theta)=\frac1n\log\frac{p(S_n,x_{\rm test}\mid\theta)}{p(S_n,x_{\rm test}\mid\theta^*)},$ which is 1 at $\theta^*$ , so it suffices that $r_n(\theta)$ has a negative limit for every other $\theta$ . The argument resembles Bernstein-von Mises, but the standard theorem does not apply: under the model one HMM chain runs through the whole prompt, and the prompt comes from a different distribution.

Step 2

The delimiters factor the likelihood into per-example terms, up to a constant per example. The delimiter token comes only from delimiter states, so they act as approximate renewal points: condition each example on the preceding delimiter state, and the bounds $c_1,\dots,c_4$ replace that conditional by a fresh start, $r_n(\theta)\le\frac1n\sum_{i=1}^n\log\frac{p(O_i\mid\theta)}{p(O_i\mid\theta^*)}+\epsilon_{\rm delim}+O(1/n),$ where $p(O_i\mid\theta)$ is the probability of one example as if it began a fresh document.

Step 3

The average converges to a difference of KLs, which the condition makes negative. Under the prompt distribution the examples are i.i.d., so the law of large numbers gives $\mathrm{KL}(p_{\rm prompt}\|p_{\theta^*})-\mathrm{KL}(p_{\rm prompt}\|p_\theta)$ ; the first is at most $\epsilon_{\rm start}$ , since prompt and $\theta^*$ -examples differ only in the start state, and the chain rule splits the second over tokens: $\lim_nr_n(\theta)\le-\sum_{j=1}^k\mathrm{KL}_j(\theta^*\|\theta)+\epsilon_{\rm start}+\epsilon_{\rm delim}\lt0 .$ The evidence grows with $k$ because every token, inputs included, is evidence, while the mismatch is a fixed number of nats per example. This is misspecified Bayes (Berk; Kleijn and van der Vaart): the posterior goes to the concept closest in KL to the prompt distribution, and the condition makes that $\theta^*$ .

Step 4

Once the concept is selected, the test example's start state must not change the argmax. With $\theta=\theta^*$ the model still mixes $p(y\mid x_{\rm test},h,\theta^*)$ over the test example's start state $h$ , with post-delimiter weights where the optimal predictor uses the prompt start distribution. The argmax is unchanged if, for every delimiter state, the two start-state posteriors given $x_{\rm test}$ are within $\Delta/2$ in total variation, since each class probability then moves by less than $\Delta/2$ . Closeness of the start distributions themselves is not enough, because conditioning on $x_{\rm test}$ reweights start states by their likelihood: a state with tiny prior mass but high likelihood can dominate one posterior and not the other.

In the start-state step, why is a total-variation bound between

p_{\rm prompt}(h)

and the post-delimiter distribution not enough by itself?

Step 5

Without distinguishability, the bound on the summed KL still bounds the last-token KL. The posterior can stay spread over the concepts that violate the condition, so the limit predicts with a mixture over them; convexity of KL bounds its excess logistic risk by the worst single $\theta$ , and $g$ converts that to 0-1 risk. What remains is the last-token KL, knowing only that the sum over positions is at most $\epsilon_{\rm start}+\epsilon_{\rm delim}$ , and the quadratic expansion of $\mathrm{KL}_j$ with Fisher information $I_j$ gives it in two steps: $\|\theta-\theta^*\|^2\lesssim\frac{\epsilon_{\rm start}+\epsilon_{\rm delim}}{(k-1)\min_j\lambda_{\min}(I_j)}\quad\Longrightarrow\quad\mathrm{KL}_k(\theta^*\|\theta)\lesssim\frac{\gamma_{\theta^*}(\epsilon_{\rm start}+\epsilon_{\rm delim})}{k-1}.$ The condition number enters when moving from the average over positions to one position. With a uniformly random test length the test KL is that average itself, and no continuity is needed.

Step 6

On a synthetic HMM mixture, trained models learn in context because of the concept structure. On GINC (5 concepts, factorial HMMs, about 10 million pretraining tokens), GPT-2-style transformers and an LSTM reach in-context accuracy that rises with both the number of examples $n$ and their length $k$ . Pretraining on random transitions, which shows every token transition but has no latent concept, gives flat curves, and prompts from concepts outside the pretraining family are not learned; this rules out "seeing the transitions" as the explanation, though it shows only that the outputs track the posterior and does not show that the network computes it. Some results are not explained by the theory: deeper transformers with the same validation loss learn better in context, the smaller LSTM does best, and accuracy varies by 10 to 40 points across orderings of the same examples. On GPT-3, five long LAMBADA examples beat five short ones (70.7% vs 69.8%), while duplicating the short ones to equal length does not help (69.6%).

Setting (Raventós et al.). A task is $w\in\mathbb R^D$ , a sequence has $x_k\sim N(0,I_D)$ and $y_k=w^\top x_k+\varepsilon_k$ , $\varepsilon_k\sim N(0,\sigma^2)$ , $k\le K$ , and a causal GPT-2 is trained to predict each $y_k$ from the preceding pairs. Pretraining tasks are uniform on a fixed set of $M$ vectors drawn once from $N(0,I_D)$ ; $M$ is the task diversity. Two posterior means bracket the behavior: dMMSE, the Bayes estimator for the discrete prior on the seen tasks, $\hat w^{\rm dMMSE}=\sum_{i=1}^M\frac{\exp\!\big(-\frac1{2\sigma^2}\sum_{j}(y_j-w^{(i)\top}x_j)^2\big)}{\sum_l\exp\!\big(-\frac1{2\sigma^2}\sum_{j}(y_j-w^{(l)\top}x_j)^2\big)}\,w^{(i)},$ which minimizes the pretraining loss, and ridge $(X^\top X+\sigma^2I)^{-1}X^\top y$ , the Bayes estimator for $N(0,I_D)$ , optimal on new tasks. Defaults are $D=8$ , $K=16$ , $\sigma^2=0.25$ .

Step 7

More data per task pushes the model in opposite directions on the two sides of a threshold. For $M$ up to about $2^6$ the transformer matches dMMSE, optimal on seen tasks and poor on new ones. Beyond that it leaves dMMSE on seen tasks, so it no longer minimizes its own objective, and moves toward ridge on new tasks much faster than dMMSE itself does. A quick simulation (my check, not in the paper) shows how far dMMSE is from ridge: on new tasks dMMSE gives 0.91, 0.63 and 0.47 per dimension at $M=2^6,2^{10},2^{15}$ , against 0.36 to 0.38 for ridge, so at the threshold the Bayes estimator for the pretraining prior is still about 25% worse than ridge, and a transformer that matches ridge there is not computing it. Raising sequences per task, by batch size or by steps with the same result, pulls the model toward dMMSE for $M\le2^{14}$ and toward ridge for $M\ge2^{15}$ . The authors extrapolate a sharp phase transition as sequences per task go to infinity, a limit inferred from finite-size trends and not observed.

Step 8

The switch is not underfitting or blurred memory, and it does not come from covering the task space. Below the threshold the distance to ridge on new tasks is minimized at an early-stopping time $t^*\propto M^{0.47}$ ; above it the distance falls monotonically through four times the base training, so the shape of the learning curve changes (a very late turn toward dMMSE cannot be excluded). A dMMSE with each point mass blurred into a Gaussian, the width tuned for new tasks, is still beaten from $M=2^{13}$ on, so storing the tasks imprecisely does not explain it. The threshold grows roughly linearly with dimension (about $2^{14},2^{15},2^{16}$ at $D=8,16,32$ ), whereas dMMSE approaches ridge only when the seen tasks cover the sphere where Gaussian tasks concentrate, which takes exponentially many; three values of the dimension are few for fitting a scaling law. Weight decay lowers the threshold, more capacity raises it for a small model, and with Laplace tasks the model again moves past a threshold to the Laplace-prior Bayes predictor (itself estimated by a model trained with unlimited tasks).

Technical notes Xie et al. assume only the prior bound

\mathrm{TV}\lt\Delta/4

, and their proof compares unnormalized vectors with the margin

\Delta

, whereas the unnormalized optimal vector has margin only

\Delta\,p_{\rm prompt}(x_{\rm test})

; since distinguishability is vacuous when

\Theta=\{\theta^*\}

, an HMM with one delimiter state breaks Theorem 1's conclusion as stated (start distributions 0.015 apart in TV, posteriors given

x_{\rm test}

0.34 apart, argmax flipped), and the paper's own simulations use a start distribution it says may not meet its assumption anyway; and Raventós et al.'s appendix "

\sigma^2

" values 0.5, 0.707, 0.866 and 1 equal

\sqrt{D/32}

and are standard deviations.

Consequences and limits. Xie et al. make "in-context learning is Bayesian inference" checkable: document-level latent structure in pretraining is enough for the exact next-token conditional to select the prompt's concept, as long as each example's evidence exceeds the log-likelihood lost to its format. Raventós et al. show where that reading stops holding for trained models: the choice is between two Bayes estimators, the one for the tasks actually seen, which the loss rewards, and the one for the population they came from, which generalizes. A statistician would recognize the difference between using the empirical distribution of past tasks as a prior and using a smooth prior fitted to it; the transformer does the second once $M$ is large, so data scale alone is not enough and diversity has to grow with it. Xie et al.'s theory is limited to the pretraining distribution itself, so it says nothing about how a finite network approximates it or why scale matters at fixed loss; it is asymptotic in $n$ with no rate, needs $\theta^*$ in the pretraining family, and is tested on five synthetic concepts. Raventós et al.'s evidence is empirical, for one model family on linear regression, and identifies neither the mechanism nor the implicit bias that selects ridge.

Terms

Concept $\theta$: document-level latent variable: the HMM transition matrix shared by a whole document
Prompt distribution: independent examples from $\theta^*$ , each restarted from $p_{\rm prompt}$ , joined by delimiters
In-context predictor: $f_n(x)=\arg\max_yp(y\mid S_n,x)$ under the pretraining mixture (posterior predictive)
$\epsilon_{\rm delim},\ \epsilon_{\rm start}$: per-example log-likelihood lost to the format at the delimiter and at the start state
Distinguishability: $\sum_{j\le k}\mathrm{KL}_j\gt\epsilon_{\rm start}+\epsilon_{\rm delim}$ for all $\theta\ne\theta^*$
Task diversity $M$: number of distinct regression vectors in pretraining, drawn once from $N(0,I_D)$
dMMSE: posterior mean under the discrete prior on seen tasks: softmax of residual likelihoods over $w^{(i)}$
Ridge, $\lambda=\sigma^2$: posterior mean under $N(0,I_D)$ ; optimal on new tasks
Threshold $M^*$: diversity where more sequences per task switch from pulling toward dMMSE to toward ridge; $2^{14}$ - $2^{15}$ at $D=8$
Dimension scaling: $M^*\approx2^{14},2^{15},2^{16}$ at $D=8,16,32$ : linear in $D$ rather than exponential

Quiz

Why do longer examples (larger $k$ ) help in-context learning in this theory?
In the continuity version (Theorem 2), why does the Fisher condition number $\gamma_{\theta^*}$ appear?
Why does the smoothed-dMMSE comparison matter?
Why is a threshold that grows linearly in $D$ evidence that the transformer is not approaching ridge the way dMMSE does?

Sources

Sang Michael Xie et al. An Explanation of In-context Learning as Implicit Bayesian Inference. ICLR, 2022.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression. NeurIPS, 2023.