Paper notes — In-Context Learning as Implicit Bayesian Inference

Xie, S. M., Raghunathan, A., Liang, P., & Ma, T. (2022). An Explanation of In-Context Learning as Implicit Bayesian Inference. ICLR 2022.

The problem the paper sets up. In-context learning has no obvious mechanism: no gradient step separates seeing the demonstrations from producing the answer. The paper asks under what pretraining conditions this behavior should be expected to emerge from a plain next-token-prediction objective, rather than treating it as an unexplained empirical curiosity.

Formal construction. The authors construct a synthetic pretraining distribution generated by a mixture of Hidden Markov Models (HMMs), where each HMM instance corresponds to a distinct latent “document concept.” A pretraining document is generated by first sampling a concept, then sampling a long sequence of tokens from that concept’s HMM — so within one document, the concept stays fixed and structures the whole sequence (this is meant to mimic how real documents have a consistent topic/style throughout). They prove that a Transformer trained to predict the next token under this generative process is, in the large-pretraining-data limit, implicitly learning to perform Bayesian inference over which latent concept generated the current document, given only the tokens observed so far.

Why this predicts in-context learning. At test time, an ICL prompt (a concatenation of demonstrations) looks to the pretrained model like a partial document. Each demonstration is additional evidence about which latent concept is active; conditioning on more demonstrations narrows the posterior over concepts, and the model’s next-token prediction under that sharpened posterior is what looks, from the outside, like “learning the task from examples.” Critically, no parameter update is needed for this — the update is Bayesian and happens entirely in the forward pass, over the model’s belief about which concept applies.

Empirical support. The theoretical argument is paired with experiments on a synthetic dataset the authors call GINC (Generative IN-Context learning dataset), built directly from an HMM mixture matching the theory’s assumptions. They show accuracy on the synthetic ICL task increases with the number of demonstrations and is fairly robust to demonstration ordering — both predicted directly by the Bayesian account — and observe qualitatively similar demonstration-count scaling trends in real pretrained LMs (GPT-2, GPT-3) on natural ICL tasks, though the natural-language results are presented as suggestive rather than as proof that real LMs satisfy the paper’s formal assumptions exactly.

Why it matters for uncertainty in ICL. If the object the model is inferring is a latent concept, then “the model is uncertain” has a specific meaning: the posterior over concepts is not concentrated on one option. That gives a principled way to ask whether an ICL failure comes from the demonstrations being ambiguous about which concept applies (a property of the evidence — aleatoric) or from the model’s posterior being poorly calibrated even given clear evidence (a property of the model — epistemic). It’s the theoretical justification for treating ICL uncertainty as uncertainty over a latent variable in the first place.