PaperlayerView exact revision on arXiv

Source-linked AI summary

An Explanation of In-context Learning as Implicit Bayesian Inference

Sang Michael Xie, Aditi Raghunathan, Percy Liang, Tengyu Ma

arXiv:2111.02080v6cs.CLcs.LG

TL;DR

The mechanism behind in-context learning is unclear, in part because real pretraining data are messy and prior study has required large datasets and models. The paper analyzes latent-concept HMM pretraining and proves in-context learning can occur under a mixture of HMMs, then demonstrates it with Transformers and LSTMs on GINC.

  • Problem

    What enables in-context learning remains unclear, with understanding limited by messy pretraining data and prior reliance on large datasets and models.

    Source: §1 · p. 2

  • Method

    The paper models pretraining documents as samples from latent-concept HMMs and analyzes when an LM can infer a shared prompt concept despite distribution mismatch.

    Source: §2 · pp. 3–4

  • Results

    The authors prove that in-context learning occurs for a mixture-of-HMM pretraining distribution and demonstrate it in GINC with Transformers and LSTMs.

    Source: §4 · p. 10, §6 · p. 12

  • Takeaways & Limitations

    GINC offers a small-scale synthetic setting for studying large-scale in-context-learning phenomena while controlling latent data properties.

    Source: §5 · p. 12

    The work leaves bridging the mismatch between fixed-length training examples and randomly sized test examples to future work.

    Source: §3.3.2 · p. 9

Abstract

from arXiv · show

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 explains in-context learning as inferring a shared latent concept from prompt examples, analogous to the concept inference needed to generate coherent pretraining documents. Its theory and GINC experiments examine when this mechanism works and what affects its performance.

Introduction

  • GPT-3 demonstrates in-context learning by inferring tasks from prompt examples, including nationality prediction, and achieves state-of-the-art results on LAMBADA (Paperno et al., 2016) and TriviaQA (Joshi et al., 2017).The paper notes that this occurs without explicit pretraining to learn from examples, even though prompts concatenate independent examples unlike natural text (Brown et al., 2020, Lieber et al., 2021, Radford et al., 2019, Wang and Komatsuzaki, 2021).Source: §1 · pp. 1–2
  • In-context learning can arise when an LM pretrained on coherent documents infers a latent HMM concept, then uses examples to infer a shared prompt concept through Bayesian marginalization.The account assumes the LM fits the pretraining distribution and treats prompt completions as posterior predictive distributions, despite prompts coming from a different distribution.Source: §1 · pp. 2–3
  • Despite prompt–pretraining distribution mismatch, prediction error is asymptotically optimal when each example’s concept signal exceeds mismatch error, and decreases as example length grows.The latter result means input information, not only the input-output mapping, can help in-context learning.Source: §1 · p. 3
  • On GINC, Transformers (Vaswani et al., 2017) and LSTMs exhibit in-context learning, with accuracy improving as example count and length increase.Ablations also indicate that latent concept structure in pretraining is crucial for in-context learning to emerge.Source: §1 · p. 3
  • Scaling model parameters steadily improves GINC in-context accuracy even at the same pretraining loss; GINC also mirrors example-order sensitivity (Zhao et al., 2021) and cases where zero-shot beats few-shot (Brown et al., 2020).These experimental phenomena raise questions beyond the paper’s theory, which studies the pretraining distribution.Source: §1 · p. 3

2 In-context learning setting

The setting models pretraining as documents generated under latent concepts that determine HMM transitions, then studies prediction from example prompts sharing a concept. Its assumptions specify delimiter behavior and constrain prompt distribution shift and model regularity.

2 In-context learning setting

  • Pretraining samples a latent concept and then a length-T document, with the concept determining the HMM hidden-state transition matrix.The concept is drawn from a family Θ, and observed tokens are generated from the document distribution conditioned on it.Source: §2 · p. 3
  • Prompts concatenate n independent training examples and a test input conditioned on a shared concept; the predictor selects the most likely test output given the examples.Examples contain an input sequence followed by its output token, and are separated by a special delimiter.Source: §2 · pp. 3–4
  • Because concatenated examples create unnatural transitions, prompts can be low-probability under pretraining, yet the framework asks whether LMs infer their shared concept despite this mismatch.The names-to-nationalities example contrasts document transitions such as name → nationality → occupation with prompt transitions such as nationality → name.Source: §2 · p. 4

2.1 Assumptions

  • The assumptions make delimiters identify a subset of hidden states without identifying a specific state, while limiting concept-dependent delimiter transitions and prompt-start shift.Delimiter emission is deterministic in delimiter states and impossible outside them; the prompt-start distribution must be within TV distance ∆/4 of transitions from delimiter states under θ∗.Source: §2.1 · pp. 4–5
  • The prompt concept must belong to Θ, and regularity conditions bound transitions, start-state probabilities, token emission, and prior support.These conditions are intended to ensure the prompt has nonzero probability under θ∗, which Bayesian inference requires to identify that concept.Source: §2.1 · p. 5

3 Theoretical analysis

In-context prediction becomes optimal as the number of examples grows when the prompt concept is distinguishable; without distinguishability, expected error still decreases with example length, with information from both inputs and input-output mappings contributing.

3.1 High-level approach

  • As n grows, Bayesian marginalization selects the prompt concept under distinguishability, making the in-context predictor optimal.The argument shows competing concepts’ likelihood terms vanish, while the prompt concept’s term remains nonzero.Source: §3.1 · p. 6

3.2 Heuristic derivation

  • Delimiter properties approximately factorize dependent examples, while KL-divergence signal must outweigh distribution-mismatch errors for competing concepts’ likelihood terms to vanish.The analysis decomposes the signal into token-level negative KL terms and accounts for O(1) mismatch errors.Source: §3.2 · pp. 6–7

3.3.1 In-context learning under distinguishability

  • Under the distinguishability condition, the in-context predictor asymptotically achieves the optimal 0-1 risk, and longer examples strengthen the signal for selecting the prompt concept.The condition requires the KL-divergence signal to exceed errors from distribution mismatch and low-probability transitions between examples.Source: §3.3.1 · p. 7, Figure 2 · p. 8

3.3.2 Non-distinguishable case

  • With continuous concepts and a KL second-order expansion, excess risk roughly decreases as O(1/k), and improves when the Fisher-information condition number is smaller.The result assumes continuity; concepts indistinguishable from the prompt concept produce similar outputs and therefore small error.Source: §3.3.2 · p. 8
  • For test examples of uniformly random length from 2 to k, the paper gives a risk bound as n grows, averaging prediction error across tokens 2 through k.The analysis leaves unresolved the mismatch between fixed-length training examples and random-length test examples.Source: §3.3.2 · p. 9

4 Simulations

On the synthetic GINC dataset, Transformers and LSTMs exhibit in-context learning, and the pretraining distribution’s in-context learning properties transfer to models that approximate it. The simulations also show that this behavior depends on the latent concept structure and varies with model scale, prompt ordering, and prompt conditions.

Simulation scope

  • The in-context learning properties of the pretraining distribution transfer to language models that approximate it, rather than fitting it exactly.Both Transformers and LSTMs trained on GINC exhibit in-context learning.Source: §4 · p. 9

GINC dataset

  • GINC uses a uniform mixture of HMMs over 5 concepts to generate 1,000 pretraining documents totaling about 10 million tokens, with prompts containing 0–64 examples.Prompt example lengths are 3, 5, 8, or 10; targets use the most likely output to set intrinsic error to zero.Source: §4 · p. 10

Main result

  • In-context accuracy improves as prompt-example count and example length increase for both Transformers and LSTMs trained on GINC, verifying the theory.Results are averaged over 5 pretraining runs across vocabulary sizes 50, 100, and 150.Source: §4 · p. 10

Ablations on the latent concept structure

  • Removing the mixture-of-concepts prior or pretraining on random data with all possible token transitions produces flat or failed in-context learning.The ablations indicate that diverse token transitions alone do not enable in-context learning.Source: §4 · p. 10

Extrapolation to unseen concepts

  • In-context learning fails on prompts generated from five random concepts outside the pretraining concept family.The test uses a 4-layer Transformer and demonstrates the effect of unseen concepts in GINC.Source: §4 · p. 10

Effect of model size and architecture

  • 81% to 85% in-context accuracy accompanies scaling from 12 to 16 Transformer layers despite both models’ 1.33 validation loss at vocabulary size 50.In-context accuracy increases with Transformer size; LSTMs also consistently outperform Transformers on GINC despite having fewer parameters.Source: §4 · p. 10, Figure 5 · p. 11, Figure 6 · p. 11

Sensitivity to example ordering

  • Permutations of the same four prompt examples yield in-context accuracy differences of 10–40% for a 4-layer Transformer on GINC.The experiment evaluates all 24 permutations for each of 10 example sets.Source: §4 · p. 10, Figure 7 · p. 11

Zero-shot is sometimes better than few-shot

  • With lower-entropy transitions, few-shot accuracy can initially trail zero-shot accuracy, then recover as examples are added.This occurs in a GINC setting with temperature 0.01, 12 concepts, and vocabulary size 100; the authors hypothesize that the prompt structure is initially distracting.Source: §4 · pp. 10–11, Figure 7 · p. 11

5 Discussion and related work

The discussion distinguishes in-context learning that emerges naturally from LM pretraining from explicit learning or inference procedures, and presents GINC as a controllable small-scale lens on large-scale phenomena.

Learning via Bayesian inference and extrapolation

  • The canonical Bernstein-von Mises theorem does not apply because prompt examples are dependent under pretraining, and existing HMM results do not cover observations from a different distribution.Future directions include sharper posterior asymptotics and misspecification or extrapolation results; factoring concepts into semantics and syntax may enable generalization to unseen pairs through operations such as copying or reordering.Source: §5 · pp. 11–12

Topic models and HMMs

  • Topic models such as LDA use document-level latent variables but typically rely on explicit inference algorithms, whereas this work studies learning as a natural result of Bayesian inference.Wei et al. (2021a) analyze masked-LM representations that improve optimization-based downstream learning, rather than in-context learning.Source: §5 · p. 12

Bridging the mismatch between pretraining and prompting

Meta-learning

  • Unlike meta-learning, where models are trained to learn, in-context learning emerges from LM pretraining.Source: §5 · p. 12

Studying large-scale phenomena at a small scale

  • GINC makes large-scale in-context learning phenomena tractable at small scale because its HMM and latent-concept complexity is low, while synthetic controls enable predictions about large LMs.The data and models are relatively larger than the pretraining distribution's complexity, and latent properties such as unseen concepts can be controlled.Source: §5 · p. 12

6 Conclusion · A Framework details

The paper frames in-context learning as implicit Bayesian concept inference and shows it can arise with mixture-of-HMM pretraining. Its framework details how prompts are sampled from hidden segments and converted into input-output examples.

6 Conclusion

  • We frame in-context learning as implicit Bayesian inference: the pretrained LM implicitly infers a concept when predicting.Source: §6 · p. 12
  • In-context learning occurs when pretraining data follow a mixture of HMMs; the authors present this as a first step toward improving pretraining and prompting.Source: §6 · p. 12

A Framework details

  • Prompts are sampled from a new distribution comprising n independent training examples and one test example.Source: A Framework details · p. 16
  • Each example's hidden segment has length k: its first state comes from the prompt-start distribution, and remaining states follow pretraining transitions for concept θ*.Source: A Framework details · p. 16
  • Each training example, except the test example, ends with a delimiter sampled from p_delim.Source: A Framework details · p. 16
  • Given hidden variables and delimiters, observed tokens are sampled from the pretraining distribution.Source: A Framework details · p. 16
  • Each input is the first k−1 observed tokens and its output is the kth; the test input is sampled through the same process using k−1 elements.Source: A Framework details · p. 16

B Propositions for Theorem 1 · C Convergence of the in-context predictor

The appendix establishes positive probability bounds for delimiter tokens and examples under θ*, then proves that under Assumption 3 the in-context predictor preserves the prompt predictive distribution’s most likely label. The convergence argument bounds predictive-distribution differences by Δ/2, which is small enough relative to the label margin Δ to preserve the argmax.

B Propositions for Theorem 1

  • The propositions establish a positive lower bound on delimiter-token probability under θ*, based on the stated assumptions.The appendix introduces these propositions as direct corollaries of the assumptions.Source: B Propositions for Theorem 1 · p. 16
  • For all examples and future hidden states, the probability of an example under θ* is bounded below by c7 > 0.The bound applies to every i and future hidden state hj,l for any l and j > i.Source: B Propositions for Theorem 1 · p. 17
  • The example-probability lower bound follows by lower-bounding numerator terms and upper-bounding the denominator by 1.The proof invokes Assumption 5 and sets c7 from the resulting positive bound.Source: B Propositions for Theorem 1 · p. 17

C Convergence of the in-context predictor

  • The convergence argument abstracts away Bayesian selection of θ* from Θ and completes that part of Theorem 1 separately.The section first proves convergence conditioned on θ* and a given prompt and test input.Source: C Convergence of the in-context… · p. 17
  • Under Assumption 3, the averaged predictive distribution conditioned on θ* and prompt Sn has the same argmax as the prompt predictive distribution.This is the stated convergence lemma for the in-context predictor.Source: C Convergence of the in-context… · p. 17
  • The convergence proof represents transitions from a delimiter state and test-label probabilities with matrices T and W, then compares these with the prompt-start distribution u.It bounds the predictive-distribution difference before applying the Δ/2 bound.Source: C Convergence of the in-context… · pp. 17–18
  • The predictor’s argmax is preserved because each output probability changes by at most Δ/2 while the top-label margin is Δ.The proof bounds the distribution difference using Assumption 3 and total-variation convexity.Source: C Convergence of the in-context… · p. 18

D Proof of Theorem 1 · E Non-distinguishable case · E.1 Proof of Theorem 2

Under distinguishability, the proof shows that marginalization selects the prompt concept and recovers the prompt’s predictive argmax. In the non-distinguishable case, it bounds the limiting predictor’s 0-1 risk through KL divergence and logistic-loss calibration, with the Theorem 2 proof supplying a KL bound under continuity.

D Proof of Theorem 1

  • When Δ′<Δ/2 and the approximation error is sufficiently small, the argmax of the predictive distribution matches arg max_y pprompt(y|xtest).The proof sets ϵ < (Δ/2 − Δ′)p(θ*) and uses a margin of at least Δ/2 − Δ′.Source: D Proof of Theorem 1 · p. 19
  • Under Condition 1, r_n(θ) converges to a negative constant for θ≠θ*, while exp(n · r_n(θ*))=1, selecting θ* through marginalization.As k grows, observations of θ* overpower distracting delimiter transitions, making the bound negative when the KL terms exceed the error terms.Source: D Proof of Theorem 1 · pp. 19, 22

E Non-distinguishable case

  • When Condition 1 is unsatisfied, Equation 14 gives an upper bound on summed KL divergences between next-token distributions for different-length histories; the task measures last-token accuracy, motivating a more precise error bound.A lemma uses results of Steinwart (2007) and Ávila Pires and Szepesvári (2016) to relate a KL-divergence bound to 0-1 loss.Source: E Non-distinguishable case · p. 23
  • The limiting in-context predictor can be an argmax of a mixture over B, and convexity bounds that mixture’s KL to the prompt concept by the KL for any θ∈B.The argument uses the mixture form of the limiting predictor and convexity of KL.Source: E Non-distinguishable case · p. 23
  • If θ* minimizes multiclass logistic risk and every θ∈B has expected predictive KL at most ϵθ, calibration converts the surrogate-risk bound into a 0-1 loss bound with zero approximation error.The lemma assumes bounded KL and uses a calibration function for multiclass logistic loss.Source: E Non-distinguishable case · pp. 23–24
  • The inverse calibration function g−1 is roughly linear in ϵ below 0.7, where the bound is non-vacuous.Source: E Non-distinguishable case · p. 24

E.1 Proof of Theorem 2

  • Theorem 2’s proof uses continuity to bound the final KL term, then substitutes that bound into Lemma 2 to obtain the result.The proof identifies KL_k(θ*∥θ) with the expected KL between prompt and θ predictive distributions.Source: §E.1 · p. 24

E.2 Proof of Theorem 3

The proof of Theorem 3 combines Condition 1’s bound on KL divergences across positions with Lemma 2, using a test-example length drawn uniformly from 2 to k.

E.2 Proof of Theorem 3

  • Condition 1 bounds summed KL divergence across all positions, and a uniformly random test-example length from 2 to k lets the predictor take credit for fitting intermediate tokens before Lemma 2 yields the result.The proof applies Theorem 1 and Condition 1, then plugs the resulting bound into Lemma 2.Source: §E.2 · pp. 24–25

F Experimental details … F.4 Varying the vocabulary size

The appendix specifies GINC’s synthetic HMM data and model-training setups, then reports that larger vocabularies generally improve in-context accuracy. It also describes GPT-3 results on LAMBADA showing improved accuracy with longer examples despite test-set distribution mismatch.

F.1 GINC dataset

  • GINC pretraining uses a mixture of factorial HMMs whose latent entities and properties evolve independently, emit tokens through a shared memory matrix, and encode concepts through property-transition matrices.Entities change slowly, while properties change more quickly; the concept and prompt start distribution determine the in-context task.Source: §F.1 · pp. 25–26, Figure 9 · p. 26
  • The synthetic corpus comprises 1,000 documents and about 10 million tokens across five concepts, with 10 entities, 10 properties, and vocabulary sizes of 50, 100, or 150.The memory matrix has 100 entries, with its first column fixed to the delimiter and other entries sampled uniformly from the vocabulary.Source: §F.1 · pp. 26, 28
  • Prompts use a uniformly sampled concept and a fixed random starting property, with 0–64 examples of lengths 3, 5, 8, or 10; labels use the most likely next token to remove sampling error.The dataset generates 2,500 prompts for each example-length and example-count pair.Source: §F.1 · pp. 27–28

F.2 Transformer details

  • The GPT-2-based Transformers have 4, 12, or 16 layers and share 12 attention heads, 768-dimensional embeddings, dropout 0.1, and a 1,024-token context window.They train for five epochs with batch size 8; results use five pretraining seeds.Source: §F.2 · p. 28

F.3 LSTM details

  • The LSTM language model has six layers, 768-dimensional embeddings and hidden states, dropout 0.2, and weight decay 1e-5.It trains for 10 epochs with batch size 8, gradient clipping at norm 1.0, and backpropagation through 1,024 steps.Source: §F.3 · p. 28

F.4 Varying the vocabulary size

  • Across 4-, 12-, and 16-layer Transformers and an LSTM, in-context accuracy generally rises with vocabulary size (50, 100, 150), because more distinct symbols ease hidden-state inference despite more output classes.The models must infer both the prompt concept and the final HMM hidden state.Source: §F.4 · p. 29, Figure 11 · p. 29, Figure 12 · p. 29
  • On a filtered LAMBADA test set with short examples, GPT-3 accuracy improves with longer examples despite test-set distribution mismatch.The table compares prompts with 5 or 10 training examples, with the latter rows using 10 examples to equalize total length; the supplied caption gives no accuracy values.Source: Table 1 · p. 30

F.5 Experiment on GPT-3

On LAMBADA test-short, GPT-3 performed better with five longer in-context examples than with five short examples, despite the train–test length mismatch. Duplicating short examples did not help, while longer examples closed part of the gap to using ten independent short examples.

F.5 Experiment on GPT-3

  • The experiment used five sampled training examples per test case, with equally sized short and long training datasets of 47 examples each.Test-short examples were up to 200–300 characters; train-short examples were 200–300 characters, and train-long examples were 500–600 characters.Source: §F.5 · p. 30
  • Using five long training examples improved LAMBADA test-short accuracy by almost 1% over five short examples and closed about 56% of the gap to ten independent short examples.The long examples came from a distribution that did not match the short test examples.Source: §F.5 · pp. 30–31
  • Duplicating five short examples to increase prompt length did not improve accuracy, unlike adding longer examples.The long examples may help the model infer a latent concept by providing additional information beyond the input-output mapping.Source: §F.5 · pp. 30–31

Related papers

Loading 2111.02080v6…