Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, Song Mei

TL;DR

This work establishes a rigorous statistical framework showing that Transformers can perform a wide range of in-context learning tasks by implementing core ML algorithms (ridge, LS, Lasso, GLMs, GD) directly in the prompt. It introduces two general in-context algorithm selection mechanisms (post-ICL validation and pre-ICL testing) that enable a single transformer to adaptively pick the best base ICL method across inputs, achieving near-Bayes or Bayes-optimal performance in challenging settings. The paper also provides end-to-end pretraining guarantees, showing that transformers pretrained on diverse ICL instances can generalize with quantified excess risk bounds, and it validates the theory with experiments demonstrating strong in-context algorithm selection in encoder and decoder architectures. Overall, the results reveal a principled route for designing transformers that act as flexible, provable statisticians capable of adaptive computation without weight updates at inference time.

Abstract

Neural sequence models based on the transformer architecture have demonstrated remarkable \emph{in-context learning} (ICL) abilities, where they can perform new tasks when prompted with training and test examples, without any parameter update to the model. This work first provides a comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, learning generalized linear models, and gradient descent on two-layer neural networks, with near-optimal predictive power on various in-context data distributions. Using an efficient implementation of in-context gradient descent as the underlying mechanism, our transformer constructions admit mild size bounds, and can be learned with polynomially many pretraining sequences. Building on these ``base'' ICL algorithms, intriguingly, we show that transformers can implement more complex ICL procedures involving \emph{in-context algorithm selection}, akin to what a statistician can do in real life -- A \emph{single} transformer can adaptively select different base ICL algorithms -- or even perform qualitatively different tasks -- on different input sequences, without any explicit prompting of the right algorithm or task. We both establish this in theory by explicit constructions, and also observe this phenomenon experimentally. In theory, we construct two general mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL validation. As an example, we use the post-ICL validation mechanism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task -- noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.

Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

TL;DR

This work establishes a rigorous statistical framework showing that Transformers can perform a wide range of in-context learning tasks by implementing core ML algorithms (ridge, LS, Lasso, GLMs, GD) directly in the prompt. It introduces two general in-context algorithm selection mechanisms (post-ICL validation and pre-ICL testing) that enable a single transformer to adaptively pick the best base ICL method across inputs, achieving near-Bayes or Bayes-optimal performance in challenging settings. The paper also provides end-to-end pretraining guarantees, showing that transformers pretrained on diverse ICL instances can generalize with quantified excess risk bounds, and it validates the theory with experiments demonstrating strong in-context algorithm selection in encoder and decoder architectures. Overall, the results reveal a principled route for designing transformers that act as flexible, provable statisticians capable of adaptive computation without weight updates at inference time.

Abstract

Neural sequence models based on the transformer architecture have demonstrated remarkable \emph{in-context learning} (ICL) abilities, where they can perform new tasks when prompted with training and test examples, without any parameter update to the model. This work first provides a comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, learning generalized linear models, and gradient descent on two-layer neural networks, with near-optimal predictive power on various in-context data distributions. Using an efficient implementation of in-context gradient descent as the underlying mechanism, our transformer constructions admit mild size bounds, and can be learned with polynomially many pretraining sequences. Building on these ``base'' ICL algorithms, intriguingly, we show that transformers can implement more complex ICL procedures involving \emph{in-context algorithm selection}, akin to what a statistician can do in real life -- A \emph{single} transformer can adaptively select different base ICL algorithms -- or even perform qualitatively different tasks -- on different input sequences, without any explicit prompting of the right algorithm or task. We both establish this in theory by explicit constructions, and also observe this phenomenon experimentally. In theory, we construct two general mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL validation. As an example, we use the post-ICL validation mechanism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task -- noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.
Paper Structure (126 sections, 74 theorems, 429 equations, 11 figures)

This paper contains 126 sections, 74 theorems, 429 equations, 11 figures.

Table of Contents

  1. Introduction
  2. Summary of contributions and paper outline
  3. Transformers as statisticians
  4. Related work
  5. In-context learning
  6. Transformers and its theory
  7. Meta-learning
  8. Techniques
  9. Preliminaries
  10. Transformers
  11. In-context learning
  12. Generalization to predicting at every token using a decoder architecture
  13. Miscellaneous setups
  14. Basic in-context learning algorithms
  15. In-context ridge regression and least squares
  16. ...and 111 more sections

Key Result

Theorem 4

For any $\lambda\ge 0$, $0\leq \alpha\leq\beta$ with $\kappa\mathrel{\mathop:}= \frac{\beta+\lambda}{\alpha+\lambda}$, $B_w>0$, and $\varepsilon<B_xB_w/2$, there exists an $L$-layer attention-only transformer ${\rm TF}^0_{\boldsymbol \theta}$ with (with $R\mathrel{\mathop:}= \max{\{ B_xB_w, B_y, 1 \}}$) such that the following holds. On any input data $(\mathcal{D},{\mathbf x}_{N+1})$ such that t

Figures (11)

  • Figure 1: Illustration of in-context algorithm selection, and two mechanisms constructed in our theory.Left, middle-left: A single transformer can perform ridge regression with different $\lambda$'s on input sequences with different observation noise; we prove this by the post-ICL validation mechanism (\ref{['sec:min-loss']}). Middle-right, right: A single transformer can perform linear regression on regression data and logistic regression on classification data; we prove this via the pre-ICL testing mechanism (\ref{['sec:pre-test']}).
  • Figure 2: Noisy linear reg with noise $\sigma_1$
  • Figure 3: Noisy linear reg with noise $\sigma_2$
  • Figure 4: Task 1 vs. task 2 at token $20$
  • Figure 6: Illustration of our main mechanism for implementing basic ICL algorithms: One attention layer implements a single \ref{['eqn:gd_iterates']} iterate (\ref{['prop:convex-gd-onestep']} & \ref{['thm:convex-gd']}). Top: the attention mechanism as in Definition \ref{['def:attention']}. Bottom: A single \ref{['eqn:gd_iterates']} iterate. Middle: Linear algebraic illustration of the attention layer for implementing a GD update.
  • ...and 6 more figures

Theorems & Definitions (110)

  • Definition 1: Attention layer
  • Definition 2: MLP layer
  • Definition 3: Transformer
  • Theorem 4: Implementing in-context ridge regression
  • Corollary 5: Near-optimal linear regression with transformers by approximating least squares
  • Corollary 6: Nearly-Bayes linear regression with transformers by approximating ridge regression
  • Theorem 7: Implementing convex risk minimization for GLMs
  • Theorem 8: Statistical guarantee for generalized linear models
  • Corollary 9: In-context logistic regression
  • Theorem 10: Implementing in-context Lasso
  • ...and 100 more