Title: Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection
Venue: NeurIPS 2023
Accepted: yes
Keywords: in-context learning

transformers

deep learning theory

learning theory
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

Abstract

Neural sequence models based on the transformer architecture have demonstrated remarkable 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 incontext 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 in-context algorithm selection , akin to what a statistician can do in real life-A 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

Large neural sequence models have demonstrated remarkable in-context learning (ICL) capabilities [12], where models can make accurate predictions on new tasks when prompted with training examples from the same task, in a zero-shot fashion without any parameter update to the model. A prevalent example is large language models based on the transformer architecture [84], which can perform a diverse range of tasks in context when trained on enormous text [12, 90]. Recent models

̂

Data

Which distribution?

dist 1

dist K

y

N

+1

= /u1D5A0/u1D5C5/u1D5C0

k

⋆

(

D

,

x

N

+1

)

TF

/u1D5A0/u1D5C5/u1D5C0

1

/u1D5A0/u1D5C5/u1D5C0

K

……

Mechanism 2: Pre-ICL Testing

Data

/u1D5A0/u1D5C5/u1D5C0

1

(

D

train

)

/u1D5A0/u1D5C5/u1D5C0

K

(

D

train

)

……

TF

/u1D5AB/u1D5C8/u1D5CC/u1D5CC

1

(

D

val

)

/u1D5AB/u1D5C8/u1D5CC/u1D5CC

K

(

D

val

)

……

Which loss is the smallest?

y

N

+1

= /u1D5A0/u1D5C5/u1D5C0

k

⋆

(

D

train

,

x

N

+1

)

Train-validation split

Mechanism 1: Post-ICL Validation

Data 1

(Regression)

Transformer

/u1D5AB/u1D5C2/u1D5C7/u1D5B1/u1D5BE/u1D5C0(

D

,

x

N

+1

)

/u1D5AB/u1D5C8/u1D5C0/u1D5B1/u1D5BE/u1D5C0(

D

,

x

N

+1

)

Data 2

(Classification)

y

N

+1

:

Example 2: Regression + Classification

Data 1

(Reg w/noise

)

σ

1

Transformer

/u1D5B1/u1D5C2/u1D5BD/u1D5C0/u1D5BE

λ

1

(

D

,

x

N

+1

)

/u1D5B1/u1D5C2/u1D5BD/u1D5C0/u1D5BE

λ

2

(

D

,

x

N

+1

)

Data 2

(Reg w/noise

)

σ

2

y

N

+1

:

Example 1: Ridge with different

λ

̂

̂

̂

in this paradigm such as GPT-4 achieve surprisingly impressive ICL performance that makes them akin to a general-purpose agent in many aspects [65, 14]. Such strong capabilities call for better understandings, which a recent line of work tackles from various aspects [49, 94, 28, 72, 15, 57, 64].

Recent pioneering work of Garg et al. [31] proposes an interpretable and theoretically amenable setting for understanding ICL in transformers. They perform ICL experiments where input tokens are real-valued (input, label) pairs generated from standard statistical models such as linear models (and the sparse version), neural networks, and decision trees. Garg et al. [31] find that transformers can learn to perform ICL with prediction power (and fitted functions) matching standard machine learning algorithms for these settings, such as least squares for linear models, and Lasso for sparse linear models. Subsequent work further studies the internal mechanisms [2, 86, 18], expressive power [2, 32], and generalization [47] of transformers in this setting. However, these works only showcase simple mechanisms such as regularized regression [31, 2, 47] or gradient descent [2, 86, 18], which are arguably only a small subset of what transformers are capable of in practice; or expressing universal function classes not specific to ICL [89, 32]. This motivates the following question:

How do transformers learn in context beyond implementing simple algorithms?

This paper makes steps on this question by making two main contributions: (1) We unveil a general mechanismin-context algorithm selection -by which a single transformer can adaptively select different 'base' ICL algorithms to use on different ICL instances , without any explicit prompting of the right algorithm to use in the input sequence. For example, a transformer may choose to perform ridge regression with regularization λ 1 on ICL instance 1, and λ 2 on ICL instance 2 (Figure 2); or perform regression on ICL instance 1 and classification on ICL instance 2 (Figure 5). This adaptivity allows transformers to achieve much stronger ICL performance than the base ICL algorithms. We both prove this in theory, and demonstrate this phenomenon empirically on standard transformer architectures. (2) Along the way, equally importantly, we present a comprehensive theory for ICL in transformers by establishing end-to-end quantitative guarantees for the expressive power, in-context prediction performance, and sample complexity of pretraining . These results add upon the recent line of work on the statistical learning theory of transformers [97, 89, 27, 39], and lay out a foundation for the intriguing special case where the learning targets are themselves ICL algorithms .

A detailed summary of our contributions is as follows.

We prove that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, convex risk minimization for learning generalized linear models (such as logistic regression), and gradient descent for two-layer neural networks (Section 3). Our constructions admit mild bounds on the number of layers, heads, and weight norms, and achieve near-optimal prediction power on many in-context data distributions.

Technically, the above transformer constructions build on a new efficient implementation of incontext gradient descent (Appendix D), which could be broaderly applicable. For a broad class of smooth convex empirical risks over the in-context training data, we construct an ( L +1) -layer transformer that approximates L steps of gradient descent. Notably, the approximation error accumulates only linearly in L , utilizing a stability-like property of smooth convex optimization.

We prove that transformers can perform in-context algorithm selection (Section 4). We construct two algorithm selection mechanisms: Post-ICL validation (Section 4.1), and Pre-ICL testing

(a) Noisy linear reg with noise σ 1 (b) Noisy linear reg with noise σ 2 (c) Task 1 vs. task 2 at token 20

0

10

20

30

40

in-context examples

0.0

0.2

0.4

0.6

0.8

1.0

square loss

TF_alg_select

TF_noise_1

TF_noise_2

ridge_lam_1

ridge_lam_2

0

10

20

30

40

in-context examples

0.4

0.6

0.8

1.0

1.2

1.4

TF_alg_select

TF_noise_1

TF_noise_2

ridge_lam_1

ridge_lam_2

0.10

0.15

0.20

0.25

0.30

noisy_reg_noise_1

0.6

0.8

1.0

1.2

1.4

noisy_reg_noise_2

TF_alg_select

TF_noise_1

TF_noise_2

ridge_lam_1

ridge_lam_2

Bayes_err_noise_1

Bayes_err_noise_2

(Section 4.2). For both mechanisms, we provide general constructions as well as concrete examples. Figure 1 provides a pictorial illustration of the two mechanisms.

As a concrete application, using the post-ICL validation mechanism, we construct a transformer that can perform nearly Bayes-optimal ICL on noisy linear models with mixed noise levels (Section 4.1.1), a more complex task than those considered in existing work.

We provide the first line of results for pretraining transformers to perform the various ICL tasks above, from polynomially many training sequences (Section 5 & Appendix K).

Experimentally, we find that learned transformers indeed exhibit strong in-context algorithm selection capabilities in the settings considered in our theory (Section 6). For example, Figure 2 shows that a single transformer can approach the individual Bayes risks (the optimal risk among all possible algorithms) simultaneously on two noisy linear models with different noise levels.

Transformers as statisticians We humbly remark that the typical toolkit of a statistician contains much more beyond those covered in this work, including and not limited to inference, uncertainty quantification, and theoretical analysis. This work merely aims to show the algorithm selection capability of transformers, akin to what a statistician can do.

Related work Our work is intimately related to the lines of work on in-context learning, theoretical understandings of transformers, as well as other formulations for learning-to-learn such as metalearning. Due to limited space, we discuss these related work in Appendix A.

2 Preliminaries

We consider a sequence of N input vectors { h i } N i =1 ⊂ R D , written compactly as an input matrix H = [ h 1 , . . . , h N ] ∈ R D × N , where each h i is a column of H (also a token ). Throughout this paper, we let σ ( t ) := ReLU( t ) = max { t, 0 } denote the standard relu activation.

2.1 Transformers

We consider transformer architectures that process any input sequence H ∈ R D × N by applying (encoder-mode 2 ) attention layers and MLP layers formally defined as follows.

Definition 1 (Attention layer) . A (self-)attention layer with M heads is denoted as Attn θ ( · ) with parameters θ = { ( V m , Q m , K m ) } m ∈ [ M ] ⊂ R D × D . On any input sequence H ∈ R D × N ,

where σ : R → R is the ReLU function. In vector form,

Above, (1) uses a normalized ReLU 3 activation t ↦→ σ ( t ) /N in place of the standard softmax activation; we remark this activation is also found to work well empirically in recent studies [78, 93].

Definition 2 (MLP layer) . A (token-wise) MLP layer with hidden dimension D ′ is denoted as MLP θ ( · ) with parameters θ = ( W 1 , W 2 ) ∈ R D ′ × D × R D × D ′ . On any input sequence H ∈ R D × N ,

where σ : R → R is the ReLU function. In vector form, we have ˜ h i = h i + W 2 σ ( W 1 h i ) .

We consider a transformer architecture with L ≥ 1 transformer layers, each consisting of a selfattention layer followed by an MLP layer.

Definition 3 (Transformer) . An L -layer transformer, denoted as TF θ ( · ) , is a composition of L self-attention layers each followed by an MLP layer: H ( L ) = TF θ ( H (0) ) , where H (0) ∈ R D × N is the input sequence, and

Above, the parameter θ = ( θ (1: L ) attn , θ (1: L ) mlp ) consists of the attention layers θ ( ℓ ) attn = { ( V ( ℓ ) m , Q ( ℓ ) m , K ( ℓ ) m ) } m ∈ [ M ( ℓ ) ] ⊂ R D × D and the MLP layers θ ( ℓ ) mlp = ( W ( ℓ ) 1 , W ( ℓ ) 2 ) ∈ R D ( ℓ ) × D × R D × D ( ℓ ) . We will frequently consider ' attention-only ' transformers with W ( ℓ ) 1 , W ( ℓ ) 2 = 0 , which we denote as TF 0 θ ( · ) for shorthand, with θ = θ (1: L ) := θ (1: L ) attn .

We additionally define the following norm of a transformer TF θ :

In (2), the choices of the operator norm and max/sums are for convenience only and not essential, as our results (e.g. for pretraining) depend only logarithmically on | | | θ | | | .

2.2 In-context learning

In an in-context learning (ICL) instance, the model is given a dataset D = { ( x i , y i ) } i ∈ [ N ] iid ∼ P and a new test input x N +1 ∼ P x for some data distribution P , where { x i } i ∈ [ N ] ⊆ R d are the input vectors, { y i } i ∈ [ N ] ⊆ R are the corresponding labels (e.g. real-valued for regression, or { 0 , 1 } -valued for binary classification), and x N +1 is the test input on which the model is required to make a prediction. Different from standard supervised learning, in ICL, each instance ( D , x N +1 ) is in general drawn from a different distribution P j , such as a linear model with a new ground truth coefficient w ⋆,j ∈ R d . Our goal is to construct fixed transformer to perform ICL on a large set of P j 's.

We consider using transformers to perform ICL, in which we encode ( D , x N +1 ) into an input sequence H ∈ R D × ( N +1) . In our theory, we use the following format, where the first two rows contain ( D , x N +1 ) (zero at the location for y N +1 ), and the third row contains fixed vectors { p i } i ∈ [ N +1] with ones, zeros, and indicator for being the train token (similar to a positional encoding vector):

We will choose D = Θ( d ) , so that the hidden dimension of H is at most a constant multiple of d . We then feed H into a transformer to obtain the output ˜ H = TF θ ( H ) ∈ R D × ( N +1) with the same shape, and read out the prediction ̂ y N +1 from the ( d +1 , N +1) -th entry of ˜ H = [ ˜ h i ] i ∈ [ N +1] (the entry corresponding to the missing test label): ̂ y N +1 = read y ( ˜ H ) := ( ˜ h N +1 ) d +1 . The goal is to predict

̂ y N +1 that is close to y N +1 ∼ P y | x N +1 measured by proper losses. We emphasize that we consider predicting only at the last token x N +1 , which is without much loss of generality. 4

Miscellaneous setups We assume bounded features and labels throughout the paper (unless otherwise specified, e.g. when x i is Gaussian): ∥ x i ∥ 2 ≤ B x and | y i | ≤ B y with probability one. We use the standard notation X = [ x ⊤ 1 ; . . . ; x ⊤ N ] ∈ R N × d and y = [ y 1 ; . . . ; y N ] ∈ R N to denote the matrix of inputs and vector of labels, respectively. To prevent the transformer from blowing up on tail events, in all our results concerning (statistical) in-context prediction powers, we consider a clipped prediction ̂ y N +1 = ˜ read y ( ˜ H ) := clip R (( ˜ h N +1 ) d +1 ) , where clip R ( t ) := Proj [ -R,R ] ( t ) is the standard clipping operator with (a suitably large) radius R ≥ 0 that varies in different problems.

3 Basic in-context learning algorithms

We begin by constructing transformers that approximately implement a variety of standard machine learning algorithms in context, with mild size bounds and near-optimal prediction power on many standard in-context data distributions.

3.1 In-context ridge regression and least squares

Consider the standard ridge regression estimator over the in-context training examples D with regularization λ ≥ 0 (reducing to least squares at λ = 0 and N ≥ d ):

We show that transformers can approximately implement (ICRidge) (proof in Appendix F.1).

Theorem 4 (Implementing in-context ridge regression) . For any λ ≥ 0 , 0 ≤ α ≤ β with κ := β + λ α + λ , B w > 0 , and ε < B x B w / 2 , there exists an L -layer attention-only transformer TF 0 θ with

(with R := max { B x B w , B y , 1 } ) such that the following holds. On any input data ( D , x N +1 ) such that the problem (ICRidge) is well-conditioned and has a bounded solution:

TF 0 θ approximately implements (ICRidge): The prediction ̂ y N +1 = read y (TF 0 θ ( H )) satisfies

Theorem 4 presents the first quantitative construction for end-to-end in-context ridge regression up to arbitrary precision, and improves upon Akyürek et al. [2] whose construction does not give (or directly imply) an explicit error bound like (6). Further, the bounds on the number of layers and heads in (4) are mild (constant heads and logarithmically many layers).

Near-optimal in-context prediction power for linear problems Combining Theorem 4 with standard analyses of linear regression yields the following corollaries (proofs in Appendix F.3 & F.4).

Corollary 5 (Near-optimal linear regression with transformers by approximating least squares) . For any N ≥ ˜ O ( d ) , there exists an O ( κ log( κN/σ )) -layer transformer θ , such that on any P satisfying standard statistical assumptions for least squares (Assumption A), its ICL prediction ̂ y N +1 achieves

Assumption A requires only generic tail properties such as sub-Gaussianity, and not realizability (i.e., P follows a true linear model); κ, σ above denote the covariance condition number and the

noise level therein. The ˜ O ( dσ 2 /N ) excess risk is known to be rate-optimal for linear regression [38], and Corollary 5 achieves this in context with a transformer with only logarithmically many layers.

Next, consider Bayesian linear models where each in-context data distribution P = P lin w ⋆ is drawn from a Gaussian prior π : w ⋆ ∼ N (0 , I d /d ) , and ( x , y ) ∼ P lin w ⋆ is sampled as x ∼ N ( 0 , I d ) , y = ⟨ w ⋆ , x ⟩ + N (0 , σ 2 ) . It is a standard result that the Bayes estimator of y N +1 given ( D , x N +1 ) is given by ridge regression (ICRidge): ̂ y Bayes N +1 := ⟨ w λ ridge , x N +1 ⟩ with λ = dσ 2 /N . We show that transformers achieve nearly-Bayes risk for this problem, and we use

to denote the Bayes risk of this problem under prior π .

Corollary 6 (Nearly-Bayes linear regression with transformers by approximating ridge regression) . Under the Bayesian linear model above with N ≥ max { d/ 10 , O (log(1 /ε )) } , there exists a L = O (log(1 /ε )) -layer transformer such that E w ⋆ , ( D , x N +1 ,y N +1 ) [ 1 2 ( ̂ y N +1 -y N +1 ) 2 ] ≤ BayesRisk π + ε .

Generalized linear models In Appendix G, we extend the above results to generalized linear models [53] and show that transformers can approximate the corresponding convex risk minimization algorithm in context (which includes logistic regression for linear classification as an important special case), and achieve near-optimal excess risk under standard statistical assumptions.

3.2 In-context Lasso

Consider the standard Lasso estimator [82] which minimizes an ℓ 1 -regularized linear regression loss ̂ L lasso over the in-context training examples D :

We show that transformers can also approximate in-context Lasso with a mild number of layers, and can perform sparse linear regression in standard sparse linear models (proofs in Appendix H).

Theorem 7 (Implementing in-context Lasso) . For any λ N ≥ 0 , β > 0 , B w > 0 , and ε > 0 , there exists a L -layer transformer TF θ with

(where R := max { B x B w , B y , 1 } ) such that the following holds. On any input data ( D , x N +1 ) such that λ max ( X ⊤ X /N ) ≤ β and ∥ w lasso ∥ 2 ≤ B w / 2 , TF θ ( H (0) ) approximately implements (ICLasso), in that it outputs ̂ y N +1 = ⟨ x N +1 , ̂ w ⟩ with ̂ L lasso ( ̂ w ) -̂ L lasso ( w lasso ) ≤ ε .

Theorem 8 (Near-optimal sparse linear regression with transformers by approximating Lasso) . For any d, N ≥ 1 , δ > 0 , B ⋆ w , σ > 0 , there exists a ˜ O (( B ⋆ w ) 2 /σ 2 × (1 + ( d/N ))) -layer transformer θ such that the following holds: For any s and N ≥ O ( s log( d/δ )) , suppose that P is an s -sparse linear model: x i ∼ N (0 , I d ) , y i = ⟨ w ⋆ , x i ⟩ + N (0 , σ 2 ) for any ∥ w ⋆ ∥ 2 ≤ B ⋆ w and ∥ w ⋆ ∥ 0 ≤ s , then with probability at least 1 -δ (over the randomness of D ), the transformer output ̂ y N +1 achieves

The ˜ O ( s log d/N ) excess risk obtained in Theorem 8 is optimal up to log factors [62, 87]. We remark that Theorem 8 is not a direct corollary of Theorem 7; Rather, the bound on the number of layers in Theorem 8 requires a sharper convergence analysis of the (ICLasso) problem under sparse linear models (Appendix H.2), similar to [1].

3.3 Proof technique: In-context gradient descent

The constructions in Section 3.1 and 3.2 is built on the following result for approximating in-context (proximal) gradient descent on (regularized) convex losses.

Theorem 9 (ICGD; Informal version of Theorem D.1 & D.2) . For a broad class of convex losses of form w ↦→ 1 N ∑ N i =1 ℓ ( w ⊤ x i , y i ) + R ( w ) , there exists an L -layer transformer that takes in any ( D , w 0 ) and outputs ̂ w L such that ∥ ̂ w L -w L { GD , PGD } ∥ 2 ≤ O ( Lε ) , by composing L identical layers each O ( ε ) -approximating a single step of GD (so that O ( Lε ) is a linear error accumulation).

Theorem 9 is established in two main steps:

Approximating one-step of ICGD using one attention layer (Proposition E.1), which substantially generalizes that of von Oswald et al. [86] (which only does GD on square losses with a linear self-attention), and is simpler than the ones in Akyürek et al. [2] and Giannou et al. [32].

Stacking L of the above layer to approximate L steps of ICGD. Done naively, the error accumulation of this stacking operation is exponential in L in the worst case. We utilize the stability of convex gradient descent (Lemma D.1) to obtain the linear in L error accumulation in Theorem 9.

In Appendix D.3, we also give results for non-convex GD on two-layer neural nets, though with a worse (exponential in L ) error accumulation as expected.

4 In-context algorithm selection

We now show that transformers can perform various kinds of in-context algorithm selection , which allows them to implement more complex ICL procedures by adaptively selecting different 'base' algorithms on different input sequences. We construct two general mechanisms: Post-ICL validation , and Pre-ICL testing ; See Figure 1 for a pictorial illustration.

4.1 Post-ICL validation mechanism

In our first mechanism, post-ICL validation, the transformer begins by implementing a train-validation split D = ( D train , D val ) , and running K base ICL algorithms on D train . Let { f k } k ∈ [ K ] ⊂ ( R d → R ) denote the K learned predictors, and

denote the validation loss of any predictor f .

We show that (proof in Appendix I.1) a 3-layer transformer can output a predictor ̂ f that achieves nearly the smallest validation loss, and thus nearly optimal expected loss if ̂ L val concentrates around the expected loss L . Below, the input sequence H uses a generalized positional encoding p i := [ 0 D -( d +3) ; 1; t i ] in (3), where t i := 1 for i ∈ D train , t i := -1 for i ∈ D val , and t N +1 := 0 .

Proposition 10 (In-context algorithm selection via train-validation split) . Suppose that ℓ ( · , · ) in (7) is approximable by sum of relus (Definition D.1, which includes all C 3 -smooth bivariate functions). Then there exists a 3-layer transformer TF θ that maps (defining y ′ i = y i 1 { i < N +1 } )

where the predictor ̂ f : R d → R is a convex combination of { f k : ̂ L val ( f k ) ≤ min k ⋆ ∈ [ K ] ̂ L val ( f k ⋆ ) + γ } . As a corollary, for any convex risk L : ( R d → R ) → R , ̂ f satisfies

Ridge regression with in-context regularization selection As an example, we use Proposition 10 to construct a transformer to perform in-context ridge regression with regularization selection according to the unregularized validation loss ̂ L val ( w ) := 1 2 |D val | ∑ ( x i ,y i ) ∈D val ( ⟨ w , x i ⟩ -y i ) 2 (proof in Appendix I.2). Let λ 1 , . . . , λ K ≥ 0 be K fixed regularization strengths.

Theorem 11 (Ridge regression with in-context regularization selection) . There exists a transformer with O (log(1 /ε )) layers and O ( K ) heads such that the following holds: On any ( D , x N +1 ) wellconditioned (cf. (5)) for all { λ k } k ∈ [ K ] , it outputs ̂ y N +1 = ⟨ ̂ w , x N +1 ⟩ , where

Above, ̂ w λ ridge , train denotes the solution to (ICRidge) on the training split D train .

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

We build on Theorem 11 to show that transformers can perform nearly Bayes-optimal ICL when data come from noisy linear models with a mixture of K different noise levels σ 1 , . . . , σ K > 0 .

Concretely, consider the following data generating model, where we first sample P = P w ⋆ ,σ k ∼ π from k ∼ Λ ∈ ∆([ K ]) , w ⋆ ∼ N ( 0 , I d /d ) , and then sample data { ( x i , y i ) } i ∈ [ N +1] iid ∼ P w ⋆ ,σ k as

For any fixed ( N,d ) , consider the Bayes risk for predicting y N +1 under this model:

By standard Bayesian calculations, the above Bayes risk is attained when A is a certain mixture of K ridge regressions with regularization λ k = dσ 2 k /N ; however, the mixing weights depend on D in a highly non-trivial fashion (see Appendix J.2 for a derivation). By using the post-ICL validation mechanism in Theorem 11, we construct a transformer that achieves nearly the Bayes risk.

Theorem 12 (Nearly Bayes-optimal ICL; Informal version of Theorem J.1) . For sufficiently large N,d , there exists a transformer with O (log N ) layers and O ( K ) heads such that on the above model, it outputs a prediction ̂ y N +1 that is nearly Bayes-optimal:

In particular, Theorem 12 applies in the proportional setting where N,d are large and N/d = Θ(1) [22], in which case BayesRisk π = Θ(1) , and thus the transformer achieves vanishing excess risk relative to the Bayes risk at large N .

This substantially strengthens the results of Akyürek et al. [2], who empirically find that transformers can achieve nearly Bayes risk under any fixed noise level. By contrast, Theorem 12 shows that a single transformer can achieve nearly Bayes risk even under a mixture of K noise levels, with quantitative guarantees. Also, our proof in fact gives a stronger guarantee: The transformer approaches the individual Bayes risks on all K noise levels simultaneously (in addition to the overall Bayes risk for k ∼ Λ as in Theorem 12). We demonstrate this empirically in Section 6 (cf. Figure 3b & 2).

Exact Bayes predictor vs. Post-ICL validation mechanism As BayesRisk π is the theoretical lower bound for the risk of any possible ICL algorithm, Theorem 12 implies that our transformer performs similarly as the exact Bayes estimator 5 . Notice that our construction builds on the (generic) post-ICL validation mechanism, rather than a direct attempt of approximating the exact Bayes predictor, whose structure may vary significantly case-by-case. This highlights post-ICL validation as a promising mechanism for approximating the Bayes predictor on broader classes of problems beyond noisy linear models, which we leave as future work.

Generalized linear models with adaptive link function selection As another example of the post-ICL validation mechanism, we construct a transformer that can learn a generalized linear model with adaptively chosen link function for the particular ICL instance; see Theorem J.2.

4.2 Pre-ICL testing mechanism

In our second mechanism, pre-ICL testing, the transformer runs a distribution testing procedure on the input sequence to determine the right ICL algorithm to use. While the test (and thus the mechanism itself) could in principle be general, we focus on cases where the test amounts to computing some simple summary statistics of the input sequence.

To showcase pre-ICL testing, we consider the toy problem of selecting between in-context regression and in-context classification, by running the following binary type check on the input labels { y i } i ∈ [ N ] .

Lemma 13. There exists a single attention layer with 6 heads that implements Ψ binary exactly.

Using this test, we construct a transformer that performs logistic regression when labels are binary, and linear regression with high probability if the label admits a continuous distribution.

Proposition 14 (Adaptive regression or classification; Informal version of Proposition I.4) . There exists a transformer with O (log(1 /ε )) layers such that the following holds: On any D such that y i ∈ { 0 , 1 } , it outputs ̂ y N +1 that ε -approximates the prediction of in-context logistic regression.

By contrast, for any distribution P whose marginal distribution of y is not concentrated around { 0 , 1 } , with high probability (over D ), ̂ y N +1 ε -approximates the prediction of in-context least squares.

The proofs can be found in Appendix I.3. We additionally show that transformers can implement more complex tests such as a linear correlation test , which can be useful in certain scenarios such as 'confident linear regression' (predict only when the signal-to-noise ratio is high); see Appendix I.4.

5 Analysis of pretraining

Building on the expressivity results in Section 3 & 4, we provide the first line of polynomial sample complexity results for pretraining transformers to perform ICL (including with in-context algorithm selection). We begin by providing a generic generalization guarantee for pretraining transformers.

Consider the pretraining ERM problem (TF-ERM), which minimizes the pretraining risk ̂ L icl ( · ) over n pretraining sequences. Let L icl ( · ) denote the corresponding population risk.

Theorem 15 (Generalization of transformers; Informal version of Theorem K.1) . The solution ̂ θ to (TF-ERM) over transformers with L layers, M heads per layer, and hidden dimension D ′ satisfies

Theorem 15 builds on standard uniform concentration analysis via chaining (Proposition B.4). Combining Theorem 15 with the in-context linear regression construction in Theorem 4 gives the following end-to-end result on the excess in-context prediction risk of trained transformers.

Theorem 16 (Pretraining transformers for in-context linear regression; Informal version of Theorem K.2) . Under Assumption A and N ≥ ˜ O ( d ) , the solution ̂ θ to (TF-ERM) with L = O ( κ log( κN/σ )) layers, M = 3 heads, D ′ = 0 (attention-only as in Theorem 4) achieves small excess ICL risk over the best linear predictor w ⋆ P := E P [ xx ⊤ ] -1 E P [ x y ] for each P :

See Appendix K.2 for similar results in several additional settings.

6 Experiments

We test our theory by studying the ICL and in-context algorithm selection capabilities of transformers, using the encoder-based architecture in our theoretical constructions (Definition 3). Due to limited space, additional experimental details can be found in Appendix M.1. Results with a decoder architecture as in [31, 47] (including the setup of Figure 2) can be found in Appendix M.2.

Training data distributions and evaluation We train a 12-layer transformer, with two modes for the training sequence (instance) distribution π . In the 'base' mode, similar to [31, 2, 86, 47], we sample the training instances from one of the following base distributions (tasks), where we first sample P = P w ⋆ ∼ π by sampling w ⋆ ∼ N ( 0 , I d /d ) , and then sample { ( x i , y i ) } i ∈ [ N +1] iid ∼ P w ⋆ as x i iid ∼ N ( 0 , I d ) , and y i from one of the following models studied in Section 3:

Linear model: y i = ⟨ w ⋆ , x i ⟩ ;

Noisy linear model: y i = ⟨ w ⋆ , x i ⟩ + σz i , where σ > 0 is a fixed noise level, and z i ∼ N (0 , 1) .

(a) Base ICL capabilities

0.00

0.25

0.50

0.75

1.00

1.25

1.50

Loss

linear_regression

noisy_reg_noise_1

noisy_reg_noise_2

sparse_reg

linear_classification

Transformer

Averaging

3-NN

ridge_lam_1

ridge_lam_2

Lasso_lam=0.01

Lasso_lam=0.001

Logistic Regression

(b) Noisy reg with two noises

0.10

0.15

0.20

0.25

0.30

noisy_reg_noise_1

0.6

0.8

1.0

1.2

1.4

noisy_reg_noise_2

TF_alg_select

TF_noise_1

TF_noise_2

ridge_lam_1

ridge_lam_2

Bayes_err_noise_1

Bayes_err_noise_2

(c) Reg + Classification

0.0

0.5

1.0

1.5

regression_square_loss

0.20

0.25

0.30

0.35

classification_error

TF_alg_select

TF_reg

TF_cls

Averaging

3-NN

Sparse linear model: y i = ⟨ w ⋆ , x i ⟩ with ∥ w ⋆ ∥ 0 ≤ s , where s < d is a fixed sparsity level, and in this case we sample w ⋆ from a special prior supported on s -sparse vectors;

Linear classification model: y i = sign( ⟨ w ⋆ , x i ⟩ ) .

These base tasks have been empirically investigated by Garg et al. [31], though we remark that our architecture (used in our theory) differs from theirs in several aspects, such as encoder-based architecture instead of decoder-based, and ReLU activation instead of softmax. All experiments use d = 20 . We choose σ ∈ { σ 1 , σ 2 } = { 0 . 1 , 0 . 5 } and N = 20 for noisy linear regression, s = 3 and N = 10 for sparse linear regression, and N = 40 for linear regression and linear classification.

In the 'mixture' mode, π is the uniform mixture of two or more base distributions . We consider two representative mixture modes studied in Section 4:

Linear model + linear classification model;

Noisy linear model with four noise levels σ ∈ { 0 . 1 , 0 . 25 , 0 . 5 , 1 } .

Transformers trained with the mixture mode will be evaluated on multiple base distributions simultaneously. When the base distributions are sufficiently diverse, a transformer performing well on all of them will likely be performing some level of in-context algorithm selection. We evaluate transformers against standard machine learning algorithms in context (for each task respectively) as baselines.

Results Figure 3a shows the ICL performance of transformers on five base tasks, within each the transformer is trained on the same task. Transformers match the best baseline algorithm in four out of the five cases, except for the sparse regression task where the Transformer still outperforms least squares and matches Lasso with some choices of λ (thus utilizing sparsity to some extent). This demonstrates the strong ICL capability of the transformer architecture considered in our theory.

Figure 3b & 3c examine the in-context algorithm selection capability of transformers, on noisy linear regression with two different noise levels (Figure 3b), and regression + classification (Figure 3c). In both figures, the transformer trained in the mixture mode ( TF_alg_select ) approaches the best baseline algorithm on both tasks simultaneously. By contrast, transformers trained in the base mode for one of the tasks perform well on that task but behave suboptimally on the other task as expected. The existence of TF_alg_select showcases a single transformer that performs well on multiple tasks simultaneously (and thus has to perform in-context algorithm selection to some extent), supporting our theoretical results in Section 4.

7 Conclusion

This work shows that transformers can perform complex in-context learning procedures with strong in-context algorithm selection capabilties, by both explicit theoretical constructions and experiments. We believe our work opens up many exciting directions, such as (1) more mechanisms for in-context algorithm selection; (2) Bayes-optimal ICL on other problems by either the post-ICL validation mechanism or new approaches; (3) understanding the internal workings of transformers performing in-context algorithm selection; (4) other mechanisms for implementing complex ICL procedures beyond in-context algorithm selection; (5) further statistical analyses, e.g. of pretraining. Besides, this work focuses on the transformer architecture; alternative sequence-to-sequence architectures (such as RNNs) are beyond our scope but would be interesting directions for future work.

Acknowledgment

The authors would like to thank Tengyu Ma and Jason D. Lee for the many insightful discussions. S. Mei is supported in part by NSF DMS-2210827 and NSF CCF-2315725.

References

A. Agarwal, S. Negahban, and M. J. Wainwright. Fast global convergence rates of gradient methods for high-dimensional statistical recovery. Advances in Neural Information Processing Systems , 23, 2010.

E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

F. Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research , 18(1):629-681, 2017.

Y. Bai, M. Chen, P. Zhou, T. Zhao, J. Lee, S. Kakade, H. Wang, and C. Xiong. How important is the train-validation split in meta-learning? In International Conference on Machine Learning , pages 543-553. PMLR, 2021.

Y. Bai, S. Mei, H. Wang, and C. Xiong. Don't just blame over-parametrization for overconfidence: Theoretical analysis of calibration in binary classification. In International Conference on Machine Learning , pages 566-576. PMLR, 2021.

J. Baxter. A model of inductive bias learning. Journal of artificial intelligence research , 12: 149-198, 2000.

A. Beck and M. Teboulle. Gradient-based algorithms with applications to signal recovery. Convex optimization in signal processing and communications , pages 42-88, 2009.

S. Bengio, Y. Bengio, J. Cloutier, and J. Gescei. On the optimization of a synaptic learning rule. In Optimality in Biological and Artificial Networks? , pages 281-303. Routledge, 2013.

S. Bhattamishra, K. Ahuja, and N. Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009.11264 , 2020.

S. Bhattamishra, A. Patel, and N. Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09286 , 2020.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

S. Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning , 8(3-4):231-357, 2015.

S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712 , 2023.

S. Chan, A. Santoro, A. Lampinen, J. Wang, A. Singh, P. Richemond, J. McClelland, and F. Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems , 35:18878-18891, 2022.

L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

K. Chua, Q. Lei, and J. D. Lee. How fine-tuning allows for effective meta-learning. Advances in Neural Information Processing Systems , 34:8871-8884, 2021.

D. Dai, Y. Sun, L. Dong, Y. Hao, Z. Sui, and F. Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.10559 , 2022.

G. Denevi, C. Ciliberto, D. Stamos, and M. Pontil. Incremental learning-to-learn with statistical guarantees. arXiv preprint arXiv:1803.08089 , 2018.

G. Denevi, C. Ciliberto, D. Stamos, and M. Pontil. Learning to learn around a common mean. Advances in Neural Information Processing Systems , 31, 2018.

J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

E. Dobriban and S. Wager. High-dimensional asymptotics of prediction: Ridge regression and classification. The Annals of Statistics , 46(1):247-279, 2018.

L. Dong, S. Xu, and B. Xu. Speech-transformer: a no-recurrence sequence-to-sequence model for speech recognition. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP) , pages 5884-5888. IEEE, 2018.

Q. Dong, L. Li, D. Dai, C. Zheng, Z. Wu, B. Chang, X. Sun, J. Xu, and Z. Sui. A survey for in-context learning. arXiv preprint arXiv:2301.00234 , 2022.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

S. S. Du, W. Hu, S. M. Kakade, J. D. Lee, and Q. Lei. Few-shot learning via learning the representation, provably. arXiv preprint arXiv:2002.09434 , 2020.

B. L. Edelman, S. Goel, S. Kakade, and C. Zhang. Inductive biases and variable creation in selfattention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread , 2021.

C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning , pages 1126-1135. PMLR, 2017.

C. Finn, A. Rajeswaran, S. Kakade, and S. Levine. Online meta-learning. In International Conference on Machine Learning , pages 1920-1930. PMLR, 2019.

S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

A. Giannou, S. Rajput, J.-y. Sohn, K. Lee, J. D. Lee, and D. Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196 , 2023.

M. Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics , 8:156-171, 2020.

S. Hochreiter, A. S. Younger, and P. R. Conwell. Learning to learn using gradient descent. In Artificial Neural Networks-ICANN 2001: International Conference Vienna, Austria, August 21-25, 2001 Proceedings 11 , pages 87-94. Springer, 2001.

N. Hollmann, S. Müller, K. Eggensperger, and F. Hutter. Tabpfn: A transformer that solves small tabular classification problems in a second. arXiv preprint arXiv:2207.01848 , 2022.

T. Hospedales, A. Antoniou, P. Micaelli, and A. Storkey. Meta-learning in neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence , 44(9):5149-5169, 2021.

J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

D. Hsu, S. M. Kakade, and T. Zhang. Random design analysis of ridge regression. In Conference on learning theory , pages 9-1. JMLR Workshop and Conference Proceedings, 2012.

S. Jelassi, M. E. Sander, and Y. Li. Vision transformers provably learn spatial structure. arXiv preprint arXiv:2210.09221 , 2022.

K. Ji, J. D. Lee, Y . Liang, and H. V . Poor. Convergence of meta-learning with task-specific adaptation over partial parameters. Advances in Neural Information Processing Systems , 33: 11490-11500, 2020.

J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. v Zídek, A. Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

S. M. Kakade, V. Kanade, O. Shamir, and A. Kalai. Efficient learning of generalized linear and single index models with isotonic regression. Advances in Neural Information Processing Systems , 24, 2011.

M. Khodak, M.-F. F. Balcan, and A. S. Talwalkar. Adaptive gradient-based meta-learning methods. Advances in Neural Information Processing Systems , 32, 2019.

L. Kirsch and J. Schmidhuber. Meta learning backpropagation and improving it. Advances in Neural Information Processing Systems , 34:14122-14134, 2021.

L. Kirsch, J. Harrison, J. Sohl-Dickstein, and L. Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprint arXiv:2212.04458 , 2022.

K. Li and J. Malik. Learning to optimize. arXiv preprint arXiv:1606.01885 , 2016.

Y. Li, M. E. Ildiz, D. Papailiopoulos, and S. Oymak. Transformers as algorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067 , 2023.

B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749 , 2022.

J. Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen. What makes good in-context examples for gpt3 ? arXiv preprint arXiv:2101.06804 , 2021.

Y. Lu, M. Bartolo, A. Moore, S. Riedel, and P. Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2104.08786 , 2021.

A. Madani, B. McCann, N. Naik, N. S. Keskar, N. Anand, R. R. Eguchi, P.-S. Huang, and R. Socher. Progen: Language modeling for protein generation. arXiv preprint arXiv:2004.03497 , 2020.

A. Maurer, M. Pontil, and B. Romera-Paredes. The benefit of multitask representation learning. Journal of Machine Learning Research , 17(81):1-32, 2016.

P. McCullagh. Generalized linear models . Routledge, 2019.

S. Mei, Y. Bai, and A. Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics , 46(6A):2747-2774, 2018.

S. Min, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Noisy channel language model prompting for few-shot text classification. arXiv preprint arXiv:2108.04106 , 2021.

S. Min, M. Lewis, L. Zettlemoyer, and H. Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv:2110.15943 , 2021.

S. Min, X. Lyu, A. Holtzman, M. Artetxe, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837 , 2022.

N. Mishra, M. Rohaninejad, X. Chen, and P. Abbeel. A simple neural attentive meta-learner. arXiv preprint arXiv:1707.03141 , 2017.

S. Müller, N. Hollmann, S. P. Arango, J. Grabocka, and F. Hutter. Transformers can do bayesian inference. arXiv preprint arXiv:2112.10510 , 2021.

T. Nagler. Statistical foundations of prior-data fitted networks. arXiv preprint arXiv:2305.11097 , 2023.

D. K. Naik and R. J. Mammone. Meta-neural networks that learn by learning. In [Proceedings 1992] IJCNN International Joint Conference on Neural Networks , volume 1, pages 437-442. IEEE, 1992.

S. N. Negahban, P. Ravikumar, M. J. Wainwright, and B. Yu. A unified framework for high-dimensional analysis of m-estimators with decomposable regularizers. 2012.

Y. Nesterov. Lectures on convex optimization , volume 137. Springer, 2018.

C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895 , 2022.

OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

N. Parikh, S. Boyd, et al. Proximal algorithms. Foundations and trends® in Optimization , 1 (3):127-239, 2014.

J. Pérez, J. Marinkovi´ c, and P. Barceló. On the turing completeness of modern neural network architectures. arXiv preprint arXiv:1901.03429 , 2019.

A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.

A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

A. Raventos, M. Paul, F. Chen, and S. Ganguli. The effects of pretraining task diversity on in-context learning of ridge regression. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models , 2023.

S. Ravi and H. Larochelle. Optimization as a model for few-shot learning. In International conference on learning representations , 2017.

Y. Razeghi, R. L. Logan IV, M. Gardner, and S. Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXiv:2202.07206 , 2022.

S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175 , 2022.

O. Rubin, J. Herzig, and J. Berant. Learning to retrieve prompts for in-context learning. arXiv preprint arXiv:2112.08633 , 2021.

A. Santoro, S. Bartunov, M. Botvinick, D. Wierstra, and T. Lillicrap. Meta-learning with memory-augmented neural networks. In International conference on machine learning , pages 1842-1850. PMLR, 2016.

N. Saunshi, A. Gupta, and W. Hu. A representation learning perspective on the importance of train-validation splitting in meta-learning. In International Conference on Machine Learning , pages 9333-9343. PMLR, 2021.

J. Schmidhuber. Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-... hook . PhD thesis, Technische Universität München, 1987.

K. Shen, J. Guo, X. Tan, S. Tang, R. Wang, and J. Bian. A study on relu and softmax in transformer. arXiv preprint arXiv:2302.06461 , 2023.

C. Snell, R. Zhong, D. Klein, and J. Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

J. Snell, K. Swersky, and R. Zemel. Prototypical networks for few-shot learning. Advances in neural information processing systems , 30, 2017.

S. Thrun and L. Pratt. Learning to learn . Springer Science & Business Media, 2012.

R. Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society: Series B (Methodological) , 58(1):267-288, 1996.

N. Tripuraneni, M. Jordan, and C. Jin. On the theory of transfer learning: The importance of task diversity. Advances in neural information processing systems , 33:7852-7862, 2020.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

R. Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

X. Wang, S. Yuan, C. Wu, and R. Ge. Guarantees for tuning the step size using a learningto-learn approach. In International Conference on Machine Learning , pages 10981-10990. PMLR, 2021.

C. Wei, Y. Chen, and T. Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. arXiv preprint arXiv:2107.13163 , 2021.

J. Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682 , 2022.

J. Wei, J. Wei, Y. Tay, D. Tran, A. Webson, Y. Lu, X. Chen, H. Liu, D. Huang, D. Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03846 , 2023.

G. Weiss, Y. Goldberg, and E. Yahav. Thinking like transformers. In International Conference on Machine Learning , pages 11080-11090. PMLR, 2021.

M. Wortsman, J. Lee, J. Gilmer, and S. Kornblith. Replacing softmax with relu in vision transformers. arXiv preprint arXiv:2309.08586 , 2023.

S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 , 2021.

S. Yao, B. Peng, C. Papadimitriou, and K. Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115 , 2021.

C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems , 34:28877-28888, 2021.

C. Yun, S. Bhojanapalli, A. S. Rawat, S. J. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and T. Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301 , 2022.

Y. Zhang, B. Liu, Q. Cai, L. Wang, and Z. Wang. An analysis of attention via the lens of exchangeability and latent variable models. arXiv preprint arXiv:2212.14852 , 2022.

Z. Zhao, E. Wallace, S. Feng, D. Klein, and S. Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning , pages 12697-12706. PMLR, 2021.

X. Zuo, Z. Chen, H. Yao, Y. Cao, and Q. Gu. Understanding train-validation split in metalearning with neural networks. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=JVlyfHEEm0k .

A Related work

In-context learning The in-context learning (ICL) capability of large language models (LLMs) has gained significant attention since demonstrated on GPT-3 Brown et al. [12]. A number of subsequent empirical studies have contributed to a better understanding of the capabilities and limitations of ICL in LLM systems, which include but are not limited to [49, 55, 56, 50, 100, 74, 72, 28, 45, 91]. For an overview of ICL, see the survey by Dong et al. [24] which highlights some key findings and advancements in this direction.

A line of recent work investigates why and how LLMs perform ICL [94, 31, 86, 2, 18, 32, 47, 70]. In particular, Xie et al. [94] propose a Bayesian inference framework explaining how ICL works despite formatting differences between training and inference distributions. Garg et al. [31] show empirically that transformers could be trained from scratch to perform ICL of linear models, sparse linear models, two-layer neural networks, and decision trees. Li et al. [47] analyze the generalization error of trained ICL transformers from a stability viewpoint. They also experimentally show that transformers could perform 'in-context model selection' (conceptually similar to in-context algorithm selection considered in this work) in specific tasks and presented related theoretical hypotheses. However, they do not provide concrete mechanisms or constructions for in-context model selection. A recent work [99] shows that pretrained transformers can perform Bayesian inference in latent variable models, which may also be interpreted as a mechanism for ICL. Our experimental findings extend these results by unveiling and demonstrating the in-context algorithm selection capabilities of transformers.

Closely related to our theoretical results are [86, 2, 18, 32], which show (among many things) that transformers can perform ICL by simulating gradient descent. However, these results do not provide quantitative error bounds for simulating multi-step gradient descent, and only handle linear regression models or their simple variants. Among these works, Akyürek et al. [2] showed that transformers can implement learning algorithms for linear models based on gradient descent and closed-form ridge regression; it also presented preliminary evidence that learned transformers perform ICL similar to Bayes-optimal ridge regression. Our work builds upon and substantially extends this line of work by (1) providing a more efficient construction for in-context gradient descent; (2) providing an end-to-end theory with additional results for pretraining and statistical power; (3) analyzing a broader spectrum of ICL algorithms, including least squares, ridge regression, Lasso, convex risk minimization for generalized linear models, and gradient descent on two-layer neural networks; and (4) constructing more complex ICL procedures using in-context algorithm selection.

When in-context data are generated from a prior, the Bayes risk is a theoretical lower bound for the risk of any possible ICL algorithm, including transformers. Xie et al. [94], Akyürek et al. [2] observe

that learned transformers behave closely to the Bayes predictor on a variety of tasks such as hidden Markov models [94] and noisy linear regression with a fixed noise level [2, 47]. Using the in-context algorithm selection mechanism (more precisely the post-ICL validation mechanism), we show that transformers can perform nearly-Bayes optimal ICL in noisy linear models with mixed noise levels (a strictly more challenging task than considered in [2, 47]), with both concrete theoretical guarantees (Section 4.1.1) and empirical evidence (Figure 2 & 3b). Complementary to these works, a line of work on 'prior-data fitted networks' [59, 60, 35] also empirically demonstrates the Bayesian optimality of transformers in various settings. Our expressivity results support these empirical findings and are applicable beyond the Bayesian setting, e.g. for providing frequentist in-context prediction guarantees for transformers.

Transformers and its theory The transformer architecture, introduced by [84], has revolutionized natural language processing and been adopted in most of the recently developed large language models such as BERT and GPT [68, 21, 12]. Broaderly, transformers have demonstrated remarkable performance in many other fields of artificial intelligence such as computer vision, speech, graph processing, reinforcement learning, and biological applications [23, 25, 51, 69, 96, 16, 41, 73, 65, 14]. Towards a better theoretical understanding, recent work has studied the capabilities [97, 67, 37, 95, 11, 98, 48], limitations [33, 10], and internal workings [28, 79, 92, 27, 64] of transformers.

We remark that the transformer architecture used in our theoretical constructions differs from the standard one by replacing the softmax activation (in the attention layers) with a (normalized) ReLU function. Transformers with ReLU activations is experimentally studied in the recent work of Shen et al. [78], who find that they perform as well as the standard softmax activation in many NLP tasks.

Meta-learning Training models (such as transformers) to perform ICL can be viewed as an approach for the broader problem of learning-to-learn or meta-learning [77, 61, 81]. A number of other approaches has been studied extensively for this problem, including (and not limited to) training a meta-learner on how to update the parameters of a downstream learner [9, 46], learning parameter initializations that quickly adapt to downstream tasks [29, 71], learning latent embeddings that allow for effective similarity search [80]. Most relevant to the ICL setting are approaches that directly take as input examples from a downstream task and a query input and produce the corresponding output [34, 58, 75, 44]. For a comprehensive overview, see the survey [36].

Theoretical aspects of meta-learning have received significant recent interest [7, 52, 26, 83, 19, 30, 43, 40, 88, 20, 5, 76, 17, 101]. In particular, [52, 26, 83] analyzed the benefit of multi-task learning through a representation learning perspective, and [88, 20, 5, 76, 101] studied the statistical properties of learning the parameter initialization for downstream tasks.

Techniques Webuild on various existing techniques from the statistics and learning theory literature to establish our approximation and generalization guarantees for transformers. For the approximation component, we rely on a technical result of Bach [4] on the approximation power of ReLU networks. We use this result to show that transformers can approximate gradient descent (GD) on a broad range of loss functions, substantially extending the results of [86, 2, 18] who primarily consider the square loss. The recent work of Giannou et al. [32] also approximates GD with general loss functions by transformers, though using a different technique of forcing the softmax activations to act as sigmoids. Our analyses of Lasso and generalized linear models build on [87, 62, 1, 54]. Our generalization bound for transformers (used in our pretraining results) build on a chaining argument [87].

B Technical tools

Additional notation for proofs We say a random variable X is σ 2 -sub-Gaussian (or SG( σ ) interchangeably) if E [exp( X 2 /σ 2 )] ≤ 2 . A random vector x ∈ R d is σ 2 -sub-Gaussian if ⟨ v , x ⟩ is σ 2 -sub-Gaussian for all ∥ v ∥ 2 = 1 . A random variable X is K -sub-Exponential (or SE( K ) interchangeably) if E [exp( | X | /K )] ≤ 2 .

B.1 Concentration inequalities

Lemma B.1. Let β ∼ N ( 0 , I d /d ) . Then we have

Lemma B.2 (Theorem 6.1 of [87]) . Let X = [ X ij ] ∈ R n × d be a Gaussian random matrix with X ij ∼ N (0 , 1) . Let σ min ( X ) and σ min ( X ) be the minimum and maximum singular value of X , respectively. Then we have

The following lemma is a standard result of covariance concentration, see e.g. [85, Theorem 4.6.1].

Lemma B.3. Suppose that x 1 , · · · , x N are independent d -dimensional K -sub-Gaussian random vectors. Then as long as N ≥ C 0 d , with probability at least 1 -exp( -N/C 0 ) we have

where C 0 is a universal constant.

Lemma B.4. For random matrix X = [ x ij ] ∈ R N × d with x ij iid ∼ N (0 , 1) and ε = [ ε i ] ∈ R N with ε i iid ∼ N (0 , σ 2 ) , it holds that

Proof. We consider u j := [ x ij ] i ∈ R N , then ∥ ∥ X ⊤ ε ∥ ∥ ∞ = max i ∈ [ d ] |⟨ u j , ε ⟩| . Notice that the random variables ⟨ u 1 , ε ⟩ , · · · , ⟨ u d , ε ⟩ are independent N (0 , ∥ ε ∥ 2 2 ) , and hence

Further, by Lemma B.1, P ( ∥ ε ∥ 2 ≥ 2 σ √ N ) ≤ exp( -N/ 2) . Taking t = √ 8 Nσ 2 log(2 d/δ ) completes the proof.

B.2 Approximation theory

For any signed measure µ over a space W , let TV( µ ) := ∫ W | dµ ( w ) | ∈ [0 , ∞ ] denote its total measure. Recall σ ( · ) = ReLU( · ) is the standard relu activation, and B k ∞ ( R ) = [ -R,R ] k denotes the standard ℓ ∞ ball in R k with radius R > 0 .

Definition B.1 (Sufficiently smooth k -variable function) . We say a function g : R k → R is ( R,C ℓ ) -smooth, if for s = ⌈ ( k -1) / 2 ⌉ +2 , g is a C s function on B k ∞ ( R ) , and

for all i ∈ { 0 , 1 , . . . , s } , with max 0 ≤ i ≤ s L i R i ≤ C ℓ .

The following result for expressing smooth functions as a random feature model with relu activation is adapted from Bach [4, Proposition 5].

Lemma B.5 (Expressing sufficiently smooth functions by relu random features) . Suppose function g : R k → R is ( R,C ℓ ) smooth. Then there exists a signed measure µ over W = { w ∈ R k +1 : ∥ w ∥ 1 = 1 } such that

and TV( µ ) ≤ C ( k ) C ℓ , where C ( k ) < ∞ is a constant that only depends on k .

Lemma B.6 (Uniform finite-neuron approximation) . Let X be a space equipped with a distance function d X ( · , · ) : X × X → R ≥ 0 . Suppose function g : X → R is given by

where ϕ ( · ; · ) : X × W → [ -B,B ] is L -Lipschitz (in d X ) in the first argument, and µ is a signed measure over W with finite total measure A = TV( µ ) < ∞ . Then for any ε > 0 , there exists α 1 , · · · , α K ∈ {± 1 } , w 1 , · · · , w K ∈ W with K = O ( A 2 B 2 log N ( X , d X , ε 3 AL ) /ε 2 ) , such that

where N ( X , d X , ε 3 AL ) denotes the ( ε 3 AL ) -covering number of X in d X .

Proof. Let α ( w ) := sign( dµ ( w )) ∈ {± 1 } denote the sign of the density dµ ( w ) . We have

Note that | dµ ( w ) | /A is the density of a probability distribution over W . Thus for any x ∈ X , as long as K ≥ O ( A 2 B 2 log(1 /δ ) /ε 2 ) , we can sample w 1 , . . . , w K iid ∼ | dµ ( · ) | /A , and obtain by Hoeffding's inequality that with probability at least 1 -δ ,

Let N ( ε 3 AL ) := N ( X , d X , ε 3 AL ) for shorthand. By union bound, as long as K ≥ O ( A 2 B 2 log( N ( ε 3 AL ) /δ ) /ε 2 ) , we have with probability at least 1 -δ that for every ̂ x in the covering set corresponding to N ( ε 3 AL ) ,

Taking δ = 1 / 2 (for which K = O ( A 2 B 2 log N ( ε 3 AL ) /ε 2 ) ), by the probabilistic method, there exists a deterministic set { w i } i ∈ [ K ] ⊂ W and { α i := α ( w i ) } i ∈ [ K ] ∈ {± 1 } such that the above holds.

Next, note that both g (by (9)) and the function x ↦→ A K ∑ K i =1 α ( w i ) ϕ ( x ; w i ) are ( AL ) -Lipschitz. Therefore, for any x ∈ X , taking ̂ x to be the point in the covereing set with d X ( x , ̂ x ) ≤ ε 3 AL , we have

This proves the lemma.

Proposition B.1 (Approximating smooth k -variable functions) . For any ε approx > 0 , R ≥ 1 , C ℓ > 0 , we have the following: Any ( R,C ℓ ) -smooth function (Definition B.1) g : R k → R is ( ε approx , R, M, C ) -approximable by sum of relus (Definition D.1) with M ≤ C ( k ) C 2 ℓ log(1 + C ℓ /ε approx ) /ε 2 approx ) and C ≤ C ( k ) C ℓ , where C ( k ) > 0 is a constant that depends only on k . In other words, there exists

such that sup z ∈ [ -R,R ] k | f ( z ) -g ( z ) | ≤ ε approx .

Proof. As function g : B k ∞ ( R ) → R is ( R,C ℓ ) -smooth, we can apply Lemma B.5 to obtain that there exists a signed measure µ over W := { w ∈ R k +1 : ∥ w ∥ 1 ≤ 1 } such that

and A = TV( µ ) ≤ C ( k ) C ℓ where C ( k ) > 0 denotes a constant depending only on k .

We now apply Lemma B.6 to approximate the above random feature by finitely many neurons. Let x := [ z ; R ] ∈ X := [ -R,R ] k ×{ R } . Then, the function ϕ ( x ; w ) := 1 R σ ( w ⊤ x ) = σ ( 1 R w ⊤ [ z ; R ]) is bounded by B = 1 and (1 /R ) -Lipschitz in x (in the standard ℓ ∞ -distance). Further, we have log N ( X , ∥· - ·∥ ∞ , ε approx 3 A/R ) ≤ O ( k log(1 + A/ε approx )) . We can thus apply Lemma B.6 to obtain that, for

there exists α = { α m } m ∈ [ M ] ⊂ {± 1 } and W = { w m } m ∈ [ M ] ⊂ W = { w ∈ R k +1 : lone w = 1 } such that

where (recalling z = [ s ; t ] )

Note that we have ∑ M m =1 | c m | = A ≤ C ( k ) C ℓ , and ∥ a m ∥ 1 ≤ ∥ w m ∥ 1 = 1 . This is the desired result.

B.3 Optimization

The following convergence result for minimizing a smooth and strongly convex function is standard from the convex optimization literature, see e.g. Bubeck [13, Theorem 3.10].

Proposition B.2 (Gradient descent for smooth and strongly convex functions) . Suppose L : R d → R is α -strongly convex and β -smooth for some 0 < α ≤ β . Then, the gradient descent iterates w t +1 GD := w t GD -η ∇ L ( w t GD ) with learning rate η = 1 /β and initialization w 0 GD ∈ R d satisfies for any t ≥ 1 ,

where κ := β/α is the condition number of L , and w ⋆ := arg min w ∈ R d L ( w ) is the minimizer of L .

The following convergence result of proximal gradient descent (PGD) on convex composite minimization problem is also standard, see e.g. [8].

Proposition B.3 (Proximal gradient descent for convex function) . Suppose L = f + h , f : R d → R is convex and β -smooth for some β > 0 , h : R d → R is a simple convex function. Then, the proximal gradient descent iterates w t +1 PGD := prox ηh ( w t PGD -η ∇ f ( w t PGD )) with learning rate η = 1 /β and initialization w 0 GD ∈ R d satisfies the following for any t ≥ 1 :

{ L ( w t PGD ) } is a decreasing sequence.

For any minimizer w ⋆ ∈ arg min w ∈ R d L ( w ) ,

and hence { ∥ w t PGD -w ⋆ ∥ 2 2 } is also a decreasing sequence.

For k ≥ 1 , t ≥ 0 , it holds that

B.4 Uniform convergence

The following result is shown in [87, Section 5.6].

Theorem B.1. Suppose that ψ : [0 , + ∞ ) → [0 , + ∞ ) is a convex, non-decreasing function that satisfies ψ ( x + y ) ≥ ψ ( x ) ψ ( y ) . For any random variable X , we consider the Orlicz norm induced by ψ : ∥ X ∥ ψ := inf { K > 0 : E ψ ( | X | /K ) } ≤ 1 .

Suppose that { X θ } θ is a zero-mean random process indexed by θ ∈ Θ such that ∥ X θ -X θ ′ ∥ ψ ≤ ρ ( θ, θ ′ ) for some metric ρ on the space Θ . Then it holds that

where D is the diameter of the metric space (Θ , ρ ) , and the generalized Dudley entropy integral J is given by

where N ( δ ; Θ , ρ ) is the δ -covering number of (Θ , ρ ) .

As a corollary of Theorem B.1, we have the following result.

Proposition B.4 (Uniform concentration bound by chaining) . Suppose that { X θ } θ ∈ Θ is a zero-mean random process given by

where z 1 , · · · , z N are i.i.d samples from a distribution P z such that the following assumption holds:

(a) The index set Θ is equipped with a distance ρ and diameter D . Further, assume that for some constant A , for any ball Θ ′ of radius r in Θ , the covering number admits upper bound log N ( δ ; Θ ′ , ρ ) ≤ d log(2 Ar/δ ) for all 0 < δ ≤ 2 r .

(b) For any fixed θ ∈ Θ and z sampled from P z , the random variable f ( z ; θ ) is a SG( B 0 ) -subGaussian random variable.

(c) For any θ, θ ′ ∈ Θ and z sampled from P z , the random variable f ( z ; θ ) -f ( z ; θ ′ ) is a SG( B 1 ρ ( θ, θ ′ )) -sub-Gaussian random variable.

Then with probability at least 1 -δ , it holds that

where C is a universal constant, and we denote κ = 1 + B 1 D/B 0 .

Furthermore, if we replace the SG in assumption (b) and (c) by SE , then with probability at least 1 -δ , it holds that

Proof. Fix a D 0 ∈ (0 , D ] to be specified later. We pick a ( D 0 / 2) -covering Θ 0 of Θ so that log | Θ 0 | ≤ d log(2 AD/D 0 ) . Then, by the standard uniform covering of independent sub-Gaussian random variables, we have with probability at least 1 -δ/ 2 ,

Assume that Θ 0 = { θ 1 , · · · , θ n } . For each j ∈ [ n ] , we consider Θ j is the ball centered at θ j of radius D 0 in (Θ , ρ ) . Then θ ∈ Θ j has diameter D 0 and admits covering number bound log N (Θ j , δ ) ≤ d log( AD 0 /δ ) . Hence, we can apply Theorem B.1 with the process { X θ } θ ∈ Θ j , then

and a simple calculation yields

Therefore, we can let t ≤ √ log(2 n/δ ) /N in the above inequality and taking the union bound over j ∈ [ n ] , and hence with probability at least 1 -δ/ 2 , it holds that for all j ∈ [ n ] ,

Notice that for each θ ∈ Θ , there exists j ∈ [ n ] such that θ ∈ Θ j , and hence

Thus, with probability at least 1 -δ , it holds

Taking D 0 = D/κ completes the proof of SG case.

We next consider the SE case. The idea is the same as the SG case, but in this case we need to consider the following Orlicz-norm:

Then Bernstein's inequality of SE random variables yields

for some universal constant C 0 . Therefore, we can repeat the argument above to deduce that with probability at least 1 -δ , it holds sup θ ∈ Θ | X θ | ≤ C ′′ ( B 0 + B 1 D 0 ) [√ d log(2 AD/D 0 ) + log(2 /δ ) N + d log(2 AD/D 0 ) + log(2 /δ ) N ] . Taking D 0 = D/κ completes the proof.

B.5 Useful properties of transformers

The following result can be obtained immediately by 'joining' the attention heads and MLP layers of two single-layer transformers.

Proposition B.5 (Joining parallel single-layer transformers) . Suppose that P 1 : R ( D 0 + D 1 ) × N → R D 1 × N , P 2 : R ( D 0 + D 2 ) × N → R D 2 × N are two sequence-to-sequence functions that are implemented by single-layer transformers, i.e. there exists θ 1 , θ 2 such that

Then, there exists θ such that for H ′ that takes form h ′ i = [ h (0) i ; h (1) i ; h (2) i ] , with h (0) i ∈ R D 0 , h (1) i ∈ R D 1 , h (2) i ∈ R D 2 , we have

Further, θ has at most M ≤ M 1 + M 2 heads, D ′ ≤ D ′ 1 + D ′ 2 hidden dimension in its MLP layer, and norm bound | | | θ | | | ≤ | | | θ 1 | | | + | | | θ 2 | | | .

Proposition B.6 (Joining parallel multi-layer transformers) . Suppose that P 1 : R ( D 0 + D 1 ) × N → R D 1 × N , P 2 : R ( D 0 + D 2 ) × N → R D 2 × N are two sequence-to-sequence functions that are implemented by multi-layer transformers, i.e. there exists θ 1 , θ 2 such that

Then, there exists θ such that for H ′ that takes form h ′ i = [ h (0) i ; h (1) i ; h (2) i ] , with h (0) i ∈ R D 0 , h (1) i ∈ R D 1 , h (2) i ∈ R D 2 , we have

Further, θ has at most L ≤ max { L 1 , L 2 } layers, max ℓ ∈ [ L ] M ( ℓ ) ≤ max ℓ ∈ [ L ] ( M ( ℓ ) 1 + M ( ℓ ) 2 ) heads, max ℓ ∈ [ L ] D ( ℓ ) ≤ max ℓ ∈ [ L ] ( D ( ℓ ) 1 + D ( ℓ ) 2 ) hidden dimension in its MLP layer (understanding the size of the empty layers as 0), and norm bound | | | θ | | | ≤ | | | θ 1 | | | + | | | θ 2 | | | .

Proof. When L 1 = L 2 ( θ 1 and θ 2 have the same number of layers), the result follows directly by applying Proposition B.5 repeatedly for all L 1 layers and the definition of the norm (2).

If (without loss of generality) L 1 < L 2 , we can augment θ 1 to L 2 layers by adding ( L 2 -L 1 ) layers with zero attention heads, and zero MLP hidden dimension (note that this does not change M 1 , D ′ 1 , and | | | θ 1 | | | ). Due to the residual structure, the transformer maintains the output P 1 ( H 1 ) throughout layer L 1 +1 , . . . , L 2 , and it reduces to the case L 1 = L 2 .

C Extension to decoder-based architecture

Here we briefly discuss how our theoretical results can be adapted to decoder-based architectures (henceforth decoder TFs). Adopting the setting as in Section 2, we consider a sequence of N input vectors { h i } N i =1 ⊂ R D , written compactly as an input matrix H = [ h 1 , . . . , h N ] ∈ R D × N . Recall that σ ( t ) := ReLU( t ) = max { t, 0 } denotes the standard relu activation.

C.1 Decoder-based transformers

Decoder TFs are the same as encoder TFs, except that the attention layers are replaced by masked attention layers with a specific decoder-based (causal) attention mask.

Definition C.1 (Masked attention layer) . A masked attention layer with M heads is denoted as MAttn θ ( · ) with parameters θ = { ( V m , Q m , K m ) } m ∈ [ M ] ⊂ R D × D . On any input sequence H ∈ R D × N ′ with N ′ ≤ N ,

where ◦ denotes the entry-wise (Hadamard) product of two matrices, and MSK ∈ R N × N is the mask matrix given by

In vector form, we have

Notice that standard masked attention definitions use the pre-activation additive masks (with mask value -∞ ) [84]. The post-activation multiplicative masks we use is equivalent to the pre-activation additive masks, and the modified presentation is for notational convenience. We also use a normalized ReLU activation t ↦→ σ ( t ) /i in place of the standard softmax activation to be consistent with Definition 1. Note that the normalization 1 /i is to ensure that the attention weights { σ ( ⟨ Q m h i , K m h j ⟩ ) /i } j ∈ [ i ] is a set of non-negative weights that sum to O (1) . The motivation of masked attention layer is to ensure that, when processing a sequence of tokens, the computations at any token do not see any later token.

We next define the decoder-based transformers with L ≥ 1 transformer layers, each consisting of a masked attention layer (c.f. Definition C.1) followed by an MLP layer (c.f. Definition 2). This definition is similar to the definition of encoder-based transformers (c.f., Definition 3), except that we replace the attention layers by masked attention layers.

Definition C.2 (Decoder-based Transformer) . An L -layer decoder-based transformer, denoted as DTF θ ( · ) , is a composition of L self-attention layers each followed by an MLP layer: H ( L ) = DTF θ ( H (0) ) , where H (0) ∈ R D × N is the input sequence, and

Above, the parameter θ = ( θ (1: L ) mattn , θ (1: L ) mlp ) is the parameter consisting of the attention layers θ ( ℓ ) mattn = { ( V ( ℓ ) m , Q ( ℓ ) m , K ( ℓ ) m ) } m ∈ [ M ( ℓ ) ] ⊂ R D × D and the MLP layers θ ( ℓ ) mlp = ( W ( ℓ ) 1 , W ( ℓ ) 2 ) ∈ R D ( ℓ ) × D × R D × D ( ℓ ) . We will frequently consider ' attention-only ' decoder-based transformers with W ( ℓ ) 1 , W ( ℓ ) 2 = 0 , which we denote as DTF 0 θ ( · ) for shorthand, with θ = θ (1: L ) := θ (1: L ) mattn .

We also use (2) to define the norm of DTF θ .

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

We consider using decoder-based TFs to perform ICL. We encode ( D , x N +1 ) , which follows the generating rule as described in Section 2.2, into an input sequence H ∈ R D × (2 N +1) . In our theory, we use the following format, where the first two rows contain ( D , x N +1 ) which alternating between [ x i ; 0] ∈ R d +1 and [ 0 d × 1 ; y i ] ∈ R d +1 (the same setup as adopted in [31, 2]); The third row contains fixed vectors { p i } i ∈ [ N +1] with ones, zeros, the example index, and indicator for being the covariate token (similar to a positional encoding vector):

(11) is different from out input format (3) for encoder-based TFs. The main difference is that ( x i , y i ) are in different tokens in (11), whereas ( x i , y i ) are in the same token in (3). The reason for the former (i.e., different tokens in decoder) is that we want to avoid every [ x i ; 0] token seeing the information of y i , since we will evaluate the loss at every token. The reason for the latter (i.e., the same token in encoder) is for presentation convenience: since we only evaluate the loss at the last token, it is not necessary to alternate between [ x i ; 0] and [ 0 ; y i ] to avoid information leakage.

We then feed H into a decoder TF to obtain the output ˜ H = DTF θ ( H ) ∈ R D × (2 N +1) with the same shape, and read out the prediction ̂ y N +1 from the ( d +1 , 2 N +1) -th entry of ˜ H = [ ˜ h i ] i ∈ [2 N +1] (the entry corresponding to the last missing test label): ̂ y N +1 = read y ( ˜ H ) := ( ˜ h 2 N +1 ) d +1 . The goal is to predict ̂ y N +1 that is close to y N +1 ∼ P y | x N +1 measured by proper losses.

The benefit of using the decoder architecture is that, during the pre-training phase, one can construct the training loss function by using all the predictions { ̂ y j } j ∈ [ N +1] , where ̂ y j gives the ( d +1 , 2 j -1) -th entry of ˜ H = [ ˜ h i ] i ∈ [2 N +1] for each j ∈ [ N +1] (the entry corresponding to the missing test label of the 2 j -1 'th token): ̂ y j = read y ,j ( ˜ H ) := ( ˜ h 2 j -1 ) d +1 . Given a loss function ℓ : R × R → R associated to a single response, the training loss associated to the whole input sequence can be defined by ℓ ( H ) = ∑ N +1 j =1 ℓ ( y j , ̂ y j ) . This potentially enables less training sequences in the pre-training stage, and some generalization bound analysis justifying this benefit was provided in [47].

C.3 Results

We discuss how our theoretical results upon encoder TFs can be converted to those of the decoder TFs. Taking the implementation of (ICGD) (a key mechanism that enables most basic ICL algorithms such as ridge regression; cf. Appendix D.1) as an example, this conversion is enabled by the following facts: (a) the input format (11) of decoders can be converted to the input format (3) of encoders by a 2-layer decoder TF; (b) the encoder TF that implements (ICGD) with input format (3), by a slight parameter modification, can be converted to a decoder TF that implements the (ICGD) algorithm with a converted input format.

Input format conversion Despite the difference between the input format (11) and (3), we show that there exists a 2-layer decoder TF that can convert the input format (11) to format (3). The proof can be found in Appendix C.4.

Proposition C.1 (Input format conversion) . There exists a 2-layer decoder TF DTF with 3 heads per layer, hidden dimension 2 and | | | θ | | | ≤ 12 such that upon taking input H of format (11) , it outputs ˜ H = DTF( H ) with

In particular, format (12) contains format (3) as a submatrix, by restricting to the { 1 , 2 , . . . , D -1 , D -2 , D } rows and { 2 , 4 , . . . , 2 N -2 , 2 N, 2 N +1 } columns.

Generalization TF constructions to decoder architecture The construction in Theorem D.1 can be generalized to using the input format (12) along with a decoder TF, by using the scratch pad within the last token to record the gradient descent iterates. Further, if we slightly change the normalization in MSK from 1 /i to 1 / (( i -1) ∨ 1) , then the same construction performs (ICGD) (with training examples { 1 , . . . , j } ) at every token i = 2 j +1 (corresponding to predicting at x j +1 ). Building on this extension, all our constructions in Section 3 and Section 4.2 can be generalized to decoder TFs.

C.4 Proof of Proposition C.1

For the simplicity of presentation, we write c i = ⌈ i/ 2 ⌉ , t i = mod( i +1 , 2) , u i = h i [1 : d ] ∈ R d +1 be the vector of first d entries of h i 6 , and let v i = h i [ d +1] be the ( d +1) -th entry of h i . With such notations, the input sequence H = [ h i ] i can be compactly written as

In the following, we construct the desired θ = ( θ (1) , θ (2) ) as follows.

Step 1: construction of θ (1) = ( θ (1) mattn , θ (1) mlp ) , so that MLP θ (1) mlp ◦ MAttn θ (1) mattn maps

For m ∈ { 0 , 1 } , we define matrices Q (1) m , K (1) m , V (1) m ∈ R D × D such that

for all i, j . By the structure of h i , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Now, for every i ,

̸

Notice that t i = 0 only when 2 | i , we then compute for i = 2 k that

Therefore, the θ (1) mattn = { ( Q (1) m , K (1) m , V (1) m ∈ R D × D ) } m ∈{ 0 , 1 } we construct above is indeed the desired attention layer. The existence of the desired θ (1) mlp is clear, and θ (1) mlp = ( W (1) 1 , W (1) 2 ) can further be chosen so that ∥ W (1) 1 ∥ op ≤ 1 , ∥ W (1) 2 ∥ op ≤ 1 .

Step 2: construction of θ (2) . For every m ∈ {-1 , 0 , 1 } , we define matrices Q (2) m , K (2) m , V (2) m ∈ R D × D such that

for all i, j . By the structure of h (1) i , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Now, for every i, j , we have

where the last equality follows from the fact that

Therefore,

Therefore, the θ (2) mattn = { ( Q (2) m , K (2) m , V (2) m ∈ R D × D ) } m ∈{-1 , 0 , 1 } we construct above maps

Finally, we only need to take a MLP layer θ (2) mlp = ( W (2) 1 , W (2) 2 ) with hidden dimension 2 that maps

which clearly exists and can be chosen so that ∥ W (2) ∥ op ≤ 1 , ∥ W (2) ∥ op ≤ 1 .

1 2 Combining the two steps above, we complete the proof of Proposition C.1.

D Mechanism: In-context gradient descent

Technically, the constructions in Section 3.1-3.2 rely on a new efficient construction for transformers to implement in-context gradient descent and its variants, which we present as follows. We begin by presenting the result for implementing (vanilla) gradient descent on convex empirical risks.

Compact notation of input We will often use shorthand y ′ i ∈ R defined as y ′ i = y i for i ∈ [ N ] and y ′ N +1 = 0 to simplify our notation, with which the input sequence H ∈ R D × ( N +1) can be compactly written as h i = [ x i ; y ′ i ; p i ] = [ x i ; y ′ i ; 0 D -d -3 ; 1; t i ] for i ∈ [ N +1] , where t i := 1 { i < N +1 } is the indicator for the training examples.

D.1 Gradient descent on convex empirical risk

Let ℓ ( · , · ) : R 2 → R be a loss function. Let ̂ L N ( w ) := 1 N ∑ N i =1 ℓ ( w ⊤ x i , y i ) denote the empirical risk with loss function ℓ on dataset { ( x i , y i ) } i ∈ [ N ] , and

denote the gradient descent trajectory on ̂ L N with initialization w 0 GD ∈ R d and learning rate η > 0 .

We require the partial derivative of the loss ∂ s ℓ : ( s, t ) ↦→ ∂ s ℓ ( s, t ) (as a bivariate function) to be approximable by a sum of relus, defined as follows.

Definition D.1 (Approximability by sum of relus) . A function g : R k → R is ( ε approx , R, M, C ) -approximable by sum of relus , if there exists a ' ( M,C ) -sum of relus' function

Definition D.1 is known to contain broad class of functions. For example, any mildly smooth k -variate function is approximable by a sum of relus for any ( ε approx , R ) , with mild bounds on ( M,C ) (Proposition B.1, building on results of Bach [4]). Also, any function that is a ( M,C ) -sum of relus itself (which includes all piecewise linear functions) is by definition (0 , ∞ , M, C ) -approximable by sum of relus.

We show that L steps of (ICGD) can be approximately implemented by an ( L +1) -layer transformer.

Theorem D.1 (Convex ICGD) . Fix any B w > 0 , L > 1 , η > 0 , and ε ≤ B w / (2 L ) . Suppose that

The loss ℓ ( · , · ) is convex in the first argument;

∂ s ℓ is ( ε, R, M, C ) -approximable by sum of relus with R = max { B x B w , B y , 1 } .

=

+

1

N

N

∑

i

=1

M

∑

m

=1

σ

(

⟨

,

⟩

)

×

h t

+1

N

+1

Qm

Vm

h t

N

+1

h t

N

+1

h t

i

h t

i

Km

a

m

σ

(

b

m

⟨

,

⟩ +

c

m

)

×

=

-

η

N

N

∑

i

=1

M

∑

m

=1

w t

+1

x

i

y

i

w t

x

i

y

i

x

i

w t

y

i

x

i

0

0

∂

s ℓ

( ⟨

,

⟩ ,

) ×

=

-

η

N

N

∑

i

=1

w t

+1

w t

x

i

w t

y

i

x

i

∂

sℓ

(

s

,

t

) = ∑

M

m

=1

a

m

⋅

σ

(

b

m

⋅

s

+

c

m

⋅

t

)

Attention

Gradient

descent

Weight

construction

Universal

approximation

h t

i

w t

x

i

y

i

=

Vm

x

i

0

0

= -

ηa

m

×

w t

x

i

y

i

,

,

,

x

i

Km

=

w t

x

i

y

i

y

i

Qm

=

w t

*

w t

1

b

m

×

*

c

m

×

Then, there exists an attention-only transformer TF 0 θ with ( L + 1) layers, max ℓ ∈ [ L ] M ( ℓ ) ≤ M heads within the first L layers, and M ( L +1) = 2 such that for any input data ( D , x N +1 ) such that

TF 0 θ ( H (0) ) approximately implements (ICGD) with initialization w 0 GD = 0 :

(Parameter space) For every ℓ ∈ [ L ] , the ℓ -th layer's output H ( ℓ ) = TF 0 θ (1: ℓ ) ( H (0) ) approximates ℓ steps of (ICGD): We have h ( ℓ ) i = [ x i ; y ′ i ; ̂ w ℓ ; 0 D -2 d -3 ; 1; t i ] for every i ∈ [ N +1] , where

Note that the bound scales as O ( ℓ ) , a linear error accumulation.

(Prediction space) The final output H ( L +1) = TF 0 θ ( H (0) ) approximates the prediction of L steps of (ICGD): We have h ( L +1) N +1 = [ x N +1 ; ̂ y N +1 ; ̂ w L ; 0 D -2 d -3 ; 1; t i ] , where ̂ y N +1 = 〈 ̂ w L , x N +1 〉 so that

Further, the transformer admits norm bound | | | θ | | | ≤ 2 + R +2 ηC .

The proof can be found in Appendix E.2. Theorem D.1 substantially generalizes that of von Oswald et al. [86] (which only does GD on square losses with a linear self-attention), and is simpler than the ones in Akyürek et al. [2] and Giannou et al. [32]. See Figure 4 for a pictorial illustration of the basic component of the construction, which implements a single step of gradient descent using a single attention layer (Proposition E.1).

Technically, we utilize the stability of convex gradient descent as in the following lemma (proof in Appendix E.3) to obtain the linear error accumulation in Theorem D.1; the error accumulation will become exponential in L in the non-convex case in general; see Lemma D.3(b).

Lemma D.1 (Composition of error for approximating convex GD) . Suppose f : R d → R is a convex function. Let w ⋆ ∈ arg min w ∈ R d f ( w ) , R ≥ 2 ∥ w ⋆ ∥ 2 , and assume that ∇ f is L f -smooth on B d 2 ( R ) . Let sequences { ̂ w ℓ } ℓ ≥ 0 ⊂ R d and { w ℓ GD } ℓ ≥ 0 ⊂ R d be given by ̂ w 0 = w 0 GD = 0 ,

for all ℓ ≥ 0 . Then as long as η ≤ 2 /L f , for any 0 ≤ L ≤ R/ (2 ε ) , it holds that ∥ ∥ ̂ w L -w L GD ∥ ∥ 2 ≤ Lε and ∥ ̂ w L ∥ 2 ≤ R 2 + Lε ≤ R .

D.2 Proximal gradient descent for regularized convex losses

Proximal gradient descent (PGD) is a variant of gradient descent that is suitable for minimizing regularized risks [66], in particular those with a non-smooth regularizer such as the ℓ 1 norm. In this section, we show that transformers can approximate PGD with similar quantitative guarantees as for GD in Appendix D.1.

Let ℓ ( · , · ) : R 2 → R be a loss function. Let ̂ L N ( w ) := 1 N ∑ N i =1 ℓ ( w ⊤ x i , y i ) + R ( w ) denote the regularized empirical risk with loss function ℓ on dataset { ( x i , y i ) } i ∈ [ N ] and regularizer R . To minimize ̂ L N , we consider the proximal gradient descent trajectory on ̂ L N with initialization w 0 = 0 ∈ R d and learning rate η > 0 :

where we denote ̂ L 0 ( w ) := 1 ∑ N ℓ ( w ⊤ x i , y i ) .

To approximate (ICPGD) by transformers, in addition to the requirement on the loss ℓ as in Theorem D.1, we additionally require the the proximal operator prox η R ( · ) to be approximable by an MLP layer (as a vector-valued analog of Definition D.1) defined as follows.

Definition D.2 (Approximability by MLP) . An operator P : R d → R d is ( ε, R, D, C ) -approximable by MLP , if there exists a there exists a MLP θ mlp = ( W 1 , W 2 ) ∈ R D × d × R d × D with hidden dimension D , ∥ W 1 ∥ op + ∥ W 2 ∥ op ≤ C ′ , such that sup ∥ w ∥ 2 ≤ R ∥ ∥ P ( w ) -MLP θ mlp ( w ) ∥ ∥ 2 ≤ ε .

The definition above captures the proximal operator prox η R for a broad class of regularizers, such as the (commonly-used) L 1 and L 2 regularizer listed in the following proposition, for all of which one can directly check that they can be exactly implemented by an MLP as stated below.

Proposition D.1 (Proximal operators for commonly-used regularizers) . For regularizer R in { λ ∥·∥ 1 , λ 2 ∥·∥ 2 2 , I B ∞ ( B ) ( · ) } , the operator prox η R : R d → R d is exactly approximable by MLP. More concretely, we have

For R = λ ∥·∥ 1 , prox η R is (0 , + ∞ , 4 d, 4 + 2 ηλ ) -approximable by MLP.

For R = λ 2 ∥·∥ 2 2 , prox η R is (0 , + ∞ , 2 d, 2 + 2 ηλ ) -approximable by MLP.

For R = I B ∞ ( B ) ( · ) , prox η R = Proj B ∞ ( B ) is (0 , + ∞ , 2 d, 2 + 2 B ) -approximable by MLP.

Theorem D.2 (Convex ICPGD) . Fix any B w > 0 , L > 1 , η > 0 , and ε + ε ′ ≤ B w / (2 L ) . Suppose that

The loss ℓ ( · , · ) is convex in the first argument;

∂ s ℓ is ( ε, R, M, C ) -approximable by sum of relus with R = max { B x B w , B y , 1 } .

R convex, and the proximal operator prox η R ( w ) is ( ηε ′ , R ′ , D ′ , C ′ ) -approximable by MLP with R ′ = sup ∥ w ∥ 2 ≤ B w ∥ ∥ w + η ∥ ∥ 2 + ηε .

Then there exists a transformer TF θ with ( L +1) layers, max ℓ ∈ [ L ] M ( ℓ ) ≤ M heads within the first L layers, M ( L +1) = 2 , and hidden dimension D ′ such that, for any input data ( D , x N +1 ) such that

TF θ ( H (0) ) approximately implements (ICGD):

(Parameter space) For every ℓ ∈ [ L ] , the ℓ -th layer's output H ( ℓ ) = TF θ (1: ℓ ) ( H (0) ) approximates ℓ steps of (ICGD): We have h ( ℓ ) i = [ x i ; y ′ i ; ̂ w ℓ ; 0 D -2 d -3 ; 1; t i ] for every i ∈ [ N +1] , where

(Prediction space) The final output H ( L +1) = TF θ ( H (0) ) approximates the prediction of L steps of (ICGD): We have h ( L +1) N +1 = [ x N +1 ; ̂ y N +1 ; ̂ w L ; 0 D -2 d -3 ; 1; t i ] , where ̂ y N +1 = 〈 ̂ w L , x N +1 〉 so that

Further, the weight matrices have norm bounds | | | θ | | | ≤ 3 + R +2 ηC + C ′ .

The proof of Theorem D.2 is essentially similar to the proof of Theorem D.1, using the following generalized version of Lemma D.1.

Lemma D.2 (Composition of error for approximating convex PGD) . Suppose f : R d → R is a convex function and R is a convex regularizer. Let w ⋆ ∈ arg min w ∈ R d f ( w ) + R ( w ) , R ≥ 2 ∥ w ⋆ ∥ 2 , and assume that ∇ f is L f -smooth on B d 2 ( R ) . Let sequences { ̂ w ℓ } ℓ ≥ 0 ⊂ R d and { w ℓ GD } ℓ ≥ 0 ⊂ R d be given by ̂ w 0 = w 0 GD = 0 ,

for all ℓ ≥ 0 . Then as long as η ≤ 2 /L f , for any 0 ≤ L ≤ R/ (2 ε ) , it holds that ∥ ∥ ̂ w L -w L GD ∥ ∥ 2 ≤ Lε and ∥ ̂ w L ∥ 2 ≤ R 2 + Lε ≤ R .

The proof of the above lemma is done by utilizing the non-expansiveness of the PGD operator w ↦→ prox η R ( w -η ∇ f ( w )) and otherwise following the same arguments as for Lemma D.1.

D.3 Gradient descent on two-layer neural networks

We now move beyond the convex setting by showing that transformers can implement gradient descent on two-layer neural networks in context.

Suppose that the prediction function pred( x ; w ) := ∑ K k =1 u k r ( v ⊤ k x ) is given by a two-layer neural network, parameterized by w = [ v k ; u k ] k ∈ [ K ] ∈ R K ( d +1) . Consider the empirical risk minimization problem:

where W is a bounded domain. For the sake of simplicity, in the following discussion we assume that Proj W can be exactly implemented by a MLP layer (e.g. W = B ∞ ( R w ) for some R w > 0 ).

Theorem D.3 (Approximate ICGD on two-layer NNs) . Fix any B v , B u > 0 , L ≥ 1 , η > 0 , and ε > 0 . Suppose that

Both the activation function r and the loss function ℓ is C 4 -smooth;

W is a closed domain such that W ⊂ { w = [ v k ; u k ] k ∈ [ K ] ∈ R K ( d +1) : ∥ v k ∥ 2 ≤ B v , | u k | ≤ B u } , and Proj W = MLP θ mlp for some MLP layer θ mlp with hidden dimension D w and ∣ ∣ ∣ ∣ ∣ ∣ θ mlp ∣ ∣ ∣ ∣ ∣ ∣ ≤ C w ;

Then there exists a (2 L ) -layer transformer TF θ with

where O ( · ) hides the constants that depend on K , the radius parameters B x , B y , B u , B v and the smoothness of r and ℓ , such that for any input data ( D , x N +1 ) such that input sequence H (0) ∈ R D × ( N +1) takes form (3) , TF θ ( H (0) ) approximately implements in-context gradient descent on risk (13): For every ℓ ∈ [ L ] , the 2 ℓ -th layer's output h (2 ℓ ) i = [ x i ; y ′ i ; ̂ w ℓ ; 0 ; 1; t i ] for every i ∈ [ N +1] , and

where ∥ ∥ ε ℓ -1 ∥ ∥ 2 ≤ ε is an error term.

As a direct corollary, the transformer constructed above can approximate the true gradient descent trajectory { w ℓ GD } ℓ ≥ 0 on (16), defined as w 0 GD = 0 and w ℓ +1 GD = w ℓ GD -η ∇ ̂ L N ( w ℓ GD ) for all ℓ ≥ 0 .

Corollary D.1 (Approximating multi-step ICGD on two-layer NNs) . For any L ≥ 1 , under the same setting as Theorem D.3, the (2 L ) -layer transformer TF θ there approximates the true gradient descent trajectory { w ℓ GD } ℓ ≥ 0 : For the intermediate iterates { ̂ w ℓ } ℓ ∈ [ L ] considered therein, we have

where L f = sup w ∈W ∥ ∥ ∥ ∇ 2 ̂ L N ( w ) ∥ ∥ ∥ op denotes the smoothness of ̂ L N within W .

Remark on error accumulation Note that in Corollary D.1, the error accumulates exponentially in ℓ rather than linearly as in Theorem D.1. This is as expected, since gradient descent on non-convex objectives is inherently unstable at a high level (a slight error added upon each step may result in a drastically different trajectories); technically, this happens as the stability-like property Lemma D.1 no longer holds for the non-convex case.

Corollary D.1 is a simple implication of Theorem D.3 and Part (b) of the following convergence and trajectory closeness result for inexact gradient descent. For any closed convex set W ⊂ R d , any function f : W → R , and any initial point w ∈ W , let

denote the gradient mapping at w with step size η , a standard measure of stationarity in constrained optimization [63]. Note that G f W ,η ( w ) = ∇ f ( w ) when w -η ∇ f ( w ) ∈ W (so that the projection does not take effect).

Lemma D.3 (Convergence and trajectory closeness of inexact GD) . Suppose f : W → R , where W ⊂ R d is a convex closed domain and ∇ f is L f -Lipschitz on W . Let sequence { ̂ w ℓ } ℓ ≥ 0 ⊂ R d be given by ̂ w 0 = w 0 ,

for all ℓ ≥ 0 . Then the following holds.

(a) As long as η ≤ 1 /L f , for all L ≥ 1 ,

(b) Let the sequences { w ℓ GD } ℓ ≥ 0 ⊂ R d and be given by w 0 GD = w 0 and w ℓ +1 GD = Proj W ( w ℓ GD -η ∇ f ( w ℓ GD )) . Then it holds that

E Proofs for Section D

E.1 Approximating a single GD step

Proposition E.1 (Approximating a single GD step by a single attention layer) . Let ℓ ( · , · ) : R 2 → R be a loss function such that ∂ 1 ℓ is ( ε, R, M, C ) -approximable by sum of relus with R = max { B x B w , B y , 1 } . Let ̂ L N ( w ) := 1 N ∑ N i =1 ℓ ( w ⊤ x i , y i ) denote the empirical risk with loss function ℓ on dataset { ( x i , y i ) } i ∈ [ N ] .

Then, for any ε > 0 , there exists an attention layer θ = { ( Q m , K m , V m ) } m ∈ [ M ] with M heads such that, for any input sequence that takes form h i = [ x i ; y ′ i ; w ; 0 D -2 d -3 ; 1; t i ] with ∥ w ∥ 2 ≤ B w , it gives output ˜ h i = [Attn θ ( H )] i = [ x i ; y ′ i ; ˜ w ; 0 D -2 d -3 ; 1; t i ] for all i ∈ [ N +1] , where

Further, | | | θ | | | ≤ 2 + R +2 ηC .

Proof of Proposition E.1. As ∂ s ℓ is ( ε, R, M, C ) -approximable by sum of relus, there exists a function f : [ -R,R ] 2 → R of form

Next, for every m ∈ [ M ] , we define matrices Q m , K m , V m ∈ R D × D such that

for all i, j ∈ [ N + 1] . As the input has structure h i = [ x i ; y ′ i ; w ; 0 D -2 d -3 ; 1; t i ] , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Consequently, | | | θ | | | ≤ 2 + R +2 ηC .

Now, for every i, j ∈ [ N +1] , we have

where the last equality follows from the bound

so that the above relu equals 0 if t j ≤ 0 . Therefore,

Thus letting the attention layer θ = { ( V m , Q m , K m ) } m ∈ [ M ] , we have

where the error vector ε ∈ R satisfies

which proves part (a).

We now prove part (b), which requires constructing the last attention layer θ ( L +1) . Recall h ( L ) i = [ x i ; y ′ i ; ̂ w L ; 0 D -2 d -3 ; 1; t i ] for all i ∈ [ N + 1] . We construct a 2-head attention layer θ ( L +1) = { ( Q ( L +1) m , K ( L +1) m , V ( L +1) m ) } m =1 , 2 such that for every i, j ∈ [ N +1] ,

Note that the weight matrices have norm bound

Then we have

This is the desired result.

E.2 Proof of Theorem D.1

We first prove part (a), which requires constructing the first L layers of θ . Note that by our precondition L ≤ B w / (2 ε ) .

By our precondition, the partial derivative of the loss ∂ s ℓ is ( ε, R, M, C ) -approximable by sum of relus. Therefore we can apply Proposition E.1 to obtain that, there exists a single attention layer θ (1) = { ( Q m , K m , V m ) } m ∈ [ M ] with M heads (and norm bounds specified in Proposition E.1), such that for any w with ∥ w ∥ 2 ≤ B w , the attention layer Attn θ (1) maps the input h i = [ x i ; y ′ i ; w ; 0 D -2 d -3 ; 1; t i ] to output h ′ i = [ x i ; y ′ i ; ̂ w ; 0 D -2 d -3 ; 1; t i ] for all i ∈ [ N +1] , where

Consider the L -layer transformer θ 1: L = ( θ (1) , . . . , θ (1) ) which stacks the same attention layer θ (1) for L times, and for the given input h (0) i = [ x i ; y ′ i ; w 0 ; 0 D -2 d -3 ; 1; t i ] , its ℓ -th layer's output h ( ℓ ) i = [ x i ; y ′ i ; ̂ w ℓ ; 0 D -2 d -3 ; 1; t i ] .

We now inductively show that ∥ ∥ ̂ w ℓ ∥ ∥ 2 ≤ B w and ∥ ∥ ̂ w ℓ -w ℓ GD ∥ ∥ 2 ≤ ℓε for all ℓ ∈ [ L ] . The base case of ℓ = 0 is trivial. Suppose the claim holds for ℓ . Then for ℓ +1 ≤ L ≤ B w / (2 ε ) , the sequence { ̂ w i } i ≤ ℓ +1 and { w i GD } i ≤ ℓ +1 satisfies the precondition of the error composition lemma (Lemma D.1) with error bound ε , from which we obtain ∥ ∥ ̂ w ℓ +1 ∥ ∥ 2 ≤ B w and

This finishes the induction, and gives the following approximation guarantee for all ℓ ∈ [ L ] :

Above, (i) uses the identity t = σ ( t ) -σ ( -t ) . Further by part (a) we have

This proves part (b), and also finishes the proof Theorem D.1 where the overall ( L + 1) -layer attention-only transformer is given by TF 0 θ with

E.3 Proof of Lemma D.1

As f is a convex, L f smooth function on B d 2 ( R ) , the mapping T η : w ↦→ w -η ∇ f ( w ) is nonexpansive in ∥·∥ 2 : Indeed, for any w , w ′ ∈ B d 2 ( R ) we have

Above, (i) uses the property ⟨ w -w ′ , ∇ f ( w ) -∇ f ( w ′ ) ⟩ ≥ 1 L f ∥∇ f ( w ) -∇ f ( w ′ ) ∥ 2 2 for smooth convex functions [63, Theorem 2.1.5]; (ii) uses the precondition that η ≤ 2 /L f .

The lemma then follows directly by induction on L . The base case of L = 0 follows directly by assumption that ̂ w 0 = w 0 GD ∈ B d 2 ( R/ 2) . Suppose the claim holds for iterate L . For iterate L +1 ≤ R/ (2 ε ) , we have

Above, (i) uses the non-expansiveness, and (ii) uses the inductive hypothesis. Similarly, by our assumption w ⋆ = T η ( w ⋆ ) ,

This finishes the induction.

E.4 Convex ICGD with ℓ 2 regularization

In the same setting as Theorem D.1, consider the ICGD dynamics over an ℓ 2 -regularized empirical risk:

with initialization w 0 GD ∈ R d and learning rate η > 0 , where ̂ L λ N ( w ) := ̂ L N ( w ) + λ 2 ∥ w ∥ 2 2 denotes the ℓ 2 -regularized empirical risk.

Corollary E.1 (Convex ICGD with ℓ 2 regularization) . Fix any B w > 0 , L > 1 , η > 0 , and ε < B x B w . Suppose the loss ℓ ( · , · ) is convex in the first argument, and ∂ s ℓ is ( ε, R, M, C ) -approximable by sum of relus with R = max { B x B w , B y , 1 } .

Then, there exists an attention-only transformer TF 0 θ with ( L +1) layers, max ℓ ∈ [ L ] M ( ℓ ) ≤ M +1 heads within the first L layers, and M ( L +1) = 2 such that for any input data ( D , x N +1 ) with

TF 0 θ ( H (0) ) approximately implements (ICGDℓ 2 ):

(Parameter space) For every ℓ ∈ [ L ] , the ℓ -th layer's output H ( ℓ ) = TF θ (1: ℓ ) ( H (0) ) approximates ℓ steps of (ICGDℓ 2 ): We have h ( ℓ ) i = [ x i ; y ′ i ; ̂ w ℓ ; 0 D -2 d -3 ; 1; t i ] for every i ∈ [ N +1] , where

(Prediction space) The final output H ( L +1) = TF θ ( H (0) ) approximates the prediction of L steps of (ICGDℓ 2 ): We have h ( L +1) N +1 = [ x N +1 ; ̂ y N +1 ; ̂ w L ; 0 D -2 d -3 ; 1; 0] , where

Further, the transformer admits norm bound | | | θ | | | ≤ 2 + R +(2 C + λ ) η.

Proof. This construction is the same as in the proof of Theorem D.1, except that within each layer ℓ ∈ [ L ] , we add one more attention head ( Q ( ℓ ) , K ( ℓ ) , V ( ℓ ) ) ⊂ R D × D which when acting on its input h ( ℓ -1) i = [ ∗ ; ∗ ; ̂ w ℓ -1 ; 1; ∗ ] gives

for all i, j ∈ [ N + 1] . Note that ∥ ∥ Q ( ℓ ) ∥ ∥ op = ∥ ∥ K ( ℓ ) ∥ ∥ op = 1 , and ∥ ∥ V ( ℓ ) ∥ ∥ op = ηλ . Further, it is straightforward to check that the output of this attention head on every h ( ℓ ) i is

Adding this onto the original output of the ℓ -th layer exactly implements the gradient of the regularizer w ↦→ λ 2 ∥ w ∥ 2 2 . The rest of the proof follows by repeating the argument of Theorem D.1, and combining the norm bound for the additional attention head here with the norm bound therein.

E.5 Proof of Theorem D.3

We only need to prove the following single-step version of Theorem D.3.

Proposition E.2. Under the assumptions of Theorem D.3, there exists a 2-layer transformer TF θ with the same bounds on the number of heads, hidden dimension and the norm, such that for any input data ( D , x N +1 ) and any w ∈ R d , TF θ maps

where

Before we present the formal (and technical) proof of Proposition E.2, we first provide some intuitions. To begin with, we first note that

where ∂ 1 ℓ is the partial derivative of ℓ with respect to the first component, and

Therefore, the basic idea is that we can use an attention layer to approximate ( x i , w ) ↦→ pred( x i ; w ) , then use an MLP layer to implement (pred( x i ; w ) , y ′ i , t i ) ↦→ 1 { i < N +1 } · ∂ 1 ℓ (pred( x i ; w ) , y i ) , and then use an attention layer to compute the gradient descent step w ↦→ w -η ∇ L N ( w ) , and finally use an MLP layer to implement the projection into W .

Based on the observations above, we now present the proof of Proposition E.2.

Proof of Proposition E.2. We write D 0 = d +1+ K ( d +1) be the length of the vector [ x i ; y i ; w ] . We also define

Let us fix ε r , ε p , ε ℓ > 0 that will be specified later in proof (see (18)). By our assumption and Proposition B.1, the following facts hold.

The function r ( t ) is ( ε r , R 1 , M 1 , C 1 ) for R 1 = max { B x B u , 1 } , M 1 ≤ ˜ O ( C 2 1 ε -2 r ) , where C 1 depends only on R 1 and the C 2 -smoothness of r . Therefore, there exists

such that sup t ∈ [ -R 1 ,R 1 ] | r ( t ) -r ( t ) | ≤ ε r .

The function ( t, y ) ↦→ ∂ 1 ℓ ( t, y ) is ( ε ℓ , R 2 , M 2 , C 2 ) for R 2 = max { KB r , B y , 1 } M 2 ≤ ˜ O ( C 2 2 ε -2 ℓ ) , where C 2 depends only on R 2 and the C 3 -smoothness of ∂ 1 ℓ . Therefore, there exists

such that sup ( t,y ) ∈ [ -R 2 ,R 2 ] 2 | g ( t, y ) -∂ 1 ℓ ( t, y ) | ≤ ε ℓ .

The function ( s, t ) ↦→ s · r ′ ( t ) is ( ε p , R 3 , M 3 , C 3 ) for R 3 = max { B x B u , B g B u , 1 } , M 3 ≤ ˜ O ( C 2 3 ε -2 p ) , where C 3 depends only on R 3 and the C 3 -smoothness of r ′ . Therefore, there exists

In the following, we proceed to construct the desired transformer step by step.

Step 1: construction of θ (1) attn . We consider the matrices { Q (1) k,m , K (1) k,m , V (1) k,m } k ∈ [ K ] ,m ∈ [ M 1 ] so that for all i, j ∈ [ N +1] , we have

As the input has structure h i = [ x i ; y ′ i ; w ; 0 ; 1; t i ] , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

A simple calculation shows that

For simplicity, we denote pred( x ; w ) := ∑ K k =1 u k r ( ⟨ v k , x ⟩ ) in the following analysis. Thus, letting the attention layer θ (1) attn = { ( V (1) k,m , Q (1) k,m , K (1) k,m ) } ( k,m ) , we have

Step 2: construction of θ (1) mlp . We pick matrices W 1 , W 2 so that W 1 maps

and W 2 ∈ R D × M 3 with entries being ( W 2 ) ( j,m ) = c 2 m 1 { j = D 0 +2 } . It is clear that ∥ W 1 ∥ op ≤ R 2 +1 , ∥ W 2 ∥ op ≤ C 2 . Then we have

In the following, we abbreviate g i = 1 { t j = 1 } · g (pred( x i ; w ) , y ′ i ) . Hence, θ mlp maps

By the definition of the function g , for each i ∈ [ N ] ,

where L ℓ := max | t |≤ KB r , | y |≤ B y ∣ ∣ ∂ 2 tt ℓ ( t, y ) ∣ ∣ is the smoothness of ∂ 1 ℓ . Also, g N +1 = 0 by definition.

Step 3: construction of θ (2) attn . We consider the matrices { Q (2) k, 1 ,m , K (2) k, 1 ,m , V (2) k, 1 ,m } k ∈ [ K ] ,m ∈ [ M 3 ] so that for all i, j ∈ [ N +1] , we have

Wefurther consider the matrices { Q (2) k, 2 ,m , K (2) k, 2 ,m , V (2) k, 2 ,m } k ∈ [ K ] ,m ∈ [ M 1 ] so that for all i, j ∈ [ N +1] , we have

By the structure of the input h (1) i , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Furthermore, a simple calculation shows that

where the summation is taken over all possibilities of the tuple ( k, w, m ) , i.e. over the union of [ K ] ×{ 1 } × [ M 3 ] and [ K ] ×{ 2 } × [ M 1 ] .

By our definition, we have | P ( s, t ) -sr ′ ( t ) | ≤ ε p for all s, t ∈ [ -R 3 , R 3 ] . Therefore, for each i ∈ [ N ] , k ∈ [ K ] ,

where L r := max | t |≤ B x B u | r ′ ( t ) | is the upper bound of r ′ . Similarly, for each i ∈ [ N ] , k ∈ [ K ] , we have

As for the case i = N +1 , we have g N +1 = 0 and | P ( u k g N +1 , ⟨ v k , x N +1 ⟩ ) | ≤ ε p for each k ∈ [ K ] by defintion. Combining these estimations and using (16) and (17), we can conclude that

Thus, to ensure ∥ ∥ ∥ η -1 g ( w ) + ∇ ̂ L N ( w ) ∥ ∥ ∥ 2 ≤ ε , we only need to choose ε p , ε ℓ , ε r as

Thus, letting the attention layer θ (2) attn = { ( V (2) k,w,m , Q (2) k,w,m , K (2) k,w,m ) } ( k,w,m ) , we have

Step 4: construction of θ (2) mlp . We only need to pick θ (2) mlp so that it maps

MLP

θ

(2)

mlp

-- - - - →

h

(1

i

= [

x

i

;

y

′

i

;

w

+

η

g

(

w

); pred(

x

i

;

w

);

g

i

;

0

; 1;

By our assumption on the map Proj W , this is easy.

Combining the four steps above and taking θ = ( θ (1) attn , θ (1) mlp , θ (2) attn , θ (2) mlp ) completes the proof.

E.6 Proof of Lemma D.3

For every ℓ ≥ 0 , define the intermediate iterates (before projection)

We first prove part (a). We begin by deriving a relation between ∥ ∥ ̂ w ℓ +1 -̂ w ℓ ∥ ∥ 2 2 and ∥ ∥ ∥ η G f W ,η ( ̂ w ℓ ) ∥ ∥ ∥ 2 2 . Let ˜ w ℓ + 1 2 := ̂ w ℓ -η ∇ f ( ̂ w ℓ ) and ˜ w ℓ +1 := Proj W ( ˜ w ℓ + 1 2 ) denote the exact projected gradient iterate starting from ̂ w ℓ . We have

Above, (i) uses the inequality ∥ a -b ∥ 2 2 ≥ 1 2 ∥ a ∥ 2 2 -∥ b ∥ 2 2 ; (ii) uses the fact that projection to a convex set is a non-expansion; (iii) uses the definition of the gradient mapping.

By the L f -smoothness of f within W , we have

.

5)

t

i

]

h

(2)

i

= [

x

i

;

y

′

i

; Proj

W

(

w

-

η

g

(

w

)); 0; 0;

0

; 1;

t

i

]

.

Above, (i) uses the property 〈 ̂ w ℓ +1 -̂ w ℓ + 1 2 , ̂ w ℓ +1 -̂ w ℓ 〉 ≤ 0 of the projection ̂ w ℓ +1 = Proj W ( ̂ w ℓ + 1 2 ) (using ̂ w ℓ ∈ W ); (ii) uses L f / 2 ≤ 1 / (2 η ) by our choice of η ≤ 1 /L f ; (iii) uses (19).

Rearranging and summing the above over ℓ = 0 , . . . , L -1 , we obtain

Dividing both sides by ηL/ 8 yields part (a).

Next, we prove part (b). Let C := 1 + ηL f . We prove by induction that

for all ℓ ≥ 0 . The base case of ℓ = 0 follows by definition that ̂ w 0 = w 0 GD = w 0 . Suppose the result holds for ℓ . Then for ℓ +1 , we have

Above, (i) uses again the non-expansiveness of the convex projection Proj W ; (ii) uses the fact that the operator w ↦→ w -η ∇ f ( w ) is (1 + ηL f ) = C -Lipschitz; and (iii) uses the inductive hypothesis. This proves the case for ℓ +1 and thus finishes the induction. We can further relax (20) into

This proves part (b).

F Proofs for Section 3.1

F.1 Proof of Theorem 4

Fix λ ≥ 0 , 0 ≤ α ≤ β with κ := β + λ α + λ , and B w > 0 , and consider any in-context data D such that the precondition of Theorem 4 holds. Let

denote the ridge regression loss in (ICRidge), so that w λ ridge = arg min w ∈ R d L ridge ( w ) . It is a standard result that ∇ 2 L ridge ( w ) = X ⊤ X /N + λ I d , so that L ridge is ( α + λ ) -strongly convex and ( β + λ ) -smooth over R d .

Consider the gradient descent algorithm on the ridge loss

with initialization, learning rate, and number of steps

By standard convergence results for strongly convex and smooth functions (Proposition B.2), we have for all t ≥ 1 that

Further, we have

It remains to construct a transformer to approximate w T GD . Notice that the problem (ICRidge) corresponds to an ℓ 2 -regularized ERM with the square loss ℓ ( s, t ) := 1 2 ( s -t ) 2 , whose partial derivative ∂ s ℓ ( s, t ) = s -t is exactly a sum of two relus:

In particular, this shows that ∂ s ℓ ( s, t ) is (0 , R, 2 , 4) -approximable for any R > 0 , in particular for R = max { B x B w , B y , 1 } .

Therefore, we can apply Corollary E.1 with the square loss ℓ , learning rate η , regularization strength λ and accuracy parameter ε = 0 to obtain that there exists an attention-only transformer TF 0 θ with ( T +1) := L layers such that the final output h ( L ) N +1 = [ x N +1 ; ̂ y N +1 ; ∗ ] with

and number of heads M ( ℓ ) = 3 for all ℓ ∈ [ L -1] (can be taken as 2 in the unregularized case λ = 0 directly by Theorem D.1), and M ( L ) = 2 . Further, θ admits norm bound | | | θ | | | ≤ 2 + R + 8+ λ β + λ ≤ 3 R +8( β + λ ) -1 +1 ≤ 4 R +8( β + λ ) -1 .

Combining (21) and (22), we obtain that

Further, we have read w ( h T i ) = w T GD for all i ∈ [ N +1] , where read w ( h ) := h ( d +2):(2 d +1) (cf. Corollary E.1), so that ∥ read w ( h T i ) -w λ ridge ∥ 2 ≤ ε/B x as shown above. This finishes the proof.

F.2 Statistical analysis of in-context least squares

Consider the standard least-squares algorithm A LS and least-squares estimator ̂ w LS ∈ R d defined as

For any distribution P over ( x , y ) ∈ R d × R and any estimator w ∈ R d , let

denote the expected risk of w over a new test example ( x ′ , y ′ ) ∼ P .

Assumption A (Well-posedness for learning linear predictors) . We say a distribution P on R d × R is well-posed for learning linear predictors, if ( x , y ) ∼ P satisfies

∥ x ∥ 2 ≤ B x and | y | ≤ B y almost surely;

The covariance Σ P := E P [ xx ⊤ ] satisfies λ min I d ⪯ Σ P ⪯ λ max I d , with 0 < λ min ≤ λ max , and κ := λ max /λ min .

The whitened vector Σ -1 / 2 P x is K 2 -sub-Gaussian for some K ≥ 1 .

The best linear predictor w ⋆ P := E P [ xx ⊤ ] -1 E P [ x y ] satisfies ∥ w ⋆ P ∥ 2 ≤ B ⋆ w .

We have E [( y -⟨ x , w ⋆ P ⟩ ) 2 | x ] ≤ σ 2 with probability one (over x

Further, we say P is well-posed with canonical parameters if

The following result bounds the excess risk of least squares under Assumption A with a clipping operation on the predictor; the clipping allows the result to only depend on the second moment of the noise (cf. Assumption A(5)) instead of e.g. its sub-Gaussianity, and also makes the result convenient to be directly translated to a result for transformers.

Proposition F.1 (Guarantees for in-context least squares) . Suppose distribution P satisfies Assumption A. Then as long as N ≥ O ( dK 4 log(1 /δ )) , we have the following:

(a) The (clipped) least squares predictor achieves small expected excess risk (fast rate) over the best linear predictor: For any clipping radius R ≥ B y ,

(b) We have P ( E cov ∩ E w ) ≥ 1 -δ/ 10 , where

Proof. We first show P ( E cov ) ≥ 1 -δ/ 20 . Let ̂ Σ := 1 N ∑ N i =1 x i x ⊤ i , and let the whitened covariance and noise variables be denoted as

Also let z i := y i -⟨ x i , w ⋆ P ⟩ denote the 'noise' variables. Note that

is exactly a covariance concentration of the whitened vectors { ˜ x i } i ∈ [ N ] . Recall that E [ ˜ x i ˜ x ⊤ i ] = I d , and ˜ x i are K 2 -sub-Gaussian by assumption. Therefore, we can apply [85, Theorem 4.6.1], we have with probability at least 1 -δ/ 10 that

Setting N ≥ O ( K 4 ( d +log(1 /δ ))) ensures that the right-hand side above is at most 1 / 2 , on which event we have

).

i.e. E cov holds. This shows that P ( E c cov ) ≤ δ/ 10 .

Next, we show (24). Using E cov , we decompose the risk as

Above, (i) follows by assumption that | y N +1 | ≤ B y ≤ R almost surely, so that removing the clipping can only potentially increase the distance in the first term, and the square loss is upper bounded by 1 2 · (2 R ) 2 almost surely in the second term; (ii) follows by the fact that E x N +1 ,y N +1 [ ⟨ ̂ w LS -w ⋆ P , x N +1 ⟩ ( ⟨ w ⋆ P , x N +1 ⟩ -y N +1 )] = 0 by the definition of w ⋆ P , as well as the fact that 1 { E cov } is independent of ( x N +1 , y N +1 ) .

It thus remains to bound E D [ 1 2 ∥ ̂ w LS -w ⋆ P ∥ 2 Σ P 1 { E cov } ] . Note that on the event E cov , we have

Therefore,

Note that E [ ˜ x i z i ] = Σ -1 / 2 P E [ x i ( y i -⟨ w ⋆ P , x i ⟩ )] = 0 . Therefore, taking expectation on the above (over D ), we get

Above, (i) follows by conditioning on x 1 and using Assumption A(5). Combining with (28), we obtain

This proves (24).

Finally, we show P ( E cov ∩ E w ) ≥ 1 -δ/ 10 . Using (29) and Σ P ⪰ λ min I d by assumption, we get

Therefore, using an argument similar to Chebyshev's inequality,

This implies that

F.3 Proof of Corollary 5

The proof follows by first checking the well-conditionedness of the data D (cf. (5)) with high probability, then invoking Theorem 4 (for approximation least squares) and Proposition F.1 (for the statistical power of least squares).

First, as P satisfies Assumption A, by Proposition F.1, as long as N ≥ O ( K 4 ( d +log(1 /δ ))) , we have with probability at least 1 -δ/ 10 that event E cov ∩ E w holds. On this event, we have

and thus the dataset D is well-conditioned (in the sense of (5)) with parameters α = λ min / 2 , β = 2 λ max , and B w defined as above. Note that the condition number of ̂ Σ is upper bounded by β/α = 4 λ max /λ min ≤ 4 κ , where κ is the upper bound on the condition number of Σ P as in Assumption A(c).

Define parameters

Note that B w ≤ O ( B ⋆ w + √ B 2 y /λ min ) by the above choice of δ .

We can thus apply Theorem 4 in the unregularized case ( λ = 0 ) to obtain that, there exists a transformer θ with max ℓ ∈ [ L ] M ( ℓ ) ≤ 3 , | | | θ | | | ≤ 4 R +4 /λ max (with R = max { B x B w , B y , 1 } ), and number of layers

such that on E cov ∩ E w (so that D is well-conditioned), we have (choosing the clipping radius in ˜ read y ( · ) = clip B y ( read y ( · )) to be B y ):

We now bound the excess risk of the above transformer. Combining Proposition F.1 and (32), we have

Above, (i) uses the approximation guarantee (32) as well as Proposition F.1(a) (with clipping radius B y ). This proves the desired excess risk guarantee.

Finally, under the canonical choice of parameters (23), the bounds for L, M, | | | θ | | | simplify to

and the requirement for N simplifies to N ≥ O ( d +log(1 /δ )) = ˜ O ( d ) (as K = Θ(1) ). This proves the claim about the required N and L .

F.4 Proof of Corollary 6

Fix parameters δ, ε > 0 to be specified later and a large universal constant C 0 . Let us set

Consider the following good events (below ε = [ ε i ] i ∈ [ N ] ∈ R N is given by ε i = y i -⟨ w ⋆ , x i ⟩ )

and we define E := E π ∩ E w ∩ E b ∩ E b,N +1 . Under the event E , the problem (ICRidge) is wellconditioned and ∥ w λ ridge ∥ ≤ B w / 2 (by Lemma F.1).

Therefore, Theorem 4 implies that for κ = α + λ β + λ , there exists a L = ⌈ 2 κ log( B w /ε ) ⌉ + 1 -layer transformer θ with prediction ̂ y N +1 := ˜ read y (TF 0 θ ( H )) (clipped by B y ), such that under the good event E , we have ̂ y N +1 = clip B y ( ⟨ x N +1 , ̂ w ⟩ ) and ∥ ̂ w -w λ ridge ∥ ≤ ε .

In the following, we show that θ is indeed the desired transformer (when ε and δ is suitably chosen). Notice that we have

and we analyze these two parts separately.

Prediction risk under good event E . We first note that

where the inequality is because y N +1 ∈ [ -B y , B y ] under the good event E . Notice that by our construction, under the good event E , ̂ w = ̂ w ( D ) depends only on the dataset D 7 . Therefore, we have ∥ ̂ w ( D ) -w λ ridge ( D ) ∥ ≤ ε as long as the event E 0 := E π ∩ E w ∩ E b holds for ( w ⋆ , D ) . Thus, under E 0 ,

and we also have

Recall that 2 BayesRisk π = E w ⋆ , D ∥ w λ ridge -w ⋆ ∥ 2 2 + σ 2 . Note that 2 BayesRisk π ≤ 1 + σ 2 by definition. Therefore, we can conclude that

Prediction risk under bad event E c . Notice that

We can upper bound P ( E c ) = P ( E c π ∪ E c w ∪ E c b ∪ E c b,N +1 ) by Lemma B.1, Lemma B.2 and the sub-Gaussian tail bound:

Thus, as long as N ≥ 8 log(12 /δ ) , we have P ( E c ) ≤ δ . Further, a simple calculation yields

Notice that y N +1 | w ⋆ ∼ N (0 , ∥ w ⋆ ∥ 2 2 + σ 2 ) , hence E y 4 N +1 = 3 E ( ∥ w ⋆ ∥ 2 2 + σ 2 ) 2 ≤ 3(3+2 σ 2 + σ 4 ) ≤ B 4 y . Thus, we can conclude that

Choosing ε and δ . Combining the inequalities above, we have

To ensure 1 2 E ( ̂ y N +1 -y N +1 ) 2 ≤ BayesRisk π + ε , we only need to take ( ε, δ ) so that the following constraints are satisfied:

Therefore, it suffices to take δ = c 0 log 2 ( N ) ( ε 2 1+ σ 2 ) 2 for some small constant c 0 , then as long as

our choice of ε and δ is feasible. Note that κ ≤ O ( 1 + σ -2 ) , and hence under such choice of ( ε, δ ) , we have L = O (log(1 /ε )) and | | | θ | | | = ˜ O ( √ d ) . This is the desired result.

Lemma F.1. Under the event E π ∩ E w , we have ∥ ∥ w λ ridge ∥ ∥ 2 ≤ O ( B ⋆ w + σ ) .

Proof of Lemma F.1. By the definition of w λ ridge and recall that λ = dσ 2 /N , we have w λ ridge = ( X ⊤ X + dσ 2 I d ) -1 X ⊤ y .

Therefore, we only need to prove the following fact: for any γ > 0 and ̂ β = ( X ⊤ X + dγ I d ) -1 X ⊤ y , we have

We now prove (34). Note that we have

where B 1 = X ⊤ X ( X ⊤ X + dγ I d ) -1 , B 2 = ( X ⊤ X + dγ I d ) -1 X ⊤ . Note that ∥ B 1 ∥ op ≤ 1 clearly holds, and under E π we also have ∥ ε ∥ 2 ≤ 2 √ Nσ . Therefore, it remains to bound the term ∥ B 2 ∥ op .

Consider the SVD decomposition of X = U Σ V , Σ = diag( λ 1 , · · · , λ d ) , and U ∈ R N × d , V ∈ R d × d are orthonormal matrices. Then B 2 = V ⊤ (Σ 2 + dγ I d ) -1 Σ U ⊤ , and hence

When N ≤ 36 d , we directly have ∥ B 2 ∥ op ≤ 1 2 ( dγ ) -1 / 2 ≤ 3( Nγ ) -1 / 2 . Otherwise, we have N ≥ 36 d , and then for each i ∈ [ d ] , λ i ≥ √ λ min ( X ⊤ X ) ≥ √ αN ≥ √ N/ 3 . Hence, in this case we also have ∥ B 2 ∥ op ≤ max i λ -1 i ≤ 3 N -1 / 2 . Combining the both cases completes the proof of (34).

G In-context learning of generalized linear models

As a natural generalization of linear regression, we now show that transformers can recover learn generalized linear models (GLMs) [53] (which includes logistic regression for linear classification as an important special case), by implementing the corresponding convex risk minimization algorithm in context, and achieve near-optimal excess risk under standard statistical assumptions.

Let g : R → R be a link function that is non-decreasing and C 2 -smooth. We consider the following convex empirical risk minimization (ERM) problem

where ℓ ( t, y ) := -yt + ∫ t 0 g ( s ) ds is the convex (integral) loss associated with g . Acanonical example of (ICGLM) is logistic regression, in which g ( t ) = σ log ( t ) := (1 + e -t ) -1 is the sigmoid function, and the resulting ℓ ( t, y ) = ℓ log ( t, y ) = -yt +log(1 + e t ) is the logistic loss.

The following result (proof in Appendix G.1) shows that, as long as the empirical risk ̂ L N satisfies strong convexity and bounded solution conditions (similar as in Theorem 4), transformers can approximately implement the ERM predictor g ( ⟨ x N +1 , w GLM ⟩ ) , with w GLM given by (ICGLM).

Theorem G.1 (Implementing convex risk minimization for GLMs) . For any 0 < α < β with κ := β α , B w > 0 , B x > 0 , κ w := L g B 2 x /α +1 and ε < B w / 2 , there exists an attention-only transformer TF 0 θ with

(where L g := sup t | g ′ ( t ) | , R := max { B x B w , B y , 1 } , and C g > 0 is a constant that depends only on R and the C 2 -smoothness of g within [ -R,R ] ), such that the following holds. On any input data ( D , x N +1 ) such that

In Theorem G.1, the number of heads scales as ˜ O (1 /ε 2 ) as opposed to Θ(1) as in ridge regression (Theorem 4), due to the fact that the gradient of the loss is in general a smooth function that can be only approximately expressed as a sum-of-relus (cf. Definition D.1 & Lemma B.5) rather than exactly expressed as in the case for the square loss.

In-context prediction power We next show that (proof in Appendix G.2) the transformer constructed in Theorem G.1 achieves desirable statistical power if the in-context data distribution satisfies standard statistical assumptions for learning GLMs. Let L P ( w ) := E ( x ,y ) ∼ P [ ℓ ( ⟨ w , x ⟩ , y )] denote the corresponding population risk for any distribution P of ( x , y ) . When P is realizable by a generalized linear model of link function g and parameter β in the sense that E P [ y | x ] = g ( ⟨ β , x ⟩ ) , it is a standard result that β is indeed a minimizer of L P [42] (see also [6, Appendix A.3]).

Theorem G.2 (Statistical guarantee for generalized linear models) . For any fixed set of parameters defined in Assumption B, there exists a transformer θ with L ≤ O (log( N )) layers and max ℓ ∈ [ L ] M ( ℓ ) ≤ ˜ O ( d 3 N ) , such that for any distribution P satisfying Assumption B with those parameters, as long as N ≥ O ( d ) , that outputs ̂ y N +1 = ˜ read y (TF θ ( H )) and ̂ w = ˜ read w (TF θ ( H )) ∈ R d (for another read-out function ˜ read w ) satisfying the following.

(a) ̂ w achieves small excess risk under the population loss, i.e. for the linear prediction ̂ y lin N +1 := ⟨ x N +1 , ̂ w ⟩ ,

(b) (Realizable setting) If there exists a β ∈ R d such that under P , E [ y | x ] = g ( ⟨ β , x ⟩ ) almost surely, then

Above, O ( · ) hides constants that depend polynomially on the parameters in Assumption B. Similar as in Corollary 5, the O ( d/N ) excess risk obtained here matches the optimal (fast) rate for typical learning problems with d parameters and N samples [87].

Assumption B (Well-posedness for learning GLMs) . We assume that there is some B µ > 0 such that for any t ∈ [ -B µ , B µ ] , g ′ ( t ) ≥ µ g > 0 .

We also assume that for each i ∈ [ N +1] , ( x i , y i ) is independently sampled from P such that the following holds.

(a) Under the law ( x , y ) ∼ P , We have x ∼ SG( K x ) , y ∼ SG( K y ) and g ( ⟨ w , x ⟩ ) ∼ SG( K y ) ∀ w ∈ B 2 ( B w ) .

(b) For some µ x > 0 , it holds that

(c) For β ⋆ = arg min L P , it holds ∥ β ⋆ ∥ 2 ≤ B w / 4 .

Applying Theorem G.2 to logistic regression, we have the following result as a direct corollary. Below, the Gaussian input assumption is for convenience only and can be generalized to e.g. sub-Gaussian input.

Corollary G.1 (In-context logistic regression) . Consider any in-context data distribution P satisfying

For the link function g = σ log and B ⋆ w = O (1) , we can choose B w , B µ , µ g , L g , µ x , K x , K y = Θ(1) so that Assumption B holds. In that case, when N ≥ O ( d ) , there exists a transformer θ with L = O (log( N )) layers, such that for any P considered above,

(a) The estimation ̂ w = ˜ read w (TF θ ( H )) outputted by θ achieves excess risk bound (36).

(b) (Realizable setting) Consider the logistic in-context data distribution

Then, for any distribution P = P log β with ∥ β ∥ 2 ≤ B ⋆ w , the prediction ̂ y N +1 = ˜ read y (TF θ ( H )) of θ additionally achieves the square loss excess risk (37).

G.1 Proof of Theorem G.1

Let us fix parameters ε g > 0 and T > 0 (that we specify later in proof).

Define R = max { B x B w , B y , 1 } and

By Proposition B.1, g is ( ε g , M, R, C ) with

Therefore, we can invoke Theorem D.1 to obtain that, as long as 2 Tε g ≤ B w , there exists a T -layer attention-only transformer θ (1: T ) with M heads per layer, such that for any input H of format (3) and satisfies (35), its last layer outputs h ( T ) i = [ x i ; y ′ i ; ̂ w T ; 0 D -2 d -3 ; 1; t i ] , such that

where { w ℓ GD } ℓ ∈ [ L ] is the sequence of gradient descent iterates with stepsize β -1 and initialization w 0 GD = 0 . Notice that Proposition B.2 implies (with κ := β/α )

Furthermore, we can show that (similar to the proof of Theorem D.1 (b)), there exists a single attention layer θ ( T +1) with M heads such that it outputs h ( T +1) N +1 = [ x N +1 ; ̂ y N +1 ; ̂ w T ; 0 D -2 d -3 ; 1; 0] , where ∣ ∣ ̂ y N +1 -g ( 〈 x N +1 , ̂ w T 〉 ) ∣ ∣ ≤ ε g .

In the following, we show that for suitably chosen ( T, ε g ) , θ = ( θ (1: T ) , θ ( T +1) ) is the desired transformer. First notice that its output h ( T +1) N +1 = [ x N +1 ; ̂ y N +1 ; ̂ w T ; 0 D -2 d -3 ; 1; 0] satisfies

Therefore, for any fixed ε > 0 , we can take

so that the θ we construct above ensures | ̂ y N +1 -g ( ⟨ x N +1 , w GLM ⟩ ) | ≤ ε for any input H that satisfies (35). The upper bound on | | | θ | | | follows immediately from Theorem D.1.

G.2 Proof of Theorem G.2

Wesummarize some basic and useful facts about GLM in the following theorem. Its proof is presented in Appendix G.3 - G.6.

Theorem G.3. Under Assumption B, the following statements hold with universal constant C 0 and constant C 1 , C 2 that depend only on the parameters ( K x , K y , B µ , B w , µ x , L g , µ g ) .

(a) As long as N ≥ C 1 · d , the following event happens with probability at least 1 -2 e -N/C 1 :

(b) For any δ > 0 , we have with probability at least 1 -δ that

where we denote ι = log(2 + L g K 2 x B w /K y ) .

(c) Condition on (a) holds and N ≥ C 2 · d , the event E r := {∥ w GLM ∥ 2 ≤ B w / 2 } happens with probability at least 1 -e N/C 2 .

(d) For any w ∈ B 2 ( B w ) , it holds that

(e) (Realizable setting) As long as w GLM ∈ B 2 ( B w ) , it holds that

Therefore, we can set

Consider the following good events

Under the event E and our choice of α, β , the problem (ICGLM) is well-conditioned (i.e. (35) holds). Theorem G.1 implies that there exists a transformer θ such that for any input H of the form (3), TF θ outputs h ′ N +1 = [ x N +1 ; ˜ y N +1 ; ˜ w ; 0 D -2 d -3 ; 1; 0] , such that the output is given by ̂ y N +1 = ˜ read y (TF θ ( H )) = clip B y ( ˜ y N +1 ) and ̂ w = ˜ read w (TF θ ( H )) := Proj B 2 ( B w ) ( ˜ w ) , and the following holds on the good event E :

(a) ˜ y N +1 = f D ( x N +1 ) , where f D = A ( D ) is a predictor such that | f D ( x ) -g ( ⟨ x , w GLM ⟩ ) | ≤ ε for all x ∈ B 2 ( B x ) .

(b) ˜ w = ˜ w ( D ) ∈ B 2 ( B w ) depends only on D (by the proof of Theorem G.1 and Theorem D.1), such that ∥ ∥ ∇ ̂ L N ( ˜ w ) ∥ ∥ 2 ≤ βε L g B w .

In the following, we show that θ constructed above fulfills both (a) & (b) of Theorem G.2. The bounds on number of layers and heads and | | | θ | | | follows from plugging our choice of B x , B y in our proof of Theorem G.1.

Proof of Theorem G.2 (a). Notice that under the good event E , we have ̂ w = ˜ w = ˜ w ( D ) depends only on D . Then we have

Thus, we can consider E 0 = E r ∩ E w ∩ E b , and then

where the second equality follows from L p ( ̂ w ( D )) = E ( x N +1 ,y N +1 ) |D ℓ ( ⟨ x N +1 , ˜ w ( D ) ⟩ , y N +1 ) . Therefore,

where the last line follows from Cauchy inequality and the fact P ( E c ) = O ( N -10 ) , and B ℓ is defined in Lemma G.1.

Notice that by Theorem G.3 (d), we have

and by Theorem G.3 (b) and taking integration over δ > 0 , we have

Also, we have inf L p = L p ( β ⋆ ) ≤ B ℓ by Lemma G.1. Therefore, we can conclude that

Taking ε 2 ≤ K 2 y ι B w K 2 x d N

Proof of Theorem G.2 (b). Similar to the proof of Corollary 6, we have

where the inequality follows from y N +1 ∈ [ -B y , B y ] on event E . For the first part, we have

where we use the fact that the conditional distribution of ( x N +1 , y N +1 ) |D agrees with P . Thus,

For the second part, we know P ( E c ) = O ( N -10 ) and

In conclusion, we have

E

(

2

̂

y

N

+1

-

y

N

+1

)

E

(

y

N

+1

-

g

(

⟨

x

N

+1

,

Taking ε 2 ≤ L g K 2 x K 2 y ι µ x µ g d N completes the proof.

β

⟩

))

2

+2

ε

2

+

O

(1)

·

L

g

K

µ

2

x

x

K

µ

g

2

y

ι

d

N

+

O

(

B

N

2

y

5

)

.

Lemma G.1. Suppose that x ∼ SG( K x ) , y ∼ SG( K y ) , and w is a (possibly random) vector such that ∥ w ∥ 2 ≤ B w . Then

Proof. Notice that by our assumption, | g (0) | ≤ 2 K y . Therefore, by the definition of ℓ ,

The proof is then done by bounding the moment by E | y | 8 ≤ O ( K 8 y ) and E |⟨ x , w ⟩| 8 ≤ B 8 w E ∥ x ∥ 8 2 ≤ O ( ( √ dB w K x ) 8 ) , which is standard (by utilizing the tail bound of sub-Gaussian/subExponential random variable).

G.3 Proof of Theorem G.3 (a)

We begin with the upper bound on λ max ( ∇ 2 ̂ L N ( w )) . By Lemma B.3, as long as N ≥ C 0 · d , the following event

happens with probability at least 1 -exp( -N/C 0 ) . By the assumption that sup | g ′ | ≤ L g , it is clear that when E w, 0 holds, we have λ max ( ∇ 2 ̂ L N ( w )) ≤ 8 L g K 2 x ∀ w ∈ R d .

In the following, we analyze the quantity λ max ( ∇ 2 ̂ L N ( w )) . We have to invoke the following covering argument (see e.g. [85, Section 4.1.1]).

Lemma G.2. Suppose that V is a ε -covering of S d -1 with ε ∈ [0 , 1) . Then the following holds:

2. For any vector x ∈ R d , ∥ x ∥ 2 ≤ 1 1 -ε max v ∈V |⟨ v , x ⟩| .

Notice that

Therefore, we can define h ( t ) := ( B µ -| t | ) + (which is a 1 -Lipschitz function), and we have

In the following, we pick a ε v -covering V of S d -1 such that |V| ≤ (3 /ε v ) d (we will specify ε v later in proof). Then for any w ∈ B 2 ( B w ) ,

By our definition of A ( w ) , we have (for any fixed B xv )

v

∥

w

sup

∥

2

B

w

By Lemma G.3, we can choose B xv = K x (15 + log( K 2 x /µ x )) , and then E [ U v ( w )] ≥ 3 B µ µ x / 8 . Thus, combining the inequalities above, we can take ε v = 128 K 2 x µ x in the following, so that under event E w, 0 ,

In the following, we consider the random process { U v ( w ) := U v ( w ) -E [ U v ( w )] } w , which is zero-mean and indexed by w ∈ B 2 ( B w ) . For any fixed v , consider applying Proposition B.4 to the random process { U v ( w ) } w . We need to verify the preconditions:

(a) With norm ρ ( w , w ′ ) = ∥ w -w ′ ∥ 2 , log N ( B ρ ( w , r ) , δ ) ≤ d log(2 Ar/δ ) with constant A = 2 ;

(b) Let f ( x ; w ) := h ( ⟨ w , x i ⟩ ) min { ⟨ v , x i ⟩ 2 , B 2 xv } , then | f ( x ; w ) | ≤ B µ B 2 xv and hence in SG( CB µ B 2 xv ) for any random x ;

(c) For w , w ′ ∈ W , we have | h ( ⟨ w , x i ⟩ ) -h ( ⟨ w ′ , x i ⟩ ) | ≤ |⟨ w -w ′ , x i ⟩| . Hence, because x ∼ SG( K x ) , the random variable h ( ⟨ w , x ⟩ ) -h ( ⟨ w ′ , x ⟩ ) is SG( CK x ∥ w -w ′ ∥ 2 ) , and the random variable f ( x ; w ) -f ( x ; w ′ ) is SG( CK x B 2 xv ∥ w -w ′ ∥ 2 ) .

Therefore, we can apply Proposition B.4 to obtain that with probability 1 -δ 0 , it holds

where we denote κ g = 1 + K x B w /B µ . Setting δ 0 = δ/ |V| and taking the union bound over v ∈ V , we obtain that with probability at least 1 -δ ,

[√

d

log(8

κ

/ε

v

) + log(1

N

where we use log |V| ≤ d log(4 /ε v ) . Therefore, we plug in the definition of ε v and B xv to deduce that, if we set

then as long as N ≥ C 1 · d , it holds max v ∈V E [ U v ( w )] -U v ( w ) ≤ µ x B µ 16 with probability at least 1 -exp( -N/C 1 ) . This is the desired result.

Lemma G.3. Under Assumption B, for B xv = K x (15 + log( K 2 x /µ x )) , it holds

Proof. For any fixed w ∈ B ( B ) , v ∈ S d -1 ,

Because x ∼ SG( K x ) , x ⊤ v ∼ SG( K x ) , and a simple calculation yields

Taking t = 15 + log( K 2 x /µ x ) gives E [( x ⊤ v ) 2 1 {| x ⊤ v | > B xv } ] ≤ µ x / 4 , which completes the proof.

/δ

)

]

∈V

∣

∣

U

v

(

w

)

∣

∣

C

′

B

µ

B

2

g

,

G.4 Proof of Theorem G.3 (b)

Notice that

In the following, we pick a minimal 1 / 2 -covering of S d -1 (so |V| ≤ 5 d ). Then by Lemma G.2, it holds

Fix a v ∈ S d -1 and set δ ′ = δ/ |V| . We proceed to bound sup w | X v ( w ) | by applying Proposition B.4 to the random process { X v ( w ) } w . We need to verify the preconditions:

(b) For z = [ x ; y ] , we let f ( z ; w ) := ( g ( ⟨ w , x ⟩ ) -y ) ⟨ x , v ⟩ , then f ( z ; w ) ∼ SE( CK x K y ) for any w by our assumption on ( x , y ) ;

(c) For w , w ′ ∈ W , we have | g ( ⟨ w , x ⟩ ) -g ( ⟨ w ′ , x ⟩ ) | ≤ L g |⟨ w -w ′ , x ⟩| . Hence, because x ∼ SG( K x ) , the random variable g ( ⟨ w , x i ⟩ ) -g ( ⟨ w ′ , x i ⟩ ) is sub-Gaussian in SG( K x L g ∥ w -w ′ ∥ 2 ) . Thus, f ( z ; w ) -f ( z ; w ′ ) is sub-exponential in SE( CK 2 x L g ∥ w -w ′ ∥ 2 ) .

Therefore, we can apply Proposition B.4 to obtain that with probability 1 -δ 0 , it holds

where we denote κ y = 1 + L g K 2 x B w /K y . Setting δ 0 = δ/ |V| and taking the union bound over v ∈ V , we obtain that with probability at least 1 -δ ,

This is the desired result.

G.5 Proof of Theorem G.3 (c)

In the following, we condition on (a) holds, i.e. ̂ L N is α -strongly-convex and β smooth over B 2 ( B w ) with α = µ x µ g / 8 and β = 8 L g K 2 x . We define

Then by standard convex analysis, we have

α ∥ ˜ w -β ⋆ ∥ 2 2 ≤ 〈 ∇ ̂ L N ( ˜ w ) -∇ ̂ L N ( β ⋆ ) , ˜ w -β ⋆ 〉 ≤ 〈 -∇ ̂ L N ( β ⋆ ) , ˜ w -β ⋆ 〉 ≤ ∥ ∥ ∇ ̂ L N ( β ⋆ ) ∥ ∥ 2 ∥ ˜ w -β ⋆ ∥ 2 .

Notice that ∥ ∥ ∇ ̂ L N ( β ⋆ ) ∥ ∥ 2 ≤ ε stat , we can conclude that

Recall that we assume ∥ β ⋆ ∥ 2 ≤ B w / 4 , we can then consider E s := { ε stat < αB w / 4 } . Once E s holds, our argument above yields ∥ ˜ w ∥ 2 < B w , which implies ∇ ̂ L N ( ˜ w ) = 0 . Therefore, ˜ w = arg min w ∈ R d ̂ L N ( w ) . Further, by Theorem G.3, we can set

so that as long as N ≥ C 2 d , the event E s holds with probability at least 1 -exp( -N/C 2 ) . This is the desired result.

G.6 Proof of Theorem G.3 (d) & (e)

We first prove Theorem G.3 (d). Notice that

Therefore, L p is ( µ g µ x ) -strongly-convex over B 2 ( B w ) . Therefore, because β ⋆ ∈ B 2 ( B w ) is the global minimum of L p , it holds that for all w ∈ B 2 ( B w ) ,

By the definition of ε stat , ∥∇ L p ( w ) ∥ 2 ≤ ε stat + ∥∇ ̂ L N ( w ) ∥ 2 , and hence the proof of Theorem G.3 (d) is completed.

We next prove Theorem G.3 (e), where we assume that E [ y | x ] = g ( ⟨ x , β ⟩ ) (which implies β ⋆ = β directly) and w GLM ∈ B 2 ( B w ) . Notice that

and hence

On the other hand, by the ( µ g µ x ) -strong-convexity of L p over B 2 ( B w ) , it holds that

Finally, using the definition of w GLM , we have ∇ ̂ L N ( w GLM ) = 0 , and hence ∥∇ L p ( w GLM ) ∥ 2 ≤ ε stat , which completes the proof of Theorem G.3 (e).

H Proofs for Section 3.2

H.1 Proof of Theorem 7

Fix λ N ≥ 0 , β > 0 and B w > 0 , and consider any in-context data D such that the precondition of Theorem 7 holds. Recall that

denotes the lasso regression loss in (ICLasso), so that w lasso = arg min w ∈ R d L lasso ( w ) . We further write

Note that ∇ 2 ̂ L 0 N ( w ) = X ⊤ X /N and thus ̂ L 0 N is β -smooth over R d .

Consider the proximal gradient descent algorithm on the ridge loss

with initialization w 0 PGD := 0 d , learning rate η := β -1 , and number of steps T to be specified later. Similar to the proof of Theorem 4, we can construct a transformer to approximate w T GD . Consider ℓ ( s, t ) = 1 2 ( s -t ) 2 and R ( w ) = λ N ∥ w ∥ 1 , then ∂ s ℓ ( s, t ) is (0 , + ∞ , 2 , 4) -approximable by sum of relus (cf. Definition D.1), and prox η R is (0 , + ∞ , 4 d, 4 + 2 ηλ N ) -approximable by sum of relus (Proposition D.1). Therefore, we can apply Theorem D.2 with the square loss ℓ , regularizer R , learning rate η and accuracy parameter 0 to obtain that there exists a transformer TF θ with ( T +1)

layers, number of heads M ( ℓ ) = 2 for all ℓ ∈ [ L ] , and hidden dimension D ′ = 2 d , such that the final output h ( L ) N +1 = [ x N +1 ; ̂ y N +1 ; w T PGD ; ∗ ] with ̂ y N +1 = 〈 w T PGD , x N +1 〉 . Further, the weight matrices have norm bounds | | | θ | | | ≤ 10 R +(8 + 2 λ N ) β -1 .

By the standard convergence result for proximal gradient descent (Proposition B.3), we have for all t ≥ 1 that

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

Collection of parameters Throughout the rest of this section, we consider fixed N ≥ 1 , λ N = √ ρν log d N for ρ ≥ 0 , ν ≥ 0 fixed (and to be determined), fixed 0 < α ≤ β , and fixed B ⋆ w > 0 . We write κ := β/α,κ s := β ( B ⋆ ) 2 /ν 2 , and ω N := ρ s log d .

Here we present a sharper convergence analysis on the proximal gradient descent algorithm for L lasso under the following well-conditionedness assumption, which will be useful for proving Theorem 8 in the sequel.

Assumption C (Well-conditioned property for Lasso) . We say the (ICLasso) problem is wellconditioned with sparsity s if the following conditions hold:

The ( α, ρ ) -RSC condition holds:

Further, λ max ( X ⊤ X /N ) ≤ β .

The data ( X , y ) is 'approximately generated from a s -sparse linear model': There exists a w ⋆ ∈ R d such that ∥ w ⋆ ∥ 2 ≤ B ⋆ w , ∥ w ⋆ ∥ 0 ≤ s and for the residue ε = y -Xw ⋆ ,

It holds that N ≥ 32 ρ α · s log d (i.e. 32 ω N ≤ 1 ).

Assumption C1 imposes the standard restricted strong convexity (RSC) condition for the feature matrix X ∈ R N × d , and Assumption C2 asserts that the data is approximately generated from a sparse linear model, with a bound on the L ∞ norm of the error vector X ⊤ ε . Assumption C is entirely deterministic in nature, and suffices to imply the following convergence result. In the proof of Theorem 8, we show that Assumption C is satisfied with high probability when data is generated from the standard sparse linear model considered therein.

Theorem H.1 (Sharper convergence guarantee for Lasso) . Under Assumption C, for the PGD iterates { w t } t ≥ 0 on loss function ̂ L lasso with stepsize η = 1 /β and starting point w 0 = 0 , we have ̂ L lasso ( w T ) -̂ L lasso ( w lasso ) ≤ ε for all

where C is a universal constant.

The proof can be found in Appendix H.4. Combining Theorem H.1 with the construction in Theorem 7, we directly obtain the following result as a corollary.

Theorem H.2 (In-context Lasso with transformers with sharper convergence) . For any N,d,s ≥ 1 , 0 < α ≤ β , ν ≥ 0 , ρ ≥ 0 , there exists a L -layer transformer TF θ with

such that the following holds. On any input data ( D , x N +1 ) such that the (ICLasso) problem satisfies Assumption C (which implies ∥ w lasso ∥ 2 ≤ B w / 2 with B w = 2 B ⋆ w + √ ν/α ), TF θ ( H (0) ) approximately implements (ICLasso), in that it outputs ̂ y N +1 = read y (TF θ ( H )) = ⟨ x N +1 , ̂ w ⟩ with

H.3 Basic properties for Lasso

Lemma H.1 (Relaxed basic inequality) . Suppose that Assumption C2 holds. Then it holds that

As a corollary, ∥ w lasso -w ⋆ ∥ 1 ≤ 4 √ s ∥ w lasso -w ⋆ ∥ 2 .

Proof. Let us first fix any w ∈ R d . Denote ∆ = w -w ⋆ , and let S = supp( w ⋆ ) be the set of indexes of nonzero entries of w ⋆ . Then by definition, y = Xw ⋆ + ε and | S | ≤ s , and hence

Combining these inequalities, we obtain

where the second inequality follows from ε ⊤ X∆ N ≤ ∥ X ⊤ ε ∥ ∞ N ∥ ∆ ∥ 1 and our assumption that 2 ∥ X ⊤ ε ∥ ∞ N ≤ λ N , and the last inequality is due to ∥ ∆ ∥ 1 = ∥ ∆ S ∥ 1 + ∥ ∆ S c ∥ 1 . Therefore, we have

where the last inequality follows from ∥ ∆ S ∥ 1 ≤ √ s ∥ ∆ S ∥ 2 ≤ √ s ∥ ∆ ∥ 2 . This completes the proof of our main inequality. As for the corollary, we only need to use the definition that ̂ L lasso ( w lasso ) ≤ ̂ L lasso ( w ⋆ ) .

Proposition H.1 (Gap to parameter estimation error) . Suppose that Assumption C holds. Then for all w ∈ R d ,

where we write gap := ̂ L lasso ( w ) -̂ L lasso ( w lasso ) , and C = 120 is a universal constant. In particular, we have ∥ w lasso -w ⋆ ∥ 2 2 ≤ 10 ρν α 2 s log d N .

Proof. We follow the notation in the proof of Lemma H.1. By (39), we have

and hence ∥ ∆ ∥ 1 ≤ 4 √ s ∥ ∆ ∥ 2 + 2 gap λ N due to ̂ L lasso ( w ) -̂ L lasso ( w ⋆ ) ≤ gap . On the other hand, by the RSC condition (38), it holds that

Therefore, we have

where the last inequality uses AM-GM inequality and Cauchy inequality. Notice that ρ 20 s log d N ≤ 2 3 α , we now derive that

Plugging in λ N = √ ρν log d N completes the proof. The corollary follows immediately by letting w = w lasso in above proof (hence gap = 0 ).

Lemma H.2 (Growth) . It holds that

Proof. For simplicity we denote w lasso := w lasso . By the first order optimality condition, it holds that

where we write R ( w ) := λ N ∥ w ∥ 1 . Then by the convexity of R , we have

Rearranging completes the proof.

H.4 Proof of Theorem H.1

For the simplicity of presentation, we write w lasso = w lasso and we denote gap t := ̂ L lasso ( w t ) -̂ L lasso ( w lasso ) .

By Lemma H.1, we have ∥ w t -w ⋆ ∥ 1 ≤ 4 √ s ∥ w t -w ⋆ ∥ 2 + 2 gap t λ N , which implies

We denote µ N = ρ 2 log d N . Using the assumption that X is ( α, ρ ) -RSC, we obtain that

Thus, as long as N ≥ 30 ρ 2 s log d α , we have

where the last inequality follows from Lemma H.2 and the definition of λ N , µ N .

We define ε stat := 640 sµ N ∥ w lasso -w ⋆ ∥ 2 2 , T 0 := 10 βν -1 ∥ w lasso ∥ 2 2 . By Proposition B.3(3), it holds that for t ≥ T 0 ,

Then for all t ≥ T 0 -1 , we have (the second ≤ below uses Proposition B.3(2))

Therefore, for t ≥ T 0 -1 ,

where the last inequality follows from Proposition B.3(2). Further, by Proposition B.3(3), we have

Hence, we can conclude that gap T ≤ ε for all T such that

Now, by Proposition H.1, it holds that ∥ w lasso -w ⋆ ∥ 2 2 ≤ 10 ρν α 2 s log d N , and hence

Plugging in our definition of

completes the proof.

H.5 Proof of Theorem 8

In this section, we present the proof of Theorem 8 based on Theorem H.2. We begin by recalling the following RSC property of a Gaussian random matrix [87, Theorem 7.16], a classical result in the high-dimensional statistics literature.

Proposition H.2 (RSC for Gaussian random design) . Suppose that X = [ x 1 ; · · · ; x N ] ⊤ ∈ R N × d is a random matrix with each row x i being i.i.d. samples from N (0 , Σ ) . Then there are universal constants c 1 = 1 8 , c 2 = 50 such that with probability at least 1 -e -N/ 32 1 -e -N/ 32 ,

where ρ ( Σ ) = max i ∈ [ d ] Σ ii is the maximum of diagonal entries of Σ .

Fix a parameter δ 1 ≤ δ (which we will specify in proof) and a large universal constant C 0 . Let us set

Similar to the proof of Corollary 6 (Appendix F.4), we consider the following good events (where ε = Xw ⋆ -y )

and we define E := E w ∩ E r ∩ E b ∩ E b,N +1 .

Furthermore, we choose ν > 0 that correspond to the choice λ N = 8 σ √ log(4 d/δ ) N , and we also assume N ≥ 32 c 2 c 1 · s log d . Then, Assumption C holds on the event E .

Therefore, we can apply Theorem H.2 with ε = νω N , which implies that there exists a L -layer transformer θ such that its prediction ̂ y N +1 := ˜ read y (TF 0 θ ( H )) , so that under the good event E we have ̂ y N +1 = clip B y ( ⟨ x N +1 , ̂ w ⟩ ) , where

In the following, we show that θ is indeed the desired transformer (similarly to the proof in Appendix F.4). Consider the conditional prediction error

Part I. We first note that

where the inequality is because y N +1 ∈ [ -B y , B y ] under the good event E . Notice that by our construction, under the good event E , ̂ w = ̂ w ( D ) depends only on the dataset D (because it is the ( L -1) -th iterate of PGD on (ICLasso) problem). Applying Proposition H.1 to ̂ w ( D ) and using the definition of ω N and our choice of λ N , we obtain that (under E 0 )

Therefore, under E 0 ,

Part II. Notice that under good event E 0 , the bad event E c holds if and only if E c b,N +1 holds, and hence

With a large enough constant C 0 , we clearly have P ( E c b,N +1 ) ≤ ( δ 1 /N ) 10 . Further, a simple calculation yields

where the last inequality is because the marginal distribution of y N +1 is simply N (0 , σ 2 + ∥ w ⋆ ∥ 2 2 ) . Combining these yields

Therefore, choosing δ 1 = min { δ, σ B ⋆ w } is enough for our purpose, and under such choice of δ 1 ,

Conclusion. Combining the inequalities above, we can conclude that under E 0 ,

It remains to show that P ( E 0 ) ≥ 1 -δ . By Proposition H.2, Lemma B.2 and Lemma B.4, we have

Therefore, as long as N ≥ 32 log(12 /δ ) , we have P ( E 0 ) ≥ 1 -δ . This completes the proof.

We also remark that in the construction above,

which would be useful for bounding | | | θ | | | .

I Proofs for Section 4

I.1 Proof of Proposition 10

We begin by restating Proposition 10 into the following version, which contains additional size bounds on θ .

Theorem I.1 (Full statement of Proposition 10) . Suppose that for

ℓ ( · , · ) is ( γ/ 3 , R, M, C ) -approximable by sum of relus (Definition D.1). Then there exists a 3-layer transformer TF θ with

that maps

where the predictor ̂ f : R d → R is a convex combination of { f k : ̂ L val ( f k ) ≤ min k ⋆ ∈ [ K ] ̂ L val ( f k ⋆ ) + γ } . As a corollary, for any convex risk L : ( R d → R ) → R , ̂ f satisfies

To prove Theorem I.1, we first state and prove the following two propositions.

Proposition I.1 (Evaluation layer) . There exists a 1-layer transformer TF θ with MK heads and | | | θ | | | ≤ 3 R +2 NKC/ |D val | such that for all H such that max i {| y ′ i |} ≤ R, max i,k {| f k ( x i ) |} ≤ R , TF θ maps

where ˜ L val ( · ) is a functional such that max k ∣ ∣ ∣ ˜ L val ( f k ) -̂ L val ( f k ) ∣ ∣ ∣ ≤ ε

Proof of Proposition I.1. As ℓ is ( ε, R, M, C ) -approximable by sum of relus, there exists a function g : R 2 → R of form

such that sup ( s,t ) ∈ [ -R,R ] 2 | g ( s, t ) -ℓ ( s, t ) | ≤ ε . We define

Next, for every m ∈ [ M ] and k ∈ [ K ] , we define matrices Q m,k , K m,k , V m,k ∈ R D × D such that for all i, j ∈ [ N +1] ,

where e s ∈ R D is the vector with s -th entry being 1 and others being 0 . As the input has structure h i = [ x i ; y ′ i ; ∗ ; f 1 ( x i ); · · · ; f K ( x i ); 0 K +1 ; 1; t i ] , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Now, for every i, j ∈ [ N +1] , we have

where the last equality follows from the bound | a m f k ( x j ) + b m y j + d m | ≤ R ( | a m | + | b m | ) + d m ≤ 2 R , so that the above relu equals 0 if t j ≤ 0 . Therefore, for each i ∈ [ N +1] and k ∈ [ K ] ,

Thus letting the attention layer θ = { ( V m,k , Q m,k , K m,k ) } ( m,k ) ∈ [ M ] × [ K ] , we have

This is the desired result.

Proposition I.2 (Selection layer) . There exists a 3-layer transformer TF θ with

such that TF θ maps

where ̂ f = ∑ K k =1 λ k f k is an aggregated predictor, where the weights λ 1 , · · · , λ K ≥ 0 are functions only on L 1 , · · · , L k such that

Proof of Proposition I.2. We construct a θ which is a composition of 2 MLP layers followed by an attention layer ( θ (1) mlp , θ (2) mlp , θ (3) attn ) .

Step 1: construction of θ (1) mlp . We consider matrix W (1) 1 that maps

i.e. W (1) 1 h is a K 2 + K dimensional vector so that its entry contains { L k -L l } k,l ∈ [ K ] and { L k , -L k } k ∈ [ K ] . Clearly, such W (1) 1 exists and can be chosen so that ∥ ∥ ∥ W (1) 1 ∥ ∥ ∥ op ≤ 2 K. We then consider a matrix W (1) that maps

̸

2

σ ( W (1) 1 h ) ↦→ W (1) 2 σ ( W (1) 1 h ) = [ 0 D -K -3 ; c 1 -L 1 ; · · · ; c K -L K ; 0 3 ] ∈ R D , where c k = c k ( L ) := ∑ l = k σ ( L k -L l ) . Notice that

̸

and hence such W (1) 2 exists and can be chosen so that ∥ ∥ ∥ W (1) 2 ∥ ∥ ∥ op ≤ K + 1 . We set θ (1) mlp = ( W (1) 1 , W (1) 2 ) , then MLP θ (1) mlp maps h i to

The basic property of { c k } k ∈ [ K ] is that, if c k ≤ γ , then L k ≤ min k ⋆ ∈ [ K ] L k ⋆ + γ .

Step 2: construction of θ (2) mlp . We consider matrix W (2) 1 that maps

and W (2) 1 can be chosen so that ∥ ∥ ∥ W (2) 1 ∥ ∥ ∥ op ≤ K +1+ γ -1 . We then consider a matrix W (2) 2 that maps

which exists and can be chosen so that ∥ ∥ ∥ W (1) 2 ∥ ∥ ∥ op ≤ 2 . We set θ (2) mlp = ( W (2) 1 , W (2) 2 ) , then MLP θ (2) mlp maps h (1) to

where u k = σ (1 -γ -1 c k ) ∀ k ∈ [ K ] . Clearly, u k ∈ [0 , 1] , and u k > 0 if and only if c k ≤ γ .

Step 3: construction of θ (3) attn . We define

Clearly, λ k ≥ 0 , and ∑ k λ k = 1 . Further,

Therefore, it remains to construct θ (3) attn that implements ̂ f = ∑ K k =1 λ k f k based on [ h (2) i ] i . Notice that

and hence we construct θ (3) attn as follows: for every k ∈ [ K +1] and w ∈ { 0 , 1 } , we define matrices Q k,w , K k,w , V k,w ∈ R D × D such that for all k ∈ [ K +1]

for all i, j ∈ [ N + 1] , where we understand f 0 = f K +1 = 0 and 1 k is the k -dimensional vector with all entries being 1. By the structure of h (2) i , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Now, for every i, j ∈ [ N +1] , k ∈ [ K +1] , w ∈ { 0 , 1 } , we have

where the last equality follows from f k ( x i ) + R ≥ 0 ∀ k ∈ [ K ] . Therefore,

i

where the last equality is due to (41). Thus letting the attention layer θ (3) attn = { ( V k,w , Q k,w , K k,w ) } ( k,w ) ∈ [ K +1] ×{ 0 , 1 } , we have

This is the desired result.

Now, we are ready to prove Theorem I.1.

Proof of Theorem I.1 As ℓ ( · , · ) is ( γ/ 3 , R, M, C ) -approximable by sum of relus, we can invoke Proposition I.1 to show that there exists a single attention layer θ (1) attn so that Attn θ (1) attn maps

for any input H = [ h i ] i of the form described in Theorem I.1, and ˜ L val ( · ) is a functional such that max k ∣ ∣ ∣ ˜ L val ( f k ) -̂ L val ( f k ) ∣ ∣ ∣ ≤ γ/ 3 .

Next, by the proof of Proposition I.2, there exists ( θ (1) mlp , θ (2) mlp , θ (3) attn ) that maps

where λ = ( λ 1 , · · · , λ K ) ∈ ∆([ K ]) and λ k > 0 only when ˜ L val ( f k ) ≤ min k ⋆ ˜ L val ( f k ⋆ ) + γ/ 3 . Using the fact that max k | ˜ L val ( f k ) -̂ L val ( f k ) | ≤ γ/ 3 , we deduce that λ is supported on { k : ̂ L val ( f k ) ≤ min k ⋆ ∈ [ K ] ̂ L val ( f k ⋆ ) + γ } .

Therefore, θ = ( θ (1) attn , θ (1) mlp , θ (2) mlp , θ (3) attn ) is the desired transformer, with

and

This completes the proof.

I.2 Proof of Theorem 11

We first restate Theorem 11 into the following version which provides additional size bounds for θ . For the simplicity of presentation, throughout this subsection and Appendix J, we denote I t = { i : ( x i , y i ) ∈ D train } , I v = { i : ( x i , y i ) ∈ D val } , X train = [ x i ] i ∈I t to be the input matrix corresponding to the training split only, and N train = |D train | , N val = |D val | .

Theorem I.2. For any sequence of regularizations { λ k } k ∈ [ K ] , 0 ≤ α ≤ β with κ := max k β + λ k α + λ k , B w > 0 , γ > 0 , and ε < B w / 2 , suppose in input format (3) we have D ≥ Θ( Kd ) . Then there exists an L -layer transformer TF θ with

such that the following holds. On any input data ( D , x N +1 ) such that the problem (ICRidge) is well-conditioned and has a bounded solution:

TF 0 θ approximately implements ridge selection: its prediction

satisfies the following.

For each k ∈ [ K ] , ̂ w k = ̂ w k ( D train ) approximates the ridge estimator w λ k ridge ( D train ) , i.e. ∥ ∥ ̂ w k -w λ k ridge ( D train ) ∥ ∥ 2 ≤ ε .

In particular, if we set γ ′ = 2( B x B w + B y ) B x ε + γ , then it holds that 8

where we denote ̂ w λ k ridge , train := w λ k ridge ( D train ) .

To prove Theorem I.2, we first show that, for the squared validation loss, there exists a 3-layer transformer that performs predictor selection based on the exactly evaluated ̂ L val ( f k ) for each k ∈ [ K ] . (Proof in Appendix I.2.1.)

Theorem I.3 (Square-loss version of Theorem I.1) . Consider the squared validation loss

Then there exists a 3-layer transformer TF θ with

such that for any input H that takes form

where TF θ outputs h N +1 = [ x N +1 ; ̂ f ( x N +1 ); ∗ ; 1; 0] , where the predictor ̂ f : R d → R is a convex combination of { f k : ̂ L val ( f k ) ≤ min k ⋆ ∈ [ K ] ̂ L val ( f k ⋆ ) + γ } . As a corollary, for any convex risk L : ( R d → R ) → R , ̂ f satisfies

Proof of Theorem I.2 First, by the proof 9 of Theorem 4 and Proposition B.6, for each k ∈ [ K ] , there exists a T = L -3 layer transformer θ (1: T ) such that TF θ (1: T ) maps

so that if (42) holds, we have ∥ ∥ ̂ w k -w λ k ridge ∥ ∥ 2 ≤ ε and ̂ w k ∈ B 2 ( B w ) .

Next, by Theorem I.3, there exists a 3-layer transformer θ ( T +1: T +3) that outputs

where ̂ w = ∑ K k =1 λ k ̂ w k , λ = ( λ 1 , · · · , λ K ) ∈ ∆([ K ]) so that

This is the desired result.

I.2.1 Proof of Theorem I.3

Similar to the proof of Proposition 10, Theorem I.3 is a direct corollary by combining Proposition I.3 with Proposition I.2.

Proposition I.3 (Evaluation layer for the squared loss) . There exists an attention layer TF θ with 2 K heads and | | | θ | | | ≤ 3 R +2 NK/ |D val | such that TF θ maps

Proof of Proposition I.3. For every k ∈ [ K ] , we define matrices Q m,k , K m,k , V m,k ∈ R D × D such that for all i, j ∈ [ N +1] ,

As the input has structure h i = [ x i ; y ′ i ; ∗ ; f 1 ( x i ); · · · ; f K ( x i ); 0 K +1 ; 1; t i ] , these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Now, for every i, j ∈ [ N +1] , we have

where the second equality follows from the bound | f k ( x j ) -y j | ≤ 2 R , so that the relus equals 0 if t j ≤ 0 . Thus letting the attention layer θ = { ( V k,w , Q k,w , K k,w ) } ( k,w ) ∈ [ K ] ×{ 0 , 1 } , we have

This is the desired result.

I.3 Proofs for Section 4.2

I.3.1 Proof of Lemma 13

It is straightforward to check that the binary type check ψ : R → R can be expressed as a linear combination of 6 relu's (recalling σ ( · ) = ReLU( · ) ):

with ∑ m | a m | = 8 /ε , max m max {| b m | , | c m |} ≤ 2 . We can thus construct an attention layer θ = { ( Q m , K m , V m ) } 6 m =1 with 6 heads such that

which gives that for every i ∈ [ N +1] ,

Further, we have | | | θ | | | ≤ 18 /ε = O (1 /ε ) . This is the desired result.

By composing the above attention layer with one additional layer (with 2 heads) that implement the following function

on the output Ψ binary ( D ) , we directly obtain the following corollary.

Corollary I.1 (Thresholded binary test) . There exists a two-layer attention-only transformer with max ℓ ∈ [2] M ( ℓ ) ≤ 6 and | | | θ | | | ≤ O (1 /ε ) that exactly implements the thresholded binary test

at every token i ∈ [ N +1] , where we recall the definition of Ψ binary in Lemma 13.

I.3.2 Formal statement and proof of Proposition 14

We say a distribution P y on R is ( C, ε 0 ) -not-concentrated around { 0 , 1 } if

for all ε ∈ (0 , ε 0 ] . A sufficient condition is that the density p y is upper bounded by C within [ -ε 0 , ε 0 ] ∪ [1 -ε 0 , 1 + ε 0 ] .

Throughout this section, let σ log ( t ) := (1 + e -t ) -1 denote the sigmoid activation, and let ̂ w log denote the solution to the in-context logistic regression problem, i.e. (ICGLM) with g ( · ) = σ log ( · ) .

Proposition I.4 (Adaptive regression or classification; Formal version of Proposition 14) . For any B w > 0 , ε ≤ B x B w / 10 , 0 < α ≤ β with κ := β/α , and any ( C, ε 0 ) , there exists a L -layer attention-only transformer with

(with R := max { B x B w , B y , 1 } , and ε depending only on ( C, ε 0 ) ) such that the following holds. Suppose the input format is (3) with dimension D ≥ 3 d +4 .

On any classification instance ( D , x N +1 ) (such that { y i } i ∈ [ N ] ⊂ { 0 , 1 } ) that is well-conditioned for logistic regression in the sense of (35), it outputs ̂ y N +1 that ε -approximates the prediction of in-context logistic regression:

On the contrary, for regression problems , i.e. any in-context distribution P whose marginal P y is ( C, ε 0 ) -not-concentrated around { 0 , 1 } , with probability at least 1 -exp( -cN ) over D (where c > 0 depends only on ( C, ε 0 ) ), ̂ y N +1 ε -approximates the prediction of in-context least squares if the data is well-conditioned:

where ̂ w LS denotes the in-context least squares estimator, i.e. (ICRidge) with λ = 0 .

Proof. The result follows by combining the binary test in Corollary I.1 with Theorem 4 and Theorem G.1. By those results, there exists three attention-only transformers θ LS , θ log , θ bin , with (below L g , C g = Θ(1) for g = σ log ( · ) )

that outputs prediction ̂ y LS N +1 , ̂ y log N +1 (at the ( N + 1) -th token) and Ψ binary thres ( D ) (at every token) respectively, which satisfy

when the corresponding well-conditionednesses are satisfied. In particular, we can make ̂ w log welldefined on non-binary data, by multiplying Ψ binary thres ( D ) onto the x i 's (which can be implemented by slightly modifying θ log without changing the order of the number of layers, heads, and norms) so that ̂ w log = 0 on any data where Ψ binary thres ( D ) = 0 .

By joining θ LS and θ log using Proposition B.6, concatenating with θ bin before, and concatenating with one additional attention layer with 2 heads after to implement

we obtain a single transformer θ with

which outputs (44) as its prediction (at the location for ̂ y N +1 ).

It remains to show that (44) reduces to either one of ̂ y log N +1 or ̂ y LS N +1 . When the data are binary ( y i ∈ { 0 , 1 } ), we have Ψ binary ( D ) = 1 and Ψ binary thres ( D ) = 1 , in which case (44) becomes exactly ̂ y log N +1 . By contrast, when data is sampled from a distribution that is ( C, ε 0 ) -not-concentrated around { 0 , 1 } , we have for any fixed ε ≤ ε 0 ∧ 1 4 C that, letting B ε := [ -ε, ε ] ∪ [1 -ε, 1 + ε ] and p ε := P y ( B ε ) ≤ Cε ≤ 1 4 , by Hoeffding's inequality,

̸

where c ′ > 0 is an absolute constant. On the event Ψ binary thres ( D ) = 0 (which happens with probability at least 1 -exp( -c ′ N ) ), (44) becomes exactly ̂ y LS N +1 . This finishes the proof.

I.4 Linear correlation test and application

In this section, we give another instantiation of the pre-ICL testing mechanism by showing that the transformer can implement a linear correlation test that tests whether the correlation vector E [ x y ] has a large norm. We then use this test to construct a transformer to perform 'confident linear regression', i.e. output a prediction from linear regression only when the signal-to-noise ratio is high.

For any fixed parameters λ min , B ⋆ w > 0 , consider the linear correlation test over data D defined as

Recall that σ ( · ) = ReLU( · ) above denotes the relu activation.

We show that Ψ lin can be exactly implemented by a 3-layer transformer.

Lemma I.1 (Expressing Ψ lin by transformer) . There exists a 3-layer attention-only transformer TF θ with at most 2 heads per layer and | | | θ | || ≤ O (1 + λ 2 min ( B ⋆ w ) 2 ) such that on input sequence H of the form (3) with D ≥ 2 d + 4 , the transformer exactly implements Ψ lin : it outputs ˜ H such that ˜ h i = [ x i ; y i t i ; ∗ ; Ψ lin ( D ); 1] for all i ∈ [ N +1] .

Proof. We begin by noting the following basic facts:

Identity function can be implemented exactly by two ReLUs: t = σ ( t ) -σ ( -t ) .

Squared ℓ 2 norm can be implemented exactly by a single attention head (assuming every input h i contains the same vector g ): ∥ g ∥ 2 2 = σ ( ⟨ g , g ⟩ ) .

We construct the transformer θ as follows.

Layer 1: Use 2 heads to implement ̂ t = 1 N ∑ N i =1 x i y i , where V (1) { 1 , 2 } h j = [ ± x j ; 0 D -d ] , Q (1) { 1 , 2 } h i = [ N +1 N ; 0 D -1 ] , and K (1) { 1 , 2 } h j = [ ± y j t j ; 0 D -1 ] = [ ± y j 1 { j < N + 1 } ; 0 D -1 ] (where we recall t j = 1 { j < N + 1 } and note that y j t j corresponds exactly to the location for y j in H , cf. (3)).

By manipulating the output dimension in V (1) , write the result ̂ t into blank memory space with dimension d at every token i ∈ [ N +1] .

Layer 2: Use a single head to compute ∥ ̂ t ∥ 2 2 : Q (2) 1 h (1) i = [ ̂ t ; 0 D -d ] , K (2) 1 h (1) j = [ ̂ t ; 0 D -d ] , and V (2) 1 h (1) j = [1; 0 D -1 ] . By manipulating the output dimension in V (2) , write the result ∥ ̂ t ∥ 2 2 into blank memory space with dimension 1 at every token i ∈ [ N +1] . After layer 2, we have h (3) i = [ x i ; y i t i ; ∗ ; ∥ ̂ t ∥ 2 2 ; ∗ ; 1] .

Layer 3: Use 2 heads to implement two ReLU functions with bias: ∥ ̂ t ∥ 2 2 ↦→ 1 B -A ( σ ( ∥ ̂ t ∥ 2 2 -A ) -σ ( ∥ ̂ t ∥ 2 2 -B )) . The two query (or key) matrices contain values A and B . In our problem we take

so that the above ReLU function implements Ψ lin ( D ) exactly. Write the result into a blank memory space with dimension 1. We finish the proof by noting that | | | θ | | | ≤ O (1 + λ 2 min ( B ⋆ w ) 2 ) .

Statistical guarantee for Ψ lin We consider the following well-posedness assumption for the linear correlation test Ψ lin . Note that, similar as Assumption A, the assumption does not require the data to be generated from any true linear model, but rather only requires some properties about the best linear fit w ⋆ P , as well as sub-Gaussianity conditions.

Assumption D (Well-posedness for linear correlation test) . We say a distribution P on R d × R is well-posed for linear independence tests, if ( x , y ) ∼ P satisfies

∥ x ∥ 2 ≤ B x and | y | ≤ B y almost surely;

The covariance Σ P := E P [ xx ⊤ ] satisfies λ min I d ⪯ Σ P ⪯ λ max I d , with 0 < λ min ≤ λ max , and κ := λ max /λ min .

The whitened vector Σ -1 / 2 P x is K 2 -sub-Gaussian for some K ≥ 1 .

The best linear predictor w ⋆ P := E P [ xx ⊤ ] -1 E P [ x y ] satisfies ∥ w ⋆ P ∥ 2 ≤ B ⋆ w .

The label y is σ 2 -sub-Gaussian.

The residual z := y -⟨ x , w ⋆ P ⟩ is σ 2 -sub-Gaussian with probability one (over x ).

The following results states that Ψ lin achieves high power as long as the sample size is high enough, and the signal ∥ w ⋆ P ∥ 2 is either sufficiently high or sufficiently low.

Proposition I.5 (Power of linear correlation test) . Suppose distribution P satisfies Assumption D with parameters λ min , λ max , B ⋆ w . Then, for the linear correlation test Ψ lin with parameters ( λ min , B ⋆ w ) with B ⋆ w ≤ B ⋆ w and any N ≥ ˜ O ( max { K 4 , λ max K 2 σ 2 ( B ⋆ w ) 2 λ 2 min } · d ) , we have

If ∥ w ⋆ P ∥ 2 ≥ B ⋆ w , then with probability at least 1 -δ over D , we have Ψ lin ( D ) = 1 .

If ∥ w ⋆ P ∥ 2 ≤ λ min 10 λ max B ⋆ w , then with probability at least 1 -δ over D , we have Ψ lin ( D ) = 0 .

Proof. For any P satisfying Assumption D, note that E [ x z ] = E [ x ( y -⟨ w ⋆ P , x ⟩ )] = 0 by construction. Therefore, by standard sub-Gaussian and sub-exponential concentration combined with union bound, the following events hold simultaneously with probability at least 1 -δ :

On the above event, we have

Therefore, in case 1, we have

In case 2, we have

The proof is finished by recalling the definition of Ψ lin in (45), so that Ψ lin ( D ) = 1 if ∥ ̂ t ∥ 2 ≥ 3 λ min B ⋆ w / 4 , and Ψ lin ( D ) = 0 if ∥ ̂ t ∥ 2 ≤ λ min B ⋆ w / 4 .

Application: Confident linear regression By directly composing the linear correlation test in Lemma I.1 with the transformer construction in Corollary 5 (using an argument similar as the proof of Proposition I.4), and using the power of the linear correlation test Proposition I.5, we immediately obtain the following result, which outputs a prediction from (approximately) least squares if ̂ ψ := Ψ lin ( D ) = 1 , and abstains from predicting if ̂ ψ = 0 . This can be viewed as a form of 'confident linear regression', where the model predicts only if it thinks the linear signal is strong enough.

Proposition I.6 (Confident linear regression) . For any B w > 0 , 0 < B ⋆ w ≤ B ⋆ w , 0 ≤ λ min ≤ λ max , ε ≤ B x B w / 10 , 0 < α ≤ β with κ := β/α , there exists a L -layer attention-only transformer with

(with R := max { B x B w , B y , 1 } ) such that the following holds. Let N ≥ ˜ O ( max { K 4 , λ max K 2 σ 2 ( B ⋆ w ) 2 λ 2 min } · d ) . Suppose the input format is (3) with dimension D ≥ 2 d +4 . Let ICL instance ( D , x N +1 ) be drawn from any distribution P satisfying Assumption D. Then the transformer outputs a 2-dimensional prediction (within the test token ˜ h N +1 )

such that the following holds:

If ∥ w ⋆ P ∥ 2 ≥ B ⋆ w , then with probability at least 1 -δ over D , we have | ̂ y N +1 -⟨ ̂ w LS , x N +1 ⟩ | ≤ ε , and ̂ ψ = 1 if D is in addition well-conditioned for least squares (in the sense of (5) with λ = 0 ).

If ∥ w ⋆ P ∥ 2 ≤ λ min 10 λ max B ⋆ w , then with probability at least 1 -δ over D , we have ̂ y N +1 = 0 and ̂ ψ = 0 .

J Proof of Theorem 12: Noisy linear model with mixed noise levels

For each fixed k ∈ [ K ] , we consider the following data generating model P k , where we first sample P = P w ⋆ ,σ k ∼ π from w ⋆ ∼ N ( 0 , I d /d ) , and then sample data { ( x i , y i ) } i ∈ [ N +1] iid ∼ P w ⋆ ,σ k as

Also, recall that the Bayes optimal estimator on P k is given by ̂ y Bayes N +1 = 〈 w λ k ridge ( D ) , x N +1 〉 with ridge λ k = σ 2 k d/N , and the Bayes risk on P k is given by

Recall that in Section 4.1.1, we consider a mixture law P π that generates data from P k with k ∼ Λ . It is clear that we have (pushing inf A into E k ∼ Λ does not increase the value) we have

i.e., the Bayes risk can only be greater if we consider a mixture of models. In other words, if a transformer can achieve near-Bayes ICL on each meta-task P k , then it can perform near-Bayes ICL on any meta-task π which is a mixture of P k with k ∼ Λ . Therefore, to prove Theorem 12, it suffices to show the following (strengthened) result.

Theorem J.1 (Formal version of Theorem 12) . Suppose that N ≥ 0 . 1 d and we write σ max = max k { σ k , 1 } , σ min = min k { σ k , 1 } . Suppose in input format (3) we have D ≥ Θ( Kd ) . Then there exists a transformer θ with

such that for any k ∈ [ K ] , it holds that

if we choose N val := |D val | ≍ N 2 / 3 [log K ] 1 / 3 .

The core of the proof of Theorem J.1 is to show that any estimator ̂ w that achieves small validation loss ̂ L val must achieve small population loss.

Throughout the rest of this section, recall that we define N train = |D train | , N val = |D val | , I train = { i : ( x i , y i ) ∈ D train } , I val = { i : ( x i , y i ) ∈ D val } , and X train = [ x i ] i ∈I train .

J.1 Proof of Theorem J.1

Fix parameters δ, ε, γ > 0 and a large universal constant C 0 . Let us set

Then, we define good events similarly to the proof of Corollary 6 (Appendix F.4):

For the good event E := E π ∩E w ∩E b, train ∩E b, test ∩E b,N +1 , we can show that P ( E c ) ≤ O ( N -10 ) . Further, by the proof of Lemma F.1 (see e.g. (34)), we know that max k ∈ [ K ] ∥ ∥ w λ k ridge ( D train ) ∥ ∥ 2 ≤ B w / 2 holds under the good event E .

For the ridge λ k = dσ 2 k N train and parameters ( α, β, γ, ε ) , we consider the transformer θ constructed in Theorem I.2, with a clipped prediction y N +1 = read y (TF θ ( H )) .

̂ ˜

In the following, we upper bound the quantity E k ( ̂ y N +1 -y N +1 ) 2 for any fixed k . Similar to the proof of Corollary 6 (Appendix F.4), we decompose

and we analyze these two parts separately.

Part I. Recall that by our construction, when E holds, we have ̂ y N +1 = clip B y ( ⟨ ̂ w , x N +1 ⟩ ) and the statements of Theorem I.2 hold for ̂ w . Thus, we have

Let us consider the following risk functional

By our construction, under the good event E 0 , we have

where we denote 2 Risk k, train = E k ∥ ∥ w λ k ridge ( D train ) -w ⋆ ∥ ∥ 2 2 + σ 2 k , and we also note that Risk k, train ≤ 1 + σ 2 k by definition. By Lemma J.1, we have

We next deal with the term ε val := max l ∈ [ K ] ∣ ∣ ∣ ̂ L val ( ̂ w l ( D train )) -L val , w ⋆ ( ̂ w l ( D train )) ∣ ∣ ∣ . Note that for the good event E train := E π ∩ E w ∩ E b, train of ( w ⋆ , D train ) , we have

Therefore, we can conclude that

Therefore, we can choose ( ε, N val ) so that N val ≤ N/ 2 as

It is worth noting that such choice of N val is feasible as long as N ≳ B 4 w σ 4 max log K . Under such choice, we obtain

Part II. Similar to the proof of Corollary 6, we have

Conclusion. Combining the both cases, we obtain

where we plug in our choice of B y . The bounds on M ( ℓ ) , D ( ℓ ) and | | | θ | | | follows immediately from Theorem I.2. This completes the proof.

J.2 Derivation of the exact Bayes predictor

Let ( D , x N +1 , y N +1 ) be ( N +1) observations from the data generating model π considered in Section 4.1.1. On observing ( D , x N +1 ) , the Bayes predictor of y N +1 is given by its posterior mean:

It thus remains to derive E π [ w ⋆ |D ] . Recall that our data generating model is given by k ∼ Λ , By Bayes' rule, we have

On k = k ′ , the data is generated from the noisy linear model w ⋆ ∼ N ( 0 , I d /d ) , and y = Xw ⋆ + ε where ε i iid ∼ N (0 , σ 2 k ′ ) . It is a standard result that E π [ w ⋆ |D , k = k ′ ] is given by the ridge estimator

(Note that the sample covariance within ̂ Σ k ′ is not normalized by N , which is not to be confused with remaining parts within the paper.) Therefore, the posterior mean (46) is exactly a weighted combination of K ridge regression estimators, each with regularization dσ 2 k /N .

It remains to derive the mixing weights P π ( k = k ′ |D ) for all k ′ ∈ [ K ] . By Bayes' rule, we have

Note that such mixing weights involve the determinant of the matrix ̂ Σ k ′ = X ⊤ X + dσ 2 k ′ I d , which depends on the data X in a non-trivial fashion; Any transformer has to approximate these weights if their mechanism is to directly approximate the exact Bayesian predictor (46).

J.3 Useful lemmas

Lemma J.1. For 2 Risk k, train = E k ∥ ∥ w λ k ridge ( D train ) -w ⋆ ∥ ∥ 2 2 + σ 2 k , there exists universal constant C such that

Proof. Recall that under P k , we have

We denote y t = [ y i ] i ∈I train , then by definition w λ k ridge ( D train ) = ( X ⊤ train X train + dσ 2 k ) -1 X train y t (with λ k = dσ 2 k /N train ). Thus, a simple calculation yields

and analogously, 2 BayesRisk k = σ 2 k E tr ( ( X ⊤ X + dσ 2 k I d ) -1 ) + σ 2 k . Therefore,

where in the above inequality we denote Σ := X ⊤ train X train + dσ 2 k I d and use the following fact:

Case 1. We first suppose that N train ≤ 16 d . Then by definition Σ ⪰ dσ 2 k I d , and hence

Case 2. When N train ≥ 9 d , then we consider the event E t := { λ min ( X ⊤ train X train /N train ) ≥ 1 16 } . By Lemma B.2 we have P ( E c t ) ≤ exp( -N train / 8) . Therefore,

Combining these two cases finishes the proof.

Lemma J.2. Condition on the event E train , we have

where we denote ̂ w l = ̂ w l ( D train ) .

Proof. We only need to work with a fixed pair of ( w ⋆ , D train ) such that E train holds. Hence, in the following we only consider the randomness of D val conditional on such a ( w ⋆ , D train ) . Recall that for any w ,

and we have E D v [ ̂ L val ( w )] = L val , w ⋆ ( w ) . For each i ∈ I val ,

Note that under E train , we have ̂ w l ∈ B 2 ( B w ) for all l ∈ [ K ] , and hence σ 2 k + ∥ w ⋆ -̂ w l ∥ 2 ≤ 5 B 2 w . We then have ( y i - ⟨ x i , ̂ w l ⟩ ) 2 's are (conditional) i.i.d random variables in SE( CB 4 w ) . Then, by Bernstein's inequality, we have

where c is a universal constant. Applying the union bound, we obtain

Taking integration completes the proof.

J.4 Generalized linear models with adaptive link function selection

Suppose that ( g k : R → R ) k ∈ [ K ] is a set of link functions such that g k is non-decreasing and C 2 -smooth for each k ∈ [ K ] . We consider the input format we introduce in Section 4.1 with |D train | = ⌈ N/ 2 ⌉ , |D val | = ⌊ N/ 2 ⌋ .

Theorem J.2 (GLMs with adaptive link function selection) . For any fixed set of parameters defined in Assumption B, as long as N ≥ O ( d ) , there exists a transformer θ with L ≤ O (log( N )) layers, input dimension D = Θ( dK ) and max ℓ ∈ [ L ] M ( ℓ ) ≤ ˜ O ( d 3 N ) , such that the following holds.

For any k ⋆ ∈ [ K ] and any distribution P that is a generalized linear model of the link function g k ⋆ and some parameter β , if Assumption B holds for each pair ( P , g k ) , then

Proof. For each k ∈ [ K ] , we consider optimizing the following training loss:

where ℓ k ( t, y ) := -yt + ∫ t 0 g k ( s ) ds is the convex (integral) loss associated with g k (as in Section 3.1).

Also, for each predictor f : R d → R , we consider the squared validation loss ̂ L val :

Fix a large universal constant C 0 . Let us set

Then, we define good events similarly to the proof of Corollary 6 (Appendix F.4):

Similar to the proof of Theorem G.2 (Appendix G.2), we know the good event E := E w ∩ E r ∩ E b, train ∩ E b, test ∩ E b,N +1 holds with high probability: P ( E c ) ≤ O ( N -10 ) .

Similar to the proof of Theorem I.2, we can show that there exists a transformer θ with prediction ̂ y N +1 = ˜ read y (TF θ ( H )) (clipped by B y ), such that (for any P ) the following holds under E :

(a) For each k ∈ [ K ] , f k = A k ( D train ) is a predictor such that ∣ ∣ ∣ f k ( x i ) -g k ( ⟨ x i , w ( k ) GLM ⟩ ) ∣ ∣ ∣ ≤ ε for all i ∈ [ N +1] (where ε is chosen as in Appendix G.2).

(b) ̂ y N +1 = clip B y ( ̂ f ( x N +1 )) , where ̂ f = A TF ( D ) is an aggregated predictor given by ̂ f = ∑ k λ k f k , such that ( λ k ) is a distribution supported on k ∈ [ K ] such that ̂ L val ( f k ) ≤ min k ′ ∈ [ K ] ̂ L val ( f k ′ ) + γ.

Similar to the proof of Theorem G.2, for E 0 := E w ∩ E r ∩ E b, train ∩ E b, test , we have

where we denote L val ( f ) := E ( x ,y ) ∼ P [ 1 {∥ x ∥ 2 ≤ B x } ( f ( x ) -y ) 2 ] for each predictor f . By the definition of ̂ f , we then have (under E 0 )

For the first term, repeating the argument in the proof of Theorem G.2 directly yields that for E train := E w ∩ E r ∩ E b, train ,

For the second term, similar to Lemma J.2, we can show that conditional on D train such that E train holds, it holds

Combining these inequalities and suitably choosing γ complete the proof.

K Analysis of pretraining

Thus far, we have established the existence of transformers for performing various ICL tasks with good in-context statistical performance. We now analyze the sample complexity of pretraining these transformers from a finite number of training ICL instances.

K.1 Generalization guarantee for pretraining

Setup At pretraining time, each training ICL instance has form Z := ( H , y N +1 ) , where H := H ( D , x N +1 ) ∈ R D × ( N +1) denote the input sequence formatted as in (3). We consider the square loss between the in-context prediction and the ground truth label:

.

Above, clip B y ( t ) := max { min { t, B y } , -B y } is the standard clipping operator onto [ -B y , B y ] , and TF R θ the transformer architecture as in Definition 3 with clipping operators after each layer: let H (0) = clip R ( H ) ,

The clipping operator is used to control the Lipschitz constant of TF θ with respect to θ , and we typically choose a sufficiently large clipping radius R so that it does not modify the behavior of the transformer on any input sequence of our concern.

Wedraw ICL instances Z := ( H , y N +1 ) = ( D , ( x N +1 , y N +1 )) from a (meta-)distribution denoted as π , which first sample an in-context data distribution P ∼ π , then sample iid examples ( x i , y i ) N +1 i =1 iid ∼ P ⊗ ( N +1) and form D = { ( x i , y i ) } i ∈ [ N ] . Our pretraining loss is the average ICL loss on n pretraining instances Z (1: n ) iid ∼ π , and we consider the corresponding test ICL loss on a new test instance:

Our pretraining algorithm is to solve a standard constrained empirical risk minimization (ERM) problem over transformers with L layers, M heads, and norm bound B (recall the definition of the | | |·| | | norm in (2)):

Generalization guarantee By standard uniform concentration analysis via chaining arguments (Proposition B.4; see also [87, Chapter 5] for similar arguments), we have the following excess loss guarantee for (TF-ERM). The proof can be found in Appendix L.2.

Theorem K.1 (Generalization for pretraining) . With probability at least 1 -ξ (over the pretraining instances { Z j } j ∈ [ n ] ), the solution ̂ θ to (TF-ERM) satisfies

where ι = log(2 + max { B, R , B y } ) is a log factor.

K.2 Examples of pretraining for in-context regression problems

In Theorem K.1, the comparator inf θ ∈ Θ L,M,D ′ ,B L icl ( θ ) is simply the smallest expected ICL loss for ICL instances drawn from π , among all transformers within the norm ball Θ L,M,D ′ ,B . Using our constructions in Section 3 & 4, we show that this comparator loss is small on various (meta)distribution π 's, by which we obtain end-to-end guarantees for pretraining transformers with small ICL loss at test time. Here we showcase this argument on several representative regression problems.

Linear regression For any in-context data distribution P , let w ⋆ P := E P [ xx ⊤ ] -1 E P [ x y ] denote the best linear predictor for P . We show that with mild choices of L, M, B , the learned transformer can perform in-context linear regression with near-optimal statistical power, in that on the sampled P ∼ π and ICL instance { ( x i , y i ) } i ∈ [ N +1] iid ∼ P , it competes with the best linear predictor w ⋆ P for this particular P . The proof follows directly by on combining Corollary 5 with Theorem K.1, and can be found in Appendix L.3.

Theorem K.2 (Pretraining transformers for in-context linear regression) . Suppose P ∼ π is almost surely well-posed for in-context linear regression (Assumption A) with the canonical parameters. Then, for N ≥ ˜ O ( d ) , with probability at least 1 -ξ (over the training instances Z (1: n ) ), the solution

̂ θ of (TF-ERM) with L = O ( κ log( κN/σ )) layers, M = 3 heads, D ′ = 0 (attention-only), and B = O ( √ κd ) achieves small excess ICL risk over w ⋆ P :

where ˜ O ( · ) only hides polylogarithmic factors in κ, N, 1 /σ .

To our best knowledge, Theorem K.2 offers the first end-to-end result for pretraining a transformer to perform in-context linear regression with explicit excess loss bounds. The ˜ O ( √ κ 2 d 2 /n ) term originates from the generalization of pretraining (Theorem K.1), where as the ˜ O ( dσ 2 /N ) term agrees with the standard fast rate for the excess loss of linear regression [38]. Further, as long as n ≥ ˜ O ( κ 2 N/σ 2 ) , the excess risk achieves the optimal rate ˜ O ( dσ 2 /N ) (up to log factors).

Additional examples By similar arguments as in the proof of Theorem K.2, we can directly turn most of our other expressivity results into results on the pretrained transformers. Here we present three such additional examples (proofs in Appendix L.4-L.6). The first example is for the sparse linear regression problem considered in Theorem 8.

Theorem K.3 (Pretraining transformers for in-context sparse linear regression) . Suppose each P ∼ π is almost surely an instance of the sparse linear model specified in Theorem 8 with parameters B ⋆ w and σ . Suppose N ≥ ˜ O ( s log(( d ∨ N ) /σ )) and let κ := B ⋆ w /σ .

Then with probability at least 1 -ξ (over the training instances Z (1: n ) ), the solution ̂ θ of (TF-ERM) with L = ˜ O ( κ 2 (1 + d/N )) layers, M = 2 heads, D ′ = 2 d , and B = ˜ O (poly( d, B ⋆ w , σ )) achieves small excess ICL risk:

where ˜ O ( · ) only hides polylogarithmic factors in d, N, 1 /σ .

Our next example is for the problem of noisy linear regression with mixed noise levels considered in Theorem 12 and Theorem J.1. There, the constructed transformer uses the post-ICL validation mechanism to perform ridge regression with an adaptive regulariation strength depending on the particular input sequence.

Theorem K.4 (Pretraining transformers for in-context noisy linear regression with algorithm selection) . Suppose π is the data generating model (noisy linear model with mixed noise levels) considered in Theorem J.1, with σ max ≤ O (1) . Let N ≥ d/ 10 .

Then, with probability at least 1 -ξ (over the training instances Z (1: n ) ), the solution ̂ θ of (TFERM) with input dimension D = Θ( dK ) , L = O ( σ -2 min log( N/σ min )) layers, M = O ( K ) heads, D ′ = O ( K 2 ) , and B = O (poly( K,σ -1 min , d, N )) achieves small excess ICL risk:

where ˜ O ( · ) only hides polylogarithmic factors in d, N, K, 1 /σ min .

Our final example is for in-context logistic regression. For simplicity we consider the realizable case. Theorem K.5 (Pretraining transformers for in-context logistic regression; square loss guarantee) . Suppose for P ∼ π , P is almost surely a realizable logistic model (i.e. P = P log β with ∥ β ∥ 2 ≤ B ⋆ w as in Corollary G.1). Suppose that B ⋆ w = O (1) and N ≥ O ( d ) .

Then, with probability at least 1 -ξ (over the training instances Z (1: n ) ), the solution ̂ θ of (TF-ERM) with L = O (log( N )) layers, M = ˜ O ( d 3 N ) heads, D ′ = 0 , and B = O (poly( d, N )) achieves small excess ICL risk:

where ˜ O ( · ) only hides polylogarithmic factors in d, N .

Remark on generality of transformer All results above are established by the expressivity results in Section 3 & 4 for transformers to implement various ICL procedures (such as least squares, Lasso, GLM, and ridge regression with in-context algorithm selection), combined with the generalization bound (Theorem K.1). However, the transformer itself was not specified to encode any actual structure about the problem at hand in any result above, other than having sufficiently large number of layers, number of heads, and weight norms, which illustrates the flexibility of the transformer architecture.

L Proofs for Section K

L.1 Lipschitzness of transformers

For any p ∈ [1 , ∞ ] , let ∥ H ∥ 2 ,p := ( ∑ N i =1 ∥ h i ∥ p 2 ) 1 /p denote the column-wise (2 , p ) -norm of H . For any radius R > 0 , we denote H R := { H : ∥ H ∥ 2 , ∞ ≤ R } be the ball of radius R under norm ∥·∥ 2 , ∞ .

Lemma L.1. For a single MLP layer θ mlp = ( W 1 , W 2 ) , we introduce its norm (as in (2) )

For any fixed hidden dimension D ′ , we consider

Then for H ∈ H R , θ mlp ∈ Θ mlp ,B , the function ( θ mlp , H ) ↦→ MLP θ mlp ( H ) is ( B R ) -Lipschitz w.r.t. θ mlp and (1 + B 2 ) -Lipschitz w.r.t. H .

Proof. Recall that by our definition, for the parameter θ mlp = ( W 1 , W 2 ) ∈ Θ mlp ,B and the input H = [ h i ] ∈ R D × N , the output MLP θ mlp ( H ) = H + W 2 σ ( W 1 H ) = [ h i + W 2 σ ( W 1 h i )] i . Therefore, for θ ′ mlp = ( W ′ 1 , W ′ 2 ) ∈ Θ mlp ,B , we have

where the second inequality follows from the 1-Lipschitznees of σ = [ · ] + . Similarly, for H ′ = [ h ′ i ] ∈ R D × N ,

Lemma L.2. For a single attention layer θ attn = { ( V m , Q m , K m ) } m ∈ [ M ] ⊂ R D × D , we introduce its norm (as in (2) )

For any fixed dimension D , we consider

Then for H ∈ H R , θ attn ∈ Θ attn ,B , the function ( θ attn , H ) ↦→ Attn θ attn ( H ) is ( B 2 R 3 ) -Lipschitz w.r.t. θ attn and (1 + B 3 R 2 ) -Lipschitz w.r.t. H .

Proof. Recall that by our definition, for the parameter θ attn = { ( V m , Q m , K m ) } m ∈ [ M ] ∈ Θ attn ,B and the input H = [ h i ] ∈ R D × N , the output Attn θ attn ( H ) = [ ˜ h i ] is given by

Now, for θ ′ attn = { ( V ′ m , Q ′ m , K ′ m ) } m ∈ [ M ] , we consider

∥ ∥

where the second inequality uses the definition of operator norm, the third inequality follows from the triangle inequality, the forth inequality is because ∥ Q m h i ∥ 2 ≤ B R , ∥ K m h j ∥ 2 ≤ B R , and σ is 1-Lipschitz. This completes the proof the Lipschitzness w.r.t. θ attn .

Similarly, we consider H ′ = [ h ′ i ] , and

By definition, we can similarly bound

where the last inequality uses | | | θ attn | | | ≤ B and the AM-GM inequality. This completes the proof the Lipschitzness w.r.t. H .

Corollary L.1. For a fixed number of heads M and hidden dimension D ′ , we consider

Then for the function TF R given by

TF R is B Θ -Lipschitz w.r.t θ and L H -Lipschitz w.r.t. H , where B Θ := B R (1 + B R 2 + B 3 R 2 ) and B H := (1 + B 2 )(1 + B 2 R 3 ) .

Proof. For any θ = ( θ attn , θ mlp ) , H ∈ H R , and θ ′ = ( θ ′ attn , θ ′ mlp ) , we have

where the second inequality follows from Lemma L.2 and Lemma L.1 and the fact that ∥ Attn θ attn ( H ) ∥ 2 , ∞ ≤ R := R + B 3 R 3 for all H ∈ H R .

Furthermore, for H ′ ∈ H R , we have

which also follows from Lemma L.2 and Lemma L.1.

Proposition L.1 (Lipschitzness of transformers) . For a fixed number of heads M and hidden dimension D ′ , we consider

Then the function TF R is ( LB L -1 H B Θ ) -Lipschitz w.r.t θ ∈ Θ TF ,L,B for any fixed H .

Proof. For θ = θ (1: L ) ∈ Θ TF ,L,B , ˜ θ = ˜ θ (1: L ) ∈ Θ TF ,L,B , we have

where the second inequality follows from Corollary L.1, and the last inequality is because B H ≥ 1 .

L.2 Proof of Theorem K.1

In this section, we prove a slightly more general result by considering the general ICL loss

We assume that the loss function ℓ satisfies sup | ℓ | ≤ B 0 ℓ and sup | ∂ 1 ℓ | ≤ B 1 ℓ . For the special case ℓ ( s, t ) = 1 2 ( s -t ) 2 , we can take B 0 ℓ = 4 B 2 y , B 1 ℓ = 2 B y .

We then consider completes the proof.

L.3 Proof of Theorem K.2

By Corollary 5, there exists a transformer TF θ such that for every P satisfying Assumption A with canonical parameters (and thus in expectation over P ∼ π ) and every N ≥ ˜ O ( d ) , it outputs prediction ̂ y N +1 = ˜ read y (TF θ ( H )) such that

where we recall that L P ( w ⋆ P ) := 1 2 E ( x ,y ) ∼ P [ ( y -⟨ w ⋆ P , x ⟩ ) 2 ] . By inspecting the proof, the same result holds if we change TF θ to the clipped version TF R θ if we choose R 2 = O ( B 2 x + B 2 y + B 2 w +1) = O ( d + κ ) , so that on the good event E cov ∩ E w considered therein, all intermediate outputs within

where Z (1: n ) are i.i.d copies of Z ∼ P , P ∼ π . It remains to apply Proposition B.4 to the random process { X θ } . We verify the preconditions:

(a) By [87, Example 5.8], it holds that log N ( δ ; B | | |·| | | ( r ) , | | |·| | | ) ≤ L (3 MD 2 +2 DD ′ ) log(1 + 2 r/δ ) , where B | | |·| | | ( r ) is any ball of radius r under norm | | |·| | | .

(b) | ℓ icl ( θ ; Z ) | ≤ B 0 ℓ and hence B 0 ℓ -sub-Gaussian.

Therefore, we can apply the uniform concentration result in Proposition B.4 to obtain that, with probability at least 1 -ξ ,

where ι = log(2 + B · LB L -1 H B Θ B 1 ℓ /B 0 ℓ ) ≤ 20 L log(2 + max { B, R , B 1 ℓ /B 0 ℓ } ) . Recalling that

TF θ has ∥·∥ 2 , ∞ ≤ R and thus the clipping does not modify the transformer output on E cov ∩ E w . Further, recall by (33) that θ has size bounds

We can thus apply Theorem K.1 to obtain that the solution ̂ θ to (TF-ERM) with the above choice of ( L, M, B ) and D ′ = 0 (attention-only) satisfies the following with probability at least 1 -ξ :

Above, ι = O (log(1 + max { B y , R , B } )) = ˜ O (1) . This finishes the proof.

L.4 Proof of Theorem K.3

We invoke Theorem 8 (using the construction in Theorem 7 with a different choice of L ) with the following parameters:

where ˜ O ( · ) hides polylogarithmic factors in d, N, B ⋆ w , κ .

Then, Theorem 8 shows that there exists a transformer θ with L layers, max ℓ ∈ [ L ] M ( ℓ ) ≤ M heads, D ′ hidden dimension for the MLP layers, and | | | θ | | | ≤ B such that, on almost surely every P ∼ π , it returns a prediction ̂ y N +1 such that, on the good event E 0 considered therein (over D ∼ P ) which satisfies P ( E 0 ) ≥ 1 -δ ,

By inspecting the proof, the same result holds if we change TF θ to the clipped version TF R θ if we choose R 2 = O ( B 2 x + B 2 y + ( B ⋆ w ) 2 + 1) = O ( d + ( B ⋆ w ) 2 + σ 2 ) , so that on the good event E 0 considered therein, all intermediate outputs within TF θ has ∥·∥ 2 , ∞ ≤ R and thus the clipping does not modify the transformer output on the good event. On the bad event E c 0 , using the same argument as in the proof of Theorem 8, we have

Combining the above two bounds and further taking expectation over P ∼ π gives

We can thus apply Theorem K.1 to obtain that the solution ̂ θ to (TF-ERM) with the above choice of ( L, M, B, D ′ ) satisfies the following with probability at least 1 -ξ :

Above, ι = O (log(1 + max { B y , R , B } )) = ˜ O (1) . This finishes the proof.

L.5 Proof of Theorem K.4

We invoke Theorem 12 and Theorem J.1, which shows that (recalling the input dimension D = Θ( Kd ) ) there exists a transformer θ with the following size bounds:

such that it outputs ̂ y N +1 that satisfies

By inspecting the proof, the same result holds if we change TF θ to the clipped version TF R θ if we choose R 2 = O ( B 2 x + B 2 y + ( B ⋆ w ) 2 + 1) = O ( d + σ 2 max ) , so that on the good event considered therein, all intermediate outputs within TF θ has ∥·∥ 2 , ∞ ≤ R and thus the clipping does not modify the transformer output on the good event. Using this clipping radius, we obtain

We can thus apply Theorem K.1 to obtain that the solution ̂ θ to (TF-ERM) with the above choice of ( L, M, B, D ′ ) satisfies the following with probability at least 1 -ξ :

Above, ι = O (log(1 + max { B y , R , B } )) = ˜ O (1) . This finishes the proof.

L.6 Proof of Theorem K.5

The proof follows from similar arguments as of Theorem K.3 and Theorem K.4, where we plug in the size bounds (number of layers, heads, and weight norms) from Theorem G.2 and Corollary G.1.

M Experimental details and additional studies

M.1 Additional details for Section 6

Architecture and optimization We train a 12-layer encoder-only transformer, where each layer consists of an attention layer as in Definition 1 with M = 8 heads, hidden dimension D = 64 , and ReLU activation (normalized by the sequence length), as well as an MLP layer as in Definition 2 hidden dimension D ′ = 64 . We add Layer Normalization [3] after each attention and MLP layer to help optimization, as in standard implementations [84]. We append linear read-in layer and linear read-out layer before and after the transformer respectively, both applying a same affine transform to all tokens in the sequence and are trainable. The read-in layer maps any input vector to a D -dimensional hidden state, and the read-out layer maps a D -dimensional hidden state to a 1-dimensional scalar.

Each training sequence corresponds to a single ICL instance with N in-context training examples { ( x i , y i ) } N i =1 ⊂ R d × R and test input x N +1 ∈ R d . The input to the transformer is formatted as

in (3) where each token has dimension d + 1 (no zero-paddings). The transformer is trained by minimizing the following loss with fresh mini-batches:

where the loss function ℓ P : R 2 → R may depend on the training data distribution P in general; we use the square loss when P is regression data, and the logistic loss when P is classification data. We use the Adam optimizer with a fixed learning rate 10 -4 , which we find works well for all our experiments. Throughout all our experiments except for the sparse linear regression experiment in Figure 3a, we train the model for 300K steps, where each step consists of a (fresh) minibatch with batch size 64 in the base mode, and K minibatches each with batch size 64 in the mixture mode.

For the sparse linear regression experiment, we find that minimizing the training objective (47) alone was not enough, e.g. for the learned transformer to achieve better loss than the least squares algorithm (which achieves much higher test loss than the Lasso; cf. Figure 3a). To help optimization, we augment (47) with another loss that encourages the second-to-last hidden states to recover the true (sparse) coefficient w ⋆ :

Specifically, the above loss encourages the first N 0 ≤ N tokens within the second-to-last layer to be close to w ⋆ . We choose N 0 = 5 (recall that the total number of tokens is N = 10 and sequence length is N +1 = 11 for this experiment). We minimize the loss L ( θ ) + λL fit -w ( θ ) with λ = 0 . 1 for 2M steps for this task.

Evaluation All evaluations are done on the trained transformer with 6400 test instances. We use the square loss for regression tasks, and the classification error ( 1 -accuracy) between the true label y N +1 ∈ { 0 , 1 } and the predicted label 1 { ̂ y N +1 ≥ 1 / 2 } . We report the means in all experiments, as well as their standard deviations (using one-std error bars) in Figure 2a, 2b, 5a, 5b. In Figure 2c, 3b, 3c 5c, all standard deviations are sufficiently small (not significantly exceeding the width of the markers), thus we did not show error bars in those plots.

Baseline algorithms We implement various baseline machine learning algorithms to compare with the learned transformers. A superset of the algorithms is shown in Figure 3a:

Least squares , Logistic regression : Standard algorithms for linear regression and linear classification, respectively. Note that least squares is also a valid algorithm for classification.

Averaging : The simple algorithm which computes the linear predictor ̂ w = 1 N ∑ N i =1 y i x i and predicts ̂ y N +1 = ⟨ ̂ w , x N +1 ⟩ ;

3-NN : 3-Nearest Neighbors.

Ridge : Standard ridge regression as in (ICRidge). We specifically consider two λ 's (denoted as lam_1 and lam_2 ): λ 1 , λ 2 = (0 . 005 , 0 . 125) . These are the Bayes-optimal regularization strengths for the noise levels ( σ 1 , σ 2 ) = (0 . 1 , 0 . 5) respectively under the noisy linear model (cf. Corollary 6), using the formula λ ⋆ = dσ 2 /N , with ( d, N ) = (20 , 40) .

Lasso : Standard Lasso as in (ICLasso) with λ ∈ { 1 , 0 . 1 , 0 . 01 , 0 . 001 } .

In Figure 2c, the ridge_analytical curve plots the expected risk of ridge regression under the noisy linear model over 20 geometrically spaced values of λ 's in between ( λ 1 , λ 2 ) , using analytical formulae (with Monte Carlo simulations). The Bayes_err_{1,2} indicate the expected risks of λ 1 on task 1 (with noise σ 1 ) and λ 2 on task 2 (with noise σ 2 ), respectively.

M.2 Decoder-based architecture

ICL capabilities have also been demonstrated in the literature for decoder-based architectures [31, 2, 47]. There, the transformer can do in-context predictions at every token x i using past tokens { ( x j , y j ) } j ≤ i -1 as training examples. Here we show that such architectures is also able to perform in-context algorithm selection at every token ; For results for this architecture on 'base' ICL tasks (such as those considered in Figure 3a), we refer the readers to Garg et al. [31].

(a) Linear regression

0

10

20

30

40

in-context examples

0.0

0.2

0.4

0.6

0.8

1.0

square loss

TF_alg_select

TF_reg

TF_cls

(b) Linear classification

0

10

20

30

40

in-context examples

0.15

0.20

0.25

0.30

0.35

0.40

error

TF_alg_select

TF_reg

TF_cls

Logistic Regression

(c) Reg vs. cls at token 40

0

1

2

3

4

regression_square_loss

0.20

0.25

0.30

0.35

classification_error

TF_alg_select

TF_reg

TF_cls

Averaging

3-NN

Setup Our setup is the same as the two 'mixture' modes (linear model + linear classification model, and noisy linear models with two different noise levels) as in Section 6, except that the architecture is GPT-2 following Garg et al. [31], and the input format is changed to (11) (so that the input sequence has 2 N +1 tokens) without positional encodings. For every i ∈ [ N +1] , we extract the prediction ̂ y i using a linear read-out function applied on output token 2 i -1 , and the (learnable) linear read-out function is the same across all tokens, similar as in Appendix M.1. The rest of the setup (optimization, training, and evaluation) is the same as in Section 6 & M.1. Note that we also train on the objective (47) for all tokens averaged, instead of for the last test token as in Section 6.

Result Figure 2 shows the results for noisy linear models with two different noise levels, and Figure 5 shows the results for linear model + linear classification model. We observe that at every token, In both cases, TF_alg_select nearly matches the strongest baseline for both tasks simultaneously, whereas transformers trained on a single task perform suboptimally on the other task. Further, this phenomenon consistently shows up at every token. For example, in Figure 2a & 2b, TF_alg_select matches ridge regression with the optimal λ on all tokens i ∈ { 1 , . . . , N } ( N = 40 ). In Figure 5a & 5b, TF_alg_select matches least squares on the regression task and logistic regression on the classification task on all tokens i ∈ [ N ] . This demonstrates the in-context algorithm selection capabilities of standard decoder-based transformer architectures.

M.3 Computational resource

All our experiments are performed on 8 Nvidia Tesla A100 GPUs (40GB memory). The total GPU time is approximately 5 days (on 8 GPUs), with the largest individual training run taking about a single day on a single GPU.

Peer reviews:
comments: SUMMARY:
1. The authors prove that transformer models can perform several machine learning tasks, such as least squares, ridge regression, lasso and learning of GLMs.
2. The authors prove that transformer models can select an algorithm from a pool of possible ones, in two scenarios, post ICL validation and pre ICL testing, and thus can perform in context learning.
3. The authors construct a transformer model with near Bayes-optimal performance on linear models with mixed noise levels
4. The authors prove polynomial sample complexity results for pretraining transformers to perform ICL.
5. Experimentally, the authors show that transformer models can indeed perform algorithm selection in context.

STRENGTHS:
1. The level of clarity and writing is very good.
2. The theoretical contributions are insightful
3. The experimental results support the theoretical arguments
4. The scope of the work is broad and deep.

WEAKNESSES:
I do not recognize any major weakness

QUESTIONS:
--

LIMITATIONS:
--
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThey take an in-depth empirical and also theoretical analysis of the in-context learning abilities of transformers on various tasks. Their theoretical analysis of the learned algorithms is especially strong. The work in my eyes presents an extension of previous works in this field to more complex ICL problems (like least squares, ridge regression, Lasso). They identify mechanisms which transformers can learn by which a number of ICL tasks can be performed.\n\nSTRENGTHS:\nThe submission is technically sound and very detailed in its explanations and supplements.\nThe work extends previous studies on in-context learning. They provide in-depth and very relevant analyses for a number of simple ICL tasks. I do, however, wonder in which ways this work can be of practical relevance. For language modelling, the tasks and architectural sizes are quite different to these evaluated toy tasks. For learning optimization methods, it would be interesting to provide more general bounds for ICL analysis which could be used to transfer to more relevant tasks. Also, it would be interesting to see if the provided analyses could be used to make modifications to the transformer architecture. I understand the scope of this work is huge already and this work might exactly enable these kinds of follow-ups.\n(Capability of transformers for ICL) In my eyes it seems rather obvious that the demonstrated capabilities can be learned through transformers, given that these architectures are able to model far more complex tasks. \n(Architectural choices for ICL) From a practical perspective the suggestions for optimal choices seem less relevant as of now (e.g. which architectural parameters work well in practice for modelling certain ICL aspects), since architectural choices are studied in isolation and only for toy problems. However, theoretical upper bounds on architectural choices seem interesting and seem to provide interesting ways in which to modify and analyse architectures for enhanced ICL.\n\nWEAKNESSES:\nI do at times wonder if the theoretical derivations on the capabilities of in-context learning do hold up in practice. Take for example Pre-ICL. Authors motivate this with an example where the hyperparamter in a linear regression varies. I would expect the transformer to be able to learn to generalize these parameters, given a number of different hyperparamters and pretraining rounds. This would not be akin to algorithm selection but rather a learned mechanism that can learn across hyperparameters. While the work might just give upper bounds here, I imagine these are quite hard to transfer into practice. \n\nThe capability of transformers to learn various tasks in in-context learning has been known for quite some time. I think some relevant prior work on \"Prior-data fitted networks\" is missing, that demonstrates in-context learning with transformers before the term \"in-context learning\" came up:\nTransformers Can Do Bayesian Inference (ICLR 2022): This work studies if Transformers can in-context learn algorithmic tasks (such as Gaussian Processes and MLPs), to my knowledge the first work to do so (e.g. prior to Garg and Ayurek) and with an eye of the Bayesian approximation learned in ICL. They show e.g. Gaussian Processes and MLPs can be approximated by transformers, demonstrating hyperparameter selection and deriving the Bayesian approximation capabilities of transformers.\nStatistical Foundations of Prior-Data Fitted Networks (ICML 2023)\nTabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second (ICLR 2023): Studies the ability of in-context learning for tabular datasets. It also evaluates task mixtures and shows that transformers are able to model complex algorithmic tasks through ICL.\n\nQUESTIONS:\n(Pretraining rounds for excess risk) the authors provide derivations of multiple tasks in terms of the pretraining rounds needed to train transformers with an excess risk to a given baseline. These results are interesting and demonstrate that for a variety of tasks such bounds can be found. I wonder weather such rules can be made up for more general classes of tasks, such that it would be easier to transfer these results to other tasks? Can polynomial guarantees, e.g. be derived given some computational complexity of the baseline algorithm and task?\n\nDid you study what happens when the pretrained transformer is applied to out-of-distribution data, e.g. linear pretraining applied to sparse linear tasks (algorithmically out of distribution) or when the input data is drawn from another distribiution (e.g. higher variance, colinearity, ..)\n\nAre there limits to the number of algorithmic tasks selected by transformers in the selection algorithm?\n\nLIMITATIONS:\nLimitations are not clearly addressed, those could include: Theoretical analyses are hard to apply to more realtistic and complex real world tasks. The analyses are difficult to perform even for the simple tasks shown in this work, probably for most complex tasks infeasable?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper theoretically investigate the in-context learning ability from the approximation view. The authors derive the error bound of approximating  least squares, ridge regression, and Lasso algorithms with transformer. The corresponding statistical properties, i.e., the convergence rate and Bayes suboptimality, are derived for transformers. In addition, the authors show that transformers are also able to implement the algorithm selection. Simulation results are provided to verify the theoretical findings.\n\nSTRENGTHS:\nThis paper presents solid error bound for transformer to approximate different kinds of algorithms, which are interesting for the approximation theory of neural networks. The corresponding convergence rates for different algorithms are relevant for explaining the in-context learning behavior of transformer. In addition, the approximation results for the algorithm selection unit are noval in the investigation of in-context learning.\n\nWEAKNESSES:\n1. Although this paper provides extensive parameters assignments to show that transformer can implement a series of algorithms, rare experimental results are provided to corroborate these constructions. For example, the authors construct transformers that implement one step gradient descent with several layers and repest this pattern. However, such periodic pattern of network parameters and intermediate results is not shown in the simulation section. Consequently, it becomes challenging to assert that transformers indeed implement the algorithms as stated in the paper.\n\n2.  The authors analyze the attention with ReLu activation function, although this particular activation function is rarely employed in practical designs of attentions. To justify the necessity of analyzing ReLU instead of softmax, it would be beneficial to provide reasoning considering that previous studies, such as [1] and [2], predominantly adhere to the softmax activation function.\n\n3. Given the approximation results in the paper of transformers with ReLu activation and the fact that Multi-Layer Perceptrons (MLP) are universal approximators of continuous functions, it is beneficial to discuss that wether the results in this paper imply that MLP can also implement the algorithms in the paper and consequently possess the in-context learning ability.\n\n4. In Corollary 6, the paper derives the estimation error of approximate ridge regression in Bayesian setting, but such Bayesian result is missing for Lasso. It will be helpful to discussion such absence of  Bayesia analysis for Lasso.\n\n5. The generalization error bound in section 5 needs more discussions. Previous works, such as [3] and [4], have derived generalization error bounds for transformers. The novalty of the generalization bound in this paper needs more discussion.\n\nI am happy to change my evaluation if the authors can answer the above questions.\n\n[1] Von Oswald J, Niklasson E, Randazzo E, et al. Transformers learn in-context by gradient descent[J]. arXiv preprint arXiv:2212.07677, 2022.\n\n[2] Akyürek E, Schuurmans D, Andreas J, et al. What learning algorithm is in-context learning? investigations with linear models[J]. arXiv preprint arXiv:2211.15661, 2022.\n\n[3] Edelman B L, Goel S, Kakade S, et al. Inductive biases and variable creation in self-attention mechanisms[C]//International Conference on Machine Learning. PMLR, 2022: 5793-5831.\n\n[4] Zhang Y, Liu B, Cai Q, et al. An Analysis of Attention via the Lens of Exchangeability and Latent Variable Models[J]. arXiv preprint arXiv:2212.14852, 2022.\n\nQUESTIONS:\nThe questions are listed in the previous section.\n\nLIMITATIONS:\nThe limitation section is missing in this paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
In context Learning, is a setting in which Transformers can learn to perform new tasks when prompted with training and test examples. This work advances the understanding of the capabilities of ICL. In particular, this paper proves that for a broad class of standard machine learning algorithms like ridge regression, lasso etc, transformers can implement these methods. Moreover ICL can perform in-context algorithm selection and select simple algorithms to learn a more complex algorithm. They also show that using their proposed method they can construct a transformer that can perform nearly Bayes-optimal ICL on noisy linear models with mixed noise levels.

STRENGTHS:
* This paper expands upon our understanding of ICL. It shows that transformers can implement a plethora of standard ML tasks and require much mild bounds on the number of layers and heads. 
* The paper extends the analysis to in-context algorithm selection and provides two algorithm selection mechanisms. They use the proposed mechanism to construct a transformer that can perform Bayes optimal ICL on noisy linear models. 
* This works opens up new directions to explore (theoretically and empirically) optimal ICL construction on other problems
* The paper is well written and easy to follow

WEAKNESSES:
* It would be good to know how results such as those in [1] compares to the one obtained in this paper.

[1] Are transformers universal approximators of sequence-to-sequence functions? a

QUESTIONS:
what was the reasoning behind choosing Bayes-optimal ICL on noisy linear models with mixed noise levels as an example to show case the usefulness of in context algorithm selection? Was that the most complex method for which something like theorem 12 could be proven?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06