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

TL;DR

This work provides a rigorous theoretical foundation for In-Context Learning in pretrained transformers by casting ICL as Bayesian Model Averaging over latent concepts. It shows that attention mechanisms implement BMA and that softmax attention approximates BMA in long prompts, yielding an $\mathcal{O}(1/t)$ regret bound for perfectly pretrained models. The paper also analyzes pretraining through a PAC-Bayes lens, deriving a decomposition of pretrained-model error into exponentially decaying approximation error with depth and sublinearly decaying generalization error with data size, and extends results to practical settings with imperfect pretraining and robustness to wrong prompt mappings. Together, these results offer a unified, quantitative account of ICL, pretraining, and transformer dynamics with concrete bounds on regret, approximation, and generalization.

Abstract

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.

Figures (1)

• Figure 1: To form the pretraining dataset, a hidden concept $z$ is first sampled according to $\mathbb{P}_{\mathfrak{Z}}$, and a document is generated from the concept. Taking the token sequence $S_{t}$ up to position $t\in[T]$ as the input, the llm is pretrained to maximize the next token $x_{t+1}$. During the icl phase, the pretrained llm is prompted with several examples to predict the response of the query.

Theorems & Definitions (26)

• Proposition 4.1: llms Perform bma
• Corollary 4.2: icl Regret of Perfectly Pretrained Model
• Proposition 4.3
• Theorem 5.3
• Proposition 5.4: Informal
• Theorem 6.2: ICL Regret of Pretrained Model
• Proposition 6.3: Informal
• Proposition E.1
• Proposition F.1
• Proposition F.2