PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2306.04637v2cs.LGcs.AIcs.CLmath.STstat.ML

TL;DR

Transformers still lack a complete account of how they learn in context beyond simple algorithms. This paper develops theory and constructions for in-context algorithms and algorithm selection, showing near-optimal prediction and strong adaptive selection across tasks and noise levels.

Abstract

from arXiv · show

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.

1 Introduction

The paper develops a comprehensive statistical theory of transformer in-context learning and introduces in-context algorithm selection, enabling one transformer to adaptively choose among algorithms or tasks without explicit prompting. It establishes theoretical constructions and empirical evidence, including nearly Bayes-optimal ICL for noisy linear models with mixed noise levels.

1 Introduction

  • The paper’s central contribution is in-context algorithm selection: a single transformer adaptively chooses different base ICL algorithms or tasks across input sequences without explicit prompting.The authors prove this capability theoretically and demonstrate it empirically on standard transformer architectures, with stronger performance than individual base algorithms.Source: §1 · p. 2

Summary of contributions and paper outline

  • The paper provides end-to-end quantitative guarantees for transformer expressive power, in-context prediction, and pretraining from polynomially many training sequences.Its efficient in-context gradient-descent implementation approximates L optimization steps with error accumulating only linearly in L for broad classes of smooth convex empirical risks.Source: §1 · p. 2, Summary of contributions and paper… · p. 3
  • Transformers can implement least squares, ridge regression, Lasso, generalized linear-model learning, and two-layer neural-network gradient descent with near-optimal prediction power.The constructions have mild bounds on layers, heads, and weight norms and achieve near-optimal prediction on many in-context data distributions.Source: Summary of contributions and paper… · p. 2
  • The theory constructs two algorithm-selection mechanisms—post-ICL validation and pre-ICL testing—for selecting among in-context procedures.Examples include choosing ridge regularization based on observation noise and choosing regression versus classification based on the input task.Source: Figure 1 · p. 2, Summary of contributions and paper… · p. 3
  • Using post-ICL validation, the paper constructs a transformer that performs nearly Bayes-optimal ICL on noisy linear models with mixed noise levels.The paper identifies this as a more challenging setting than those considered in existing work and provides both theoretical guarantees and empirical evidence.Source: Summary of contributions and paper… · p. 3, In-context learning · p. 4

1.1 Related work

  • The related-work context emphasizes that this work extends prior transformer ICL theory from simple regression or gradient-descent mechanisms to broader algorithms and explicit algorithm-selection constructions.Prior work studied model selection experimentally but did not provide concrete mechanisms or constructions for in-context model selection.Source: In-context learning · p. 4

2 Preliminaries

The paper formalizes transformer layers and the in-context learning setup used throughout its theory. Inputs are encoded as token sequences, processed by fixed transformers, and evaluated by predictions at the missing-label test token.

2 Preliminaries

  • The preliminaries represent sequences as matrices of token vectors and use ReLU activations throughout the model definitions.Each h_i is a column token in H ∈ R^{D×N}, and σ(t)=ReLU(t)=max{t,0}.Source: §2 · p. 5

2.1 Transformers

  • A transformer is an L-layer composition of self-attention followed by token-wise MLP layers.Attention uses normalized ReLU weights instead of softmax for technical convenience, while each MLP applies a two-layer ReLU map to each token.Source: §2.1 · p. 5

2.2 In-context learning

  • In-context learning gives a fixed transformer N labeled examples and one unlabeled test input drawn from potentially different task distributions.The goal is to predict y_{N+1} from the encoded dataset and test input without parameter updates.Source: §2.2 · p. 6
  • The ICL input sequence stores features, labels, a zero test label, and train-token indicators, then reads the prediction from the test token’s designated output entry.The hidden dimension is chosen as D=Θ(d), and the theory focuses on prediction at the last token.Source: §2.2 · p. 6

3 Basic in-context learning algorithms

This section constructs transformers that implement standard machine-learning algorithms in context, achieving near-optimal prediction under standard assumptions. A shared in-context gradient-descent mechanism supports ridge regression, GLMs, Lasso, and two-layer neural-network training.

3 Basic in-context learning algorithms

  • Transformers implement diverse standard machine-learning algorithms in context with mild size bounds and near-optimal predictive power on standard in-context distributions.The constructions cover linear, generalized linear, sparse linear, and two-layer neural-network settings.Source: §3 · p. 6

3.1 In-context ridge regression and least squares

  • Transformers approximately implement ridge regression and least squares to arbitrary precision with constant heads and logarithmically many layers.Under standard least-squares assumptions, the resulting excess risk is O~(dσ^2/N), the rate-optimal linear-regression rate; Bayesian linear models achieve nearly-Bayes risk.Source: §3.1 · p. 7

3.2 In-context learning of generalized linear models

  • Transformers implement convex empirical-risk minimization for GLMs, including logistic regression, with O(log(N)) layers and near-optimal O(d/N) excess risk under standard assumptions.The guarantees apply when the empirical risk is strongly convex with a bounded solution and include realizable GLM settings.Source: §3.2 · pp. 8–9

3.3 In-context Lasso

  • Transformers approximate in-context Lasso and achieve O~(s log d/N) excess risk for s-sparse linear models, optimal up to logarithmic factors.The implementation assumes bounded spectral norm and solution norm, while the statistical guarantee holds for Gaussian sparse linear models with N ≥ O(s log(d/δ)).Source: §3.3 · p. 9

3.4 Gradient descent on two-layer neural networks

  • An attention-only transformer implements inexact and exact gradient descent on two-layer neural networks under suitable assumptions, with explicit size bounds but weaker guarantees due to non-convexity.The construction uses a 2L-layer transformer for L inexact-gradient steps and is more quantitative and arguably simpler than [32, Algorithm 11].Source: §3.4 · p. 10

3.5 Mechanism: In-context gradient descent

  • The core mechanism uses one attention layer per in-context gradient-descent iterate, so an (L+1)-layer transformer approximates L convex updates with linearly accumulating error.MLP layers additionally implement proximal operators for regularized losses such as the ℓ1 penalty in Lasso.Source: §3.5 · pp. 10–12

4 In-context algorithm selection

Transformers can perform in-context algorithm selection by adaptively choosing among base ICL algorithms through post-ICL validation or pre-ICL testing. These mechanisms yield nearly Bayes-optimal prediction for mixed-noise linear models and adaptive selection between regression and classification.

4.1 Post-ICL validation mechanism

  • Post-ICL validation splits data into training and validation sets, runs K base ICL algorithms, and combines predictors whose validation loss is within γ of the minimum.A 3-layer transformer can produce such a predictor and achieve nearly optimal expected loss when validation loss concentrates around expected loss.Source: §4.1 · p. 12
  • The post-ICL mechanism supports in-context ridge regularization selection, using validation loss to choose among K fixed regularization strengths with O(log(1/ε)) layers and O(K) heads.The result holds on well-conditioned data for all selected regularization strengths.Source: §4.1 · p. 12

4.1.1 Nearly Bayes-optimal ICL on noisy linear models with mixed noise levels

  • A single transformer can achieve nearly Bayes-optimal ICL under a mixture of K noise levels, with O(log N) layers and O(K) heads.The guarantee applies for sufficiently large N and d, including the proportional setting N/d = Θ(1), and approaches the individual Bayes risks for all K noise levels simultaneously.Source: §4.1.1 · p. 13
  • The framework also enables adaptive link-function selection for generalized linear models and extends to tests such as linear correlation for confident linear regression.The generalized-linear-model construction selects the link function for each ICL instance.Source: §4.1.1 · p. 13, §4.2 · p. 14

4.2 Pre-ICL testing mechanism

  • Pre-ICL testing computes simple summary statistics to identify the input distribution and select an appropriate ICL algorithm before running it.A binary-label test yields logistic regression for labels in {0, 1} and least squares with high probability when the label distribution is not concentrated around {0, 1}.Source: §4.2 · pp. 13–14

5 Analysis of pretraining

This section establishes finite-sample pretraining guarantees for transformers and combines them with expressivity results to obtain end-to-end guarantees across several in-context regression problems. The constructions include adaptive algorithm selection while requiring no problem-specific structure beyond sufficient architectural capacity.

5.1 Generalization guarantee for pretraining

  • Theorem 20 provides a high-probability excess-loss guarantee for constrained empirical-risk minimization over clipped transformers trained on iid ICL instances.Clipping controls the transformer’s Lipschitz constant with respect to its parameters and is chosen large enough not to alter relevant inputs.Source: §5.1 · pp. 14–15

5.2 Examples of pretraining for in-context regression problems

  • Pretraining yields near-optimal in-context linear regression, with excess risk reaching e O(dσ2/N) once n ≥e O(κ2N/σ2), up to logarithmic factors.The result uses L = O(κ log(κN/σ)) layers, three heads, attention only, and mild norm bounds; the pretraining term is e O(κ2d2/n).Source: §5.2 · p. 15
  • The same generalization argument extends pretrained transformers to sparse linear regression, mixed-noise linear regression with adaptive ridge selection, and realizable logistic regression.The mixed-noise construction uses post-ICL validation to adapt regularization to each input sequence, while the overall results require only sufficiently large depth, heads, and weight norms rather than encoded problem structure.Source: §5.2 · p. 16

6 Experiments

Experiments show that the theoretical transformer architecture performs strong in-context learning and can select algorithms across heterogeneous tasks, while decoder-based transformers exhibit the same selection ability at every token.

6.1 In-context learning and algorithm selection

  • The experiments evaluate base and mixture training on linear, noisy linear, sparse linear, and classification tasks against standard in-context machine-learning baselines.Mixture modes include linear model plus linear classification and noisy linear models with σ ∈ {0.1, 0.25, 0.5, 1}; evaluation uses d = 20.Source: §6.1 · p. 17
  • Transformers match the best baseline on four of five base tasks, while sparse-regression transformers outperform least squares and match Lasso for some λ choices.This supports strong in-context learning for the encoder-based architecture used in the theory.Source: Results · p. 17
  • A single mixture-trained transformer approaches the strongest baseline simultaneously across noisy regression with different noise levels and across regression versus classification.Single-task transformers perform well on their training task but suboptimally on the other, supporting in-context algorithm selection.Source: Results · p. 18, Figure 5 · p. 18, Figure 4b

6.2 Decoder-based architecture & details for Figure 2

  • At every token, decoder-based TF_alg_select nearly matches the strongest baseline on both tasks simultaneously, including ridge regression with optimal λ and task-specific least-squares or logistic-regression baselines.Single-task transformers remain suboptimal on the other task, and the effect persists across all evaluated tokens.Source: §6.2 · p. 18

7 Conclusion … B.2 In-context learning with decoder-based transformers

This section concludes the paper’s study of transformer-based in-context learning and algorithm selection, while providing the theoretical and architectural foundations supporting these results. It also defines decoder-based transformers and their tokenized ICL training setup.

7 Conclusion

  • The work establishes strong in-context algorithm selection through explicit theoretical constructions and experiments, while identifying several directions for extending complex ICL procedures.Suggested directions include new selection mechanisms, Bayes-optimal ICL on other problems, interpretability, and further analyses of pretraining.Source: §7 · pp. 18–19

A Technical tools

  • The concentration appendix defines sub-Gaussian and sub-exponential variables and vectors, while recording Gaussian singular-value and covariance-concentration lemmas used in the proofs.These are presented as standard technical tools, including results attributed to [84, Theorem 6.1] and [82, Theorem 4.6.1].Source: A Technical tools · p. 26, §A.1 · p. 26

A.2 Approximation theory

  • Smooth functions can be uniformly approximated by finite sums of ReLUs, with neuron count and coefficient bounds controlled by dimension, smoothness, and approximation accuracy.Proposition A.1 gives M ≤ C(k)Cℓ^2/εapprox^2 and C ≤ C(k)Cℓ for approximation on [−R,R]^k.Source: §A.2 · p. 28

A.3 Optimization

  • The technical appendix records standard gradient-descent and proximal-gradient convergence results for smooth strongly convex objectives and convex composite objectives.The gradient-descent rate depends on the condition number κ = β/α, while proximal-gradient objectives and gradient mappings are decreasing.Source: §A.3 · pp. 28–29

A.4 Uniform convergence

  • Uniform concentration is established for sub-Gaussian and sub-exponential random processes under covering-number and increment-control assumptions.The result uses generalized Dudley entropy bounds and yields probability at least 1 − δ in the sub-exponential case.Source: §A.4 · pp. 29–30

A.5 Useful properties of transformers

  • Parallel transformer sequence-to-sequence functions can be joined into a single transformer while controlling heads, layers, hidden dimensions, and parameter norms.The one-layer and multi-layer constructions combine component transformers with additive norm bounds and no more than the larger layer count.Source: §A.5 · p. 31

B Extension to decoder-based architecture

  • The decoder-based extension adopts the same sequence-of-input-vectors setting, representing N vectors in R^D as an input matrix H.The passage introduces decoder transformers and retains the standard ReLU activation.Source: B Extension to decoder-based… · p. 32

B.1 Decoder-based transformers

  • Decoder-based transformers replace encoder attention with causal masked attention, so each token’s computation cannot access later tokens.Each transformer layer combines a masked attention layer with an MLP layer; the construction uses normalized ReLU attention weights.Source: §B.1 · p. 32

B.2 In-context learning with decoder-based transformers

  • For decoder ICL, covariates and responses occupy alternating tokens, preventing each covariate token from seeing its corresponding response before prediction.The decoder reads the final missing-label position to predict y_N+1, while the alternating format avoids information leakage because losses are evaluated at every token.Source: §B.2 · p. 33
  • Decoder ICL training can sum losses over all N+1 missing-label predictions, potentially requiring fewer pretraining sequences, with generalization-bound analysis supporting this benefit [46].The prediction at each odd token supplies a response estimate for its corresponding example, including the test position.Source: §B.2 · p. 33

B.3 Results · B.4 Proof of Proposition B.1 · C Proofs for Section 3.5 and additional results

The appendix extends encoder-transformer in-context gradient descent constructions to decoder transformers by converting input formats and preserving the relevant algorithmic computation. It also proves the stated two-layer decoder input-conversion proposition.

B.3 Results

  • A 2-layer decoder transformer converts input format (15) into format (3), enabling encoder-based ICL constructions to transfer to decoder transformers.The converted format contains format (3) as a submatrix after restricting specified rows and columns.Source: §B.3 · pp. 33–34
  • Proposition B.1 guarantees a 2-layer decoder transformer with 3 heads per layer, hidden dimension 2, and |||θ||| ≤12 for this conversion.Source: §B.3 · p. 33
  • The decoder construction generalizes Theorem 13 by using the final token’s scratch pad to store gradient-descent iterates.With modified MSK normalization 1/((i −1) ∨1), it performs ICGD at every prediction token and extends the constructions in Sections 3 and 4.2 to decoder transformers.Source: §B.3 · p. 34

B.4 Proof of Proposition B.1

  • The proof of Proposition B.1 constructs the decoder parameters in two steps and verifies that the resulting matrices satisfy the required norm bounds.Source: §B.4 · pp. 34–35
  • The proof exploits the alternating structure of the input sequence, where odd positions contain x_(i−1)/2 and even positions contain zero vectors.Source: §B.4 · p. 34
  • Combining the two parameter-construction steps yields the decoder transformer asserted by Proposition B.1.Source: §B.4 · p. 35

C.1 Proximal gradient descent for regularized convex losses … D.1 Proof of Theorem 4

The paper shows that transformers can approximate proximal and gradient-based optimization procedures in context, including regularized convex losses and ridge regression. The constructions rely on attention-layer implementations, non-expansive error control, and explicit regularizer-handling mechanisms.

C.1 Proximal gradient descent for regularized convex losses

  • Transformers approximate proximal gradient descent for regularized convex losses in both parameter and prediction space under approximable loss derivatives and proximal operators.The construction uses (L + 1) layers and provides explicit weight-norm bounds.Source: §C.1 · pp. 36–37
  • For ℓ1, ℓ2, and ℓ∞-constraint regularizers, the corresponding proximal operators are exactly implementable by MLPs with explicit hidden dimensions and norm bounds.The stated hidden dimensions are 4d for ℓ1 and 2d for both ℓ2 and ℓ∞ projection.Source: §C.1 · p. 36

C.4 Proof of Lemma 14

  • The PGD error-composition lemma follows from the non-expansiveness of the proximal-gradient operator, provided η ≤ 2/L_f and the iteration count satisfies L ≤ R/(2ε).The proof proceeds by induction on the iterates.Source: §C.1 · p. 37, §C.4 · p. 40

C.2 Approximating a single GD step

  • A single attention layer approximates one gradient-descent step when the loss derivative is approximable by a sum of ReLUs, with parameter norm bounded by 2 + R + 2ηC.Stacking such layers yields the iterative constructions used later.Source: §C.2 · p. 37

C.3 Proof of Theorem 13

  • The proof of Theorem 13 stacks the same attention layer for L iterations and uses induction to control accumulated approximation error in parameter and prediction outputs.The construction adds a final two-head layer for prediction, while preserving explicit weight-norm control.Source: §C.3 · pp. 38–39

C.5 Convex ICGD with ℓ2 regularization

  • With ℓ2 regularization, adding the regularizer contribution to each layer’s output exactly implements the regularizer gradient, and an attention-only transformer approximates L steps of ICGD.The transformer admits norm bound |||θ||| ≤2 + R + (2C + λ)η.Source: §C.5 · pp. 40–41

D Proofs for Section 3.1

  • For ridge regression, strong convexity and smoothness give standard gradient-descent convergence, while Corollary C.1 supplies an attention-only transformer implementing the required regularized iterations.The square-loss derivative is exactly representable by two ReLUs, enabling zero approximation error for the loss derivative.Source: §D.1 · pp. 41–42

D.2 Statistical analysis of in-context least squares … E.1 Proof of Theorem 7

The paper establishes statistical guarantees for in-context least squares and transfers them to transformer constructions, then extends the proof framework to ridge regression, generalized linear models, and gradient-descent-based GLM prediction. These results rely on well-conditioned data, concentration, clipping, and explicit transformer implementations with controlled approximation error and complexity.

D.2 Statistical analysis of in-context least squares

  • Under Assumption A, clipped in-context least squares achieves a fast excess-risk guarantee once N ≥ O(dK^4 log(1/δ)).The proof uses covariance concentration of whitened sub-Gaussian features and requires only second-moment control of the noise through clipping.Source: §D.2 · pp. 42–43

D.3 Proof of Corollary 5

  • The least-squares guarantee implies a transformer construction on well-conditioned datasets, with dataset condition number at most 4κ; under canonical parameters, the requirement simplifies to N ≥ O(d + log(1/δ)).Source: §D.3 · pp. 45–46

D.4 Proof of Corollary 6

  • The ridge-regression construction achieves nearly Bayes-optimal prediction risk by controlling its approximation to the population ridge solution on a high-probability good event.The proof separately bounds prediction risk on good and bad events, with P(E^c) ≤ δ when N ≥ 8 log(12/δ), and analyzes the ridge operator using singular-value decomposition.Source: §D.4 · pp. 46–48

E Proofs for Section 3.2

  • The generalized-linear-model proof assumes a uniformly positive derivative g′(t) ≥ μ_g on a bounded interval and sub-Gaussian covariates, responses, and model outputs.It also assumes independent samples and a lower-bound condition on the feature distribution.Source: E Proofs for Section 3.2 · p. 48

E.1 Proof of Theorem 7

  • Theorem 7 is proved by implementing T steps of gradient descent with an attention-only transformer, provided the iterates remain within the parameter bound 2Tε_g ≤ B_w.The construction first invokes the gradient-descent simulation theorem and then adds one attention layer for prediction.Source: §E.1 · pp. 48–49
  • After T gradient-descent steps, the optimization error contracts as exp(−T/(2κ))‖w_GLM‖² ≤ exp(−T/(2κ))·B_w².This contraction defines the intermediate error ε_o used to select T and the gradient approximation tolerance ε_g.Source: §E.1 · p. 49
  • The final transformer predicts g(⟨x_{N+1}, w_GLM⟩) within ε on every input satisfying the stated conditions, while inheriting the norm bound from Theorem 13.The last layer appends the predicted response to the transformer state and combines optimization and gradient-approximation errors.Source: §E.1 · p. 49

E.2 Proof of Theorem 8 … F.1 Proof of Theorem 10

The proofs establish the GLM concentration, optimization, and realizable guarantees underlying Theorem 8, then construct a transformer implementing proximal gradient descent for Lasso in Theorem 10.

E.3 Proof of Theorem E.1 (a)

  • For N ≥ C1 · d, with probability at least 1 − 2e−N/C1, uniformly over w ∈ B2(Bw), 1/8µgµx ≤ λmin(∇2bLN(w)) ≤ λmax(∇2bLN(w)) ≤ 8LgK2_x.Source: §E.2 · p. 49

E.2 Proof of Theorem 8

  • Theorem 8’s transformer outputs a predictor with uniform GLM prediction error at most ε on the good event while producing a bounded parameter estimate.The construction uses the transformer output to encode both the prediction and projected parameter estimate, with guarantees established under the stated good events.Source: §E.2 · p. 50

E.4 Proof of Theorem E.1 (b)

  • The proof of Theorem E.1(b) applies covering and sub-exponential-process concentration to control the empirical GLM gradient uniformly over the parameter ball.It verifies sub-exponential tails and Lipschitz increments before applying the concentration proposition and a union bound.Source: §E.4 · p. 54

E.5 Proof of Theorem E.1 (c)

  • Under strong convexity and N ≥ C2d, the empirical minimizer lies within the prescribed parameter ball with probability at least 1 − exp(−N/C2).The argument conditions on the Hessian event and uses a statistical-error event to show the unconstrained empirical minimizer is inside the ball.Source: §E.5 · p. 54

E.6 Proof of Theorem E.1 (d) & (e)

  • The population GLM loss is strongly convex, yielding the non-realizable gradient guarantee and, in the realizable setting, the corresponding guarantee at wGLM.The realizable argument uses E[y|x] = g(⟨x, β⟩), β⋆ = β, and ∇bLN(wGLM) = 0.Source: §E.6 · p. 55

F.1 Proof of Theorem 10

  • The resulting proximal-gradient iterate satisfies the standard convergence guarantee for all t ≥ 1 under the stated Lasso setup.This final step invokes Proposition A.3 after establishing the transformer implementation of proximal gradient descent.Source: §F.1 · pp. 55–56
  • For Lasso, the proof implements proximal gradient descent with η = β−1 and T steps using a transformer with (T + 1) layers, 2 heads per layer, and hidden dimension 2d.The construction approximates the square-loss gradient and ℓ1 proximal operator by sums of ReLUs, with weight norm bounded by 10R + (8 + 2λN)β−1.Source: §F.1 · pp. 55–56

F.2 Sharper convergence analysis of proximal gradient descent for Lasso … F.5 Proof of Theorem 11

These appendices establish sharper proximal-gradient convergence for Lasso under deterministic well-conditioning, derive supporting estimation inequalities, and use them to prove transformer in-context Lasso guarantees. Theorem 11 follows by verifying the conditions with high probability for Gaussian designs and controlling prediction error on a good event.

F.2 Sharper convergence analysis of proximal gradient descent for Lasso

  • Under Assumption C, proximal gradient descent with stepsize η = 1/β and initialization w0 = 0 satisfies the sharper convergence guarantee stated in Theorem F.1.Assumption C combines restricted strong convexity, approximate generation by an s-sparse linear model, and a sample-size condition; it is deterministic and is later shown to hold with high probability in the standard sparse model.Source: §F.2 · p. 56
  • Theorem F.2 converts the sharper proximal-gradient guarantee into a transformer that approximately implements in-context Lasso whenever Assumption C holds.The construction uses a transformer whose output is an inner product with an approximate Lasso iterate and applies to arbitrary N, d, s and fixed parameter choices covered by the theorem.Source: §F.2 · p. 57

F.3 Basic properties for Lasso

  • The basic Lasso lemmas relate objective suboptimality to parameter and prediction errors, including an ℓ1–ℓ2 error relation for sparse targets and a quadratic growth bound.The gap-to-parameter result uses the universal constant C = 120, while the growth lemma bounds 1/(2N)∥X(w − w_lasso)∥2^2 by the objective gap.Source: §F.3 · pp. 57–58

F.4 Proof of Theorem F.1

  • The proof of Theorem F.1 combines the sparse-model inequality, restricted strong convexity, growth, and proximal-gradient descent arguments to drive the iterate gap below ε after sufficiently many iterations.The analysis separates an initial phase from the later contraction regime and bounds the Lasso solution norm using the target norm and the sparsity-dependent term.Source: §F.4 · pp. 58–59

F.5 Proof of Theorem 11

  • For Gaussian random designs, the proof of Theorem 11 establishes restricted strong convexity and the remaining parts of Assumption C on a jointly defined good event.It selects α, β, ρ, and ν to match the Lasso regularization and invokes the Gaussian-design RSC proposition together with auxiliary concentration lemmas.Source: §F.5 · p. 60
  • Theorem 11’s proof applies the transformer Lasso result on the good event and decomposes conditional prediction error into estimation and bad-event contributions.The constructed prediction is clipped from the inner product with the approximate Lasso iterate, and the gap-to-parameter bound controls the estimation component.Source: §F.5 · pp. 60–61
  • The good event has probability at least 1 − δ when N ≥ 32 log(12/δ), completing the high-probability proof of Theorem 11.The argument bounds the bad-event probability using the marginal Gaussian distribution of y_N+1 and combines it with the event-probability estimates.Source: §F.5 · p. 61

G Gradient descent on two-layer neural networks … H Proofs for Section 4

The paper extends in-context gradient descent beyond convex problems by constructing transformers that approximate multi-step gradient descent for two-layer neural networks. The accompanying proofs decompose one update into attention, MLP, gradient-aggregation, and projection operations, then establish convergence and trajectory-closeness guarantees for inexact projected gradient descent.

G Gradient descent on two-layer neural networks

  • Transformers can implement in-context gradient descent on two-layer neural-network risks, extending the framework beyond the convex setting.The construction assumes C4-smooth activation and loss functions and a bounded parameter domain with an MLP-implementable projection.Source: G Gradient descent on two-layer neural… · p. 62
  • Under the theorem’s assumptions, each transformer layer pair approximately performs a projected gradient-descent update from the in-context data.The bounds hide constants depending on network width, radius parameters, and the smoothness of the activation and loss functions.Source: G Gradient descent on two-layer neural… · p. 62
  • The resulting (2L)-layer transformer approximates the true multi-step gradient-descent trajectory for two-layer neural networks.This follows as Corollary G.1 under the same setting as Theorem G.1.Source: G Gradient descent on two-layer neural… · p. 62
  • Trajectory error can accumulate exponentially with the iteration count because nonconvex gradient descent lacks the stability property used in the convex analysis.A small error at each step may therefore produce substantially different trajectories.Source: G Gradient descent on two-layer neural… · p. 62

G.2 Proof of Lemma G.1

  • Lemma G.1 proves convergence and trajectory closeness for inexact projected gradient descent on a convex domain with a Lipschitz gradient.Its trajectory bound follows from non-expansive projection and the (1 + ηL_f)-Lipschitz gradient-update operator.Source: G Gradient descent on two-layer neural… · p. 63, §G.2 · p. 67

G.1 Proof of Theorem G.1

  • The proof of Theorem G.1 implements one gradient step by combining attention for prediction and gradient aggregation with MLP evaluation of loss derivatives and projection.The sequence is prediction, loss-gradient computation, attention-based update w → w − η∇L_N(w), and projection into W.Source: §G.1 · p. 63
  • The proof constructs the transformer by approximating the activation, loss derivative, and derivative-product computations, with errors controlled through smoothness and bounded-domain assumptions.These approximations are assembled into the attention and MLP layers required by the single-step proposition.Source: §G.1 · p. 64

H.1 Proof of Proposition 15 … H.4 Linear correlation test and application

This section develops and proves transformer mechanisms for adaptive predictor selection, regression, and classification, including validation-based selection and linear correlation testing. It also establishes exact or efficient constructions and statistical guarantees under stated approximation, conditioning, and distributional assumptions.

H.1 Proof of Proposition 15

  • A 3-layer transformer selects a convex combination of candidates whose validation loss is within γ of the best candidate, with the construction split into evaluation and selection layers.The evaluation layer computes approximate validation losses, while the selection layer assigns weights only to candidates within the γ tolerance.Source: §H.1 · pp. 68, 69, 72

H.2 Proof of Theorem 16

  • Under well-conditioned inputs with bounded solutions, an L-layer transformer produces ridge estimators approximating the ridge solutions for any specified regularization sequence.The construction uses an in-context gradient-descent component followed by the predictor-selection transformer, with layer depth and width bounds stated in the theorem.Source: §H.2 · pp. 72–73

H.2.1 Proof of Theorem H.3

  • For squared validation loss, a 2K-head attention layer computes each candidate’s validation loss and appends these values to every sequence representation.This evaluation layer is combined with the selection construction to obtain the squared-loss version of the predictor-selection theorem.Source: §H.2.1 · pp. 73–74

H.3.1 Proof of Lemma 18

  • The binary type-check function ψ can be represented as a linear combination of 6 ReLUs and implemented by an attention layer.The construction uses six attention heads with bounded coefficient parameters.Source: §H.3.1 · pp. 74–75
  • The thresholded binary test is exactly implemented by a two-layer attention-only transformer with at most 6 heads per layer and parameter norm O(1/ε).This binary test is the component used to select between the logistic-regression and least-squares predictors in Proposition H.4.Source: §H.3.1 · p. 75, §H.3.2 · p. 76

H.3.2 Formal statement and proof of Proposition 19

  • A single attention-only transformer adaptively approximates in-context logistic regression on well-conditioned binary data and least squares on well-conditioned non-binary regression data with probability at least 1 − exp(−cN).The regression guarantee assumes the label marginal is (C, ε0)-not-concentrated around {0, 1}; the classification output ε-approximates in-context logistic regression, while the regression output ε-approximates in-context least squares.Source: §H.3.2 · pp. 76–77

H.4 Linear correlation test and application

  • The linear correlation test Ψlin is exactly implemented by a 3-layer attention-only transformer with at most 2 heads per layer and parameter norm O(1 + λ2_w)^2.The construction operates on input dimension D ≥ 2d + 4 and outputs Ψlin(D) at every token.Source: §H.4 · p. 77
  • Under Assumption D, Ψlin has high power: with probability at least 1 − δ, it outputs 1 for sufficiently large ||w⋆_P||2 and 0 when ||w⋆_P||2 ≤ λmin/(10λmaxB⋆_w).Assumption D includes bounded covariance eigenvalues and sub-Gaussianity conditions; the test distinguishes sufficiently high and sufficiently low linear signal.Source: §H.4 · p. 78
  • Composing Ψlin with least-squares prediction yields confident linear regression: the transformer predicts within ε of least squares when the signal is high and abstains with output 0 when it is low.In the high-signal case, bψ = 1 under the additional least-squares well-conditioning condition; in the low-signal case, by_N+1 = 0 and bψ = 0 with probability at least 1 − δ.Source: §H.4 · p. 79

I Proof of Theorem 17: Noisy linear model with mixed noise levels … J.2 Proof of Theorem 20

This section develops the theory of adaptive in-context prediction for mixed-noise linear and generalized linear models, including the Bayes predictor and adaptive link-function selection through validation. It also establishes transformer Lipschitzness and the concentration and covering-number tools used to prove the relevant results.

I Proof of Theorem 17: Noisy linear model with mixed noise levels

  • Under N ≥ 0.1d and D ≥ Θ(Kd), a transformer achieves the strengthened near-Bayes guarantee for mixtures of K noisy linear models with mixed noise levels.The proof reduces mixture performance to near-Bayes performance on each component and uses validation to select among candidate predictors.Source: I Proof of Theorem 17: Noisy linear… · pp. 79–80

I.1 Proof of Theorem I.1

I.2 Derivation of the exact Bayes predictor

  • For mixed-noise linear models, the exact Bayes predictor is a data-dependent weighted combination of K ridge estimators, with weights requiring nontrivial determinant-based computations.Each component posterior mean is a ridge estimator, while the mixture weights are posterior probabilities depending on the observed design matrix.Source: §I.2 · pp. 82–83

I.3 Useful lemmas

  • The supporting lemmas establish the needed training-design and validation-loss concentration bounds by splitting sample-size regimes and applying Bernstein’s inequality with a union bound.These arguments condition on suitable good events and control candidate prediction errors before aggregation.Source: §I.3 · pp. 83–84

I.4 Generalized linear models with adaptive link function selection

  • For generalized linear models with one of K non-decreasing, C2-smooth links, a transformer aggregates candidate predictors so validation loss is within γ of the best candidate.The construction uses separate training losses for each link and squared validation loss to select among them, with N ≥ O(d) and input dimension D = Θ(dK).Source: §I.4 · pp. 84–85

J.1 Lipschitzness of transformers

  • A single MLP layer is (BR)-Lipschitz in parameters and (1 + B^2)-Lipschitz in inputs, while an attention layer is (B^2R^3)- and (1 + B^3R^2)-Lipschitz, respectively.Consequently, the full transformer is B_Θ-Lipschitz in parameters and L_H-Lipschitz in hidden states, with explicit bounds given in Corollary J.1.Source: §J.1 · pp. 86–88
  • J.1 bounds transformer Lipschitz behavior using parameter and attention-norm controls for fixed head count and hidden dimension.The proof repeatedly invokes Lemmas J.1–J.2 and Corollary J.1, with attention outputs bounded by R := R + B3R3.Source: §J.1 · pp. 88–89

J.2 Proof of Theorem 20

  • J.2 proves Theorem 20 for a general ICL loss by verifying entropy-covering conditions and applying Proposition A.4’s uniform concentration result.The argument uses a covering-number bound from [84, Example 5.8] and concludes concentration with probability at least 1 − ξ.Source: §J.2 · p. 89

J.3 Proof of Theorem 21 … J.6 Proof of Theorem 24

These sections prove Theorems 21–24 by invoking preceding transformer constructions and Theorem 20, establishing predictions and high-probability TF-ERM guarantees under specified size bounds. The arguments also show that clipping preserves outputs on the relevant good events, with Theorem 24 following by analogous reasoning.

J.3 Proof of Theorem 21

  • Theorem 21’s construction yields predictions for every distribution satisfying Assumption A and every N ≥e O(d), via Corollary 5.The proof then applies Theorem 20 to derive a high-probability guarantee for the corresponding TF-ERM solution.Source: §J.3 · p. 90
  • For Theorem 21, replacing the transformer by its clipped version preserves its output on the good event because all intermediate outputs have bounded norm.The construction also retains the stated size bounds, with logarithmic factor ι = e O(1).Source: §J.3 · p. 90

J.4 Proof of Theorem 22

  • For almost surely every P ∼ π, the construction yields a transformer returning the specified prediction on the good event E0, which satisfies P(E0) ≥1−δ.Applying Theorem 20 gives the TF-ERM guarantee, and clipping does not modify the transformer output on E0 because all intermediate outputs satisfy ∥·∥2,∞≤R.Source: §J.4 · pp. 90–91
  • Theorem 22’s proof combines the good- and bad-event bounds and takes expectation over P ∼π before applying the TF-ERM result.The construction uses bounded logarithmic factor ι = e O(1).Source: §J.4 · pp. 90–91

J.5 Proof of Theorem 23

  • Theorem 23 uses Theorem 17 and Theorem I.1 to construct a transformer with input dimension D = Θ(Kd), M = O(K) heads, and D′ = O(K2).Theorem 20 then gives the stated high-probability TF-ERM guarantee, and clipping preserves outputs on the relevant good event.Source: §J.5 · p. 91
  • Theorem 23’s proof establishes the required prediction property before transferring it to the TF-ERM solution through Theorem 20.The clipped construction has the same output on the good event, with ι = e O(1).Source: §J.5 · p. 91

J.6 Proof of Theorem 24

  • Theorem 24 follows by the same proof strategy as Theorems 22 and 23, substituting the layer, head, and weight-norm bounds from Theorem 8 and Corollary 9.Its proof therefore reuses the preceding construction and guarantee arguments rather than introducing a separate mechanism.Source: §J.6 · p. 92

K Experimental details · K.1 Additional details for Section 6.1 · K.2 Computational resource

The experiments use a 12-layer encoder-only transformer with standardized optimization and evaluation procedures, plus a specialized objective for sparse linear regression. Training compares the learned models against classical baselines and required approximately five GPU-days across eight A100 GPUs.

K.1 Additional details for Section 6.1

  • The transformer has 12 encoder layers with 8 attention heads, hidden dimensions 64, ReLU activations, Layer Normalization, and trainable affine read-in/read-out layers.The read-in maps inputs to 64-dimensional states, while the read-out maps hidden states to scalar predictions.Source: §K.1 · p. 92
  • Training uses fresh minibatches, square loss for regression, logistic loss for classification, Adam with learning rate 10^-4, and generally 300K steps with batch size 64.Mixture-mode training uses K minibatches per step, each with batch size 64.Source: §K.1 · p. 92
  • Sparse linear regression additionally minimizes a coefficient-recovery loss, targeting the first five of 10 in-context tokens and using λ = 0.1 for 2M steps.This augmentation was introduced because the standard training objective alone did not outperform least squares, whereas Lasso achieved lower test loss.Source: §K.1 · p. 92
  • Evaluation uses 6400 test instances, reporting square loss for regression and classification error for binary prediction, with standard-deviation error bars shown where variability is appreciable.Means are reported throughout; several figures omit error bars because standard deviations are small relative to marker widths.Source: §K.1 · p. 92
  • Baselines include least squares, logistic regression, averaging, 3-nearest neighbors, ridge regression with λ values 0.005 and 0.125, and Lasso with λ in {1, 0.1, 0.01, 0.001}.The ridge values are Bayes-optimal for noise levels 0.1 and 0.5 under the specified noisy linear model.Source: §K.1 · pp. 92–93
  • The ridge_analytical curve evaluates expected ridge risk over 20 geometrically spaced regularization values, while Bayes_err_{1,2} report task-specific risks for λ1 and λ2.These quantities are computed under the noisy linear model using analytical formulae with Monte Carlo simulations.Source: §K.1 · p. 93

K.2 Computational resource

  • All experiments ran on 8 Nvidia Tesla A100 GPUs with 40GB memory, totaling approximately 5 days of GPU time; the largest run took about one day on one GPU.Source: §K.2 · p. 93

Related papers

Loading 2306.04637v2…