Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection
Authors: Yu Bai
Fan Chen
Huan Wang
Caiming Xiong
Song Mei
Venue: NeurIPS 2023 oral (2023)
Keywords: in-context learning
transformers
deep learning theory
learning theory
TL;DR: Transformers can learn in context like statisticians --- A single transformer can select different algorithms for different data at hand.

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
Yu Bai
Fan Chen'
Salesforce Research yu. bai www.essllsfrre.commmmmmm com
Massachusetts Institute of Technology
fanchen@mit.edu
Huan Wang
Caiming Xiong
Song Mei UC Berkeley
Salesforce Research
Salesforce Research
huan. wwwa.aaww@@aleesfrree..oo com
cxiong@salesforce..c com
songmeiiberkelley
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 comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement 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 netr-oltimal 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 in-contexi algorithm se- lection, akin what statistician adaptively select different base ICL algorithms or even perform qualitatively different tasks different imput 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 gen- eral mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL. validation. As an example, we use the POst-ICL validation mecha- nism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task task-noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.
do in real life-A single transformer
Introduction
Large neural sequence models have demonstrated remarkable in-context learning (ICL.) capabili- ties [12], where models can make accurate predictions on new tasks when prompted with training examples from the same task, in zero-shot fashion without any parameter update to the model. prevalent example is large language models based the transformer architecture [84], which can perform diverse range of tasks in context when trained on enormous text [12, 90]. Recent models
"Equal technical and directional contributions.
Code is available at https: :// //gi thub. com/ all www.eembtttttaaafrreeeeeee as- statisticians.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Example 1: 1:iidge with different
Mochanism Rost ICL Validation
Example 2: Regression Clasafication
Mechaniam 2: Pre- Texting
1aN Ridgo, (1,, Kaus! Ridge, (0, reus)
Free Ale,, (Pray Apual
Lineeg(D.Roa)
LogRog(D.Xeul
Sane Ale, (2 Soy!
Which smailear??
Ale,
Transfommer
Lossy!t..
Lo
Transforme
Jack
Alg (Das)
Which distribution?
Duta
Data 2
Tra n valiciaioon split
Data (Ragression)
Data 2
(Rog winolso oj (Rog wcholso ei
Duta
(Clazs lication)
Dota
Figure 1: Illustratoon of in-context algorithm selection, and two mechanisms constructed in our theory. Left, middle-left: A single transformer can perform ridge regression with different A's on input sequences with different observation noise; we prove this by the post-ICL validation mechanism (Section 4.1). Middle- right, right: A single transformer can perform linear regression on regression data and logistic regression on classification data; we prove this via the pre-ICL testing mechanism (Section 4.2).
in this paradigm such GPT-4 achieve surprisingly impressive ICL performance that makes them akin to a general-purpont agent in many aspects [65, 14]. Such strong capabilities call for better understandings, which 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 ICI. experiments where input tokens are realevalued (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 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 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 general mechanism-ine in-context algorithm selection- which a single transformer can adaptively select different "base" ICL algorithmss 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 X on ICI. instance 1, and X2 on ICL instance 2 (Figure 2); or perform regression on ICL. instance and classification on ICL instance (Figure 5). This adaptivity allows transformers to achieve much stronger ICL performance than the base ICL algorithms. We both prove this theory, and demonstrate this phenomenon empirically on standard transformer architecture.. (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, 391, and lay out a foundation for the intriguing special case where the learning targets are themselves ICL algorithms.
detailed summary of our contributions 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 in- context gradient descent (Appendix D), which could be broaderly applicable. For broad class of smooth convex empirical risks over the in-context training data, we construct an (L (L+J) 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 d1 (b) Noisy linear reg with noise d2 (c) Task I vs. task 2 at token 20
1.0
1.4
+ sr.ag
Th. 1Puroeic
Tr. sotect
TT. rotee 2
rige lam nior_any
1.2
.... rége Jam
17 nokee
very nobee
rope, Jem.1
TF. raton
0.6
1.0
*.- rage_jem?
0.4
0.6
Rayex ar noles
0.2
0.6
0.0
0.4g
20
30
10
20
40
0.30
0.15
0.20
0.25
0.30
in-conteat examples
i-cornlec examples
norey hw____gg__nn noinnsese
Figure 2: In-context algorithm selection on two separate noisy linear regression tasks with noise (a), 02)) (0.1, 0.5). (a,b) A single transformer TF TF_alg_select simultancously approaches the performance of the two individual Bayes predictors ridge_ ____amm_____ 1 on task and ridge_ ____a___2 task 2. (c) At token 20 (using example (o...., ]9} for training), TF TF_ FF_alg_selll select approaches the Bayes error on two tasks simultancously, and outperforms ridge regression with any fixed A. (a,b.c) Note that transformers pretrained on a single task (TF noise_1. TF __noise___ perform near-optimally on that task but suboptimally on the other task. More details about the setup and training method can he found in Appendix M.2.
(Section 4.2). For both mechanisms, we provide general constructions as well as concrete examples. Figure 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 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 shows that a single transformer can approach the individual Bayes risks (the optimal risk among all possible algorithms) simultaneously on noisy linear models with different noise levels.
Transformers statisticians We humbly remark that the typical toolkit of 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.
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-tor toamiiggto--earn such as meta- leaming. Due to limited space, we discuss these related work in Appendix
Related work
2 Preliminaries
We consider sequence of N input vectors [ C R,,, written compactly as an input matrix H [h,.... hw] E IR DxN N where each h, is column of H (also token). Throughout this paper, we let (1) == (t) = max {1, 0} denote the standard relu activation.
2.1 Transformers
We consider transformer architectures that process any input sequence H E DDDN by applying (encoder-mode-) attention layers and MLP layers formally defined follows.
Definition 1 (Attention layer). A jself-jattention layer with M heads is denoted as Attme()) with parameters 0 = {(Vm. Qm, K,,) ( me[M] C IR. DxD. On any input sequence H E D
H = AU) := H N EM (V H) (((QH)))((K)
E DDN
(1)
where a R + I. is the ReLU function. In vector
h [Attno HH),, = h, M Σ hi (,nkm, Vnhy.
Many of results can be generalized decoder-base based architectures; see Appendix C for discussion.
3

---

Above, (1) uses normalized ReLU activation + att)NN in place of the standard softmax activation; we remark this activation also found work well empirically in recent studies [78, 93]. Definition 2 (MLP layer). (token-wise) MLP layer with hidden dimension D' is denoted as (+) with parameters (W1: W2) E DDD x xRDxD'
On any input sequence H RRD
H = ( ( (W
where a R. R is the ReL. function. In vector form, we have h, = h, W2
We consider transformer architecture with I, 2 transformer layers, each consisting of self- attention layer followed by an MLP layer.
Definition 3 (Transformer). An L-layer transformer, denoted as TFa()), is composition of L self-attention layers each followed by an MLP layer: H(L) -FF (H()), where H()) l2 DxN is the input sequence, and
H MLP
Attnn
H 1))
E fl......,
0 = (0 (0 (1:L)
O (1:b)) consists of the attention layers (0)
Above, the parameter
{(V Qu k! Mem(srto) C IR DxD and the MLP layers o (w()) w!") E R D x
DD D(O) We will frequently consider "attention-only transformers with w( w = 0, which we denote as (.) for shorthand, with (1:L) ((12))
We additionally define the following norm of a transformer TF
|||| := max max EE ME[M]
{ {10Q
(
llops
K
op
+
Σ m=1
V llop
| 1lop 1 Hop
(2)
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;; ())))s P and new test input XN+1 Px for some data distribution P, where (x; S Fed are the input vectors, [y.)eeN S I are the corresponding labels (e.g. real-valued for regression, or {0, 1}-valued for binary classification), and X is the test input on which the model is required to make prediction. Different from standard supervised learning, in ICL, each instance (D, xN+1) )) general drawn from a different distribution Pj. such as linear model with new ground truth coefficient W..j Rd Our goal is to construct fixed transformer to perform ICL on large set of P,'s.
We consider using transformers perform ICL, in which encode (D, (D,xN+1) into an input se- quence H E IR (X1)) In theory, use the following format, where the first two rows contain (D, xx++1 ) (zero at the location for +1+ and the third row contains fixed vectors (pileeN44) with ones, zeros, and indicator for being the train token (similar to a positional encoding vector):
X1
X2
XN
XN+1
0 (D-(4)3) (d+3)
H
1/1
y2
VN
E DX(N11)
E R D-(+1)
(3)
P1
P2
P.V
PN+1
1 1i N
We will choose D (d), so that the hidden dimension of H is at most constant multiple of d. We then feed H into a transformer to obtain the output = TF, (H) E Dx with the same shape, and read out the prediction +1 from the (d + N + 1)-th entry of H = [h.).e[N+ +1] (the entry corresponding to the missing test label): ~N+1 = read, (H) = (hN+1)041- The goal is to predict
Foor each query index i, the attention weights (o.(Q K.. hr))))))))((p) is also set of non-negative weights that sum to (1) (similar as a softmax probability distribution) in typical scenarios. Also, our approxi- mation results potentially be generalized to softmax attention e.g. using the technique of [32].

---

UN+1 that is close to UN+1 ~ P measured by proper losses. We emphasize that we consider predicting only at the last token XN+1, which is without much loss of generality.
Miscellaneous setups We assume bounded features and labels throughout the paper (unless oth- erwise specified, e.g. when x; is Gaussian): ||x,|12 2 Ba and 1ys S B, with probability one. We use the standard notation X [ :...: x x E NRd and y = [y::...; E IR 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 ØN+1 = ready (H) := clip ((N+4))44) where clip == Proji-R.K R.R (l) is the standard clipping operator with (a suitably large) radius R 2 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 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 À 2 0 (reducing to least squares at À = 0 and N 2 d):
ridge arg min weRd 2N Ei=1 (ww,x;) yx ||
(ICRidge)
We show that transformers can approximately implement (ICRidge) (proof in Appendix F.I). Theorem 4 (Implementing in-context ridge regression). For any À 2 0, 0 < S B with k: 2+ Bu > 0, and B. Bw/ there exists an L-layer attention-only transformer TF. with
a+
L = log B2 / (2=)) +1, maXtE[L] M 3, |01 S 4R + 8(3 + A)-
(4)
(with R == max Bw. By 1}) such that the following holds. On any input data (D,xN xN such that the problem (ICRidge) is well-conditoneed and has a bounded solution:
S Amim (X X/N)
(xx/N) S Amax (X (XNN S B, [| S Bww/2.
(5)
TF. approximately implements (ICRidge): The prediction ØN = ((F,, (H) satisfies
DN+1
(wadge
x S.
(6)
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 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 > Ö(d), there exists OK log(k.) KN (o)o/)))luyr transformer O, such that any P satisfying standard statistical assumptions for least squares (Assumption A), its ICL prediction ÿN achieves
E E(DxxN+110N=)==P((00 - 1 < inf E(x.g)~PP (1 (w, (w,x)) x) ((do2/N)
Assumption A requires only generic tail properties such as sub-Gaussianity, and not realizability (i.e., P follows a true linear model); K, above denote the covariance condition number and the
*Our constructions may be generalized to predicting at every token, by using a decoder architecture and potentially different input formats correspondingly (cf. Appendix C). Our theory focuses on predicting at the last token only, which simplifies the setting, Our experiments test both settings.
5

---

noise level therein. The O(do2/N) excess risk is known to be rate-optimal for linear regression [38], and Corollary 5 achieves this in context with transformer with only logarithmically many layers.
Next, consider Bayesian linear models where each in-context data distribution P = Pin, is drawn from a Gaussian prior
N(O, La/d,, and (x,y) v) ~ Pür. is sampled as x ~ N(O, y (W., x) + N(O, o²). It is standard result that the Bayes estimator of yN+1 given (D,xx+) xx+1) is given by ridge regression (ICRidge): VN+1 == (w ridge XN+1) with À do? /N. We show that transformers achieve nearly-Bayes risk for this problem, and we use
BayesRisk, (w,xx 1 12N11))PP.
(are WN+1)
to denote the Bayes risk of this problem under prior m.
Corollary 6 (Nearly-Bayes linear regression with transformers by approximating ridge regression). Under the Bayesian linear model above with N 2 max {d/ O (log(1/=))), there exists a L = O (1og(1/65 (1g) /e))-layerr transformer such that Ew,. Www.(P.xwi.... 2 (xN+1 UN+1)² - BayesRisk, t..
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 l--reullarzeed linear regression loss Lassso over the in-context training examples D;
1 Wiasso == arg min T (w) = 2N Σ ((w, x.) - yi)² 3= AN ||//1
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 (Implementing in-context Lasso). For any AN 2 0, 3 > B > O, and > 0, there exists I--layer transformer TFa with
== [ /c] +1, maxte(L] M() 2, maxte(L] D 2d, |10| S O + (1 AN)
(where R == max {B, Bw, By, 1}) such that the following holds. On any input data (D, xN+1) such that Amax [ X/N)= < 3 and [[aasss I / TFe approximately implements ( (ICLasso), in that oulputs + = (XN+1, W) with Lasss W) LLasso (Winsso) S E.
Theorem 8 (Near-optimal sparse linear regression with transformers by approximating Lasso). For )² (1 (d/N))) transformer
any d, N Z 1,5 > O, B > O, there exists such that the following holds: For any 8 and N 2 (s (ogodd/))), suppose that P is an 8-sparse linear model: ~ N(O, Ia), y/ {Ww, x:) N(O, 2) for any (w=ll2 S B and w S then (over the randomness DD), the transformer output UN achieves
with probability at least
E Ex (X [ØN+1 - [
) S o²lL + O(s log(d/6)/N]]
The Oog d/N) excess risk obtained in Theorem 8 optimal up to log factors [62, 87]. We remark that Theorem 8 is not direct corollary of Theorem 7; Rather, the bound the number of layers in Theorem 8 requires a sharper convergence analysis of the (ICLasso) problem under sparse lincar models (Appendix H.2), similar to [1].
3.3 Proof technique: In-context gradient descent
The constructions in Section 3.1 and 3.2 built on the following result for approximating in-context (proximal) gradient descent on (regularized) convex losses.
Theorem (ICGD; Informal version of Theorem D.1 & D.2). For a broad class of convex losses of form w + & Eie f(w" x;, y:) + R(w), there exists an L-layer transformer that takes in any (D. w and outputs wh such that | w (GD,, l|2 S O(Le), by composing L identical layers each O(E) aaprrxxaiii single step of GD (so that ((Le) is linear error accumulation).

---

Theorem 9 is established in two main steps:
A Approximating one-step of ICGD using 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 al. [2] and Giannou et al. [32].
Stacking L of the above layer to approximate steps of ICGD. Done naively, the error accumu- lation of this stacking operation is exponential in L in the worst case. We utilize the stability of convex gradient descent (Lemma D.I) to obtain the linear in L error accumulation in Theorem 9,
In Appendix D.3, we also give results for non non-convex GD two-layer neural nets, though with worse (exponential in L) error accumulation as expected.
In-context algorithm selection
We now show that transformers can perform various kinds of in-context algorithm selection, which allows them 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 (Drrain, D.ai). and running K base ICL algorithms Derain Let [fk]ke[K] C C(R. - R) denote the K learned predictors, and
Lval( == Piit Exx.n ((f(x)) 9:)
denote the validation loss of any predictor f.
We show that (proof in Appendix I.1) a 3-layer transformer can output predictor f that achieves nearly the smallest validation loss, and thus nearly optimal expected loss if Lval concentrates around the expected loss L. Below, the input sequence H uses a generalized positional encoding p; [0 (d+-a)+)) 1; in (3), where t; == 1 for i E Derain t; == -1 for E Dval, and tN+1* := 0.
Proposition 10 (In-context algorithm selection via valid-valiattion split). Suppose that E(..-) in (7)
is approximable by sum of relus (Definition D.1, which includes all C smoooh bivariate functions).
Then there exists 3-layer transformer TF, that maps (defining yf = y.ili < N + 13)
h = (x * fi(x)) ... 5x((x,)) 0k+15 1:4] + =x=[x; f 1;4 iE [N + 11,
where the predictor J I IR a convex combination of Usx Lva((fa) S mink, E|K] (va))
7}. As corollary, for any convex risk L (Rd - IR) + R. f satisfies
L(f) < minn..ekk L(fa.) + max [val(fa) -((x) +7
Ridge regression with in-context regularization selection As example, use Proposition 10 to construct a transformer to perform in-context ridge regression with regularization selection ac- cording to the unregularized validation loss (w) TTP.a (x..,)) ((w, - v1) (proof
in Appendix I.2). Let A AK 0 be K fixed regularization strengths.
Theorem 11 (Ridge regression with in-context regularization selection). There exists transformer with (log(1/e) layers and O(K) heads such that the following holds: On any (D, (Dxx+1) well- conditioned (cf. for all {Ak ( it outputs ØN+1 (W,NN+)), where
dist wr conv W Ar traie Lyail Si
rrdge.trainn < ming, E[K] Lvail (W ridge, ridge.train) + ) S S..
de.
Above, W.
trdge,train denotes the solution to (ICRidge) the training split Derain-
7

---

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 mixture of K different noise levels ......: OK > O.
Concretely, consider the following data generating model, where we first sample = ow,,,, from A E A([K]), w. ~ N(O, Ia/d), and then sample data {(x;, ys) b))
iid P Pw..or as
P Pw..ax X; N(O, Ia); 1/1 (x;, w.) + N(0, of)
For any fixed (N, d), consider the Bayes risk for predicting #N+1 under this model:
BayesRisk_ == inf A E, (xA(())(xx)) - 1
By standard Bayesian calculations, the above Bayes risk is attained when A is a certain mixture of K ridge regressions with regularization A do LN; 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, construct transformer that achieves nearly the Bayes risk.
Theorem 12 (Nearly Bayes-optimal ICI, Informal version of Theorem J.1). For sufficiently large N, d, there exists transformer with (log N) layers and (K) heads such that on the above model, it outputs a prediction +1 that is nearly Bayes-optimal:
E -
)N+1)
BayesRisk, + +0 ((log K/N)//3).
(8)
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 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 ~ A 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'.. 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 transformer that can learn 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-ICI. testing, the transformer runs 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 co-context regression and in-context classification, by running the following binary rype check on the input labels {ı4 {y]]eeN])
y E {0, 1}, y $ 8- -
y binary (D) = N Σ (n.),
((x)
0.
linear interpolation,
otherwise.
'y the Bayes risk decomposition for square loss, (8) implies that E [ ) 2 < O(llo 133)
8

---

Lemma 13. There exists single attention layer with 6 heads that implements Tibinary 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 continuous distribution.
Proposition 14 (Adaptive regression or classification; Informal version of Proposition 1.4). There exists transformer with (glog(1/e) layers such that the following holds: On any DD such that ye E {0, 1), it outputs UN+1 that e-appreximates the prediction of in-context logistic regression.
By contrast, for any distribution whose marginal distribution of y is not concentrated around {0, 1}. with high probability (over D), UN+1 e-approximates the prediction of in-context least squares.
The proof's can found in Appendix I.3. We additionally show that transformers can implement more complex tests such as 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 Lia(-) over n pretraining sequences. Let (-d(-) denote the corresponding population risk.
Theorem 15 (Generalization of transformers; Informal version of Theorem K.1). The solution θ to (TF-ERM) over transformers with T. layers, M heads per layer, and hidden dimension D' satisfies
(ia S inf Lia (0) + LP(MDD + DDD)
(V
71
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 The- orem K.2). Under Assumption A and N 2 O(d), the solution 0 to (TF-ERM) with L = ((wlog(nN/o)) layers, M heads, D' (attention-only as in Theorem 4) achieves small excess ICL risk over the best linear predictor wp == Ep [xx]] -Pppppx for each P:
(( (wp. x))}
so(v
R2d
do
Lia(0) - -E ((x))
S
See Appendix K.2 for similar results in several additional settings.
Experiments
We test our theory by studying the ICI. and in-context algorithm selection capabilities of transformers. using the sacoder-based architecture in our theoretical constructions (Definition 3). Due to limited space, additional experimental details can be found in Appendix M. Results with a decoder architecture as in [31, 47) (including the setup of Figure 2) can found 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 following base distributions (tasks), where we first iid Pw. by sampling w, N(O, L,/d), and then sample ((x;, v:)) [e([N+) Pw. as
sample P
N(0, I2), and from one of the following models studied in Section 3: iid
x;
1. Linear model: y x));
2. Noisy linear model: y1 (w.,x,) + o24 where > 0 is fixed noise level, and 2 N(O, 1).

---

(a) Base ICL capabilities
Noisy reg with two noises
(c) Reg + Classification
Lean Savem
nosy. n
3-M
6,50
ede len
18g Sen
ndny, pose
Leaoo Jam
0.23
sparst
LAIN Jam.JJam-0000
0.
0.00 0.25
0.50
0.35 1.00 125 1.50
0.0 0.10 0.15
6.30
0:30
0.5
Las
Figure 3: ICI. capabilities of the transformer architecture used in our theoretical constructions. (a) On five representative hase tasks, transformers approximately match the best baseline algorithm for each task, when pretrained on the corresponding task. (b,c) single transformer TF ___aa_g_____ select simultaneously approaches the performance of the strongest baseline algorithm on two separate tasks: (b) noisy linear regression with two different noise levels a E (0.1, and (c) adaptively selecting between regression and classification.
3. Sparse linear model: y/ = (W.: x;) with ||w S 8, where < is fixed sparsity level, and in this case we sample w. from special prior supported on s-sparse vectors;
4. Linear classification model: y/- sign((w.,
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-bass architecture instead of baseder-based, and ReIUU activation instead of softmax. All experiments use d = 20. We choose E (2) {0.1, 0.5) and N = 20 for noisy linear regression, 8 = 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 a E 0.25, 0.5,
Transformers trained with the mixture mode will be evaluated on multiple base distributions simulta- neously. 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 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, 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_alge 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 the other task as expected. The existence of TF_ada sa__aalgeeeell select showcases a single transformer that performs well on multiple tasks simultaneously (and thus has 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, architectures
this work focuses on the transformer architecture; alternative (such RNNs) are beyond our scope but would be interesting directions for future work.
10

---

Acknowledgment
The authors would like to thank Tengyu Ma and Jason D. Lee for the many insightful discussions. and NSF CCF-2315725.
Mei is supported in part by NSF
References
[1] 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.
[2] 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.1560 2022.
[3] J. L. J., R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv: 2016.
[4] F. Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-611, 2017.
[5] Y. Bai, M. Chen, P. Zhou, T. Zhao, 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.
[6] Y. Bai, S. Mei, H. Wang, and C. Xiong. Don't just blame perrparraettizztioon for over- confidence: Theoretical analysis of calibration binary classification. International Conference on Machine Learning, pages 566-576. PMLR, 2021.
[7] J. Baxter. model of inductive bias learning. Journal of artificial intelligence research, 12: 149-198, 2000.
[8] A. Beck and M. Teboulle. Gradient-based algorithms with applications to signal recovery. Convex optimization in signal processing and communications, pages 42 88, 2009.
[9] S. Bengio, Y. Bengio, J. Cloutier, and J. Gescei. On the optimization of synaptic learning rule. In Optimality in Biological and Artificial Networks?, pages 281-303, Routledge, 2013.
[10] S. Bhattamishra, K. Ahuja, and N. Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009. 11264, 2020.
[11] S. Bhattamishra, A. Patel, and N. Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.00 09286. 2020.
[12] 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:187 1877-19999, 2020.
[13] S. Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends0 in Machine Learning, (33-4) 2015.
[14] S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.. 127 2023.
[15] 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-1889 2022.
[16] 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, 1e.15584-15999 2021.

---

[17] 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.
[18] D. Dai, Sun, L. Dong, Y. Hao, Z. Sui, and F. Wei. Why can gpt learn in-context? language models secretly perform gradient descent meta optimizers. arXiv preprini arXiv:2212.10559 2022.
[19] G. Denevi, C. Ciliberto, D. Stamos, and M. Pontil. Incremental learning-to ta--g--tt--aamm with statistical guarantees. arXiv preprint Akkv:/8803.0800 2018.
[20] G. Denevi, C. Ciliberto, D. Stamos, and M. Pontil. Learning to learn around common mean. Advances in Neural Information Processing Systems, 31, 2018.
[21] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1810.0081004800 2018.
[22] E. Dobriban and S. Wager. High-dimensional asymptotics of prediction: Ridge regression and classification. The Annals of Statistics, 46(1):247-279, 2018
[23] L. Dong, S. Xu, and B. Xu. Speech-transformer: a no-recurrence squence-to-sequene model for speech recognition. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), pages 5884-5888. IEEE, 2018.
[24] 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 W.Xiv:2301.00234 2022.
[25] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. De- hghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11923 2020.
[26] S. S. Du, W. Hu, M. Kakade, D. Lee, and Q. Lei. Few-shot learning via learning the representation, provably. arXiv preprint arXiv:2002.09434 2020.
[27] B. L. Edelman, S. Goel, S. Kakade, and Zhang. Inductive biases and variable creation in self- attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[28] N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Bai, A. Chen, T. Conerly, al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
[29] C. Finn, P. Abbeel, and S. Levine. Model-agnostie meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1 1126-1133: PMLR, 2017.
[30] C. Finn, A. Rajeswaran, S. Kakade, and S. Levine. Online meta-lcarning In International Conference on Machine Learning, pages 1920-1930. PMLR, 2019.
[31] S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30555 2022.
[32] A. Giannou, S. Rajput, J.-y, Sohn, K. Lee, D. Lee, and D. Papailiopoulos. Looped trans- formers as programmable computers. arXiv preprint arXiv:2301. 2023.
[33] M. Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 1715 2020.
[34] S. Hochreiter, A. S. Younger, and P. R. Conwell. Learning to learn using gradient descent. In 2001: International Conference Vienna, Austria, August
Artificial Neural Networks 21-25, 2001 Proceedings 11, pages 87-94, Springer, 2001.
[35] Hollmann, S. Müller, K. Eggensperger, and F. Hutter. Tabpfn: transformer that solves small tabular classification problems in a second. arXiv preprint S20i:2207.00844 2022.
12

---

[36] 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):5144-5110 2021.
[37] J. Hron, Bahri, Sohl-Dickstein, and R. Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning, pages PMLR, 2020.
[38] D. Hsu, S. M. Kakade, and T. Zhang. Random design analysis of ridge regression. In Conference learning theory, pages 9-1, JMLR Workshop and Conference Proceedings, 2012.
[39] S. Jelassi, M. E. Sander, and Y. Li. Vision transformers provably learn spatial structure. arXiv preprint arXiv:2210.09222,, 2022.
[40] K. J. D. Lee, Y. Liang, and H. V. Poor. Convergence of meta-learning with task-specific adaptation over partial parameters. Advances Neurul Information Processing Systems, 33: 11490-11500, 2020.
[41] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. vZídek, A. Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):833-88, 2021.
[42] 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.
[43] M. Khodak, M.-F. F. Balcan, and A. S. Talwalkar. Adaptive gradient-based meta-learning methods. Advances in Neural Information Processing Systems, 32, 2019.
[44] L. Kirsch and J. Schmidhuber. Meta learning backpropagation and improving it. Advances in Neural Information Processing Systems, 11144222-1114,, 2021.
[45] L. Kirsch, J. Harrison, J. Sohl-Dickstein, and L. Metz. General-purposs in-context learning by meta-learning transformers. arXiv preprint urXiv:22/2.04445, 2022.
[46] K. Li and J. Malik. Learning to optimize. arXiv preprint arXiv: 2016.
[47] Y. Li, M. E. Ildiz, D. Papailiopoulos, and Oymak. Transformers as algorithms: General- ization and implicit model selection in in-context learning. arXiv preprint 2xXi::3001.0006,, 2023.
[48] B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.107499 2022.
[49] J. Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen. What makes good in-context examples for gpt-3? arXiv preprint arXiv:2101.06800 2021.
[50] Lu, M. Bartolo, Moore, S. Riedel, and P. Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2104.088 08786, 2021.
[51] 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. 034497, 2020.
[52] A. Maurer, M. Pontil, and B. Romera-Paredes. The benefit of multitask representatiin learning. Journal of Machine Learning Research, 17(81):1-32, 2016.
[53] P. McCullagh. Generalized linear models. Routledge, 2019.
[54] S. Mei, Bai, and A. Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics, 46(6A):274727744 2018.
13

---

[55] Min, M. Lewis, H. Hajishirzi, and L. Zettlemoyer. Noisy channel language model prompting for feo-shot text classification. arXiv preprint arXiv:2108.04100 04066, 2021.
[56] S. Min, M. Lewis, L. Zettlemoyer, and H. Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv:2110.15943, 2021.
[57] 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 preprini arXiv:2202. 12837, 2022.
[58] N. Mishra, M. Rohaninejad, X. Chen, and P. Abbeel. A simple neural attentive meta- leaa-eearmer. arXiv preprint 170.v11707.03141, 2017.
[59] S. Müller, N. Hollmann, S. P. Arango, J. Grabocka, and F. Hutter. Transformers do bayesian inference. arXiv preprint arXiv:2112.10510, 2021.
[60] T. Nagler. Statistical foundations of prior-data fitted networks. arXiv preprint arXiv:2305. 11097, 2023.
[61] D. K. Naik and R. 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.
[62] S. N. Negahban, Ravikumar, M. Wainwright, and B. Yu. unified framework for high- -ineesiimenional analysis of m-estimators with decomposable regularizers. 2012.
[63] Y. Nesterov. Lectures on convex optimization, volume 137. Springer. 2018.
[64] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, Askell, Y. Bai, A. Chen, al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1199, 2022.
[65] OpenAL. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[66] N. Parikh, S. Boyd, et al. Proximal algorithms. Foundations and trends@ in Optimization, (3):127-239, 2014.
[67] J. Pérez, Marinković, and P. Barceló. On the turing completeness of modern neural network architectures. arXiv preprint arXiv: 1901.03290 2019.
[68] A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, al. Improving language understanding by generative pre-training. 2018.
[69] A. Radford, 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 supervi- sion. In International conference on machine learning, pages 8748-8763, PMLR, 2021.
[70] 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.
[71] S. Ravi and H. Larochelle. Optimization model for few-shot learning. In International conference on learning representations, 2017.
[72] Y. Razeghi, R. L. Logan IV, M. Gardner, and S. Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXivv2202.07200, 2022.
[73] 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.
[74] Rubin, J. Herzig, and Berant. Learning to retrieve prompts for in-context learning. arXiv preprint 2rXiv:2112..88 2021.
14

---

[75] A. Santoro, S. Bartunov, M. Botvinick, D. Wierstra, and T. Lillicrap. Meta-learnin with memory-augmentte neural networks, In International conference on machine learning, pages 1842-1850. PMLR, 2016.
[76] N. Saunshi, A. Cupta, and W. Hu, A representation learning perspective the importance of train-validation splitting in meta-learning. In International Conference on Machine Learning, pages 9333-9343. PMLR, 2021.
[77] J. Schmidhuber. Evolutionary principles in self-referential learning, on learning how to learn: the meta-meta-... hook. PhD thesis, Technische Universität München, 1987.
[78] K. Shen, Guo, X. Tan, Tang, R. Wang, and J. Bian. A study on relu and softmax in transformer. arXiv preprint 2Xiv:2302.064611 2023.
[79] C. Snell, R. Zhong, D. Klein, and J. Steinhardt. Approximating how single head attention learns. arXiv preprint ar 103.07601, 2021.
[80] J. Snell, K. Swersky, and R. Zemel. Prototypical networks for few-shot learning. Advances in neural information processing systems, 30, 2017.
[81] S. Thrun and L.. Pratt. Learning to learn. Springer Science & Business Media, 2012.
[82] R. Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society: Series B (Methodological), 58(1):6672288 1996.
[83] 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.
[84] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, Jones, A. N. Gomez, E.. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[85] Vershynin. High-dimensional probabiltyy: An introduction with applications data science, volume 47. Cambridge university press, 2018.
[86] 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.06677, 2022.
[87] M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge university press, 2019.
[88] X. Wang, S. Yuan, C. Wu, and R. Ge. Guarantees for tuning the step size using a learning- to-learn approach. In International Conference Machine Learning, pages 10981-10990) PMLR, 2021.
[89] C. Wei, Y. Chen, and T. Ma. Statistically meaningful approximation: a case study on approxi- mating turing machines with transformers. arXiv preprint 1rx.v:2107.13:66 2021.
[90] 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 2022.
arXiv:2206.0766
[91] J. Wei. Wei, Tay, D. Tran, A. Webson, Lu, X. Chen, H. Liu, D. Huang. D. Zhou, et al. Larger language models do in-context learning differently. arXiv preprint 2rXiiv2303.038.4 2023.
[92] G. Weiss, Y. Goldberg, and Yahav. Thinking like transformers. International Conference on Machine Learning, pages 11080-11090. PMLR, 2021.
[93] M. Wortsman, J. Lee, J. Gilmer, and S. Kornblith. Replacing softmax with relu in vision transformers. arXiv preprint w.Xiv:2309.08588 2023.
[94] S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint 2rXiv:2111.02080. 2021.
15

---

[95] S. Yao, B. Peng. C. Papadimitriou, and K. Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115. 2021.
[96] 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.
[97] C. Yun, S. Bhojanapalli, A. S. Rawat, J. Reddi, and S. Kumar. Are transformers universal approximators of to-suueene-- functions? arXiv preprint arXiv: 2019.
[98] Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and T. Wagner. Unveiling trans- formers with lego: synthetic reasoning task. arXiv preprint arXiv:2206. 04301, 2022.
[99] 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.
[100] 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.
[101] X. Zuo, Z. Chen, H. Yao, Y. Cao, and Q. Gu. Understanding train-validation split in meta- learning with neural networks. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: pettps:/oopeee peeperreview.net/torrree 1dJJI1ytHEEOOO
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.
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. 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 thcoretical 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- -fomm ridge regression; it also presented preliminary evidence that Icarned transformers perform ICL similar to Bayes-optimal ridge regression. Our work builds upon and substantially extends this line of work by (1) providing 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 a theoretical lower bound for the risk of any possible ICL algorithm, including transformers. Xie et al. [94], Akyürek et al. [2] observe
16

---

that learned transformers behave closely to the Bayes predictor on variety of tasks such as hidden Markov models [94] and noisy linear regression with fixed noise level [2. 47]. Using the in-context algorithm selection mechanism (more precisely the post-ICI. validation mechanism), we show that transformers can perform nearly-Bayes optimal ICI. 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 (normalized) ReLU function. Transformers with ReLU activations is experimentally studied in the recent work of Shen 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 177, 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 downstream learner [9, 46), learning parameter initializations that quickly adapt to downstream tasks [29, 711, 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 query input and produce the corresponding output [34, 58, 75, 44]. For a comprehensive overview, see the survey [36].
Theoretical aspects of meta- veca-learning have received significant recent interest [7, 52, 26, 83, 19, 30, 43, 40, 88, 20, 5, 76, 17, 101]. particular, [52, 26, 83) analyzed the benefit of multi-task learning through representation learning perspective, and [88, 20, 5, 76, 101] studied the statistical properties of learning the parameter initialization for downstream tasks.
Techniques We build 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] 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 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 chaining argument [87].
B Technical tools
Additional notation for proofs We say a random variable X is o'ssub-Gaussian (or SG( is interchangeably) if (eexp(X2/oo)) < 2. A random vector x E IR is P-suu--Gaussiee if (v.x)
for all ||v||2 = 1. A random variable X is K-suub-Exponentall (or SE(K)
interchangeably) if [[cxp([X]] K)) < 2.
B.1 Concentration inequalities
Lemma B.1. Let 3 ~ N(O, Ia/d). Then we have
P (1 2 (1 +1++00 (8+) Se -d02//1
17

---

Lemma B.2 (Theorem 6.1 of [87]). Let X = [X,]] Raxd be a Gaussian random matrix with
Xi N(O, 1). Let min and m (X) be the minimum and maximum singular value of X, respectively. Then we have
max (x)/ 1
Vd/n+ Se-nó112
P(min
d/n 8) S n2
The following lemma is a standard result of covariance concentration, see e.g. 185, Theorem 4.6. I. Lemma B.3. Suppose that X1.". .NN are independent d-dimensional Sassuu-Gaussiee random vectors. Then as long N 2 Cnd, with probability at least 1 - exp(-N/Co) we have
N
1 N
X
S SK2
1=1
where Co is universal constant.
Lemma B.A. For random matrix [mig]] E Nxdd with Viey nd N(O, 1) and [8i] E R. with idd N(O, )²), holds that
&
PP ((| E /8No 2 exp(-N/22.
Proof. We consider uy [migle RN,, then || random variables .... (Ua, E) are independent N(0, [lel and hence
el|
max (u,, Notice that the
(- t
P
max (u,, a) 2 2(e) sEld
< 2dexp
2
Further, by Lemma B.1, P([ee12 20VN)= s exp(-NN/2 Taking 1 v 18No 18No=log(22/0 completes the proof.
B.2
Approximation theory
For any signed measure μ over space W. let TV)) == Sw Idu(w)] E [0, oo] denote its total measure. Recall o(-) ReLU(.) is the standard relu activation, and B (R) = [-R, Rt denotes the standard foo ball in # with radius R > O.
Definition B. (Sufficiently smooth k-variable function). We say a function R&* - RR is (R, C²)- smooth, for = [(k 1)/21 + 2, g is Cs function on B (R), and
sup
V g() 10=0
sup ZEB (R) J.EK)
max 10,, S L,
ZEB (R)
for all i E 1......... s). with maxesiss L,R' < CC.
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 func- tion g Rk - I is (R, Ce) smooth. Then there exists signed measure P over W
{w pk+1
||w]] 1} such that
g(x) = J. oo [x; RD)du(w,
Vx
and (p) C(k)C2, where C(k) < co is constant that only depends k. Lemma B.6 (Uniform finite-neuron approximation). Let X' be a space equipped with a distance function dx(;)) x X Rzo Suppose function g X I. is given by
g(x)= = S. d(x; w) w)),(w),
18

---

where o(::): X x W + ii is L-Lipschitz (in dx) in the first argument, and is a signed measure over W with finite total measure (p) cc. Then for any E O, there exists 01,"" .KK E [+1}, WI,".", Wx E W with K = O (X dx, IAI) / such that
A
sup x
g(x)
4 w:) SE
K
where (X, dx, 3AL
denotes the 3AL
covering number of X in dx.
Proof. Let a(w) == sign(d((nn))) E {+1} denote the sign of the density dp(w). We have
(xx) = =A / a(w)o |d((w]
(9)
A
Note that [du(w)l/A is the density of a probability distribution over W. Thus for any x E X, lid
as long as K 2 O(A B² 1og(1/8)/E") we can sample W1; WK Idy(-))/A, and obtain by Hoeffding's inequality that with probability least 1 - 5,
K
g(x)
(aws)o((;;;)) wi) Se.
K
Let N (AAL) = NN dx, SAI) for shorthand. By union bound, as long as K 2 O (VV( (/))) set corresponding to No( (ZĀZ))
we have with probability at least ô that for every & in the covering
9(x)
A K
Σ
a(w:)o({2: w:) SE/3.
Taking 5 1/2 (for which K O(A²B² log N 3AI )/²)), by the probabilistic method, there exists deterministic set [Wihie[k] C W and {a; := al (o(w:)) ((wi)lie(k) E {+1} such that the above holds.
Next, note that both (by (9)) and the function x + Therefore, for any x E X, taking x to be the point in the covereing set with dx(xx, &) I 3.I. we have
Σ (lw,)))) ))(;)x w.) are (AL)-Lipschitz
K
A
9(x)
a ((W)))k
Jo(x:
K
1=1
A k
(1((x)
((x) g(x)
A
Σ
a(w,)o(X:1)
+
A K
Σ
a(w
(
Σ ((w)))(x;w.)
K
K
i=1
AL
3AL + AL.
3AL
This proves the lemma.
Proposition B.1 (Approximating smooth k-variable functions). For any Eapprox
0, R 1,
Ct > O, have the following: Any (R, Oz) Smoohh function (Definition B.1) g is (Capprox: R M, -)-pproxximable by sum of relus (Definition D.I) with M & C}
R* + R
log(1 +
Cs/ St//Epprox)// 2/2pp))xx and < C(k)Cr, where C(k) > a constant that depends only on k. In other words, there exists
M
f(z)
Σ
Cm (am with
|mml S C,
max ||mmly 1 S 1,
mE [M]
such that supze[- RRR]* [f(z) - g(2)) < Eappeox-
19

---

Proof. As function a (R) - R is (R, Ct) -mmooth, we can apply Lemma B.5 to obtain that there exists signed measure over W == {we E R
||w
S 1) such that
g(z) = S/ 1 ( z ddw) R
Vz E (-
and A TV())= & C(k)C2 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 : [t] E X == [-R, R2 x {R). Then, the function e(x; w) := 7 x) = a [z: Rt) is bounded by B and (1/ L/is-hipsc in x (in the standard fo -distance) Further, we have
lo,.V(XX, ||--- - o1.00 3A/R S OO log(1 + A/Eapprox)). We can thus apply Lemma B.6 to obtain
Sapprox
that, for
M = log(1 + A/ //Sapprox)/uppprl upprox) C(k)C2 log(1 + + Esp/sappox) approxx
there exists {am (am)mem mE(M] C {+1} and W- (Wmlme(M) C W {w RR++ llonww = 1} such that
sup R
[g(z) fa.w(z)) s Eapprox:
where (recalling z - [s;
A
Aam
fa.w(z)=
Σ amd
Wm 22 R
Wm, Wm [2; 11).
M
R
M
m=1
co
a
Note that we have ΣMe result.
Cml A S C(k)Ce, and ||amll S
Wmlll This is the desired
B.3
Optimization
The following convergence result for minimizing 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 RR R is a-strongly convex and 3-smooth for some < a I B. Then, the gradient descent iterates WGD == W L/W with learning rate 1) 1/B and initialization w&D E Rdd satisfies for
.++
any 2
WWO w w'11 S exp (-t/k) //// - -w-ll2- B
I(w(D) L(w") S 2 exp (-t/ks) ||w&D w w
where x :== B/o the condition number of L, and w* == arg min L(w) is the minimizer of L.
The following convergence result of proximal gradient descent (PGD) on convex composite mini- mization problem is also standard, see e.g. [8].
Proposition B.3 (Proximal gradient descent for convex function), Suppose L = f + h, f Rd + R convex and 8-smooth for some 3 > 0, h Rd R. is simple convex function. Then, the proximal gradient descent iterates WPGD == prox,h (WPGD V fV with learning rate = 1/B
[+1
and initialization w E IR satisfies the following for any t > 1:
1. (L(WeGD)) is a decreasing sequence.
2. For any minimizer w" arg margmin weRd L(w).
T( (DD T.(w*)==+ 3
(IWPGD -
|/wttb
w w*1l2).
and hence { {IlwbG - w(]] is also a decreasing sequence.
3. For k 2 1,62 0, it holds that
L(w L(w*) <
2k | w=
20

---

B.4 Uniform convergence
The following result shown [87, Section 5.6].
Theorem B.1. Suppose that to [0, +00) - [0, +06) is convex, non-decreasing function that
satisfies (f + v) 2 (x))(())) For any random variable X, we consider the Orlicz norm induced by y: |X = inf {K > O: ( 1.
Suppose that {Xala is zero-mean random process indexed by such that || p(B, () for some metric p on the space O Then it holds that
X 11p. S
X S (B(J+0)
sup 0,0'00
X
((/DD)
A2 0,
where D the diameter of the metric space (O, p). and the generalized Dudley entropy integral is given by
J J== . (N(6
where N(0; P) is the 5-covering number of (9. p)-
As corollary of Theorem B.1, we have the following result. Proposition B.4 (Uniform concentration bound by chaining). Suppose that (Xoloce is a zero-mean random process given by
N Xo = N Σ f(z;0) [x];=) -B)(:
where 21.-.. AN are i.i.d samples from distribution F& such that the following assumption holds:
(a) The index set θ is equipped with distance p and diameter D. Further, assume that for some constant A, for any ball @' of radius in θ the covering number admits upper bound log (5;; O' p) & dd drlog(2Ar (5) for all 0 < 8 2r.
(b) For any fixed 0 E and sampled from P 2, the random variable f(z;0) SG(BB)-sull Gaussian random variable.
(c) For any 0,0' E ④ and sampled from P.. the random variable f(z; f(z;0') is pGO
random variable.
Then with probability at least 1 - 8, it holds that
V |loog(AA)) + log(1/0)
sup Xoo S C... 8
N
where C a universal constant, and we denote x = 11 =
Furthermore, if we replace the SG in assumption (b) and (c) by SE, then with probability at least 1 5, it holds that
dlog(2Ak) log(1/8) dlog(2Am) log( (og((11) (1/8)
sup X S CBD 6ee
N
Proof. Fix Do () be specified later. We pick (Do/2) (///22)-ovvrrngggg O2 of so that log (0pl < dlog(2AD/Doi Then, by the standard uniform covering of independent sub-Gaussian random variables, we have with probability at least 8/2,
sup Xol S CB |log(2AD/ Do) + log(2/8) 0-00
N
21

---

Assume that O, = {81. .. Do in p). Then e has diameter Do and admits covering number bound log (( ;))))))))) - d dloog 66). Hence, we can apply Theorem B.1 with the process (Xotoee, then
} ). For each j E [n], we consider 2 is the ball centered at 8, of radius
Ba 1 Xo S VN P(0,0)
= 1/2
and simple calculation yields
(Valostt / N
+t eex((-N Vt 0. ))
IP
sup 0,
Ixo - Xo X CCB''D
Therefore, we can let L S /log(2n/o)/N in the above inequality and taking the union bound over j E [n]] and hence with probability least 1 6/2, it holds that for all [n].
sup X X Nol CCB' 2d 2o (2g((AD/ Do) (4/(4 0,
N
Notice that for each E θ there exists j E [ri] such that 0 E 9, and hence
[xol s |eee + | - Xo,1
Thus, with probability at least - 5, it holds
Xe Xo,1 x C - B²Do) dlog(2AD/Do) + log(2/6)
sup dee
Xel I sup [Xel + sup sup see,
BE
N
Taking Do D/K 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:
Nt
( exp
1.
++1
Then Bernstein's 's inequality of SE random variables yields
||xo No fww I CoB C p(0
for universal constant Co. Therefore, probability least 1 - 8, it holds
repeat the argument above deduce that with
dlog((AD//oo) + log(2/6)
d dlog(2AD/D) log((/6)
sup [aol s CB + B' B') lEe
N
Taking Dg 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 P1 : Do-D-D) D xN Pa RDDo+D.)xN - RD R. are two sequence-to-se se-quuenee-eeee functionss that are implemented by sangle-layer transformers, i.e. there exists 0,, O2 such that
TF :H
h
(0)
E E (Do+Di)N
())
h
((
ISKSN
(0)
())))
TF 0, .H.
R (Do+DD))
((22)
(2(H))
1SISN
22

---

Then, there exists O such that for H' that takes form h [h, ;h 1, with h (0) E ERDo h,
R h E P2, we have
00)
H (
TF6 HH'
RRDDODD+DDAxN
(2
ISKEN
Further, 0 has at most M < MM + M2 heads, D' S D2 D2 hidden dimension in its MLP layer, and norm bound [0] |10:1 + 102ll.
Proposition B.6 (Joining parallel multi-layer transformers). Suppose that P1 Di)x+ IR. XN P2 RDDo+D.)xN + D:xN are two sequence-tor toacccettt--eeeeeeee to-seeu-ncee fumctions that are implemented by multi-layer transformers, i.e. there exists 01, 02 such that
h
H
TFo. :H
RDoDDDi)xN
h
11)
(H(H))
ISKSN
(0)
H)))
TFo. :H2
DDDD+
h
(2
ISSSN
Then, there exists 0 such that H' that takes form h = [ (h h2 2 with h, (0) D h (2)
E R h(")
E D2, we have
(
TFo :H'
E DD)
Pi (H) (2
h ISISN
Further, 0 has most I. S max {I,, L2) layers, max M(O maxxe(L) (M - -M))
heads, max D < max (D + D(.)) hidden dimension its MLP layer (understanding the size of the empry layers 0), and norm bound |@l S ||0,| |02
Proof. When L, = L2 (0, and 0g have the same number of layers), the result follows directly by applying Proposition B.5 repeatedly for all L1 layers and the definition of the norm (2).
(without loss of generality) L1 < L2, we can augment 0 to Le layers by adding (L2 - L1) layers with zero attention heads, and zero MLP hidden dimension (note that this does not change M,. Dí. and |0 I). Due to the residual structure, the transformer maintains the output P1 (HH throughout layer L1 + .,.... L2, and it reduces to the case L1 L2-
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 sequence of N input vectors (h, } i=1
input matrix H [hj.... hy] E DDNN Recall that (t) == Re) max (t, 0} denotes the standard relu activation.
N RDD written compactly
C.1 Decoder-basee transformers
Decoder TF's are the same encoder TFs, except that the attention layers are replaced by masked attention layers with specific decoder-based (causal) attention mask.
Definition C.1 (Masked attention layer). A masked attention layer with M heads is denoted as MAttno()) with parameters 0 = {(Vms Qm: km)]Km))mmp DDD On any input sequence
H DxN' with N' < N,
( H MAttn (H) == H + DM, (V H) x (MSK o () (K H))H E DxN
(10)
23

---

where denotes the entry-wise (Hadamard) product of two matrices, and MSK E RNNN is the mask matrix given by
1/2
1/3
1/N
1/2
1/3
1/N
MSK
0
0
1/3
1/N
0
0
0
1/N
In vector form, we have
h [tttn (H)) h, Em= 1 Eド a(Q Kmh))) Vnhhy.
Notice that standard masked attention definitions use the pre-activation additive masks (with mask value -oo) [84]. The post-activation multiplicative masks we use equivalent to the pre-activation additive masks, and the modified presentation is for notational convenience. We also use nor- malized ReLU activation + (()/i in place of the standard softmax activation to be consis- tent with Definition 1. Note that the normalization 1/i is to ensure that the attention weights (o((Qm kyh / is a set non-negative weights that sum to O(1). The motivation of masked attention layer is to ensure that, when processing sequence of tokens, the computations at any token do not see any later token.
We next define the decoder-based transformers with I, 2 transformer layers, each consisting of masked attention layer (c.f. Definition C.I) followed by MLP layer (c.f. Definition 2). This definition similar to the definition of encoder- baser transformers (c.f., Definition 3), except that we replace the attention layers by masked attention layers.
Definition C.2 (Decoder-basse based Transformer). An I.-layer decoder-based transformer, denoted as (-) is composition of L self-attention layers each followed by MLP layer: HLL) DTFo(H(()) where (0)) DxN is the imput sequence, and
H MLP
MMtttn
000 (H(- E L}.
(l:1)) 0 ((0
Above, the parameter
is the parameter consisting of the attention layers K!n)) me(asros C DxD and the MLP layers 0.1 = (w(9, w!") E
mip
1 Q = {(V
RD xD x 18 Dx DxDx D We will frequently consider "attention-only" decoder-based transformers with
((O) (
which we denote (TFg()) (-) shorthand, with O (1:L)) 00:22
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, xN+1). which follows the
generating rule as described in Section 2.2, into input sequence H E R Dx(2++11) In our theory, we use the following format, where the first two rows contain xD1x)11 which alternating between [x;; 0] RR++1 and (0xx1; vi] E RR++1 (the same setup adopted in [31, 2]); The third row contains fixed vectors {Pi [PiheeNN+1) with ones, zeros, the example index, and indicator for being the covariate token (similar to a positional encoding vector):
Op (d+4) i
X1
XN
O
XN+1 0
51
J/N
p ==
E ER (d+1))
(11)
P2N-1
P2N
P2N+1
mod + 1,2)
(11) different from out input format (3) for encoder-basee TFs. The main difference is that (x,, vi) are in different tokens in (11), whereas (x; ys) are in the same token in (3). The reason for the former (i.e., different tokens in decoder) that we want to avoid every [x,; 0] token seeing the information of since 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 the last token, not necessary to alternate between [x;; 0] and [0; y.] to avoid information leakage.

---

We then feed H into decoder TF to obtain the output H D E R Dx(2++1) with the same shape, and read out the prediction ØN+1 from the (d 1, 2N + 1) entry of H [h. [ei] (the entry corresponding to the last missing test label): 1N+1 = read, (H) = (h2N The goal is to predict #N+1 that is close to UN+1 Poxxx41 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 {ūz Jj([N+1], where Yj gives the (d+ 1, 1)- th entry of H = [bi]se/an++] for each j E [N + 1] (the entry corresponding to the missing test label
of the 2j 1'th token): Dj ready.; (H) == (h2)-1)a+1 Given a loss function é RR x + TR associated to a single response, the training loss associated to the whole input sequence can be defined by E {HH) ンø= N-1 {(yj. 02). This potentially enables less training sequences in the pre- 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 TF's can be converted those of the decoder TFs. Taking the implementation of (ICGD) (a key mechanism that enables most basic ICL algorithms such as ridge regression;
Appendix D.1) as an example, this conversion enabled by the following facts: (a) the input format () of decoders can be converted to the input format (3) of encoders by 2-layer decoder TF; (b) the encoder TF that implements (ICGD) with input format (3), by slight parameter modification, can be converted to a decoder TF that implements the (ICGD) algorithm with 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 |00|| < 12 such that upon taking input H of format (11), outputs H = (T)(H)) with
X1
X1
XN
XN
XN+1
51
UN
0
(12)
P2
P2N-1
P2N
P2N+1
In particular, format (12) contains format (3) as a submatrix, by restricting to the {1, (1.2.....0)
1, D 2, D} rows and {2, ,,, 2N
2, 2N + 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, we slightly change the normalization in MSK from 1/2 to (/( 1) v 1), then the same construction performs (ICGD) (with training examples (1....,j) at every token = 2j + 1 (corresponding to predicting at Xj+1). Building on this extension, all our constructions in Section and Section 4.2 can be generalized to decoder TFs.
C.4 Proof of Proposition C.1
For the simplicity presentation, we write c; = [1/2] 4 mod(i + +1,2), u, = h1/11 d Rd+1 be the vector of first d entries of h; and let v; h,(d +1] be the (d + 1)-th entry of h, With such notations, the input sequence H [h:] can be compactly written as
h, [n oy; 0p-d-iic;; 1;
In the following, we construct the desired = (0(1), 222) as follows.
Step 1: construction of (1)) = (0 0
( ) so that MLP ( MAttn ((1) maps
MAtte
h
h = [u; Uj 0 D ( 40;;;
'In other words, when 2 i, u X(e- 1)/21 when i, u Oa-
25

---

MLP
h h [u 0i; 0p- bo (e:
For m {0, 13. we define matrices Q!! Kmn, vn RR
DxD such that
0 h h; = Q t.
K. h K| = 0
0
0
3c,
v hy
0pp-
2
O,
02
for all j. By the structure of h,, these matrices indeed exist, and further it is straightforward to check that they have norm bounds
max
S 1,
max m
K lp
1
Σ
5.
T
op
Now, for every i,
Σ Σ ((O!!) k )
i j=1 mE(U.1)
i 1=1
Notice that * 0 only when i, then compute for 2k that
Σ 3e = 3 k(k 1)(2k - 1) + 3k 2k - k,
Σ 2c, 2- -2k + 2k =2k?
3
Therefore, the 0 (!1) natta
construct above is indeed the desired attention layer. The existence of the desired 0ul is clear, and 0!l (W(1) w()) can
((Qm) K. vO E RRDD)) Jmm (0.1)
further be chosen so that [ llop I 1, 1 llop 1..
Step 2: construction of ((2), For every m {-1,0, 1}. we define matrices Q m K
2 v
(2)
DDD such that
tf
Q
2 2
= 0
2 h = Q
t
k0"h (1)
-Cj
K
k(2n000 h
(1)
(cj + 1)
K (2) (1)
h
- (cp -1)
v n
(1)
4u
h
(1)
1
(1)
211 0D-d
OD-d
for all j.j. By the structure of h (1) check that they have norm bounds
these matrices indeed exist, and further it is straightforward to
max 771
1
max m
K (2) op
2,
Σ
8.8
op
Now, for every j, we have
Σ o ((o 1 ( (.o0""
( (1)
mm(-11,01)
{ 1-20 (1c) 40ge) + (ag halg 1)) o o - 40l6 1))} 2uu; On-al
=1220(c c,) + aalaa-o000 11) +ola c,) 1)} 20,t/ O
= =I(a -;) 2c,t, (u: Op-al:
where the last equality follows from the fact that
0, x x2 1 or r < -1,
-20(x) olx 1) ox + +1) 1)=
++1, Ie 1,0),
1-x, x, xE [01]
26

---

Therefore,
Σ Σ j=1 m
Σ
((
) v ,
Σ 21 = e)) Op
[ 0 0p-dls i= 2
OD.
otherwise
Therefore, the (22)
[(Qm V E RDDD)) mE{--10,1) we construct above maps
nattn
h
x1//2] 12 0 - L0f 40g 1;4]]
Finally, we only need to take MLP layer 0 (W (2) w!()) with hidden dimension that maps
h" h 2 [ Uy; 0p DD 0; c 1;4))
which clearly exists and can be chosen so that w (2) op - 1, W 1lop S 1.
Combining the two steps above, complete the proof of Proposition C.I.
D Mechanism: In-context gradient descent
Technically, the constructions in Section 3.1-3.2 rely on new efficient construction for transformers 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.
will often use shorthand s/. E R defined as y/ = yı for E [N] and - = 0 to simplify our notation, with which the input sequence H R Dx (N+1) can be compactly written as x,; v/: p.] [x; 1/3 -p-d-p--ss- 1; for E [N + 1], where t == 1 < N + 1) is the indicator for the training examples.
Compact notation of input
D.1 Gradient descent on convex empirical risk
Let ((:,)) ) ?R? + R be loss function. Let L (w) == N Em l(w Nw, X; denote the empirical risk with loss function dataset {(x,, 3:) [ and
w WD w'an - V (w(n)
(ICGD)
denote the gradient descent trajectory on IN with initialization wan E IRd and learning rate > 0.
We require the partial derivative of the loss 0,6 (s,t) + a,l(s,t) (as bivariate function) to be approximable by a sum of relus, defined as follows.
Definition D.1 (Approximability by sum of relus), function g R* R (Eapproxs R, M, C)- approximable by sum of relus, if there exists a "(M, C')sum of relus" function
fM,C(Z) EM C
(a Am [z; 1)) with ΣM |cml SCC, max |am S 1, am E RR+11 Cm E R,
such that supzE(-R.nn R. (g(z) fM,c(z)] S Eapprox-
Definition D.1 is known to contain broad class of functions. For example, any mildly smooth k- variate function is approximable by sum of relus for any (Capprox, R), with mild bounds (M, C)
(Proposition B.I, building on results of Bach [4]). Also, any function that is a (M, C') C')-sum of relus itself (which includes all piecewise linear functions) is by definition (0, co, M, appppooxiaa by sum of relus.
We show that L steps of (ICGD) be approximately implemented by an (L 1)-layer transformer. Theorem D.I (Convex ICGD). Fix any B., > 0, L L, 0, and s < B,, / (2L). Suppose that
1. The loss f(.,)) is convex in the first argument;
2. 0,6 is (s, R, M, C)apppooxeelble by sum of relus with R max (B,Ba, By 1}.
27

---

Altention
Qac
K
V
We ght commiattion
N
Σ an
b
Universa
a Ey By
approsimatiio
Gradient
x s
descent
Figure 4: Illustration of main mechanism for implementing basic ICL algorithms: One attention layer implements single (ICGD) iterate (Proposition E.I & Theorem D.I). Top: attention mechanism as in Definition 1. Bottom: A single (ICGD) iterate. Middle: Linear algebraic illustration of the attention layer for implementing a GD update.
Then, there exists an attention-only transformer TF. with (L 1) layers, max [1] M S M
heads within the first L layers, and (L++) = 2 such that for any input data (D, xN+1) such that
sup Amax L (w)) S 2/m, Ew E arg min LN(w) (w) such that ||w" 1|2 S Bw/2, [|w B
wER"
TF. O)) approximately implements (ICGD) with initialization w W&D = 0:
1. (Parameter space) For every E [L]. the E-th layer's ' output H(O) TF gang (H))) approx-
imates l steps of (ICGD): We have h [x w Op 2d-3i 1; (4] for every [N 1]. where
| w
(Bn
Note that the bound scales as Ol)), linear error accumulation.
2. (Prediction space) The final output H (L++) TF. (H)) approximates the prediction of steps
of (ICGD): We have h (I+1) so that
N+1 [xN+1 +1 w. 0p-20-s2:: 1;t,), where UN+1 -( xN
ŪN+1 woDDXN++1 5 (L.
Further, the transformer admits norm bound ||0|| < + R 2nC.
The proof can be found in Appendix E.2. Theorem D. substantially generalizes that of von Oswald et al. [86] (which only does GD on square losses with linear self-attention), and is simpler than the ones in Akyürek et al. [2] and Giannou et al. [32]. See Figure for pictorial illustration of the basic component of the construction, which implements a single step of gradient descent using a single attention layer (Proposition 1.
Technically, we utilize the stability of corvex 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 I. in the non-convex case in general; see Lemma D.3(b).
Lemma D.1 (Composition of error for approximating convex GD). Suppose f Rd + R. convex function. Let w E arg min f(w), R 2 2||w*]|2, and assume that V f is L -smooth on ( Let sequences [w" Ye20 c. Rd and {weDle2o C. Rd be given by w = = = 0,
c++1 w /VW +
4+1 w v fW W
for all & 2 0. Then as long as $ 2/ L,, for any $ & R/(2e), it holds that | - worll2
Le and [ 1|2 < + Le R.
28

---

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 non-smooth regularizer such as the ε norm. In this section, we show that transformers can approximate PGD with similar quantitative guarantees for GD in Appendix D.1.
Ee Let ((.,-) IR² + R be a loss function. Let [w(w)= == N -i=1 <(w x, Vb) + R(w) denote the regularized empirical risk with loss function E dataset ((x,, 3e) and regularizer R.
x,,
To minimize LN. we consider the proximal gradient descent trajectory on Ln with initialization WGD = 0 E Rd and learning rate 7) > O:
W
H+1 PGD == prox RR WPGD 1/ T.N (weeD)),
(ICPGD)
where we denote (w) Ei= e(w x.,
To approximate (ICPGD) by transformers, in addition the requirement on the loss é in Theo- rem D.I, we additionally require the the proximal operator , (-) to be approximable by an MLP layer (as vector- valued analog of Definition D.I) defined as follows.
Definition D.2 (Approximability by MLP). An operator P Fd + Rd is (E, R. D, C. approximable by MLP, if there exists a there exists a MLP Om1- (W1, W2) E Dxd RdxD with hidden dimension D, || Illp + ||22llpp S C", such that sup w 1 SR P(w) MLPel, (w)ll2 S..
The definition above captures the proximal operator prox, for broad class of regularizers, such as the (commonly-used) Lı and Le regularizer listed in the following proposition, for of which one can directly check that they can be exactly implemented by MLP stated below. Proposition D.1 (Proximal operators for commonly-used regularizers). For regularizer R in [A & I1-112 (B)((,)())), the operator prox,R Rd IRd is exactly approximable by MLP. More concrelely, we have
1. For R > 11- , (0, +00, 4d, + 2nA)aapprozimable by MLP.
2. For R= 2/2 11.112 p is (0, too, 2d, 2n approximale by MLP.
3. For
(B)()) prox, Proje (B) is (0, +00, 2d, 2 by MLP.
Theorem D.2 (Convex ICPGD). Fix any B L 1, 17 and e e B2w(2 Suppose that
1. The loss ((.,)) is convex in the first argument;
2. 0.6 is (s, R, M, C)-approximable by sum of relus with R = max {Ba Bur, By
3. R convex, and the proximal operator prox, (w) is (nE', m R', , D', C²) C")-approximabl by MLP with R' sup//w/l,SB.. ( + n7..
Then there exists a transformer TF' with (L+ 1) layers, maxte[[] (()) M heads within the first layers, (L++)) = 2, and hidden dimension D' such that, for any input data (D, (D,))11) such that
sup Amax L (w)) < 2/n, Aw* E arg min T. (w) such that || 2| E B weRd
||w B
TF (H()) upproximately implements (ICGD):
1. (Parameter space) For every f. E [T.,, the E-th layer's 's output H(O) TT
Paa. (H()) approx-
imates l steps of (ICGD): We have h,? = [x;;1/: w 20 Op a5 1;44 (4] for every i E [N + 1]. where
|| W w 1 < (s + (LrB2).
2. (Prediction space) The final outpur H (1++1) = TF. (H) approximates the prediction of steps (L+1)
of (ICGD): We have h N [xx+1;; 11 1 0p-24-3 1;4;), where UN+1 (w xN+1) so that
[N41 WWFCD.XN41) ( (2LmB?).
29

---

Further, the weight matrices have norm bounds |0]| & 3 R 2n/C + C".
The proof of Theorem D.2 is essentially similar to the proof of Theorem D.I, using the following generalized version of Lemma D.I.
Lemma D.2 (Composition of error for approximating convex PGD). Suppose f Rdd R is convex function and R. is convex regularizer. Let w E arg min f(w) R(w), R 2 and assume that V f is L -s-ooohh on (R). Let sequences Je20 C Rdd and {WGD)E20 C Rd be given by wD WGD = 0.
et proX (w - -VV[(WW)) + +e',
11e'1125e,
[+1
proxYR (WGD /VVS(WG))
GD
for all E2 O. Then as long as S 2/ for any 0 L S R/(2s), it holds that 1 - woDll2 v
LE and || S + Le R.
The proof of the above lemma is done by utilizing the non-expansiveness of the PGD operator w + prox,R R (w - 1/ f(w)) and otherwise following the same arguments as for Lemma D.I.
D.3 Gradient descent on two-layer neural networks
We now move beyond the convex setting by showing that transformers can implement gradient descent two-layer neural networks in context.
Suppose that the prediction function pred(x; w) k== u is given by a two-layer neural network, parameterized by [Vx:"A]N<(K] (R((dt)) Consider the empirical risk minimization problem:
IN (w) WEW
min
2N Σ (pred(xx; w).
Σ (or(vxx)), 81
(13)
2N
where W is bounded domain. For the sake of simplicity, in the following discussion we assume that Projw be exactly implemented by MLP layer (e.g. W (R..). for some R. 0).
Theorem D.3 (Approximate ICGD two-layer NNs). Fix any Bu, Ba 0, L 1 0, and
E > O. Suppose that
1. Both the activation function r and the loss function E is C+-smooth;
2. W is a closed domain such that W C {w- [VI; <RK(d+1), 1| S Buy |ugl s B,}, and Projw MLP for some MLP layer Ouap with hidden dimension Dw and
||| 1mppp || S Cw:
Then there exists (2T.)-layer transformer TFa with
max M < 2ō (c-2) EE(2L]
((9) S (E-2) Dw, EE[21]
max
||0| s (1 + 1) + Cu.
where O()) hides the constants that depend on K, the radius parameters Bz, By. Bu: Bi and the smoothness of r and such that for any input data (D, xNx4+1) such that input sequence H DxN+1+1 takes form (3), TF (H())) approximately implements in-context gradient descent on
risk (13): For every E [L], the 26-th layer's output h (26) = [x; w 0; 1; for every E [N N++1 and
w = Projw (w m/( (W )
+++++++ +))-
w = =0
(14)
where || is error term.
As direct corollary, the transformer constructed above can approximate the true gradient descent e+1
trajectory (wáDlezo on (16), defined wop 0 and w w Vn for all e > 0.
30

---

Corollary D.1 (Approximating multi-step step ICGD on two-layer NNs). For any L 2 1, under the same setting Theorem D.3, the (2Ly-layer transformer TF, there approximates the true gradient descent trajectory (WeDle20' For the intermediate iterates {w") Jeets considered therein, we have
W woDll2 S (1+
where T, = SuPwEW
2 T w))
denotes the smoothness of T, 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 expected, since gradient descent 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 the stability-like property Lemma D.1 no longer holds for the non-convex case.
Corollary D.1 is 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. any function f W > IR, and any initial point w E W, let
W (w) == w - Projw(w vVs(w)
7
denote the gradient mapping at w with step size 7, a standard measure of stationarity in constrained optimization [63]. Note that GW.,( (w) = V f((() when w - f())(w) E W (so that the projection does not take effect).
Lemma D.3 (Convergence and trajectory eloseness inexact GD). Suppose f W + R. where W C. Idd is conver closed domain and f is Ly Lipschhti on W. Let sequence {w Je20 IRd be given by W
w
Projw (w - ( VVV[(((WWW +E)).
| S,
for all l 120 0. Then the following holds.
(a) As long as 1j S 1 I/T. all 1
EEL-11 G/ ((w,,("))) (
G (W ( 8(f(w") - new f(w))
min
+ 10sa.
nL
(b) Let the sequences (waplezo C Rd and be given by w&D w and WGD 4+1 Projw(W'p - (VV((w'D)) Then it holds that
w s 1 + nT Ve 0.
E Proofs for Section D
E.1 Approximating a single GD step
single GD step by a single attention layer). Let E(.:-) 2 + IR be loss function such that 0,l is (a, R, M, C)-approximable by sum of relus with R = max Bu: By, 1}. Let [w(w) := E f(w x,, vk) denote the empirical risk with loss function on dataset ((x,, y:) ieeee
Proposition E.1 (Approximating
Then, for any e > 0, there exists an attention layer 0 = {(Qm- Km, Vm) mE(M] with M heads such that, for any input sequence that takes form h; [xi; 1/4 w; On 24-3; 1; Li) with ||w]]2 S B gives output h, = [Attng (H)); = [xx; 1/; W: 1D-00--22-00
1; 1;1 for all [N 1), where
w - (w [WN)())]
S s (n..).
Further, ||0| 2 R + 2nC.
31

---

Proof of Proposition E.I. As i,L (E, R, M, C) approximable by sum of relus, there exists func-
tion f -RR R| R of form
M
[s,
Σ m=1
Cm (am 8 + bmt + dm) with Σ
cmal S C, |am| + |bml Idml 1, Vm E [M],
such that (uP().))E(RRR R, Rj2 (s, 1) - 0, 1) S
Next, for every m E [M], we define matrices Qm Km, Vn E RR DxD such that
am
0a
bm
Vmhj (N + 1)nem
0
Qmh =
dm
Km hj
-2
R(1 - 4) 0
N
Xj
OD-22-1
for all i,j E IN + 1]. As the input has structure h, [x6881 Op-00222 20-30331 1;t,l. these matrices indeed exist, and further is straightforward to check that they have norm bounds
max M
1Qmllap
S 3,
max ME(M]
K ll
2 R
1| S 2nCC. mE[M]
Σ
Consequently, |0| 2 R Cn.. Now, for every i,j [N + 1), we have
a(Q@m h; Km (am w xj Pm 1z))y3 dm 2Rl,
(am x, bmyl + dm) {L, 1};
where the last equality follows from the bound 0m W x, + 6m (1 t,)y, + dml S am BgB + R S 2R.
(15)
so that the above relu equals if t2 I O. Therefore,
M
a ((mm,, kmh Vnh,
m=1
M
(N + 1)m7
C (am w Xj + bmi + dm))
Tx,
114, 03 xy; 02]
N
11=1
(N + 1)m 11, = 01)==0000aa 10++ Xy; On 02-2a-1
s(w Xj,
N
Thus letting the attention layer ((Vm. m; Qm, Km) me(M) we have
N+1 M
h [Attno h; +
1
Σ
a((Qm k,)
N +
m=1
N
= -h;- 1J
Σ
f(w Xj, 1z: (0)4))0041;;;;0000 Ozl
N
j=1
=x: w: 1:t+]
17 N
Σ
a x5. 1p-2- + Op-2a-1]
= [x:06 w + e; Op-24-p---- 24-3: 1;4,],
where the enror vector e E Id satisfies
N
||l||
Σ
([(wTx,,y)) x,,3;) 0,l(w") O.((wTx,,.)))) xx,,))))))
1=1
32

---

N
f(w T x3. - -Og((w' x,, (5.15)1 ||x;|] l12
N g. Br = = (1)B)).
N
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 O. Note that by our precondi- tion L S / (2s).
By our precondition, the partial derivative of the loss 0,E is (e, R, M, C*) approximaalee by sum of relus. Therefore we can apply Proposition E. to obtain that, there exists a single attention layer (()) = {QQms Km, Vm)) with M heads (and norm bounds specified in Proposition E.1), such that for any w with ||w|] 12 S Bu, the attention layer Att maps the input h, = [x, 8/3 w; Op -p-20-32- 1;
to output h = [x: W; OD-2d-3: 1; ta] for all E [N + 1], where
w - -ww-nVLA 7 (w) |1. S E (n.,) =:E'. w
Consider the L-layer transformer 0² ( )()) which stacks the same attention layer ((1) for L times, and for the given input h, (0) [x, f/i w 0D-20-3: 1; its E-th layer's output h = xx; s/: w²: 1:-22-3::;;:)
We now inductively show that || [| S B and w w s ls for all E [L]. The base case of f 0 is trivial. Suppose the claim holds for Then for f s L B (2s), the sequence [w N5e+1 and (WGDliser satisfies the precondition of the error composition lemma (Lemma D.1)
with error bound e, from which we obtain
|W++||
S B and
f+1
W
ll S (e + 1)e'. CD
++1
This finishes the induction, and gives the following approximation guarantee for all E [L]:
|| W - Se' S e (L.)B.)
which proves part (a).
We now prove part (b), which requires constructing the last attention layer ((L+1), Recall h (1)
[x v/: wh; 124;0p-24-0 20-24-3i1:44 1;44] for all i E [N + 1J. We construct 2-head attention layer 0(++1)
1(Om (L)1)) (((+1) v((+I)) )mm=1, such that for every i,j E [N + 1].
(L+1) (L) =xx; d-d))
d). K k (2) ( v(L+1) (2) [0.: 1; 0D-d-1)
(++1) (L) [x; -d). K
(++
1- v! bft
Note that the weight matrices have norm bound
max
< 1,
max 1=1.2
K
[++1)
$ 1,
Σ v L++
22
1=
op
op
Then we have
N+1 2
h (1++1 =
1
Σ Σ ( ((Q(2+1) (1) ((+t))n())) (1 v
N+1 1N+1
+
N ++
m=1
[x 1W: Op-2d-3: ii 1; 1] + (a(( W/ xx+1)) a(- (W xx+1))) (0x: 1: 1:0
=
(x: 0; : 3: 1: 1] + 00x: (W xN+1) 0D-J-1]
[x; (wL. xN+1): 4 0p 1;1]
NN+1
33

---

Above, (i) uses the identity t = o(t) - (--)). Further by part (a) we have
ØN+1 - (WGD: xN+1>| = (w wen, xN+1)| S E (Ln
This proves part (b), and also finishes the proof Theorem D.1 where the overall (L + 1) layer attention-only transformer is given by TF2 with
0 (0 0(1) e(L+))).
times
E.3
Proof of Lemma D.1
As f is convex, Ly smooth function on (1)) the mapping To + -)/f(w) is non- expansive in I|:11:: Indeed, for any w, w' E (I) we have
[[T)(w) (w')l2 [ww T/W (w' (/((')))1
ww w'l12 2n (w w. Vf((w)) +1 [VV][I - Vf(w')lī (i))
w/12 (2n/L 7/) IV V/(w")ll S || w'l
Above, (i) uses the property (w - w ff(w)- - vf(()) 2 ,, lV/(w) - Vf((w)l for smooth convex functions [63, Theorem 2. 2.1.5]; (ii) uses the precondition that 1) $ L/
The lemma then follows directly by induction on The base case of L follows directly by assumption that w w&D (R/2). Suppose the claim holds for iterate L. For iterate L & R/(2),, we have
1+1 w J++1
112 = Th - (ww
GD
(W (T 112 + ||o413
(i))
- w6oll2 + S (L + 1)e.
Above, (i) uses the non-expansiveness, and (ii) uses the inductive hypothesis. Similarly, by our assumption w' = (w")).
R + (L + 1)s S 2
| w 12 [T, + Tn(w))l2 S 1 +1112+11==11
This finishes the induction.
E.4 Convex ICGD with l regularization
In the same setting as Theorem D.1, consider the ICGD dynamics over an lz-regularized empirical risk:
w GD = WaD - X (W(D)
(ICGD-&2)
with initialization w Rd and learning rate 17 > O, where { == [N((w) + |w|2 denotes the fz-regularized empirical risk.
Corollary E.1 (Convex ICGD with l regularization), Fix any B > O, L n > 0, and < B,B Suppose the loss e(.,-) is convex in the first argument, and 0,6 is (e, R, M, C)- approximable by sum of relus with R max (B Bar By: 1}.
Then, there exists an attention-only transformer TF0 with (L+1)) + 1) layers, maxre(r] M S M + heads within the first L layers, and ML+1) = 2 such that for any input data (D, xN+ with
sup Amax 2 (w)) < 2n Iw* E arg min E (w) such that || 1|2 S B
|w B..
weRd
TF. ())
approximately implements (ICGD-&):
34

---

1. (Parameter space) For every & E [L], the E-th layer's outpur H = TF on (H approxi- mates i steps of (ICGD-E2): We have h,d [x w Op D0-00------- 24-3; 1; for every E [N + 1). where
[w WGDll2 E (2Lm
2. (Prediction space) The final output ((+1) = TF H approximates the prediction of L steps of (ICGD-82): We have h hN+1
(L+1) [xN+1] 1 w 30: 1; 0), where
1N+1 (WeD, xN+1 s (2L
Further, the transformer admits norm bound ||||| - + R + (2C + X)1--
Proof. This construction is the same as in the proof of Theorem D.I, except that within each layer e [L]. we add one more attention head (Q(9, K(O), v(0) c IR DxD which when acting its input
h (É-1) [+;* :;
w 3; 1; #]
gives
04+1
(
k(Onf-))
viOn ((-))= =
-//W" --1
0D-1
0D-1
02
for all i.j E [N + 1). Note that | llop = |
= 1, and | lop
ny Further, it is
straightforward check that the output of this attention head every h is
0a+1 -m/W"
N N Σ (1 (1-)) (Ko)""" viont-
02
Adding this onto the original output of the E-th layer exactly implements the gradient of the regularizer w + // The rest the proof follows by repeating the argument of Theorem D.I, and combining the norm bound for the additional attention head here with the norm bound therein.
E.5 Proof of Theorem D.3
only need to prove the following single-step version 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 and any E Rdd TF maps
h; = [xx;0//1 w; 0; 1;4]
=/= [xi ( w 0; 1;44,
where
+ Projw (w- v + =(w)).
||e(w)]22 S mE.
Before we present the formal (and technical) proof of Proposition E.2, we first provide some intuitions. To begin with, first note that
1 Σ 0 ((((pre((x;; w).x). V V pred(x;; w), V L (w)
(16)
where 0,e f is the partial derivative of & with respect to the first component, and
U1 r' ('v xi)) x;
r((vixx))
V wwerdd(x; w)
E kid++))
(17)
uk ((VK, xx)) X;
r( r((VK.
35

---

Therefore, the basic idea is that we can use attention layer to approximate (x,, w) + pred w), then use MLP layer to implement (pred(x,; w). V.. t,) 1 N + 1] (p((pred(xx; w) Vs). and then use attention layer to compute the gradient descent step + w 1/ IN (w), and finally use an MIPP 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 Do d 1 K(d+ 1) be the length of the vector [x5566 w]. We also define
B max 1d
|r(t),,
Bg
max [are(t,))) IIKKBBlulSB,
Let us fix Er: Ep. Eg > 0 that will be specified later in proof (see (18)). By our assumption and Proposition B.1, the following facts hold.
(1) The function r()) is (Er, R1, M1, Ci) for Ri max (B, 1J, Mi S Ō (C9e, where C1 depends only on R1 and the C2 -m-sootthsssss of Therefore, there exists
M
r(t)
Σ m=1
Cn 1}) with Σ
cm
S C1.
|| 1, Vm E [M]]
m=1
such that sUPrE(-RI,8) [r()) =(()] S Er.
(2) The function (t, 0,((t, is (Ee, R2, M2, C2) for R₂ max (KB,, By, M2 I Ô (C3e,?), where C2 depends only on Rg and the C3 smoothhhss of 0.e. Therefore, there exists
| 11 S 1, Vm [M2], gt. y) o (t:;;:1)) with Σ cm S C2.
such that suP(ta))E- R2. lo v) - 0,((t, v)1 S Ee.
(3) The function (s,t) + s s. r') is (Ep. R3, M3. C3) for R3 max [B,Bu, Bu 1}.
Ma I Ô (C3ep)) where C3 depends only on R3 and the -mmooothnesss of r'. Therefore, there exists
M P(s,t) = Σ
M
en c (a [s;t; 1])) with Σ
c S C3,
| ||, I 1, Vm E [Ma].
m=1
such that suP(s.))el R. P((,
(1))) 6p-
In the following, proceed construct the desired transformer step by step.
Step I: construction of (1) We consider the matrices
{Q k. K()) v(I) kke so that
for all i,j E [N 1]. we have
a [1] X; am
Vk 1
(1) k.m
h,
[2]
K (1) k.m hy
v(I) k,m h, = ct, UpCD.+1
As the input has structure h; = [x; y/: w; 0; 1; fz). straightforward to check that they have norm bounds
these matrices indeed exist, and further it is
max
11,
max k.m
k.m
11,
(
k.m
S
C
op
op
up
kmm
A simple calculation shows that
K Σ
Σ
((0.2 h, (()) hy)
k.m
v(l) )
kim
h
up 1 x,)) -pptt--
ME[M]]
k=1
36

---

For simplicity, we denote pred w) E k= 2x7((vk, in the following analysis. Thus, letting attention layer 0 {(V()) kim Qhim km (k.m)
we have
o h, h; (0.5) [xA 3/: w; p w); 0; 1,
Attn
Step 2: construction of 0 mlp' We pick matrices WI, 1, W2 so that W1 maps
RR,, and W2 E Ms with entries being (W (22) = en { Do + 2). It is clear that [ llpp S
W2hg (0.5)
a [1] (x; w) + a (2) y/ + a [3] R2(1 t.) mE[Ma]
R2 + || I C2. Then we have
W2, (0.5)) Σ
((aa... [pred(x;; w): u/: 11) -
- eD-2
ME[Ma]
ift, = 1} (prrdd(x; w), v/) CDo42.
In the following, we abbreviate gi 1[t 1} g(pred(xx; w), vt) Hence, Omp maps
MLP am. h (0.5) h (1) [x; w: px; w); 0; 1;
By the definition of the function g, for each [N].
0 (pred(xx;; w), vi)l E BuLiEr.
where Le = max | LS (t, y) is the smoothness of 0, Also, 9N+1 0 O. definition.
Step 3: construction of 0 oee... We consider the matrices
{Q
K (() ( k.l.m [em so
km
that for all i E [N + 1], we have
(2) 1.
h
(1)
Ta [1]
amm [2] Vk
am [3]
K R.
h
(L)
X
gj
I
v k. h
(N + 1)mc
N
1k(d+1)
X
We further consider the matrices we have
{Q
K k.2mm
(22)
(E(()),M((())) so that for all i.j E [N+1].
k.2.m
mm [1] a/m [2]
x
(N + 1 1)7ch.
Ob(d+1)+d
k22mm
K 2m
h
=
1
V
k.2mm
N
0
0
By the structure of the input h. these matrices indeed exist, and further it is straightforward to check that they have norm bounds
max (1100)117))
@k.
11,,
max (k,,.,m)
K R. k.w.m
S
1,
v 2)
k.w,m
2 + 27/33.
op
lo
(k,w,)
up
Furthermore, simple calculation shows that
Oa+1 P(uig (v1,x;)) X3 F((v1,83)) 9j
N+1
N+1
g(w)
Σ (k,w.))
h, (2) k. h,
v(()
k.u.m h,
N
N
P( (VK: xx)) X, F((vk,,,)) 93
where the summation is taken over all possibilities of the tuple (k, w. m). i.e. over the union of [K] x {1} x [Ma] and [K] x {2} x [M1].
37

---

By our definition, we have [P(s,t) - st)( S Ep for all 8, [-R3, Rs). Therefore, for each i [N]] & [K],
[P(wgs (Va, x;)) 01 (1e(pro)(xx; w), UU) up r' ('v x,))1 - Ep 1gj O1 (x((prd(xx; w), vi)l |uel ((vx:)
S + ( + B, Leer);
where L, == max (t)) the upper bound rf. Similarly, for each i [N], [K], we
have
Ir((vk, x,)) Oj r((vk((vvvvv x,)) O {(pred(x;; w). v3)) & 2 + 2B, + Bu [.ger)
As for the case = NN + 1, we have 9N+1 = 0 and [( (V,:NN+1)l S Ep for each & E [K] by defintion. Combining these estimations and using (16) and (17), we can conclude that
| 1n'gg( g(w)= VEx(w) S SVKB, (ep + (e, + B,LL,,)] 22 Br(er + Bubr,,)]
|'(((w) + VIN(w)// S, we only need to choose Ep, Et, &r as
Thus, to ensure
E
E
Ep
Et
3
9vK max (rr, L Bu
15v max (B,, Le Buy Ly LBB,
(18)
Thus, letting the attention layer 0 {(V(2) km
(2)
k,oom K (2)
km) (k,))))) have
Attnace te : :h
h (1.5)
[xx D4i
+ 1/ W): pred Xx; w) gi 0; 1:4].
Step 4: construction of 0 mlp We only need to pick 0 0
so that it maps
MLP (22)
h (1.5) = [x: w + rg(gw) (x; w): 9:: 0: 1;
(2)
[x s/s Projw( 0::
By our assumption on the map Projw. this is easy.
2) Combining the four steps above and taking = (e!!) atta mips atto completes the proof.
(
E.6 Proof of Lemma D.3
For every e 2 0, define the intermediate iterates (before projection)
w ++ W 77 (V ( + E²),
e+ WGD
== W -VWV((WDD)
so that w+1 Projw( (W and w GD Projw (
e+ w ).
G (we) ||
We first prove part (a). We begin by deriving relation between | - wel1 and
Let w == we - v(W*)) and w = Projw(W(th) denote the exact projected gradient iterate starting from w. We have
f++ w011
(8)
2
w
w
w
191 - | w
11
(iz)
|| w well2
2
2
(si)
1
G
W. (w
(w)
11
-11vel12
>
1
G (w)) 113 1/202
2
2
(19)
Above, (i) uses the inequality ||a - b1l > 1 ||a|| - |bb12: (ii) uses the fact that projection to a convex set is non-expansion; (iii) uses the definition of the gradient mapping.
By the L smoothness of ff within W. we have
f( f(w(t1) -f(W'= < (Vf(WW) w0t1 w²)
Ls 2
wl+11 1
38

---

W
w
E W
we
+
Ls
w wet w0l12
2
W
we-1
wet1 we
(E",
wetl)
+ 1ig
|
- w
2
1
1
|| - -w'11
1
||
1011 17 >||=1|1
+
17
2
4n
(++1
w011 ||
4n
(60)
1
n
G. (w+) 113 1 |I1ve(16) + 1||5811
4n
2
S " Gw. (w0) ||/
57 37702
+
4
Above, (i) uses the property ( (w wet w w²) - 0 of the projection w
Projw (wol)(W(+))) (W(+)) (using we W): (ii) uses L s 1/(2n) by our choice of S 1/by; (iii)
uses (19).
Rearranging and summing the above over e = L..... 1,, we obtain
7 Gh ww)l || S f(ww) f(wh) + 5n 2
Σ
8
4
Dividing both sides by n//.8 yields part (a).
Next, we prove part (b). Let C 1 nLg We prove by induction that
C²
wanl1 S C-1
(20)
for all e 2 0. The base case of e = 0 follows by definition that w. = wGp = w Suppose the result holds for É. Then for l 1, we have
| 2+1
w l12
(6) 2
W
W
8+ GD
||
=
w -n([ - - -WW
(1)) $ C W
(iit)
c6-1
C++1 -1
woDll2 + ne
S
C C-1
Ne +E nE
C-1 nE.
Above, (i) uses again the non-expansiveness: of the convex projection Projw: (ii) uses the fact that the operator w w w-r/Vf(w is (1 +nLb) C-Lipschitz; and (iii) uses inductive hypothesis. This proves the case for E + 1 and thus finishes the induction. We can further relax (20) into
O
|
w 1 S 1nmL -1 1ec L, (1 + nI,
This proves part (b).
F Proofs for Section 3.1
F1 Proof of Theorem 4
Fix 2 0,0 a 3 with x == 8+A and B,, > and consider any in-context data such that the precondition of Theorem 4 holds. Let
o+
N
((w,,)) - v2)² + A
Lridge (w) := 2.V
Σ
/|
1=1
39

---

denote the ridge regression loss in (ICRidge), so that wiidge = arg min bridge(W) It is standard result that V Lridge (w) X X/N Ala, that Lridge is (a + A)-strongly convex and (B A)-smooth over R4.
Consider the gradient descent algorithm on the ridge loss
W WGD WGD - 7/ Lridgo ( GGD
with initialization, learning rate, and number of steps
Bg B.
]
W = Od: = ++
T: 2z log
2=
By standard convergence results for strongly convex and smooth functions (Proposition B.2), we have for all 2 1 that
WW'D
w wridgells S
= exp
(-²) /wOD waiderl12 exp
(-x)
w ridge 1
Further, we have
WGD Wridge
2 S exp
T 2k
S
2 B
E
(21)
B.B.
2
B
remains to construct a transformer to approximate WED Notice that the problem (ICRidge)
corresponds to an fz-regularized ERM with the square loss ((8 == [s 1)2 whose partial is exactly
derivative 0,((s,
of two relus:
0,t(s. t) = 2ar(( -t/22 20(-(s--t) t)/2).
In particular, this shows that (,((s, (0, R, appprxxxiee for any R > 0, in particular for R = max { B2 By: By. }}.
Therefore, can apply Corollary E.1 with the square loss learning rate n,, regularization strength À and accuracy parameter
to obtain that there exists attention-only transformer TF. with
(T + 1) == L. layers such that the final output h (I. N = INN+1:NN+1 +] with
UN+1 (w xN+1 0,
(22)
and number of heads ( for all - 1] (can be taken as 2 the unregularized case =
directly by Theorem D.1), and (L) = Further, O admits norm bound [0] s 2 3R + 8(3 + A) + 4R 8(3 A)-
R +
S+A
S
3+X
Combining (21) and (22), we obtain that
|N++1 ( Wridge:-NN
= (WOD Wridge xN+1) (e/B2) BN =E
Further, we have read (h )= weD for all E [N +1], where read (h) := h(d-2) (cf. Corol- lary 1) so that ||read,,, (h²") w Wridgell l2 S e/ B, shown above. This finishes the proof. I
F.2 Statistical analysis of in-context least squares
Consider the standard least-squares algorithm ALs and least-squares estimator WLS E Rd defined as
ALs Ats(D)(xx++)
XN+1), Wis (X X) x ERR
(CCLS)
For any distribution over (x, IRd and any estimator w E Rd let
w/)²| Lp(w) == Exx wP~ ((W x") -
denote the expected risk of w over new example (x', ./'' P.
Assumption (Well-posedness for learning linear predictors). say a distribution P on TR. R well-posed for learning linear predictors, if (x. y) ~ P satisfies
40

---

(1) |xx|2 S B& and [yl S By almost surely;
(2) The covariance Σp == Ep xx satisfies Amin 3 Ep 3 Amaxta. with < Amin S Amax. and == Amax/ Amin-
(3) The whitened vector Ep 1/2 is -ss--u---ssssae for some K2 2
(4) The best linear predictor w Ep |x (xxxxx]"pxxx satisfies |wpll2 B
(5) We have El(y - (x, x S o² with probability one (over x)..
Further, say P is well-posed with canonical parameters if
B₂ ((vd), By (1), _ = @11, a O(1,, Amax = @(1), K e(1), (23)
where ())) and O-) only hides absolute constants.
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 only depend on the second moment the noise (cf. Assumption A(5)) instead of e.g. its sub- Gaussiaai,,,,,, and also makes the result convenient be directly translated to result for transformers.
Proposition F.1 (Guarantees for in-context least squares). Suppose distribution P satisfies Assump- tion Then as long as 2 logK1 (1/g(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 2 By
do²
ED.
(clip (WIIs.x++))) - UN+1)² S inf Lp(w) +0 R²8 + 2
weRd
N
(e(W))
(24)
(b) We have P(Ecor Ew) 2 1 66/10, where
Ecov Ecor (D) = { La 3 Ep
1-122 1122'p 1/2 2Ld
(25)
80do
Em. ED) := |Wrsll2 S B
(26)
8N Xmii
Proof. We first show P(E 211 1 5/20. Let Let E: = 1 D. and noise variables be denoted as
xxx, and let the whitened covariance
N
N
x, = =
1/2 X4,
X Σp
1/2
N
Also let 2, := y: (x; wp) denote the "noise" variables. Note that
Ecor 1
exactly covariance concentration of the whitened vectors {xi)ie|N)- Recall that E[x,X,] La- and x; are with probability at least 5/10 that
by assumption. Therefore, we can apply [85, Theorem 1J. we have
+ log(1/))
log(1/6) })
Σ
- L op
K2 max
N
N
Setting N' 2 O(K'(d event we have
log(1 ((g(((+))((1)))) ensures that the right-hand side above is at most 1/2, on which
1 I = S = I = 21
(27)
2
41

---

Lec
Ecov holds. This shows that P(E & 6/10.
Next, show (24). Using Foov, we decompose the risk
E
(clipa(( 1N+1)211 ( +E
(III -yx 1
(2)
S
E
- - + +2R² (6/20)
2
(1)
Exxx++
[
(WWLS - wp.xN+1)1 ((
+ Exx-
2 ((Wp, XN+1) - ( [Boor)]]]
+O(R'0)
2
$ ED
2 WL.
will E + Exx+I 2 (WD, - -N+I +0(R28).
Lp(w))
(28)
Above, (i) follows by assumption that lyN+1| < By R almost surely, that removing the clipping can only potentially increase the distance in the first term, and the square loss is upper bounded by 1 (2R)² almost surely in the second term; (ii) follows by the fact that Ex [[ wyp.XN' ((wp, XN+1) yN+1)] = 0 by the definition of wp, as well as the fact that { [Fev)] is independent of (NN+1;UN+1).
|Ww wp E. 1(Ecov] }}. Note that on the event Ecov, we have
It thus remains to bound D
D 1D = "p,:/"E>''
2L
Therefore,
1
|WLS w 1(Ecov)
(xxx) -xx y wP) Ep Ex -X y- 1
2
2
XX(XXX
X
-XX
2 1 (Fooo
2"
1
1/2
--
E2 1/
1/2 1-1/2x
2. 1(Eoov]
2N2
N
2
Dp
1/2
x
2
1
1(Ecov1 }
2
Σ
x,4
1(Ecov
2
Σ
2
N2
N2
N2
i=1
Note that E[x;2;] Σp 1-1/22[xx (over D), we get
(1/
(wp. x;) = Therefore, taking expectation on the above
ED
1
WL.S
well
1 (FFar
5
2
E
Σ
x
2 R
2
1
2
[=
E
-
2
N2
N
N
(29)
20 N
2d
E
x
Ep x1]
(30)
N
Above, (i) follows by conditioning on x, and using Assumption A(5). Combining with (28), obtain
) < LP(wp) + 0 R28+ 0 [
do? do²) N
E
(cip
This proves (24).
Finally, we show Ecov n Ew) 2 1 - -8/10. Using (29) and Ep 2 Amin-a by assumption, we get
4do
E
WLs - w 1 [Ecor) } 5 N NAmin
42

---

Therefore, using an argument similar to Chebyshev's inequality,
20
4do2
P(Ecor O Fee) = E 1(Feor] 1fll
V
+ B
N NAmin
20
4do²
< E
Bov } x 1i/WLS - well2
8
N Amii
|WLS wpll2
$ E
(rrco) x
S 6/20.
2 N Amin
This implies that
P(Erov n E = ((Eew)) ) - P (Eeov n (Ew) 2 1 6/20 6/20 2 8/10.
This is the desired result.
F.3 Proof Corollary
The proof follows by first checking the well-conditionedness of the data D (cf. (5)) with high probability, then invoking Theorem (for approximation least squares) and Proposition F.1 (for the statistical power of least squares).
First, P satisfies Assumption A, by Proposition F.1, long as N (d log(1/6))), have with probability at least - 6/10 that event Evor O E., holds. On this event, we have
1 Ep = X X/N 3 2Ep I 2dmax
2 Amia Is 3
2
d
Wrssl22 S B=/2:: O
B +
JNAmin
and thus the dataset D is well-conditioned (in the sense of (5)) with parameters Amin/ /2 2Amaxx and B defined above. Note that the condition number of Σ is upper bounded by 3/a = 4 S 4k, where the upper bound on the condition number of Ep as in Assumption A(c).
Define parameters
d
do²
8
A1.
(31)
N
B&N
Note that B S O( + / BB//minn by the above choice of
We can thus apply Theorem 4 in the unregularized case (A 0) to obtain that, there exists
transformer 0 with maxte ( < 3, |||| < 4R + 4/Amax (with R max ( (B, Bm, By, 1 }), and number layers
L 0 log
BBB E
< 0
klog
B
N
B* +
B²
do²
V
such that on Eeov Ev (so that is well-conditioned), we have (choosing the clipping radius in ready = clip,, (r)ady(-)) to be By):
(TF2 (H) clipa, (WLS xN+1)) < (TF3 (H) (WLS. XN+1)| S / do?
(32)
We now bound the excess risk the above transformer. Combining Proposition F. and (32), we have
ready read, (TF)(H)) yN-1)
43

---

(ead, (H)) )N+1) 1fEcor Ew! + +E (TF) UN+1) 11(Ecor n 2
E
Eww
21E
read, (TF (H)) c ((WLS)XN+1))) 1fEoor n Eww]
+ 2E (clipe, (.LLS.XN+1) wN+1) 2 1BEcor n Eut + +2B 8/10
(
do
$ 28 + Lp(wp)) + O
5 +
+ 0 (
N
do² O
da
< Lp(wp) + o Bg
N
N
Above, (i) uses the approximation guarantee (32) well Proposition F.1(a) (with clipping radius By). This proves the desired excess risk guarantee.
Finally, under the canonical choice of parameters (23), the bounds for L, M, |0| simplify to
L. S o
llg
NK
max ( S 3, |0| s O(Vkd). EE[L]
(33)
and the requirement for simplifies to 2 Old + log(1/5)) (d) (as = ((1). This proves the claim about the required N and L.
F44 Proof of Corollary 6
Fix parameters E > 0 be specified later and large universal constant Co. Let us
a = max 0.112-vady v
3 = 25,
log(4/0) d
B = 1+
By Co + +0),
B, = Cu dlogggN(N/)
By Co o) +o)ooollgg Vlog(N/8).
Consider the following good events (below [CIKEEN N is given by e, = y1 (W., x;))
E || // S Big /|El22 S 2V
en = {a S. Amin (X X/N) S Amax (X XX/N) s 8}).
E = (Vi E [N], |xil2 S Bz, lusl S By),
No, {||xN+1ll ll2 B,, 1yN+1| S By),
and we define E = & n En n & n E1, Under the event E, the problem (ICRidge) is well- conditioned and ||w wwidgell S B 2 (by Lemma F.1).
Therefore, Theorem 4 implies that for at 3FX there exists T. = [2k log Bw/e)] + 1-layer transformer with prediction TN+1 read, (TF (H)) (clipped by By). such that under the good event E, we have UN+1 clip,, ((x+11, w)) and | widgell S
In the following, we show that 0 is indeed the desired transformer (when g and ò is suitably chosen). Notice that we have
E(ūN+1 9N+1)² = uN-1) E 9N+1)?),
and we analyze these two parts separately.
Prediction risk under good event E. We first note that
((N(E)(NN+ - wN+1) = (cEpp ((XN+1, W)) - 1)++))
<E(1(E)(XN W) -VN+1)21
44

---

where the inequality is because UN+1 E By Byl under the good event E. Notice that by our construction, under the good event &, w = w (D)) depends only on the dataset D'. Therefore, we have (D) - w (ndge(D)] S e as long as the event &o == & n & N& holds for (w., D). Thus, under Eo.
E [1EE)( w) UN++)) W+; DJ
E [1(E)((xN+1 (())) UN+1)² w+, D]
E [(xX++1, (D)) UN+1) W., D)
E [(XN++1, ((D) (XN+1. W., w+)) w., (w. + +
[D(D) -w w.1 +02
and we also have
[W( - -w.12 - w - w-ll2 + Wridge -w. (| - w Wrigells + ((() - W widge l1
w
ridge
W. 11
2g
w Wrido 1 +
Recall that 2BayesRisk, Ew,, D pl definition. Therefore, we can conclude that
w.11 + o². Note that 2BayesRisk, S + o² by
[[L{EHON+1 - 2 < 2BayesRisk 2 Ee
Prediction risk under bad event E Notice that
([veE*(N++ yN+1) & VP(E)E((N++ MN+1)9..
We can upper bound P(E*) P(E; U & U U8 U (gN41) by Lemma B.I, Lemma B.2 and the sub-Gaussian tail bound:
PPE U ( S 8 4'
(PE))) S + exp(-N/8), P(E)) S 2exp(-
Thus, as long as 2 12g((1//)), we have P(89) < 6. Further, a simple calculation yields
E[ +1+ )d & SE0 + 8EyA - 8B3 + 8Eyh+1
Notice that #1+1 |w N(0, w. 112 + hence Eyi++1 3E([Iw, +0)+ = +0t) By. Thus, can conclude that
E 11E UN+1) S 4V8By.
Choosing E and
Combining the inequalities above, we have
E(ŪN+1 yN+1)² < BayyssRisk, + [2Ev RisyssRkk,,,, + + + 4V5B,
To ensure JE +1 1N+1)² 2 < BayesRisk, +s, we only need to take d) so that the following constraints are satisfied:
1
min (E, VE).
AVOB S 2'
N > 8log(12/6).
2
Therefore, it suffices to take 5 (N)
22 I+o
2 for some small constant co. then long
o2 +
N C log
+ C.
our choice of and d is feasible. Note that < 01 (1(++-3)) and hence under such choice of (c. 6).
have L = (log()/))) and |0| = δ (Vd). (vd) This is the desired result.
"e need this, as on E the transformer output at this location could in principle depend additionally on 41 as (15) may not hold due to the potential unbounededness of its input. A similar fact will also appear in later proofs (for generalized linear models and Lasso).
45

---

Lemma F.1. Under the event Ex N Eww, we have | - (B +
Proof of Lemma F.I. By the definition of w ridge and recall that do2/N, we have Wrdge
(X + d
Therefore, we only need to prove the following fact: for any > 0 and 3 = (X x +d +d we have
| - B +
(34)
We now prove (34). Note that we have
| [X X + drL,) x (Xw, + + - BIlllp l/ + [ llop |
where B x X X + da B2 = (X X + dyla) Note that [ llap S 1 clearly holds, and under & we also have ||||| - 2v No. Therefore, it remains to bound the term |B2lop
Consider the SVD decomposition of X = UEV, , Σ = A1. Aa). and U E R V E Rdxd are orthonormal matrices. Then B2 V (Iº + dyI,) - and hence
A ||<= + d>Iaa) 1 max X + d
B2 lop
When N S 36d, directly have [B lop $ (d7) 1/2 S 3(N7) 1/2
Otherwise, we have
N 2 36d, and then for each i E [d]. X; 2 v Amin X X) 2 VaN 2 VN/3. Hence, in this case
also have |B2llop S max; A s 3N 1/2 Combining the both cases completes the proof of (34).
G In-context learning of generalized linear models
As 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 link function that non-decreasing and C²-smooth. We consider the following convex empirical risk minimization (ERM) problem
WGLM := arg min T.x wERd
1
Σ
El {(xx, .)
(ICGLM)
N
where (tt, == -tt+ (s)) is the convex (integral) loss associated with g. A canonical example of (ICGLM) is logistic regression, in which (t) = aug (t) == (1 + e is the sigmoid function, and the resulting (t.. y) (t. v) = yyt + log(1 + e²) is the logistic loss.
The following result (proof in Appendix G.1) shows that, as long as the empirical risk Σy satisfies strong convexity and bounded solution conditions (similar in Theorem 4), transformers can approximately implement the ERM predictor ((xN+1, WGLM)), with WGLM given by (ICGLM).
Theorem G.1 (Implementing convex risk minimization for GLMs), For any 0 < < B with == & B > 0, Br > O, Now == Lg B2/2/ a + and < B2/2, there exists an attention-only transformer TF2 with
L [2x log L Bz Bz/e]] +,
max M <Ō(CR
||o|| < o (R+
BE(T]
== (where = sup, I (t)), R == max (B. By, 1). and Cg > 0 is a constant that depends only on R and the C smoothhesss of g within [-R, Rl), such that the following holds. On any input data (D, (N+1) such that
< Amin (v2 < Amax [x & B for all w E B2(Bw),
|WGLM]l2 I Bw/2, 12 (35)
46

---

TF (H upproximately implements (ICGLM); We have h (L+1) N+1 = [xN+1: W; 1; 1;11, where
1N++ - [( S.
In Theorem G.I, the number of heads scales as (1)/e") as opposed to ( as in ridge regression (Theorem 4), due to the fact that the gradient of the loss general a smooth function that can be only approximately expressed as 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 con- structed in Theorem G.1 achieves desirable statistical power if the in-context data distribution satisfies (w,x,,,. , y)] denote the corresponding population risk for any distribution P of (x, y). When P is realizable by generalized
standard statistical assumptions for learning GL.ss Let T.p (.(w)) == (x) [f((w, x),
linear model of link function 9 and parameter / in the sense that Ep yy |X] = ((3, x)). it is a standard result that (3 is indeed a minimizer of Lp [42] (see also [6, Appendix A.3)).
Theorem G.2 (Statistical guarantee for generalized linear models). For any fixed set of param- eters defined in Assumption B, there exists transformer with L S O (log( (N)) layers and maxte(1] M( < Ö (dºN). such that for any distribution P satisfying Assumption B with those pa- rameters, long N O (d). that outputs UN+1 (TFoH)))) and w (TFo(HH) Id (for another read-out function readw) satisfying the following.
(a) w achieves small excess risk under the population loss, i.e. for the linear prediction DN+1 = (XN+1. W).
E E(D [[[ WN+1)] Lp()) 5O)
(36)
(b) (Realizable setting) If there exists a B E Rdd such that under P. E(ylx] g((3,x)) x)) almost surely, then
E (DDNN-INHH) [[N+1 - 3N+1) s S0xx:1aw=n)~pPll((t xw41)) - yN+1) + Od(/N),
(37)
equivalently, [ B[yv+1 S 0 (d/N).
Above, O (-) hides constants that depend polynomially on the parameters in Assumption B. Similar as in Corollary 5, the (d/NN) 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 E - B,, Bul. g' (t) 2 pgg > O.
We also assume that for each i E [N + 1]. (xi, v1) is independently sampled from P such that the following holds.
(a) Under the law (x. y) (G(Ky) Vw E (B2(BB.) (.
~ P, We have x ~ SG(K2). y ~ SG((2) and g((w,x)) ~
(b) For some pz 0, it holds that
[[[ w| S Bp/22xxx]] Mnla Vw E B2(B..).
(c) For 3 = arg min Lp, it holds 1,,1 S Bw/
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. uub-Gaussian input.
Corollary G.1 (In-context logistic regression). Consider any in-context data distribution satisfying x~ N(), y (0,
arg min Lp (B) (2(
BER
For the link function Tlog and B = (1), we can choose Bwg Bys Mg: Log Bay Kn, K,= (1) so that Assumption B holds. that case, when N 2 O (d), there exists a transformer 0 with L O (log(N)) layers, such that for any P considered above,
47

---

(a) The estimation w (TFa(H)))(H)) outpulted by O achieves excess risk bound (36).
(b) (Realizable setting) Consider the logistic in-context data distribution
plog X N(O, La), y|x (e
B
Then, for any distribution P = P with || S B the prediction ÜN+1
(Tady(TFo)) (H)) of 0 additionally achieves the square loss excess risk (37).
G.1 Proof of Theorem G.1
Let us fix parameters En and > 0 (that we specify later proof).
Define R max By, 1) and
g (s)())
Cg
max im0,1,2
R' max
sE
B,
By Proposition B.I, g is (Egs M, R, C*) with
C OC₂)
M 2 log(1 + Cge
Therefore, we can invoke Theorem D. to obtain that, long 27 S B 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
[x; v WT, 10p-00p2--0 1; 4x), such that
T GD
Eg (LB Bz),
where (EEL) is the sequence of gradient descent iterates with stepsize 9 and initialization W&D = O. Notice that Proposition B.2 implies (with ń := B/aa)
B 2
|woD WoLM|l, exp(-T/(2x)) ||WLM/|2 exp(-T/(2k))
Eo.
Furthermore, can show that (similar to the proof of Theorem D.1 (b)), there exists single attention
layer ((T+1) with M heads such that outputs hx+1 NEN+i:NN41 wT: 2p-00pp-:33l 1;0 where ØN+1 - 9((xN41, wT)) I Eg.
(T+1)
In the following, we show that for suitably chosen (T, (0(1:T), the desired transformer. First notice that its output h (T+1) MNN+1iNN+1 W :0 1; satisties
[iN+1 o((XN+1, W I 10N+1 g((XN++1 w ) Lg ( W (XN+1, WGLM))
$ Eg + Lg | @ l1 + Ly wha WGLM||2
< e + Lg B2 T3- Ba) + L. Byc..
Therefore, for any fixed > 0, we can take
T [2x (Lg. Bw) Eg 21 + B23-1))
so that the 0 we construct above ensures [ūN+1 - ((xN+11 WeLM)) I for any input H that satisfies (35). The upper bound on ||0|| follows immediately from Theorem D.I.
G.2 Proof Theorem G.2
We summarize 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 Co and constant C₁ C2 that depend only the parameters (Kz, Ky, y, B,s B Pezy Lis Mg).
48

---

(a) As long as N 2 C1 d. the following event happens with probability at least 1 2c 22--2~/00 N/C1;
Ew=
Vw E B2(Bw) glglte S mmi (w)) S mmaxx ( (W())(w)) S 81 K2
1
(b) For any > we have with probability least 8 that
V TwT.N (w) V B(.(((N())
- CoKKKa max {
de ydu+ (og(5)
de + log 1/0)
Estat
sup (BB)(BB))
N
N
where we denote = log(2 + Lig LgK2Bu.. BwyKy)
(c) Condition (a) holds and N 2 C2 d, the event & {llWelMll2 S Bw/2] happens with probability at least NNC,2
(d) For any w E B2(Bw), it holds that
4 (estat + [[]
Lp(w) Lp(3) S
Mglh
(e) (Realizable setting) As long as WGLM B2( holds that
La Mxllg
(x WGTNN) 0((x,3))) S
stat
Therefore, set
MgH4
B 8Tog K?
8
B₂ = Co (log(N/8)
By= = Co v log( (N/8.
Consider the following good events
E, -VVi E [N], |xxll2 S Bz, luil < By): E, [||xN+1||2 Bz, lyN+1| By): E = E, E n N& N+,
Under the event & and choice of a, 3, the problem (ICGLM) is well-conditioned (i.c. (35) holds).
Theorem G.1 implies that there exists a transformer & such that for any input H of the form (3), TF, outputs h' [xN+1; +1; w: 10-00-2- 1; 0]. such that the output is given by fN+1
(TF (H)) = clip,, (iN+1) and w = read (TF (H)) == (B.) (W), and the following holds on the good event x.
(a) WN+1 fD(xN+1), where fo A(D) is predictor such that [D((x) 9((x, WGLM)) S E for all x E (2(
w (DD) E B2 (B.)) depends only on D (by the proof of Theorem G.I and Theorem D. such that VŻN I1, B.
A5
In the following, we show that 0 constructed above fulfills both (a) & (b) of Theorem G.2. The bounds on number of layers and heads and llll] follows from plugging our choice of B, By in our proof of Theorem G.
Proof of Theorem G.2 (a). Notice that under the good event E, we have w w (DD) depends only on D. Then we have
E(D (N++))
-((..v+I (1 #N+1]] + (ED
[1(8))((XX+11 E(D.NNIL. L (1 (x))((((x+1, UN+1)] 1) E (OO
Thus, we can consider & = E, E n Eg. and then
E (p (4E)) w(D)) BN+1)]
49

---

= xDxx IN [1(6))(((xx+1 W(D) - 111f - E)e((xN+1, ,(D)),
= I ( - 1 - ((D)),
where the second equality follows from (,(w)D) B: W(D) ,N4+1). Therefore,
E E(D ((N++: [,(1(E))L,
= EEDENNLLNNN
-B B - E]e((xN+1 W(D) vN+1)l
-
B
- 2 / P(E9) max { (1 (((( UN+1) = 0
N
where the last line follows from Cauchy inequality and the fact P(ES) O (N-1 and Bg is defined in Lemma G.I.
Notice that by Theorem G.3 (d), we have
(L inf Lp]] [Bl=2 + +E
[
Myfte
and by Theorem G.3 (b) and taking integration over O, we have de
de +
Elestal S (1) K
Also, we have inf Lp - (3()) ) 2 Br by Lemma G.I. Therefore, we can conclude that
[K2K3 d Mglt N
Kt Pglts Bu
B² N5
E(D (1UN+1::N++)) S inf Lp +0
+²+
+
K2 d completes the proof. B N
Taking 8² S
Proof of Theorem G.2 (b). Similar the proof of Corollary 6, we have
E(ūN+1 IN+1) [[ MN+1 ) E[IEE" MN+1)
< <E[L<E)ON=1 DN+1) + EP(ES)E(NN+ 1N+1)4
where the inequality follows from #N+1 E [-By. By] on event E. For the first part, we have
E B[1[8)Cn+1
-NN+1)² 1
E ((E)(D(xX+)) - -NNH1)) )
Be)(((())-)))]] < <ED 1(6)} (x [111 S B,)(D((x) - =]
where we use the fact that the conditional distribution of (XN+1. BN+11|D agrees with P. Thus,
E 1[1(E)(N+4 )N+1)²) E B(x,)-(((((xx33) y)²
<ED 118)) (E(x.)) [|xx| S Ba)(f((x)) y)² ((xxxi)(((((xx3))- w)²]
Ep (188) Ex1 | I (x)(f(()) ,((x, [3))81 ,((x,,))))
< 2D ((E8} 1 < (x)(((x) (((x, WGLN)))) 2ED [((8} Ex(g((x, woLM))) - g(x,,(((,)))) (x,,3)))²
2hg 2 < 2 + (1) KgK3K3. d
< 26 +
Pzfy
|xH/ N
For the second part, we know P(E 0 (N- 10) and
E 1NN-1199N+1) < SE + 1 -0
In conclusion, we have
L K3K3K d
Ba N5
IE@N+1 2 E6yN+1 ( 2s +
+ O
Taking e² < Ly K2K d B2ll N completes the proof.
50

---

Lemma G.1. Suppose that x ~ SG(K2), ,G(Kg), and w is (possibly random) vector such that [w]] I B. Then
[[[((x, w), 1/4 < (Lg (.KBBBl) =: By-
Proof. Notice that by our assumption, la(0)) S 2 Therefore, by the definition of E,
[e(t, v)1 - -yt yt + S/" g(s)ds s (((((()) - v)1 + 1., (o(s)) s(0))ds S |M] (2K, + (ul) + g..
-
The proof is then done by bounding the moment by I lvl s O (K)) and R || w) S || O (VāB...)) which is standard (by utilizing the tail bound of Jub-Gaussian/sub Exponential random variable).
G.3 Proof of Theorem G.3 (a)
We begin with the upper bound on Amax A (V2LN(w) following event
By Lemma B.3, as long as N 2 Co d, the
Ew,o :
x
S 8K2
happens with probability at least 1 - exp( (lcxpp-NN/C) By the assumption that sup Ig'1 S Lg. it is clear that when & holds, we have Amax (V - KL Vw E Rd.
In the following, we analyze the quantity Amax L (w)). We have to invoke the following covering argument (see e.g. [85, Section 4.1.1)). Lemma G.2. Suppose that V is e-covering Sd with E [0, 1). Then the following holds:
1. For any x d symmetric matrix A, ||Allop S I-2e maxvev /v T Avl and
Amin (A) min v Av 28 ||A ||
2. For any vector E Rd, ||x||2 I=g maxvev xv,
Notice that
1 V2ĪN(w)= = N
Σ g'((,,,,)))xx
My (,,([(w,x.)) S By)x xxxx
N
i=1
N
(1
Mg
x,X,
N
B₁
1=1
Therefore, we can define h(t)= = (B, -18)+ (which is 1-Lipschitz function), and have
Hg 1 By
v²īx(w)=
h((w b((,,x))xxxx
N
=:A(w)
In the following, we pick a e--overing V of gd-1 such that |V] < (3/sw)d d (we will specify &v later in proof). Then for any w Bz(B,))
Amin (A(()) 2 min T A(w)v - 25, |1A(w))lop
vew
By our definition of A(w), we have (for any fixed Brv)
N
min A(w)v= = min VEV
A((w, (v, (v)x
VEV
N
51

---

N h((w, xz)) min { xi) B3 ,
min vEV
=:li(w))
min B + min (Ur(w) B[Uv(w)]] veV
vEV
By Lemma G.3, we can choose Bre = K2 (15 + log ( and then [[Uv(w)] 2 B2/ Thus, combining the inequalities above, we can take Ev event Em,
128 in the following, so that under Mx
Xmin V2LN(w))
MgHte
Mg ( Bulle
max (B]Uv(w)] (w)))
-Uv)))
8
Ba
16
VEV
In the following, we consider the random process {Uv(w) U.(w) -(U which is zero-mean and indexed by w E B2(B.). For any fixed v. consider applying Proposition B.4 the random process (Uv(w)) We need to verify the preconditions:
(a) With norm p(w, w') //wl///l-- Log.N (w. d (2Ar/)) with constant A 2;
(b) Let f(x; w) == h((w, x:)) min { xv) BB BB.)) then [f(x s By and hence in SG(CB,, for any random x;
(c) For w, w' W, we have [/((w x,)) - h((w/((ww)))))) < (w w', xx)x. Hence, because SG(K2), the random variable h( (w, x) h((w' x)) is SG(CK, | w (12), and the random
variable f(x; w) - f(x; w') is (G(CK, K2 B2
|| //
w l12).
Therefore, we can apply Proposition B.4 to obtain that with probability 8g. holds
dlog(2ng) + log log (1/6g)
sup Uv(w) S C' B,
N
where we denote Kig = 1+ + Bw/ Bu, Setting 6g 5/ [V] and taking the union bound over v E V, obtain that with probability at least 5,
Uw(w) S B, dlog(Srig/E)) + log(1/8)
max vEV 1l
sup
where we use log | I dllog/4///// (4/8)). Therefore, we plug in the definition of ev and B.v to deduce that, if we set
(
166C'BB
B2 2
Ey = 128K2
C1 =
log(Sng/Ev),
Bes = Kg + log(K(/H22)
Pz
then as long as N 2 Ci- dd, it holds maxvev B[Uv(w]] - Uv(w) < 1 - exp(-N/C.). This is the desired result.
BaBr
with probability at least 16
Lemma G.3. Under Assumption B. for Brd = K2 + (og(K2/Hun
holds
inf wEB.(B...VEs--1
[[[1[x T w| S B,,/2((x v)2 [[llxxv]] |vl S Bav H] 2 3pz/4.
Proof. For any fixed v
E[1[|x w| < (x 1) vl & Bzu)l
= =E[ii[x w| < (x,/2 - w < Bx v, vl > Bru
2 Mr B[(x 1) [vl Bru.l. H.
Because x ~ SG(K)), ), x T SG(K.), ) and a simple calculation yields
E[(x T 11 T > [K,)] < 21 11)ex)(((""
Taking 1 + log( K2/kz) gives [l(( T v T vl > Bavil S un/4, which completes the proof.
D
52

---

G.4
Proof of Theorem G.3 (b)
Notice that
VLN(w)= - E (g((w, 11)xx
In the following, pick minimal 1/2-covering of gd-1 (so | S 55).. Then by Lemma G.2, holds
VLN -EL LL (w)]
max VEV
(VLN(W).V) v) B((LN(()).))
(w)(ww)
Fix v E gd-1 and set = =6/VV. We proceed to bound supw (w)) by applying Proposition B.4 the random process (X) w We need verify the preconditions:
(a) With norm p(w, w') = ||w - w 112- log (6; (r),, p) & & (loo(2Arr0)) with constant A 2: (b) For z [x we let [z; w) == (o((w))))))))) x)) v) (x, v), then (2; SE(CK, Ky) for any w by our assumption on (x, y):
(C) For w, w'w'e E W. we have (g(( x)) g((w'x x)) S by ((w - w', x) Hence, because x ~
SG(K2), the random variable ( x:) )- g((w'( xx) sub-Gaussian in SG(K,Lyl|w w'll2).
Thus, f(z; w) f(z; w') is sub-exxonential in SE(CK w' 112).
Therefore, we can apply Proposition B.4 obtain that with probability 8g. holds
d (log(2xg) + log(1/ōn)
d + log (1/6)) +
sup Xv(w)] C',,,,
N
N
where we denote Ky = = + L. Kg. Setting 60 8/ [ and taking the union bound over v E V. we obtain that with probability at least òd
/ (dllg((0nn)) + log(1/5) d llog(10sy)) + log(1/5)
[x((w) =WWW)
sup SB
VEV
N
N
This is the desired result.
G.5 Proof of Theorem G.3 (c)
In the following, condition on (a) holds, ie. ΣN is a-strongly- and 8 smooth over B₂ with a = /nlgg and 3 = KL We define
W arg min (w()). wEB2(B.)
Then by standard convex analysis, we have
3 < [V..N W) ViN( w- ( S-vixx() w- -w
|| -W3-1 11-
Notice that ( I Estat, conclude that
Iw < 1,31112
Estat
Recall that we assume || 12 & Bw/4, we can then consider &, == {Estat < aB../4). Once
E. holds, argument above yields || Bw, which implies VLv(W) 0. Therefore, arg wind (w). Further, by Theorem G.3, we can set
32aK,Ky
32aK,
C2 max 2c
2t
B
B
so that as long as N 2 C2d, the event &, holds with probability least 1 exp(-NNCC). This the desired result.
53

---

G.6 Proof of Theorem G.3 (d) & (e)
We first prove Theorem G.3 (d). Notice that
(w)) E[g' ((x, xx)xx | = [(x, w)l S B,)xx "] = Hylald Vw (2(Bw).
Therefore, Lp is -strongy connexx
B2 (B.). Therefore, because 3 B2((,.) is the
global minimum of Ly, it holds that for all w E B2 Bw))
1
Ip(w)= (x( -
2p
By the definition of Estat, LL (w)ll I Estat + II (w) and hence the proof of Theorem G.3 (d) is completed.
We next prove Theorem G.3 (c), where we assume that Ely|x| = g 3)) (which implies 3 = B directly) and WOLM (B.).. Notice that
(w) E VVN(wW) (w)
[l(o((x y)x( [x( w)) (((w,3)))xx)
and hence
WVL,(WLLMI WGLM B] [x(((x WGLM)) ((w,3)))) ((x, WGLM) (ww,B)))
[(x((x WGLM)) o (ww, B)
Lg
On the other hand, by the strrog))--ooooooneeeee of Lp over B2 B..) it holds that
(VLL(WGLM). WGLM B)
Hgl2z
Finally, using the definition of WGLM, we have (WNLM)) 0, and hence | L (WGLM))l2 Estat, which completes the proof of Theorem G.3 (e).
H Proofs for Section 3.2
H.1 Proof of Theorem 7
Fix Av 2 0, 8 0 and B 0, and consider any in-context data D such that the precondition of Theorem 7 holds. Recall that
Liasso (w)
Σ ((w, v1)² AN ||/]
2N
denotes the lasso regression loss in (ICLasso), that Wlasso arg min (w). We further write
N
L (w) =
Σ (w, x.) - vi) R(w) := AN ||w|]]
2N
Note that (0 X X/N and thus Σg 8-smooth over Rd. Consider the proximal gradient descent algorithm on the ridge loss
1-1 PGD prox,, RR (WeGD 7/ D3 (wPG))
with initialization wPGD : Og, learning rate 77 = 8-1, and number of steps T be specified later. Similar to the proof of Theorem 4, we can construct a transformer to approximate woD- Consider ((s, () = (s - 1)2 and R(w) = An [|]] then 0,((s, 1) is (0, +ox, 2, approxiaaalee by sum of relus (cf. Definition D.I), and prox, nR is (0, +oc 4d, 4 - 21jA aaproxibpio by sum of relus (Proposition 1. Therefore, we can apply Theorem D.2 with the square loss f, regularizer R. learning rate 7 and accuracy parameter (I to obtain that there exists transformer TFa with (T + 1)
54

---

layers, number of heads M = for all & E [L], and hidden dimension D' = 2d, such that the (2) final output h N xx++: N1+ wFGn: *] with ØN+1 (wPOD. X XN+1). Further, the weight matrices have norm bounds ||O]| I 10R + (8 2
By the standard convergence result for proximal gradient descent (Proposition B.3), we have for all t 2 that
Linsso ( (WPGD) - Lissso (Wlasso)
S
W l1
21
Plugging in ||Wasssoll 1l2 S B and T 1 -- [BBL/=
finishes the proof.
H.2 Sharper convergence analysis of proximal gradient descent for Lasso
Collection of parameters Throughout the rest of this section, we consider fixed N 2 1, An =
V pu for p 2 0, 2 0 fixed (and to be determined), fixed 0 < S B, and fixed B > 0. We
write = 3/a, == (( and WN slogd N
Here we present sharper convergence analysis the proximal gradient descent algorithm for Tassso 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 well- conditioned with sparsity if the following conditions hold:
1. The (a,, P--RSC condition holds:
Xw 1
log d
2 a wwl2
[w]
Vw E R
(38)
N
N
Further, Amax (X X
X/NN 8
2. The data (X,y) is "approximately generated from s-sparse linear model": There exists w. E Bdd such that ||w=ll2 I Big |wwllo I and for the residue E = Xw.,
x |T& S - NAN.
3. It holds that N 2 32 s slog (i.e 32WN $ 1j.
Assumption Cl imposes the standard restricted strong convexity (RSC) condition for the feature matrix X E xd and Assumption C2 asserts that the data is approximately generated from sparse linear model, with bound on the L norm of the error vector X "e. Assumption C is entirely deterministic in nature, and suffices to imply the following convergence result. In the proof of Theorem 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')azo on loss function Trasso with stepsize 1) 1/3 and starting point w O, we have Lrasso (wT) IAassso (inseooWnsseee (Winsso) S f.r all
3(B + K. log C.
3(BB) 1
T C
+ K
where C' is a universal constant.
The proof be found in Appendix H.4. Combining Theorem H.1 with the construction Theorem we directly obtain the following result as corollary.
Theorem H.2 (In-context Lasso with transformers with sharper convergence). For any N, N,d, 2 1, there exists a L-layer transformer TF, with
0 8, >
L = ((,,, + n(log(Crn/!! +
max M (6) < 2, EE[L]
max ( < 2d,
[)
55

---

||8| - 3+ R < (8 2
such that the following holds. On any input data (D, (D,xN+1) such that the (ICLasso) problem satisfies Assumption C (which implies Wassoll 2 S B2/2 with B 2 vv/a). (H) approximately implements (ICLasso), in that it outputs UN+1 = ready (eady(TFF(H)) = (XN-1, 11:1. W) with
Laasoo W Il (Wlnsso) S.
H.3 Basic properties for Lasso
Lemma H.l (Relaxed basic inequality). Suppose that Assumption C2 holds. Then it holds that
2 (Iinsso (w) Linesso [issso((.)) (w.)).
|w w.lli S ///3/ w. AN
Vw E 12d
As a corollary, ||Wlasso w- S 4 Wlasso w-ll2-
Proof. Let us first fix any w E Rd. Denote A = w w.. and let S -uup(ww,) be the set of indexes of nonzero entries of W.. Then by definition, y = Xw, + E and |S| S s, and hence
Xw y12 |ww, y12
|XA
=1 ||e||
| 2 2 XA.
w|| - 1|wall1 =<( (/w[i] [w,l + Σ
wwiill
jES
j&S
2 w[]] w-[ill Σ |w]]]= 1As=
A slll
JES
igs
Combining these inequalities, obtain
XA
5
X
AN(I - ( + Liasso Liaoso (W..
2N
N
An
||| ||1 + Ax All - ||A s- (12) + P Telassoo ( (W,) (39)
2
AN
(3
A A ||As=|1)) + Elasso(w) Diasso (w.);
2
[x
where the second inequality follows from XA
N || and our assumption that S AN. and the last inequality is due |||| Aslli + As Therefore,
2 ||
have
2 (Einss(w) Erassso (w.))
||| || + 11 5411Asl11 + AN
2 (Itasoow)) (w) P.lassso Coossoowws (w.)).
I <</s| + An
where the last inequality follows from ||sll S /sl/Asll2 S /s || This completes the proof of our main incquality, As for the corollary, we only need to the definition that Liasso (Wlasso) <
L (w.).
Proposition H.1 (Gap to parameter estimation error). E d
Suppose that Assumption C holds. Then for all
s + 'v''aap
||w - w.112 s C
"gap"" + gap
a
where we write gap = Liavso (w) L (Wasso)), and C 120 is a universal constant. In particular, have ||ass w.ll2 < 10
logg a3
56

---

Proof. We follow the notation in the proof of Lemma H.1. By (39), we have
AN
0 S 2N
3 A - | 1 + Diaano (w) Lasso (oosso(W,))) (w.),
2
and hence || I
4/s A
l +
2gap due to (lasso(w) Liasso (wosso(w,, (w.) gap. On the other hand, by
the RSC condition (38), it holds that
XA 2
logd
a
N
N
Therefore, we have
log d N
a P
SAN
A || + 2gap
log d
2gap
S 3ANvV
P
4 ||
+ 2gap
AN
5
20s log d X N
20 log d gap + 2gap:
a
where the last inequality uses AM-GM inequality and Cauchy inequality. Notice that p 20s N < ga, now derive that
30s
120 log N gap" + 12gap.
A
Plugging in An Wlasso in above proof (hence gap 0).
par log N
completes the proof. The corollary follows immediately by letting
Lemma H.2 (Growth). It holds that
X(ww
Wlass)
>1² Liasso (w) Diasso Wlsaso);
Vw.
2NV
Proof. For simplicity we denote Wlasso Wlasso that
By the first order optimality condition, it holds
OE N
(Wiissso y) + OR( Wisss),
where we write R(w)= = AN |w/] Then by the convexity of R, we have
R(w)- - R( (Waass))
2 ( Wi w Winsso? = T (XWlssso - y), w Wlasso
1
X Wlasss y, (Xw -y) (X XWiasso y))
N
1
Xw y12 + 2N
XWlasso
y11 + 2N
[X(w Wlasso 11
2N
Rearranging completes the proof.
H.4 Proof of Theorem H.1
For the simplicity of presentation, we write Wiasso WInsso and we denote gap" : Lassso (w²)
Liasso S (Wlasso).
By Lemma H.1, we have l/w' walli I 4/o//w" w.llz + 2gap" Av
which implies
2gap
ww' Wiasso || w' w=ll, + Winsso walls 5 4/8 |w' Winss ||2 + 8/8 Wassso -.
An
57

---

We denote nn = log.d =22logd Using the assumption that X is (c. we obtain that
N ||x(w' Winsso) )| 2 /w' Wiasso 1|2
J. /w' Waass 1?
w' Wiasso 112 MN 20s Iww" Wiasso l1
640s W
w.
40
(gap')
X
Thus, as long as N
30p" log d we have
w Winsso
1
1 N
I((w'
Winssn 11 640spuN || w.112
40pI (gap")²
X
S 2gap' + 40v
(gap")² "gap")" + 640spM Wiasso - wall2
where the last inequality follows from Lemma H.2 and the definition of AN PAN-
We define Estat := 640spN | holds that for & 2 To.
w. 11 To 103 | 112. By Proposition B.3(3),
B
gap $ I 2
Wlass l1 I
W/ssso
12
27b
20
Then for all t > To 1, we have (the second
below uses Proposition B.3(2))
w r+1
Wlasso
12
4gap"+1
+ Estat
2,8 (||w" Wlasso
1
r+1 Wlasso 2 /))
Estat
3
+ >||w/t/- Winsso 112
35atat
1
a
|w'
Winsso 112
35stat
633
Therefore, for t 2 To 1 1,
l|w² W Winsso 1/2 S cxp
123 (t [To] + +11) |w
a
| Tol
3t
W Wiasso
a
a (t Tu)
||Wlasso|| 11
3estat
x cxp
83
a
where the last inequality follows from Proposition B.3(2). Further, by Proposition B.3(3), we have (t - To) ||
3 w' - Wissso
1 -
8
(-
3Estat
gap
<
exp
To
1, 2
2k
2k
83
Hence, we can conclude that gap
< 5 for all T such that
3 |
12 3testa + 1.
T 2 103v
| 1|2 + 8x log
Now, by Proposition H.1, it holds that [|Wlasso - w.ll2 112 < 10 pv llogd
and hence
N
20pus log
W Wasss
< [|w 112 + Wlasso
w.l < 2( 2
a²N
Plugging in our definition of
nN plog d
p s log WN
Estat
400sJN
Winsso
w
a N
completes the proof.
H.S
Proof Theorem
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), classical result in the high-dimensional statistics literature.
58

---

Proposition H.2 (RSC for Gaussian random design). Suppose that X
[x x E R
random matrix with each row x, being i.i.d. samples from N(0, ).
Then there are universal
constants g 50 such that with probability at least 1 I- N752
N/82
Xxwl2 N 2 e [! - -2p(p)
log d N
[wlli
Vw E R
(40)
where p(E) max Ea the maximum of diagonal entries of E.
Fix a parameter 61 S å (which we will specify proof) and large universal constant Co. Let us set p= o
a = c1 = (1), 3 = 8(1 (d/N)), B Cov dl log (N/6), By Co +og
Similar to the proof Corollary 6 (Appendix F.4), we consider the following good events (where c = Xw. y)
& Amax X X N) S 3 and X is ().
3, {|XT E Ell= 40 log( 4d 4d/6)
& = (Vi E [N], ||xill2 S Bz, lui| S By).
& = (1|x++1 l12 S B2: 1 S By): and we define & == E n E, n & n N,,++1
Furthermore, we choose v > that correspond to the choice AN = S0 log(Ad/6) assume N 2 32c2 s log d. Then, Assumption C holds on the event E.
and we also
N
Therefore, we can apply Theorem H.2 with pwN, which implies that there exists a L-layer transformer such that its prediction UN+1 == (TH so that under the good event & we have 7N+1 clip,, ((xN+1, w)), where
Lisson (W) Linsso (Wasss))
MWN-
In the following, we show that θ is indeed the desired transformer (similarly to the proof in Ap- pendix F.4). Consider the conditional prediction error
E. (ØN+1 + |] E ((N 11 - ynv 1 [1{E9 ([v+) - )
and we analyze these two parts separately under the good event Eg == En n E, n Df
Part We first note that
E ([NE) - UN+1)2 D E ((E)(llip,, ((XN-1. W)) - yN+1)² D)
EE [11E)((XNNEL w) }
where the inequality is because UN+1 E B., By] under the good event E. Notice that by our construction, under the good event E, w = (DD depends only on the dataset D (because it is the (L-1)-t iterate of PGD on (ICLasso) problem). Applying Proposition H.1 to ((D) and using the definition of wN and our choice of AN, obtain that (under Ex)
>> +wwy + VWN 0 (22s (0% log
log(d/8) (0'sloog(/0))
(D) w.11 C
a²
Therefore, under Eo. E [1{8))(XN+11 UNN+)22
(1 W - +N+1)² D]
< E [[(XN+1, (D)) 3N+1)
E [(xN++1, wD)) (XN+1. w.)) +
= [[)(D) w. w.11 00²
s $log(d/0)
0² 1 +
N
59

---

Part II.
Notice that under good event Eú, the bad event E" holds if and only if E holds, and
hence
IE [1(8") - MN+1)² P] = EE [1
S PVC 1 9N+1)*
With a large enough constant Co. we clearly have P(E),N+1) & (61/N)) Further, a simple calculation yields
B(ÜN+I UN+1)* S SE(UN+1 + v\+1) S 8 8EyN+1 16
y.
where the last inequality is because the marginal distribution of yN+1 is simply N(O, 0² + l| 112). Combining these yields
8
B
oi((BB))) + o²) log((161)
E HON+ UN+1)
so
S
N5
N
Therefore, choosing 61 = min Br
is enough for our purpose, and under such choice of 81.
a²
E (1EE))N D 5o
N
Conclusion. Combining the inequalities above, conclude that under &o.
log(d/0)
[[NN+I - yN++ ) DJ - 1 1
remains to show that P(Ep) 2 - 6. By Proposition H.2, Lemma B.2 and Lemma B.4, we have
8 P(ε,) S
P(E.) < 3 exp((N//32)
P(E,) 4
2'
Therefore, as long N 2 32 log(12/8), we have P(E.) 2 1 5. This completes the proof.
We also remark that in the construction above,
R= R= (B +o)Volo (1+ BB/a)))
which would be useful for bounding |00||
I Proofs for Section 4
L.1 Proof of Proposition 10
We begin by restating Proposition 10 into the following version, which contains additional size bounds on 0.
Theorem 1.1 (Full statement of Proposition 10). Suppose that for
Σ ((f(x;) U:).
[uual
((.) is (x/3, R, M, C)appproximmble by sum of relus (Definition D.I). Then there exists a 3-layer transformer TF, with
KN
max (²) < (M + 3),, [E[3]
max ( < K + K1,
|0| <
+ +3 + +7KR.
KE[3]
DDval
that maps
h, = [+; (x(i); ...
where the predictor F _8d - a convex combination of [fxx (fx) S mink, E|K Livll
7. As corollary, for any convex risk L. (Rd IR) R, ʃ satisfies
L() S mink, E[K] L(Jr.) + max4e[k] [Tva((()) - -[((k) +
60

---

To prove Theorem I.1, we first state and prove the following two propositions.
Proposition I.1 (Evaluation layer). There exists a I-layer transformer TFe with MK heads and ||o| & 3R + 2NKC/ (Dval such that for all H such that maxxxllVl] R. f R. TFa maps
h; [x: #5 (xx); (xx(x);;
h [x: v/3 fi(x;;; (x,); Lual (fi):... (f:
e +1,
where (-a(.) is functional such that maxk [Laal (fa) - Lvai(Sa) Se.
Proof of Proposition L.I. As is (E, R,M, M, apprroxmiaale by sum of relus, there exists a function
9: IR? - R of form
M
((, = Σ
(am + bmt + dm) with Σ
1mml S C, |aml + [6ml + Idml S 1, Vm E (M),
m=1
m=1
such that Sup(e.))([-R.1) 1g(3, f(8, 1)1 S s. We define
[val(f) = [Dvall
Σ
s(f(x))
(X
Next, for every m E [M] and [K], we define matrices Qm,k. Km,k, V mmk E RDD such that for all i.j E [N + 1].
am
fx(x,) y1
6m
(N + 1)cm [Dval
Qm =
dm
Km,thi=
V mn
-DD((K-)-3
2
R(1 12)
where e, E IR is the vector with s-th entry being 1 and others being 0. As the input has struc- ture h. = xx y/: #; fi (xx:::..... fK (x,): 0K+1; 1;t,), these matrices indeed exist, and further is straightforward to check that they have norm bounds
K(N 1)C
max ME(MM].E([
1Qm, S
max [E(K]
K Km,llopp 2
Σ
IV S
D
mc ME
Now, for every i,j E [N + 1], we have
a((Qm, kh, km,)h)) = fim Mj + dm 2R(1 + tz))
=0 (am w X, + bm + dm 1(t, = -1},
where the last equality follows from the bound lam Jam (x,, ( bm + dm S R[(aml + 1m + dm S 2R, so that the above relu equals 0 if < 0. Therefore, for each i E [N + 1] and & [K].
a (Qm,thh; km,k V Mm,
M
TTx Cm (am w X, + bm + dm))
(N+1) llt, 1)en-(x-x)) DDval
(N 11+1)
(ff(xx))
lft, -1e (K-
[Dval
Thus letting the attention layer 0 {(V m,k: Qm,k, Km,k) (m.N)((M)))) have
N+1
h. (H))) = h, N+1
Σ
Σ Co(1mm Km) Vm,
J=1
61

---

N+1 K
hy +
g(f((x));)) 8/) iit -K-))
DDval j=1
h;
+
{(I((x),1s)) -D-(K--)--
(Dval (x
k=1
h;
T (fk) (K-
= xx v/s #3 fi (xx) ).. ... fk (x;); 0 1:4] + [0D-K- 100-K-3:440 Lva (fi): (x)((fx);; [x; v/s #5 fi (x:); ... f (x:); Lvai (fi); [[xa((kk);0;1;;)
[E[N+1]
This is the desired result.
Proposition 1.2 (Selection layer). There exists 3-layer transformer TFa with
max (0 2K +2
max ( < K2 + K K1 |10|| -- + 3K + rE(3
[E[3]
such that TFg maps
h, |#; (x.)): ... if (x,): L1j Lk;0;1;tt)
h = (x.)); (x(xx))
* f(x)) 1:4],
i [N +1:
where f= F=EK A& an aggregated predictor, where the weights , Ax 2 are functions only on -L, Lb such that
Σ NE 1
A only if S min L..- + [[([K]
Proof of Proposition 1.2. We construct 0 which is composition of 2 MLP layers followed by an attention layer (0 012)
mlps ((3)
Step construction of ( We consider matrix w that maps
h= h= [ Lij... ILxi;*;;
h [L, Lo:... [T - Lgj - ;L - LL In; 1a:::-ap -L;;.. Lk; -L,],
Le. w( h is 8² K dimensional vector that its entry contains (The
T1 L.),EK[ and
(LL,, -Loke Clearly, such w(') exists and can be chosen so that (w())
S 2K. We then
lp
consider a matrix ( that maps
o( h) (WWo (0n-K k -33; Lij... ;Cx Lxx Oz] E RD.
where C c (T.) Σizx (L. T.). Notice that
Ca L (TLx)) a--Tx) + Σ
(TL - -T4),
1z
and hence such w!() exists and can be chosen that |W
We set 00!) dlp
S +
(W()), wi("), then MLP a maps h; to
h h(th (x; (x():: (x;) ;ck; 0; 1:
The basic property of (ck]ke(K) is that, if cs I 7, then L2 I ming E(K] L + 7.
Step 2: construction of 0% (2) alp'
We consider matrix w!!) that maps
h [ K [CK; 1;+]
62

---

w h- - -
qiqi 1 -i:
cK cki -ck] E IR
and w (2)
can be chosen that w
K + + x We then consider matrix w
(2) that
up
maps
o (W 2) 1 h)
o b)
(-01 - -ci :0 - -xx) cx. 0gl E ERD
which exists and can be chosen so that |
2. We set ! = (w(2) w then MLP e
op
maps h (())
h (x:);: ... ffx (x;); :. u 0; 1;
where -a & [K]. Clearly, E (0, 1]. and > 0 if and only if ck S.
Step 3: construction of 0 0 We define
A = 1 - a(1 - u1); Clearly, & 2 O, and Ex A 1. Further,
A = a(1 - u1 - uk-1) o(1 u1 vk 2.
A > 0 uk > 0 ch S LL S min Li +7. k+[[K]
Therefore, it remains to construct 0 ((3) that implements f Ex=l Asda based on [h,"),, Notice that
=xx)=o(1) [1xx) + Σ 0 41 ux-1) (fa(x,) - fx-1(x;))
(41)
(01 01
uk) fx((x)),
(3) and hence we construct Datin as follows: for every ls E [K - 1] and E {0, 1), we define matrices
Qk,e Kk,ur, kw IR. DXD such that for all &o E [K + +1]
Qk (2) (Ja(x;) + R)
1
Qu.hh
(2)
( (x:) )xx)+ + R) 1s
0
0
1
U1j
Kg,hh kz
VE.c (2) en - 1v
uk-1
for all i,j E [N+ + 1), where we understand fo fK+1 = and 1 the k-dimensional vector with all entries being 1. By the structure of h (2) these matrices indeed exist, and further it is straightforward check that they have norm bounds
max we(
1Qu S KR,
max KE[K+ (I.0E
|K S1,
Σ
I < 2K + 2.
kEIK (001
Now, for every i,j IN 11,w w. (0, 1}. we have
of ((Ohmuh)))) (2) K Kg (2011 =0((1 ij
Uu ()x v (x:) + R))
=0 U1
uz-1) (fr-w (xx) + R),
where the last equality follows from fa(x;) + R 2 00k E [K]. Therefore,
Σ
((Am,dli Km, (2)) ) mmhhh?"
kEIK =E10111
Σ 0 u1 uk-1) (xs( + R) a(1 un uz-1) (fa-1(x;) + +R] CD-2
63

---

= f(xx)) cD-2:
where the last equality is due to (41). Thus
letting the attention layer 0 0
{(Vi,m Qu K ) [1))))// we have
N+1
hg 3) = Attno ( (H()) h N
, bn2))))))))))) V (2))
k,
h (2) f(x:)) -e-22
b+; f (x:): : fx(x));u : 1xx;;;;;
This is the desired result.
Now, we are ready prove Theorem I.I. Proof of Theorem L11 As :) is (7/3, R, M, -aaappooximmale by sum of relus, we can invoke
Proposition 1.1 to show that there exists single attention layer l so that Attn (()) maps
h2 h = [xs s/i #; fxx):; : fx (x:); (fi): Lial (fk);;
i [N 1],
for any input H = [h;]; the form described in Theorem I.1, and (-) is a functional such that max; [<va(fi)) (fis $ 7/3.
Next, by the proof of Proposition 1.2, there exists (Omip Oup 0 that maps
hg b
x 1/i #; fxx):: fR(x));* *
ie +
where À = (A1.-.. Ax) A([K] and d > 0 only when Eval(fa) S ming. Lvai( (fk-) + 7/3. Using the fact that max Lval (fx) [va((fi) s 7/3. we deduce that X is supported on {K Lival (fx) S ming,e(k] Lfa((fx.) + 7 Therefore, 0 (0 0upp ol2) ( is the desired transformer, with max
M < (M + K, EE()
max (()
< <
[E(3
and
2N
|0] & max
3R +
+ 3K + 1, K + +3 +3+++4 7
KR + 2K +
D
2NKC
7KR+
+y-1
[Dvall
This completes the proof.
1.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, (i (xs; yi) E D.rain) I, (i (x) E D.al), Xtrain [x.]ieT, to be the input matrix corresponding to the training split only, and Nirain = | is Nyal = (Deall-
Theorem 1.2. For any sequence of regularizations (hx ]RE[K, 0 < S B with == max B+A
B,, > 0, O, and < Bw/2, suppose in input format (3) we have D O(Kd). Then there exists an L-layer transformer TFg with
L 12 (.. (2=))
max (0) < 3K +1, EE[L]
max ( < <² K 1
[E[L]
64

---

||l|| S O KR
N
R max B, By 13,
Nval
such that the following holds. On any input data (D, XN+1) such that the problem (ICRidge) is well-conditioned and has bounded solution:
a < Amin X (X Xtrain/ (train) I Amax (Xaai Xtrain Ntraii) S B,
|| ridge (Pirain) S Bw/2.
max kE(K]
(42)
TF3s approximately implements ridge selection: its prediction
ØN-1 = read, (w, (w,))
W
A W
k=1
satisfies the following.
1. For each k [K], We: (Prrain))) approximates the ridge estimator An
ridgo (Drrain); i.e.
We wridoo (D (Drrii)ll2 S.
2. \ (A1, Ak) E A((K)) so that
A > only if Trual Trwal (wk) S min Twai( (w<)) k-E[K]
In particular, if we set set B( + B2e + 7 then it holds that
dist w, conv (
Lval (w ridge,,train)
min E(K)
Lval (WWe traie + 7/3 S E,
ridge,train
where we denote w
dv tiage,train
== w Gooige(Draal (Derain).
ridge
To prove Theorem I.2, we first show that, for the squared validation loss, there exists 3-layer transformer that performs predictor selection based on the exactly evaluated (val(fk) for each k E [K]. (Proof in Appendix I.2.1.)
Theorem I.3 (Square-loss version of Theorem I.1), Consider the squared validation loss
(f) 2 (E..)))DDall
1
Σ
(f(x,) - y1)
Then there exists 3-layer transformer TFo with
2N Duall
max
M (0) 2K + +2,
max
( < K2 +K+,, FE(3)
||0|| < 7KR+
+ +
[E[3]
such that for any input H that takes form
h xx v/s *; fi (x:); ... fK (x;); 0g: 1:4],
where TFa outputs hN+1 = [xN+1: [XN+1); 1; 0), where the predictor J R.d. - a conver
combination of [fk Lvai(fe) I mink, [v.)) + 7.. As corollary, for any convex risk
L: (RRd R) 1R f satisfies
L(f) < mink, L(fe.) + maxke[K]
[ua L(fx) +
Proof Theorem I.2 First, by the proof Theorem 4 and Proposition B.6, for each E [K]. there exists T L layer transformer (1:7) such that TF FFgur) maps
h
|x 1/1 # (W1.x)); i.... -.... (WK, x,) 0x; 1;t],
'This because (wa((w) is B... "Technically, an adapted version where the underlying ICGD mechanism operates on the training split (with
Lipschitz
B (B......
t 1) with size Nuaa instead of on all N training examples, which only changes ||0|| by at most a constant factor, and does not change the number of layers and heads.
65

---

so that if (42) holds, we have
Wx
w
ridge
2 S
and We (2
Next, by Theorem 1.3, there exists 3-layer transformer 0(++1:++3) that outputs
[xN+1i (W, (NNN ; 1;4], N+1
(T+3)
where w = -k=I d. A (AI,... ; E A((K]) so that
A > 0 only if Lvai( (wa)) S min [wai(Wi-)) +7- k*E(K]
This is the desired result.
L221 Proof of Theorem I.3
Similar to the proof of Proposition 10, Theorem I.3 is direct corollary by combining Proposition I.3 with Proposition I.2.
Proposition 1.3 (Evaluation layer for the squared loss), There exists an attention layer TF' e with 2K
heads and |O]| < 3R + 2NK/ [Dvail such that TFg maps
h; = (xx); .. 5 (x:); 0x; 1;44]
h (x:); ... fK (x+); (Ji) 11
i [N + 1].
Proof of Proposition I.3. For every E [K], we define matrices Qm,k Km, V Vm, E R DxD such that for all i [N + 1].
Sx(x;) Dj
-1
Qc,hh;
Qu
Kt,oh, Kr,hh,
-2
-
R11+ + tz)
0
(N + 1).
Va,hh = Vv,ih;
(fx(x)) U/ U;)eD-(KKAA--
2 [Dvall
As the input has structure h, = 3/: *; fi (x,);; ... ; fK (xx): 0k+1; t,), these matrices indeed exist, and further it is straightforward check that they have norm bounds
K(N + 1)
max wE(
1Q 3,
max kE(K]
K Kg < 11+1 R.
Σ
V llop <
Duail
kE[K),
Now, for every i,j [N + 1], we have
Σ a(QQe, V
wEf0,11
(N (N+1) + 1)
=[o(f((X;)) - yj -2 + t2))
o01; f4(2)) 2R(1 ty))]
(fa(x,) yn-
D
(N + 1) Dvall
= =1st, -1} (x())( - vi) - oly; fx(xx))
(fa(x))
(N 1)
= 1 11 2 D.ail
(fa(x;) 2e))
where the second equality follows from the bound (x,, - Myl < 2R, so that the relus equals 0
< 0. Thus letting the attention layer 0 = {(V Qk KE (1. we have
N+1
h. (At)) = h N+1
Σ
o1Q KR Va
k.
66

---

N+1 K
h; + vall
- -2)) 1ft = -IJeD-(K-A)-)
j=1
h;
+
Σ (fa(x)) yj)² (D-(K-A)-3
Daa
k=1
(K):D))EDal
h;
T (fk)) -D-(KK-))3
= xx uls #3 (x:): 3. .. fx((x;); 10000k1:11 1it] + (0D-K-35 Lval ;
[N[++] [x v/s #5 fi (x:); (o ;[xa(([k):;;1;44] ie IN + +1].
This is the desired result.
L. Proofs for Section 4.2
L3.1 Proof of Lemma 13
is straightforward to check that the binary type check vp : -> can be expressed as a linear combination of 6 relu's (recalling o-) ReLLU()))::
(1)
y
20
y
o
y
8)
20
y/ yy1
+ 0
y (1
8
am o (bmi) + cm))
m=1
with Em laml = 8/E, max, max [16ml: |cml) S
We can thus construct an attention layer
0 ((Qm. Km, mm)) mm with 6 heads such that
N +
Qmh; [bmi Cmi OD-2). Kmhy- bys 1:0p2),
Vmh;
am tzi 15:10111
N
which gives that for every i E [N 1].
Σ
a((Qmin (1, K
N
+
m-
JE[N+1]
N
σ (bm. + Cm cm)a
Σ ((5)) = y binary (D).
N
N
m=1
1=1
Further, we have ||0|
< 18/= (1/E). This is the desired result.
By composing the above attention layer with one additional layer (with heads) that implement the following function
a (2 1/2)) o(2(t 1)).
on the output (D)) we directly obtain the following corollary.
Corollary L.1 (Thresholded binary test). There exists a two-layer attention-only transformer with maxtetg] ( and [|o|| < O(1/e) that exactly implements the thresholded binary test
1,
if binaary (D) 2 1,
T binary (D) =
O,
if T binary (D)
< 2'
(43)
linear interpolation,
O.W.
every token i E [N + 11, where recall the definition of y binary in Lemma 13.
67

---

133.2 Formal statement and proof of Proposition 14
We say a distribution y on R is (C, notcoonnccneee around {0, 1) if
Pu([=e,e]U[L - E, +]) S Ce
for all E en] A sufficient condition is that the density is upper bounded by C within [-Sn, So] U En, + so].
Throughout this section, let (tog(t) == (1 + -- denote the sigmoid activation, and let Wlog denote the solution to the in-context logistic regression problem, i.c. (ICGLM) with g(t) Gog( Proposition I.4 (Adaptive regression or classification; Formal version of Proposition 14). For any B S B2 Bww/11 0 S B with attention-only transformer with
B/a, and any (C, Eo), there exists I-llaee
Bi E
max M(9) S so <0((1 [E[T]
)
(1+
B4 a2 )=))
R 1 -++++
L L 100 1:
101 S O
(with := max (B, Bw, By, 1,, and depending only (C, such that the following holds.
Suppose the format (3) with dimension D 2 3d +
On any classification instance (D, xN+1) (such thut AUiheENI C (0,1)) that is well-conditioned for logistic regression in the sense of (35), outputs UN+1 that s-approximates the prediction of in-context logistic regression:
ON-1 - -log ((XN+1, Mogi))
On the contrary, for regression problems, i.e. any in-context distribution P whose marginal Py is (C, (oo)=-tconncetrrutee around {0, 1}, with probability at least 1 - exp(-cN) over 'D (where e >> O depends only on (o), WN+1 e-approximates the prediction of in-context least squares the data is cell-conditioned:
[N++1 - (XN++,,WL)] S whenever DD satisfies (5) with À O,
where WLs denotes the in-context least squares estimator, i.e. (ICRidge) with =
Proof. The result follows by combining the binary test in Corollary I.1 with Theorem 4 and Theo- rem G. By those results, there exists three attention-only transformers oLS, Ouog Osin. with (below Ly:: Cg (1) for Olg(.))
LAS < 0 log
B3 B
M S 3, FE[hus
max
|10.s|l
5 R
1) 8)
Lg B.B
L2B4 a
110 s 00(
Cg.
Toog S O
log
max
log < 0 02
1
2
R
L Liog]
Lbin 2, max Mll) bin
S 6,
10 S ((1/E),
Jog that outputs prediction 4 NN respectively, which satisfy
(at the (N + 1)-th token) and y binaary (D) (at every token)
lng UN
Tlog((XNa WoRL Se,
(XN+1, WiS) S..
when the corresponding well-conditionednessss are satisfied. In particular, we can make Wlog well- defined non-binary data, by multiplying binarry (D) onto the x, 's (which be implemented by slightly modifying Olog without changing the order of the number of layers, heads, and norms) so that Wlog on any data where v binanyy (D) 0.
By joining LLS and Dgo using Proposition B.6, concatenating with 0 before, and concatenating with one additional attention layer with 2 heads after to implement
binary thres
(D) +
vi binary (D)
DN
(44)
68

---

we obtain single transformer 0 with
B Bi
50((1
B a²
|0L - s0 R 1 1 B x+ 1
L LS
1: log
max (0) RE[L]
c
which outputs (44) as its prediction (at the location for (N+1).
It remains to show that (44) reduces to either one of olog UN++1 or JK-11 When the data are binary (y) {0. 1)), we have b 1 and yp shinesy (D) = 1, in which case (44) becomes exactly glg
UN+I- By contrast, when data is sampled from a distribution that is ( -oot--ooocncettttteee
around {0, 1}. we have for any fixed S Eo IC that, letting Bg == [-,, s| U [1 e, + 5) and
Pe == (Be) s Ce s 1. by Hoeffding's inequality,
1 =
N N
My E Be 2 5
P( binary (D) * 0) = IP y binarry (D)
2 2
S cxp -
S ) cx)(-c'N),
where d' is an absolute constant. On the event binaary (D) 0 (which happens with probability
least 1 c'p(-c'N)), (44) becomes exactly 7 This finishes the proof.
1,4 Linear correlation test and application
In this section, we give another instantiation of the pre-ICI. testing mechanism by showing that the transformer can implement a linear correlation test that tests whether the correlation vector E[xy] has a large norm. We then use this test to construct a transformer to perform "confident linear regression", i.c output a prediction from linear regression only when the signal-to-noise ratio high.
For any fixed parameters Amins B 0, consider the linear correlation test over data defined as
V lin
(D)
||012 (Amin /4)² [10] (3Amin (3 4
X min (22)
0, where E T(D)=
| S (Amin / | 2 (3Amin /4)
1,
(45)
linear interpolation,
o.w.,
1
X,Yii
N
Recall that (--) R-LU(-) above denotes the relu activation.
We show that winn can be exactly implemented by 3-layer transformer.
Lemma 1.1 (Expressing glin by transformer). There exists 3-layer attention-only transformer TF F with at most heads per layer and ||o|| S O(1 + Xann((")) such that on input sequence of the form (3) with 2 2d + 4, the transformer exactly implements ii,, it outputs H. such that Xx yet #; y (D); for all i E [N + 1].
h
Proof. We begin by noting the following basic facts:
Identity function can be implemented exactly by two ReI.Us: L = a(1) - --()).
Squared f2 norm can be implemented exactly by a single attention head (assuming every input h; contains the same vector g): ||ll = (gg.g)
We construct the transformer 0 as follows.
Layer 1: Use 2 heads to implement f 1 MN x where v(1) (1.21 h, [=xy; Op-a), Q/2 h (1,2) h, = [ = b < N + 1}:0 (where we recall t, Hj < N + 1} and note that y,ts corresponds exactly to the location for yj in H. cf. (3)).
[ N+1 -p--), and K (1)
69

---

By manipulating the output dimension in V write the result è into blank memory space with dimension every token E [NV + 1].
Layer 2: Use a single head to compute || Q P = [t; Op-al. K h2) = ft: Op-al- and
V (1) = [1;0p-1] By manipulating the output dimension in V write the result |ft 12 into
blank memory space with dimension at every token i E [NV + 1]. After layer 2, we have h [x;; Us4;; || 1]
Layer 3: Use 2 heads to implement two ReLU functions with bias: [
-- A)
<[[0]]2 B)). The two query (or key) matrices contain values and B. our problem we take
B)).
A = (Amin /4)2, B = (3Aainn B
so that the above ReLU function implements Tlin (D) exactly. Write the result into blank memory space with dimension 1. We finish the proof by noting that ||0|| - O(1 + A (B
Statistical guarantee for Tlin We consider the following sell-posedness assumption for the linear correlation test t 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 wp, well as sub- Gaussianity conditions.
Assumption D (Well-pun poseeessss for linear correlation test). We say a distribution P on Id x R is well-posed for linear independence tests, if (x, y) P satisfies
(1) ||x||2 S B2 and lyl S By almost surely;
(2) The covariance Ep := Ep (xx"] satisfies Lain 3 Ep 3 AmaxLa. with < Amii S Amaxx and k := Amax/ Amin-
(3) The whitened vector
1/2 Ep!/2 x is K2-sub--Ce Suub-Gaussianor for some K 2 1.
(4) The best linear predictor w == Ep Xxx [ [1-EE[[xj) satisfies |wpll S B"
(5) The label y is a2- -subbGGausiinn
(6) The residual == - (x, wp) is FF-sub-Gausiian with probability one (over x).
The following results states that in achieves high power as long as the sample size is high enough, and the signal l/ is either sufficiently high or sufficiently low.
Proposition 1.5 (Power of linear correlation test). Suppose distribution P satisfies Assumption D with parameters Amins Amax, B* Then, for the linear correlation test yiin with parameters (Amins B
with B < B and any 2 Ö maxx {K*,
Amar } d) we have
(Bu)aa X2
I. If w 2 B then with probability at least over D, we have v (D) = 1.
2. If w <
Bi then with probability at least 5 over DD have Tlin(D) 0. 10Amax
Amin
Proof. any P satisfying Assumption D. note that E[x²] Exx(y- (wp. 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 6:
N
0.9E
3
E
x
I
1.1Ep
as N 2 O(dK*) by (27),
N
1=1
1
N
X124
S
1/2
N
Σ P
1/2
X
< 1122
Kovd Kod)
VN + N
))
N
N
1=1
1=1
d
Amin 1"
δ Ammam dk2
<
max
Koy
N
8
a8
N >
BB 22
70

---

On the above event, we have
N
N Σ
|
1
Σ
x
(x,,,
wp) + 2,)
= Ewp +
X
N
N
1=1
2
Therefore, in case 1, we have
| 2
Ewp
N
x;
2 0.9Amin ||wll12
Amin B 8
m
30min B
N
4
2
In case 2, we have
Amin B
Amni B
Amiu B
f - Ewpl
+
N
X
Amax
10A
+
8
2
The proof is finished by recalling the definition of qlin in (45), so that qplin (D) 3mmin B//4 44, and (D) = if S Amin B
if ||t|12 2
Application: Confident linear regression By directly composing the linear correlation test Lemma I.1 with the transformer construction in Corollary (using an argument similar as the proof of Proposition I.4), and using the power of the linear correlation test Proposition I.5, we imme- diately obtain the following result, which outputs a prediction from (approximately) least squares T == //nn (D) = 1, and abstains from predicting if & = O. This can be viewed as form of "confident linear regression", where the model predicts only it thinks the linear signal is strong enough.
Proposition I.6 (Confident linear regression). For any B > 0, < B S B 0 OS Xmin S >maxx 0 < a with := B/a, there exists a L-layer attention-only transformer with
< B,. 100,
BrBw
1 max (0 < 0(1), 101 = 0 R+ Xmin (B²)) EE[L]
L S O Kllo
(with R = max BB, By, 1}) such that the following holds.
Let N
δ mmxx {K*, Amnr K20"
d). Suppose the input format is (3) with dimension D 2 2d + 4. Let ICL
(BB)) instance (D1xN+1) be drawn from any distribution P satisfying Assumption D. Then the transformer outputs 2-dimensional prediction (within the test token hN+1)
(UN+1, 3) E {0,1}
such that the following holds:
1. If w 1| 12 2 Big then with probability at least 1 8 over D, we have 1@N+1- -WILs XN+1) S 6, and v/ 1 D is in addition well-conditioned for least squares (in the sense of (5) with X 0).
2. If |Iwlle 0|12= S
Amin B then with probability least 1- over we have UN41 and 10Amax
Proof of Theorem 12:
Noisy linear model with mixed noise levels
For each fixed l E [K]. we consider the following data generating model PA.. where we first sample and then sample data ((x,, y1 bi)hee(N41)
P Pw..o ~ T from w, ~
iid P P as
P Pw.. : X; N(0, Ia), y: = (xi, w.) + Et, & ~ N(,,of))
Also, recall that the Bayes optimal estimator on P& is given by UN+1
Bayes
w Wridge (D), (N+I) with
ridge X odd/N, and the Bayes risk on Ps is given by
Bayes == inf E. (A(D)(xx= UN+1 ) = Ex
Bayes
(UN+1 -N+1))
71

---

Recall that in Section 4.1.1, we consider mixture law Pr that generates data from Ps. with & ~ A. It clear that we have (pushing inf into ER-A does not increase the value) we have
BayesRisk 2 R [BayesRiskg],
i.e. the Bayes risk can only be greater if we consider mixture of models. In other words, if transformer can achieve near-Bayes ICL on each meta-task Ps. then it can perform near-Bayes ICL on any meta-task which is mixture of Ps with k ~ A. Therefore, to prove Theorem 12, it suffices to show the following (strengthened) result.
Theorem J.1 (Formal version of Theorem 12). Suppose that N 2 0.1d and we write max max; (kk 1}, min = minx(ox: 1}. Suppose in input format (3) we have 2 d(Kd). Then there exists transformer 0 with
I, S O (a (omin (og(N//mmn))
max M S (K)), RE[L]
max D S O(K2), EE
||O|| < O (ogg())
such that for any & [K], it holds that
o22 max
log K 1/3
Ex
(uN+1 - DN+1)² - BayesRiski + +
2/3
N
0 min
if we choose Nval = (Dvall x N2/3 [log KJ2/3.
The core of the proof of Theorem J.1 is to show that any estimator w that achieves small validation loss Eval must achieve small population loss.
Throughout the rest of this section, recall that we define Nirain P .Nvl [Dvall Itrain (xa, y1) E Derii Ival = {i: (x4, y1) E Dvalh, and Xtrain = [X]]iEInan'
J.1 Proof Theorem J.1
Fix parameters 5, 0 and large universal constant Co. Let set
max 0. 0/1
dd Ntrain
3 25,
log(N) d
B" = 1 Co
B = Co + Fmax minn)
B₂ Cov dog((N,
= Co + maax) / log(M) N);
Then, we define good events similarly to the proof of Corollary (Appendix F.4):
E = f|/wall2 S B ||E|| S 20 max VN}.
E {a S Amin X Extrnnn X NNraii) S Amax (Xirain Xtrain/ /Neaii) S 3).
Et {V(x,, 8:) Dtrain ||x:112 I Ba, 1yvl S By}:
Ea [V(x,, v1) Dval. | - Bz, 1Bil S By).
& {||x~+1ll < Bz, WN++1 By).
For the good event ε Ex n & Eb,train O Eb,test O E, Further, by the proof of Lemma F.1 (see e.g. (34)), we know that maX || |w (Dtrain)) s Bw/2 holds under the good event E.
show that IP())) 5
doj
For the ridge λ Nuvis and parameters (a,, 3, 7,e), we consider the transformer 0 constructed in Theorem 1.2, with a clipped prediction ΦN+1 read, (HF)
In the following, we upper bound the quantity (ØN+1 1N+1)² )² for any fixed k. Similar to the proof of Corollary (Appendix F.4), we decompose
Eg (ūN++ - ) = E. (1 1 + Eg (18))(N+ =
and we analyze these two parts separately.
72

---

Part I. Recall that by our construction, when & holds, we have ON-1 - clip (w+,))+1) and
the statements of Theorem 1.2 hold for w. Thus, have
Es (1NE) - UN+1)
Es ((xN),
BN++)
< EE <E&[ W) UN+1)2].
Let us consider the following risk functional
- - = =111 (Iww- w- + +
Twall,w (w) = P(x.,)>Pw.. Pw.. on
Then, under the good event Eo- = E
n n n Es,train n Edttest of (w.: D),
Ex [188((xN+1) W) 11)-+++] 1N++)² w., D Ex [1(E)((NN+1, ((D)) - yN 1) w., D)
I Ex [[(xN+I, (DD)) -NN+11 W., D]
= =B(x.v)~Poooe [(Xx++,, 1N-
-vval.w. (w(D))
By our construction, under the good event Eg. we have
( (D() < Lva..w. (Wk (P (Ptrain) + max
Lval (wr Lvll,w. ( (Drrinn)))) +%
[E(K]
where
w/ (Perainn)) w weidge (Derain) 1| S for each 1 E [K]. Clearly,
2E [1 (1 Lval,w. (wr (Derain) )) = Ex [1E0} [W (Derain)) - -. 12: + oî
< Ex ({Eb} (1) [[Wwridge (Derain) - w.112 1122 + 2s w
ridge (Dtrain) -w.ll2 + +5²)
+ +
< E&
Wride (Dtrain) - w.ll2 + 22 2:|l ridge (Derain) w.ll2 + +=² + to
< 2Risk + 2Risk train +
where denote 2Risk = E6 w whidge (Dtrain) w. w.l12 1/22++ + oz and we also note that Riskg S 1 + of by definition. By Lemma J.1, we have
Risk I
BayesRisk; + O (0 1)
Nval
N
(Wi We next deal with the term Eval == max(E[A] Evvll W/ (Derain) ( (Wi((Drain)) the good event Etriin & n Ew tt,train of (w., Derain), have
Note that for
[x I Ex [etrin I 1Ew, Desin [1[Etrain] ED.. [Evall W., Derain]].
Thus, Lemma J.2 yields
log K Nial
log K Nval
(16) S (B(B))
+
Therefore, we can conclude that
o o Nual
log K Nval
B log K Nval
E. (1N8) ) < 2BayesRisk + + E0 + +
+ B²
+
N
Therefore, we can choose (e, Nval) so that Nial S N/2 as
B2 {
2/3
Nyal = max
N)
log (K), log
Fmax
2
max
N
It is worth noting that such choice of Noai is feasible as long as N > we obtain
B
log K. Under such choice,
mmax
log K"
1/3
E. [1(E)ON++
U < BayesRisk, + +0 0 p2/3
p2/3
2
N
73

---

Part II. Similar to the proof of Corollary 6, we have
B y
(EN 1 So
0
max
N
N4
Conclusion. Combining the both cases, we obtain
Ex [ DN+1) v BayesRisk: + +
a Tmax 22/3
log (NN"
1/3
BayesRisk O
Amax 273 min
log N
1/3
+ Fmax
log log"/3 (N) (
d875N175
log 1/3 (K)
1/3
BayesRisk;
log
2/3 min
NV
where we plug our choice of By. The bounds on M(O) D and |0] follows immediately from
Theorem 1.2. This completes the proof.
J.2 Derivation of the exact Bayes predictor
Let (D, (D, MN+1) be (N + 1) observations from the data generating model considered in Sec-
tion 4.1.1. On observing (D, xN+1), 1), the Bayes predictor of UN+1 is given by its posterior mean:
E, ( xN11 E, [xN+1: w.) + #N++ xw = (w.+1...(..(D))
It thus remains to derive Exx [w. [D]. Recall that our data generating model is given by ks ~ A, By Bayes' rule, we have
E, W. [D
E MElK]
P= E.D,, R7).
(46)
On 1 = K', the data is generated from the noisy linear model w. N(O, Ia/d), and Xw, +e
iid where E; N(O, of.). It is standard result that E, /., [D, k = k') is given by the ridge estimator
[, [D, k 16'1
X
X + do
x x =: We
X
X
do
X y
(Note that the sample covariance within Σx is not normalized by N, which is not 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 do
remains to derive the mixing weights P, = k'(D) for all k' E [K]. By Bayes' rule, we have
P. (k (K)) IDD) Ock P* = =k)= / p(w.) Pk',w. (D) (D)w.)dre wd)w.
x Av. [ (2 N/2 exp
Xww y1l2
dw
2
20
X X
X y
|| 12
x Av.
(2.T02)) JN/2 exp
2 w
dIp
w +
w,
dw
O
20
x Av
(2TT02)) )N12
exp
20 (w
wx))
T (w
wx)
20
1
(1x/11 - y "'X X xTy)) x)) dw
det 1 ) -1/2
[|y112 - - x
Av-
exp
NN
20
74

---

Ag- oe det + I.). exp
1 2
[1x112 (y. (pxxxx)))
Note that such mixing weights involve the determinant of the matrix Σn- X + do Ia. which depends on the data X a non-trivial fashion; Any transformer has to approximate these weights their mechanism is to directly approximate the exact Bayesian predictor (46).
J.3
Useful lemmas
Lemma J.1. For trisn = Eg w wht such that
ridgo (Dtrain) - -wall2 + of there exists universal constant C
Nual
train S BayesRisk,
+ +C(ok + N
Proof. Recall that under P&,, we have
W. IaDdda/d).
yo (x,, w.) + E; ~ N(0.002))
We denote yt [yilieZeaain then by definition w ridge (Derain) = (X Xtrain + +do))= -dxa))"Xtrainny" Xt (with
A = do /Nriaaiin). Thus, simple calculation yields
trisk E& w ridge (Dtrain) w. w.112 o of (X (Xai train + do oh.
and analogously, 2ayes = of T X + do 11 + o%. Therefore,
tris 2BayesRisk; = Btr((((XX train + do Ia)- 11) of + do 11
S Lo..NvalBy ( (1)
where in the above inequality we denote Σ = X
train Xtrain + do and use the following fact:
(r(Σー)) - ur(( + + X. ir Σ
1/2 (Ia Σ 1/2 x x ,,,,,
-r
Σ
11112 [I + Σ 1/2 X 1/2
1/, X
((a+ Σ 1/2 X, Σ 1/2)
X 1/2 Σ
=
<rank(E
x X
1/3 Amax E
S Nval (
Case 1. first suppose that Nirain S 16d. Then by definition Σ > do and hence
of Nual (2)-4
16Nval Nerain
32.N
Manin (Mmin
do
Case 2. When Ntrain > 9d, then consider the event & Dmmin X train Xtrain Nerain) > 1 By Lemma B.2 we have TP( I exp( exp(-Nrrain/ (train/8). Therefore,
of A Ammin [Amin E) - of Na Ex [1 Amin (E)-1) En Amin (E) -1
16 Nval P(E.) Nval P(E()
Ntrain
N
Noal exp(-N/16) @0
(o 1).Nal
d
N
Combining these two cases finishes the proof.
Lemma J.2. Condition on the event Etrain- we have
Lvaal wii) - Lval, (wi))
log(2K)
log(2K) Nval
Denin EDag /a/..///ww....... Denin
max [E[K
CCB
Nval
where denote (P
75

---

Proof. We only need to work with a fixed pair of (w., Dtrain) such that Etrain holds. Hence, in the following we only consider the randomness of Dval conditional on such a (w., Drriin).
Recall that for any w, [w) 2 Davll (e
1
Σ
((x,, w) -V:)²
and we have ED, [[val(W) [w] Lval.. (w). For each E Ival+
y1 - (xi, Wi) = Ei (xi, w. - Wi) ~ SG + l|w, - -wil1)).
Note that under Etainn have Wi B2( for all / [K]. and hence of + ||w, w/ll S 5B We then have (yi (xi, Wi) 2 are (conditional) i.i.d random variables in SF(CB..) Then, by Bernstein's 's inequality, we have
PD.. ((Ina) Lval.w. (wi) 2t) S 2 exp (ceNval min
{ 2 1
1
B B
where e is a universal constant. Applying the union bound, we obtain
(
Lval(Wr) (wi) Lval.w. (wi) 2t s K exp (-cNial min
8 t
D max
[E(K]
12 Bw
Taking integration completes the proof.
J.4 Generalized linear models with adaptive link function selection
Suppose that (gk IR + R) is set of link functions such that is non-decreasing and Cº -m-oohhh for each & E []] We consider the input format we introduce in Section 4.1 with [Derain| = [ [Dvall = LN/2].
Theorem J.2 (GLMs with adaptive link function selection). For any fixed set of parameters defined in Assumption B, as long as N 2 O (d), there exists a (ransformer O with L S O (log(N)) layers, input dimension D = θ (dK) and maxte(L) M(0 < Ö Nd, such that the following holds.
For any k- E [K] and any distribution P that is generalized linear model of the link function gk" and some parameter B, if Assumption B holds for each pair (P,9k), then
log K)
((DDXN+L [ØN+1 1 < ((x.))~ [(n ((x, -)) +0
N
N
or equivalently, (ID:XHH))~P((N:: E|yn+1/ (d/N Vlog ()NN)
Proof. For each & E [K]. we consider optimizing the following training loss:
(k) WGLM
arg min train (w) Ntrain
Σ
4x((x,, w) 14).
(X..V)) Drran
where (t, == yt So ( (s) is the convex (integral) loss associated with gk (as in Section 3.1).
Also, for each predictor f Rd R. we consider the squared validation loss Lvat:
(() 2N
1
Σ (f(x,) 31
Fix a large universal constant Co. Let us set
a MyMu//8
B 8LyK2,
B, = Co , (og)
B, Co V (og(N),
76

---

Then, we define good events similarly to the proof of Corollary (Appendix F.4):
Ew { Vkk [K], vw E (2(Bw), a S Amin ( (V2I(K) "train (w)) Σ Amaxe (V2 (()
ttain (w) - B,
&, { Vk [K], [| _ B2/22
Eaarrain = [V(x,, 96) E Dtrain |xx||2 S Buy |vil - Byh,
Eg,val (V(x,, w/) E vall: ||x|||2 < Bz, |//l S By),
E = [||xN+1ll2 _ Bz |N++1 - By).
Similar to the proof of Theorem G.2 (Appendix G.2), we know the good event & := Ew N
Ei,rrain O Esstest n N+, holds with high probability: P(E*) S O (N-10).
Similar to the proof of Theorem I.2, we can show that there exists a transformer 0 with prediction #N+1 ready (TF (H)) (clipped by By), such that (for any P) the following holds under E:
(a) For each E [K], fa = A& (Deran)) is a predictor such that fx(x,) g& ((x, w WGLM)) SE for all E IN + (where e is chosen as in Appendix G.2).
(b) IN+1 clip,, (f(xN+1)), where f ATF(D) an aggregated predictor given by f
Σ. AeJh. such that (Ax) is distribution supported k [K] such that (fa)(()) S minkg[[]] (ffe)fe)) +
Similar to the proof of Theorem G.2, for Eo = E n n Esttrain n Ed we have
B2
BeD (ON UN+1)² < ED>p ((a))/)a(()) + o
N5
where we denote Lval(f) = E( [||lx|1 Bx)(f(x) - y)² for each predictor f. By the definition of f, we then have (under Eg)
Lval(f) S [va((fir)) + max [fal(f)) Lval(fi) +
For the first term, repeating the argument in the proof of Theorem G.2 directly yields that for Eerain == En O &, n Ex,rainn
ED (Mrrra)
S E ((xx)))P(a ((x, B)) 1)²² + O (d/Ntrain).
For the second term, similar Lemma J.2, we can show that conditional on Derain such that Evaii holds, it holds
(viegti K Nyall
1(80 max ([aal(i)) - Inanl() s Ok
log K Nval
+
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 finite number of training ICL instances.
K.1
Generalization guarantee for pretraining
Setup At pretraining time, each training ICL instance has form Z = (H, (N+,),),, where
H(D, x Dx++1) denote the input sequence formatted in (3). We consider the square loss between the in-context prediction and the ground truth label:
4a(0;
UN+1
clip,,, (read, (TFR (H))
read,
77

---

Above, clip,, (t) = max {min (t, B.y.Byy... is the standard clipping operator onto [- Byl- and TF the transformer architecture as Definition with clipping operators after each layer: let H = (H),
H (c) clipR (MLP oH
(Attn en
H (H())))) for all [4], clipR (H) == [Projjpil (hh))]
The clipping operator is used to control the Lipschitz constant of TFa with respect to 0, 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.
We draw ICL instances Z = (H, UN+1) = (D, (XW+1- UN+1)) from (meta- )distribution denoted as T, which first sample an in-context data distribution P ~ и, then sample iid examples (xi, Mi Mi)i=1 N+1 iid P&(N+1) and form D = {(x,, 31) HENN Our pretraining loss is the average ICL loss on n pretraining
instances (1m)) lid n, and we consider the corresponding test ICL, loss on a new test instance:
(0) Σ (a: Z'),
(0) == Ep Z-PoCN+1) (0;
Our pretraining algorithm is to solve standard constrained empirical risk minimization (ERM) problem over transformers with layers, M heads, and norm bound B (recall the definition of the ||-Il norm in (2)):
arg min (0).
0
LLMD = 10=(0attz) (0
01pp)
(1:22)) max (0 S M, max D D', |1l|| - EE[L]
SE[L]
(TF-ERM)
Generalization guarantee By standard uniform concentration analysis via chaining arguments (Proposition B.4; 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) )e(m)) the solution 0 to (TF-ERM) satisfies
12(MD2 DD), + log(1/E)
(θ) S inf LE0L,M.D'S.
(0) + O B²
n
where: L = =log(2 + max {B,R, By)) is log, factor.
K.2 Examples of pretraining for in-context regression problems
In Theorem K.1, the comparator infoce (0) is simply the smallest expected ICL loss for ICL instances drawn from т, among all transformers within the norm ball LL,M, D...B. Using our constructions Section Xdistribution I T's, by which we obtain end-to-end guarantees for pretraining transformers with small ICI. loss at test time. Here we showease this argument on several representative regression problems.
4, we show that this comparator loss small on various (meta-
Linear regression For any in-context data distribution P, let wa Ep [xx -)-'Epp(x)) denote the best linear predictor for P. We show that with mild choices of L,, M, B, the learned transformer perform in-context linear regression with near-optimal statistical power, in that the sampled
P~s x and ICL instance (1); v:)) P. it competes with the best linear predictor wp for this particular P. The proof follows directly by on combining Corollary 5 with Theorem K. and can be found in Appendix L.3.
Theorem K.2 (Pretraining transformers for in-context linear regression). Suppose P a is almost surely well-posed for in-context linear regression (Assumption A) with the canonical parameters. Then, for N 2 Ö(d), with probability least 1 (over the training instances Z (1m)) ), the solution
78

---

of (TF-ERM) with L = On (K)(//))) layers, M = 3 heads, D' = 0 (attention-only) and
B O(VRd) achieves small excess ICL risk over wp:
12 log(1/E)
do²
(ic( E. 2 yy
(wp. x S
+
n
N
where (.) only hides polylogarithmic factors in N, 1/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 Ö /R2dd /n) term originates from the generalization of pretraining (Theorem K.1), where as the O /N) term agrees with the standard fast rate for the excess loss of linear regression [38]. Further, as long as N/o2 the excess risk achieves the optimal rate Ó /N) (up log factors).
Additional examples By similar arguments as the proof of Theorem K.2, we can directly turn most of our other expressivity results into results on the pretrained transformers. Here 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 and Suppose N Os log((d v N)/o)) and let K B///o.
Then with probability at least & (over the training instances Z (1:²)), with L (x2 + d/N)) layers, M 2 heads, D' 2d, and B Oppoly(d, Bw, o)) achieves small excess ICL risk:
). the solution 1 of (TF-ERM)
o²² N d/N)² log(L/E)
+ +0" 2 logd
n
N
where Ō(-) only hides polylogarithmic factors in d, N 1/0.
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-ICI. 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 selec- tion). Suppose is the data generating model (noisy linear model with mixed noise levels) considered Theorem J. I, with maxx S (1). Let N 2 d/10.
Then, with probability at least 1 - (over the training instances Z(mm)), ), the solution θ of (TF- ERM) with input dimension D ((dK), L O D' O(22), and B = Opooyy o min: d, N)) achieves small excess ICL risk:
layers, M = O(K) heads,
Lici 0 BayesRisk, SO O Kz + log(1/E)
log K
1/3
max
2/3
N
min
where (-) only hides polylogarithmic factors in d, N, K, K,1/ 1/0min
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 n, P is almost surely realizable logistic model (i.e. plag with ||31|2 < B as in Corollary G.1). Suppose that B O (1) and N 2 20(d) (d).
Then, with probubility at least 1 - (over the training instances (1:m), ), the solution Ô of (TF-ERM) with L = (log((N) layers, M = Ô Nd heads, D' = O, and B = O(poly (d, N)) achieves small excess ICL risk:
1 (y oog (13, x1)x olog
<ō N + (1E)
d
(ia(0) Eppa E (
N
where ((-) only hides polylogarithmic factors in d, N.
79

---

Remark on generality of transformer All results above are established by the expressivity results in Section 3 & for transformers 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.I). However, the transformer itself was not specified encode any actual structure about the problem at hand in any result above, other than having sufficiently large number of layers, number heads, and weight norms, which illustrates the flexibility the transformer architecture.
L Proofs for Section K
L.1 Lipschitzness of transformers
For any p E co], let H 12.p = (E/=
ii=1 h 12) denote the column-wise (2, pnorm of H. For
any radius R > O, we denote HR == (H:
|H 2,00 I RJ be the ball of radius R under norm I1-ll,,e
Lemma L.1. For single MLP layer Onip
(W1, W2). introduce its norm (as in (2))
|| W| llop
Wellop
For any fixed hidden dimension D', we consider
Qup, == {0.1p's 01 || S B}.
Then for H E HR. aip E mpp, B. the function (0mip: H) + MTP Pat (H) is (iR)-Lipschitzz w.r.t. mml and (1 + B²) L-ipschitt w.r.t. H.
Proof. Recall that by our definition, for the parameter Oapp (W1, W2) E O.., B and the input H [h.] E DDNN the output MLP (H) H ( [b. + W2o (W h:)];-
Therefore, for 62p- (W,, W22) Ompp. have
MLP Omp HH) MLP %a
H
2,00
max
W2 2 W,h,)
W (W(()
= max
< max
W2 Willop ||o( + / (Wo - oWl
max
W2 Willop [W W.hhllp W 2l. W,,,, W/h,112
BR| W2 Wallop + BR| W Willop
where the second inequality follows from the I-Liipchitzzees of a = [+ Similarly, for H' = =hh]] E D
MLP 0m (H)- -MLP 0m1p (H') (H)
2,00 max h. + Wio(W,h,) h W12(W2
H H' 1120000 + max [W[o - o o(Wz
H H' + max B (W2))) - (WWhk)l
H H' 2,,00
B² H H' 2,00
Lemma L.2. For single attention layer Oarnn {(Vm m. Qms Km)) its norm (as in (2))
CRPx, we introduce
| := max mE(M)]
{11 {IQQmllopp
| (smllop) Σ
Vmllap
m=1
For any fixed dimension D, we consider
O == (Ourta ||0 || < B}.
Then for H HR. Outta Ostra,B, the function (0attay H) Attne w.r.t. Oatta and (1 3 Lipsshhit w.r.t. H.
(H) is
80

---

Proof. Recall that by our definition, for the parameter aatnn {(Vm, m. Qm, Km) Im)M E aetta.B
and the input H [h.] E RxN,
the output Attn (H) = [h;] is given by
M
N
h. h
Vmhy.
m==
Now, for 0 Oatta = {(V O K mm [M]
we consider
M
N
h = [ttnm [H)] = h +
a( be, Km,) V E [N].
N
m=1
Clearly
||ttuoue
H)
Attug (H)
112.00 max,
h, - h For any [N], we have
M
h4 h
[a((Qmh; kmb))) -00---o(Q000 ol(Q h, K h,) hy
N
m=
1
==1
|ao h; Kmhj)) 100mmmo(1 o((Q'm hi, K hy)) ol [bb
N
M
N
bbاll2 10 (00(@mm Kmhj))] |/ Villlop
N
+ +0 h,),)) x((Q'm h. Kmhj)) [| llop
+o(1Q'm h, kmby)) - a((Q'm h, K hy)) || llop }
N
Σ
R B²R². IV V llop 1m@h, hillz 1kmh | mmmllop
N
m=1
+ || Q' hilla Km,h K hyll |Vmllap)
M
R B²R² V
Villoop + BR² Qmq O'mllop || 1/ + BR² |mmm
Kn
II Vmllop) }
< R
V
V Vll max ||Qm Qmmllo max Kmn
K llop
RB2 0aten
dattall
where the second inequality uses the definition of operator norm, the third inequality follows from the triangle inequality, the forth inequality is because |/mm S BR, ||Kmh BR, and is 1 1-Lipschitz. This completes the proof the Lipschitzness w.r.t. 6atta-
Similarly, we consider H'
[h]], and
M
N
Attng HH), h
Σ
((mmm Vnhh,,
Vie [N].
N
m=1
1=1
By definition, similarly bound
hi ht) - [n. - hn:)
M
N
[a((2m, h; km, - a((Qm K kmh)) V hr]
N
m=1
V hm hr Km bb a((Qmii Kn,h)) b2ا
m=1
81

---

M
Mmllop [0(@m Km km)) [b - -hil12
m=1
==
+ 1o((Qmh Kn hy)) o((Q h/, Km my)))] [[blll
+ 1o((( Km - -0((-0m1mt))
Kmhj))
Km [bb)
M
N
m op 3 311Qmmll 1mp K la R² R2||hg h1yl2
N
m=1
M Σ
R²
H
H'
2,00
max mE[M]
1mmllop
Km
Ivm
lp
m==
B
2,00
where the last inequality uses 10attall B and the AM-GM inequality. This completes the proof the Lipschitzness w.r.t. H.
Corollary L.I. For a a-fixed number of heads M and hidden dimension D', we consider
TTF,IB = {0- -fattn, 0mpp) M heads, hidden dimension D', |0| s B}. Then for the function TFR given by
TER (0, H) ~ clipR (MLP Bai, (Allno (H)));
O E OTF,I,.. H HR
TFR Lipscchtii w.n. 0 and Lg Lipschitz w.r.t. H. where Be := BR(1 + BR + BRR²) and
Bи = (1 B (1 )
Proof. For any 0 (0attn: Ontp). H HR. and 0' = (Oactar (asp), we have
TF a (H) TF (H) 1,,00 I ||LLPan an (Attng
(H) MLP thiy
(Attngga (H) 2,00
+ |MPP Bui (Attng cita
(H) MLP
(Attngg 2ite
(H))
Kalp
2,00
< (1 BP) Alttn (H)
aeer. (H) + +BR]] Ompp 02p
$ (1 B² attn
0
Oattal
BR Omip 0appl
$ 0 0'1l,
where the second inequality follows from Lemma I.2 and Lemma I.1 and the fact that ||ttno (H) 1,,00 $ R : R B for all H E HR.
Furthermore, for H' HR. we have
(F) - TF. H 112,00 S (1 + B²) |Attno
H Attno H 112,00
S1 B²) BBRR) H')
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
OTF,L f0=
(0 (I:L) atta
o (1:))) M M, D', 1101 < B)
Then the function TFR is (LB Be) Lipschit wint O E O for any fixed H.
Proof. For 0 -(:LL) E OTF,L,B. Ø = @(:L)) E OTF,L.B. we have
T (H) - H
2,00
82

---

TF 0((+1:))
TE pir:
(TER
(H))
TE p(e+sa))
(TF (10)
(TER
m0ue-1 (H))
2,00
BB- 1
TF 010)
(TRR
gisee11) (H)) - TTR
TFR
o(ee-1 (H)
oie)
t=1
Σ
B
Be
Be
|
T())
LB
Be
0
where the second inequality follows from Corollary I.1,, and the last inequality is because By
L.2 Proof of Theorem K.1
In this section, we prove slightly more general result by considering the general ICL. loss la (0: 2) == L( E(red, (H)) WN+1).
We assume that the loss function satisfies sup | - B and sup [a & B/. For the special case {(s, t) {ss t)², we can take B. = 4By B: 2B
We then consider
Xe == Σ (0; 22 - Ez(Lia(O; ();
where (lmm) are i.i copies of Z P. x. It remains to apply Proposition B.4 to the random
process (Xo}. We verify the preconditions:
(a) By [87, Example 5.8], it holds that log N(;; B (.). I.D) s L(3MD2 + 2DD' log(1 2r/8), where B (r) is any ball of radius under norm ||-|]-
(b) Mial (): Z) By and hence BB-sub-Gaussian.
B (LB BB) ||e /|| by Proposition L.1.
(c) 4e(( (0: Z) 4a( Z)
Therefore, we can apply the uniform concentration result in Proposition B.4 obtain that, with probability at least E.
sup Xol S CB.y L(MD* + DD')+ log(1/E)
n
where log(2 + B LB Ba B) BB) - 20L log(2 + max B, R, /BB)). Recalling that
0) S inf (ia (0) + sup Xoo
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 with canonical parameters (and thus in expectation over P ~ тT) and every N 2 Ö(d), it outputs prediction #N+1 read, (TF (H)) such that
(0) (p (oN+1 2N+1)² < Ep [Lp(W)) + O do
N
where we recall that Lp(wP) = (B (1 - {)].. By inspecting the proof, the same result holds if we change TF, to the clipped version T if we choose RD O +B BB +1) Old + k), so that on the good event Ecov n En considered therein, all intermediate outputs within
83

---

TFg has 11-112.00 R and thus the clipping does not modify the transformer output on Ecov n Ew. Further, recall by (33) that & has size bounds
Nk
L S O log
max
M RELL]
33
|0| S ((Vkd).
We can thus apply Theorem K.1 to obtain that the solution f to (TF-ERM) with the above choice of (I, M, B) and D' (attention-only) satisfies the following with probability at least 1 &:
(0' +0
L2MD2 + log(1/5)
(0) S
inf 0'E0L,M.D',
$ (0) + L2MD* + (1/E)
R2d log(1/E)
+ do²
n
N
Above, 1 = O(log(1 + max (Bg, R, B\)) = Õ(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:
((B x (1+ = ((x2 + d/N)):
M ==
D'= 2d,
B₂ Ö(Va),
By = O +0), 8= 0² ByN
10| < B B O (1+ s 5 R log I Öpooly(d, B o)),
where Ō(.) hides polylogarithmic factors in d, N, R
Then, Theorem 8 shows that there exists a transformer 0 with layers, maxte[L] M < M heads, D' hidden dimension for the MLP layers, and ||0| I B such that, on almost surely every P ~ n, returns prediction +1 such that, on the good event En considered therein (over D P) which satisfies P(Eo) 2 1 5,
B(xx=1 (MN+1 ) o O(s log(d/6)/N)]-
By inspecting the proof, the same result holds we change TFo o to the clipped version TF. if choose RD = OB3 B + (B + 1) O(d + (B,,)² r²), so that on the good event E considered therein, all intermediate outputs within TF. has ||-I12. S R and thus the clipping does not modify the transformer output on the good event. On the bad event Efj using the same argument in the proof of Theorem 8, we have
o2 N
Ep, ( ( (ŪN+1 ) (P((E)) (SE
B v) 1/2 <
Combining the above two bounds and further taking expectation over π gives
Lia (0) = (Dxx.
(TN+1 - UN+1)² - 0² 1og0>
We can thus apply Theorem K.I to obtain that the solution ① to (TF-ERM) with the above choice of (L, M, B, D') satisfies the following with probability at least &:
L + DD', + log(1/E)
(i) S
inf O'E0L.M.
Lia (0') + +
L2 + DD') + log(1/E)
< L1a
S 5²² 0o² 1td d/N)² + log(1/E)
s log d
0²
N
Above, O(log(1 max (By, R, B})) Ō(1) This finishes the proof.
84

---

L.5 Proof of Theorem K.4
We invoke Theorem 12 and Theorem J.1, which shows that (recalling the input dimension (Kd)) there exists a transformer 0 with the following size bounds:
O (a min
M = -M))K)) EE
max
max EC
D D' O(K3),
|@| s O [oaax log()))
such that outputs UN+1 that satisfies
_ maxx 2/3
log K
1/3
T - - 0v+1)²
BayesRisk, Ō
min
By inspecting the proof, the same result holds if we change TFg to the clipped version TF. if choose R² = O + B + (B%) + 1) Odd + omax), so that on the good event considered therein, all intermediate outputs within TF, has 11-12,,0 S R and thus the clipping does not modify the transformer output on the good event. Using this clipping radius, we obtain
0 max
log K
1/3
(0) (D, (DNNIINN (ON+1 - WN+1)²
< BayesRisk, + Ô
2/3
N
min
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 - E:
Lia (0") + (
L? + DD') + (lg()/))
Lia(0) S
inf M,,,,,
71
L2(M22 + DD') + log(1/E)
< (0) + +
< BayesRisk, + +0(1
0 min
K3d + log(1/E)
a max
log K
1/3
2/3
N
min
Above, 1 = O(log( + max (By, R, B})) Ö(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 the size bounds (number of layers, heads, and weight norms) from Theorem G.2 and Corollary G.1. O
M
Experimental details and additional studies
M.1 Additional details for Section 6
Architecture and optimization We train 12-layer encoder-only transformer, where each layer consists of an attention layer as in Definition 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 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 D-dimensional hidden state, and the read-out layer maps D D-dimensional hidden state to I-dimensional scalar.
Each training sequence corresponds to single ICL instance with N in-context training examples ((x, 1 C gd x IR and test input XN+1 e gd The input the transformer is formatted as
85

---

in (3) where each token has dimension d + 1 (no paddingd.. The transformer is trained by minimizing the following loss with fresh mini- mini-batches:
L( (Hxx=))) (F
(f (H)), UN+1)]
(47)
where the loss function 4p R 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 fixed learning rate 10 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 (fresh) minibatch with batch size 64 in the base mode, and K minibatches each with batch size in the mixture mode.
For the sparse linear regression experiment, we find that minimizing the training objective (47) alone was not enough, c.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-lass hidden states to recover the true (sparse) coefficient w.:
No. Σ Ep-P (H Ep No
1
||| TF (F((L-1)
w* m/1
Lfit-w (0)
HH)
(48)
J(DDd41):
Specifically, the above loss encourages the first No S N tokens within the second-to-last layer to be close to w We choose No 5 (recall that the total number of tokens is N 10 and sequence length is N + = 11 for this experiment). We minimize the loss L(O) + ALit- (0) 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 yN+1 E {0, 1} and the predicted label 1 [WN++ 2 1/2). We report the means 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. 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 valid algorithm for classification.
Averaging: The simple algorithm which computes the linear predictor w predicts UN-1
N 1i= Vix, and
(W+1)+
3-NN: 3-Nearest Neighbors.
Ridge: Standard ridge regression as in (ICRidge). We specifically consider two X's (denoted as lam_1 and lam_2): λ1,dy = (0.005, 0.125). These are the Bayes-optimal regularization strengths for the noise levels (0), o2) (0.1, 0.5) respectively under the noisy linear model (cf. Corollary 6), using the formula X do /N, with (d, N) = (20, 40).
Lasso: Standard Lasso as in (ICLasso) with À E (1, 0.01, 0.001).
In Figure 2c, the aadge_analyticee curve plots the expected risk of ridge regression under the noisy linear model over 20 geometrically spaced values of X's in between (Ar, A2), using analytical formulae (with Monte Carlo simulations). The Bayes_err _____rrr_1___________________11_ 2} indicate the expected risks of λ on task (with noise σ)) and X2 on task 2 (with noise σ2), respectively.
(1,
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; using past tokens (()),))),e --- as training examples. Here we show that such architectures also able to perform in-context algorithm selection at every token; For results for this architecture on "base" ICL tasks (such those considered in Figure 3a), we refer the readers to Garg et al. [31].
86

---

(a) Lincar regression
(b) Linear classification
(c) Reg vs. cls at token 40
1.0
0.40
select res
0.35
0.8
0.35
dlassiiiattonn enter
leas Seuares
Logibii Ragression
0.30
0.6
0.30
S-NN
0.4
0.25
0.25
0.2
0.20
0.20
20
40
0.15
100
20 30
40
In context examples
examples
regression Square, loss
Figure 5: In-context algorithu selection abilities of transformers between linear regression and linear clas- sification. (a, On these two tasks, a single transformer "F_alg_seled simultaneously approaches the performance of the strongest baseline algorithm Least Squares on lincar regression and Logistic Regressio on on linear classification. (c) At token 40 (using example ......, 399 } for training), TF TF_alg select ..,,,....
matches the performance of the best baseline algorithm for both tasks. (a,b.c Note that transformers pretrained on a single task (TF_reg, _F_cls) perform near-optimally on their pretraining task but suboptimally on the other task.
Setup Our setup is the same the "mixture" modes (linear model linear classification model, and noisy linear models with two different noise levels) as in Section 6, except that the architecture GPT-2 following Garg et al. [31], and the input format is changed to (11) (so that the input sequence has 2NV + 1 tokens) without positional encodings. For every i E [N + 1], we extract the prediction Y; using a linear read-out function applied on output token 2i and the (learnable) linear read-out function is the same across all tokens, similar as in Appendix M. The rest of the setup (optimization, training, and evaluation) is the same in Section 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 TF_alg ael_sellett nearly matches the strongest baseline for both tasks simultaneously. whereas transformers trained
the other task. Further, this phenomenon consistently shows up at every token. For example, in Figure 2a & 2b, TF selectt matches ridge regression with the optimal on all tokens i E (4....,N) (N 40). In Figure 5a & 5b, TF_ala TF_alg_ alg_select matches least squares on the regression task and logistic regression on the classification task on all tokens [N]. This demonstrates the in-context algorithm selection capabilities of standard decoder-based transformer architectures.
single task perform suboptimally
M.3 Computational resource
All 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.
87


Reviews:
[{"review_id": "H1VESy6ERc", "reviewer": "Reviewer_pWto", "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. ", "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", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "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. \n* 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. \n* This works opens up new directions to explore (theoretically and empirically) optimal ICL construction on other problems\n* The paper is well written and easy to follow \n", "weaknesses": "* It would be good to know how results such as those in [1] compares to the one obtained in this paper.\n\n[1] Are transformers universal approximators of sequence-to-sequence functions? a", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback on our paper. We respond to the specific questions as follows.\n\n> How results such as those in [1] compares to the one obtained in this paper.\n\nOur approximation results and [1] are very different. [1] shows that transformers are universal approximators for a large class of sequence-to-sequence functions. However, their result is for generic continuous seq-to-seq functions, and their number of layers is exponential in the worst-case (cf. Section 4.4 of [1]). \n\nBy contrast, **our results are for specific target functions (ICL algorithms) but much more efficient**—The number of layers, heads, and weight norms only depend polynomially on relevant problem parameters and the desired approximation accuracy. This happens as our constructions utilize the special structure of the ICL algorithms (our approximation target). \n\n> The reasoning behind choosing Bayes-optimal ICL on noisy linear models with mixed noise levels.\n\nThere is no fundamental reason behind our choice of noisy linear models with mixed noise levels. Rather, we picked this setting as it is solvable by transformers via algorithm selection, and the setting itself is a harder one than that of Akyurek et al. 2022 which studies a single noise level.\n\n> Was that the most complex method for which something like theorem 12 could be proven?\n\nIf we understand correctly, you mean the most complex “setting” (data generating model)? We believe results like Theorem 12 hold as long as 1) the model is a **mixture model**; 2) Bayes-optimal ICL can be done for each component of the mixture; and 3) $N$ is sufficiently large so that the validation losses are accurate estimations of the true population losses. Under these conditions, algorithm selection (post-ICL validation) can be used to do nearly Bayes-optimal ICL. An example is a mixture of generalized linear models with different link functions.\n"}, {"review_id": "DOX5w4fiRO", "reviewer": "Reviewer_oV7P", "summary": "This 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.", "questions": "The questions are listed in the previous section.", "limitations": "The limitation section is missing in this paper.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This 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.", "weaknesses": "1. 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback on our paper. We respond to the specific questions as follows.\n\n> …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 repeat 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\nWe did not mean to show that empirically learned transformers implement gradient descent (GD) in its intermediate layers, and we suspect that alternative mechanisms could exist. GD is rather **one possible mechanism** we use in theory for constructing the transformers.\n\n\n> 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\nFor the purpose of our theory, we believe there is no essential difference between ReLU and softmax. We believe our constructions can be generalized to standard softmax attention with some additional technical treatments. Our choice of ReLU was merely to simplify certain arguments.\n\nExperimentally, we tried out exactly the ReLU architecture we used in our theory (cf. Line 289-291 & Appendix M.1) and it performs well. Other recent studies (e.g. Shen et al. 2023) have also found transformers with ReLU attentions to perform well.\n\nK. 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.\n\nRe [1,2], we remark that [1] only used **linear attentions** instead of softmax, and the constructions in [2] used softmax in its “saturating regime” where it is approximately a **hard max** (cf. their arXiv version Appendix C.4.1, Page 19). Therefore, neither paper should count as \"predominantly using the softmax activation\" in a strict sense.\n\n> whether 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\nBy “MLP”, did the reviewer mean an MLP over the concatenated input (a vector in $\\mathbb{R}^{dN}$)? In that case, by the universal approximation property, MLPs can also approximate the ICL algorithms we considered. \n\nHowever, such an MLP would be **significantly more inefficient** (with much larger depth and width) compared with our constructions, as i) the input dimension (and thus the number of parameters) already scales linearly with the sequence length, whereas the size of our transformers do not; and ii) we used the attention layers in efficient ways to implement various operations such as in-context gradient descent (see e.g. Proposition E.1 for an example where the attention structure is crucial), which are unclear how to implement efficiently by MLPs. We will add a discussion about this point in our revision.\n\nFeel free to let us know if we understood the meaning of the \"MLP\" correctly.\n\n> 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 Bayesian analysis for Lasso.\n\nA Bayesian result for Lasso analogous to Corollary 6 holds directly for linear models with the Laplacian prior (by applying Theorem 7). We stated the result for ridge regression for convenience only (instead of for any fundamental reason), as our main focus was on the later results on adaptive algorithm selection building on these “base” algorithms.\n\n> 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\nThe sample complexity results in Section 5 (& Appendix M) can be viewed as corollaries of our efficient transformer constructions in Section 3 & 4, where we combine the constructions with standard generalization analysis to derive the sample complexities for pretraining. \n\nWe did not claim technical novelty in the generalization analysis part. Our techniques are indeed standard (by controlling Lipschitz constants + chaining arguments), and alternative techniques like [3,4] may also work here.\n\nThe novelty is rather in the efficient transformer constructions in Section 3 & 4, which ensures that the final sample complexity in Section 5 is mild. Results like [3,4] did not provide such concrete constructions for ICL algorithms, and thus they themselves cannot deduce our sample complexity results.\n\nWe will add a discussion of these points (space permitted) in our revision.\n\n> The limitation section is missing in this paper.\n\nWe appreciate the suggestion, and will properly discuss the limitations of our work in our revision.\n\n---\nWe thank the reviewer again for reading our response. We appreciate it if the reviewer could consider raising the score, if our rebuttal has addressed your concerns.\n"}, {"review_id": "PfBAlwSkvV", "reviewer": "Reviewer_RiCR", "summary": "They 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.", "questions": "(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?", "limitations": "Limitations 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?", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback on our paper. We respond to the specific questions as follows.\n\n> …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\nWe would appreciate it if the reviewer can clarify this question or provide some concrete examples, as we were unsure if the following interpretation was correct. \n\nBy “generalizing these hyperparameters”, did the reviewer mean e.g. generalizing from a finite set of $\\lambda$ (as in Theorem 11)  to a continuous range of $\\lambda$’s? In other words, can transformer pretrained on finitely many $\\lambda$’s perform near-optimal algorithm selection over continuous $\\lambda$.\n\nIn that case, we believe it may be possible to construct a transformer that achieves this. The technical difficulty would be in the statistical analysis of ridge regression about best $\\lambda$ selection (rather than the transformer construction). While we did not pursue this, we believe this could be an interesting direction for future work.\n\n> …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. …\n\nWe thank the reviewer for pointing out the missing references on Prior-data Fitted Networks (PFNs), and will properly cite them in our revision. \n\nOverall, these work on PFNs indeed demonstrates the in-context learning capability (Bayesian optimality) of transformers in various settings. We believe the experiments in these work and the theory in our paper complement each other: Our results essentially show that transformers can efficiently approximate a broad class of ICL algorithms, and we give theoretical constructions for concrete ICL algorithms, such as linear regression, Lasso, and gradient descent on neural networks. Such results for the expressive power and the Bayes-optimality of the resulting transformers were not established in the PFN literature.\n\nAlso, our results are not restricted to the Bayesian setting and can be broaderly applicable for providing frequentist in-context prediction guarantees for transformers.\n\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… whether 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\nWe believe there may indeed be some general conditions for tasks that are learnable in-context by transformers. For example, for any task that i) can be efficiently learnable by gradient descent, and ii) the gradient is approximable by attention layers, we should be able to construct transformers to do ICL and obtain similar polynomial sample complexity guarantees as in our paper. Concretely spelling out such general conditions, and identifying alternative conditions (different from GD) would be an important question for future work.\n\n> … 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\nMost of our approximation results (e.g. Theorem 4, Corollary 5, Theorem 7, Theorem 11) apply to any such OOD scenario, since the transformers we construct work on any in-context dataset with mild boundedness and non-degeneracy assumptions, with no distributional assumptions required. On the other hand, our pretraining results (Section 5) and experiments focus on the in-distribution setting. \n\nExperimentally, Garg et al. (2022) have conducted extensive empirical studies of ICL in OOD scenarios, with mixed results. Theoretically, we believe characterizing the OOD behavior of pretrained transformers is a major open problem and may require new insights and techniques.\n\n> Are there limits to the number of algorithmic tasks selected by transformers in the selection algorithm?\n\nOur transformer constructions do not have a hard limit on the number of tasks. A larger number of tasks $K$ would require a larger hidden token dimension $D$ to store all the intermediate results, as well as a larger $N$ in order for the validation losses to be faithful estimates of the true losses. See Theorem 11 (formal version in Theorem J.1) for an example.\n\n> Limitations are not clearly addressed, those could include: Theoretical analyses are hard to apply to more realtistic and complex real world tasks…\n\nWe agree that complex real-world tasks could be much more challenging than our setting, especially if the task distributions are more complicated or even consist of language data rather than real-valued data (as in our setting).\n\nWe will add a discussion on this and other limitations of our work in our revision."}, {"review_id": "ste4SLIPch", "reviewer": "Reviewer_Utbg", "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.\n2. 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.\n3. The authors construct a transformer model with near Bayes-optimal performance on linear models with mixed noise levels\n4. The authors prove polynomial sample complexity results for pretraining transformers to perform ICL.\n5. Experimentally, the authors show that transformer models can indeed perform algorithm selection in context. ", "questions": "--", "limitations": "--", "rating": 9, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The level of clarity and writing is very good.\n2. The theoretical contributions are insightful\n3. The experimental results support the theoretical arguments\n4. The scope of the work is broad and deep.", "weaknesses": "I do not recognize any major weakness", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the strong positive feedback on our paper!"}]

Author rebuttal:

Meta-review: In-context learning is an intriguing emerging ability of large language models. This paper makes important and practically relevant contributions towards understanding in-context learning abilities of transformers. There are more than one solid contributions. First, the authors provide efficient constructions for transformers that can learn regression tasks in-context. Their implementation is significantly more efficient than prior art. Notably, they can use a single attention layer with ReLU activation to approximate one GD step whereas prior works required multiple layers. This is insightful for practical performance of ICL. Secondly, the paper makes good contribution towards understanding model selection ability via in-context learning. For instance, prior work provided insights into Bayes optimal linear regression. This work manages to extend such findings to a broader class of algorithms and provide rigorous insights into the algorithm selection ability of TFs. I should note that there are some shortcomings of the paper that reviewer points out such as the use of ReLU attention rather than the canonical attention. But I find these fairly acceptable. I believe it will make a great addition to NeurIPS 2023 program.
Decision: Accept (oral)