Transformers are almost optimal metalearners for linear classification
Authors: Roey Magen
Gal Vardi
Venue: NeurIPS 2025 poster ()
Keywords: metalearning
in-context learning
linear classification
transformers
Primary area: theory
TL;DR: We prove that, under appropriate conditions, linear attention is an almost optimal metalearner for linear classification.

Abstract: Transformers have demonstrated impressive in-context learning (ICL) capabilities, raising the question of whether they can serve as metalearners that adapt to new tasks using only a small number of in-context examples, without any further training. While recent theoretical work has studied transformers' ability to perform ICL, most of these analyses do not address the formal metalearning setting, where the objective is to solve a collection of related tasks more efficiently than would be possible by solving each task individually. In this paper, we provide the first theoretical analysis showing that a simplified transformer architecture trained via gradient descent can act as a near-optimal metalearner in a linear classification setting. We consider a natural family of tasks where each task corresponds to a class-conditional Gaussian mixture model, with the mean vectors lying in a shared $k$-dimensional subspace of $\mathbb{R}^d$. After training on a sufficient number of such tasks, we show that the transformer can generalize to a new task using only $\widetilde{O}(k / \widetilde{R}^4)$ in-context examples, where $\widetilde{R}$ denotes the signal strength at test time. This performance (almost) matches that of an optimal learner that knows exactly the shared subspace and significantly outperforms any learner that only has access to the in-context data, which requires $\Omega(d / \widetilde{R}^4)$ examples to generalize.

Transformers are almost optimal metalearners for linear classification
Roey Magen Weizmann Institute of Science roey. mwg.maawiiiaaaaaammmmmmm.. wmmannnneeiznnn...m ac
Gal Vardi Weizmann Institute of Science
gal. var.iweeidieaman ac. i1
Abstract
Transformers have demonstrated impressive in-context learning (ICL) capabilities, raising the question of whether they can serve as metalearners that adapt to new tasks using only small number of in-context examples, without any further training. While recent theoretical work has studied transformers' ability to perform ICI., most of these analyses do address the formal metalearning setting, where the objective is to solve collection of related tasks more efficiently than would be possible by solving each task individually. In this paper, we provide the first theoretical analysis showing that a simplified transformer architecture trained via gradient descent can act as near-optimal metalearner in linear classification setting. We consider a natural family of tasks where each task corresponds to a class-conditional Gaussian mixture model, with the mean vectors lying in a shared k-diimensional subspace of Id. After training on a sufficient number of such tasks, we show that the transformer can generalize
new task using only Ökk R²)) in-context examples, where R denotes the signal strength at test time. This performance (almost) matches that of an optimal learner that knows exactly the shared subspace and significantly outperforms any learner that only has access to the in-context data, which requires (dd R examples to generalize. Importantly, our bounds on the number of training tasks and examples per task needed achieve this result are independent of the ambient dimension d.
1 Introduction
Transformer-based models are the dominant architecture in both natural language processing (NLP) and computer vision. Since their introduction by Vaswani al. [1], transformers have been scaled to produce remarkable advances in language modeling [2], image classification [3], and multimodal learning [4]. Their strength lies in their ability to model complex dependencies through attention mechanisms and to generalize across diverse tasks with minimal task-specific supervision.
One of the most intriguing emergent capabilities of large transformer models is in-context learning (ICL). In ICL, a model given short sequence of input-output pairs (called a prompt) from a particular (possibly new) task, and is asked make predictions without any explicit parameter updates. This ability to rapidly adapt to new tasks from a small number of examples, solely by conditioning the prompt, has been observed in large language models [5], and is central to the ongoing shift toward prompt-based learning paradigms.
test examples from that task
The ICL phenomenon is closely connected to the broader framework of metalearning, "learning to learn" [6-8], which has been widely studied before. metalcarning, a learner is trained to perform well across a distribution of related tasks, thereby acquiring representations that allow for rapid adaptation to new tasks. It is often helpful to think of tasks as corresponding to individual users. For instance, generating personalized email completions. While each user provides limited task- ass-speccfiic data, such as writing style and personal preferences, there is a rich shared structure across users that
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

can be exploited. Metalearning aims to leverage this structure to improve performance on each task beyond what would be possible if learned independently.
In this paper, we consider binary classification, where we model each task as a distribution DD over labeled examples y) in Rd x {+1). During training, the learner has access B datasets, where the j-th dataset consists of N samples drawn i..d from task Dj. We further assume that each task labeled by a classifier that relies on a common representation h : R + RR, where typically k is smaller than the ambient dimension d. For each task D, there is a classifier ÍD= : RR - {+1} such that /p o h has high accuracy on DD. Our interest lies in studying families of tasks for which knowing this shared representation h substantially reduces the number of samples required to learn each task separately.
Metalearning Objective. Assuming that the tasks Di, Di...., DB are themselves drawn i.i.d. from an unknown metadistribution Ω2 the goal is output representation h that can be specialized to new unseen task D ~ Q. In the modern view, the model is first trained on tasks Di...., DB, each task contain N samples. Then, at test time, the learner is given M in-context labeled samples drawn i.i.d. from some new task D, and needs to classify a new sample from D, without further training. To evaluate performance, mostly the required number of in-context labeled samples M that required to ensure (with high probability) small error, we consider two benchmark baselines:
Single-task optimal learner: An optimal algorithm, in terms of the number of required samples, that has access only to samples from the new task D)
Optimal learner with access to the ground- representation: An optimal algorithm that has access both samples of the new task and to the ground-truth representation h.
The above discussion naturally motivates the study of transformers from the lens of metalearning. In particular, we are interested in understanding whether, and under what conditions, transformer architectures can act as (optimal) metalearners. Specifically,
Can a transformer outperform a single-task optimal learner (that only have access to the in-context data), and potentially approach the performance of an optimal learner with access to the ground-truth representation?
1.1 Our contributation
To address the question of whether transformers can serve as effective metalearners, we analyze their behavior in the well-studied Gaussian mixture framework. In this setting, each task a random instance of a class-conditional Gaussian mixture model in IR with identical spherical covariance and opposite means. The aus--specffic signal vector f sampled from shared low-dimensional subspace of dimension k < d. Formally, for some semi-orthogome matrix P strength R > 0, each task T 1...., defined by:
E R and signal
μ1, i.d. P Unif(R- -
Vr,i iid. Un)f((+ Z7,5 ild N(O. Ia), Ir,i Dr,ffly + Zr.
At test time, in-context examples are also drawn from a class-conditions Gaussian mixture model, but potentially with a different signal-to-noise ratio determined by a test-time signal strength R. That is. the meta-distribution may shift between training and testing. We consider simplified transformer with linear attention, a setup similar to many prior works on ICL [9-12]. The transformer is trained via gradient descent (GD) on the logistic exponential loss over the above random linear classification tasks.
Our main contributions are as follows:
We prove that if the transformer is trained on a sufficiently large number of tasks, then the number of in-context samples required at test time to achieve a small constant error on a new task (without parameters update) is Ö(k/ ((k/Fi fet), where R denote the test-time signal strength. In contrast, any single-task learner, which has access only to samples from the new task, require at least (2dd/ fe*) samples (see our discussion Remark 2.1 on the information- theoretic lower bound by Giraud and Verzelen [13]). To our knowledge, this is the first theoretical result establishing such a guarantee for metalearning with transformers in a linear classification setting.
2

---

Using the lower bound from Giraud and Verzelen [13], we show that even an optimal learner with access to the ground-truth representation, one that knows the true shared low-dimensional subspace P. requires least (2(k/ It*) samples to achieve a small constant error. This implies that transformers, in our setting, are nearly optimal metalearners for linear classification.
Our analysis also yields a characterization of the number of pretraining tasks required to generalize effectively at test time. Specifically, we derive an explicit relationship between the number of tasks B and the signal strength R, which controls the signal-to-noise ratio (SNR) during training. We find that it is sufficient to train the transformer on B = (A/SNR²) tasks and N = (11/ V 1 samples per task, to (almost) match the performance of an optimal metalearner, without any dependence on the ambient dimension d.
Finally, while Frei and Vardi [11] analyze similar setting without assuming a shared representation (i.e., they assumed & = d) and require strong assumption I ((vd), we show that it suffices to have R 22(1) for achieving in-context generalization, even when h: = d. We note that a single-task optimal learner needs only O(1) samples when the signal strength already equals (dd//a,, whereas it is information- information-thhorrttiaall theoreticalll impossible achieve small error when the signal strength is o(1), regardless of how many samples are available (see again Remark 2.1). Thus, learning a Gaussian mixture is challenging in the regime where the signal strength is between (2(1) and O(d!/4) and we cover the case where both R and R are in this regime.
1.2 Related Work
In context learning. Following the initial experiments of Garg et al. [14], which demonstrated empirically that transformers can perform in-context learning of various function classes, such as linear functions, two-layer neural networks, and decision trees, a number of works have explored what types of algorithms are implemented by trained transformers. Akyürek et al. [15], Bai et al. [16], Von Oswald et al. [17] provided expressivity results showing that transformers can implement wide range of in-context algorithms such as least squares, ridge regression, Lasso and gradient descent on two-layer neural networks. Wies et al. [18] provided a PAC framework for in-context learnability, and established finite sample complexity guarantees. Huang et al. [19] investigated the training dynamics of one-layer transformer with softmax attention trained by GD in regression setting. Focusing on linear regression, Wu et al. [10] established a statistical task complexity bound. Ahn et al. [20] and Mahankali al. [21] demonstrated that one-layer transformer minimizing the pre- pre-training loss effectively implements a single step of gradient descent. Zhang et al. [22] additionally developed guarantees for the convergence of (non-convex) gradient flow dynamies.
In the linear classification setting, Shen al. [12] showed that linear transformer trained via gradient descent is equivalent to the optimal logistic regressor, whenever the number of training tasks B, the number of samples per task N, and the test prompt length M, are all tend to infinity. Li et al. [23] showed that a single-layer linear attention model can learn the optimal binary classifier under the squared loss, with a focus on semi-supervised learning. The work most closely related to ours is Frei and Vardi [11], who studied the behavior of linear transformers via an analysis of the implicit regularization of gradient descent (similar to our approach). As we already mentioned, they analyzed a setting with a strong signal R = 2(vā, dd) while we allow R = 2(()). Moreover, we emphasize that none of the above papers addresses metalearning in the sense studied in work.
Metalearning. There is a large body of research on metalearning, often associated with related concepts or alternative names such as multitask learning, transfer learning, learning to learn, and few-shot learning (See Thrun and Pratt [24] for an early overview). Baxter [6] provided distribution- free sample complexity bounds for metalearning, A long line of works [25-30] has developed computationally efficient metalcarning algorithms, such as MAML AMosel-Agnostic Meta-Learning) under various settings, primarily for regression tasks. Several works consider metalcarning with shared low- oow-dimensional linear representation, which resembles our setting, albeit their metalearners not related to transformers [31-33,7 7, 34, 29]. To our knowledge, the only work that explores form of metalearning in transformers is Oko et al. [35], which studies a linear transformer architecture augmented with a nonlinear MLP layer. For target functions of the form fax) = a ((H,,))) where e _R lies in k-dimensional subspace, they show that the model can learn in-context with
3

---

prompt length that scales only with k. The key differences from our work are: First, their focus on regression rather than classification. Second, they employ somewhat artificial two-step optimization procedure first applying gradient descent on the MLP layer, and only afterwards performing empirical risk minimization (HRM) on the attention layer. contrast, we consider standard end-to-end gradient descent. Third, they require that the number of tasks and samples during training scales with the ambient dimension
Implicit Regularization in Transformers. Our theoretical analysis begins by examining the implicit regularization induced by gradient descent in transformer models. We refer readers the survey by Vardi [36] for broader overview. The convex linear transformer architecture study linear in the vectorized parameters, which, following Soudry et al. [37], implics that gradient descent converges in direction to the max-margin classifier. More general transformer architectures are non-convex, but many subclasses exhibit parameter homogeneity and thus converge (in direction) to Karush- Kuhn- (KKT) points of max- max-margin solutions [38, 39]. Another line of work investigates the implicit bias of gradient descent in softmax-based transformers [40-43], typically under stronger assumptions about the structure of the training data.
2 Preliminaries
Notations. We use bold-face letters to denote vectors and matrices, and let [n] be shorthand for Let La be the d identity matrix, and let O₂ (or just if is clear from the context) denote the zero vector in RRd. We let ||- denote the Euclidean norm. The Frobenius norm of matrix denoted || We use a V = max(a, b) and a A == min(a, ((a,)) b). We use standard big-Oh notation, with O (,)),, (()) hiding universal constants and ((-), Ö(-) hiding constants and factors that are polylogarithmic in the problem parameters.
n). (1,
2.1 Data Generation Setting
We consider the following metadistribution during training:
Assumption 2.1 (training timm task distribution). Fix some k S d and let P E Rdxk be a semi- Ig. Let B, N 2 and signal strength R. > 0. For any task T [B]. orthogonal matrix, i.e. P N+1 in lask т satisfy the following:
the input-label pairs (xy,5s Ur 3r
1. Let μ', be sampled ii.d from the distribution Unif gk-1), i.e., the uniform distribution on the sphere of radius R in l dimensions.
2. Set = -p',, to be the isometric embedding of pl, in Rdd under P.
3. Let Zτ,i 1hd N(O, Ia), and yr. iid independent.
Unif((+1)), where pr., Zri and yr are mutually
4. Conditioned on the lask parameter pry, set Z7, == Vr.i/47 + Zr,
Thus, the above assumption states that each pretraining task is class-conditional Gaussian mixture, with two opposite Gaussians, where the direction of the cluster means for each task is drawn randomly from k-dimensional subspace in Rd. Next, we introduce the test-time distribution, which may generalize the pretraining distribution by allowing in-context examples to have a different cluster mean size and sample size (denoted by R and M + 1) than those during training (R and N + 1).
Assumption 2.2 (test-time task distribution). Let M 2 be the number of in-context examples and M+1 in the test task R be the signal strength during test time. The input-label pairs (I,, yi y:) satisfy the following:
1. Let μ' be sampled from the distribution Unif (R gk-1), i.e. the uniform distribution on the sphere of radius Rt in k dimensions.
2. Set μ = -Pp', where P is from assumption 2.1, be the isometric embedding of μL in IRd.
3. Let zi N(O, and ys Unif((=13), where μ,, z; and ys are mutually independent.

---

4. Conditioned on the task parameter f4, set Dift + 2,.
Since the cluster means our training and test distributions have norms R and R, and the deviation from the cluster means has standard Gaussian distribution and hence norm of roughly vd, we call the ratios and A the signal-to-noise ratios (SNR for short).
va Vd
2.2 Attention Model & Tokenization
In our setting, each example is a task that consists of a sequence of (a, y) pairs. In order to encode such a task and provide it as an input to the transformer (i.e., to tokenize it), we use the following embedding matrix:
E
32
EN
ØN+1 R(d=1)x(N+1)
y1
1/2
YN
That is, each of the N + examples is given in separate column, and for the (N 1)-th example we do not encode the label and place 0 instead. The single-head transformer with softmax attention [1] is parametrized by query, key, and value matrices: W Rdx W. W E Rd xd., Then, softmax attention is defined by
(W E ( NOTWEE
(EE W Wº, W E W E softmax
P
where is a fixed normalization that may depend on N and de, but not learned. We focus on linear transformers, where the softmax is replaced with the identity function. Following many prior works (e.g., Von Oswald et al. [17]. Zhang et al. [22], Ahn et al. [20]), we consider parameterization where the key and query matrices W Wº are merged into W == (W T Our objective is to use the first N columns of Eq. 1 to predict #N+1- Similar to prior works, we consider a convex parameterization the linear transformer, obtained by fixing some the parameters to (see details in Appendix A), which results in the following prediction for the label of æN
N
i E: W) =
Σ y.
WØN+I-
(3)
N
Here, the trained parameters are W. We note that the model in Eq. has become common toy model for analyzing in-context learning both for both regression [10, 9] and classification [11, 12).
2.3 Gradient Descent & Implicit Bias
Given task т. define the embedding matrix E, using the labeled examples (wr, Fr. ) from Assumption 2. tokenized according to Eq. 1. We consider linear transformers (Eq. 3) trained to minimize the prediction loss on the final token ær, N+1, i. with ground-true label yr,N41- Formally, for training dataset {(Er, Nr,N NN1 +1) define the empirical loss:
L(W) E (y 100,,N+1-((E::) ((E;;
where i(E; W) is the prediction function from Eq. 3 and E is either the logistic loss f(z) log(1 + cxp(-2)) or the exponential loss {(2) exp( 2)) We train on this objective using gradient descent: W++1 W. aVL(W,) where a > 0 is fixed learning rate. Since j( E,; W) is linear in W, gradient descent has an implicit bias towards maximum margin solutions, as formalized the following theorem:
Theorem 2.3 (Soudry et al. [37]). Let MM denote the solution to the max-margin problem:
WMM := arg min I s.t. (yN V=. UyH.N+12,,N43 1T, = 11 B.
(4)
If the above problem is feasible and the learning rate sufficiently small, then W; converges in direction to WMM. that is Wa/liwsll - cWMM, as t - oo, for some constant > 0.
Thus, the max-margin solution WMM characterizes the asymptotic behavior of gradient descent for our model. In the remainder of this work, we analyze the ability of this max-margin solution to perform in-context learning and metalearning.
5

---

2.4 Ceermatiin-Thhooeeee Lower Bounds for Single-Task and Metalearners
We begin by characterizing the minimax test error for algorithms that only have access to the M labeled examples from the test task (I1, 31),.... (Za MMr) but not to the underlying subspace P. and need to predict the label of M-+-- This models the performance of an optimal single-task learner. Importantly, this lower bound holds at random from the space of all semi-orthogonal matrices in Rd We emphasize that our main generalization result (Theorem 3.1) applies in the worst-case setting-i.e. holds for any fixed P. However, when P is sampled uniformly at random, the induced task mean μe becomes uniformly distributed on the sphere of radius in R4, In this setting, we recover the following lower bound:
when P. defined in Assumption 2.1, is drawn uniformly
Remark 2.1 (Giraud and Verzelen [13], Appendix B). The minimax test error for Gaussian clas- sification with identical spherical covariance and opposite means, as defined in Assumption 2.2
with k= d, is least exp (-d 1-mmmin { R2, MR' }} } .for some absolute constants c, 0. In particular, when R (11). the number of samples required to achieve small constant error must satisfy M = (d/ R²). While for R = o(1 it is impossible to learn with small error.
We now consider the setting where the learner is granted full access to the underlying subspace P. in addition to the in-context labeled examples and the test point (æ11 vi....... (IM, Yar), #M+1- This models the performance of an optimal learner that knows the shared representation.
Remark 2.2. Consider the same Gaussian classification model, where the mean vectors lie in low-dimensional subspace P C Rd, as described in Assumption 2.2 with k S d. Then, the minimax
test error for any algorithm with access to is at least c. exp (-d' -eiiminn/ min [2² MR² ) some
absolute constants c, > O.
The proof of this remark follows by direct reduction to the bounds established by Giraud and Verzelen [13] and is included in the appendix for completeness.
3
Main Result
We make the following assumptions:
Assumption A. Let d be desired probability of failure. There exists sufficiently large universal constant C (independent in d, B, k. N and 8), such that the following conditions hold:
(Al) The signal strength R satisfies: C Cloog S R2 (1/8_
(A2) Dimension d should be sufficiently large: d 2 Clog (l/8).
(A3) Number of samples per task N should be sufficiently large: N 2 C(d/RP) v1
Assumption (A.) provides explicit bounds on the signal-to-non ratio (SNR). We emphasize that when R o(1), learning with small error is information-theoreee theoreticlly impossible (see Remark 2.1).
In contrast, when R 32(v)), the learning problem is solvable even by a trained transformer that observes only a single example per task, i.e., N M = 1 (see Frei and Vardi [1I]). Assumptions (A2) and (A3) technical conditions introduced to guarantee certain concentration inequalities involving sub-exponential random variables.
Recall that the transformer with parameters W makes predictions by embedding the set {(x,, y1 J.M U {(*M+1,0) 0)} into matrix E (as in Eq. 1), and then predicting VM+1 as sign(ū( E; W)). Our objective is to characterize the expected risk of the max-margin solution, namely, the probability that the transformer misclassifies the test example (æ UM+1) when parameterized by W
We now state our main result, which shows that transformers can serve as near-optimal metalearners:
Theorem 3.1. Let E 10,1) be arbitrary. There are absolute constants C > 1, > such that if Assumption A holds (w.r.. C). then with probability at least 1 - 5 over the draws of { (Zr.s i Vr.a) JT=b when sampling a new task {M, (xi, M 1,,,, the max-margin solution
M+118

---

from Eq. 4 satisfies
IP (m..))) 1 (ign(((( E; Wuns)) # yas+1)
M -
(IA)
BR
S 6exp
log (og5)
dk
Let us make a few observations on the above theorem:
Assume that the number of tasks B during training is sufficiently large, specifically +(1). To achieve an arbitrarily small BR
B ((dk//(R)) R2), so that the term (LAV
di
constant test error (e.g., at most 0.001), it suffices to have k = Õ(1), R = Ó(1) and M = Ok/ (Ok/k))) in-context examples. By Remark 2.1, an optimal algorithm that only has access to the samples from a new task needs (2(d/ R*) samples to achieve small constant eиror. Therefore, trained transformer can learn the small subspace during training and enjoy a better in-context sample complexity than such an optimal algorithm whenever k < d. In particular, if k S d², for some a < 1, we obtain a polynomial improvement in the sample complexity whenever R = o(d
In fact, a trained transformer is almost an optimal metalearner: given enough tasks during training. it suffices to have M = O (+)/ in-context examples to achieve constant error. This matches the lower bound for optimal learning algorithm that has full access to the sub- space P (See Remark 2.2). To achieve an error at most e (with probability at least - over the training data), the trained transformer will need O ((k/R*) log (1/e) log bog((B/)) in-context samples, while the lower bound for optimal learner that has access to the ground truth subspace is Z (k/ Rt*) log(1/e)) samples.
logo
A common assumption in the metalearning literature is that the training and test tasks are drawn from the same metadistribution. our setting, this corresponds to the case where R ft. Under this assumption, our analysis shows that a trained transformer can generalize as long as R = 2(11. & â(1), and the number of pertaining tasks and in-context samples M are sufficiently large. This improves upon the result of Frei and Vardi [11], which required the stronger condition of R (2dd/22). Note that when R is at least ((d4/4) a single-task optimal learner needs only O(1) samples, and for R = o(1) learning is impossible by Remark 2.1. Hence, our weaker condition on R allows us to cover the challenging regime where 2111= S S Odd/4).
Moreover, our analysis yields a tighter dependence on the number of training tasks B required for generalization in the high signal regime compared to Frei and Vardi [11]. As conerete example, suppose R R ((vd), M @(1), and d = k. Then our analysis shows that suffices to train on B Ô(1) tasks, whereas the result of Frei and Vardi [11] implies that in this case should be Ö(V)).
At first glance, it may seem that the number of tasks B and the number of samples per task during training N must scale with the ambient dimension d. This impression arises because the noise terms in the data (i.e., Zr,i from Assumption 2.1) have norm ||=+,م|| vd. Consequently, it is natural to ask how B and N relate to the signal-to-noise ratio (SNR) during training, which is defined as || / a| ~ R/ dd. More importantly, can we eliminate any dependence of B and N on d, which may be very large? Perhaps surprisingly, the answer is positive. By substituting R² d SNR? into Theorem 3.1 and Assumption (A3), we find that it suffices to train the transformer on B O(k/SNR²) tasks and N O(1 (1//SRR2) V1 samples per task, to achieve performance equivalent to optimal metalearner, without any dependence on d.
We emphasize that this assumption is without loss of generality. Indeed, if 2,, N(O, 21) for some # I, we can rescale A by factor of a and plug it into results, since the dynamics of GD remain the same.
7

---

Proof Sketch
Let W-
WMM be the max-margin solution (Eq. 4), i.e.
W = arg (in P Uy 2 1, VT = 1,....B).
(5)
where fú, N
Σ i= &y and (z,, yr) (I+,N+1, yr -1). For notational simplicity let us Ei= ,, and let us drop the M + subscript such that (IM+1.
denote ft := M
1
= (x. y).
Then the test error given by P(sign(j( E; W)) * y) P(ujay ya < 0). Using the identity ES y:u; = 1 + Eis Miz; and by properties of the Gausian, get: d l + M M-1/2z and ya μ + z', where 2, z' Lid N(O, Ia). Thus, using the transformer prediction rule (Eq. 3):
P(O E; W) + VM+1) = (
M (=++21/22) W( + <0)
=P PP u WH -u W=' M T M 1- z')
< P u WH Wz + M WH M- x W x// ) (6)
Recall that μ Pp', where P E dx is semi-orthogonal matrix and μ iid. we use concentration inequalities of quadratic forms to show that with high probability
~ Unif R-sh-R 1h- Then,
R2 T tr (PP)
R W 1|
R RRWI LE
)
1 WH
1
Wz So
1z 1 ([
k:
Vi
Then, our goal becomes: (1) Establish a lower bound on tr ( W). ensuring that μa W fu is large and positive. (ti) Establish upper bound on [W]] [[|.. Assuming the number of tasks B is large enough,
show that tr( (P') W P) Qkk R²) and [W[];= O(VT/RR) Substituting these bounds into Eq. 7, and then plugging the result into Eq. 6, yields the desired conclusion.
Lower bound on (r WP). Using the KKT conditions for the max-margin optimization problem and the fact that Vi(E;; W) x., we obtain that there exist AI....,AB AB 2 O such that
W
Σ Ur x.
(8)
We first show that tr(P 2 R2 D Then, by substituting the expression W (Eq. 8) into the margin constraints (Eq. 5), we obtain that for any T E [B]:
-=1 >r which means it suffices lower bound B X
B
B
1 1 AT
Σ
Ao B, = A, 12
Σ 9デT
Xy (o) (yg x x.yr).
1=1
Averaging and rearranging gives:
B
1 1
\y j
11
1|2+
Σ
(pr (ag) (q
9=1
|
||
B
1
Σ
(pt, (ay) (Byy x.Vr)
B
B
1=1
To derive lower bound on E >r. we aim first to upper bound || | 112 B, and second upper
bound the cross-term average B Er.r, (pr,, (Ug ) While the first is straightforward, the second requires a more delicate argument. Since [[pp) = E = Mr. the cross-term contains terms like (pr, p,y2 which can be bounded by ( /k), well zero-mean noise terms, whose
8

---

average is small when B is large. Together. these imply that both terms are at most O for large enough B, yielding Σ x 2 R/R* R$ and thus (PP) 2 k/R
Upper bound || WF... We derive an upper bound on [ by constructing a matrix U (up to scaling) that satisfies the constraints of the max-margin problem. Since W is the minimum Frobenius norm matrix that separates all training examples, this implies that ||W]| S |||| When the signal vector l lies in a low-dimensional subspace, natural candidate for U is the projection matrix PPT. We show that setting U == PP gives margin مa Uy, I, 6(R²). ). so that U/R satisfies the margin constraints. This yields the bound ||W| S || | | Vk R².
When the number of tasks B is small (i.e., B = o(dk/ R²), we use an alternative construction: We let U = 0 - Va for parameter 0 = (()/(R Then we can show that || r $ / BR vk/RR RR. This approach improves the dependence on the number of tasks B, compared to Frei and Vardi [11].
Remark 4.1. Interestingly, our analysis indicates that W MM exhibits properties similar to those of the projection matrix PPT, up to scaling factor of R2. Specifically, letting U := PPT, we observe that UUP
v and tr(P UPP) k. If the matrix W defined by the learning rule in indeed corresponds this projection, then the transformer effectively carries out the following procedure: it first projects the duta onto the ground-truth subspace; next, it performs maximum likelihood estimation (MLE) of the signal f by averaging the in-context examples (cf. Example 9.11 in [44]); and finally, it uses this estimate for prediction. Since this is the Bayes classifier under Gaussian prior (cf. Giraud and Verzelen [13, Appendix BJ), this procedure gives an optimal learner with access to the ground truth representation.
5 Experiments
We complement our theoretical results with an empirical study on metalearning with linear attention. We trained linear attention models (Eq. (3)) data generated as specified in Section 2.1 using GD with a fixed step size and the logistic loss function. In Figure 1, we compare the in-context sample complexity of linear attention against three baseline algorithms: (i) Support Vector Machines (see Section 15 in Shalev-Shwartz and Ben-David [45]); (ii) The maximum likelihood estimator (MLE); which estimate μs under a Gaussian prior by averaging the in-context examples (see Example 9.11 in [44]), and then uses this estimation for prediction; (iii) MLE with access to the ground-true matrix P. which first projects the data using P. and only then applies MLE. We see that the linear transformer outperforms both SVM and MLE, which lack access to P, and nearly match the performance of the MLE with projection. Additional experiments and details are provided in the appendix.
Accuracyy vs Context Size (R = R d= 500. 30)
Accuracy vs Context Size IA == A== 3, d= 500 k= 30)
1.
Test Accuracy
Test Accuracy
Tinnsorrassfrmme
LinearTransformer
MLE (full)
MLE (full)
MLE (proj 30)
MLE (proj 30)
SVM
SVM
108
168
Context Size (M)
Context Size (M)
Figure 1: Test accuracy versus the number of in-context examples M, where each plot represents different signal strength R R. We compare the performance of the trained linear transformer model against three baselines: full MIE, projected MLE (with access to the true subspace), and SVM. The transformer closely approaches the performance of the projected MLE and outperforms both the full MIFE and SVM, which lack access to the subspace. Accuracy improves as the signal strength
R = R increases. d = 500, k = 30, B = 20000.

---

Conclusion and Future Direction
We study the sample complexity of metalearning for the Gaussian mixture framework using pretrained linear transformer. By analyzing gradient descent, establish generalization bound that provably competes with any metalearner and outperforms any algorithm that only has access to the in-context examples. Importantly,
bounds do not depend on the ambient dimension, highlighting the transformer's ability to leverage low-dimensional task structure efficiently.
Our findings underscore the potential of transformers to extract shared representations across diverse but related tasks. This opens several future directions, and encourages extending the metalearning analysis to additional data distributions and transformer architectures, such as deep and multi-head softmax attention.
Moreover, our proof suggests that trained transformer is closely related to specific optimal learner with access the ground-truth representation, namely to a learner that first projects the data onto the ground-trutt subspace, and then performs maximum likelihood estimation (see Remark 4. 1).. However, it remains open whether the transformer can exactly mimic this procedure the number of training tasks approaches infinity.
Finally, although results indicate that transformers can implement effective metalearning using only a relatively small number of tasks and examples task during training, independent of the ambient dimension d, an interesting open question is precisely characterize the minimal sample requirements for successful metalearning. In particular, it remains unclear how many tasks and examples per task are sufficient for training an optimal metalearner, what the exact tradeoff is between the number of tasks and examples per task, and how these requirements may differ between transformer- baasdd architectures and more general metalearning frameworks.
Acknowledgments and Disclosure of Funding
GV is supported by the Israel Science Foundation (grant No. 2574/25), by a research grant from Mortimer Zuckerman (the Zuckerman STEM Leadership Program), and by research grants from the Center for New Scientists at the Weizmann Institute Science, and the Shimon and Golde Picker - Weizmann Annual Grant.
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[2] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, al. Palm: Scaling language models with pathways. arXiv preprint arXiv:2204.02311 2022.
[3] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
[4] Alec Radford, Jong Wook Kim, M Hallacy, et al. Learning transferable visual models from natural language supervision. In Proceedings of the International Conference Machine Learning (ICML), 2021.
[5] Tom B. Brown, Benjamin Mann, Nick Ryder, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33, 2020.
[6] Jonathan Baxter. A model of inductive bias learning. Journal of artificial intelligence research, 12:149-198, 2000.
[7] Nilesh Tripuraneni, Chi Jin, and Michael Jordan. Provable meta-learning of linear repre- sentations. In International Conference on Machine Learning, pages 10434-10443. PMLR, 2021.
10

---

[8] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnosto meta-learning for fast adapta- tion of deep networks. In International Conference on Machine Learning (ICML), 2017.
[9] Juno Kim, Tai Nakamaki, and Taiji Suzuki. Transformers are minimax optimal nonparametric in-context learners. Advances in Neural Information Processing Systems, 37: 17:106667-106713. 2024.
[10] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprini 1rXiv:2310.08391, 2023.
[11] Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. arXiv preprint arXiv:2410.01774, 2024.
[12] Wei Shen, Ruida Zhou, Jing Yang, and Cong Shen. On the training convergence transformers for in-context classification. arXiv preprint arXiv:2410.11778, 2024.
[13] Christophe Giraud and Nicolas Verzelen. Partial recovery bounds for clustering with the relaxed k-means. Mathematical Statistics and Learning. 1(3):317-374, 2019.
[14] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598. 2022.
[15] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211. 15666, 2022.
[16] 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:571 125 57211,, 2023.
[17] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[18] Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. Advances Neural Information Processing Systems, 36:36637-36633 2023.
[19] Yu Huang. Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 2023.
[20] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 56:45614-45650, 2023.
[21] Arvind Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self attentttenion. arXiv preprint wrXiv:2307.03572 2023.
[22] Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
[23] Yingeong Li, Xiangyu Chang, Muti Kara, Xiaofeng Liu, Amit Noy-Chowdhury, and Samet Oymak. When and how unlabeled data provably improve in-context learning. arXiv preprint arXiv:2506. /5329, 2025.
[24] Sebastian Thrun and Lorien Pratt Learning to learn: Introduction and overview. In Learning to learn, pages 3-17. Springer, 1998.
[25] Weihao Kong, Raghav Somani, Sham Kakade, and Sewoong Oh. Robust meta-learning for mixed linear regression with small batches. Advances neural information processing systems, 35:468346696 2020.
11

---

[26] Nikunj Saunshi, Yi Zhang, Mikhail Khodak, and Sanjeev Arora. A sample complexity separation between non-convex and convex meta-learning. In International Conference on Machine Learning, pages 8512-8521. PMLR, 2020.
[27] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. On the convergence theory of gradient- based -gdes-ag meta-learning algorithms. In International Conference on Artificial Intelligence and Statistics, pages 1082-1092. PMLR, 2020.
[28] Qi Chen, Changjian Shui, and Mario Marchand. Generalization bounds for meta-learning: An cetprmattion-thhoeete analysis. Advances Neural Information Processing Systems, 34: 25878-25890 2021.
[29] Liam Collins, Aryan Mokhtari, Sewoong Oh, and Sanjay Shakkottai. Maml and anil provably learn representations. In International Conference on Machine Learning, pages 4238 -338- PMLR, 2022.
[30] Konstantina Bairaktari, Guy Blanc, Li- LinYang Tan, Jonathan Ullman, and Lydia Zakynthinou. Multitask learning via shared features: Algorithms and hardness. In The Thirty Sixth Annual Conference Learning Theory, pages 747-772. PMLR, 2023.
[31] Maryam Aliakbarpour, Konstantina Bairaktari, Gavin Brown, Adam Smith, Nathan Srebro, and Jonathan Ullman. Metalearning with very few samples per task. In Proceedings Thirty Seventh Conference on Learning Theory, pages 46-93, 2024.
[32] Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes. The benefit of multitask representation learning. Journal of Machine Learning Research, 17(81)11-32, 2016.
[33] Andreas Maurer. Transfer bounds for linear feature learning. Machine learning, 75(3):327-350, 2009.
[34] Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei. Few-shot learning via learning the representation, provably. International Conference on Learning Representations, 2021.
[35] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Pretrained transformer efficiently learns low-dimensional target functions in-context. Advances in Neural Information Processing Systems, 37:77316-77365 2024.
[36] Gal Vardi. On the implicit bias in deep-learning algorithms. Communications of the ACM, 66 (6):86-93, 2023.
2023.
[37] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19 170):1-57, 2018. 2018.
[38] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint 1rXiv::906.0055 2019.
[39] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning, Ad- vances in Neural Information Processing Systems, 33:17176-17186 2020.
[40] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Advances in neural information processing systems, 36: 48314-48362, 2023.
[41] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transform- ers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023.
[42] Christos Thrampoulidis. Implicit bias of next-token prediction. arXiv e-prints, pages arXiv- 2402, 2024.
[43] Bhavya Vasudeva, Puncesh Deora, and Christos Thrampoulidis. Implicit bias and fast conver- gence rates for self-attention. arXiv preprint arXiv:2402.057338 2024.
[44] Larry Wasserman. All of statistics: a concise course statistical inference. Springer Science Business Media, 2013.
12

---

[45] Shai Shalev-Shwarde and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[46] Roman Vershynin. Hiihh-dimensionll probability: An introduction with applications data science, volume 47. Cambridge university press, 2018
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:[Yes
Justification:
The abstract and introduction describe the concrete contributions of this paper, as indicated throughout the rest of the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: [Yes]
Guidelines: Assumptions are compared with those in prior work, and the limitations of the results are discussed throughout the paper, including the final section.
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: The assumptions and proof sketch are stated explicitly in the main text, while the formal proofs are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: In the experimental sections (both in the appendix and the main text), we pro- vide detailed explanations of our experiments, with the exact parameters for each experiment given in the plot captions.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
15

---

reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification:
We plan to release the code in the future
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (attps://nips.. co/
onoolicy))yy) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips. cc/ wtc/public//ee httllii//gides///e/eee /Coeee/modSSubiic iundbbiissse ieiinonooiccy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We state this explicitly in the experimental sections. We use standard gradient descent, so each batch includes all examples.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
We state this explicitly in the experimental sections.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
All the experiments used a CPU.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https : htt////euurpppeeeeeeeeee
wwwpurips.ccppubiiie wwtpcccccppcccceeeeeeeeeccccccce htt.ai/e/eehhiiuieeeeeeeeeeee ines?
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and confirm that our research complies with its guidelines.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is theoretical work that explains metalearning in transformers.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer[[NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwode.hcoe...c. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
18

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer:[NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer:[NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.M policy (https: ftttps:/nneuree http//rrppss.cccofeeeeeeeeeeeeee /LLM) for what should or should not be described.
20

---

Contents
1
Introduction
1.1
-Ou contributation
1.2
Related Work
2
Preliminaries
2.1
Data Generation Setting
2.2
Attention Model & Tokenization
2.3
Gradient Descent & Implicit Bias
2.4
Information-Theoreti Lower Bounds for Single- Task and Metalcarners
3
Main Result
Proof Sketch
Experiments
Conclusion and Future Direction
10
Convex Parameterization the Linear Transformer
21
Proof of Remark 2.2
22
Proofs for Section
23
C.1 Notations
23
C.2 Properties of the training set
23
C.3 Analysis of WMM
28
C.4 Concentration inequalities of quadratic forms
31
C,5 Proof of Thm. 3.1
32
Additional Lemmas
34
Further Experiments & Additional Details
36
A Convex Parameterization of the Linear Transformer
In this section, provide detailed explanation of why our prediction model (Eq. represents a convex parameterization of a linear transformer, formulation that has been explored in several prior works. Following Eq. The linear transformer with key-query matrix KK is defined by
J(E:0) = E VVE E. (E'WKO E
(
PN.d.
We use the first columns of that model formulate predictions for ØN-1. whereby the bottom-right corner of the output matrix of fEE 0) serves this prediction.
21

---

Writing WA
W
wiz W22
for A (V, KQ). for the linear transformer architecture, this
results in the prediction
WKQ (w2)) KQ
g(E:0) ((
W22)
EET
))T NN+1-
N
Due to the product of matrices appearing above, the resulting objective function is non-convex, which
makes the analysis of its training dynamics complex. We instead consider a convex parameterization KQ of the linear transformer [10, 9, 11, 12], which results from taking w21 = w = 0 and setting wh2 = 1. This leads to the following prediction for the label of æ--1-
N
(E W) = N Σ yi
WIN+1,
as defined Eq. 3.
B Proof of Remark 2.2
Let c, 0 be the absolute constants from Remark 2.1. Fix sufficiently large integers V' d' such that k' S d', let P' gd''k'' be a semi-orthogonal matrix, and let R 0 be the signal strength. Consider the dataset S ((x1,11)..... vi)...., ... (æM, 9M)} S Rd x {=1} and the test sample (#M+1. MM++1) E Rdd {+1) sampled according to Assumption 2.2 (with d = d', k W, P P iid Unif({#1)), have iid. N(O, (g). yi
itd. and R R). That is, for μ Unif(( - 1)), for any i E [M + 1]:
Viu
(9)
Let AP any algorithm with access to P Our goal is to show that
Pr [AP (5)(2 = yas+il 2 c. cxp (-d cd' min { (MM
R2, MR k
})
(10)
Assume by contradiction that this is not the case, i.e,. there exists algorithm AP whose error smaller than c. Then we can construct an algorithm B for Gaussian classification with identi- cal spherical covariance and opposite means, defined in Assumption 2.2 (with d k K'), that achieves error smaller than in contradiction to Remark 2.1. Indeed, given a data set (#'s b(r) C _R x { and a test sample (4'm4393s44)) E RK x {+1}.
s' = {<<x'.>)) sampled according to 2.2 (with h = d = N, R R). That is, for μ Lid. Unif(R gk ()) Ltd. N(O, 1.).% lid Unif((+1)), we have for any i E [M + 1}:
lid. x/ ~ d'p + z'.
The algorithm chooses some semi-orthogonal matrix E dd'' and sampled independent N(O, La P. ). Then B simulates AP on the training set {(P'r
vectors 81, ..... 8M+1 lid $1,3))...., (P'x +8M (.MMMMM)) VM)} and the test sample P'MM41+8M++ (144))))) The key observation is that for any E [ + 1]:
P'ay + 8, u + Pz| + xi
where x; is from Eq. 9. Indeed, since Pz N(O, PPT), we have Pz' + N(O +IN PPT) ) N(O, Iv). Therefore, we can conclude that
yM+1l where is defined in 10. Contradiction.
Pr [B(5)(( (1 M
Pr S
(8)( yM5+1] <6
22

---

C Proofs for Section
C.1
Notations
First, we introduce useful notation for the remainder of the proof.
Assumption C.1. For some parameter cB have CB dk/ R².
The notation eg from Assumption C.I specifies how the number of training tasks B scales function of the SNR. We will later show that the generalization error on in in-context examples depends partly the quantity 1 A /CB, where larger value implies better generalization. Importantly, allow Cg to non-constant; for instance, settings where B oa(d) are permitted even for constant R.
Table I: Notation used throughout the paper.
Symbol
Description
ambient dimension
Shared subspace dimension
P
Shared low-dimensional subspace & semi-orthogonal matrix P E Bdxk
u'
Cluster mean, μ4' E I
Isometric embedding of the cluster mean, μrr Pp' E IRd
Features, E RRd ,t,i Vr,iH žr,i
Labels, y E {+1}
The noise vector, # E R
Probability of failure
R
Norm of cluster means during pre-training
R
Norm of cluster means test time Number of pre-training tasks
B
CB
Quantity such that B = CB RR
N
Number of samples per pre-training task Number of samples per test-time task
E
Data tokanization
E
[1
32
ON
EN+1
g(++1)x(++))
31
y2
UN
0
Mean predictor: 1 Σ
30
1=1 y:
3 ((E()): W)
Transformer output: 1
M Em y:
Wx
A Ww
Logistic loss or exponential loss
C.2 Properties of the training set
Lemma C.1. Let å E 10 be arbitrary. There is an absolute constant co > such that with = for all T. [B] such
probability at least 1 - 8 over the draws of fpt, (xx, Mr), (Zr NN 1
23

---

that q # T the following hold:
CoRlog(B// B/o)
2d V co log B/6)
||121 R21 v
+
VN
N
log B/o)
|| 112- -dl d) - R
Co
+ R (Bg(B/8).
vd
R²
Vd
KAy pb,)1 S
(B()),))
Vh
VN
N
RC2
1a0, mz)) 00
+
R
+ vil
log( B/6),
d
(A- (-H+T) - R² S (1+
R +
log (B/8)
k
(PT ür. PTP. y+Σ+T) - R² I cD
1
R
log( (B/5)
Proof. By definition of fú, and properties of the Gaussian distribution, there is 2 N(O, Ia) such
that
1 Σ Mr, Vi Zr,i) pr Dr po +
1
μ.,
VN
Thus for r # there are z/,, 2/1 hid N(O, Id) such that
(pg. (()) (HT + N 1/2 2/- M4 + N NN11/2
= ( Mgg + N
1/2 (z+ Mgg + NN 1/2 (2/g (4) + +N- (2t. zz))
(11)
We first derive an upper bound for this quantity when * т. Remember that μ4, = PH' for some semi-orthogonal matrix P and n¹, (S-
Key, payl Analysis. since pq pl, are independent and uab-Gaussian random vectors with sub-Gaussian norm most cR/ v/k (Remark D.1) for some absolute constant by Vershynin [46, Lemma 6.2.3] with A Ig, we have for some c' it holds that for iid. N(O, Ig).
any R. g.9
E p. )] E3exp(3H1 x4)1 Elexp(é 33 [gg g')l-
By Vershynin [46, Lemma 6.2.2], for some provided d/ /k: c1, it holds that Elexp(d 3g g')] s exp(c, 62 R 2 exp(e R*
Since pay bh, are mean-zero, by Vershynin [46, Proposition 2.7.1] this implies the quantity μ4, is sub-exponential with [|u ll S 12 VT for some absolute constant 62 > We therefore have by definition of sub- sub-exponential [46, Proposition 2.7.1, first item] and union bound, that for some absolute constant C3 > 0, w.p. at least 1 - 5, for all 7 # q.
KMTT S R2 1/2 (B/6).
(12)
p. z, Analysis. We have that u1 = TTz Since μt, has nub-Gaussian norm at most cR/ VT (Remark D.1) and PT 2 has sub-Gaussian norm at most e (Remark D.2), by using again Lemmas 6.2.2 and 6.2.3 from [46], have for any B E R if iid. N(O, Ix), then
Elexp(3H, x1] < Elexp(d Rk 1/2 -1/2 3/ T "')).
24

---

and thus provided (/2111/28) S we have
Elexp(c' R B/ g S R
= ex
In particular, the quantity μ, x, is sub-exponential with sub sub-exuuooeee eee-eeepoonnaall norm 1| 24 lly, _ 02 R2R, and so for some absolute constant C3 we have with probability at least 1 6, for all q, 7 E [B] with o T,
(gg x++|| S c3 (log((/))).
(13)
( 2/} Analysis. For (2 24) with + q we directly use the MGF of Gaussian chaos [46, Lemma 6.2.2]: (2 z/²) g Iag' for i..d g. g.g' N(O, Ia) so that for 3 s c/ Lallz.
Blexp(B(Zz 2(,i- ك exp(ca [ali)) exp(ca3"d).
In particular, (| x4) 11 S cs d so that sub-exponential concentration implies that with probability at least for any 4, [B] with g # T,
k x S co
(14)
Putting Eq. 12, Eq. 13, and Eq. 14 into Eq. 11 we get for q # T.
R²
R
d
Kin i - C7
log 1B8(
(15)
N
As for | 11? from Eq. 11 we have
| 1|2 = || + L2N-1/2(6) (2+ -
(16)
From here, the same argument used to bound Eq. 13 holds since that bound only relied upon the fact that 24 and z,, are independent, while fl, and z', are independent as well. In particular, with probability at least 1 - ó, for all E [B].
(1,,=+)1= S log())))))
(17)
Each coordinate of z has sub-exponential norm most some constant and ||] d. Therefore by by Bernstein's inequality [46, Thm. 2.8.1], we have for some constant co > 0, with probability least 1 5, for any [B].
(/o) d
| 2| 4 S eov
(18)
Putting Eq. 18 and Eq. 17 into Eq. 16 and using that 1 28,
||,||| = R², we get with probability at least
Co (Rogg)(2//)) 2d eg lBg(
-2 R2| -R <
VN +
N
As for | 12, by definition,
x 112 = || + 2( 2イ) + 1 R22 2( 2+) + 1|21112
Since 2, N(0, Ia) has the same distribution as z",, the same analysis used prove Eq. 17 and Eq. 18 yields that with probability at least 1- - 28, for all 7 E [B].
(1,, z,)l log BB0),
co log B/6) Vd
|| 2 d
Substituting these into the preceding display we have that
c3 log || ||2 d < R22 log(B/8)
Vd
25

---

Thus provided d is sufficiently large, then we also have |2=112 = d.
Next we bound Ke+, Rg)l: There are z/, 2/ sid N(O, Ia) such that
(U- Fg
(+e + 2,,, My + z/y)-
is clear that the same exact analysis we used to analyze Eq. 11 leads to the claim that with probability at least 1 - 6, for all #
K 1 S co
R (Ti
+ R
vd
log(B/6).
(19)
Finally, we consider y, fi x,. Just as in the previous analyses, there are 2r, br, N(O, La) such that
(pr: Øトト (++, + N 1/2 Zr Mr x²t)
= |||1 + N N/1/2 (zT, 147) + (2/, 147) + 1/3 (21 12),
Again using an analysis similar to that used for Eq. 11 yields that with probability at least 8, for all E [B].
(1++
1 VN
vd
(pt, vr.r)) - 1221 S c10
R
log(2B/6).
(20)
Moreover, note that (PT ür, PT y, (pl + 2,/ /n ut, + x for 2,, 2 itd. N NO, Ix) (see
Remark D.2). Then we can use the same argument as Eq. 20 with k = d and P = I3 to conclude that with probability least
8 we have:
Vh
(PT
p., PT y- "..,P
R2 S cia
R
log(28/6).
(21)
VN
VN
Taking union bound over each of the events shows that all of the desired claims of Lemma C.I hold with probability least 1 105. also easy to verify that the Lemma holds with probability at least 1 8, for a different choice of constant co-
Lemma C.2. Let 8 E 10 be arbitrary. There is absolute conslant co > such that with probability at least 8 over the draws of (1r (ær, 17), (Zr,i Vr,a) 10)))= then we have
([
([+
Vd
B
tr(P' PM 2
R2
R+
log (B/))
X
Proof. By Remark D.2
WMM P = Σ Dg PT P. P
Σ Ag(p'y + /g/ 2g/VN)(( (NL + 2
ltd. where 2q, 2 ~
N(O, Ix). Next, we lower bound tr( P PMMP) as a function of Σg Agi
B tr(P WMMP) Σ Xo tr (ni + +z///NN (nn) + =
(22)
9=1
Observe that
tr ((1' + / N/N (Ng + +2(T))
= ((K + +z²
(w/g + z /.///N)
R ([++
Vd
R +
log (B/8).
VN
where the last inequality follows from Lemma C.1 (last item), and since for h d and P = Is, we have that (p.. ).gg.) & (H,T + 2 (14, + /g/VN).
Substituting the displayed Eq. into Eq. 22 yields the desired result.
26

---

The events of Lemmas C.I and C.2 hold with probability at least - 6, independently of Assump- tion Combining Lemmas C.I, C.2 and Assumption A we can conclude that:
Lemma C.3. Suppose that Assumption A holds for sufficiently large C. Then there exists some
constant co such that with probability at least 56:
A 1| R² < R2//4.
(23)
||,0,11 - d| S d/4,
(24)
R2
((gg fT)) S c0
log (B/8),
(25)
vhl
R²
(00 x20) - C0
V
Vd
log(B/6g
(26)
Vk
(AT V- R²1
Σ
R2/4
(27)
R²
B
> Σ
X
(28)
R
(P PT. Ptr,PTT vr
R22
(29)
4
Moreover, recall that B = CB dk/ R² (Assumption C.I). Then for any T E [ we have
p 1
2R
(30)
B
cBk
R¹ vi) h
log? log) s (1 v1)
Ra
Σ
(pr, (g) (yg x+U+)
Sc Rv/B
log 2
(B/6).
B
cB
k
qE(B]uFT
(31)
for some constant e 0. We also conclude that if B S dk, i.e. ycB < $1/R, then have that
B
Rd
Σ
(Pr, (Mg x+V+)
Sc
(B/8)
R /
k
R²² BB
Bd
(B/6)
(32)
VT
Proof. The first part (Eqs. 23-29) occurs with probability at least 1 - 25, and follows directly by substituting Assumption into Lemmas C.1 and C.2, and by applying the union bound. Regarding the last part, observe that
[ x B
1
2Rd
Rd
2
(33)
B
Cak'
where the first inequality holds by the first part of the lemma, i.e. the upper bounds on ر 12, || ??? The last equality holds for by definition of B. This proves Eq. 30. Moreover, observe that
1
Σ
(Dr, (g) (Mg xr)
<
(pr: Pag 14)
B
B
qE ():
qE (
1
(pt, (ag) (Hag 2nd
+
Σ
(Pr, Pay) (Zag pe)
B
B
qe Bl
qE gE(BB
1
+
(jtrs Pay) (Zgg Zz)
(34)
B
(E(B)
NNote, however, that the quantities appearing on the right-han hand sides of each inequality in the lemma are only small when these assumptions hold; this is the reason for these assumptions.
27

---

Next, we analyze each of the above terms seperatly, Regarding the first term,
1
Σ (itr Pa) (pas ue))
S
R² max ((rr;,))))) (1)1 uc)1 S co log (B/O). r.g
(35)
B
k
where the last inequality holds by the first part (third item) of this Lemma and Eq. 12. Regarding the second term of Eq. 34, let's fix fr: Py and by and observe that (jtr (Ra) MMg 21),.., (ir, P zB} we currently treat only 21....,. 2B as random variables, while the other terms are fixed. In other words, the following analysis holds for any realization of fers Mr and Mag conditioning on the events of the first part (Eqs. 23-29). By General Hoeffding's s inequality (Thm. 2.6.3 from Vershynin [46]),
independent random variables. We emphasize that
we have
ct
Σ
(prr, (ba) (Ma: 2x) 2 S 2xx
B
K2
Σ Eappe r (er Hg)
qE(B)
where e is some constant and K == maxgua77 18 z7) 11 By properties of the Gaussian distri- bution we have that (pg: zr) ~ N(O, [| 2) which means that [1 z,) I1lp = c /|1اا = cR, for some constant Moreover, by the first part of this lemma (item 3) we have that Eq (p+ (M)
co B log" (B/8)/k. By choosing t = v R log( B/6) log2 (2/0)//BB,, we obtain that with proba- bility at least 8, for any [B],
1 B
Σ (E(B)
(irT (ag (Wags 2+) 5
R log" (B/8) log (2/6)
1
R log (B/5),
cBk
c VCBd
k
(36)
where the last inequality holds by definition of B. The same bound also holds for the third term Eq. 34. Regarding the last term of Eq. 34, use again the General Hoeffding's inequality to obtain,
(
(- K2 {= grote (pr, pmg
cf B
Σ gE(B
(pr, (gg) (Zgi z7)
21 S 2exp
(-K2)
B
where now K = maxqqarr 1k(2 Zx) 11,pz. Now we have that (Z 2r) N(O, || 112), which means
that (==g 2x) 1l,s || 2cvd. d. By choosing t = Vdle* log" (B/8) log(2))) obtain that with probability least 8.
1 B QE(MT
R² dR log" (B/))))o((//) < cR /cn k
(it, (g) Zn) S
log (B/6),
cBk
(37)
where the last inequality holds by definition of B. Substituting Eqs. 35,36 and 37 into Eq. 34, and observe that 1 V 1/cB 2 1/ / V 1/Vdcg, for any R, d = 2211 (Assumptions (A2) and (A3)), Eq. 31 follows. We note that by union bound, all the events of this lemma occur with probability at least 1 - 56.
Our results will require this event to hold, so we introduce the following to allow us to refer to it in later lemmas:
Definition C.1. Let us say that a good run occurs if the events of Lemma C.3 hold. By that lemma, this happens with probability at least 56 over the draws of the training sets.
C.3 Analysis of WMM
Next, we introduce the following lemma regarding the max-margin solution W = WMM, as defined in Eq. 4:
28

---

Lemma C.4. On good run for C > 1 sufficiently large under Assumption A, the max-margin
solution W of Problem
W Σ 2
is such that the ' 2 0 satisfy the following:
B
(cB A A1
k
log (B/8)
R
where co > 1 is some constant. Further, have the inequalities
Vi R
IP'WPlle P | S (1AVEB)
and
(CB A 1) k log² (B/6) R
((P)WP) 2 c.
for some constant c. We recall that = cB dk/ R (Assumption C.1).
Proof. In this part, and cu represent some constants that change from line to line.
Tr(PWPP) lower bound. By the feasibility conditions of the max-margin problem, we have for any [B).
1 S yr f W
a
Σ No yr
B
Ng (Vr (r yorta)
9=1
||
112 112+
E
(p) (Mn x.
The above equation also holds if we average over 7 i.e.
B
S
X
1-110-1111
dy (gg) (Wg at
B
q: qata
(prr (V (r
B
B
B
|
|
1|
B
+
Ag
Σ
(jers Pag) (Wl 2+
B
B
7m
q
T:T#2
2R
B
R²a log
B (B/5)
Σ
X +
C
v1)
cgk
CB
k
where the last inequality holds by the second part of Lemma C.3. We can conclude that
B
(cp 1)
Σ X 2 C
log" (B/O) R
which proves the first part of the lemma. The last part of the Lemma regarding ir (P(P following directly from the displayed Equation and Lemma C.3 (last item).
29

---

WilF upper bound. We first derive an upper bound on | by showing the existence of some matrix U (up to some scaling) that satisfies the constraints of the max-margin problem. Frei and Vardi [11] used U la so natural candidate in our case (where the signal la is sampled from some low-dimensional subspace) the projection matrix PPT. Let U PP to be the projection matrix into the columns of P. We first derive an upper bound on W 1 by showing that the matrix (up to some scaling) satisfies the constraints of the max-margin problem (Problem 4). Indeed, by Lemma C.3 (last item), for any [B],
p. PP V- ÚQR²).
Thus the matrix U/ R separates the training data with a margin of at least for every sample. Since W is the minimum Frobenius norm matrix which separates all of the training data with margin 1,thi implies
Vk |i U/ F R²
(38)
By Remark D.3 the same bound also holds for | WPP
[ upper bound, when number of tasks B < dk The above approach does not yield tight bound when the number of tasks is small. Instead, we use different approach and define
U 0
yg aa
(39)
The idea is choose specific value for 0 and show that for this specific value, U satisfies the constraints of the max-margin problem, then by definition of W (Thm. 2.3) we can upper bound [ p by ||U|||>- We start with a thechnical calculation that will be usefull later in the proof. Since good run holds, by Lemma C.3 (last part), if B s dk, then for any T E
[B] we have,
RV
Σ (pt, Pag) (Feg xT) S Co quattr
VB.
llog2 (B/8)
(40)
Vi
Next, we move to show that U (Eq. 39) satisfies the constraints of the max-margin problem. Indeed, for any E [B].
V
0- 11/12-111
|0
+
Σ
(fr, Pay (Wn Vr
9:
2 0. |1 ||,,11
Σ
(pr. (M y 2.
R&d
R2 Vd
2 0
0. cov
/
loog (B/8))
Vk
(8)
(41)
Inequality (i) uses Lemma C.3 and Eq. 40. Inequality (ii) holds by choosing = (R22d) and for small enough i.e. S kd kd/(C ood/"logg (B/5)) Thus, the matrix U (with 3/(R2) separates the training data with margin at least 1 for every sample. By Lemma
30

---

D.2 we have that
|U|| - 0²
Σ 1 1 || + 0².
Σ
(isg
(1 wewes)
976
B
02.
Σ
|1/
1
||
1| 0
Σ
Σ
(ig (i yeare)
9=1
q=1
E:ETy
WV
3Rdd
R
B
B CoVB.
loog (B/5)
R
2
R
vk
(3R'd).
(42)
R
Inequality (1) uses Lemma C.3, Eq. 40 and 3/(Rdd). Inequality (ii) holds for small
enough B i.c. whenever
2cgdl
d d 2c, V (B/8)
B
log (B/5)
Since W is the minimum Frobenius norm matrix which separates all of the training data with margin 1, and together with Eq. 42 we have that
6VB VaR
Vi
W - |0,, S
6VCB
(43)
R
Combine with Eq. 38, this proves the upper bound of | ||||. |.. By Remark D.3 the same bound also holds for || WPIlF.
C.4 Concentration inequalities of quadratic forms
In this section, we derive concentration inequalities for quadratic forms, which can be viewed as variants of the Hanson Hanson-Wriiht Wright incquality, tailored to the following terms:
μ W μt, where μ Py' for some matrix P E Rdxk and μı' Unif( R.
f1 W z, where N(O, Ia).
z Wz, where z' N(O, Ia).
Wp:
Lemma C.5 (Hanson- Wright for uniform on the sphere of a subspace). Let R > 0 and W E ddd be a matrix. If μ Pp' for some matrix P E TRdxk and μ³ ~ Unif R Sk-1), then for any 2 0,
R² k
1
W μ
tr ( WP) 2t =P )
W Elu
W pi] 2t
12
tk PT WP
S 2exp
-emin
R+
W WPI
12
tk pelppTWPlle
$ 2xp
-cmin
P: W PId
where c is some constant. We also can conclude that
R²
ctk PT WP
P
1
W 1
tr
WP)
2 < 2exp
31

---

Proof. Write Q := P and observe that μ W p² Qp' for p ~ Unif( S Then the first inequality follows directly from Lemma C.2 Frei and Vardi [11]. The second inequality holds since || W I [| for any matrix W. Regarding the last part, write = td/ R ||P WPll,, if 2 1, we have that exp(- min(a, a²)) < exp(-a). If a 1, the above bound still holds for small enough e and since P(.) s 1 is a trivial inequality.
J Wz:
Lemma C.6. Let μ.. Pp' for some semi-orthogonal matrix P E Raxk and pa' ~ Unif R -h--).. Moreover, let g ~ NO, Ia) (independent of u') and W E IRdxd be a matrix. There is an absolute constant c 0 such that for any O,
dvv
P ( Wg 2t) 2exp
Proof. By Remark D.I, has sub-Gaussian norm at most cRt V for some constant 0. From this point, the proof follows similarly to the proof of Lemma C.4 from Frei and Vardi [11], and provide it here for completeness. Observe that ||g v22 for some absolute constant > 0. By Vershynin [46, Lemma 6.2.3], this implies that for independent g 92 ~ N(O, Ia) and any B E R,
Vk Wg
E exp
B
s Eexp (c W Wg)).
R
Then using the noment-ggnneating function of Gaussian chaos [46, Lemma 6.2.2], for B satisfying |8] S W have
Vk Wg
Eexp
BBv p R
B
s Exxp (c1 W Wg2)
expl(ea (
That is, the random variable R- 1 Ky
Wg is mean-zero and has sub-exponential norm at most max ca)]l There therefore a constant c4 such that any 2 O,
Ru Vk
P( Vkp W 2 u) PP
[ Wg 2
S 2 2exp(-ou/||W]]p
Setting u tvk/ R we get
ctv/
P P(H Wg 2 2t) < 2xxp
R|
<W22: 2:
Lemma C.7 (Lemma C5 from Frei and Vardi [1I]). Let 5, (id N(O, Ia) and W E IRdxd matrix. There is constant 0 such that for all t 0.
(- ct
P(IST WC' 2 1) < 2exp
(-IWITF)
C.5 Proof of Thm. 3.1
For notational simplicity let denote W WMM. f subscript that denote (æM+1, MM+11 C.I), l), and let's denote as the quantity
M ΣM ya and let us drop the M + 1 (x. v). We recall that B cB dk/ R2 (Assumption
(cB A 1)
p ::=c log (B/o) E (0. 1).
32

---

Then the test entor is given by the probability of the event,
(sign()(E; W)) + ym++1 = {ir Wy yz S 0}.
By standard properties of the Gaussian, we have for 2, z' lid N(O, Ia),
A 4 + M 1-1222
d
jæ n + z'.
d
We thus have
P(O(E; W) + MM+1)
= =P u+ =P(u+
M 1/2)) W(H +')+)++++- <0)
P(
W -H W== - 1- W 1- TT Wz))
W 1 T W=' + M 1 WH + 1- 2 W=!l)
SP
pf²
N
V 1/ T W11 > pp²²
<P
1 T
W
U
| W='
>
U
M 1/ '='
2
2R
SR?
8R2
8R²
pR²
pli
> pff
< P
μ WH <
+P
1
Wz'
>
+ P
M- x T W μ
2R
RR?
8RR
pRp
+P
M 1-1 N T x
(44)
SR
where the last two inequalities hold since given events IB; B2 we have that I(B < ( and by the union bound. Next, we assume that good run occurs, which indeed happened with probability at least 1 56 over the draws of the training (see Definition C.I). We now proceed by bounding each of the remaining terms in the inequality above:
IP (P(1 W
e 21
For the first term can use Lemma C.5 with 212
to obtain
R2
-1)
R WP)
pli²
P
11 T W μ
Ir
=P P p
W p
S
(r
k
k
2R
ck
pf2
< 2cxp
(45)
R2 WPlle
2R
In Lemma C.4, we show that tr(P WP) > equation, we obtain
pk Rz.
By substituting that into the displayed
RR²
ck
R
TP
W 1 S
< 2exp
2R
RR P W P
F
2R
(6)
ck
R²
pft
2 2exp
R
(LA
VK
2R
PVT.
= 2exp
IAVCB
Inequality (i) uses the upper bound
| WP S VCB) 1
from Lemma. C.4.
33

---

IP F Wz
e SR
and I (IM~ 1/22 WH 2
pft 8R2
The second and the third term
in Eq. 44 can be bounded by Lemma C.6. Indeed,
pR²
cVT.
pR
P 11
W='1
2exp
8
R I|
8Ra
(1) 2 exp
cvk
R²
RR
R
(1 (LAVCB)VE
SR2
cpR
exp
(46)
8(1A VCB
Inequality (1) uses the upper bound W S (1 /CB)
vi R
from Lemma. C.4.
(IM 1/2 1122 Wz' 2
OA²
For the final term in Eq. 44 we use Lemma C.7,
svkdd
P |MM-1/2 W=
= P
Wz'
8R2
SR²
PM 122/2
C
S 2exp
|
8R2
(
cRZ
PM p/222 R22
S 2exp
(1 VCBKK
8R2
cpM-/22
2xx
(47)
(1 VCB)Vk
Again Inequality (1) uses the upper bound of ||
P W P|l from Lemma. C.4.
Putting together Eqs. 45, 46 and 47 we get
PVk IA CBB
cpR 8(L
c2M
P(O(E:W) VM+1) s 2 exp
2 exp
+ eexp
CC)
By plugging p = c. Tog (B/8)
(CBAl)
we can upper bound the displayed equation by:
(VRARA
1/2
6 exp
(LAVEB)-
log (B/8)
Vk
Recall that B = dk/ R² (Assumption C.1), which means that cB = BR² /dk. We can conclude that
(VEARA M M
Vk )
BR
P(J(EW)) # yM+1) 6exp
1
(oog
dk
D Additional Lemmas
Remark D.1. Let P E dd be an semi-orthogonal matrix i.e. PTP= Ik. Let p' Unif Unif( (R-S*-1) and set μ Pn'. Then μ and p' are uub-Gaussian random vectors with sub- Gaussian norm most cR/ W for some absolute constant > 0 i.e. || A || cl/ Vk.
Proof. μí' is sub- Gaussian random vector with nab-Gausiian norm at most RR/Vk [46, Theorem 3.4.6] 6] for some absolute constant e 0. particular, we have that [46, Definition 3.4.1]
||
= sup [| 11 = sup
P 2>
=
me3d
[1 x) 11 = || 1
< sup
-
where the second equality holds since
IP_ x
< | for any x E IRd
34

---

Remark D.2. Let P Baik be an semi-orthogone matrix (i.e. PTP= = I&) and let 2 ~ N Ia). then 2~ (0,
Proof.
E[P
E Coo(P) 2)) = E 2)
E
P
[P= P
E[zz P P
Lemma D.1 (Hoeffding's theorem). Let 21,22...., be independent random variables such that E[ and di E [-a, a) almost surely. Consider the sum of these random variable S, 11 +... Tn- Then Hoeflding's theorem states that
Pr[[SA] 2 n 0.75 a] S 2exp( 2.5 a² 32/4ann) = 2 2exp(-n' /2),
Moreover, if 2 42og(2/)) then
Pr[[Sp] 2 n.76% S
Lemma D.2. Let W E yofix T. Then
B
W
Σ
A 11 |,,
Eip (ee) (g years)
g=1
Proof.
B
W = Σ
D4 (0 35.1
9=1
2
( Evo Vol ( Dessee Das
W
F
Σ By (ig )
=E
Σ
+ Σ
By (Pe )is (Perze Dis
>
| T |
Σ
MaWte (ey p2xx)11
F
Recall that || = A),
B) , tr = tr(BAT), which means that
T _ F | 1 | and ( fa )F = (pg. pu)) (xx,, xe). Substituting that into the
displayed equation gives us the desired result.
Remark D.3. Let P E dx be an semi-orthogonal matrix and let W E Red be matrix. Then
P PPWPl S [|
Proof. Given 2 matrices A E Raxm B E pmmxd it is will known that |AB|, < |||]| [B]r Therefore,
P PWPP
< ||P]]= || w < W |
35

---

E Further Experiments & Additional Details
In this section, we provide further experiments and additional details. We trained linear attention models (Eq. (3)) on data generated as specified in Section 2.1 using GD with fixed step size a = 0.01, = 40 and the logistic loss function. Training was performed for 200 300 steps from a zero initialization, implemented in Py PyTorch. In all figures, the x-axis represents the number of in-context examples M, starting from M = 1, while the y-axis corresponds to the test accuracy. The accuracy is computed by checking whether the prediction EE; Wt) equals the true label VM++. where E denotes the tokenization of the sequence (at, v1),..., (Ir M: 3M), (#+ 1, 0), used to predict +1++1- We average the accuracies over Buest 1200 tasks, and we plot this average, optionally including error bars representing 95% confidence intervals (CIs). These intervals are computed using the
standard error of the mean and assuming normal distribution:
CI acc - 2ō, acc + 20),
where ace := E, alc; Brest is the accuracy mean, := E, (ace; acc)"/ 1) is the standard error of the mean and = 1.96. All computations can be completed within an hour on CPU.
In Figure 2, we examine the effect of the shared subspace dimension k; with each plot corresponding to a different signal strength. Accuracy improves as the signal strength It = R increases and as the subspace dimension k decreases. Notably, in some plots, the accuracy remains far from even for large values of M. This can be attributed to the term Vk A R A M /k that appears in the generalization bound of Theorem 3.1. Indeed, when M is sufficiently large, the dominant term becomes v A Rt indicating that perfect accuracy (i.c., zero error) is unattainable in this regime.
Accurecy Contect (2
femeary Ve Cortlest Siew
200000
foriaay Contrett 2000)
1-500
1-500
oc longh' (M)
Segurco Jorgit (5)
Seauende Jergh1 (9)
Figure 2: Test accuracy versus the number of in-context examples M for various subspace dimensions k. Accuracy improves as the signal strength R R increases and as the subspace dimension decreases. B d 2000.
36


Reviews:
[{"review_id": "Mo48W96NR8", "reviewer": "Reviewer_PK2N", "summary": "This paper studied the in-context learning of one-layer linear self-attention models to do binary classification on a Gaussian mixture dataset. After training such a transformer, it has learned the ground truth projection, so it only requires fewer test samples to achieve a small test error. It reduces the sample complexity compared with the optimal binary classifier on the test dataset. This result improved previous results of [11] by Spencer Frei and Gal Vardi. This paper, employing a straightforward linear model, theoretically demonstrated that transformers can function as metalearners, capable of adapting to novel tasks using only a limited number of in-context examples, without the need for additional training.", "questions": "1. Line 179: there should not be index $\\tau$. Line 305: Gausian\n\n2. In section 2.2, I understand the Tokenization is normal in many theory paper, but can we consider the embedding different from (1)? For instance, $\\mathbf{E}=(x_1^\\top,y_1,\\ldots,x_n^\\top,y_n,x_{n+1}^\\top,0)$.\n\n3. In Theorem 2.3, if we consider nonlinear softmax attention, can we still get some similar solution for the gradient descent training? \n\n4. In theorem 3.1, what is $\\mathbf{E}$? Is it training or test dataset?\n\n5. In the last remark in Section 3, the authors can eliminate the dependence of $B$ and $N$ on $d$ in the analysis. Can you compare such results with the classical PCA method for Gaussian mixture models, e.g. [1]? In [1], the threshold for exact recovery is also governed by sample size and SNR. Also, it is worth mentioning [2] as another direction for studying transformers using Gaussian mixture models.\n\n-----------------------------------------------------------------------------------------------\n[1] An $l_p$ theory of PCA and spectral clustering\n\n[2] Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures", "limitations": "The paper should have a detailed discussion of the limitations in the main text. The limitation of the theory may come from the assumption of the Gaussian mixture model and the linear self-attention model.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: This paper is well-organized and states the theorems in a very clear way. It is technically sound with sufficient discussions on the results and the proof is clear. \n\nWeaknesses: The main concern is the assumptions of the dataset and attention model in the theory. This paper only considers the linearly separable binary classification and a convex parameterization of the linear transform. These may loose many crucial properties of the transformer both in pretraining and test process. Additionally, it should be more impactful if the authors could extend the results to in-context multi-class classification like reference [12]. And although getting theory for more general setting may be hard, the paper could present more general simulations to show its result can be extended to more general settings. Finally, it is not clear what kind of message a machine learning practitioner can obtain from this theoretical work.", "was_revised": true, "final_justification": "The authors partially addressed my concerns. I greatly appreciate their effort and detailed discussion. However, I cannot agree with their first response: Eq. (3) serves as a standard toy model for theoretical analyses of in-context learning, so we can focus on this toy model and do not need to enhance our analysis to make it more realistic for practical transformer models. I believe Eq. (3) does not fully capture the behavior of the genuine Transformer attention mechanism. As for the impact on practitioners, I think this paper uses this toy model to elucidate some ideas that may already be familiar to practitioners. Based on the novelty of the mathematical methods and the final results presented in the paper, I will maintain my current score.", "rebuttal": "We thank the reviewer for their feedback. Below, we respond to each point raised. \n\n**Simplified Architecture:** While we analyze a 1-layer linear transformer, this simplification allows us to rigorously characterize fundamental properties of in-context learning that may extend to more complex architectures. Similar to how early theoretical works on neural networks focused on simplified architectures to establish foundational results, our work provides a theoretical framework that can guide future analyses of more complex transformers.  \n\nMoreover, our architecture is exactly identical to the architecture considered in at least four prior works on in-context learning in linear classification or linear regression [1-4 below] (all of them appeared in top tier ML conferences), and other variants of 1-layer linear attention models have appeared in many other works on in-context learning in linear regression tasks (5-8, and much more). Thus, this model has become a standard toy model for theoretical analyses of in-context learning, and it clearly attracts much interest in the theoretical deep-learning community. Hence, we feel that recommending rejection based on this reason would be unfair.\n\n[1] Juno Kim, Tai Nakamaki, and Taiji Suzuki. Transformers are minimax optimal nonparametric in-context learners. NeurIPS 2024. \n\n[2] 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? ICLR 2024. \n\n[3] Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025. \n\n[4] Wei Shen, Ruida Zhou, Jing Yang, and Cong Shen. On the training convergence of transformers for in-context classification. ICML 25. \n\n[5] Ruiqi Zhang, Spencer Frei, Peter L. Bartlett . Trained Transformers Learn Linear Models In-Context. JMLR 2024. \n\n[6] Oswald  ET AL. Transformers Learn In-Context by Gradient Descent. ICML 23. \n\n[7] Ahn ET AL. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 23. \n\n[8] Mahankali et al. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. \n\n \n**Data Setting:** We believe our assumptions are reasonable for a theoretical work that gives a precise analysis of the effects of all of the core ingredients to the transformer training pipeline: the implicit regularization effect of gradient descent in pre-training, the number of pre-training tasks, the number of examples per pre-training task, the signal-to-noise ratio of the pre-training tasks, the number of in-context examples at test time and the signal-to-noise ratio of the test-time tasks. Moreover, Gaussian mixture models is a well-studied setting in the literature, which enables us to derive a tight lower bound on the performance of an optimal metalearner (see Remark 2.2). Perhaps surprisingly, our analysis shows that the trained transformer nearly matches this optimal performance. We also note that Frei and Vardi [3] studied in-context learning in transformers using a very similar setting. To the best of our knowledge, there are currently no theoretical works on transformer-based in-context learning for classification that use significantly more realistic or complex datasets and provide such comprehensive results.  \n\n[3] Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025. \n\n**Impact for Practitioners** \nWe showed that simple linear attention can efficiently encode task structure when the tasks lie in a low-dimensional subspace, and that transformers can approach optimal performance with substantially fewer in-context examples, assuming sufficient pretraining. These insights may guide architectural decisions in real-world few-shot applications. That said, the primary goal of this work, like many theoretical studies, is to deepen our understanding of how transformers perform metalearning and to clarify the relationships between key factors such as the number of pretraining tasks, the number of examples per task, and the SNR to enable metalearning. We are aware that theoretical lines of research often take time to translate into significant practical impact.  \n\n\n**Respond for Questions:**  \n**Q1:** Thank you for pointing this out. We will fix that in the revised version.\n\n\n**Q2:** We are not entirely sure we understand the proposed tokenization. In our setup, each column of $E$ corresponds to a different token. As a general principle, each token fed into the transformer must have the same dimension. If the idea is to tokenize each $(x_i,y_i)$ pair as two separate tokens, this raises several issues. First, it would require padding $y_i$  to match the dimension of $x_i$, which adds unnecessary complexity. More importantly, this design leads to attention interactions between label tokens and input tokens within the same example, rather than between different examples. We therefore suspect that such a tokenization would impair generalization and increase the in-context sample complexity. \n\nWe agree that analyzing different tokenization is an interesting question. \n\n\n**Q3:**  As noted in the discussion of our “simplified architecture”, our analysis relies on the linear attention model, where the implicit bias of gradient descent is well understood. We believe it should be possible to extend our analysis to non-convex linear transformer architectures which are homogeneous in the parameters, since the KKT condition analysis approach can work in that setting as well. Unfortunately, the implicit bias of gradient descent in the softmax attention setting is not yet fully characterized (see references in line 151), which makes it challenging to extend our proof techniques to that case. Developing such a theory for softmax-based transformers remains an important and open direction, as we mentioned in Section 6. \n\n\n**Q4.**  Thank you for pointing this out. $E$ refers to the test data. Specifically, the tokenization of the sequence $(x_1,y_1), \\ldots, (x_M,y_M), (x_{M+1}, 0)$, used to predict $y_{M+1}$. We agree that this was not clearly stated, and we will add an explicit clarification in the final version. \n\n\n**Q5:**  In PCA  the data matrix $X \\in \\mathbb{R}^{n \\times d}$ is given, and the goal is to compute a rank-$k$ approximation that minimizes the Frobenius norm error. The rank $k$ is known and fixed, and the objective is purely reconstruction-based. \n\nIn contrast, in our setting, the rank $k$ is not provided in advance, and the transformer must implicitly infer a relevant subspace through the attention mechanism in the course of solving a different task: namely, in-context linear classification. That is, the attention head learns to separate labeled data points according to class, rather than recover the underlying data geometry. This difference in objective makes a direct compression from one setting to the other nontrivial. \n\nMoreover, the guarantees in [9] rely on an asymptotic regime, where the number of samples $n \\to \\infty$ (see Assumption 2.1 therein). Translating to our setting, this would require the number of tasks to tend to infinity. In contrast, our bounds hold for a finite number of tasks and, notably, are independent of the ambient dimension $d$. \n\nWe will add a citation to [10], which takes a complementary approach by showing that transformers are capable of performing unsupervised learning in the Gaussian mixture model (GMM) setting.  \n\n [9] An $\\ell_p$ theory of PCA and spectral clustering\n\n [10] Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures"}, {"review_id": "XM6rZCX9aQ", "reviewer": "Reviewer_5r5s", "summary": "The paper considers the ICL ability of linear transformers on a Gaussian mixture of linear classification.\nIt shows that a trained transformer is nearly an optimal metalearner.\nThe result also shows the number of pretraining tasks needed to generalize at test time, where the number is not dependent on the ambient dimension.", "questions": "I have a single question on the validity of the data generation setting, which will affect whether my score is accept:\n\nIn Assumption 2.2 (and similarly in Assumption 2.1), is it true that the $x_i$ with $y_i=1$ is of the same distribution as the $x_i$ with $y_i=-1$? Since $\\mu'$ is uniform on the sphere and $P$ is isometric, it looks to me that the answer is yes. In that case, how is classification possible?\n\nThis question is particularly intriguing to me since in Line 261 and Line 272, the paper is claiming that the algorithm can achieve \"an arbitrarily small constant\" test error. I am not exactly sure how it is possible in the current setting, and would be eager to hear from the authors.", "limitations": "N/A on societal impact.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "This is a technical solid paper that studies a rather important problem in ICL theory.\nIt considers the mixture of classification tasks, which is an important theoretical setting that I do not think has been properly addressed in previous work.\nThe theoretical result is informative and demonstrates the ICL ability of linear transformer in such setting.\n\n---\nI will detail my observation as follows.\n\n- Methods And Evaluation Criteria:\n\nThe data distribution setting, namely the Gaussian mixture of linear classification task, is a rather novel setting and an important step toward understanding the ICL ability of linear transformers.\nICL on single-distribution linear regularization has been investigated in e.g. [1], and a potential setting on the Gaussian mixture of linear regularization has been proposed in a recent ICML submission [2].\nThe Gaussian mixture of linear classification task is a step forward.\nSince it is a novel setting, I would like to further check with the authors on whether the setting is reasonable, i.e., it does not refer to some degenerate case or some cases where metalearning is unachievable.\nThe definition itself in Section 2.1 is mathematically rigid, but I would like to seek some clarification on its implication.\nTherefore, I will specify my questions on the problem setting in later sections. \n\n- Theoretical Claims:\n\nI skimmed through the proofs including the proof sketch in Section 4 and the Appendix.\nThe proof uses several standard tools in analyzing linear transformer, namely Eq. (3) which has been present in [1, 2].\nThe proof involves some standard concentration lemmas, and the concentration lemmas are rather established in the textbooks, so it looks to me that the proved result stands as it is.\n\nThe main result reduces the bound on in-context examples so that it no longer depends on the ambient dimension $d$, which is a good improvement.\n\n- Experimental Designs Or Analyses:\n\nThe numerical experiment replicates and justifies the theoretical findings.\nIt is not the main contribution in this paper anyway, so it does not affect my evaluation.\n\n\n- Relation To Broader Scientific Literature:\nIt is closely related to the theoretical in-context learning literature. It studies the in-context learning ability on a Gaussian mixture of linear classification models, which is a step forward from the non-mixed linear regression/classification function class that is present in the literature.\n\nIf the mixture of linear classification setting makes sense, this is surely a good paper for this venue. However, I have some doubts about the setting, as stated before. Therefore, I have a single question on the setting which will greatly impact whether my score is accept or reject, which is listed in the questions below.\n\n---\n\n[1] Zhang, R., Wu, J., & Bartlett, P. (2024). In-context learning of a linear transformer block: Benefits of the mlp component and one-step gd initialization. Advances in Neural Information Processing Systems, 37, 18310-18361.\n\n[2] Jin, Y., Balasubramanian, K., & Lai, L. (2024). In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamics. arXiv preprint arXiv:2410.14183.", "was_revised": true, "final_justification": "Thank the authors for their rebuttal. It addresses my concern. I believe this work is quite important for the community since it theoretically demonstrates a very important and strong practical feature (meta-learning) of a very influential model (linear transformer). I will raise the score from 4 to 6 and recommend this paper for wider exposure, like oral or spotlight, depending on how the venue is organized.\n\nI recommend that the authors somehow reflect their response into the paper in Assumptions 2.1, 2.2, and Ln 186-188. It looks to me that my initial misunderstanding is quite natural, since it is currently not explicitly mentioned in Assumptions 2.1 and 2.2 that they are the *task* distribution rather than the *data* distribution. For example:\n\nAssumption 2.1' (training-time **task** distribution). ... For any task $\\tau \\in [B]$, the input and label pairs $(x_{\\tau,i}, y_{\\tau,i})_{i=1}^{N+1}$ in task $\\tau$ satisfies the following:\n1. Let $\\mu'_\\tau$ be sampled from a prior distribution $\\text{Unif}(\\tilde{R}\\cdot S^{k-1})$, i.e., the uniform distribution on the sphere of radius $\\tilde{R}\\$ in $k$ dimensions.\n...\n\nThis and something similar for Assumption 2 should work. Also, in Ln 186, you might want to specify that \"In our setting, each example is a task $(x_i, y_i)_{i=1}^{N+1}$ that consists of a sequence of $(x,y)$ pairs\". \n\nI believe something like these edits can prevent other people from suffering the same misunderstanding as mine. This can also keep the notation consistent with Ln 257 in Theorem 3.1.\n\nPlease feel free to comment if this suggestion makes sense.\n\nAlso, I disagree with Reviewer uuw9's evaluation on Eq. (3). As mentioned in my review, it is a standard tool that has been present in prominent papers in this field.", "rebuttal": "We thank the reviewer for their thoughtful feedback. \nRegarding your question, we believe you are referring to Assumptions 2.1 and 2.2 (there is no Assumption 3.1 or 3.2). The answer is no,  but it is essential to also condition on $\\mu$. That is, the distributions of $x_i$ conditioned on $y_i = 1$ and $y_i = -1$ are not the same, when we also condition on $\\mu$. \n\nIn both assumptions, the task-specific vector $\\mu$ is sampled once for each task and fixed. Then, \n\n$x_i = y_i \\cdot \\mu + z_i,$ \n\nwhere $z_i \\sim \\mathcal{N}(0, I_d)$. Thus, conditioned on $\\mu$, the positive and negative classes are Gaussian distributions centered at $\\mu$ and $-\\mu$, respectively. This separation makes classification possible despite the symmetry in the prior over $\\mu$. Thus, since our in-context sequence consists of many examples $(x_i,y_i)$ which are drawn from a Gaussian mixture that corresponds to the same $\\mu$, then given enough in-context examples we can correctly classify fresh examples which are drawn using this $\\mu$. \n\nWe hope this clarifies the data generation process and addresses the concern about class distinguishability."}, {"review_id": "8Udp1WSOSB", "reviewer": "Reviewer_ibKw", "summary": "This paper analyzes in-context learning (ICL) with a single-layer linear attention model on Gaussian mixture tasks. It provides theoretical results showing that such an architecture is an almost optimal metalearner for binary classification. The analysis carefully treats the finite-sample setting, which makes it distinct from many prior theoretical ICL works that assume a population training loss.", "questions": "1. Do Remarks 2.1 and 2.2 rely on Assumption A? If so, could the authors clarify this connection?\n2. A recent paper shows that a single-layer linear attention layer can indeed learn the optimal binary classifier. Could the authors clarify it?\n - Li, Yingcong, et al. When and How Unlabeled Data Provably Improve In-Context Learning. arXiv:2506.15329 (2025).\n3. What specific model architecture is used in the experiments? Are the authors implementing Eq. (2) by replacing softmax with linear attention, or are they directly simulating Eq. (3)?\n4. The experimental curves are quite noisy. It would be helpful to run more trials and report the averaged results with confidence intervals to ensure the trends are robust.\n5. Have the authors tried comparing to the simplest baseline: estimating the task mean directly using $\\sum_iy_ix_i$? How does its performance compare to the linear attention model?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper is clearly written and provides rigorous results in a finite-sample regime, addressing a realistic scenario where the number of tasks and samples per task is finite, which is an important distinction from many previous ICL theory works that assume infinite training data.\n\n2. The main results, particularly Theorem 3.1, are accompanied by detailed discussion and intuition, helping the reader understand the key contributions and their implications.\n\nWeaknesses:\n1. Theorem 2.3 holds under the assumption that the data is linearly separable so that an SVM solution exists. In real in-context learning, demonstration samples are typically sampled randomly, which may not always yield separable sets. The paper would benefit from clarifying how $W_{MM}$  in Theorem 2.3 relates to practical training, as well as the feasibility  of $W_{MM}$ for Theorem 3.1.\n2. The result in Theorem 2.3 align with optimizing a single weight matrix $W$ as in Eq. (3). However, the actual attention architecture (Eq.(2)) is more complex. it involves separable $W_q$, $W_k$ matrices, which can lead to different optimization behaviors. Prior work has shown that separately optimizing $W_q$ and $W_k$ can implicitly minimize the unclear norm of $U$ in Eq. (4) rather than Frobenius norm. Thus, the paper's theoretical focus on Eq. (3) does not fully capture the behavior of the true Transformer attention mechanism.\n3. The paper lacks a concrete comparison between the bound in Theorem 3.1 and the Bayes optimal loss. Additionally, the Remarks 2.1 and 2.2 are not consistent with Theorem 3.1 since they do not appear to account for Assumption A, which seems critical for making the comparison meaningful. It would be better to show the main results as the excess risk between the \"single-task optimal learner\" and \"optimal learner with access to the ground-truth representation\". Additionally, simply showing that each error is constant in [0,1] is less informative. It would be more compelling to show whether the error decays exponentially with increasing samples, context length or dimension.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns. I suggest incorporating more discussion on comparison to the Bayes error. I have decided to raise my score.", "rebuttal": "We thank the reviewer for their thoughtful feedback. Below we respond to each concern and question in order. \n\n**W1:**  \n- **Feasibility of $W_{MM}$:**  We note that this is explicitly addressed in our analysis. See Lemma C.4 and the proof sketch (line 325) for details, where we show a specific construction (i.e. the projection matrix) that separates the data (with high probability over the data distribution). While extending the analysis to non-separable data is an important direction for future work, we believe that studying the separable setting is a meaningful and informative starting point. \n\n- **Data Setting Justification:** We believe our assumptions are reasonable for a theoretical work that gives a precise analysis of the effects of all of the core ingredients to the transformer training pipeline: the implicit regularization effect of gradient descent in pre-training, the number of pre-training tasks, the number of examples per pre-training task, the signal-to-noise ratio of the pre-training tasks, the number of in-context examples at test time and the signal-to-noise ratio of the test-time tasks. Moreover, Gaussian mixture models is a well-studied setting in the literature, which enables us to derive a tight lower bound on the performance of an optimal metalearner (see Remark 2.2). Perhaps surprisingly, our analysis shows that the trained transformer nearly matches this optimal performance. We also note that Frei and Vardi [3] studied in-context learning in transformers using a very similar setting. To the best of our knowledge, there are currently no theoretical works on transformer-based in-context learning for classification that use significantly more realistic or complex datasets and provide such comprehensive results.\n\n[3] Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. ICLR 2025.\n\n- **Relation between $W_{MM}$ and practical training:** Formally, the model is guaranteed to converge to $W_{MM}$ asymptotically as the time $t$ tends to infinity.  In our experiments, we find that gradient descent trained for a finite number of steps already achieves strong generalization, closely matching the behavior predicted by the max-margin solution. We note that analyzing the asymptotic solution rather than a finite-time solution is common in many theoretical papers. \n\n**W2:** We agree that analyzing a more general transformer architecture with separate key and query matrices is an important and interesting direction for future work. In our work, we adopt the simplified attention model with a merged key-query matrix, which is commonly used in theoretical studies of transformers and in-context learning (e.g. [3-7], and many more). This setting allows a tractable analysis of the optimization dynamics and implicit bias. \n\nRegarding the reviewer’s comment on nuclear norm minimization, we believe this refers to [2]. As we mentioned in the paper, their analysis relies on strong and somewhat artificial assumptions, such as the existence of a single “optimal” token, and all other tokens receiving identical “scores”. These assumptions are not satisfied in our setting and, in fact, are violated by almost all data distributions with more than two tokens drawn randomly. \n\nWhile our model does not capture the full complexity of attention with separate key and query matrices, we believe it serves as a meaningful and interpretable proxy that sheds light on how gradient descent interacts with attention in metalearning tasks. Also, as we already mentioned, this is a common setting in theoretical works on transformers.  \n\n[2] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. \n\n[3] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ICML 23. \n\n[4] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 23. \n\n[5] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. JMLR 2024. \n\n[6] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in the attention mechanism. \n\n[7] Vasudeva, Bhavya, Puneesh Deora, and Christos Thrampoulidis. ”Implicit bias and fast convergence rates for self-attention.” \n\n**W3:**  \n\n **Error is not a constant** : We are not sure why the reviewer mentioned that the error in our result is constant. Theorem 3.1 shows that the error decays exponentially when several parameters increase simultaneously, specifically, when the subspace dimension k, the test-time signal strength \\Tilda(R), and the quantity \\sqrt{M\\Tilde(R)^4/k}  all increase, assuming the number of training tasks B is large enough to recover the shared representation.  \n\n**Comparison with Bayes optimal:** \n- In lines 260-268, we compare the transformer against a Bayes optimal algorithm that has access just to the samples from the new task. While the latter requires \\Omega(d/\\tilde{R^4}) examples to achieve small constant error (d is the ambient dimension and \\Tilde(R) is the signal strength during test), the transformer needs only O(k/\\tilde{R^4}) in-context examples, as long the number of tasks is large enough.  \n\n- In lines 269-275 we compare the transformer against a Bayes optimal algorithm that also has full access to the shared sub-space, which is a very strong baseline, since metalearners don’t have access to the sub-space, and they need to learn it during training. We showed that both have the same in-context sample complexity, up to polylogarithmic factors. \n\n**Relation between Assumption A to Remark 2.1 and 2.2:**  We clarify that Remarks 2.1 and 2.2 do not rely on Assumption A. These remarks present information-theoretic lower bounds that hold independently of any specific algorithm or training setup. In contrast, Assumption A is used specifically to prove our main result (Theorem 3.1) and ensures that the max-margin solution generalizes with high probability. \n\nAs discussed in lines 248--253: \n\n- Assumption A1: requires that the signal strength $R$ is at least $\\widetilde{\\Theta}(1)$. This is the interesting regime, since Remarks 2.1 and 2.2 show that learning becomes impossible when the test-time signal $\\widetilde{R}$ is $o(1)$. \n\n - Assumption A2 requires that the ambient dimension $d$ is at least polylogarithmic in the number of tasks. This modest assumption, while not needed in Remarks 2.1 and 2.2, is standard in high-dimensional analysis and ensures concentration bounds for sub-exponential random variables. \n\n- Assumption A3 requires the number of examples per task during training to be at least $N \\geq d / R^2$, or equivalently $1/\\text{SNR}^2$. This is necessary for learning the shared representation during pretraining. However, Remarks 2.1 and 2.2 analyze test-time learners that either do not have access to the shared sub-space $P$ (Remark 2.1) or have oracle access to $P$ (Remark 2.2), so their bounds are independent of $N$. \n\n**Q1.** See the previous part i.e. Relation between Assumption A to Remark 2.1 and 2.2. \n\n**Q2.** \n**Comparison to Li et al.:** We thank the reviewer for pointing us to this relevant reference. While the setting in Li et al. differs from ours in important ways, it is still conceptually related and warrants citation in our paper. \n\n- First, their work does not address the metalearning setting. In particular, the tasks in their framework do not share a common representation (e.g., a low-dimensional subspace), which makes their setting fundamentally different and not directly comparable to ours. In fact, their focus is closer to semi-supervised learning and naturally their bounds depend on the ambient dimension $d$, while our work avoids that. \n\n- Second, they analyze the squared loss, which is more natural for regression and allows for closed-form solutions. In contrast, we study the logistic loss, which is standard for classification and does not admit a closed-form solution. To handle this, we leverage the implicit bias of GD to characterize the behavior of the trained model. \n\n- Third, their results do not clearly describe how key parameters influence generalization. For example, how the number of tasks during training affects the generalization error. In our work, we provide explicit relationships between generalization error and quantities such as the number of training tasks, the number of samples per task, the ambient dimension, shared subspace dimension, and the SNR both in training and test. \n\n  \n\n**Q3.** We are directly simulating Eq. (3) \n\n**Q4.** We have included this in the experimental section of the appendix (see the supplementary materials, last section). \n\n**Q5.** Yes, this is precisely the Bayes optimal algorithm, that is, the maximum likelihood estimator (MLE) that we use as a baseline in our experiments (see the discussion at the beginning of Section 5 in the main paper, as well as the experiments presented both in the main text and in the appendix). We evaluate two versions of MLE: one that only uses the average of the in-context test examples, and another that has access to the ground-truth subspace matrix P, which first projects the data using P and then applies MLE. Our results show that the linear transformer significantly outperforms the MLE baseline that lacks access to P, and closely matches the performance of the MLE with projection."}, {"review_id": "4qOhB9UmMB", "reviewer": "Reviewer_uuw9", "summary": "This paper theoretically investigates that  ICL with transformer can match the optimal meta-learner, on linear regression tasks with class-conditional Gaussian mixture model, where all tasks share the same input space.", "questions": "Please refer to weakness.", "limitations": "NA", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The main conclusion in this paper that with sufficient training, ICL with transformer acts like optimal meta-learners, is reasonable. \nHowever, it seems that [1] also gives such result, in more general settings, empirically without theoretical guarantee. But [1] is not mentioned at all. \n\nThough the reviewer generally agrees with the main conclusion, there are some critical concerns:\n1. It seems that the results have nothing to do with ICL with transformer. It seems that without Sec 2.2, by replacing it with any model parameterized by W, the sequel still holds. As the sequel only rely on the setting/assumptions, rather then the model.\n2. The formal result (Theorem 3.1) is confusing. The review understand the way to decomposition of model into task-specific $f$ and shared $h$, which is common in multitask-learning and meta-learning [2]. In the particular setting (Sec 2.1) where the problem is linear and $h$ is linear (the step 1&2 in assumption 2.1/2.2 is no other than replacing the input from uniform distribution on unit sphere in $R^k$ with $R^d$), as far as the reviewer could understand, useful information can only been brought when d<k, as the noise $z$ would has less dimensions with fixed $I_d$. However, the main result (theorem 3.1) shows that the meta-learning risk would increase with d decreases, which is contradictory.\n\nThe reviewer would consider to raise the score if related works and above concerns can be discussed properly.\n\n[1] Wu S, Wang Y, Yao Q. Why In-Context Learning Models are Good Few-Shot Learners?[C]//The Thirteenth International Conference on Learning Representations.\n[2] Maurer A, Pontil M, Romera-Paredes B. The benefit of multitask representation learning[J]. Journal of Machine Learning Research, 2016, 17(81): 1-32.", "was_revised": true, "final_justification": "After rebuttal, my second concern was addressed, so I raised my score. But my first concern remains critical, and seems not could be addressed by discussion or minor revision.", "rebuttal": "We thank the reviewer for their thoughtful feedback. Below we respond to each concern and question in order.\n\n**Missing Citation to [1] (Wu et al.):**  Wu et al. (2025) present an empirical study suggesting that transformers can approximate optimal meta-learners in few-shot classification tasks. While their results are encouraging, the work is primarily empirical and does not offer formal guarantees. The only theoretical result in the paper is Theorem 2.1, which establishes an expressivity claim, that is, the existence of transformer weights that can simulate a metalearning algorithm. However, they do not analyze whether such a solution can be obtained through optimization, which is a central focus of our work. \n\nIn addition, their analysis relies on the universal approximation property of multi-layer perceptrons (MLPs), without specifying the required network size. In contrast, our model does not include an MLP component. Finally, it is unclear whether the algorithms studied in their paper achieve optimality in our Gaussian mixture setting. \n\nOur contribution complements this line of work by providing the first formal guarantee that a simple transformer trained by gradient descent can achieve near-optimal meta-learning performance under well-defined assumptions.  \n\nWe agree that Wu et al. is a relevant and timely reference, and we will cite their work in the revised version. \n\n\n**W1. Is the Result Specific to Transformers?**\n\nThe answer is yes. While our theoretical guarantee in Theorem 3.1 is expressed in terms of the max-margin solution $W_{mm}$ , this solution is directly tied to the specific transformer architecture presented in Eq. 3. Section 2.2 introduces a linear attention-based transformer model, which has been studied in several prior works, and our analysis closely follows the behavior of this model when trained via gradient descent. \n\nTheorem 2.3 defines the max-margin solution $W_{mm}$ in terms of this attention model, as noted in line 215. Specifically, the constraint in Eq. (4) reflects the prediction structure of the model described in Eq. (3). **If a different model were used, the constraints would differ, leading to a different optimization problem and, in general, a different solution.** As a result, our analysis would no longer apply in that case. \n\n**W2. Clarification on Theorem 3.1 and the Role of d and k:** \n\nWe believe there may be a misunderstanding: \n\n**Clarification of the setting:** In our setting, each task's signal vector $\\mu$ lies in a $k$-dimensional subspace of $\\mathbb{R}^d$, and the noise is isotropic in $\\mathbb{R}^d$. Therefore, **the data is only defined for $k \\leq d$** (see line 162). Indeed, Assumption 2.1 requires that the matrix $P \\in R^{d \\times k}$, whose columns span the shared subspace, is semi-orthogonal. That is, $P^TP=I_k$, which is only possible when $k \\leq d$, as stated in the assumption. \n\nEach task's signal $\\mu \\in \\mathbb{R}^d$ is generated by embedding a vector $\\mu' \\sim \\text{Unif}(R \\cdot \\mathbb{S}^{k-1})$ into $R^d$ via $\\mu = P \\mu'$. Importantly, **$\\mu$ is not uniformly distributed over the unit sphere in $\\mathbb{R}^d$**; with probability 1, it lies entirely within the shared subspace and has no component in directions orthogonal to it. \n\nIn standard metalearning terminology, the optimal shared representation in our setting is given by $h: \\mathbb{R}^d \\rightarrow \\mathbb{R}^k$, where $h(x) = P^T x$. For a signal vector $\\mu = P\\mu’$, this gives: $h(\\mu) = P^T\\mu = P^TP\\mu’ = \\mu’$, which recovers the original signal in the k-dimensional subspace.  \n\nGiven a task with signal vector $\\mu \\in \\mathbb{R}^d$, constructed from  $\\mu’ \\in \\mathbb{R}^k$, the task-specific classifier $f: \\mathbb{R}^k \\rightarrow (-1, +1)$ is defined as $f(x) = \\text{sign}(\\mu'^T x)$. \n\n**Risk would increase with d decreases.**  Our result in Theorem 3.1 shows that the meta-learning error decreases as the ambient dimension $d$ decreases, not the other way around. In particular, as $d$ becomes smaller, the SNR during training increases, which facilitates the learning process and leads to better generalization.."}]

Author rebuttal:

Meta-review: This submission studies in-context learning in a transformer as meta-learning, for the restricted case of linear classification. The submission presents theoretical contributions that reviewers found well-motivated and interesting. The primary concern raised across reviews centers on the restrictive nature of the theoretical assumptions. While these assumptions may limit the immediate applicability of the results, the reviewers assessed that the motivation and theoretical insights are sufficiently compelling to warrant acceptance.
Decision: Accept (poster)