Source-linked AI summary
An Explanation of In-context Learning as Implicit Bayesian Inference
Sang Michael Xie, Aditi Raghunathan, Percy Liang, Tengyu Ma
TL;DR
The paper studies how in-context learning can emerge even though prompts concatenate independent examples unlike natural pretraining documents. It models pretraining as a mixture of HMMs with latent concepts, proves when Bayesian-style concept inference succeeds under distribution mismatch, and validates the account on synthetic GINC experiments showing in-context learning and related real-world phenomena.
Problem
Existing language models perform in-context learning without explicit training for it, but the mechanism is unclear because real pretraining data are messy and prompts differ from natural sequences.
Method
The paper models pretraining documents as a mixture of HMMs with latent concepts, analyzes prompt-conditioned Bayesian inference under distribution mismatch, and evaluates the theory using the synthetic GINC dataset.
Results
In-context learning occurs when the model infers a shared latent concept across prompt examples; GINC experiments reproduce improved performance with more examples and longer examples in Transformers and LSTMs, plus order sensitivity.
Takeaways & Limitations
The paper frames in-context learning as implicit Bayesian inference and offers a small-scale synthetic setting for studying large-scale in-context phenomena.
Takeaways & Limitations
The analysis leaves bridging fixed-length training examples with randomly sized test examples to future work.
Abstract
from arXiv · showhide
Large language models (LMs) such as GPT-3 have the surprising ability to do in-context learning, where the model learns to do a downstream task simply by conditioning on a prompt consisting of input-output examples. The LM learns from these examples without being explicitly pretrained to learn. Thus, it is unclear what enables in-context learning. In this paper, we study how in-context learning can emerge when pretraining documents have long-range coherence. Here, the LM must infer a latent document-level concept to generate coherent next tokens during pretraining. At test time, in-context learning occurs when the LM also infers a shared latent concept between examples in a prompt. We prove when this occurs despite a distribution mismatch between prompts and pretraining data in a setting where the pretraining distribution is a mixture of HMMs. In contrast to messy large-scale datasets used to train LMs capable of in-context learning, we generate a small-scale synthetic dataset (GINC) where Transformers and LSTMs both exhibit in-context learning. Beyond the theory, experiments on GINC exhibit large-scale real-world phenomena including improved in-context performance with model scaling (despite the same pretraining loss), sensitivity to example order, and instances where zero-shot is better than few-shot in-context learning.
1 Introduction
The paper asks how in-context learning can arise despite prompts differing from natural pretraining sequences. It proposes that long-range coherence lets language models infer latent concepts, enabling Bayesian-like task inference from examples.
- The motivating example uses wiki-bio structure, where name-to-nationality examples reveal a document-level concept shared across otherwise independent prompt examples.
- In-context learning lets language models infer a downstream task from input-output examples without explicit pretraining for that capability.
- The paper introduces a pretraining distribution in which each document is generated from a latent concept parameterizing an HMM.The model must infer the concept across sentences to generate coherent continuations.
- At inference time, in-context learning occurs when the model infers a shared concept across the examples in an otherwise unnatural prompt.The paper frames prediction as marginalizing over concepts and selecting the prompt concept as evidence accumulates.
- The theoretical analysis shows that in-context prediction can be asymptotically optimal despite prompt-pretraining distribution mismatch when per-example signal exceeds mismatch error.It also shows that prediction error decreases as example length increases.
2 In-context learning setting
The setting models pretraining documents and prompts with latent concepts and HMM-generated token sequences. Prompts concatenate independent examples sharing one concept, creating a distribution mismatch that the analysis controls through delimiter, regularity, and shift assumptions.
- Pretraining distribution: Each pretraining document samples a latent concept θ and then generates a token sequence whose distribution is defined by an HMM.The concept determines the HMM transition probability matrix over hidden states.
- Prompt distribution: A prompt concatenates n independent training examples and one test input, all conditioned on a shared prompt concept θ∗.Each example contains an input sequence followed by an output token, with delimiter tokens separating examples.
- Distribution mismatch: Concatenating independent examples creates low-probability transitions under the pretraining distribution, producing a mismatch that the model must overcome.The paper illustrates this with transitions such as “German” followed by “Mahatma Gandhi” in a name-to-nationality prompt.
- In-context predictor and task: The predictor chooses the most likely output token conditioned on the prompt and test input, and its performance is measured by expected 0-1 error.
- Assumptions: The framework assumes delimiter hidden states emit the delimiter deterministically and that delimiter probabilities constrain concept identification.
- Assumptions: The example length k is fixed for simplicity, while extending the analysis to variable-length examples is left for future work.
- Assumptions: The prompt start distribution is assumed close in total variation to hidden-state transition distributions under θ∗, with the bound scaled by the output margin Δ.
- Assumptions: The analysis assumes the prompt concept belongs to the concept family and imposes lower bounds on transitions, start states, and token emissions.
3 Theoretical analysis
The analysis explains in-context learning as implicit Bayesian inference: enough prompt evidence causes the predictor to select the shared latent concept despite prompt–pretraining distribution mismatch. Under distinguishability, this yields asymptotically optimal prediction; without it, error still decreases with example length under additional continuity conditions.
- Bayesian inference and optimality: As the number of examples grows, Bayesian marginalization suppresses all concepts except a distinguishable prompt concept, making the in-context predictor asymptotically optimal.The proof analyzes likelihood ratios and shows the prompt concept is selected when the distinguishability condition holds.
- Bayesian inference and optimality: The distinguishability condition requires each competing concept’s KL signal to exceed the error caused by prompt–pretraining distribution mismatch.The mismatch contributes O(1) error terms, while token-level KL contributions grow with example length k.
- Bayesian inference and optimality: Increasing example length k strengthens distinguishability because more input and input-output tokens provide evidence about the latent concept.The paper explicitly attributes in-context signal to tokens in both the inputs and the input-output mapping.
- Non-distinguishable concepts: When distinguishability fails, the expected 0-1 error still decreases with k under a second-order KL continuity assumption and is lower for well-conditioned Fisher information.The stated bound gives roughly O(1/k) excess-risk decay, with the rate depending on the worst-case Fisher-information condition number.
- Scope and limitations: The evaluation averages prediction error over tokens 2 through k and leaves the mismatch between fixed-length training examples and randomly sized test examples for future work.This is an explicit scope boundary of the non-distinguishable-case analysis.
4 Simulations
Simulations on synthetic GINC show that Transformers and LSTMs exhibit in-context learning, with accuracy shaped by prompt examples, model scale, architecture, and example ordering.
- Main result: Transformers and LSTMs trained on GINC improve in-context accuracy as the number and length of prompt examples increase.This verifies the theoretical prediction across three vocabulary sizes and averages results over five pretraining runs.
- Ablations on the latent concept structure: Removing the mixture-of-concepts structure or using random pretraining data causes in-context learning to fail.The ablations indicate that simply exposing models to diverse token transitions is insufficient.
- Extrapolation to unseen concepts: In-context learning also fails when prompts are generated from concepts excluded from the pretraining family.GINC permits direct testing on five randomly selected unseen concepts.
- Effect of model size and architecture: 81% to 85%: in-context accuracy rises from 12- to 16-layer Transformers despite identical 1.33 validation loss for vocabulary size 50.LSTMs consistently outperform Transformers on GINC despite having fewer parameters.
- Sensitivity to example ordering: 10–40% difference: permutations of the same four examples produce substantial variation in in-context accuracy.The experiment evaluates all 24 orderings for each of ten example sets.
- Zero-shot is sometimes better than few-shot: Zero-shot performance can exceed few-shot performance in some GINC settings, especially initially under lower-entropy transition probabilities.The authors hypothesize that the prompt structure is distracting, with accuracy recovering as more examples are added.
5 Discussion and related work
The discussion situates the work as a small-scale, controllable study of Bayesian in-context learning and distinguishes it from related inference and meta-learning approaches.
- Learning via Bayesian inference and extrapolation: The Bernstein-von Mises theorem does not directly apply because prompt examples are dependent and sampled from a different distribution than pretraining observations.Future work includes sharper asymptotic results under misspecification and extrapolation.
- Bridging the mismatch between pretraining and prompting: Prompt-format finetuning and prompt-template optimization are presented as approaches for reducing the mismatch between pretraining and prompting distributions.The cited prior work reports improved zero-shot or few-shot performance under such interventions.
- Meta-learning: Unlike meta-learning, in-context learning in this account emerges from language-model pretraining rather than explicit training to learn from examples.The comparison concerns the training objective, not whether both systems can process example sequences.
- Studying large-scale phenomena at a small scale: GINC enables small-scale study of large-scale phenomena because its latent concepts and HMM state complexity are directly controllable.Synthetic control also allows experiments on properties such as unseen concepts.
6 Conclusion
The conclusion frames in-context learning as implicit Bayesian inference over latent concepts and identifies mixture-of-HMM pretraining as a setting where it occurs.
- 6 Conclusion: The pretrained LM implicitly infers a latent concept while making predictions, providing a Bayesian interpretation of in-context learning.The conclusion presents this as the paper’s central framing.
- 6 Conclusion: In-context learning occurs when the pretraining distribution is a mixture of HMMs.This is the paper’s stated theoretical setting for the phenomenon.
- Prompt distribution details: The prompt consists of independent training examples and a test input that share a latent prompt concept.The construction samples hidden segments, delimiters, and observed tokens for the examples before forming the prompt sequence.
- Prompt distribution details: Each training example uses the first k−1 observed tokens as input and the kth token as output, with delimiters separating examples.The test input is formed similarly but contains k−1 elements.
B Propositions for Theorem 1
The propositions establish probability bounds needed to analyze how the prompt concept is selected and how the in-context predictor converges. Under the stated assumptions, the averaged predictive distribution and prompt predictive distribution have the same argmax.
- Probability bounds: Proposition 1 bounds delimiter-hidden-state probabilities above and below under the target concept.These bounds are used as direct inputs to later probability-ratio arguments.
- Probability bounds: Proposition 2 lower bounds the probability of each example under the target concept across future hidden states.The bound supplies a uniform positive lower bound for example probabilities.
- Predictive distributions: Lemma 1 shows that conditioning on the target concept and prompt preserves the argmax of the prompt predictive distribution.The result follows when the predictive-distribution difference is within half the label margin.
- Predictive distributions: The proof combines delimiter-state and example-probability bounds to control predictive-distribution differences.The argument uses matrix representations for transitions from delimiter states and test-output probabilities.
D Proof of Theorem 1
The proof shows that distinguishability makes the target prompt concept dominate marginalization as examples accumulate, while calibration converts predictive-distribution bounds into classification-error bounds.
- Concept selection: For every incorrect concept, the likelihood ratio converges to a negative constant, whereas the ratio equals 1 for the target concept.Marginalization therefore selects the desired prompt concept when the distinguishability condition holds.
- Predictor convergence: Under the distribution-shift assumption, the limiting in-context predictor has the same argmax as the prompt predictive distribution.The proof combines the concept-selection limit with the margin condition and Lemma 1.
- Concept selection: The prompt-start distribution and delimiter transitions contribute error terms that are eventually overpowered by accumulated observations of the target concept.The proof expresses the relevant bound using KL divergences and notes that increasing the number of examples strengthens the target-concept contribution.
- Non-distinguishable concepts: When the distinguishability condition fails, the limiting predictor can mix predictive distributions over concepts that are not distinguishable from the target.The analysis bounds the resulting KL divergence and uses convexity to control the mixture.
- Error bounds: Calibration converts a multiclass logistic-risk bound into a bound on the limiting predictor’s 0-1 loss.The argument has zero approximation error because the target concept belongs to the model’s concept class.
E.2 Proof of Theorem 3
The proof extends the KL-based argument to intermediate positions in variable-length test examples. Randomizing the test-example length allows the predictor to receive credit for fitting those intermediate tokens.
- KL control: Condition 1 bounds the sum of KL divergences across positions within a k-length input, not only at the final position.This supplies the multi-position control needed for the theorem’s error analysis.
- Theorem 3: The theorem follows by combining the KL bound with Lemma 2 when the test-example length is uniformly random from 2 to k.The proof explicitly plugs the bound into Lemma 2.
F.1 GINC dataset
GINC is a synthetic mixture-of-HMM dataset in which hidden entities and properties generate tokens, with the property transition matrix defining the latent concept. Experiments show that increasing observable vocabulary generally improves in-context accuracy, including across several Transformer depths.
- Dataset construction: GINC pretraining documents come from a mixture of factorial HMMs whose hidden states pair entities with properties.Entities and properties evolve as independent Markov chains, and emissions are produced from their combined hidden state.
- Dataset construction: The property transition matrix is the concept parameter, while the entity transition matrix is shared across mixture components.The dataset uses five concepts and a shared memory matrix to map hidden states to observed tokens.
- Prompt generation: Prompts concatenate independently generated examples conditioned on a shared concept, with 0–64 training examples and example lengths 3, 5, 8, or 10.The target token is chosen as the most likely prompt-distribution output to remove intrinsic sampling error.
- Dataset details: The dataset contains 1000 training documents and 100 validation documents, while prompts are generated for each example-length and example-count setting.Training documents contain 10240 tokens and validation documents contain 1024 tokens.
- Experimental findings: Increasing the number of observable symbols generally improves in-context performance despite increasing the number of output classes.More symbols make hidden-state inference easier because hidden states are more likely to emit distinct symbols.
F.5 Experiment on GPT-3
The GPT-3 experiment tests whether longer in-context examples help despite a train-test length mismatch. Longer examples improved accuracy, whereas duplicating short examples did not.
- 5-shot learning with long training examples improved LAMBADA test-short accuracy by almost 1% over short training examples.The comparison used 47-example training datasets and five sampled examples per test prompt.
- Duplicating five short examples to equalize prompt length did not improve accuracy.
- Five long examples closed about 56% of the gap between five and ten independent short examples without adding examples or task-related information.
- The results suggest that longer examples provide extra information that helps infer a latent concept beyond the input-output mapping.