PaperlayerView exact revision on arXiv

Source-linked AI summary

What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization

Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, Zhaoran Wang

arXiv:2305.19420v2stat.MLcs.LG

TL;DR

Theoretical accounts have not explained how attention supports ICL or how pretraining shapes it. The paper formulates ICL in a Bayesian latent-concept model and analyzes its estimator, regret, and transformer mechanisms. It shows that perfectly pretrained ICL implements BMA and establishes bounds on regret and pretraining error.

  • Problem

    Theoretical understanding remains limited regarding why LLMs can perform ICL, attention’s role, and how pretraining influences ICL.

    Source: §1 · p. 2

  • Method

    The paper formulates ICL as Bayesian prediction in a latent-variable model and analyzes the estimator, regret, attention mechanism, and pretraining error.

    Source: §1 · p. 2

  • Results

    Perfectly pretrained ICL has O(1/t) regret up to pretrained-model statistical error; attention approximately encodes BMA, and pretraining error splits into approximation and generalization terms with exponential-depth and sublinear-token decay.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The results provide a unified theoretical account of ICL regret, approximation, generalization, and attention’s role in ICL.

    Source: §1 · pp. 2–3

    The BMA characterization first assumes perfect pretraining; pretraining error is analyzed separately.

    Source: §4 · p. 6

Abstract

from arXiv · show

In this paper, we conduct a comprehensive study of In-Context Learning (ICL) by addressing several open questions: (a) What type of ICL estimator is learned by large language models? (b) What is a proper performance metric for ICL and what is the error rate? (c) How does the transformer architecture enable ICL? To answer these questions, we adopt a Bayesian view and formulate ICL as a problem of predicting the response corresponding to the current covariate, given a number of examples drawn from a latent variable model. To answer (a), we show that, without updating the neural network parameters, ICL implicitly implements the Bayesian model averaging algorithm, which is proven to be approximately parameterized by the attention mechanism. For (b), we analyze the ICL performance from an online learning perspective and establish a $\mathcal{O}(1/T)$ regret bound for perfectly pretrained ICL, where $T$ is the number of examples in the prompt. To answer (c), we show that, in addition to encoding Bayesian model averaging via attention, the transformer architecture also enables a fine-grained statistical analysis of pretraining under realistic assumptions. In particular, we prove that the error of pretrained model is bounded by a sum of an approximation error and a generalization error, where the former decays to zero exponentially as the depth grows, and the latter decays to zero sublinearly with the number of tokens in the pretraining dataset. Our results provide a unified understanding of the transformer and its ICL ability with bounds on ICL regret, approximation, and generalization, which deepens our knowledge of these essential aspects of modern language models.

1 Introduction

The paper addresses limited theoretical understanding of ICL by analyzing its estimator, performance, and the transformer’s role in prompting and pretraining. Under a Bayesian latent-variable formulation, it connects Bayesian model averaging, ICL regret, attention, and pretraining error.

1 Introduction

  • The analysis formulates ICL as predicting a response for the current covariate from prompt examples drawn from a latent-variable model with a hidden concept.The paper studies transformers pretrained on a large dataset and then prompted for ICL.Source: §1 · p. 2
  • Perfectly pretrained LLMs perform ICL as Bayesian model averaging, inferring a posterior over the hidden concept from prompt examples and aggregating predictions over it.This characterizes the estimator learned by ICL without updating the network parameters.Source: §1 · p. 2
  • ICL regret after t prompt examples is O(1/t), up to the statistical error of the pretrained model.The paper defines regret as averaged prediction error over a sequence of covariate-response examples.Source: §1 · p. 2
  • A variant of attention encodes BMA, and softmax attention approximately encodes BMA; the variant coincides with standard softmax attention as prompt length tends to infinity.The paper relates this attention mechanism to an extension of linear attention.Source: §1 · p. 2
  • Pretraining error in total variation is bounded by approximation plus generalization error; approximation decays exponentially with transformer depth and generalization sublinearly with pretraining tokens.The paper also gives sufficient conditions for ICL robustness to perturbed responses and establishes a finite-sample statistical error for that misspecified case.Source: §1 · p. 3
  • Together, the results provide a unified theoretical account of ICL regret, approximation error, generalization error, and attention’s role.The stated scope addresses the estimator learned by ICL, its performance metrics and error rates, and the transformer’s role in pretraining and prompting.Source: §1 · pp. 2–3

2 Related Work

Prior work studies ICL’s empirical behavior and theoretical mechanisms, including algorithmic and Bayesian accounts. The paper distinguishes its analysis by relating attention to BMA and connecting transformer pretraining error to ICL regret.

2 Related Work

  • Prior work boosts ICL through training, demonstrations, instructions, and chain-of-thought prompting, while empirical studies examine data distributions and example mappings (Wei et al., 2021; Liu et al., 2021; Wei et al., 2022c; Chan et al., 2022; Min et al., 2022).Other studies examine approximable function classes and the role of task demonstrations.Source: §2 · p. 3
  • Several theoretical analyses cast ICL as gradient descent or least squares, but some show only approximation capacity and another omits softmax (Aky¨urek et al., 2022; von Oswald et al., 2022; Bai et al., 2023; Dai et al., 2022).The paper notes these choices may not align with pretrained models or practical implementation.Source: §2 · p. 3
  • Earlier Bayesian analyses use restrictive HMM or pretrained-distribution assumptions; this paper instead links attention-encoded BMA to transformer pretraining error and ICL regret (Xie et al., 2021; Wies et al., 2023).The related-work discussion also notes that Bayesian analyses do not explain attention’s role in ICL.Source: §2 · pp. 3–4

3 Preliminary

The preliminaries establish notation and describe attention, transformers, autoregressive language modeling, and the latent-concept setup used to formalize ICL.

3 Preliminary

  • The notation defines total variation distance between distributions as the supremum over events of their probability difference, alongside conventions for indexing vectors and matrices.It also introduces probability-measure and matrix-norm notation.Source: §3 · p. 4
  • Attention maps a query and key-value pairs to a weighted sum of values, with weights determined by softmax; multi-head attention combines parallel attention modules.The transformer is described as attention modules concatenated with fully connected layers.Source: §3 · p. 4
  • Autoregressive LLMs predict each next token from previous tokens, are pretrained to predict nominal next tokens, and are then prompted for downstream tasks without parameter tuning.The paper distinguishes this in-context learning from explicitly task-trained fine-tuning.Source: §3 · p. 4
  • The ICL setup draws prompt examples and a query from a hidden concept, then predicts the response for the query conditioned on the prompt and concept.The paper’s example illustrates inferring the response “fungi” from examples classifying cats and pineapples.Source: §3 · p. 5

4 In-Context Learning via Bayesian Model Averaging

Under a latent-concept model, the paper characterizes perfectly pretrained ICL as Bayesian model averaging and bounds its online regret. It then shows how an attention variant encodes that averaging and approaches softmax attention in long sequences.

4 In-Context Learning via Bayesian Model Averaging

  • Perfectly pretrained LLMs implement BMA by inferring a posterior over the hidden concept from prompt examples and averaging concept-conditioned predictions for the query.The posterior is based on the observed examples, and the resulting prediction aggregates likelihoods over that posterior.Source: §4 · pp. 5–6
  • With a finite concept space and uniform prior, regret is at most log |Z|/t, and enough examples make ICL nearly match an oracle that knows z*.The oracle also knows the likelihood function under the true concept.Source: §4 · pp. 6–7
  • ICL regret is bounded by log(1/P_Z(z*))/t, comparing predictions with the best hidden concept in hindsight.The bound makes regret decrease as the number of examples grows; the bound is larger for concepts with lower prior probability.Source: §4 · p. 7

4.1 Attention Parameterizes Bayesian Model Averaging

  • In a Gaussian linear model, attn† estimates the posterior concept from key-value examples and combines it with the query to encode BMA.The paper derives this attention variant by relating the posterior-mean concept to query-key similarity.Source: §4.1 · pp. 7–8
  • With iid unit-norm key-value pairs and a Gaussian RBF kernel, attn† converges to a constant-scaled softmax attention as sequence length grows.The convergence result is stated for the assumptions in Proposition 4.3.Source: §4.1 · p. 8

5 Theoretical Analysis of Pretraining

The paper analyzes pretraining on sequences generated from sampled hidden concepts, then bounds the pretrained model’s error as approximation plus generalization error. Approximation error decreases exponentially with transformer depth under stated conditions.

5.1 Pretraining Algorithm

  • Pretraining data comprise independent trajectories, each generated after sampling a hidden concept; the transformer is trained with cross-entropy to predict the next token.The model uses stacked multi-head attention and feed-forward modules, with a softmax output.Source: §5.1 · pp. 8–9

5.2 Performance Guarantee for Pretraining

  • Under Assumptions 5.1 and 5.2, the pretraining guarantee holds with probability at least 1−δ.The stated bound uses the mixing time of the Markov chains induced by the nominal model.Source: §5.2 · p. 10
  • Theorem 5.3 bounds pretrained-model error by approximation and generalization terms, with the generalization term not increasing with sequence length.The bound is established via PAC-Bayes; its approximation terms vanish when the nominal model is exactly representable by transformers.Source: §5.2 · p. 10
  • The paper contrasts its pretraining analysis with prior work that assumes an arbitrarily precise model or cannot capture approximation error for maximum-likelihood training.It reports that Li et al. (2023) yield an O((NpTp)^−1/4) total-variation rate even in the realizable case.Source: §5.2 · p. 10
  • Under certain smoothness conditions, approximation error decreases exponentially as transformer depth increases.The proposition also imposes conditions on feed-forward width and parameter bounds; the analysis derives the bound for one-dimensional token inputs and says it carries over to higher dimensions.Source: §5.2 · p. 11

6 ICL Regret under Practical Settings

For imperfect pretraining, the paper relates ICL regret to pretraining error and analyzes prompts with incorrect input-output mappings under distinguishability conditions. It concludes that its theory combines regret and pretraining guarantees, while focusing on example-based rather than instruction-based prompts.

6.1 ICL Regret with an Imperfectly Pretrained Model

  • Expected ICL regret for the pretrained model is bounded by the underlying true-model ICL regret plus pretraining error.The prompt distribution is assumed to be covered by the pretraining distribution; the authors note that inference outside pretraining support may be imprecise.Source: §6.1 · pp. 11–12

6.2 Prompting with Wrong Input-Output Mappings

  • When the task concept is distinguishable, the concept-identification error for prompts with wrong mappings decays exponentially in t.The bound also includes pretraining error, which is related to pretraining-set size and neural-network capacity.Source: §6.2 · p. 12

A Conclusion

  • The conclusion combines ICL regret, attention’s parameterization of Bayesian model averaging, and pretraining approximation and generalization errors to describe pretrained-model ICL.The analysis focuses on prompts containing examples and leaves instruction-based prompts for future work.Source: A Conclusion · p. 18

B More Related Works

  • The paper reports sharper channel-number dependence in its transformer generalization bounds than Edelman et al. (2021).It frames its analysis as covering maximum-likelihood estimation with a transformer function class.Source: B More Related Works · p. 19

D Discussion About The Experimental Results in Existing Works

The paper relates its theory to reported ICL behavior in existing experiments and introduces conditional mean embeddings and their empirical operators.

D Discussion About The Experimental Results in Existing Works

  • Akyürek et al. (2022) train a transformer from scratch on the ICL distribution, and their reported Bayesian-predictor match is the behavior proved in this paper.The paper describes that setup as a special case of its pretraining setting.Source: D Discussion About The Experimental… · p. 19
  • Garg et al. (2022) report significant ICL error when T≤d; the paper’s discretization argument yields regret d log(C/δ)/T, matching that figure.Here d is the hidden-variable dimension and C is an absolute constant.Source: D Discussion About The Experimental… · p. 19

E.1 Introduction of Conditional Mean Embedding

  • The conditional mean embedding represents a conditional distribution in an RKHS, and its empirical operator is defined using sampled feature mappings and covariance operators.The passage also states that the population conditional operator equals the cross-covariance operator times the inverse covariance operator.Source: §E.1 · p. 20

E.2 Proof of Proposition 4.1

The proof applies Bayes’ theorem to express the predictive distribution through posterior averaging over hidden concepts. The accompanying corollary gives an O(1/T) convergence rate for Bayesian posterior averaging.

E.2 Proof of Proposition 4.1

  • Bayes’ theorem rewrites the predictive distribution as an average over hidden concepts weighted by their posterior given the prompt.The proof then invokes Bayesian model averaging for density estimation.Source: §E.2 · pp. 20–21, §E.3 · p. 21

E.3 Proof of Corollary 4.2

  • O(1/T) is the stated statistical convergence rate of Bayesian posterior averaging.The proof considers q in the class of all Dirac measures.Source: §E.3 · p. 21

E.4 Proof of Proposition 4.3

The proof relates attention to conditional mean embeddings: attn† recovers the empirical embedding, while softmax attention approaches a scaled conditional expectation as the number of examples grows.

E.4 Proof of Proposition 4.3

  • The newly defined attn† recovers the empirical conditional mean embedding.The proof identifies this relationship as one of its two main steps.Source: §E.4 · p. 22
  • As L grows, softmax attention attnSM(q, K, V) approaches a constant multiple of E[V | K = q].Source: §E.4 · p. 25

F.1 Supplemental Definitions for Markov Chains

  • The mixing time tmix(ε) is the earliest t for which the chain’s distance d(t) is at most ε; tmix uses ε = 1/4.Source: §F.1 · p. 26

F.2 Proof of Theorem 5.3

The proof of Theorem 5.3 uses a PAC-Bayes error decomposition, while Proposition 5.4 analyzes transformer approximation and the pretraining setup uses squared loss on input-target pairs.

F.2 Proof of Theorem 5.3

  • The proof decomposes error using PAC-Bayes into fluctuation, approximation, and KL-divergence terms.It treats the best approximation as the projection minimizing KL divergence and controls generalization with a uniform-convergence bound.Source: §F.2 · pp. 27, 28, 30, 31

F.3 Formal Statement and Proof of Proposition 5.4

  • For finite alphabets, the next-word distribution can be represented as ρ* applied to the average of feature values φ*(x_i).The proof constructs an approximator for this representation and assumes sufficient smoothness of the target functions.Source: §F.3 · pp. 32–33

F.4.1 Pretraining Algorithm with ℓ2 Loss

  • The squared-loss pretraining setup uses i.i.d. input-target pairs, and its population-optimal predictor is the conditional expectation E[x | S].Source: §F.4.1 · p. 34

F.4.2 Performance Guarantee for Pretraining with ℓ2 Loss

The section presents a pretraining error guarantee, explains its approximation and generalization components, and states assumptions for analyzing prompts containing wrong input-output mappings.

F.4.2 Performance Guarantee for Pretraining with ℓ2 Loss

  • Theorem F.6 bounds pretrained-model error by approximation and generalization terms, with the generalization term independent of token sequence length.The guarantee is stated under Assumption F.5.Source: §F.4.2 · p. 35

F.4.3 Proof of Theorem F.6

  • The proof decomposes excess population risk into generalization and approximation error.It uses the empirical-risk minimizer to relate these two terms to the excess risk.Source: §F.4.3 · p. 36

F.4.4 Proof of Proposition F.8

  • The proof of Proposition F.8 follows Proposition F.4’s first two steps and sets the final-layer matrix to the identity.Source: §F.4.4 · p. 39

G.2 Assumptions and Formal Statement for Prompting With Wrong Input-Output Mappings

  • The wrong-mapping analysis assumes independent input-output pairs, positive prior probability for the true concept, and sufficient KL separation from other concepts.The text notes that pair independence can be relaxed when the token-generation process has additional structure, such as a hidden Markov model.Source: §G.2 · pp. 40–41

G.1 Proof of Theorem 6.2

The proof of Theorem 6.2 applies the pretraining bound under a prompt-distribution comparison; the proof of Proposition G.4 uses Bayesian model averaging to analyze wrong mappings.

G.1 Proof of Theorem 6.2

  • Theorem 6.2’s proof compares the ICL prompt distribution with the prompt distribution using Assumption 6.1, then invokes Theorem 5.3.Source: §G.1 · p. 40

G.3 Proof of Proposition G.4

  • For wrong mappings, the proof writes the output distribution using Bayesian model averaging and rewrites the concept weights with Bayes’ theorem.Source: §G.3 · p. 41
  • The proof bounds wrong-concept likelihood ratios with Hoeffding’s inequality and uses the pretraining error bound Δpre(Np, Tp, δ).Source: §G.3 · p. 42

H Proof of Supporting Propositions

The appendix proves supporting propositions using sphere symmetry, concentration arguments, layerwise network-error bounds, and technical lemmas.

H.1 Proof of Proposition E.1

  • Sphere symmetry makes the constant C1 independent of b in the proof of Proposition E.1.The proof begins with vectors on the unit sphere and uses the symmetry of that sphere.Source: §H.1 · p. 43

H.2 Proof of Proposition F.2

  • The proof of Proposition F.2 combines concentration and a variational argument to obtain a probability statement holding simultaneously for every distribution P.The stated probability is at least 1 − δ.Source: §H.2 · p. 44

H.3 Proof of Proposition F.1

  • The proof bounds layerwise network-output differences recursively, using layer normalization and a feedforward-network perturbation bound.The proof also builds a recursion relating output differences between successive layers.Source: §H.3 · pp. 44–45

H.4 Proof of Proposition F.9

  • The proof of Proposition F.9 follows Proposition F.1’s layerwise procedure, with layer normalization used in the final-layer bound.The passage states that the remaining proof follows the procedures for Proposition F.1.Source: §H.4 · p. 45

H.5 Proof of Proposition H.1

  • The feedforward-network bound separates input differences from parameter differences, while the proof uses ReLU’s non-expansiveness.The bound includes terms involving both the input and the first-layer parameter difference.Source: §H.5 · p. 46

I Technical Lemmas

  • The technical lemmas include a width-at-most-16 network approximation guarantee and an expectation-based representation of KL and total variation.The approximation guarantee also bounds network depth and maximal weight.Source: I Technical Lemmas · p. 48
Loading 2305.19420v2…