Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context
Authors: Spencer Frei
Gal Vardi
Venue: ICLR 2025 Poster (2025)
Keywords: in-context learning
generalization
benign overfitting
implicit regularization
Primary area: learning theory

Abstract: Transformers have the capacity to act as supervised learning algorithms: by properly encoding a set of labeled training (''in-context'') examples and an unlabeled test example into an input sequence of vectors of the same dimension, the forward pass of the transformer can produce predictions for that unlabeled test example.  A line of recent work has shown that when linear transformers are pre-trained on random instances for linear regression tasks, these trained transformers make predictions using an algorithm similar to that of ordinary least squares.  In this work, we investigate the behavior of linear transformers trained on random linear classification tasks.  Via an analysis of the implicit regularization of gradient descent, we characterize how many pre-training tasks and in-context examples are needed for the trained transformer to generalize well at test-time.  We further show that in some settings, these trained transformers can exhibit ''benign overfitting in-context'': when in-context examples are corrupted by label flipping noise, the transformer memorizes all of its in-context examples (including those with noisy labels) yet still generalizes near-optimally for clean test examples.

Published as a conference paper at ICLR 2025
TRAINED TRANSFORMER CLASSIFIERS GENERALIZE AND EXHIBIT BENIGN OVERFITTING IN-CONTEXT
Spencer Frei
Gal Vardi
UC Davis" I Fre @gcogge.ccom
Weizmann Institute of Science gal. va rardi @we! zmann ac
ABSTRACT
Transformers have the capacity to act as supervised learning algorithms: by prop- erly encoding a set of labeled training ("in-context") examples and an unlabeled test example into an input sequence of vectors of the same dimension, the forward pass of the transformer can produce predictions for that unlabeled test example. A line of recent work has shown that when linear transformers on random instances for linear regression tasks, these trained transformers make predictions using an algorithm similar that of ordinary least squares. In this work, we investigate the behavior of linear transformers trained on random lin- ear classification tasks. Via an analysis of the implicit regularization of gradient descent, we characterize how many pre-training tasks and in-context examples are needed for the trained transformer to generalize well at test-time. We further show that in some settings, these trained transformers can exhibit "benign overfit- ting in-context": when in-context examples are corrupted by label flipping noise, the transformer memorizes all of its in-context examples (including those with noisy labels) yet still generalizes near-optimally for clean test examples.
pre-trained
INTRODUCTION
A key feature of transformer-based large language models (LLMs) is their ability to perform in- context learning: by providing few labeled examples to trained transformer with an unlabeled example for which the user wants prediction, LLMs can formulate accurate predictions without any updates to their parameters (Brown et al., 2020). This ability to perform in-context learning is in- fluenced by the interplay between tokenization of inputs (i.c. how to feed data into the transformer), pretraining datasets, optimization algorithms used for pre-training, the particular architecture of the transformer, well as the properties of the in-context examples that are provided test time.
A number of recent works have sought to develop a deeper theoretical understanding of in-context learning by investigating transformers trained on classical supervised learning tasks. This line of work was initiated by the experiments of Garg et al.) (2022), who showed that when transformers are trained on random instances of supervised learning problems, they learn implement supervised learning algorithms at test-time: e.g., when training GPT2 architectures (Radford et al.] 2019) on random linear regression tasks with weights sampled from a Gaussian prior, the transformer learns to implement an algorithm similar to ordinary least squares. This led to a series of follow-up works which analyzed the dynamies of gradient descent/flow over simplified linear transformer architec- tures when trained on linear regression tasks (von Oswald et al. 2022: Akkurek et al.) 2022), some providing guarantees for how many tasks or samples were needed to generalize well (Zhang et al.) 2024 Ahn et al. 2023 Wu et al. 2024).
In this work, we analyze the behavior of linear transformer architectures which are trained by gradient descent on the logistic or exponential loss over random linear classification tasks, We consider a restricted linear attention model, a setting considered in prior works (Wu et al. 2024] Kim et al. 2024). We assume that pre-training tasks are sampled from random instances of class- conditional Gaussian mixture model data, i.e. for some R > 0, covariance matrix A and for each
NNow at Google DeepMind

---

Published as a conference paper at ICLR 2025
Test Examples
Training Data (In-Context Examples)
France: Peris
Transformer
UK:
Glasgow
UK: Giosgew
USA: ?
Canada: Ottawa
Figure 1: Benign overfitting in-context: after pre-training (not shown), when given a sequence of in-context examples, the transformer memorizes noisy labels yet still generalizes well to unseen test examples.
task I,..., we have
Lid. Unif(R (nif(R-s"-1), sd-1), M7
Vo. 1id- Unif({=1)),
Lid N(O,
27,5
2 = Vräfft + 27,5
We assume that test-time in-context examples also come from class-conditional Gaussian mixture models as above but with two important differences. First, the signal-to-noie (determined by R) test time can differ from those seen during pre-training, Second, we allow for label-flipping noise to be present test time, i.e. we flip labels yi + yy; with probability p.
Our first contribution is an analysis of how many pre-training tasks are needed in order to generalize well at test-time. We find that after pre-training the transformer tasks which it was trained on. Moreover, we find that at test time the transformer can tolerate label- flipping noise, even though the pre-training data does not have label-flipping noise. Thus, even when pre-training on simple and easy-to-learn datasets, the transformer can generalize on more complex tasks. Our proof follows by an explicit analysis of the implicit regularization due gradient descent during pre-training. To our knowledge, this is the first theoretical result in-context learning for linear classification.
tolerate smaller SNR than those
Our second contribution is the finding that the trained transformer can exhibit "benign overfitting in- context": namely, if the test-time sequence is subject to label-flippng noise, then in some settings the transformer memorizes the in-context training data yet still generalizes vear-optimally. The precise setting where this occurs requires the features to lie a high-dimensional space and for the SNR to be relatively small, set of conditions observed in prior works benign overfitting in classification tasks (Chatterji & Long 2021] Frei et al.) 2022 2023a). Figure iillustrates what this phenomenon would look like in the language setting. To the best of our knowledge, no prior work demonstrated that transformers could exhibit benign overfitting.
RELATED WORK
In-context learning for supervised learning tasks. Following the initial experiments of Garg et al. 2022), a number of works sought to understand what types of algorithms are implemented by transformers when trained supervised learning tasks. [Akyurek al. (2022) and von Oswald et al. (2022) used approximation- of transformer architectures could implement linear regression algorithms. Bai al. (2024) showed that transformers could implement a variety of classical algorithms used for different supervised learning tasks, e.g. ridge regression, Lasso, and gradient descent on two-layer networks. /aavorrs.//hnngggg al. 2024), Ahn et al. (2023) and Mahankali et al. (2024) examined the landscape of single-layer linear transformers trained on linear regression tasks, with Zhang et al. (2024) additionally developing guarantees for convergence of (non-conved) gradient flow dynamics. There are other works which examine training of transformers on data coming from hidden Markov chains (Edelman et al. 2024) or nonparametric function classes (Kim et al.) 2024). The most closely related work this one is Wu et al. (2024), who focused on a convex linear transformer architecture, as we do this work, trained on linear regression tasks with SGD. They provided task and (in-context) sample complexity guarantees when training by SGD. contrast, we develop guarantees for the classification setting, and we develop guarantees via analysis of the implicit regularization of GD.
and mechanistic approaches to understand how a variety
2

---

Published as a conference paper at ICLR 2025
Implicit regularization in transformers. The starting point for our analysis of the task complex- ity of pre-training and the sample complexity of in-context learning is an analysis of the implicit regularization of the optimization algorithm. We outline the most related works here and refer the reader to the survey by Vardi (2022). The convex linear transformer architecture we study is linear in (the vectorization of) its parameters, which by |Soudry et al. (2018) implies that gradient descent converges in direction to the max-margin solution. The more general class of linear transformer architectures are non-convex in general, but certain subclasses of them are homogeneous in their parameters and hence converge (in direction) to KKT points of max- max-margin solutions (Lyu & Li, ?). Another line of work seeks to understand the implicit bias of GD for softmax- based trans- formers (Ataee Tarzanagh et al. 20233 Tarzanagh et al. 2023; Thrampoulidis [2024] Vasudeva et al.)
2020
2024
typically with more stringent assumptions on structure the training data.
Benign overfitting. The ability for neural networks to overfit to noise yet still generalize well (Zhag et al.) 2017) led significant line of work on reconciling this phenomenon with classical learning theory. We focus on the most directly relevant works here and point the reader to the surveys Bartlett et al. (202T): Belkin (20T)) for a more extensive review. Prior works have shown that kernel-based based methods (Belkin et al. 2019) and the ordinary least squares solution (Bartlett et al., 2020) can exhibit benign overfitting in regression tasks, and that the max max--magiin solution can in clas- sification tasks (Chatterji & Long 2021 Frei et al. 2023a). Our proof technique comes from the works Frei et al. (2023c a). where it was shown that the max-margin solution over 2-layer neural nets behaves similarly to nearly uniform average of the training data and that this average can exhibit benign overfitting in class-conditional Gaussian mixture models. We similarly show that the forward pass of pre-trained transformer behaves similarly an average over the in-context training examples, and use this to show benign overfitting.
3 PRELIMINARIES
In this section we introduce the pre-training distribution, the particular transformer architecture we consider, and the assumptions on the pre-training data.
3.1 NOTATION
We denote matrices with capital letters W and vectors and scalars in lowercase. The Frobenius norm of matrix is denoted [[W]|, and its spectral norm as || We use A b == min(a, and a V b-= max(a, b). We use the standard ((-).((-). and θ notations, where Ō,Ω, ignore logarithmic factors. We use (æ) as the indicator function, which is x is true and otherwise. We denote [M] (1, ..., M) for a positive integer M, and we denote Unif sd-1) as the uniform distribution on the sphere of radius R d dimensions. In Table [ in the appendix, we collect all notation used throughout the paper.
3.2
SETTING
In the linear regression setting, there is natural distribution of pre-training tasks that serves as basis for number of theoretical works on in-context learning in that setting (von Oswald et al., 2022 Akyurek et al. 2022: Zhang et al. 2024 Ahn et al. 2023), namely, a fixed distribution over the features x and Gaussian prior over the weights w, for each task and labels generated as w. x per task. We are unaware of any prior works in the linear classification setting, so we propose the following. We assume each task is a random instance of a class-conditional Gaussian mixture model with cluster means μ4, Unif(R (d-1). In particular, we assume the following.
Assumption 3.1. Let A > be a symmetric positive-definite d d matrix and let B, N and R > O. Let 1.... B and = I....NN
N N +11.
iid. Unif(R 1) be uniform on the sphere of radius in d dimensions.
1. Let py
2. Let 27, ltd. N(O, A), and yo. Unif((+1)), where Pr. 27,5 and Fr are mutually independent.
Lid.
3. Set Nr,i == Vralll + Zr,i-
3

---

Published as a conference paper at ICLR 2025
We tokenize samples to be fed into the transformer in the following way: for sequence of N labeled examples (x,, [=1 which we use to formulate a prediction for an unlabeled test example 2N+1, concatenate the (Zt, Va) pairs into d + dimensional vector and tokenize the test example as (I,+110),
21
22
EN YN
EN+1 0
E Rld+1)x(N+1)
y/1
3/2
In the standard softmax- attention with single head (Vaswani et al. 2017), the transformer is defined in terms parameters W E IRd,xd, W W gdaxd. IfO (W W9, WV) then softmax attention computes
(WWKE)TOR)
(E;0)=E< BEEEWWER E softmax
PN.d.
where PN.d. is normalization factor which is not optimized (but which may depend on N and de). In linear transformers, the softmax is replaced with the identity function. Following prior work (von merge the K and Q matrices into and we consider an objective function where the aim to use the first N columns of ① to formulate predictions for IN+1, whereby the bottom-right corner of the output
Oswald al. 2022 Zhang et al., 2024 Ahn et al. 2023).
W (W
W
w
matrix of f(E; f) serves as this prediction. Writing W
for A E (V,KQ),
(w
W22-'
for the linear transformer architecture, this results the prediction
( W.k ((wt))) IN+1-
O E: 0) ((wX)
W22
.ET.
N
Due to the product matrices appearing above, the resulting objective function is non-convex, which makes the analysis of its training dynamics complex. In this work, follow Wu et al. (2024); [iim et al. (2024) and instead consider convex parameterization of the linear transformer
KQ which results from taking W21 = w 0 and setting w = 1. This results in the following
prediction for the label of 2N-1+
N
T
TE W) =
yi
WøN+1-
(2)
N
That is, [E; W) corresponds to the predictions coming from linear predictor trained by one step of gradient descent on the logistic or squared loss (initialized at 0) but where the parameter W is a (learned) matrix preconditioner. Prior work by Zhang al. (2024) showed that when single-layer linear transformer architecture is trained by gradient flow on linear regression tasks, it learns a function of the same type, where the preconditioner was found to correspond the inverse covariance matrix of the pretraining data.
For each task т, denote the embedding matrix E, as the one formed using labeled examples (Ir.is Mr, )a from Assumption 3.1 per ( We consider linear transformers which are trained on the prediction the last token using the logistic or exponential loss, namely
B
Ż(W)
1
Σ
EVV.N+1 ()):
i {g + log(1 + exp(-q)), a ++ exp(-q)}-
We are interested in the behavior of gradient descent on this objective,
W Wr V...W..
Note that (FF; know that gradient descent has an implicit bias towards maximum-margin solutions, as we recall in the following theorem:
F:; W) homogeneous in W (linear in vec(W)) and so by (ouudyy et al. (2018) we
Theorem 3.2 (Soudry et al. (018)). Define the tz-max margin solution
WMM argmin ||0||2 (I
Uyr 2 11 BB.
(3)
If there exists which satisfies the constraints in B3, then provided a sufficiently small, then (t)) converges in direction to WMM. iee. for some constant > O we have W./lw.ll - cWMMs.
4

---

Published as a conference paper at ICLR 2025
This theorem shows that the max-margin solution characterizes the limiting behavior of gra- dient descent. In the remainder, we will focus on the max- max-margin solution and derive all of the generalization guarantees consequence of the properties of the max-margin solution.
We next introduce the assumptions we need to ensure that the max max-margie margin solution is well-behaved enough that (0, 1), and we shall show that all of our results hold with probability at least 1 - 6 provided the absolute constant C > 1 appearing below a large enough universal constant (independent of d, B, A, and 8).
later derive in-context learning guarantees. We fix a probability threshold ő
(AI) R is large enough so that
tr(A) d
tr() V IlAll2/ log(2B/6)
R 2 C (dtt VC
N
(A2) For some cB > O, we have B 2 cgd, and
(A3) d 2 (2B2/8)
The first assumption (AT) guarantees that the signal-to-noise ratio is sufficiently large (recall that A is the covariance of Z in the identity Zr,i VVr,ill + zr,); provided / is polynomial in d, this assumptions is satisfied for A = 7 when R >>> vd. The second assumption (A2) specifies how many pre-training sequences we (pre-)train on; we shall see later on that the generalization error achieved for in-context examples is determined in part by how small cp is (cg larger ensures better generalization). Note that we allow to be non-constant (so that oz(d) tasks as permitted), The assumption (A3) needed for technical reason, namely to ensure certain concentration bounds regarding sub-exponential random variables. Finally, note the near-independence of these assump- tions on the number of examples per task during training: in many situations (e.g., A = I), we only require single demonstration (i.e., N 1) of the form (Zr,1,
Finally, we introduce the distribution on the test-time in-context examples. We allow for a more general distribution at test-time than the one used during pre-training, where the in-context examples can have subelffippiig noise and a potentially smaller cluster mean, corresponding to a potentially smaller signal- sollltt--oooo to-n ratio. We assume, for I,,
1. Y/ Unif ({=1)).
2. y/ = 5: w.p. - p and = - D1 w.p. p (label flipping noise).
3. μ R. Unif((d-1)
4. 2 ~ A() and x; Dil + 24.
We assume the random variables Ü: ,1, 2 are all mutually independent. We allow for the size of the cluster mean R at test-time to potentially be different than the size of the cluster mean R from the pretraining data, and indeed we shall see in Theorem 4.1 below that we can permit a smaller it than what require for R. Also note that we allow for potentially different sequence length test- time (M) than was observed during training (N); prior work in the linear regression setting showed that the transformer's behavior can depend quite differently on M vs. N (Zhang al. 2024).
4
MAIN RESULTS
4.1 GENERALIZATION
Recall that the transformer with parameters W makes predictions by embedding the set ((x,, M. U [(TM+11:00 0)} into a matrix E and then using bign(7(E))) (TE: as our prediction for 3M++1- Our goal is to understand the expected risk the max-margin solution, i.e. the probability of misclassification of the test example (MM+1,94+11 using (
RMWMM) = IP (a..))) 11, (sigu(M E; W MMM) # 3M+1)
1 M
PP
WMMUM++1MMH <0)
5

---

Published as a conference paper at ICLR 2025
Our main result regarding generalization is the following theorem.
Theorem 4.1. Let d E (0, 1) be arbitrary. Suppose that 1/2 Cp where Cp (0. 1/2] is
an absolute constant. There are absolute constants > 1,c > 0 depending only on Cp such that
assumptions (AT)
(A2)
and [A3] then with probability least 108 over the draws of M++1
Apr, (Ir Dr k= Jr=1 when sampling new task {u, (2i, V6) 2 }, the max-margin solution satisfies
B
P (M (iE) WMM)) # yM+
cpR
cpM ||Al/2vd
S (p+ (x++cxx((cM))) (pp 0) 2exp (-opva) 4exp
(T
2exp
where p := Tog" (2FF(8)
CaAl
Let us make few observations on the above theorem. For simplicity let us assume that the covari- ance matrix A = Id. so that ||A ||A|| 1, and let us assume we want the results to hold with
fixed probability threshold of 6 0.001. Then assumptions [A1]] through (A3) are satisfied for
any N 2 1 examples per task so long as R || > CVd.. For the test error, there are a number of different regimes which require different considerations.
Due to the dependence on p. the test eπror degrades when eø gets smaller, and we require cg = 1/22 to achieve non-vacuous generalization bound, i.c. B = (2d pre- training tasks suffices if R and M are large enough
While we require pre-training signal = [2( (), to achieve test error near the noise rate at test-time the signal it can be as small as π = Ộ1(1) provided M is large enough and P (1)). This means the hignal-tonoois ratio at test-time can be significantly smaller than what was observed during pre-training
When = 0, there is no label-flipping noise. Here it is possible to have as few as one example per task (M 1) provided ft = (1dd/4), and p -= ((1).
When > O, there is label-flipping noise. The same analysis holds as in the p = 0 setting except that we always require M = 22(1), which makes intuitive sense (with label noise, one must see more than one example guarantee generalization w.h.p.). Notably, there was no label noise during pre-training, yet the transformer can generalize under label noise at test time.
4.2 BENIGN OVERFITTING
In this subsection we shall assume that the covariance matrix A = Ia- This is needed for technical reasons of the proof. We could more accommodate a covariance matrix of the form A where < < 02 but we just present the case A = I for simplicity.
Our goal will be show that the max-margin transformer exhibit benign overfitting in-context. By this we mean that in its forward pass, the transformer memorizes all of the in-context examples, yet still generalizes well for new test examples. To make this more precise, for sequence of train- ing examples {(,,, J.M, denote E(z) as the embedding matrix corresponding to the sequence (I;, yı M (z, 0); recall that the examples (x,, yı W are used to formulate predictions for x. We say that the transformer exhibits benign overfitting in-context if,
All training examples memorized: with high probability (over l and (xt, 9: M for each example (Ik, 3k), k S M, we have sign( E(zx); W)) yk.
Test examples are classified near-optimally:
TP P (xaa)) (sign(((((M()))) =
yM+1) S + for some vanishing
We also require B = (dpolydd)) this comes from using union bound over the pre-training tasks to ensure the max-margin solution is well-hehaved.

---

Published as a conference paper at ICLR 2025
Performance Dimension gi-1) [5eq. Length 20, abel Flip 0.2)
Nerformance (8 d233) (Seq. Length 20, Label Flip - 0.2)
Performance Dimansion d2.6)
(Seg. Length 20, Label Flip 0.2)
1.0
10
1.0
0.9
0.9
0.9
0.88
Optimal Best 50) context
0.3
0.5
0.7
0.7
0.7
0.6
0.6
Optimal Test 80)
0.5
Optimal In-context
In-context
0.5
0.5
0.5
15a
10
Input Dimansion
Inpur: Dimansion
Figure 2: Cluster separation R affects both (in-context) train and test accuracy, with small R leading to overfitting, and large R. leading to better test accuracy; in between, benign overfitting occurs.
In Figure [ we provide a schematic of what this phenomenon would look like in language tasks, where user provides a sequence of country-capin pairs, some of which have errors, and the trans- former repeats these errors at test time for those countries which it has seen with noisy labels, yet still generalizes well for never-before-seen countries.
To prove this phenomenon holds, we need to show both memorization and generalization. Theo- rem emm4.11in in the previous subsection proved the generalization part, thus all that remains is to show that memorization can co-occur with generalization. The next theorem demonstrates that this is possible if the dimension d is sufficiently large with respect to the number of in-context examples M.
Theorem 4.2. Assume A = I and let d E (0, 1) be arbitrary. Suppose Cp (0, 1/2) is an absolute constant. There are absolute constants C > 1, 0 depending only on Cp such that if assumptions [AT] [A3] hold, then the following holds. With probability at least 1 - 106 over the draws of {ur, (Ir, Dr i=1 1 for all T E [B]. when sampling new task
1/2 cp where
{M, (1,, (1) M))) }. the max-margin solution satisfies,
In-context training examples are memorized: for p log 12
P F(ELN1)
(3k E [M] s.I. dign(α(E((za)) WMM)) + yk)
cpvd
epd
S 4M exp
+ SM
exp
(-M(Pe v Ri)
VM
In-context test example achieves test error close to the noise rate:
P (1. M+1 (S(M)()(((( WmM)) + yss+2)
/2
S p 2exp((xx-MM + 2xxp (-epVd) + 4exp (-cpft) + 2exp
Vd
The claim regarding generalization in the above theorem is the same in Theorem 4.1 (we fo- cus here on the case p > 0 since label noise is essential for overfittin), and the same comments following that theorem apply here as well. But a natural question whether one can satisfy both memorization and near-noise-rate test eиror, i.c. benign overfitting. This is indeed possible, and can be seen most easily in the following setting: let d = 0.001, M be large constant satis- so p = 1/ log" (2d2/8), If R d
fying p- + 2 exx(-cMM S p + 0.0001, and assume B for 33 (1/4, 1/2), then w.p. at least 99.9% over the pre-training data, since / < 1/2 memoriza- tion occurs with probability - o1))), and since 33 > 1/4 the test error for the fresh test example JM++ is at most + 0.0001 + (1). More generally, one can see that memorization occurs when d > R v R) and d > M, i.e. when the dimension is large relative to the number of samples and when the signal of the test-time sequence is not too large relative to the dimension divided by the number of in-context samples.
While Theorem 4.I and Theorem 4.2 apply for the infinite-time limit of gradient descent, a natural question is whether gradient descent trained for finite steps has similar behavior, and indeed find
7

---

Published as a conference paper at ICLR 2025
Performance Number Tasks (& d0I)
Performance Number Tacks (f- oody
Performance Number Tacks (6 go0)
10-1000, Seo. Length 20, Label Fiip 0.2)
(d- 1--100 1000, Seq. Length 20, Label 0.21
1--000, Seq. Length 20, Label Flip 0.2)
1.0
10
Opcimal Tes: (0.80)
in-contect train
0.9
0.9
0.9
Test
0.88
Optimal Tes:: (0.80)
0.3
0.8
incoontect trais
0.7
Test
2.0
0.7
0.6
0.6
Optimal Tec: (2.80)
0.6
In-contect
0.5
0.3
0.1
Tasks 10
10*
Tasks
Tasks
Figure 3: In-context test accuracy improves as the number of pre-training tasks (batch size) B approaches the ambient dimension d, provided is large enough, with fewer needed for large R.
that it does. Figure 2 examines the role of high dimensionality and R. illustrating three qualitatively different phenomena: overfitting without generalization, when R is small and d/M large; benign overfitting, when At is in the sweet spot for generalization: and optimal test performance without overfitting, which occurs when R is very large. In Figure 3 we show the effect of the number of pre-training tasks B on the in-context test accuracy, again showing improved performance for more pre-training tasks with fewer pre- training tasks needed for optimal performance when R is large, as suggested by Theorem 4.2 In both figures, we have GD-trained networks with step size v) 0.01 for 300 steps. Details experiments and link to our codebase are provided in Appendix B
PROOF SKETCH
We will first discuss a sketch of the proof of generalization given in Theorem 4.1 when A = J. For M notational simplicity let denote W WmM: ÿ == 2 Σ 1=1 y:24, and let drop the M
subscript so that we denote (1M+1, 3M = (æz, D) Then the test error is given by the
probability of the event,
(sign(((E; + yM+1) (ī TayWyy yor S 0, y = ú) ( T yo SO, y --
The test error can thus be bounded as
((O(F;; F;; W) + ym+1) s War S 50) + Ply = - P( Whx S 50) +
Let us denote the examples (I,, 34) in the test-time sequence which have clean labels y = y: as i E C, while those with noisy labels ys Di as i E N. Then ΣM V: = (IC) - [W1) + AM WF; = (M- - 2(A1)) + Em y:zz, If we denote P = WI/M then by standard properties of the Gaussian, we have for 2, N(O, I),
lid.
n _ (1 - 2p) μ + M M- 1/2
j0 d
μ t'z z'.
We thus have
(O((;W) FVM+1)
SptP (1122)) + M-1/2 T "(" 2²) <0)
= =p+P((1 =p+P((1-5D) ((112ppji
TW -(1 - 2 Wz 1-122 Wp 1-1 Wz/ TT2
This decomposition allows for us to show the test error is near the noise rate if we can show that (1 2 W TWp) is large and positive while all of the other terms involving μ W'z, , 2 Wpp and 2 are small in comparison. Now if д were standard Gaussian,
if it were sub- Gaussinn with independent components, then j Wμ would be close to its mean tr(W). with fluctuations determined from its mean determined by [[]][,, via standard Maanson-Wright inequality. However,
8

---

Published as a conference paper at ICLR 2025
since j Unif(R Unif(R-S(-1) Sd-1), the components of ja not independent, and derive modified version of Hanson- Wright (Lemma C.2 which applies in this setting where we show μ W p is close to its mean 2 tr(W) with fluctuations from its mean determined by |[W]| Thus, in order to show that the max-margin solution achieves (in-context) test error close to the noise rate p, we must show the following:
&
A lower bound for tr tr(W), so that the mean of μ W is large and positive.
An upper bound on [W][,, so that we can control the fluctuations of μ T W from its mean.
An upper bound on ý so that (1 - 2p) positive.
Upper bounds on the quantities [[1 - 2p) W |M-1/22 W and |M-1/2; W='| so that these quantities are smaller than the positive term coming from the mean of μ W μ. These, in turn, will require upper bounds on ||
Thus provided we have control over tr(W) and | l.. we can apply straightforward concentration inequalities to show that the (in-context) test error is near the noise rate. So let us now describe how to derive such bounds.
The starting place for analysis is the definition of the max-margin solution For notational simplicity let us define ū. == N Σi=1 Vr,1"r,i and let us denote 2 Vr, as 27, Ur so that the max-margin solution can be written
W argmin([]U]]] : UT 2 1T l....,
(5)
Since Vi W) n. the KKT conditions imply that there exist A....., , An such that
W Σ A-
(6)
and moreover we have A, = 0 whenever F. Wx, * 1. By a careful analysis of the KKT upper and lower bounds for Σ, A
conditions (see Lemma B.Ip. we can derive nearly-matchinn tr(W), and [[ if we assume that cB > O is an absolute constant (i.e. the number of tasks B is greater than a constant multiple of the input dimension d) then these bounds take the form
~(==((
tr(W) = ((d/)))(RRR)) ||] - -
These bounds, together with the concentration inequalities outlined above, suffice for proving The- orem 4.2
As for the possibility of benign overfitting in-context shown Theorem 4.2, since Theorem 4.1 guarantees generalization, the only task that remains to be shown is that overfitting occurs: namely, we want to ensure that for every example (æx, 8k), if we denote E(zx) as the embedding matrix with input sequence (I,, y4) M but with test example Tg, then want to show that
sign (#(E(x)); W)) yk = WyRTK > 0.
Now, since tr(cf) = ed and |c/||F evd. the identities in 7 suggest that W has properties similar that of scaled identity matrix. If W were indeed scaled identity matrix, then the goal would be to show that
Mī. yk
Σ 3eo0s
yk || 2
Σ
(War 2 || Eg
Σ Kar xx)) >
That is, would need to ensure that the examples {x; )MM are nearly-orthogonal in particular sense (Frei et al.) 2023c). This occurs provided the ambient dimension d is much larger than the number of examples M and the signal strength 8² |||||2, and has been shown in prior work on benign overfitting in neural networks (Frei et al. 2022 2023a). In our setting, W is not multiple of the identity matrix so this proof technique does not directly apply, but at a high level the proof
ideas are similar, and appear in Section

---

Published as a conference paper at ICLR 2025
6 CONCLUSION
In this work we developed task complexity and sample complexity guarantees for in-context learning class-conditional Gaussian mixture models with a single-layer linear transformer architecture. We analyzed the implicit regularization of gradient descent characterize the algorithm implemented by the transformer after pre-training. This allowed for us to quantify how many in-context samples needed in order achieve small test error, which the best of our knowledge has not been explored in the classification setting prior to this work. We also showed how the trained transformer can exhibit benign overfitting in-context, i.e. in its forward pass the transformer can memorize noisy examples yet still achieve near-optimal test error.
There are number of natural directions for future research. We relied upon convex linear trans- former architecture which allows for us to identify the pre-trained transformer as the global max- margin solution in parameter space. More general linear transformer architectures are not convex but are often homogeneous in their parameters. In this setting we thus know (Lyu & Li, 2020; Ji Telgarsky 2020) that GD has an implicit bias towards first-order stationary (KKT) points of the max-margin problem in parameter space. It may be possible to analzye the consequences of the KKT conditions to develop generalization guarantees (Safran et al. 2022 Frei et al. 2023b a). For softmax-based transformer architectures, it would be interesting to
prior works implicit
regularization GD over these architectures (Ataee Tarzanagh et al. 2023 Tarzanagh et al. 2023
Thrampoulidis 2024 Vasudeva al., 2024) can be used understand in-context learning.
Finally, we assumed that the signal-to to-noise ratio the pre-training data was quite large (see (A1) and that the pre-training data did not have noisy labels. It would be interesting understand if pre-training on more difficult or noisy data would result plemented by the pre-trained transformer. We believe the theoretical analysis of pre-training with noisy labels would be significantly different: if we aimed to use the implicit regularization approach based on analyzing solutions our equation (3), in this setting there may be which satisfies the constraints of (3), and even if such a U does exist then we would be investigating the behavior of a pre-trained transformer which has memorized noisy labels. In this setting it is not clear that generalization is possible.
qualitatively different algorithm im-
ACKNOWLEDGEMENTS
Part of this work was completed while SF was a visitor at the Modern Paradigms in Generalization Program at the Simons Institute for the Theory of Computing. GV is supported by research grant from Mortimer Zuckerman, the Zuckerman STEM Leadership Program, and by research grants from the Center for New Scientists at the Weizmann Institute of Science, and the Shimon and Golde Picker Weizmann Annual Grant.
REFERENCES
Radoslaw Adamczak. A note on the Hanson- Wright incquality for random vectors with dependen- cies. Electronic Communicationn in Probability, 20:1 13, 2015.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. Preprint, arXiv:2306.00297, 2023.
Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learning? investigations with linear models. arXiv preprint arX6v.2211.15661, 2022.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin to- ken selection in attention mechanism. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
10

---

Published as a conference paper at ICLR 2025
Peter L. Bartlett, Philip M. Long, Gabor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117748):30033330077 2020.
Peter L. Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: a statistical viewpoint.
Acta Numerica, 30:87-201, 2021.
Mikhail Belkin. Fit without fear: remarkable mathematical phenomena of deep learning through the prism of interpolation. Acta Numerica, 30:203-244 248, 2021.
Mikhail Belkin, Alexander Rakhlin, and Alexandre B. Tsybakov. Does data interpolation contradict statistical optimality? In AISTATS, 2019.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamers. Advances in neural information processing systems, 33:1877- 2020 2020.
Niladri S. Chatterji and Philip M. Long. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22(129)) 2021.
Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis.
The evolution of statistical induction heads: In-context learning markov chains. Preprint, 24X1v:2402.11004 2024.
Spencer Frei, Niladri Chatterji, and Peter L. Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory (COLT), 2022.
Spencer Frei, Gal Vardi, Peter L. Bartlett, and Nathan Srebro. Benign overfitting linear classifiers and leaky relu networks from kkt conditions for margin maximization. In Conference Learning Theory (COLT), 2023a.
Spencer Frei, Gal Vardi, Peter L. Bartlett, and Nathan Srebro. The double-edged sword of im- plicit bias: Generalization robustness in relu networks. In Advances in Neural Information Processing Systems (NeurIPS), 2023b.
Spencer Frei, Gal Vardi, Peter L. Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu networks trained on high-dimensional data. In International Conference Learning Represen- tations (ICLR), 2023c.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. arXiv preprint arXiv:2208.01065 2022.
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Juno Kim, Tai Nakamaki, and Taiji Suzuki. Transformers are minimax optimal nonparametric in- context learners. Preprint, arXiv:2408. 12186, 2024.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations (ICLR), 2020.
Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One step gradient descent prov- ably the optimal in-context learner with layer of linear self-attention. In International Con- ference on Learning Theory (ICLR), 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Itay Safran, Gal Vardi, and Jason D Lee. On the effective number of linear regions in shallow univariate relu networks: Convergence guarantees and implicit bias. Preprint, arXiv:2205.0900 09072, 2022.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(1), 2018.
11

---

Published as a conference paper at ICLR 2025
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308. 168988 2023.
Christos Thrampoulidis. Implicit bias of next-token prediction. Preprint, arXiv:2402.188535 2024.
Gal Vardi. On the implicit bias in deep-learning algorithms. Preprint, arXiv:2208. 12090,
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.057..5 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- sion processing systems, 30, 2017.
Roman Vershynin. /ligh-dimensional probability: An introduction with applications in data science, volume 47. Cambridge University Press, 2018.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.0r1 02.77777777 2022.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In Interna- tional Conference on Learning Representations (ICLR), 2024.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. International Conference on Learning Rep- resentations (ICLR), 2017.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in- context. Journal of Machine Learning Research, 25(49):1-55, 2024.
12

---

Published as a conference paper at ICLR 2025
Table 1: Notation used throughout the paper.
Symbol
Description
Cluster mean, Rd
Features, r E IRd Tri Vr,ill + Zr,
Labels, y {+1}
y
Noise variables, 2 Rd
d
Feature dimension
A
Cluster noise covariance matrix
8
Probability of failure
Norm of cluster means during pre-training
R
Norm of cluster means at test time
B
Number of pre-training tasks
CB
Quantity such that 2 cgd
(eB A 1)/ 1//loog"22 (2 /5)). appears generalization bounds Number of samples per pre-training task Number samples per test-time task Label noise flipping rate
p
N
E
Data matrix E
T2
IN
IN+1
I((++))(N+1)
91
y/2
UN
Mean predictor:
37
i=
M
DAJz
((E()):
Neural net output:
Li=1 y:
M
Wx
W
L(W)
Logistic loss
A PROPERTIES OF THE PRE-TRAINING DATASETS
We begin by developing guarantees for various properties of the pre-training datasets, which will form the basis for understanding properties of the max-margin solution Q
Lemma A.1. There is an absolute constant co such that with probability least 108 over the draws of (Hr, (Ir: Vz), (Ir,s, Vr,:), N
for all E [B] and q + T,
RoR / tr(A) (2B(2
tr(A) coll A (2B(2)
|| R21
VNd
R tr(A) log(2B/0)
|| - 1291
4(t((A) V C log(///)))(22
Vd
R² RV/U(A)
(r(PP)
(()s 5 co
+ log(2B²/0)
Vd
VNd
N
R2
R//t((A
(1)r sco
+ v
log(2B2 /8),
Vd
Vd
[+
1
R RVTr(A)
V tt((2)) VN
(PT. R²
v Co
+
log(2B/8)
VN
Vd
Proof.
By definition of e. and properties of the Gaussian distribution, there is o4 ~ N(O, A) such
that
1 N N
N
1
P.
Σ Vr.i + 2T,s) = pe, +
Σ
Br.aZo,t = fr +
N 2
N
Thus for 7 + g there are z 2/ iid. N(O, A) such that
(T _ (pt + N-1/22/ Pg + N-1/22/)
(11::1g) + N N-1 (2/ Mgg + N (2, (+)) + -N-1 =))
(8)
13

---

Published as a conference paper at ICLR 2025
We first derive upper bound for this quantity when q * Now, since Ma: Mr independent and sub- Gaussian random vectors with Lub-Gaussian norm at most cRt V (Vershynin 2018 Theorem 3.4.6) for some absolute constant > 0, by Vershynin 2018 Lemma 6.2.3), we have for some
> 0 holds that for any 3 IR, if g.g i..d
Na),
Blexp(311 MT)] & Elexple R2d Bg "g)g')}
By Vershynin 2018 Lemma 6.2.2), for some e1 > 0, provided e/a) S 01, holds that
E (exp(cR2dd113 g')l I exp(e1 R ( = exp(e1 (
Since Ma Mr are mean-zero, by Vershynin (2018 Proposition 2.7.1) this implies the quantity My Pr sub-exponential with I|u Mrllv, I R² AARR/VV /vd for some absolute constant 0o 0. We therefore
have, for some absolute 03 O, w.p. least 1 8, for all [B].
(MT:Pg)) S cgRl log (2B/6).
(9)
Again using Lemmas 6.2.2 and 6.2.3 from (Vershynin 2018), since H4g has sub-Gaussian norm at most cR/ Vd and of 1/2 g" where g" has sub-Gaussian norm at most c, we have
13 (3 T x4) < Elexp(cHdd11333 (2))'),
and thus provided (RRd-1/210 - / A/ 2| we have
= E/exp(cRdd~/23g) / < " 18 tr()).
In particular, the quantity M' T = is sub-exponential with sub-exponential norm llug テ 11 S ca Rd Rdd-//2 /tr(A), and so for some absolute cs 0 we have with probability at least 1 8, for all q, 7 E [B] with *
Kpg ++ح1 S og 1/2 1 tr(A) log(2B²16).
(10)
For (zg: 2) with 5 * q we can directly use the MGF of Gaussian chaos (Vershynin, 2018 Lemma 6.2.2): (2g 24) = g Ag for iid. g.g ~ N(O, Ia) so that for 3 $ c/1|A|12,
(3 41) S exp(ce3² ||A) = exp(ce62 tr(3)).
In particular, [|<=,,= S 0 / tr(A2) that sub-exponential concentration implies that with probability at least 1 - 8, for any g. [B] with +
(2/ 2
07 / (r(A2) log (2((BB/)))
(11)
Putting 9 10 and
into 8
we get for a + T,
R²
R tr(r
+ v tr()?)
(iq. pT))
C8
Vd
VNd
log(2B² (8).
(12)
As for |1?, from 8 we have
[2 1 1 1|
2N 1LN-11212)) 1/2 (24. 14) 1
(13)
From here, the same argument used to bound (1O) holds since that bound only relied upon the fact that Py and 2 are independent, while H4, and of are independent as well. In particular, with probability least 1 - 6, for all E [B].
Ku,, x S eg 1/2 log(B//))
(14)
Next, have that - d g T Ag g|A for g N(O, Ia). Therefore Vershynin (2018 Theorem 6.3.2) implies
|||| || ||A//2|| & ell |/ for some absolute constant 0. Note
14

---

Published as a conference paper at ICLR 2025
that ||/ (t() And thus by nub-Gaussian concentration, have for some constant Cg > 0, with probability at least 1 8, for all T E [B),
|| - v(r(A)l - 2 IA /2 log((B/0))
| & 2(v(r(A) V V //
In particular,
| < 4 (tr(A) co log(2/)))
(15)
Putting 15 and 14 into
and using that || 112 R², we get with probability at least
28,
SgRyung (r(A) log (2B/) (r(A) Vall cal AlI2 log(2B/o)
| 112 R2
VNd
As for |,,1| by definition,
|| 2 = || 1| 2(U 27) -|=-111 R² 2(17 z+) 1|=-1|1
Since 2, ~ N(0, A) has the same distribution as =/, the same analysis used to prove 14 and 15 yields that with probability at least 26, for all [B].
Ky4, 2,)1 S cg Rd 1/2 (2gB//)))) || v (tr(AA) v eg /g (28/))))).
Substituting these into the preceding display we get
||| 1| R² 20g RegRvtr((an (t))) log + (tr(A) V |a (a log(2B/6))
dd
Thus provided R is sufficiently large, then also have ||> 112 @
Next we bound (xr, 1y)l: There 24, 24 i.dd N(O, such that (V-Ir, Fg d (ur + 24 My 24-)
It is clear that the same exact analysis we used to analyze (8) leads the claim that with probability at least 1 - 6, for all # T:
R Rvtr(A)
vtr(Ae) log(2B²/0) /8)
(iig xT)) S 0g
(16)
Vd
vd
Finally, we consider V- xr. Just as in the previous analyses, there are 2, N(O, such that
(p. vー) _ (ut + NN 1/2 2r, Hr + 2/4)
1 1/2 (z), 1+) + (2/, + N (2,,
Again using analysis similar to that used for (8) yields that with probability at least 8. for all T E [B).
(1+
1 R/tr(A)
V(T(A2)
(pr.V.xT) R² S c10
log(2B/8).
(17)
Vd
VN
Taking a union bound over each of the events shows that all of the desired claims of Lemma A.1 hold with probability at least 1 106.
The events of Lemma A.I hold with probability at least 1 - 108, independently of the assump- Our results will require this event to hold, and so we introduce the refer this in later lemmas.
tions (AI) (A2) and (A3)
following to allow for
Definition A.2. Let us say that a good run holds if the events of Lemma A. hold.
Noote, however, that the quantities appearing on the right-hand sides of each inequality in the lemma only small when these assumptions hold; this the reason for these assumptions.
15

---

Published as a conference paper at ICLR 2025
B ANALYSIS OF MAX-MARGIN SOLUTION
In this section we derive a number of properties of the max-margin solution (3)
Lemma B.1. On good run, for any CB 0 and for C > 1 sufficiently large under Assump- tions (A1) and (A2) the max-margin solution W of Problem
B
W
Ax Ur 2,
such that the \ 2 0 satisfy the following:
(CB A 1)d
B
4d
\
Sc RA log (2B²/6)
R4
==1
where > 1 is the constant from Lemma A.1
Further, we have the inequalities
(cB V)
2/d R²
1 [ -
16 R log ( (2B²/5)
and
(CB A 1)d
6d
Rt log (2B² /0) <tr(W) S
R²
Proof. We first derive upper bound | by showing that the matrix U Ia satisfies the constraints the max-margin problem
V. FT 1x,
(fr,, yr
(1)
V (t((A) Rvd
tr(A²)
R2
1-200
1 - 200
log(2B/8)
N
(40)
4cu
R²
1
C
1
R2.
(18)
2
while inequality (Hi) holds for C > Sco sufficiently large via As- sumption (A1) Thus the matrix 21a/R² R2 separates the training data with margin at least 1 for every sample. Since W is the minimum Frobenius norm matrix which separates all of the training data
Inequality (2) uses Lemma A.I
with margin 1, this implies
| S 121a/ 2/d
(19)
R22
On the other hand, by the variational definition of the norm, since |Ta/ = 1 we know ||]| 2 (W, Ia/ Vd), and hence
|
(,, Ia)
Vd
B
A par a. la
Vd
B
1
Σ
Xy (Tr v+.+r)
F=1
1
1
R²
(20)
Vd
2
16

---

Published as a conference paper at ICLR 2025
where the last line
18 Putting this and the preceding display together, get
E+
(21)
R
Next, we note that by the feasibility conditions of the max-margin problem, we have for any 7,
S yr Fi, Wx+
B
17,
Σ
Aolnttinl your
9=1
Xa (ir: Fig (yr yosta)
x ||
| |
Σ
Mg Paa) (e x-H-)-
(22)
q: qat
Now by Lemma
A. I
we have
[ 11
Go (r((A) log(288)
tr(A) V co collAl12 (2B62))))
R²
RVNd
NR
and
|| 1|?
co / (r(A) log (2B/0)
tr(A) V co Alll log (2B/6)
122
RVd
R02
Using Assumption (AI) we see that for C sufficiently large we have
7 || R²d
co
4cD
(1+
cD
4co
s Rd
(23)
NC
NC
C
2
where the inequality uses Assumption (A1) Likewise, we have
| |1 | 1| N R4
(24)
As for the cross terms, again using Lemma A.l for all + q.
Kpr, payil
R²
R tT/t((A)
(T((A2) N
R RV(A(A)
+
+ (r(?))
log l2 18)
VNd
Vd
R'
14 Vi((A
Vdir(A3)
V VI(A))
(dtr(A2)
14
+
log" (2B2/8)
d
RVN
RºN
R
R
( 2RR
1
1 (1+
1+ CVN
log² (2(22//))
d
4N
(8) Rd log2 (2.B
(25)
Inequality (i) uses Lemma A. and (ti) uses that C' is large enough. Putting this together with (23
and substituting these into the consequences of the feasibility condition (22) we get for any
3 R' + 2c Rs 12B²(2B111) 18)
1 5 2
Ag.
(26)
d
We now show that this implics E/= Ag 2 B Re log" (2 Towards this end, let first consider
(cgAlld
the case that there exists some with X, -
egd
Then the preceding display implies
< R²
CBd +
20 R 2 (2B2/8
Σ Ag
2
4R4B
d
T.
2c Rd log"
(2B2 18)
Σ Ay
<
4
d
q 9FT
17

---

Published as a conference paper at ICLR 2025
The final inequality uses Assumption (A2) i.. B cgd. Rearranging that
d
(cB A 1)d 8eg log (2B2/5)
8eg R* log (2B2/8)
9: g=f
Let us now consider the only remaining case, whereby for each we have X >
cgd
Then,
48
CBd
CBd 4R
(CB A 1)d
Σ
4R
RC log (2B2/5)
where the final inequality uses that co 2 We therefore have Σ Xy 2 Beg R log (2) which
(cgAl)d
together with (21) completes the proof for the upper and lower bounds of E Xg. The upper
bound for the Frobenius norm follow by (19). while the lower bound follows by Σq= Ag 2
(ea Al]d 8rg 124 log logg (27 and (2)).
For the trace term, we have.
B
tr (r(M)M) = ir
Σ
B
Σ
A tr(yr(urz))
T=1
Σ
Yr - Prr
T=1
(8)
Σ
X₂ R? /2
(1)
(en A 1)d
(27)
16c log" (2B
Inequality (i) uses (18) while inequality (ii) uses the lower bound for E, )r. Similarly, by
Lemma A. have
(1
|+
(r
(A²)
3
1r. (+r)) R2 1
co
log(2B/8)
S R²
VN
Rvd
VR
where the last inequality uses Assumption (Al)
Therefore,
(MM)
=
Ary 2. Po
r=1
3
X
R2
2
6d
R2
where the last inequality uses the upper bound for ΣB X proved earlier.
C PROOF OF THEOREM 4.
Our goal is to bound
R(WMM) = IP (a.
WMM)) # yM+1)
1
PP
W <0)
(28)
M
18

---

Published as a conference paper at ICLR 2025
The probability above is over the draws of j and of (I,, Ur) M+1 =1 and μs. The matrix WMM is the max-margin solution when training on the in-context tasks (3) To prove the risk is close to the label-flippig noise rate p. it suffices to show that the error on clean examples is small:
R(WMM) IP (yM++1 sign (FF;; E; WMMM) < 0, MM++1 # Das+1)
+ P (VM+1 (ign( E: WMM) < O, WM+1 = iM+1)
S (yM+1 sign 7 E; WMM) 0, yM+1 (M+1)).
(29)
For notational simplicity, let us denote 7 = 77( E; W MM.. Let us also denote the event Å as Å = (ynr+13 19 0) Vv Elym4+ 01
-Elys-1011
so that R(W (M)) P(A). Further, let us introduce the sets C,N C [M] as the clean and noisy test examples respectively (IC) + [ M) so that
= - E. Yi U ViE
Then we have the identities
yos; μ + yizii E
Vrz; = - + Yili, iEN.
Therefore there exist independent SMM41 N(O, A) such that
1 Σ y: (1
2/A7
1 -G, c.
ŪM+13M-L = f CM+1:
μ+
M
M
VM
where we have used that |C|- W = M 2 In partieular, if we define
= ( + M-1125) W(u+ (M+1) Elyv++10] < -Bly
then we have
P PUMM+iŪ 1[uM+10] < - - UM+1 = 1M+1)=
= =((A)
Continuing from
29 this means for any oo: on 02, 03 0,
R(WMM)
Sp+P P (000(((111)))))0 (OI-AN//))++//)
T W(u + (M++)) - EluM+ -E(yM+13)
S p + P(WI/M S an) -P ( > a1)
1 //CTWMMGG-11>> 02) PP((1 2/N1/M))μ² WMMMS++1>> as)
[n an) (1M-1/2 WWMMAl and
+P
WMMS S a2) n flı WMMSS44l s oss)
(30)
In particular, to derive an upper bound on the risk suffices to derive a lower bound on Wmm and upper bounds each of the absolute values of the four quantities [W/M,, WMMu, T WMMS141 and μ WCN+1. We shall do so what follows.
C.1 WI/M
Lemma C.I. There is some constant e > O such that for any 2 0, have
W = 0, P
0,
W
21) < 22exp(-22 -262 > 0.
(31)
p
M
19

---

Published as a conference paper at ICLR 2025
Proof. If p 0, there is label flipping noise and so [ = 0 deterministically. For the p case, by definition, [W] = Σ L(y: A V/)) a sum of M independent random variables bounded between 0 and with mean By Hoeffding s inequality, for any u 2 0 we have
2u²
P( W1 Mp 2 u) = P(( WW/M - 12 u/M) S exx
M
Setting a = ML completes the proof.
C.2 H WW
The quantity j T WW a quadratic form, and we can provide a concentration inequality this in the following lemma.
Lemma C.2 (Hanson-Wright for uniform the sphere). Let R >> 0 and E Rddx be matrix. μ Unif(R- 3d-1), then for any t 0.
R Qu d (r()) 21 =P(
td
Elu TT 2 2) Qu -P(TTQQr
exxp -emin
RU1Q1 R 12110112
Since Tu Qu is a quadratic form, we would like to use the Hanson- Wright inequality here. But the trouble is that for μ ( the components of μ are not independent, so the standard Hanson-Wright Wright inequality is not directly applicable, but requires some additional work. Our proof instead leverages the following:
Theorem C.3 (Adamczak (2015), Theorem 2.5). Let X be mean-zero random vector in Rd. If X satisfies the K-convex concentration property, i.e. there exists K such that for every I-Lipschitz convex function RRd TR we have EX)(X) oc and for every > O,
P((((X) EØ(X) 2 t) < 2 2exp( 4 K2),
then there exists such that for any matrix Q dxd and every 0,
t
P(XXTQX E[X E[XQQX]]
2 < exxp
-emin
2K' K2 K²1Q
Proof of Lemma C.2
By Adamezak (2015), sufficient condition for Hanson-Wright--yn in-
equality to hold is that the random vector le satisfies Lipschitz concentration. By Vershynin (2018) Theorem 5.4. we have that for some constant e 0, for any 1-Lipschitz function on the sphere vd- gd-1 and for X Unif( (Vddd-1)..
P()((X) E[f(X]I 2 t) s 2 2exp(-ct))
Thus if we consider g R. gd-1 + R which is 1 L-Lipschitz and is defined on the sphere of radius R, then Y
R X where X is uniform the sphere of radius R. In particular, the function Vd
R g(Y) = g( vd
X) the composition a I-Lipschitz function and a R/ /d/iiiiht function defined on the sphere of radius Vd, and we therefore see that
P(lg(Y) Eg(Y) 2 t) - 2f d/R)).
That is, any 1-Lipschitz function the sphere of radius R satisfies the K-convex concentration property with K R/ vd. Thus Theorem C.3 implies for any matrix Q.
12
ld
P(ITTQu Qu E Qnl 2 < 2exp
-cmin
Re Re pe11Q11
Finally, note that by rotational symmetry we have E[p]] = 0 and E[n'] = Ñ2 /d for each i E [d].
Additionally, for distinct components j. ElugHy] (014((14;]] = 0 and hence E[p Qu] =
Σd Eb ae tr(Q)-
20

---

Published as a conference paper at ICLR 2025
C.3 WWC AND WSM+1
Lemma C.A. Let μ R Unif(Sd-1), g N(O,T, Ia) (independent of p)) and let E I be matrix. There is an absolute constant c >> O such that for any 2 O,
ctvd
P (Iu Qgi 2) S 2xxp
RI|QUIF
Proof. The uniform distribution on the sphere of radius Rt is sub-Gaussian with sub-Gaussian norm satisfying || < cRt vd by Vershynin (2018 Theorem 3.4.6), while | I c, for some absolute constant c By Vershynin 2018 91: ~ NO0, Ia) and any B R
Lemma 6.2.3), this implies that for independent
Vd 3 1 R
Eexp
Qg
Eexp (c1 Q92)-
Then using the noment-generating function of Gaussian chaos satisfying [3| s c2/1|Q1|2 we have
/Vershynin 2018 Lemma 6.2.2), for
vd BBY p R
E Exxp
B
TQg
S Exx (c1,0] Q92)
S <exp(c3321Q113)
That is, the random variable R-1 'd'v Qg is mean-zero and has sub-exponential norm at most
max ca))llQI There is therefore constant > 0 such that for any 2 0,
P(IR-Lvap Qol 2 u) P (ITQ9l QgQ 2
Nu Vd
s
Setting LVd/Rt we get
(_ava]
clVd
P (1 T Qgl 2 2t) 2xe
(TRIQQFF)
C.4 WCC
Lemma C.5. Let S, Lid. N(O, A) and let Q E Rddd be matrix. There is a constant c 0 such that for all t 2 0,
d < 2exp
(- at
[
Proof. We can write C WC as 96 A1/2WA2/2 where 91.92 N(0, Ia) are independent,
ie. is Gaussian chaos random variable. By Vershynin (2018) Lemma 6.2.2) this random vari- able is sub-exponential with sub-exponential norm at most c/[A//WW 1////: Then standard sub- exponential concentration ( Vershynin 2018 Proposition 2.7.1)) completes the proof.
C.5 Ely WMMp4]
Lemma C.6. On good run, for any eg > 0 and for C > 1 sufficiently large under assump- tions (Al) through (A3) the max-margin solution WMM satisfies
(cB A 1) 16c log (2B2/8) S Bly WMMMj]]
6
S
R
21

---

Published as a conference paper at ICLR 2025
Proof. By definition,
R.2 tr tr(Wwm).
E
[p WMmpl
(32)
d
We therefore need only upper and lower bounds on Lr(WMm). Using the bounds on Lr Lr(WMM) MM) from Lemma B. we get the desired inequalities.
C.6 PUTTING IT ALL TOGETHER
Let's denote RMWMM) the test error for clean example (yM+1 = (M+1). Returning to the decom- position (30) and using (29p, we have
R(WMM) Sp (1
E (1-2W1/11))" Wp
W -Elvöl
< Elyy - 1/2 T WH - (1 2////a) T WCM++1 M- WSAs-1
(33)
We'll consider two cases, depending upon whether the label-flipping noise rate = 0 or > 0. For notational simplicity let's denote p as the quantity
CB A1
p:
16c log
l
(2B²
E (0,
Then Lemma C.6 states
Rt P: R
6 R²
R
S
W
W S
(34)
Noiseless case (p = 0). Since = 0 we know W = 0. From
33
we have
R(Wmm)
P
W Wp] < EEE WA]
M 1- WH p W Tta.://TsW... WSM+1)
< P
< TW]
Ww]
M M-1 WH 11 WSM+1 M-112CTTCCMH
R2
n&²
pR
< P
W El4 T Wpl
+P |
1/- W >
+P | [NTWCM+2|
2
2R²
8R²
8.R22
+P (IN- S.TWCMM T
Rス²
(35)
8
Inequality (i) uses Lemma C.6 inequality above.
We now proceed by bounding each of the remaining terms in the
For the first term we can use Lemma C.2 with = pR2/2R2 To do so we need to examine the which we have
quantity R']lW|F"
td
pd
pd
PVd
R
2
W
2R 2/
4
22

---

Published as a conference paper at ICLR 2025
Let's assumee for the moment that pvd/4 4 > 1. We then apply Lemma C.2 lemma holds if we replace the || with [[Q[[))
(noting that this
pfte 2R
d2
p² ft
pf
P Wp - Elu T Wo] <
S exx
-emin
4R
W2llWil
2R
cdp
S exp
2 | R
(
.cpV
- 2 exp
(36)
4
Note that if pvd/4 < then the above bound still holds since P(.) is trivial inequality. This completes the bound for the first term in (35).
The second and third terms in 35 can be bounded using Lemma C.4 g N(0, Ia).
since SM-1 = A2/2g g for
pff
pl
P 1 WCM+1l
P
8R²
8R
cVd
pR²
2xxp
R
8R
cVd
RR²
2xxp
R|AAV2 | F
8R2
cplivd
< 2exp
8 2V
cpR
exp
(37)
|6||
The incquality (i) uses the upper bound ||]
S 2/ R² from Lemma B.1
For the final term in 35) we can use Lemma
C.5
RR
PM1/2
IP
M WCM+1 2
= P
|s WCM+1
8R2
8R²
PM 1/2 R2
2exp
W
SR2
(6)
C
1M2/2
< 2exp
|A 2VRRR
8R2
cpM-/2 R2
2exp
(38)
16]|All/vd
Again (1) uses Lemma
C.5
Putting together 336
(38 get
cpvd d
cpft
[cppMt/2 12 (-16]AA]2v)
R(WMM) I exxp
4 exp
2exp
16|
Noisy case. Returning to 33) before, the 'signal' in the problem comes from the term (1 2.Vi/M)p TWP. WW, which ideally is large and positive. It is natural that we should require more samples M the closer the noise p gets to 1 (namely, the smaller ep is), since otherwise with
23

---

Published as a conference paper at ICLR 2025
nontrivial probability will noisy examples than clean and learning should be Cp for some absolute constant Cp E (0, 1/2).
impossible. To this end, let's assume that p $ Continuing from 331 which shows that
1 2
P 1
1 2/M//
P
W M
+ 2cp
6 S
(
s 22xp
18
22exp(-c2.)
2
16
(39)
Let's call the event
W|
=
Cp
M
Continuing from 33 since 1 2/W1/M & have
R(WMM) SP
E yy
E [9]
M --1/22 WH (1 2//1/MM) WCM+1 WCM+1- M -
WCM--1
+
(1 2/7
(1 - 2WW1/a))
Spt PPE
Elyū]
Ely0l
M -- W (1 22/N///))
WCM+1 T M-1/2(TWCCC WCC++
+P E, WW
(1 2M/MM)
(1 - -21N1/
3 pp 2 2exp( -2-M)
Ely
Bbyūl
M M- WH TW W M-1/25 T WCr++1
WCM+1
+P E, 4¹ T. WW (1 21/MM
(1 - -21N1/2)
Spt + 2exp(-eppppp -2,))
E[
+P E, TW
(1 2/ 2//1//a))
El
[ TW + +11 11- 2///MM 11 WC + IM 1-2 WTWWMMM (Aamel))
+ 205
W
(40)
The inequality (i) uses (39). We want to deal with the quantity H TTp. on the event E. We have,
Ei4 TW 1-2
Bly
(1-2(N//M)
1 W R² 2p) (r( d(1 2112/W1/M
R R (r((W) TW
1 2p
d
2//1/m
RC tr(W) d
R22 tx(W))
(1- - 2p
=p Twp
d
1 2//1/m
Re ir(W)
72 (r(W)
-2p + +pp 2/N//NN
=
d
1 1-2NN/M
24

---

Published as a conference paper at ICLR 2025
Continuing from (40) this means
R(WMM) S p 2 exp M2 + P E, M W
R tr(W)
R² tr(W) -2p + 2//1/M
d
1-2/N/MM
1ie(1 2p) Ir(W) + 20pp |M- Wnl + 2////rl |n WSM++l |M- WGrHHil) [WCM/1
d(1 2/MM/MM
=p+ 2 2exp(-ce.
+P 3, n1 WH Elu W
R2 (r(W)) 2cp
M - Wo] + +11- 2w/1/a 1p WCTWS4 + 1M- (M-V/PCTWGMHL) WWMAHal)
d
(1)
$ p + 2 (xxp-cc.M
+P E, p "TW - 1'E'" W
1e2
20p11 [M-1/2/ + 1111// 1n WCar+1l | W (arril]]
Pp
R2
Rt
pf
S p 1exp(-<2,M) P( T Wp Blu W < 20
+
20, 20-2//1 Wpl
R²
8R²
pR SR²
pff
+ P (IW1MM 1p WCM+11
+ F
| T 2
(41)
8
Inequality (i) uses 34 From here the proof is exactly the same as in the clean case. For the first term, we use similar arguments used derive (36). To apply Lemma C.2 we need to examine the quantity [/
id when t = pf2 /2R²: we have,
Id =
pd
pd
PVd
(42)
R2]llll 2R
2R2 2/RR-2
4
Again if pvd/4>>///// 1 then we have by Lemma C.2
pR²
d2
p² RR
R
P u WPp Elu Wp]
S 2 exx
-cmin
212
R+I|WI? ?
4R
R R
2
edp
< 2exp
2 | R?
(8)
cpvd
$ 2exp
(43)
Incquality (i) uses 42 a trivial inequality.
Note that if dd/ < then the above bound still holds since P(.) is
As for the u WCa++1 and <WW, terms, similarly as to the analysis used to derive 37 via
Lemma we have,
pft²
Spplie
TP
2cp 1n WCM+112
P 11 W/2/2
(g N(O, T))
8R²
16R²
ocpplt
< 2 2xxp
(44)
32/
25

---

Published as a conference paper at ICLR 2025
Note that / /M| < 1 always so this takes care of the term involving μ W WSM+1. Since the analysis used to derive (37 only relies upon upper bounds for ||W||= | Fr and since M 2 and Cp (0, 1/2), the same analysis holds for the term TWp..
Wpl.
For the < WSM+1 term, the same analysis used for 38 applies here as well, since
pR²
/ R22
P
MM | (
W 2
= IP 1s WCM+11
8/
8112
12M
< exxp
(45)
1611A|laad
Plugging
43
and (45) into (11)
we get
.cvVd
CCppR
cpM R2
R(WMM) S p 2 exp(-c + 2exp
4
+ 4exp
32 |2
+ +
1661A
By adjusting the constant c to depend on Cp. this completes the proof of Theorem 4.1
D PROOF OF THEOREM 4.2
Recall that we assume that A = I in this section, and for simplicity denote W WMM- Our goal is to show that for each k;
n W Wy
>
We have:
M
W
yx =
W WyLTk
M
i=
W
Wo
+
Σ y:
Wxy.
M
Our goal here will be to show that the first quantity vi wg is large and positive and dominates the second term involving x. W Wxx.T By definition we have
2 W Ig
(n - ) W (a + 3k
= T + 2 W + WW + =1. W 2k.
(46)
In particular we have the identities
(p T W y < 0}
(xT War Vx < - Eiff Vilnar, Wor}
( WEb mか - -tr(W)- < - -r((W) - yu Wp - Wa - Wark
Wy
C ( Wat tr(W)- < ~ (r( 21 Wpl> 1 [[(W))
U UlNT W (r(W)) Wp
Tw < - (r(W))
[ ( Vinn. D Wark
> (ww) (47)
8
We will therefore proceed by bounding the probability each of these events. We will use the same notation from the proof of Theorem 4. presented Appendix C whereby E refers to clean examples (xt, y:) with ū: yt, while i E V means y -
26

---

Published as a conference paper at ICLR 2025
D.1 Σisr 2, War TERM.
We have,
Eve
E
M. +
iEC, igtl
iEN, Egehe
Σ
(1 + U1.k) +
E
(-p (1)2k)
iEC, if
iEN, Gp
(IC \ [b}] W (1)D)+
Σ
DKZK
Σ
Un 2k-
VEC, itt
VEN, ifth
Now since the label-flipingg noise is independent of zg. the random variables dn are i.i.d. standard normals. Therefore there is ~ N(O, Ia) such that
Σ
U
Σ
ÚKZZ = VM 1gk-
KEC, 1
SEN, BEER
If we denote N₂ : |C { W \ {K} then clearly NN S M and thus we have
Σ yaz Nu + VM - Igk.
where |& < M and is standard normal and is independent of z and yk- Therefore
Σ Man; W yu
(NNAH VM Igx) T W (UNDRR + 3R
S MIK W + VM W + Ml T W + VM ON Wz
(48)
In particular we have
Σ yo W WyLkk
tr(W)
tr(W),
IP
< PPM Wp W
8
32
tr(w),
tr(w),
tr(W),
(P(V M WM)
+ P(M
W
(49)
32
32
32
We'll proceed by bounding each of these terms in sequence.
Wh term. Let L be an integer (we will take L = M in this proof but in a later proof we will take 1). Since tr(W)= > by Lemma B.I have
tr(W),
} f1p TW R² (r))
R trt((W)
tr(W),
(LL W
32
d
32L
This means
P (ILUTWAL Tp
tr(W)
1ta (r((W)
tr(W)
32
P Ww
1
32
32L
d
Assume for now d > 128p
R2 L so that 32 L/d 1/2. We then have by Lemma
C.2
tr(W)
R (r(W)
(r(W)
<
H Wp
32
64L
de
tr(W))
2
tr(W)
< 2 exp
-emin
64L
64L
As for the quantity appearing in the minimum we have by Lemma B.I
dlr(W)
pd
R²
pd
pd
64L
R2
64L
R
2R
128LR
128LR
27

---

Published as a conference paper at ICLR 2025
We therefore have
tx(W)
epd
P [Lu T Wal
exp
(50)
128LR
We see that the case d v 128p-1 R2 L results in the same inequality to hold (albeit vacuously). This completes the proof for the μ W j term in 49 by taking == M.
Ok W and Zzz terms. have
The terms P Wg. gb and 2 W can be controlled using Lemma C.4
For an integer
(- eVa
tr(W)
tr(W) 8L
P
Llu
W [LlTT
S 2exp
[Ri]W]l
32
We have
Va
tr(W)
Revd
pd
pd
R|
W
32L
R R-2vd
32R
64RL
Note that the above argument applies to W as well as W, so that we have
tx(W)
tr(W)
cpd
IP
L W
VP
L/u W
S 2
(51)
32
32
64RL
particular, we have
tr(W)
cpd
IP
V M
W
S 2xpp
32
64RvM M
tr(W)
cpd
IP
Mlp
W
< 2 2xp
(52)
32
64RM
9k W term. The last term 9 W 22 is Gaussian chaos random variable. By Vershynin 2018 Lemma 6.2.2) this random variable is sub-exponential with sub-exponential norm at most W Therefore we have
(r(W), = P(19 Wa
Lr(W)
(r(W)
P( PPv M 9h Wzd
I 2exp
-c
32
32VM M
322MM/WWFF
Since Ir( lr(W)/|W]lr 2 pvd/2 by Lemma B.1 this implies
tr(W),
cpvd
W
9h
s 2exp
(53)
32
64v M
Putting it all together.
Continuing from
49
we can use
50
52 and
53 to get
(r(W)
cpd
Σ yr
WYRKK
< 2exp
128M
(- cpd
cpd
cpvá
+ 2 2exp
+ 2exp
+ +2exp
(54)
BARYM
64MR
64vM
D.2 25 War TERM.
It remains bound the probability of the first three events in (47
28

---

Published as a conference paper at ICLR 2025
2k W2 term. First note that since 2 is standard Gaussian, by standard Hanson- Wright in- equality (Vershynin 2018 Lemma 6.2. 1) we have for some constant 0 < c < 1, for any 2 0,
P( Was zΙ. - tr(W) 2 S 2exp (-emin
| ))
[
So for - (r(W) have
(r(W)
pd
R²
pVd
W
2 W|
R
4vd
This means that if v/d4 / 2 1 we have by Lemma
B.1
P
22 W2. tr(w) tr(w)
- 2 exp (-cpVd/4).
(55)
Since the same inequality trivially holds if pvd/4 4 (by potentially taking a to a smaller constant), this completes this term.
Wir term. This is covered by (50) with = 1:
tr(W)
apd
IP
14 Wal
< 2 exp
(56)
32
128R2
2g Wa terms.
Since gl and 2μ i..d., this is covered by (11 with L = 1:
tr(W)
tr(W)
(- (_opd cpd
W Fk
VP
W 2x
2exp
(57)
32
64R
Putting it all together.
Continuing from 47 we have
(p W yK < 0)
24 W2k tr(W) 2 tr(1)) +
22 WH] (
P
+ PP
[ WEW
(r(w)) +P
1 WW < - - (r(W)) + P
ViNKe, ir(W))
cpvd
cpd
cpd
$ 2eep
+ 4expp
+ 2exp
64/
128M
cpd
cpd
cpva
+ +expp
+ +exxp
+ +xexp
64RVMM
64MR
(6AVMM
cpvd
cpd
<exp
+ 8 exp
64VM
(R v R)
Using union bound this allows for to conclude
[<pVV
cpd (T128M(FRV R)
P B3k
[M]
s.t. ū W WYLTR < <0) 4M exp
8M exp
64/M
as claimed in the theorem.
E EXPERIMENT DETAILS
We describe here the experimental setup for Figures 2 and 3
code is available on Github B
httg:://ghhhh..ccmmmmmeeefffefee classification
29

---

Published as a conference paper at ICLR 2025
We pre-train models using standard full-batch gradient descent the logistic loss with R 5vd, N = 40, learning rate = 0.01, for 300 steps from a zero initialization, using PyTorch, The in- context training accuracy is measured using the definition of training accuracy from Theorem 4.2 namely, we look at what proportion of the in-context examples (training data) that is accurately predicted with the model 77( (.F.MMM W), where W is the trained transformer, for single task 7, i.c. averaging (1% = sign(y( L7 (rx); W)) over = I..... M. The in-context test accuracy is computed by measuring whether sign( (#M+1); W)) MM+1- We then average over 2500 tasks, and we plot this average with enror bars corresponding to one standard error over these 2500 numbers. All computations can be run within an hour on high-quality CPU, although we used an
NVIDIA RTX 3500 Ada which helped speed up the computations for the B d z 1000 setting.
30


Reviews:
[{"review_id": "EO0nBiuJOf", "reviewer": "Reviewer_dEuQ", "summary": "This paper studies the ability of linear transformers to perform in-context learning for linear classification tasks. The authors consider a simplified, convex linear transformer architecture trained on a series of random, class-conditional Gaussian mixture models.\n\nThe paper investigates two key aspects:\n\n1. the number of pre-training tasks required for the transformer to generalize well at test time, even when the test data has lower signal-to-noise ratios (SNR) and label-flipping noise not present in the pre-training data.\n\n2. the phenomenon of “benign overfitting in-context,” where the transformer memorizes noisy in-context training examples while still achieving near-optimal generalization for clean test examples.\n\nThe authors claim this work to be the first theoretical analysis of in-context learning for linear classification using linear transformers, and also the first demonstration of benign overfitting capabilities.\n\nThis paper makes valuable theoretical contributions to the understanding of in-context learning in transformers, highlighting their generalization capabilities and the intriguing phenomenon of benign overfitting. While the simplified model and assumptions might limit direct practical implications, the paper opens up interesting avenues for future research on the theoretical foundations of in-context learning.", "questions": "Practical Implications: Recognizing the limitations of the simplified setting, can the authors discuss any potential practical insights for realistic problems?\n\nExploration of Alternative Architectures: How would the results generalize to more complex transformer models? \n\nPre-Training with Noisy Labels: The study assumes clean input data and labels during pre-training. How would the result change once noise is introduced in the pre-training data? \n\nNumerical Validation of Theoretical Results: Have the authors tried conducting numerical experiments to validate your theoretical findings? It would be highly insightful to see how well the theoretical predictions align with empirical observations in both the simplified setting used for analysis and in more realistic settings involving complex transformer architectures and real-world datasets. For instance, could the authors present experimental results demonstrating the relationship between the number of pre-training tasks, the signal-to-noise ratio, and the generalization performance, as suggested by your theoretical bounds? Furthermore, showcasing empirical examples of benign overfitting in-context would be compelling evidence to support this key finding.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Rigorous theoretical analysis: The paper offers a detailed analysis of the implicit regularization of gradient descent during pre-training and leverages the max-margin framework to derive generalization guarantees.\n\nNovel insights: The paper reveals the intriguing ability of the trained transformer to generalize under lower SNR and label-flipping noise at test time, even when such conditions are absent during pre-training.\n\nDemonstration of benign overfitting: The paper provides theoretical and experimental evidence for the previously unobserved phenomenon of benign overfitting in transformers.\n\nThis paper presents novel theoretical insights into the in-context learning capabilities of transformers for linear classification tasks. The demonstration of benign overfitting is particularly noteworthy. While the simplified assumptions and focus on a convex architecture limit the direct applicability of the findings, the paper offers a valuable foundation for future research in understanding in-context learning in more realistic transformer settings.", "weaknesses": "Simplified architecture: The analysis relies on a 1-layer linear transformer model, which is a simplification of standard multi-layer transformer models with softmax-based attention. It is unclear whether about the applicability of the results to more complex transformer architectures.\n\nNoise-free pre-training: The requirement of clean data during pre-training restricts the practicality of the approach in noisy environments.\n\nStrong assumptions on data distribution: The paper focuses on class-conditional Gaussian mixture models, which might not be representative of the complex data distributions encountered in natural language processing tasks.\n\nAbsence of Numerical Validation: While the theoretical analysis is valuable, the paper lacks numerical experiments to validate the theoretical findings. Providing empirical evidence to support the theoretical claims, and demonstrating how the results generalize to more realistic settings, would enhance the paper's impact and persuasiveness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "374tJrmWiL", "reviewer": "Reviewer_5w9f", "summary": "This work presents a sample complexity analysis for in-context learning of class-conditional Gaussian mixtures with a restricted linear attention transformer model (single layer). It uses the results established on how SGD on convex-linear transformer models has a bias toward maximum-margin solution (in direction). The work thus assumes the KKT conditions are satisfied at the SGD solution and uses them to quantify the sample requirement to achieve a small test error. It shows how --by choice of specific model and data parameters-- one can achieve benign overfitting for in-context training, a phenomenon in which noisy training examples are exactly memorized, yet the test accuracy remains near-optimal.", "questions": "- Could tighter bounds be established for stochastic (Gibbs) classifiers? This would be stochasticity in the loss that is not a result of label-flips on the test set, but is a result of stochastic classifier decision.\n\n\n- Line 377 breaks a formula.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The work references and uses prior work to great effect. It builds up on ideas established around the implicit regularization effect of the optimization algorithms used to train transformer models. SGD-based algorithms applied to convex-linear transformers have a directional bias towards maximum-margin solutions. Hence once can apply techniques used to study such solutions to derive generalization bounds such as the one presented in 4.1 and 4.2.", "weaknesses": "- Readability: The paper is very dense and in parts hard to follow, requiring multiple revisits to earlier sections for (non-standard) notation and technical definitions. This, unfortunately, severely cuts into the readability of the paper. Adding a table of notation either in the paper or the appendix can be helpful.\n\n- Analysis: There is no empirical evaluation of the tightness of derived bounds. IMO the derived bounds are to be understood in the limit and to prove the theoretical possibility of benign over-fitting. They might be too lose in practice. Adding an study in which $M$ or $d$ can be varied can help establish the empirical tightness of the bound.  \n\n- Scope: The work only focuses on a specific restricted linear model. It defers analysis on non-convex linear or softmax-based attention models to future work.\n\n- Assumptions: Even with the SGD-based solutions being directionally biased towards max-margin solutions, it doesn’t necessarily follow that the KKT conditions are satisfied in real implementations with early stopping or partial convergence. It’s also a leap to assume this (line of) theoretical analysis can be extended to other settings beyond the restricted or convex linear case addressed in the paper. IMO these assumption should be studied/validated under realistic training settings. Adding such analysis to the paper can clear some of these points.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RwRODOGvZj", "reviewer": "Reviewer_kNQi", "summary": "The paper investigated the behavior of linear transformer models trained on random classification tasks, extending the existing work on linear regression tasks. It studied how linear transformers can generalize and exhibit _benign overfitting_, meaning that even if a model memorizes noisy training examples, it still performs well on test examples. This paper is completely theoretical, and no empirical evidence was given.", "questions": "In Figure 1, the author used a natural language example to show the phenomenon of benign overfitting. However, I don't think it satisfies the data assumption. Is there an example/illustration that precisely satisfies the data/model assumptions?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This work theoretically analyzed how many pre-training tasks and in-context examples are needed for linear transformers to effectively generalize in linear classification tasks.\n- This work studied the phenomenon of benign overfitting in this context and bounded the generalization error.\n- The paper is solid and well presented. All the symbols are properly defined. The assumptions and theoretical results are clearly stated. (I did not carefully check the proofs.)\n- The paper is well contextualized. The author mentioned related work (that I'm not familiar with) and discussed future work directions, which helped me position this work.", "weaknesses": "My biggest concern is the same for all theoretical work on simple models (e.g., single-layer linear transformer) with strong data assumptions (e.g., class-conditional Gaussian mixture). It is perfectly fine to start with a simple model with strong assumptions, as long as it is a good proxy for real-world problems or a stepping stone to understanding more complex problems. However, we need either theoretical or empirical evidence showing it is the case. In the current paper, there is a lack of discussion (or references) on whether or to what extent real-world data satisfies the data assumptions. It is unclear whether the theoretical results can be generalized to more complex models. Overall, I'd like to know how the theoretical results can guide our practice, which is not completely clear in the current version.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nH3ZY6YUYf", "reviewer": "Reviewer_F6y8", "summary": "This paper studies in-context learning of linear classification tasks on linear transformers. The authors characterize how many pre-training tasks and in-context samples are needed for generalization, with an analysis of the implicit regularization of gradient descent. They show that after pre-training the transformer can tolerate smaller SNR than those tasks on which it was trained. Additionally, they show these trained transformers can memorize all in-context samples yet still generalize near-optimally when the in-context samples’ labels are flipped.", "questions": "1. At row 72, 73 → ‘even when pre-training on simple and easy-to-learn datasets, the transformer can generalize on more complex tasks.’ What does easy and complex task mean in this context? Does it mean the signal-to-noise ratio?\n\nMinor:\n1. At line 331, a typo of ‘sufficiently’.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The study of in-context learning is important, and the investigation on linear classification setting is novel.\n2. The results of benign overfitting of transformer, to the reviewer’s knowledge, is new.\n3. The paper is well-written and easy to follow.", "weaknesses": "1. It would have been interesting to discuss the connection of task complexity results under linear classification to Wu et al. [1]. In addition, briefly mentioning the importance of investigating linear classification instead of regression in the paper would be beneficial.\n2. Can you explain why it would be interesting to consider a different signal-to-noise ratio R\\tilde at test time? Analyzing the difference between M and N seems to be more interesting. \n3. The reviewer personally would like to see consistent empirical observation with simple experiments (such as a simplified version of Raventós et al [2]) but understand this might be out of the scope of this work. Maybe the author can consider mentioning it as a limitation.\n\n[1] How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? Wu et al. ICLR 2024   \n[2] Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression. Raventós et al. NeurIPS 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)