On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Authors: Wei Shen
Ruida Zhou
Jing Yang
Cong Shen
Venue: ICML 2025 poster (2025)
Keywords: In-context learning
Transformer
Primary area: theory->learning_theory
TL;DR: Transformers trained via gradient descent can provably perform in-context classification of Gaussian mixtures.

Abstract: Although transformers have demonstrated impressive capabilities for in-context learning (ICL) in practice, theoretical understanding of the underlying mechanism that allows transformers to perform ICL is still in its infancy. This work aims to theoretically study the training dynamics of transformers for in-context classification tasks. We demonstrate that, for in-context classification of Gaussian mixtures under certain assumptions, a single-layer transformer trained via gradient descent converges to a globally optimal model at a linear rate. We further quantify the impact of the training and testing prompt lengths on the ICL inference error of the trained transformer. We show that when the lengths of training and testing prompts are sufficiently large, the prediction of the trained transformer approaches the ground truth distribution of the labels. Experimental results corroborate the theoretical findings.

On the Training Convergence of Transformers for In- In-Context Classification of Gaussian Mixtures
Wei Shen 11
Ruida Zhou *2
Jing Yang' Cong Shen
Abstract
abilities of transformers. A more comprehensive literature review can be found in Appendix A. Garg et al. (2022) was the first to study the ICL abilities of transformers for various function classes. They empirically showed that transformers can learn linear regression models in context. Later on, a line of research was developed to theoretically explain how transformers perform in-context linear regression. For ex- ample, Akyürek et al. (2022); Von Oswald et al. (2023); Bai et al. (2024); Fu et al. (2023); Giannou et al. (2024) showed by construction that, some specially-desinnet transformers can perform linear regression in context. Moreover, some re- cent works such as Zhang et al. (2023a); Huang et al. (2023); Chen et al. (2024) studied the training dynamics of a single- layer transformer for in-context linear regression. They proved the convergence of certain single-layer transformers during training and showed that the trained transformers able to perform linear regression in context.
Although transformers have demonstrated impres- sive capabilities for in-context learning (ICL) in practice, theoretical understanding of the under- lying mechanism that allows transformers to per- form ICL is still its infancy. This work aims to theoretically study the training dynamics of trans- formers for in-context classification tasks. We demonstrate that, for in-context classification of Gaussian mixtures under certain assumptions, single-layer transformer trained via gradient de- scent converges to a globally optimal model at linear rate. We further quantify the impact of the training and testing prompt lengths on the ICI. inference eтror of the trained transformer. We show that when the lengths of training and test- ing prompts are sufficiently large, the prediction of the trained transformer approaches the ground truth distribution of the labels. Experimental re- sults corroborate the theoretical findings.
Building on the earlier works that largely focus on linear regression problems, several recent papers have started to investigate the ICL capabilities of transformers for classi- fication problems. For instance, Bai al. (2024) showed that, by construction, multi-layer transformers proximately viewed as multiple steps of gradient descents for logistic regression. Giannou et al. (2024) further show- cased that the constructed transformers can approximately perform Newton's 's method for logistic regression. Lin & Lee (2024) studied the dual operating modes for in-context classification of Gaussian mixtures. However, their analy- ses were based on the idealized Bayes-optimal next-token predictor for ICL tasks and did not consider the training dynamics of transformer models. Some recent works have started to study the training dynamics of transformers for cer- tain classification problems. For example, Li et al. (2024b) proved that a single-layer transformer can be trained to learn one-nearest neighbor in context. Li al. (2024a) studied the training dynamics of a single-layer transformer for some binary classification tasks with finite, pairwise orthogonal patterns. Frei & Vardi (2024) analyzed linear transformer model for in-context classification of Gaussian mixtures, assuming additional conditions such as a suffi- ciently large signal-to-nol noise ratio. However, all these works (Li et al., 2024a;b; Frei & Vardi, 2024) only considered the binary classification problems. The training dynamics
be ap-
1. Introduction
Large language models (LLMs) based on the transformer architecture (Vaswani et al., 2017) have demonstrated re- markable in-context learning (ICL) abilities (Brown et al., 2020). When given a prompt consisting of examples of a learning task, these models can learn to solve this task for new test examples without any parameter updating. This be- havior has been empirically demonstrated in ofte-off-the-rt models on real-world tasks (OpenAI, 2023; Touvron et al., 2023).
This impressive capacity of transformer- based models has inspired many recent works aiming to understand the ICL
'Equal contribution 'Department of Electrical and Com- puter Engineering. University of Virginia, Charlottesville, USA 'Department of Electrical Engineering, University of Cali- foriaa, Los Angeles, USA. Correspondence to: Cong Shen <cong@ virginia.cdu>>
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
of transformers for more general in-context classification problems beyond the specifie settings and assumptions in Li et al. (2024a:b); Frei & Vardi (2024) remain largely under- explored.
2.1. Single-layer transformer
Given an input embedding matrix E E Rd.xde a single head self-attention module FSA with width de will output
Es W W
In this work, we study the training dynamics of a single- layer transformer for both binary and multi-class classi- fication of Gaussian mixtures during ICL, a fundamental problem in machine learning. Our main contributions can be summarized as follows:
=F+ E. san (W E) (OREE
(I)
where WV, WK, WQ E Rddedde are the value, key, and query weight matrices, respectively, P > 0 a normaliza- tion factor, and fan is an activation function for attention. There are different choices of famo: for example Vaswani et al. (2017) adopts softmax.
We prove that with appropriately distributed training data (Assumptions 3.2, 4.2), single-layer transformer trained via gradient descent will converge to its global minimizer at a linear rate (Theorems 3.3, 4.3) for both in-context binary or multi-class classification problems. To the best of our knowledge, we are the first to prove the training convergence of transformers for in-context multi-class classification. Moreover, our analysis reveals that the trained single-layer transformer can be viewed as approxi- mately implementing linear discriminant analysis (LDA).
In this work, similar to Zhang et al. (2023a); Wu et al. (2023), we set fam = I and define W (WK)TT W E (edexds, We use F to denote this simpli- fied model. Then, the output of F with an input embedding matrix E Rd,xd, can be expressed as
F(E:WV, wK9)= = E + E E. EWKQ
Due to the non-linearity of our loss function, we cannot directly find the closed-form expression of the global minimizer. Instead, we prove an important property that the global minimizer consists of constant plus error term that is induced by the finite training prompt length (N). We further show that the max norm of this error term is bounded, and converges to zero at rate of O(1/N). With properly distributed testing prompts (Assumptions 3.5, 4.4), we establish an upper bound of the inference error (defined in Equation (3)) of the trained transformer and quantify the impact of the training and testing prompt lengths this error. We further prove that when the lengths of training prompts (N) and testing prompts (M) approach infinity, this error converges to zero at a rate of O1/N + 1/VM), and the prediction of the trained transformer has an identical distribution to that of the ground-truth label (Theorems 3.6, 4.5).
(2)
P
In the following theoretical study and the subsequent experi- ments (Section 5.2), we show that this simplified transformer model has sufficient capability to approach the optimal clas- sifier for the in-context classification of Gaussian mixtures.
2.2. In-context learning framework
We adopt a framework for in-context learning similar to that used in Bai et al. (2024). Under this framework, the model receives a prompt P (uuuqy) comprising set of demonstrations DD ((x,, [e)) Lid. P and query Equery Pa, where P is the joint distribution of (z, y) and Pa is the marginal distribution of x. Here, E IR is an in-context example, and yi is the corresponding label for Tye For instance, in regression tasks, y/: E R is a scalar. In this paper, we focus on classification tasks. Thus, the range of ys can be any set containing e different elements, such as f1,...,c). for classification problems involving c classes. The objective is to generate output Mquery that approximates the target Iquery ~ P Pulsquer'
2. Preliminaries
c).,
Notations. We denote [n] = 21. .n.. For a matrix A E Rm , we denote its Frobenius norm as ||Al|,, and its max norm as ||A||max = max{e[m]]j{fr) |Ayll We use Aa,b (or Aat) to represent the element of matrix A the a-th row and b-th column, and use Ause,b represent vector a + 1 whose i-th element is A A(a+i-1).6- We denote the 12 norm of a vector as || l12. We denote the all-zero vector of size n as O,, and the all-zero matrix of size m X n as 0mxn. We use (xx) == 1/(1 + exp(-x))
Since Iquery is a discrete random variable, we use the total variation distance to measure the difference between Qquery and Jquery:
of dimension
A MVquerry Iquery)
sup
|IP (Mquery = 2) - (Nquery = 2) lb
to denote the sigmoid function. We define softmax(() :
26 R(Sqqur)))
R + (0, 1)*, and its i-th element as softmax(-),,. where softmax(x)) = exp(x;)/ exp(x;)).
where H(yquery)
is the range of yquery. When A (Vqueryy Jquery) = 0. Dquery has the same distribution as Vquery, which means the output of the model perfectly ap- proximates Yquery.
2

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Unlike standard supervised learning, each prompt P, can be sampled from different distribution P, in ICI. We say that a model has the ICL capability if it can approximate gr.query for a broad range of P, with fixed parameters.
Definition 3.1. We say a data pair (x, pd (uo 111 A) if y follows Bernoulli distribution with TP (y -1) P (v = 1) = 1/2 and f(zly -1) N(uo, A), f(x 1) = N(H1, A), where Hoo P1 E I' and A E Raxd is a positive definite matrix.
3. In-context binary classification
We consider the case of B training tasks indexed by r E [B]. Each training task T is associated with a prompt P. (Iy,1;
In this section, we study the learning dynamics of single- layer transformer for in-context binary classification. It is special case of the general multi-class classification. As result, the analysis is more concise. The general in-context multi-class classification problem is studied in Section 4.
Vr,N, 24 query) and a corresponding la- bel qreyy We make the following assumption in this sec- tion.
Assumption 3.2. For each learning task T E [B], we as- sume
We first introduce the prompt and the transformer struc- ture we will use for in-context binary classification. The prompt for in-context binary classification is denoted as
i.d.
(1) (Er, Mr,ali=1 and {Ir,querys Mr Mr.query (Mr,0, 1,
P (Ii1, y ,UN Xquery)) where I; R and yi E {-1,1}. We can convert this prompt P into its correspond- ing embedding matrix E(P) in the following form:
(2) Pr,,0 is randomly sampled from N(0, Ta), Mr,11
U7,A47,0 where UTAA = - -1/2, and U, is uni- formly distributed over the closed set of real unitary matrices such that U,U7 = Ld-
E= = E(P) =
x1 1/1
32
EN
Equery 0
(4)
We denote the distribution of (Pr.00 (4,11) as Pa (A)). Note that U.A = A 1/1,,AA 1/2 can be viewed as a linear trans- formation that preserves the inner product vectors in A-1. weighted norm, and we have pr,0AA ",,,0- -7,1 -1 "r,1 0.
3/2
UN
Similar to Huang et al. (2023); Wu et al. (2023); Ahn et al. (2024), we set some of the parameters in our model to or to simplify the optimization problem, and consider the parameters of our model (WV, WKK) in the following sparse form:
Let Tr.
(FF(((P));; W) ()(++))(++)) be the output - Tr.
of our transformer for task т. We define the empirical risk over B independent tasks as
WV (axd 0d) 0g
).).
0d) 0 0 0)
WW
WKQ_
(5)
B
where W E 12dxd We set the normalization factor equal to the length of the prompt N. Let F(E(P);; W' be the output matrix of the transformer. We then read out the bottom-right entry of the output matrix through a sigmoid function, and denote this output as Yout- The output Nout of the transformer with prompt P and parameters W can be expressed as
i(W)=
(1 + Urr,yerr)) (or(
2B
Urr log(1 (1 -rr,out.
(6)
Taking the limit of infinite training tasks B - oo, the expected training loss can be defined
1 L(W) = lim I(W) = - B[(1 + yr query) log (og(Tr,
Dout = ([F(E(P); W)LHDD)N+H)
B
+ (1 Vr.query) log(1 (r,,ur)],
N
T
W Wxquery)
=a
N
where the expectation is taken over (H us My,1) ) ~ Pa(A), i.i.d
(Ir,s Vr.a W (ur, Vr.query ~ pb Mr,,, A).
We denote the prediction of our model for Tquery as Jquery- which is random variable depending on Yout- Consider generating a random variable u uniformly on [0, 1]. If u < Taut. we output Squery = 1; if u > Dout. we out- put Jquery- = -1. Then, we have IP (Mquery = 1) = Hout. (Dquery = 11)
Applying gradient descent over the expected training loss (7), we have the following theorem.
Theorem 3.3. Under Assumption 3.2, the following state- ments hold.
Dout.
(1) Optimizing the training loss L(W) in Equation (7) with training prompt length N via gradient descent W++1 W -WW", we have that for any t 1,
3.1. Training procedure
We study the binary classification of two Gaussian mixtures and use the following definition.
|W'--*133 - exp(-t/k))II" W' (8)
- w
3

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
where po is the initial parameter and W is the global minimizer of L.W), and K. = 1/a. Here a, I are con- stants satisfying
T 2x 2x log(N: | -I* we have || mmax S 1/N Denoting W- = WT we have W = 2 =
2(A 2(A-1 + 3) where G = G+ HT /2, || S ||||| mmx + Imax = O(1/N). Thus, we have the following corol-
a S Amia (V2 S Xmax (Vnx((22((W) 54, (9)
lary.
where O: > 0, < o0, W E Rw; Rw- = ( E
Corollary 3.4. If we optimize the expected loss L(W) in Equation (7) via gradient descent with training prompt length N, initial parameters W and learning rate 1j = 1/1,
Raxd| || W7 || < 11 1l
(2) Define G 1 - 9 = I guery,, p = pr.1 pr,,0.
u 2(Hr,1 Mr,0), and = μ A 'qq for simplicity.
then, under Assumption 3.2, after T 2 2x (N | W=llp) steps, the updated model W satisfies
Then we have
| max
1 (Flo'(a))
+ 1 - AA-1991
W = 2(A + +G),
(1I)
N
+0"(w>ls
+
1 (u ((
"a==~1)²pp + =2 A "qpqq T M))
where ||G] [[lmmax = O(1/N) A = 1/aa, and ax are constants defined in (9).
8
+o(1/N);
(10)
Theorem 3.3 and Corollary 3.4 show that training single- layer transformer with properly distributed data (Assump- tion 3.2) for binary classification via gradient descent can linearly converge to its global minimum W* = 2((A GG). Furthermore, when the prompt length N grows, this global minimum W will converge to 2A 11 at rate of O(1/N),
where S = 4V-L(2A-1) I(2A-1)
so L(2A-1), d'() and d"(-) are the first- and second-order derivatives of (.), respectively, and the expectation is taken over (Hr,0 1T,1) (A), queyy Pt (Pr.u.Hr.1A)) A))
The detailed proof of Theorem 3.3 can be found in Appendix D. In the following, we provide a brief proof sketch to highlight the key ideas.
3.2. In-context inference
Next, we analyze the performance of the trained transformer in Equation (11) for in-context binary classification tasks. We make the following assumption.
Proof sketch for Theorem 3.3. As a first step, we prove in Lemma D.2 that the expected loss function (.W in Equa- tion (7) is strictly convex with respect to (w.r.t)) W and is strongly convex in any compact of R Moreover, we prove L(W) has one unique global minimizer W Since the loss function L(W) we consider is highly non- non-linear, we cannot directly find the closed-form expression of W*, as is often done in the prior literature. This poses a significant challenge to our analysis.
Assumption 3.5. For an in-context test prompt Prest (r1,3),,- 1i,.M,,.... 2M VMr Qquery), we assume
(I) [z MM Lid. ph (pos M1, A1, Equery E RRd. (2) Mo A A A 11-
With this assumption, for Vquery Pb yls (No, Ml, A), ac- cording to the Bayes' theorem, we have
We address this technical challenge via the following method. First, in Lemma D.3, by analyzing the Taylor expansion of L(W), we prove that as N co, our loss function L(W) converges to T(W) pointwisely (defined in Equation (25)), and the global minimizer W converges to 2A Thus, we denote W* = 2(A-1 + C), and prove || is bounded and scales as | = O(N-1/2) Next, in Lemma D.4, by further analyzing the Taylor expan- sion of the equation VL(W*)= = 0 at the point 2A -1, we establish tighter bound ||G|| = O(N-1). In Lemma D.5, we prove that our loss function is (-smooth and provide an upper bound for I. Thus, in compact set Rw, our loss function is a-strongly convex and (-smooth. Finally, lever- aging the standard results from the convex optimization, we prove Theorem 3.3.
P (Uquery = ([Iquery)
[ Mqueer = 1)IP (Vquery =
Exe(di) I (Iquery Bquery 2)P (Nqquery 2)
=0((w1 -mo)² A "Iquery).
If we test the trained transformer with parameters W in Equation (11) and Prest, by simple calculation, we have
Dout = o ( M y:
(A (x++G)azque) ). (12)
Intuitively, when the training prompt length N + 00, we have G - 0, and when the test prompt length M - oc, we have 2
According to Theorem 3.3, we have W = W' + H² where |H max S exp(-t/(2<))) - W'Il.. If we set
35 E. y:2.T + (11 - -Mo) T. Thus, when N, +
00, P (Uquery = 1) Dour -> a(11 - -Ho) To -xxqurr))
4

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
IP (Vquery querry), and the prediction of the trained trans- former Vquery perfectly matches with the distribution of the ground truth label Mquery-
requirement of A in training and testing had also been observed for in-context linear regression in Zhang al. (2023a). Here, we discuss the consequences when As- sumption 3.5 does not hold. For example, suppose the labels of our data in test prompts are not balanced where P (y = 1) = pr, (y = -1) = Po. Besides, Po: P1 do not have the same A A weighted norm, and the covariance ma- trix of test data satisfies I * A. Then, as N, M - oc, we
By analyzing the Taylor expansion of Yout at point ((11 ( ): we formally present the aforementioned intuition in the following theorem, which establishes an upper bound of the total variation distance between Jquery and Equery.
Iquerr));
have
Theorem 3.6. Consider a test prompt Prest satisfying As- sumption 3.5, and let Vquery ~ PP) Wll (Ho. P1,A), A). Let Jquery be the prediction of the trained transformer with pa- rameters W in Equation (11). Then, for the inference error defined in Equation (3), we have
M Σ M 1=1
2
you, - 2(p1 - -PoRt)
and
EE A (yquery Jquery)) )
IP (Tquery = 1) - 02(p1 - -Popto) A TA "query)
T A S <
10)
|| max
Σ
lmggy
On the other hand, the distribution of the ground truth
[EE(d]
label is IP (Uquery =) = ((H - HO) T r Equery + (N A 1o 21 - 1=)/2 + (og(p1/Po)), Define P (m1 - nO) T Equery + (MIA- |m Mg A > /2 + log(p1/Po) and & - 2(p1/11 Po/mn) T A Then, we can see that unless ; or (a( a=)) is suffi- ciently small, the transformer cannot correctly perform the in-context binary classification.
1/2 Øiو ]
+
1
A
2V2
Σ |M,3"
VM
2
VTT
SJEld]
1
+
+
N
V
M
where μ Pn - po. u- = 2(m + +a)) q == Equery and the expectation is taken over {z,, ye 1id. P' P1, A).
Remark 3.8. Another important insight of our analysis is that the pre-trained single-layer transformer can be viewed as approximately implementing linear discriminant analysis (LDA). For example, suppose we are given { 3 M and Tquery, and we need to predict the label Jquery for Tquery- LDA assumes that {z,, D: 3 and {T Vquery } are i.i.d. samples, with IP (リ 1) = (y: = - -1), and the con- ditional probability density functions f f = 1) and (a:lly 11) are Gaussian with means /41, 1t-1 and same covariance E. Under these assumptions, it can be derived that the optimal decision criterion for Tquery is to predict Equery = 1 if (μ1 - p-1) D'''xquen + _ [ 1 4-1 A E ' > O and Vquery = 11 otherwise. LDA can estimate jú as the average of x, with y/ = 1, estimate P-1 as the average of x; with y4 1. and estimate the covariance Σ from the within-class variances. For the single- layer transformer, it can compute the in-context estimate Á 7ー = x ΣM Vir,, however, it is hard for the single- layer transformer to estimate E in context. Thus, in our pa- per, we make the following assumptions (Assumptions 3.2 and 3.5). We assume the pre-train data and test data have the same covariance matrix A so that the transformer can learn an approximation of A during pre- pre-training. Moreover, we as- sume the two class means Po- 14 have the same A A-weighted norm so that μ Σ Ho /4 Σ "11 = 0. Under these as- sumptions, the quadratic term cancels out, and the estimated decision criterion simplifies to (L MM 9:2.)" T "A" --=xeuyey- which is very close to Equation (12) in our paper. When we use W to approximate 2A -1, the estimated decision
The proof of Theorem 3.6 can be found in Appendix E. Since ||x||max = O(L/N), Theorem 3.6 suggests that if we ignore the constants regarding Ho: /1, A, Aquer, the ex- pected total variation distance between Jqueey and Jquery is at most O(1/N- + M/) On the other hand, for data pair (I, Pd (Ho. the distribution of y. P (y- = 1/z) = ((1 - no) A-1x), can be characterized by a logistic regression model o(w" x + b) with param- Therefore, when N. + so. the prediction of the trained transformer is equivalent the optimal logistic regressor for binary classi- fication problems with distribution ph ((o,, A).
31
eters w = A (μ1 (4)) and b
Note that different from Assumption 3.2 which states that 1r,,0 1r Tuqquryy are sampled according to some specific distributions during training, Assumption 3.5 does not im- pose strong distributional constraints on Po- M4 and querry which shows the strong generalization ability of the trained transformer. Moreover, even if M + oo, the distribution variation between Vquery and Mquery does not disappear unless N + 00. Thus, the ICL ability of trained transformers for binary classification is limited by the finite length of train- ing prompts. Similar behaviors have also been observed in Zhang et al. (2023a) for in-context linear regression.
Remark 3.7. Theorem 3.6 requires Assumption 3.5 to hold. For example, we need the covariance matrix A in train- ing and testing to be the same. A similar consistency
5

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
criterion becomes exactly Equation (12). Therefore, when W = 2A and the in-context examples are balanced across classes, the transformer's 's decision criterion is the same as that of the LDA with exact knowledge of A. Our exper- iments also corroborate this theoretical findings. For ample, in Figure 3, since the pre-trained transformer has already learned a relatively good approximation of A- while LDA must estimate A in context, the trained trans- former significantly outperforms LDA when the number of in-context examples small. As the context length in- creases, LDA's performance approaches that of the trained transformer.
[0, 1]. If E E)ji((out) Ej=u()ou)s)) where (Nout) is the j-th element of Yout, we let Uquery e;- Thus, P (Fquery = e;) = (Nout)
4.1. Training procedure
We focus on the multi-class classification of Gaussian mix- tures and use the following definition.
Definition 4.1. We say data pair (I, v) ~ pm (u, A) if P (v e.) = 1/c and f(rly = e;) N(p;, A) for E [c]. where μ = (e)) E Rdxe and A E Rdxd is a positive definite matrix.
We consider the case of B training tasks indexed by r E [B]. Each training task T is associated with a prompt Pr (Iy, Vr.1....., Zy,N . Vr.N query) and corresponding la- bel Vr.qeery We make the following assumption in this sec- tion.
4. In-context multi-class classification
We now extend the study of the learning dynamics of a single-layer transformer to in-context multi-class classifi- cation, generalizing the results of the previous section. We will present the detailed formulation and then focus on the main differences to binary classification.
Assumption 4.2. For each learning task
[B],
sume:
We first introduce the prompt and the transformer struc- ture that will be used for in-context multi-class classi- fication. The prompt for in-context multi-class classi- fication involving c 2 2 classes can be expressed as P (I1, Si.yy...... NN UNN (qqerr), where 2; E Rd yss E (el, .2,..., ) and e; is the i-th standard unit vector of R. Its embedding matrix can be formulated
(1) ( Vr, LN and {ār,query, Vr, Uurqu) (Ur.1 (..
i.dd ~ pm
(2) 1r, sampled from N(O, Ta);
Ta), f7. = UTk,AA k
3, .3...... where UT,kAA = =A 1122 and Ur,k are uniformly distributed over the closed set of real unitary matrices such that U la.
We denote the distribution of fr as (A). Note that Ur..AA = A A1/2U,,A-1/2 can be viewed as lin- ear transformation that preserves the inner product of vectors in the A weighted norm, and we have AT Mr, = -T Mrg for i.j E [c]. Let Tr. = softmax ([F(E(P,); W W)
E E(P)
=
21 D2
EN
Equery 0g
(13)
y/
1/2
UN
Similar to the binary case, we set some of the parameters in our model as 0 and 1 simplify the optimization problem and consider the parameters of our model (WV, WKQ) in the following sparse form:
be the output of the transformer for task r. We define the empirical risk over B independent tasks as
WV (Oaxd 0dxc) Te
WKQ= =
W
Odxe
i(W)
1
ΣΣ
B c
(Ur (og(((7,out)k)) (15)
Ooxd
Ooxd
Oexc
(14)
where W Rdxdd We set the normalization factor p equal to the length of the prompt N. We read out the bottom-right c-dimensional column vector from the output matrix with softmax function as the output, denoted as Tout- With param- eters W and a prompt P = (#1, (1,,,1,-0 ,N,, UN, Tquery): ) the output can be expressed as
Taking the limit of infinite training tasks B co, the expected training loss can be defined as
L(W) = lim I(W) B
E
(yr.qu queryu)))) log( (urt
(16)
Hout =softmax (IF(E(P); WilH
=softmax ( 1
where the expectation is taken over Mr ~ (A),
N
yay
W Wreuery
{Ir,s br. {zy query Vr Aurry lid. pm (Mr,A).
Applying gradient descent over the expected training loss in
Equation (16), we have the following theorem.
We denote the prediction of the model for Equery as Jquery- which is a random variable depending on Tout Randomly sample a random variable u that is uniformly distributed on
Theorem 4.3. (Informal) Under Assumption 4.2, the follow- ing statements hold.
6

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Optimizing training loss L(W) in Equation (16) with training prompt length N via gradient descent W++1_ = W - L(W) for any 2 L, we have | - S exp((-\/<) - l (17)
If we test the trained transformer with parameters W in Equation (19) and prompt Ptest, by a simple calculation, we have
(( M
M Σ 1=1
where Wo is the initial parameter and W is the global minimizer of L(W), n 1/a. Here, a, I are constants
Jout = softmax
yi
(A (A-1+
+ G)equery
(20)
such that
a S Amin (2L(W)) I Xmaxx ( 11, (18)
Note that, when the training prompt length NN - 00, we
have G + 0, and when the test prompt length M + 00.
where a > 0, < oo, W E Rw, Rw = { E
we have M Ei= yar, - n== Thus, when N,M M + 0o, P (ïquery = Cx) - (Nout) + softmax(w' A A querr)) ) P (yquery = ek (qqurr)) i.e., the prediction of the trained transformer Jquery matches the ground truth label Uquuery
Rdxd| ||W - W*|| < ||10 - -I
(2) Defining G = /c/c- A-1, we have W = c(A-1+ G) and ||a[|max = Occ/N).
(3) After T 2 2k log || - (F) steps, denoting
By analyzing the Taylor expansion of Yout at point softmax(1' A quuury),, we crystallize the aforementioned intuition in the following theorem, which establishes an up- per bound of the total variation distance between Uquery and Jquery-
the updated model W satisfies
W W=dA-++0
(19)
where ||&]|ma O(c/N).
The formal statement and complete proof of Theorem 4.3 can be found in Appendix F. Technically, the proof of Theo- rem 4.3 builds on that of Theorem 3.3, but the more compli- cated cross terms in the Taylor expansions of the softmax functions, which are due to the nature of multi-class classi- fication, bring new challenges to the analysis. To address these issues, we derive new bounds on the expected errors of the cross terms in Lemma F.1, F.2, which may be of independent interest to other similar problems.
Theorem 4.5. (Informal) Let Prest satisfy Assumption 4.4 and Vquery ~ P Viisqum (H, Denote Nquery as the predic- tion of the trained transformer with parameter W in Equa- tion (19). Then, for the inference error defined in Equa- tion (3), we have
E [yqqurr,. = + 1112-11/2)
where the expectation is taken over M Lid. pm (n, A).
Theorem 4.3 shows that training a single-layer transformer with properly distributed data (Assumption 4.2) for in- context multi-class classification via gradient descent can linearly converge to its global minimum W = c(A + GG). When the prompt length N grows, this global minimum W will converge to cA at a rate of (c/N). Compared to the binary case, the new results establish the scaling behavior w.r.t. the number of classes c.
The formal statement and proof of Theorem 4.5 can be found in Appendix G. We can see that the convergence rate of the inference error in multi-class classification w.r.t. N and M is similar to that in the binary classification, except for the constant coefficient c. This suggests that classification tasks with more classes may have higher er- rors than those with fewer classes. On the other hand, for data pair (z; v) ~ pm (H,, A), the distribution of y, P (y = ex [z) = softmax(H" A -'=z),, can be characterized by softmax regression model softmax (WannaxxW=>>>> a: + b) with pa- rameters W pu A- and b = O. When N, - co, the prediction of the trained transformer is equivalent to the optimal softmax regressor for multi-class classification problems with distribution pm (14, A). Note that different from Assumption 4.2 which states that Jur; query are sam- pled according to some specific distributions during train- ing, Assumption 4.4 does not impose strong distributional constraints on μ or l'uery, which shows the strong general- ization ability of the trained transformer. We also discuss the consequences when Assumption 4.4 does not hold in Remark G.2, which highlights the necessity of Assumption 4.4. Moreover, even if M + oc, the distribution varia- tion between Jquery and Tquery does not disappear unless
4.2. In-context inference
Assumption 4.4. For an in-context test prompt Piest (x1, Vi....,......................... TMr VM, Tquery). we assume
(1) (Ii, yi M
L.d. (m, p = (1i,-.., E
Rdxe Iquery E Rd.
(2) M. A- u1 = 11 A- Mg. for for [c).
With this assumption, for Vqueery ~ 5) Rewery (H, A), according to the Bayes' theorem, we have
(Uquery = (x((equ)ry)
f (Jquery lyquery = ex (Yqquery ex) Σj=1 I (Jqueeyy liquery = e, (Uquerr = =e, =softmax(u]) T A-- -qquery)) k.
7

---

On the Training Convergence of Transformers for In-Contes Classification of Gaussian Mixtures
0.6
1-ayer, N=10
1-layer, c=2
0.7
3-layer, N=10
3-layer, c=2
0.5
1-layer, N=20
0.6
1-layer, c=4
3-layer, N=20
3-layer, c=4
Inference error
1-layer, N=40
0.5
1-layer, c=6
0.4
3-layer, N=40
Inference
0.4
3-layer. c=
0.3
0.3
0.2
0.2
0.1
0.1
0.0
40
Test prompt length (M)
Test prompt length (M)
(a) c 10
(b) 80
Figure 1. 'I-layer': single-layer transformer defined in Section 4, '3-layer': 3-layer transformers with softmax attention. N: training
prompt length. c: number of Gaussian mixtures.
NN - Thus, the ICL ability of the trained transformers for multi-class classification is limited by the finite length of training prompts. Similar behavior has also been observed Zhang et al. (2023a) for in-context linear regression and in Section 3.2 for in-context binary classification.
{Ur, i E [c]} have the same A weighted norm. In Re- mark 3.7, G.2, we also discuss the situation when Assump- tions 3.5, 4.4 do not hold and show the necessities of them. In this subsection, we consider training transformers with data of varying covariances A and with Gaussian component means unequal A A-1 weighted norms, and examine how these factors affect the ICL abilities of transformers. Results are shown in Figure 2.
5. Experiments
In this section, we report the experiment results on multi- layer, nonlinear transformers to investigate their similarities and differences to the single-layer, linear transformer we theoretically analyzed in the pervious sections. Detailed experimental settings and additional results can be found in Appendix H.
From Figure (a). we can see that both models perform better when their μor, have the same A A-1 weighted norm ('same norm'), ) however, in the 'different norms' setting, the performance of 'Illayer' deteriorates more significantly, while transformers with a more complex structure ('3-layer') show better robustness under this distribution shift. Similar situations also happen in Figure 2 (b), where '3-layer' also shows better tolerance to the covariance shifts than '1-layer'.
We train single-layer and multi-layer transformers for in- context classification of Gaussian mixtures with different numbers of Gaussian mixtures c, different lengths of train- ing prompts N, and test them with different test prompt lengths M. The results are reported Figure We can see that for both single-layer and multi-layer transformers, the inference errors decrease N and M increase, and they increase theoretical claims but also shows that, the simplified model we have studied indeed exhibits behavioral similarities to the more complex multi-layer, nonlinear transformers, and some of our observations for this simplified model also hold for more complex transformers.
Experimental results in Figure 2 show the necessities of As- sumptions 3.2, 4.2, 3.5, 4.4 for the single-layer transformers considered in this paper, and also demonstrate the better ro- bustness of multi-layer, nonlinear transformers. Developing better understanding of the robustness of more complex transformers is an intriguing direction for future research.
increases, which not only verifies our
5.2. Comparison of transformers with other ML algorithms
Additionally, we conduct experiments comparing the ICL performances of the transformers with other machine learn- ing algorithms for the classification of three Gaussian mix- tures. From Figure 3, we can see that all three transformer models significantly outperform the classical methods (soft- max regression, linear discriminant analysis), demonstrating the strong ICL capacities of transformers.
5.1. Varying covariances and norms
Note that in Assumption 3.2, 4.2, 3.5, 4.4, we assume that the covariance A during pre-traine pre-training and during inference
the same, and the means of all Gaussian components
8

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
0.30
0:35
1-layer, same norm
1-layer, same covariance
0.25
3-layer, same norm
0.30
3-laeer, same covariance
1-layer, different norms
1-layer, different covariances
Inference error
3-layer, different norms
Inference error
3-layer, different covariances
0.15
0.10
0.10
0.05
Test prompt length (M)
Test prompt length (M)
(a) Norm
(b) Covariance
Figure generate diag(Ai,........ .., )id). i {0, 1,,33, where Ag | and Xiς 1.d. N(3, 1). All models are trained with prompt length N = 100, tested with prompts satisfying Assumption 4.4 with Ag. c 3. (a) 'same norm': pre- pre-training data are sampled according Assumption 4.2 with Ag. 'different norms': For each r, with probability (l:
1/10, frra ~ N(k, J = 0, .,....5. 9. (b) 'same covariance': pre-training data sampled according Assumption 4.2 for the fixed Ao. 'different covariances': Sample additional A1, A2. Ag. Then, generate pre-training data according Assumption 4.2 with Ao, A2,
amined single-layer transformers with linear attention and sparse parameters. The learning dynamics of multi-layer transformers with nonlinear attention (e.g., softmax) for in- context classification problems remain an interesting area for future investigation.
0.18
1-layer, sparse
1-layer, full
0.16
3-layer
Inference error
0.14
softmax
LDA
0.12
Acknowledgements
0.10
The work of Wei Shen and Cong Shen was supported part by the US National Science Foundation (NSF) under awards ECCS-203367L, ECCS-2143559, CPS-23131ID 10, and CNS- 2002902. The work of Ruida Zhou was supported in part by NSF grants 2 2139304. 2146838 and the Army Research Laboratory grant under Cooperative Agreement W91 INF- 17-2-0196. The work of Jing Yang was supported part by NSF award ECCS-2531023.
0.08
0.00
0.04
Test prompt length (M)
Figure 3. '1-layer, sparse': single-layer transformer defined in Sec- tion l-layer, full': single-layer transformer with full parameters (59), '3-layer': a 3-layer transformer with softmax attention, 'soft- max': softmax regression, 'LDA': lincar discriminant analysis. All three transformers are trained with prompt length N = 100.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
6. Conclusion
We studied the learning dynamics of transformers for in- context classification of Gaussian mixtures, and showed that with properly distributed data, single-layer transformer trained via gradient descent converges to global mini- mum. Moreover, we established the upper bounds of the inference errors the trained transformers and discussed how the training and test prompt lengths influence the per- formance of the model. Experimental results also corrob- orated the theoretical claims. There are some directions worth further exploring. One potential avenue is to inves- tigate whether the assumptions regarding the training and test prompts can be relaxed. Additionally, we have only ex-
References
Ahn, K., Cheng, X., Daneshmand, H., and Sra, Trans- formers learn implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, What learning algorithm is in-context learn-

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
ing? investigations with linear models. arXiv preprint arXiv:2211.1566, 2022.
Giannou, A., Yang, L.,, Wang, T., Papailiopoulos, D., and Lee, J. D. How well can transformers emulate in-context newton's method? arXiv preprint XXXiv:4403.03188 2024.
Bai, Y., Chen, F., Wang, H.. Xiong. C., and Mei, S. Trans- formers statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural infor- mation processing systems, 36, 2024.
Guo, T., Hu, W., Mei, S., Wang. H., Xiong. C., Savarese, S., and Bai, Y. How do transformers learn in-context beyond simple functions? case study learning with representations. arXiv preprint arXiv:2310. 10616, 2023.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33; 1877-1901, 2020.
Hahn, M. and Goyal, N. A theory of emergent in-context learning as implicit structure induction. arXiv preprint 0rXiv:2303.0791 2023.
Bu, D., Huang, W., Han, A., Nitanda, A., Suzuki, T., Zhang, Q. and Wong, H.-S. Provably transformers harness multi- concept word semantics for efficient in-context learning. arXiv preprint arXiv:2411.02199, 2024.
Huang, Y., Cheng. Y., and Liang, In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Jelassi, S., Sander, M., and Li, Y. Vision transformers
provably learn spatial structure. Advances in Neural In- formation Processing Systems, 35:37822-37885 2022.
Bubeck, S. Convex optimization: Algorithms and complex- ity, 2015.
Karimi, H., Nutini, J., and Schmidt, M. Linear conver- gence of gradient and proximal -gaadiet methods under the wojsieweijaiiiwiiz condition. In Machine Learn- ing and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16, pp. 795- 811. Springer, 2016.
Chen, S., Sheen, H., Wang, T., and Yang, Z. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. /9442, 2024.
Cheng, X., Chen, Y., and Sra, S. Transformers implement functional gradient descent learn non-linear functions in context. arXiv preprint (oXiv:231.006520 2023.
Kim, J. and Suzuki, T. Transformers learn nonlinear fea- tures in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258, 2024.
Dasgupta, I., Lampinen, A. K., Chan, S. C., Creswell, A., Kumaran, D., McClelland, L., and Hill, Language models show human-like content effects on reasoning. arXiv preprint arXiv:2207.070555 2022.
Li, H., Wang, M., Liu, S., and Chen, P.-Y. theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.000/5 2023a.
Frei, S. and Vardi, G. Trained transformer classifiers gen- eralize and exhibit benign overfitting in-context. arXiv preprint ar 210:07440 2024.
Fu, D., Chen, T.-Q., Jia, R., and Sharan, V. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint 10.. 11.0./70000 17086, 2023.
Li, H., Wang, M., Lu, S., Cui, X., and Chen, P.-Y Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024a.
Gao, C., Cao, Y., Li, Z., He, Y., Wang, M., Liu, H., Klu- sowski, J., and Fan, J. Global convergence in training large-scale transformers. Advances in Neural Information Processing Systems, 37:29213-29288 2024.
Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. Transformers as algorithms: Generalization and stability in in-context learning. International Conference on Machine Learning, pp. 19565- 19594. PMLR, 2023b.
Garg, S., Tsipras, D., Liang. P. S., and Valiant, G. What can transformers learn in-context? a case study of sim- ple function classes. Advances in Neural Information Processing Systems, 1:30558-30598,, 2022.
Li, Y., Li, Y., and Risteski, A. How do transformers learn topic structure: Towards a mechanistic understanding. International Conference on Machine Learning, pp. 19689-19729, PMLR, 2023c.
Giannou, A., Rajput, S., Sohn, J.-y., Lee, K., Lee, J. D., and Papailiopoulos, D. Looped transformers as pro- grammable computers. International Conference on Machine Learning, pp. 11398-11442. PMLR, 2023.
Li, Z., Cao. Y., Gao, C., He, Y., Liu, H., Klusowski, J. M., Fan, J., and Wang, M. One-layer transformer provably learns one-nearest neighbor in context. arXiv preprint arXiv:241 1. 2024b.
10

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Lin, L., Bai, Y., and Mei, S. Transformers decision makers: Provable in-context reinforcement learning via supervised pretraining, arXiv preprint arXiv:2310.0856, 2023.
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.0760 2022.
Lin, Z. and Lee, K. Dual operating modes of in-context learning. arXiv preprint 2rXiv2240.118819 2024.
Wu, J., Zou, D., Chen, Z., Braverman, V., Gu, Q. and Bartlett, P. L. How many pretraining tasks needed for in-context learning of linear regression? arXiv preprint 2023.
Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprini arXiv:2307. 03.66., 2023.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit Bayesian inference. arXiv preprint arXiv:2111.000
Nichani, E., Damian, A., and Lee, J. D. How transform- ers learn causal structure with gradient descent. arXiv preprint wwx.i::2422114755 2024.
2021.
Yang, T., Huang, Y., Liang, Y., and Chi, Y. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv:2408.100 101477 2024.
Nye, M., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D., Lewkowyzz, A., Bosma, M., Luan, D., et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2 /2.001 14, 2021.
Zhang, C., Meng, X., and Cao, Y. Transformer learns optimal variable selection group-sparse classification. arXiv preprint arXiv:2504 08638, 2025.
OpenAl. GPT-4 technical report, 2023.
Zhang, R., Frei, S., and Bartlett, L. Trained trans- formers learn linear models in-context. arXiv preprint arXiv:2306. 0et.22366.9997, 2023a.
Pathak, R., Sen, R., Kong, W., and Das, A. Transformers can optimally learn regression mixture models. arXiv preprint arXiv:2311 1.08622, 2023.
Zhang, R., Wu, J., and Bartlett, L. In-context learning a linear transformer block: benefits of the MLP com- ponent and one-step GD initialization. arXiv preprint wwXh:2240.149551 2024.
Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. Transformers as support vector machines. arXiv preprint Stfi:2300.16896 2023.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Ilama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. OPT:
Open pre-trained transformer language models. arXiv preprint wwxiv:2205.01065 2022.
Zhang, Y., Zhang, F., Yang, Z., and Wang, Z. What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. /9420 2023b.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention all you need. Advances in neural information processing systems, 30, 2017.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp- 35151-35174 PMLR, 2023.
Wang, X., Zhu, W., Saxon, M., Steyvers, M., and Wang, W. Large language models are implicitly topic mod- els: Explaining and finding good demonstrations for in- context learning. In Workshop Efficient Systems for Foundation Models@ ICML2023, 2023.
Wang. Z. Wei, S., Hsu. D., and Lee, J. D. Transform- ers provably learn sparse token selection while fully- connected nets cannot. In International Conference on Machine Learning, pp. $1854-51912. PMLR, 2024.

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Appendix
The Appendix is organized as follows. In Section A, we provide a literature review of the related works that studied the ICL abilitiies of transformers. In Section B. we introduce the additional notations for the proofs in the Appendix. In Section C, we introduce some useful Lemmas we adopt from previous literature. In Sections D. E. F, G. we present the proofs of Theorem 3.3, 3.6, 4.3, 4.5 respectively. In Section H, we provide additional results and details of our experiments.
A. Related work
has been observed that transformer-basee models have impressive ICL abilities in natural language processing (Brown et al., 2020; Nye et al., 2021; Wei et al., 2022; Dasgupta et al., 2022; Zhang et al., 2022). Garg et al. (2022) first initiated the study of the ICL abilities of transformers in a mathematical framework and they empirically showed that transformers can in-context learn linear regression, two-layer ReLU networks, and decision trees. Subsequently, numerous works have been developed to explain the ICL capacities of transformers in solving in-context mathematical problems. These works mainly use two approaches: constructing specific transformers capable of performing certain in-context learning tasks, and studying the training dynamics of transformers for such tasks.
Constructions of transformers. Akyürek et al. (2022); Von Oswald et al. (2023) showed by construction that multi-layer transformers can be viewed as multiple steps of gradient descent for linear regression. Akyürek et al. (2022) also showed that constructed transformers can implement closed-form ridge regression. Guo et al. (2023) showed that constructed transformers can perform in-context learning with representations. Bai et al. (2024) proved that constructed transformers can perform various statistical machine learning algorithms through in-context gradient descent and showed that constructed transformers can perform in-context model selection. Lin et al. (2023) demonstrated that constructed transformers can approximate several in-context reinforcement learning algorithms. Fu et al. (2023); Giannou et al. (2024) further proved that constructed transformers can perform higher-order optimization algorithms like Newton's 's method. Pathak et al. (2023) showed that transformers can learn mixtures of linear regressions. Giannou et al. (2023) proved that looped transformers that can emulate various in-context learning algorithms. Cheng et al. (2023) showed that transformers can perform functional gradient descent for learning non-linear funetions in context. Zhang et al. (2024) showed that a linear attention layer followed by linear layer can learn and encode a mean signal vector for in-context linear regression.
Training dynamics of transformers. Mahankali et al. (2023); Ahn et al. (2024) proved that the global minimizer of the in-context learning loss of linear transformer can be equivalently viewed as one-step preconditioned gradient descent for linear regression. Zhang et al. (2023a) proved the convergence of gradient flow on a single-layer linear transformer and discussed how training and test prompt length will influence the prediction error of transformers for linear regression. Huang et al. (2023) proved the convergence of gradient descent on single-layer transformer with softmax attention with certain orthogonality assumptions on the data features. Li et al. (2023c) showed that trained transformers learn topic structure. Wu et al. (2023) analyzed the task complexity bound for pretraining single-layer linear transformers on in-context linear regression tasks. Tarzanagh et al. (2023) built the connections between single-layer transformers and support vector machines (SVMs), Nichani et al. (2024) showed that transformers trained via gradient descent can learn causal structure. Chen et al. (2024) proved the convergence of gradient flow on a melti-head softmax attention model for in-context multi-task linear regression. Kim & Suzuki (2024); Yang et al. (2024) proved that trained transformers can learn nonlinear features in context. Bu et al. (2024) studied the training convergence of transformer for tasks encoded with multiple cross-concept semanties. Li et al. (2024b) proved that a one-layer transformer can be trained to learn one-nearest neighbor for binary classification in context. Li et al. (2024a) studied the training dynamics of single layer transformer for in-context classification problems. However, they only studied the binary classification tasks with finite patterns. They generated their data as x = P4 + KVk, where {,,) are in-domain-relevane patterns and {2x J-Ma are in-domain inreevvannee pattems,, M, M2 and these patterns are all pairwise orthogonal. Thus, the possible distribution of their data is finite and highly limited. In contrast, our work explores the ICL capabilities of transformers for both binary and multi-class classification of Gaussian mixtures. Specifically, our data is drawn according to P'(Ho, 141, A) or pm (n,, A). and the range and possible distributions of our data are infinite. Recently, (Frei & Vardi, 2024) also studied the the ICL of a linear transformer model for classifying Gaussian mixtures and showed that trained transformers can exhibit benign overfitting in-context. However, their analysis relies on additional assumptions, such as a sufficiently large hggal--t-noiss ratio. In contrast, our work does not rely on such assumptions. Additionally, compared to (Frei & Vardi, 2024), we consider more general multi-class setting. We believe these distinctions highlight the independent contributions of our work.
Some works also studied the ICL from other perspectives. To name a few, Xie et al. (2021) explained the ICL as implicit
12

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Bayesian inference; Wang al. (2023) explained the LLMs latent variable models; Zhang et al. (2023b) explained the ICL abilities of transformers implicitly implementing Bayesian model averaging algorithm; and Li et al. (2023b) studied the generalization and stability of the ICL abilities of transformers. Hahn & Goyal (2023) showed that ICL can arise through recombination of compositional structure found in linguistic data. Lin & Lee (2024) studied the dual operating modes for in-context classification of Gaussian mixtures. However, the analyses of both Hahn & Goyal (2023) and Lin & Lee (2024) are based on the idealized optimal next-token predictor for ICL tasks and did not discuss the training dynamics of transformer models.
Apart from studies focusing on the training convergence properties of transformers in-context learning (ICL), many other works have investigated the convergence behavior of transformers when learning different tasks. For example, (Jelassi et al., 2022) and (Li et al., 2023a) theoretically analyzed how Vision Transformers learn and converge. (Wang et al., 2024) showed that transformers can provably learn sparse token selection, whereas fully-connected networks cannot. (Zhang et al., 2025) demonstrated that transformers can learn optimal variable selection in group-sparse classification. (Gao et al., 2024) studied the global convergence of large-scale transformers under gradient flow dynamics.
B. Additional notations
We denote X ~ Bin(n, p) if random variable X follows the binomial distribution with parameters N and p nd
[0, 1], which means IP ((X- = k) = =(n-x);Pp (1 - p)* We denote X ~ Multin(n. p) if random variables X = X2,... Xx) follow the Multinomial distribution with parameters e N and p1 P2 means P ((X- = (I1, F2, "2.,,,,......... xx)) = I m2! h kn.
p3 1/k, which _ We denote (((x)) = softmax(I); = ( exp(x;)) for simplicity. We define 6, = 1, 8:g- = 0, i + j. For N, we define (x)) = [(x - 1)/dj +1,12 = ((x- ((x-1) 1) mod d) + 1.
nd
C. Useful lemmas Lemma C.1 (Karimi et al., 2016)). If : - R is p-strongly convex, then
[(x) - min f(x) 2 ~ー~ 11 - x113
where x²' = argmin, fxx).
Lemma C.2 ((Bubeck, 2015)). Suppose f: Rd + E is c-sttrongly convex and 3-smooth for some 0 < S. B. Then, the gradient descent iterating w = w² 1/ (wo" with learning rule 11 1/3 and initialization w E Rd satisfies that for any 2 1,
|| - w0 S exp(-(/<>)] - w
where K = =B/ox is the condition number of f. and w² = argmin f(w) is the minimizer of f.
D. Training procedure for in-context binary classification
In this section, we present the proof of Theorem 3.3.
D.1. Proof sketch
First, prove in Lemma D.2 that the expected loss function L(W) in (7) is strictly convex w.r.t. W and is strongly convex in compact set of dxd, Moreover, we prove L(W) has one unique global minimizer W*, Then, in Lemma D.3, by analyzing the Taylor expansion of L(W), we prove that as N + oc, our loss function L(W) point wisely converges to L(W) (defined in (25)), and the global minimizer T converge to 2A We denote W = 2(A 1 + G), and prove || max -1/2). Next, in Lemma D.4, by further analyzing the Taylor expansion of the equation VL(W*) = 0 at the point 2A -1, we establish a tighter bound || O(N-)). In Lemma D.5, we prove that our loss function is I-smooth and provide an upper bound for I. Thus, in a compact set Rw, our loss function is a-strongly convex and (-smooth. Finally, leveraging the standard results from convex optimization, we prove Theorem 3.3 in subsection D.4.
In this section, we use the following notations.
13

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
D.2. Notations
Recall the expected loss function (7)
1 I(W)-
[1 + Vr.query) log(Ur.out) + (1 Vr.query) log(1 Fr,out)
(21)
2
where
Trtt = (( Vr T
W query
the output of the transformer, and the label of the data follows the distribution
IP (Ur.query 1 queery) o((Hr,1 (+.00" A TA IIrqqeery)).
In this section, we introduce the following notations to analyze (7). We denote Mr.1 Pr.0- = Pr,1, /r.,0 and
q = Ly query- Then with probability query = 1) = 1/2 we have = μi +u, and with probability IP (yrr.query = 0) = 1/2 we have 140 + v, where v N(O, A). We define
N Li=1 Vr,jiiir,i- Since with probability IP (Vr,i = 1) 1/2 we have T = + Us and with probability P (yr,) = 0) = 1/2 we have œ Hf + U/ where U, N(O, we known
2
p 2N pm /N 2NopoNN + g. where g = 2 Σ> V/ g N(O, 4A/N), Ni (in,N,1, 1/2). Defining / 1/2.
2( + no). we have No/N = 1/2 h and
p= hu +
(22)
Then, the expected loss function (7) can be expressed as
L(W)= E[-o(p" - log(o(p Wg/2)) (1 olp A 9)) log(1 o(p T Wg/2))]
(23)
The gradient of the loss function (7) can be expressed as
1 VL(W) (lo(pp Wg/2) - ol4 A - "))pq' A-)g))
(24)
Moreover, define function T.WW)=.
I(w)= = E|- A A 1- 4)) log(a(n ( Wg)))) (1- a( T A- 9)) log(1 o(u T W Wq/2))].
(25)
In Lemma D.3, we show that as N + co, L(W) will point wisely converge to L(W).
D.3. Lemmas
Lemma D.1. Suppose N₁ ~ NN,( Defining = N1/N - 1/2, we have
[[A] = 0
B[A?] 1
4N
[[AA] = O(N 2), for n 2 4
[[M]] S 2N1/2
[[43)] O(N-3/2)
Proof. Since N1 Bin(N, 1/2), the moment-generating function of Ni is N (1 2 exp(t)
MN. (t)
14

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
We compute the moment-generating function of h follows:
N
N
(1(1) = exp
(-2
MN1
=
exp
2N
exP
2N
cosh
N
2
2N
N
12
8
22r
8N2
(2
Thus, we know the coefficients of t, 2,, 6 are 0, 1/(8N), respectively, and the coefficients of t", n 2 4 are O(1/N3). We have
E[h] 0
E[h²] 1 4N
| E[h" = O11/N),, for n 2
Moreover, according to the Jensen's incquality, have
|[[A|] & (B(4²)) 1/2
2N1/2
B| < (B(14))3/4
3/4 3323)
Lemma D.2. For the loss function L(W) (7), have L(W( > O. For any compact set Rw of [dddd, when W e Rw we have 2 L(W) > :Iaafor some 7 >> O. Additionally, L(W) has one unique global minimizer on Rdxd
For L(W) defined in (25), also have VLL(W) O. For any compact Rw of Raddd when W E Rw, have (2T(W) > Ia for some > O. Additionally, T(W)) has one unique global minimizer on Rddxd
Proof. We vectorize W Vec(W) RRd where Vec(W); Wol0s20)) 4(x) [(x ((1d] 1,12(1) ((x 1) mod d) + 1. Then, we have
(VL())) Ep. [1 [lllTTWa/2) Wq/2) olu TA - (-10)10(00000))
(26)
The Hessian matrix of the loss function (7)
1 ol Wq/2)(1 - a(p W (122))) (V2L(W))s = Epg
Considering od such that + 0, we have
zn ta Wq/2)) TVLL(W)= =By.p 1 ol Wq/2)(1-o03 (()
Σ
ab
= oo o(p Twq22)
2n SaPt, (12aa)
fm (p. 9)dppq.
RE
where (p. q) are the probability density function (PDF) function of p. Since for any p.g. o(p ((TWW//2 Wg/2)(1 aep (Wg/))) > we have 2 T V(L) 2 0. Thus, LW) ł 0 and L(W) is convex.
15

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Moreover, for any + O, we denote 2zj =((1-1)4+)); E [d]. Suppose a, a arg max4,j j lzgg/e we consider a set of constants {Clpi: (ppi), ; (Cigi, C2p.), i [d], where Clpa d, C2pa d 1, Clgb d, C2gb d + 1, and cıpi
1/16, C2pi 1/8, + Clas = 1/16, C2aj 1/8,j + b. Then, for any Cpi E [Clpis C2pil, Coj [Crgi C2ail. We have
Σ
- 2(d ( - (d 1)2 max |zell 2 dl max 22sl/2.
CijCpiCoj
FEEL
Then, we define region l(a, b) _ E, Spili E, Cpppi [Cipi <2pi), Caj E [Cigj (2nr)
We have
2
min ((ab))
Σ nE(P")
(cP(1
2 d max 12y 2//2//2/ 1|211//4.
Defining
/ ac|d] (4] M(a.)
C(Q)
min
(p 2)dpdq.
SS2, W) = min min aE[d] M(a.))
we have SS2, W) > 0. Since with probability P (U+.query = 1) = 1/2, 1= = # with probability TP (Vr.query = 0) = 1/2, q= no + v, where v ~ N(O, A) and p + hu +g where
N(0, 4A/N), N(O, A), Ho ~ N(0, La), the covariance matrices of p. are positive definite and we have (p.(P.q) q) > 0 for all p. 9 E Rd. Moreover, Ma b) are non-zero measures on Rdddd Thus, we have ((1))= > Then, for any = + O, we have
Σ 2Phh ) fpalD. 9)ppdq
2 TV=L(W)>> 2 /
1 L01D Wq/2) op T Wq/2))
((a,)) (0())((2, 13//4//44
>0.
Thus, we have V2L(W) > 0. L(W) is strictly convex.
Moreover, for any compact set Rw of dxd for any W E Rw, we have
Wq/2)(1 olp Tw(/)))) > 0. min min WQ/TW///2
S(?)
= min we Rw Le(d ((a
Then, for any W E Rw, for any + O, we have
Σ aan ) frolP, 9)ppdq
z LVW)2 2 / M(aM)
towo Wq/2)(1 Wg/2))
((1)s
Thus, when W E Rw, where Rw is a compact set, we have V2L(W) > ((1))5()))444 and the loss function L(W) is y-strongly convex, where ((2))
Because our loss function is strictly convex Rdxd has at most one global minimizer in Rdxd Next, we prove all level sets of our loss function are compact, i.e. Va (W e d L(W) < a} is compact for all a. We prove by contradiction. Suppose V. is not compact for some a. Since our loss function is continuous and convex, V, is unbounded convex set. Since the dimension of Va is d², consider point Wa e Va. there must exists W + Oaxd such that
{Wo + [W* |t= ]o)} E V.. For this W * Oaxd, there must exist set of constants O < C3pi < Capi, 0 < C3gj < Cagi such that for any Cyp [Capss Capul, Caj E [C3gh Crail, we have
Σ 1
CpcCoj W + O.
16

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Thus, have
lim Σ
cp tW = 20.
We define Qg {p- Σ; Spseis == Ej Cyjees cps E [espe Capel cas E [C3gg Cagil, 11al12 S E, caps + +Zgg]. Then, defining
/ Jo (P q)dpdg,
C(Qa)
((2) min {min(o(n) T A (1-1-o(m Ro
we have S(D)) > Since Pg are non-zero measures for p, q. we have C(Qo) > Then, we have
lim L(Wa + +
1-000
lim E[-o(1 A 9) log(o(p" (W" (W*)//2) (1 ol A A) log(1 olp (W + +/W*)//2))
2 lim |-o A x log( CpiCaj W u3 + / 22)] (v. 4)dpd4
1-000 So
lim
T [-(1 oF A 11=1 loggllg11111111111 - o E SpsCoi (WS tw x 22)] 9)dppq
Da
2C()))
min
lim | log a Σ 1+00
CpiCog (W + #
+C(()))S(Lo)
min
lim [ log(1
0 Σ
CpiCoi (W + +1W)
Ro
1000
=00
This contradicts the assumption L(W" + ++**)= < a. Thus, all level sets of the loss function L(W) are compact, which means there exists a global minimizer for L(W). Together with the fact that L(W) is strictly convex, L(W) has one unique global minimizer on Raxd
Similarly, we can prove the same conclusions for L(W).
Lemma D.3. Denoting the global minimizer of the loss function (7) as W have W 2(A-1 G), where |['|Imax O(N-1/22)
Proof. Let a = = -A 8= p Wg/2,r (hu g) WW/2. Performing the Taylor expansion on (7), we have
L(W)- =E (
E (==tt))oooooo))---------=--1110- (1-
((
E
(1
+ E [
(()) + E -
where r)) are real numbers between and According Lemma D.I. we have E [r] E [(uu + g) [Wq/]] 0.
Thus, have
[(((()(1-0(s)) - - -(-))))))) Enag (((()) (s)) (1 (s))) Ey [r]] 0.
17

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Moreover, have
E (
< [r2]
= =E(h-u Wqq Wq
lequ-aq--t + 9 Wag Wal
1 Wqu Wa/(4N) + 4(AW T Wg/N]
[ /N,
where (a) is due Lemma D.1, 9 W qg" Wq Eij.eei 9e W. Wygg Warg Exs.ieei Ban Weigh Waya (gg) Wq) T Wq and Elgg" 4A/N. Ct is a constant independent of N and W. Thus, we have
[[w) - L(W) - C/IW /N.
This shows that L(W) point wisely converges to I(W). According to Lemma D.2, L(W) has one unique global minimizer. Consider the equation:
vī(W) E[o(( Wq/2) o T - A'' O.
We easily find that vī(2A-1) 0 and W = 2A -1 is the global minimizer of L(W).
Considering a compact set Rw {W ||W - 2A-11|1 pw), J. we have || mmax S Ow for W E Rw. Here pw, Cw are some positive finite constants. Then, we have
L(W) L(W) & C/N, W E Rw,
where Cj C,CW is constant independent of N and W. This shows that, for W E Rw, our loss function L(W) uniformly converge I(W).
Denote W as the global minimizer of the loss function L(W) with prompt length N. Then, we show that, when N is sufficiently large, W* E Rw. first denote oRw = { |[ 2A || = pw) and A = minwearw L(W)
L(2A-1) > O. Then, for N 2 4 and for any W E Rtw, we have
L(W) - L(W) < A/4,
This means
min WEORW
L(W) min L(W)
WERW
2 min WEORW
I(W) I.(2A-) 1)
Σ(W) Z(2A-1) - 2/2
min WEORW
A/2 1/220 0.
Since L(W) is strictly convex, we have W argming L(W)& E Rw-
Then, we have
[(W**) - L(W")) C//N
[[22)-)) -(2A-))) CJ/N
I(W*) < L(W") L(2A-1) C//N S <(2A-))+2C//
18

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
According to Lemma D.2, for W Rw. have (WI > yld. where y is positive constant independent of N. Thus, L(W) 7-strongly convex in Rw. According to Lemma C.I, we have
L(2A | -1 2 2 -
4C² ?N
Thus, when N oo, we have W 2A-1 Denoting W 2(A 1 + G). we have || O(L/VN).
Lemma D.4. The global minimizer of the loss function is) W 22(A G), where
||G]|mux I N "loss T + un uu A /4)
( T + 29 A " max + o(1/N),
= - == 122L(2
Proof. According Lemma D.2, the loss function L(W) has unique global minimizer W We have
VL(W") E [1r(p] W W* o(u² T A -- ")) 0. Let W* = =(A-=((A+++))) 2() T A ---,-th,b b- (u + luu + g) Gq + (lau + 9) A A-g.. We have
(27)
p² W' 9/q
=( + hu 9) T)T(())))) 3A-
=(u+ + hu + g) Gq + (hu + TA-1 -
The Taylor expansion of aaa b) at point with an Lagrange form of remainder is
aaa + b)pd = aa))a (a)bppp + o"(")
p²pqq" + o (E((E(a b))
13 pd
2
3!
where E(a,b) b) are real numbers between a and a + b. Thus, our equation (27) become
E Ey.u.ghh. d'(a)bogT o"(a) 6² b²pqT pq + o" (E(a, b))
6² 0.
(28)
3!
Note that [[o'(a)bpqT] = Epnig [o) [bpqT] T]]. For (bpq J. according to Lemma D.1 and g N(O, 4A/,4 we have
B [bpq '1
"quqq + hhuu 1-2uqTGquqTT Blu Gqua T + T 166 +g GQgq +h2
TA A- qqqq
up Gqq T 49q /N 4AGqg /N uu Gqg /(4N) uu." T. A A1 199 /(4N).
(29)
Then, we have
lo (a)() T /N + uu G ((4N))]] S |max//,,
where 61 maxij Σ Ebr kr 4a' (a) (Auxqiq;) En (a) (uauqiq;/44i constant independent N. According Lemma D.3, || max = 0(1 N) = o(1), we have
E g l (4) /N uu Gqq /(4N)) o(1/N),
(30)
19

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Similarly for (lo"(a))²1)) T /2], have
Bg (0p Elu Gqp" Gqua" uuuu u Gquq" g
Gag +
"quu" T G + 2h2 uhqu Gquq + 2g Gqu "Gq9
(4)
222h2u "al''qu" A quqq + 22 Gqg T A qnp + 22h2uu Gqu T A 'uqqqq" + 2p T Gqg A 999 1
(16)
E Elh T A qu 1
qug" T + TA - q9 T A qpq"
(618)
For each term in (i). it contains two G. Thus, their max norms are at most smaller than max)) For each term in (wi), contains one G and 1t2 contains one G and two g. According to E[ 1/(4N) in Lemma D.1, the max norm of terms with one G and h are smaller than O[]]G]|maxx/N) Defining
we have N(0, Id) and g 2N-1/2 //2g. Thus, converting two g to g. we have a coefficient of N-1_ Therefore, the max norms of terms with one G' and two are also smaller than O max/ Maa Therefore, for terms (f), (it), we have
/ 101/2-11////
(a))"( max I O Imax) = of G maxx,
(31)
(a(a"(a)(t)) max I [[ m /N) = o
(32)
For term (ii)), according Lemma D.I and g N(O, 4A/N), we have
B (i)) 22
o (()))(h²u A -= A A + T A==
"gg TA A
a=-qqqq)) ll max
(33)
E To" (a) 9) //
/8+ + 29 A apu" )] mm
(34)
For ([(" b))bb))))ppp))))))))))) /3!, we have
||E(o" b))b pq T //3]]lmm max
< max sER
la (=)l//3 /3. max E pxpll
max <0(1)
Σ
NOT Gqoz Gqog Gypeaj
Σ
| Gqoz GqØs A- qpeb)
11.01 (p. (hu.g)
Σ
( GqoJ 19933 A - 9
(s0e)
A
900 A -
900 A 9p:4j
$1.0 {bw.g)
(#we)
For terms in (+) containing two or three G, these terms" expected absolute values are at most smaller than [[ 2
For terms in (**) containing one G, these terms must contain n, number of h and n2 number of elements of g. where
71 + n2 = 2,3, 4, n1, n2 E N. According to Lemma D.1, we know that for n; = 1,2,3, 4, E| I O(N-m12) Defining g = /2 1/2 /2, we have ğ ~ N(O, Ia) and g = - 1/2 1/2g. Converting g to g. we have coefficient of 22 Thus, for terms in (**), these terms' expected absolute values are at most smaller than O[I] Imaa N (1 S
20

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
[ mmxx I N- For terms in (* *) without G, these terms must contain #1 number of and ng number of elements of g. have n1 n2 3, 4, 11, n2 N. Similarly, these term's expected absolute values are most smaller than 2 O(N-3/2) Therefore, we have
[E(a, b))b"pp pq /3 max
max E ||6 p:g; max (o'' (z)1/33
=[[ =o[[G]]|max + o(1/N).
((ax)) max) mmxx NN O(1/N 8/2)
(35)
Moreover, we have
lo "aunTGuaT]]
Σ kI
SislaGRI,
(36)
where stylı (d) We vectorize G as Vec(G), (1, (110)))0(0 (100 Then (36) can be expressed as
Define S xd where Stj
(Euvlo'(a) (a)pp' Gqq Gqq } SG.
(37)
Note that S = 42=1(2 According to Lemma D.2, is positive definite. Thus, combining (28), (29), (30), (31), (32), (34), (35), (37), we have
G max
N S-- (B)o((a)(AqaT< + wu 1A'''q- + +o"(a)(( ((=())()T>>))>>) + 2q -1 1=''quq4) max
+ o11/N)
Lemma D.5. The loss function (7) is I-smooth, where < \ Eieta (((
Proof. The Hessian matrix of the loss function
LV2 1
Considering 2 c such that + 0, we have
L(W)= =
Za Pa (p W ( Σ
ab
=E
op (1 (W/22)
(a)
aE|d?
SE
CaPt
aE
(a)
||
sE[d?
where (a) is due to the incquality. Thus, V2L(W) = HL and L(W) is L-smooth, where a constant
smaller than 1 Eiejayy (())())))))))))
21

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
D.4. Proof of Theorem 3.3
Proof. According Lemma D.4, the global minimizer of L(W) is W 2(A 11A1-4 G), where
G max I N Il (Elo'(a)
+
1
uu T A
4
+ +0(( ( A 1-=))/9//001 nq" + +24 A ' "))) (11NN)
(38)
- W I || w } Rw is compact set. Then, according Lemma D.2, for W E Rw, we have L(W) ala. Here > 0 is a positive constant number. Thus, L(W) is a-strongly convex in Rw Moreover, according to Lemma D.5, L(W) is I-smooth. Then according to Lemma C.2, applying gradient descent with 7 = 1/1, for any t 2 1, we have
Define Rw { |
| W 13 I exp( t/6) || -WW-- 12
where & = l/a.
E. In-context inference of binary classification
E.1. Notations
In this section,
the following notations. We denote μ μ사 Po. 2( + Mo), 9 Iquery- Define
p 2/ M Ei=1 ye Since with probability P (y) 1) = 1/2, x, = f1 + 2s with probability (y = 0) 1/2, 2, Mo + t, where U, ~ A), have p 2M P1/M M- 2Mouo/M 9. where g 3/2 E %/ N(0, 4A/M),
M1 ~ 1/n(,,1/2) Defining h = MI/N - 1/2 u 2(11 po), we have Mo/N = 1122 h and
p μ hu g.
(39)
E.2. Proof of Theorem 3.6
Proof. The output of the trained transformer is
((3 M Σ
yout
you
(a=+G)))raaa = =
(A-1 (A-+++)))))))
(40)
The probability of Vquery
= 1 given Equery is
IP (yquery = 1 [ = =0((u1 - -o) T A- xq-'eyy o A-19).
Defining 1r A
(n+ + hu + g)TTT Gg + (hu+ + (g) AA-1) 'g. we have
p (A +G4
=(1 + hu g) T(A(A= -A-1+
g) = + hu 9) T Gg + (hu + g) - 1 + T T - + b.
and
E (A-- = [a((a b)] (B((() + o'( o" ()((a 2)))2/22 /2).
where& & are real numbers between a and a + b. Thus, we have
Bloo(+ + - a(a)l -(())l
( 6))6² /2||
[a)(E)][b]] Ebb3]
22

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
We first consider the term ('(4) Defining y /22 have
o' E'(b]
(a((a)
1 Ggl + 1i [[]] Ggll Bilg [g]] - - + Ellg
A '=1
(0)
(a)
# TGgll= 2M1/2
Iu T Ggl +
2
E
G// +
1
lu
A
ggl
+
2
Blgg" A - 1
M2/2
2M
M2/2
1/3 q3 |
(a)
|1 max
|mggg + M1/2
|
2v2
Σ
A
+
1
V
VM
ijEld
where (a) is due to B[]A]] s 1/(2M1/22 in Lemma D.1. (b) is because that 9: N(0, 1) and E = v2/ /T, for i E [d].
For E|b²), we have
B6b S
B [/ hu + TG] + [[hu + g)) >>-1 + 2E (u hu + 9) Tq(hu + +9) + A
Notice that terms in E [(u + hu + 9) g)ggG)G])
Gg]²] contain two G. Thus, they are at most smaller than O | O(1/N3).
|13 1 /2 contain two h, or two 9, or one h and one g. According to Lemma D.I, we have Terms in E [(hu + T) A
E||h|] O(L/VM), M), E[h²] 1/(4M). Moreover, g = 2M 2/21 //2g. Converting one g g. we have coefficient of g. one h and one g are O(1/M). Terms in E (u+hu++)) (alluugggTGGlli (g(hu+gg) - A-19 contain least one G and one h or one G and one g. Thus, they are most smaller than Therefore, we have E 011 MM) = =01/N+1/MM
M 1 Thus, terms in E [[(hu + A 4/² /2 contain two h,
O(I max// MMM =
Finally, we have
E (Mquery Jquery)]
B[[Tout
TP (Yquery [Iquery)
Ello(a b) - a(aall S (a) B(b*]
<a)(a)
||
max
|pggil
1
A 11-1g +
2v2
Σ
|
1/2
+o N
VM
11/2
F. Training procedure for in-context multi-class classification In this section, we present the proof of Theorem 4.3.
F.1. Proof sketch
First, we prove in Lemma F.3 that the expected loss function L(W) (16) is strictly convex w.r.t W and is strongly convex in compact set of Rdxd, Moreover, we prove L(W) has one unique global minimizer W*. Then, in Lemma F.4, by analyzing the Taylor expansion of L(W), we prove that as N + oc, our loss function L(W) point wisely converges to I(W) (defined in (44)), and the global minimizer W converge 2 Thus, denote W 2(A G), and prove ||G O(N-1/4) Next, in Lemma F.5, by further analyzing the Taylor expansion of the equation VL(W*) = 0 at the point 2A 11, we establish a tighter bound ||
OccN-1). In Lemma F.6, we prove that our loss function is !-smooth and provide an upper bound for I. Thus, in a compact set Rw, our loss function is a-strongly convex and (-smooth. Finally, leveraging the standard results from the convex optimization, we prove Theorem 4.3 in subsection F.3.
In this section, we use the following notations.
23

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
F.2. Notations
Recall the expected loss function (16) is
L(W) =
(rr query & log log((oo,oooaoooool (ir )
(41)
where
(Ur,out): softmax (2
( N
E
Wa Wurrqurry
i=1
the output of the transformer, and the label the data follows the distribution
(Ur.query ex (e(eryy) softmax (n, A Aalp:_AAA--aaaa "Frrqqurry))
In this section, we introduce the following notations to analyze (16). We denote Mr = Mr = (21, ,2....,, E Rdxe and g r.quurr.. Then with probability (Ur.query = ex) = 1/c, = Pe + u, where v ~ N(O,A). We define Pk N Ei=((1 (Br,)) Idd and P P21.P2.....PP. Pe) ) Raxe We have PT N Ei= yor E Rexd Since with probability P (V+.1 ex) = 1/c we have Ir,i = 14 + 04, where cNeuu/N + Ok. where = (He(rae. 1/ ~ N(O, N&A/N22 and (N1, ,,, ~ Multin(n, 1/c). Defining hig = N2/N - 1/c, we have Nx/N = 1/c + h2 and pk = μ4 + chylt + Jk- Defining 9b. A-1 gs. we have OK N(O, e NaN"). Defining 1h = (h1 hyg......, hxjls) E Rdxe and g = (g), 22.09...... . 9k) Rdxo we have p cμb + g.
N(O,A), we known Ph = N E (3r.))
Then, the expected loss function (16) can be expressed as
L(W)=
E
softmax(j1
4)x T A- log(softmaxP" Wg/c)))
(42)
The gradient of the loss function (16) can be expressed as
VL(W)= = E
[ssottmx(PP (/a//)) softmax(p" TA -11 /d]
(43)
Moreover, we define function I(W)+s as
I(W)= = E
softmax(x* T A A-'g),) ")x, log(soft T Wg/c)a)l
(44)
In Lemma F.4, we show that as N + xo, L LWW will point wisely converge to i(W). Lemma F.I. Suppose (N1, No,..., No)
(2,...,,N))
~ Multin(N, 1/c). Defining b N 1 1/c, we have
B|ha]
Elhil 1
N
c 2²
B[gghg]
1 $j
Ne
1 O (N-3) -), Σ nhp 23
k=1
||[h 5 1/22 -1/2 1/2
E[[h,h;]] =[[h:,,]]=[[[[]]]
B[hghthhl]] = 0 3N-3
Blhbhhy/hhul1 0 (2)
where i,j. k, E [c].
24

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Proof. Since (N1, Ng...., Ne) Multin(N,,/o) the moment-generating function of N2,.N....... No) is
Mx(1)= = ( i=1
N
exp(ts)
We can compute the moment -ennertignn function of h (h1, hez....... hee as follows:
Ma(1) = exp
- Σ
4/e
Mn((/N)= =
exp
1
11 -
Σ
1
N
1=1
1=1
j=1
2
+ Ne
Σ
t - Σ
ty/c
+
1
Σ
c
t, -
Σ
c
ty/c
2N2c
i=1
j=l
1=1
j=1
Σ
t;:
Σ
ty/c
k
h=3
j=1
1
1/12-1/2)1f - Σ
2N lity
+
Σ
ti -
Σ
t
2N
KINKc
*jjlc
k=3
,=1
Observing the coefficients of h, we have
E [[hs] 0
1
1 1
TE
[M
N
c 2²)
1
[/gh
=
N i $
E
] - 0 (N- Σ
nu 2 3,
k
where i.j. k E [c).
Iteratively applying the Hölder's inequality, we have
E [ > (E() 1/2 = N 1/2 -1100/1111
1/2)1/2
B[ <
||,3/ hy / |11/3/110010011111 = O (A-3122
Elhishghahil] < BiI || 1| 0
where i.j.k, 1 E [c].
Lemma F.2. Suppose Gk- ~ N(O, NNA/NN) and (Ni, N2,,N,,..... .2,...,,N ~ MNl 1fc), define Ih = A gb and
N+/N- = 1/c + hy, we have
B(( -0
([())<(()))] = S
E( ):. (0ks) ), (Ti.)) )ial = 0
B(()) = E|3e2 N + = Elh... (gx) /
hl =0
25

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
where i, 81, 12, E (d 1 m, Ki, higs hig [c]. For any nik. niki satisfying Exelq Rik
Ekelelaejd
12a4 2, we have
E
h =
ke kE|d]
Moreover, we huve
E| (Tx) S N-11111
B (T)) - ((x)) O( -
where E i E [c].
For any 1ik. n2ki satisfying Ekeld Aık + ExElepeed n2ki n, n 2, 4, we have
B kE(c]
(Tx) [ O(N
Proof. Since gh N(0, and 9h
N(O. (c/N &haa/)aaa,, we have
Bl( =
[((x)) Swsdecc/1
[[ ). (0ka) )iz (Ting )ia) =
((()) B(3e² 2 cha = El (gx): (5)3s)
O(N-2)
he =
where i,5, [d],k, m, k1, k2, ka E [c]. Thus, with the results from Lemma F.I, for any Tike 7224 satisfying Enele) M1k + EKE(CheeEd N2ki 1,2,3, 3, we have
kE [],
Moreover, according to the Jensen's inequality, we have
E| [ | < E (0x)) 1/1/2
E| (7x+) a $ < ((( 3//4 ||3/4
O(N
where E [d], A E [c]. Thus, with the results from Lemma F.1, for any ni.. now satisfying Exe nix+ Eke(o) 12x4 n, n = 1,2, 3, 4. we have
Ih (gk) PaRe
S
II kE(c].iE(M]
"[|b|/////B[( 112//// = O(N2
RE [[]
Lemma F.3. For the loss function L(W) (16), have V2L(W)> > O. For any compact set Rw, when W Rw, we have L(L(W) > aa for some O. Additionally, L(W) hus one unique global minimizer on Iadxd
For L(W) deftned in (44), also have V2i(W) > O. For any compact set Rw, when W RRw. we have v2Î(W) 7Ia for some 7 > Additionally, I.(W) has one unique global minimizer on Rdxd
26

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Proof. We vectorize W Vee(W) IRd where Vee(W); W610.000 (xx L(x 1)/dj 1,,2(2) ((s
1) mod d) Then, we have
(VL(W)) = (softmax(P Wq/c)))
T - (0000p0))))00000000
(45)
Note that
softmax(P = a(a))
Vsoftmax Wg/c), = o(aa)) Val.
where - -log( Iith exp ((D (p:)-//)))))) For Vau. we have
lllk exp ((p: Ph) P Wg/c) (px pi)q" /c
Vak-
Ei=t al7k exp ((p - pa) (Wq/c)
Σ Ei=l,...c,44 exp (pi W q//)) (px -pi)a// /c
Ei exp (p) Wg/c)
Then we have
(enttmaxP"TWq//ee =softmax(PP Wagce) Σi Ei-immniy eiyk exp (p) Wq) (px - -iia /c
En=l.... exp (p,, Wg/c)
Σ
softmax(P P T Wofomasof (e (p/c)(Pkk - pi)a /cc
and
(Vsoftmax(P" PPo/)ha
Σ
softmax(P
Wo/c)ssottaxxp poati
We express the Hessian matrix of the loss function with the following form:
(VDL(W));
Ssthri//WWw///sttttt/////////hhh -D)aciaann 1=11=1....,,.44
=E
PT (ot)s (ecetmaxx(TT///((ee PP - p))
Considering ddd such that + 0, we have
T V(LL()> =E
softmax(P) Pfttmax(P"W//)eell)l (P//))sootmax((P Wg/c); Σ
Zazolpu Po) (px (0)
ab
k-1
=E
softmax(P" Wg/c) PTT Wg/c)
za(pk - (a)s Σ
aE(d
Since for any P. g. k,,, softmax(P" (t/pp/W////a///e W//ccusoftmax((P Wq/c); > 0. we have V2L(W)> 2 0. Thus, V2L(W) 2 and (.W is convex.
27

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Defining P1 P2, we have
L(W)=
27 za - (2)e
E
softmaxx Wg/c))softmax(P Wg/c/
Σ aE|d4
softmax(P PT W (eatmaa(""Wg//issooeei ///c)
Wg/c)2
Σ aE(19)
ZaPt (a)
(P
where (P,9), the PDF function of P. For any * 0, denote Zij ((li-1)d+() suppose a, b E arg maxi,j |2il we consider set of constants {Cipi: C2pi ), (Cige, i. E [d]. where Clpa = d, 02pa = d + 1, Clgb d, C2gb d d 1, and Cips = 1/16, Caps = 1/8, # Cigf = 1/16, C2gs = 1/8,j # b. Then, for any cps E [Cipt C2pil Cus E [eiqs) C2gil, we have
Σ (JE(d
(d (d-1)d 1)2/64 max lay 2 d max 12,y1/2. [d - 1)dd 1)/8 -
if
Then, we define region (((aa))) b) = [pp- E, Cpiti q q=E, E3 Cy [cipé. Oppll. Cgs [Cigs c205l. |P| (g }.
c
O2 +
We have
min 0(a,b)
Σ [E[da]
2 d max |za 2/42 1
17
Defining
/
C( = min aE bE(d) ((a
fPa(P. g)dPd
S(2, W/
min min DE bE (d] ((ah)
softmax(P W (aatnax(PPW/////fffftfioo htop///oooaa((PPPaaaaaaa Wafc))
we have S(2,, W) 0. Since have (P. for all P.g g and M(a, b) are non-zero measures for P.g. q. Thus, we have C(R) > 0. Then, for any O, we have
2 L(W)z
2 fPo(P. 9ddPd
21 21,,, (0 (E(UF)
Σ
M(a
(() W |) 0
Thus, we have V LWL > L(W) is strictly convex.
Moreover, for any compact set Rw of Iddd for any W Rw, have
S(2)=
min min WeRt at DE|4 R(a.))
min
soltmax soltiss Wa/c)s 0.
Then, for any W E Rw, for any + O, we have
z L(W222
2
Poftmax(PTWa//)ssoo) PPWa///
Wa//o)
Σ 2 (r (E(da)
((abb) 2
28

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Thus, when W Rw, Bw is compact set, we have VLL(W) (()))4 our loss function is y strongly convex, where 7 (S)))
Because our loss function is strictly convex in Rdxd, it has at most one global minimizer in Rdxd, Next, we prove all level sets of our loss function are compact, i.e. Va {W E æd L(W) I a} is compact for all a. We prove it by contradiction. Suppose Va not compact for some a. Since our loss function is continuous and convex, Va is an unbounded convex set. Since the dimension of Va is d,, consider a point W Va. there must exists W of Oaxa that C3pi < Cipis C3qj < Cigs
{W + [0, [o)] Va. For this W* * Odxd, there must exist of constants [C3gh. Cigil. have
such that for any Cyp [Capss Ctpal: Cas
Σ f
Cp W #
Thus, have
lim Σ 1-00
CpcCog (W + :W = o0.
We define Gn = Li=Eiom q= E3 ei Gpt E [Capt Capil Caj E [Cags Crasl, 1| S <3 Sppi + (ag
e² (E Egp (.gj)]. Then, defining
C(Qo) [ (ra(P.g)Pul
5(Qb) min fo
{ min (softmaama (waeee////////////o Wg/e),, (aaotma Wg/c)2)}
have S(Qo) > Since Qn non-zero measures for P, have C(Q.) 0. Then, have
lim L(Wa + :W* t-00
lim
softmax((' T A 1-1)))) log(softmax(( (Wa + (/W))//))))
t-000
2 lim
softmax(H1 T A 4-) log(softmax(P" (F""://"""""""""))))/p/pppp.pppp
t+00
Bo
lim
[-softmax(j" A (Wwp://twn///T----/o///////oooo/
1-000 So
lim 1000
[-softmax(H" T A 9)) log(o (W"
(W" + (
lim
softmaxx A 9)2 log(o(-7 -
(Wa t + /
10000
2(()))
min
lim log(o Σ CprCoy +tW3)/)
Bg
+C())) min
lim log(o(- Cpicos (W]] [V+))(((
This contradicts the assumption L(W" + +WW) < a. Thus, all level sets of the loss function I((W) are compact, which means there exists a global minimizer for L(W). Together with the fact that L(W) is strictly convex, L(W) has one unique global minimizer on ddxd
Similarly, we can prove the same conclusions for L(W).
Lemma F.4. Denoting the global minimizer of our loss function (16) as W*, we have W = c(A -1 + G), where || max ONN11/1)
29

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Proof. Let p A Wg/cr +)) T Wg/c. on me A A Wa/e. (chaul +g) W Wa/c. Performing the Taylor expansion (16), have
L(W) =E
Ca(a) (o(o))og((((s+r1))
=E
(a(a) log((a(s))
(a(a)Ru((sr))
=L(W) E
(e(
where rs)( S sup,
@ (og(Ca(u)) 0y
supy y (z(0)
0Cr(2) Dy
sup, 1ó ((()) 1. Thus, we have
L(W)- ((W)-L(W)) L(W)
E ||m|]]
1=1
cE [Muip]] Wa] + E [19] Wal]
1=1
<0(
Im + max (or):
< N-1/2
where the last inequality is due to Lemma F.1, F.2. Ct is constant independent of N and W. This shows that L(W) point wisely converge to I(W). According to Lemma D.2, L(W) has one unique global minimizer. Considering the equation:
vī(W)
softmax(p' A A-'9) log(softmax(( T Wq/c)&)]
We can easily find that Vi(cA-1) = 0 and W cA-1 1 is the global minimizer of T(W)..
Considering compact set Rw {W] ||W 22 - pw3. we have WIlmax S Cw for any W Rw. Here pw, Ow some positive finite constants. Then, we have
L(w) L(W) 1/2-1/2 W Rw
where C/ = CrCw is a constant independent of N and W. This shows that, for any W E Rw, L(W) uniformly converge to T.(W).
Denote W as the global minimizer of L(W) with prompt length N. Then, we show that, when N is sufficiently large, W Rw. We first denote dRw { ||W - cA - pw) A -minwearw I(W) L(cA-) )1) 0. Then, for
N 2 (4C/ and for any W Rw, we have
((w)-1A L(W) A/4
min L(W)- min L(W) > min
L(W) - L(cA-1) A/2
W ERRW
WeRw
W WeORw
Since T.W) strictly convex, we have W argminw L(W) E Rw.
30

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Then, have
[T(W*) - I(
I(W*) S L(W*)+ C//N & L(CA CI/N & +C/ 1/2
According to Lemma D.2, for W Rw, have V2T(W) > yla. where y is positive constant independent of N. Thus, L(W) 7-strongly convex in Rw. According Lemma C.I, we have
2 (T(W*)- 1(cA-1))
4C²
|| cA 1
N21/2
Thus, when N oo, we have W cA !! Denoting W c(A -1 G). we have || |mxx O(4-1/4).
Lemma F.5. The global minimizer of the loss function (16) is W c(A 11 + G). We have
0ça(a)
(Ca(a) caq" dak
||]]
S
(cōsa 1) 1-'"9"
N
da:
h111
(C((a)
Sc((a)
(cōin - 1)14 A "un A
quaq /2 +
o4 T A "uunn" /2
daidan
dag
k,,m==1
k1=1
+ o(1NN),
where = - a 1o A A S. S (V Ignoring constants other than N, we have ||G max O(c/N).
Proof. According to Lemma F.3, the loss function L(W) has unique global minimizer W We have
VL(W") = E
[CA celu A 9 /c] = 0.
(46)
k=
Let W' c(A + G). TA au
- b- (u+ cun g) Gq (cun - by (ue + chapk + Gq Tx) + (ch+/4 TA A1g.. The Taylor expansion of (a + b) at point a is
T)
c OCela)
8 (a bzby//2
b) = (a) +
Σ I,,m=11
Retio (a, byb
Day
daydan
where R (a, b)1 sup,
(c(())
Thus, our equation (46) become
aCa(a)
g(((a) bybnpaq T 2221 +
E
bipea
Σ
Ritno (a, by T /33
(47)
day
Daydan
k.hn=1
k.
31

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
For the first term Ek
dccla) dai
bipky T, according to Lemma F.1, have
e
(a(
E
bipua
da
(a)
E
[m] Gquaq" e Chhap Gquaq" Chuhary A-
a "quuq"" gi A A''qgag' ai Gaged |
Oau
kl==1
(a)(a) (Hany) Gqq" + (cōka 1)Paet Gqq G/ (cdεt 1)pm A- 2p94 /N)
=E
dap
(a)
(aqq T /N + cAGqu (N)
(48)
Ook
According to Lemma F4, [[I]]||max) = OONN o(1), we have
(Cx(a)
0cxla) dak
E
(cdws 1) Gqq /N
cAGqq /N
k,1=1
maxx
50( /N) = o(1/N)
(49)
(c((a) bebnpan /21, we have
For the second term Exd.n==
Baydan
(c((a)
E
bnbnpua /2!
Dajdan
1
(a)
E
Σ
oe Gqoz Gapua'
2
dajdan
A,,,
po.
(6)
Σ
2 A
A-1 9pp9d
di E/piccuuuu www.eeeooooee.........eoo fchh
(i)
Σ
o A 900 A-1 A''pppa
(184)
For terms (i) having two G, their max norms are at most smaller than G For terms (ii) having one G, define
Øi = 1/2 gt. these terms must contain nı number of hy and n22 number of (ūy);, we have Eje(o N1j + n2js n, n4 = 2.3. According Lemma F.2, we know that for The 1,2,3,
E II
1
]E(C),EE[
Thus, the max norm of expectations of terms in (ii) are at most smaller than O[I Therefore, for terms (1), (M))
we have
E [+]]] max I O[ = o
(50)
[[() mmx <
((1NN
(51)
32

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
For terms (iii) without G, we have
[(ii]]
PC((aa) Chhbaa N "qup A
{Ga(a) Oaf
E
".pTA--uu)] "qua"""" /2+
- ~''99 A A''que /2
da,dan
k
k11=1
(G((a) Bajdak
8 (c(aa) 6 Nuhhnhaa A -- r'qn - "ggugg /2
chapi A 1-'11pp- 99g A - "qgba"
da,dan
k.,m=1
max
(C(( (cdin - 1)ll A "u, "uuuq"
(C((a)) duf
E
+
o4 T A - "'qua""
2N
uu,dan
k,Lm=1
k1==1
+ ((1N22)
(52)
where the last inequity is due to Lemma F.1, F.2.
For the third term Lx Ratam (a, byb /31, we have
E
Σ
Ratuom (a, bybbbabb /3!
k.l.mm=11
<0(1)
max E
Σ
bb bkg (pk. )19ml
(mmed [d]
20(1)E
Σ
Σ
GqØz Gqoz Gq(pk, 1mmm
bn ka El(ee
LoEE (11xg when May ane E (Mig chen Mag 32gg dds E (1ig chin (4kg (91g)
(*)
o Gqoz Gooz A "((kk 11m
E (Mey ,che, Sey tute (uagiehh why Sig AASE (chh Wlkg :
(+)
Gqo A - 903 A A (9p.
Eipe, ,ch,, Mh, .h, d. E fehky Mbg 2 l. May
(x44)
+
Σ
TA 902
900 A -1-----l 1g(. 1mm
di Elechh, Bb, E Mag 2ag J. Jo (che 9kg
(*0**)
For terms in (+) having two or three G, these terms' expected absolute values at most smaller than [(1 For terms in (**) having one G. these terms must contain nı number of hy and n2jj number of (ū));, we have Eye(chee|4 M1j +2jj na, = 2, 3, 4. According to Lemma F.2, for nu 2,3, 4, we have
E
II
(T)) R234 = ONN
jE(c]aeld
Thus, these term's expected absolute values are at most smaller than O[I ||a [ N-1.. For terms in (* + without G. these terms must contain R1 number of hy and n2fs number of (ū,),. we have Eje(s]ie Aiy + n2js = 72g 72 = 3, 4. According to Lemma F22, for n = 3, 4, we have
Ih 123 = -1/ O(N-3/3)
jE|c]aE[d]
33

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Thus, these term's expected absolute values smaller than O Therefore, we have
E
Rkinm (a, b)bbbr PmP /3!
m.l.n.ml
max
<0(1) max
El
E
bbx, bug (pka)iqml
(ME
Na Ra
<0[ mmaz)
O[I maa N- ONN-3/12
+ o(1/N).
(53)
Moreover, we have
dCa(a) day
Gqq T }
Papi
k,1=1
(1 CC((a)uuuun Gqq
Σ
(a(a)Giaa)umn) Gqa
k Roptt
(o( T
(p) - mi) -1) Gqgq
1m1
Sijnm Gram:
(54)
where Sijnm = E Ex= Ei= ( -
- p)ngmos. We vectorize G as Vec(G), = (z, Define
E Pf where Sij = (110)00000 = Ei- Σt (e(a)G((a)(H) - -M) - -4)1
(54) can be expressed
(c)(a) Mada Gqq SG.
(55)
Da,
Note that S c(A"I(cA-1) According to Lemma F.3, positive definite, Thus, combining (47), (48), (49), (50), (51),
(52), (53), (55), we have
| max
0Ckla)
(a)(a) ogq
E
(cdst - 1)pupT A- "q
+
N
day
Oaz
k11=1
1=1
(a) (cōin - 1) A "qu,T""" quay" 1/2
Σ
oq A "undd" /2
tardan
da?
k.
k./=1
max
+ o11NN)
Ignoring constants other than c, N, we have ||G]|max S {(ccNN).
Lemma F.6. The loss function (7) I-smooth, where S 1 Ei Et-1
Eieter
E[((pE (ta)))))
34

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Proof. The Hessian matrix of the loss function is
c
(V-)(W)) = E
coftmax(P"W///)eee weeppP/W////eee Wg/c)ssoltaxx(1 Warp////ee Wg/c)((p= - Pouoos - mmnoiaas
Considering 2 R& such that + O, we have
2
L(W)2
k-1
1
=E
softmax " W
Wg/C)ssoftmax( (PT (W////)
Za za(Px - (1) aE|dP)
Σ
2=
(0)
E[((p> (1))
2=
1=1
where (a) is due to the Cauchy-Schwarz Schwarz inequality. Thus, L() N Ul and L(W) is I-smooth, where / is a constant
smaller than Ei_ Ei Eietey B[((pk ())))))(())))
Theorem F.7 (Formal statement of Theorem 4.3), The following statements hold.
(1) Optimizing training loss L(W) (16) with training prompt length N via gradient descent W++ W -VLWW),((W), we have for any t
| - w' S <exp(-</s)lWW- W'
where 8 is the initial parameter and W the global minimizer of L(W), k l/a. a, are constants such that
0<a S Amia (V2 L(W)) & Amax V22((W) s 1, for all W E Rw,
(56)
where Rw (W Rdd IIW W < I - w 1}}.
(2) Denoting W -1--1+ + G), we have
1
0Cala) day
0Gala)
| max S
S-1E
(cōμt 1)paml A- 199 T+
+
eqq
N
Sax
(Ca(aa)
(cōin 1)14 A -1 A -1 A'qugq
1
Σ
0 (a)
cq T A --
2
Dardan
2
dat
k.I.n=1
k11=1
+ ((L/N)
Oc/NN)
where (V2((2--1) 1(2A-1) limm- L(2A The expectation is taken over Py P(L(A) Ir.query
Pm ())
(3) After T 2 2k log(N | W*llF) gradient steps, denoting W as the final model, we have
W c(A + +G)
(57)
where || Occ/N).
35

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
F.3. Proof of Theorem 4.3 Proof. According Lemma F.5, the global minimizer of L(W) is W c(A - + G), where
||] max
aCc(a)
(cōn D
(a) eqq
E
1)paMT A-' +
dai
Oax
k=1
(a
80 (x
(cdin - 1)MT - AA,, "u
- pubba +
Σ
oq
T A 9uud A-'quu /2
da,dan
duj
k,
k11=1
max
+ o(//N).
Ignoring constants other than c, N, we have ||]] mmax S ((cNN).
Define fw (W E Rdxd || - W < | w F)] and Rw is a compact set. Then, according to Lemma F.3, for W Rw. we have LL ł alg. Here a > 0 a positive constant number. Thus, L(W) is a-strongly convex in Rw. Moreover, according to Lemma F.6, L(W) is I-smooth. Then according Lemma C.2, applying gradient descent with 1) 1/l, for any 2 have
I - W* 1 I exp(-1/n) - W'
where & = l/a..
After T 2 2k 2ogl | - -W*l gradient steps, we have W = WT = + H /c)
+ GG), where -((A-1
2(A-++G),
G = G G HT HHTH///, /c, H mmax exp((TT/s)- - 1/N. Thus, || max S |G] max H max O(c/N). n
G. In-context inference of multi-class classification
G.1. Notations
e. In this section, we use the following notations. We denote = (11 (2,..., Equeey- Define 1 E (vi) and define P P2,,P2,...,,38
Raxe, We have PT A), we have Pk Ex (Vr,)aar cMauu/M + gk. where 9h = (ie(i,,. U4 9k ~ N(O, M,A// and (M1, MM,,M,.... M2,.... Mc) ~ Multin(M,1/c), Defining hu
ΣM yio Rexd, Since with probability i- MR/M- 1/c, we have Ms/M = 1/c + hy and pk = M4 chypA + gk-
M
P (yr.s ex) 1/c, Izz M4 + vi. where
Theorem G.1 (Formal statement of Theorem 4.5). Let Jquery be the prediction of the trained transformer with parameters
W in (19) and Prest satisfying Assumption 4.4, and let yquery ~ pmm we have
Faaary (pt, A). Then, for the inference error defined in (3),
E [[Aayurry:: lquery)]
max kEle]
(ça(a)
G max
Σ
[V(1-1/0lm1 A- 1
Ve Σ 1A,71/2 A.jE|d]
Dai
M1/2
1
+ +
VM
where a = TTA -
as -TA 9. The expectation is taken over {2,, M Lig pm pm A).
G.2. Proof of Theorem 4.5
Proof. The output of the trained transformer is
Hout- softmax ((s) L yoar
(A (A- + @)rquerr
softmax (A-1 + +Ggg)
(58)
36

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
The probability of lqueey = C2 given J'uuery i8
P (yquery
e/ [Querr) softmax(4" T A -=1xeuury)) ) softmax((1 T A 19)x
TA Defining a - - (n+ 4) Eq+ + A A- au 11 A -, bu - (IN charue + gn) "Gg + (chause +
9h) A.'A-'g 9, we have
c DGo(a)
softmax PT P (A-) + +)))a b]] +
by +
Σ
Rus (a, b)b/bn/2)
day
1m==1
where Ruin (a, b]] sup=
(x)
Thus, have
aCk(a) Da
El - SE
bi +
Rath (a, b)b.ba. /2
1,70l
We first consider the term IE
1==
(C((a) Da:
by
Defining ği A A-1 we have
(Gx(a)
E
bt
Das
[Ca(a)
(IuT Ggl + Bllchuui Gyll Bl197 [g B(cch144A A-1 /6 + El19/11 A 'rll)
day
(Gr(a)
vol1-116 1 1/c)
ccl-1/c)
|m Ggl
Ggl
Elg
A
G
111
gg Bllg A -1/22/1)
Oay
M1/2
M1/2
(C)(a)
||G] max
Σ
VC(L-1/c)
Vc Σ
day
M1/2
(JE(d]
<JE(d
N
VM
where (a) is due to Lemma F.1 that B[[h]] & M-11/2-112(1 - 1/2 (b) 110)0/ (g) < [[ (c/ for / [c], E [d].
because that Ør ~ N(O, M/T/MM)
Eim== ReIn (a, b)brb,, //]
For
E
we have
Ruin b)b:bn//
O(LE
Σ
Σ
o GGo Gg
1,=1
[(n=1
(iE(Michuuroo) J o2 ( cha
Σ
20 Goo A-1
Σ
ToTA T0TA-'g0 900 A A-19)
d Ear, dos E am
di (cE on) d. (chnnt ond
(16)
(1i)
For terms (i) having two G. they are at most smaller than O max) O(1/22). For terms (ii) having one G. these terms must contain n1 number of hy and naji number of (ūj)is we have Σ Ejetc]eeja M15 n2ja ne, ne 1,2. According to Lemma F22, we know that for ne 1,2,
E
h (g,) m2ye = OMM-113
jE jE(c)aE
37

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
Thus, terms in (ii) are most smaller than |l]/"llmaxx////// must contain mij number of h and naji number of (Øz)ı, have Σ jE(c],KE[d] Ri5 + naji Re, Re According to Lemma F.2, for 2, we have
M- = (01/(N For terms (iii) without G. these terms
II
10% (g));
1 = -
- O )
selc),e[d
Thus, these term are (MM Therefore, we have E [i Ruin (a, b)bbbmn = O(1/N22 + 1// 1/(NVM))
o1/N I/V
Finally, we have
E E[A (yquery: ]query]]
max((Elsofftmax(a + b)x softmax(a))l]3
k
((x(a)
(VdL-1/o)leT A gl Ve Σ
max kEle
|
maax
Σ
((1)a1 M1/2
| 1/2 9i
da
1.jE/d]
1
VM
Remark G.2. We note that Theorem 4.5 requires Assumption 4.4 hold. For example, we need the covariance in training and testing to be the same. A similar consistency requirement of the covariance in training and testing had also been observed for in-context linear regression in Zhang et al. (2023a) and for in-context binary classification in the previous section 3.2.
Here, we discuss the consequences when Assumption 4.4 does not hold. For example, suppose the labels of our data in test prompts are not balanced P (y- = ex) = Pk- l do not have the same A-1 weighted norm juر A A- 14 Vx and the covariance matrix of test data r * A, then as N, M + oc, we have
M
C
Σ
yu - (P1/1,P2222 "PPlee)
M
and
(Jquery = 1) + softmax(e(p1/1), P2//2 Pofte) T A A- =xquery)):
Denote V/- = ...,, M T 4 = (log(pı) ..... log(pe)) and x = T I Tquery 4/2 4. Then distribution of the ground truth label is
P (Uquery = e/) = softmax(2)s:
Define = c(p1 Pl, P2H2, Pefte) T A Tquery. Then, unless 2 = or ||sotmmax -0ooosof((((]]]) is sufficiently small, the transformer cannot correctly perform the in-context multi-class classification.
H. Additional Experiments
In this section, we provide additional experimental results and the detailed experimental settings.
H.1. Single-layer transformers
We train single-layer transformers for in-context classification of Gaussian mixtures with different numbers of classes c. different lengths of training prompts N, and test them with different test prompt lengths M. The results are reported in Figure 4. We can see from Figure 4 (a,b) that the inference errors decrease N and M increase, and they increase increases. In Figure 4 (c,d), we first fix the training prompt length (test prompt length) to large number 2000, and then vary the test prompt length (training prompt length) from 20 to 2000. The results show that, as M and N become sufficiently large, the inference error, which is an approximation of E[A (Jquery- Jquerry)l (see Appendix H.2 for detailed definitions), decreases to near-zero. This indicates that the prediction of the trained transformer approaches the Bayes-optimal classifier. All these experimental results corroborate our theoretical claims.
38

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
0.:0
N-40
c-2
N=60
c=4
N=80
Inference error
c=6
N H-120
C=8
0.30
100
150
250
300
300
Test prompt length (M)
Test prompt length (M)
(a) 10
(b) N 80
0.23
c=6, N=2000
c=6, M=
0.20
Infe referenceee error
Inference error
0:35
0.04
6m
180
2000
2000
Test prompt pt length (M)
Training prompt (N)
(c) log-log axes
(d) log-log axes
Figure 4. Inference errors of single-layer transformers. (a): Models trained different training prompt lengths N on classification tasks involving c = classes. (b): Models trained on different classification tasks involving classes with a fixed training prompt length N = 80. (c): Relationship between the inference error and the test prompt length M in log-log axes. Training prompt length N 2000 and number of classes e 6. (d): Relationship between the inference error and the training prompt length N log-log axes. Test prompt length M 2000 and number of classes 6.
39

---

On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
H.2. Experiment Details
For all tasks, we set 20 and we randomly generate covariance matrix A diag(A Ad), where A = |As] and & 1.d. N(3, For each training dataset with different training prompt lengths N, and different class numbers randomly generate B training samples. Training prompts P [B] and their corresponding labels Vr.qqerry are generated according to Assumption 4.2. Moreover, we also generate testing datasets. For example, for each testing dataset, we first randomly generate 20 pairs of (#), queery: Vj.prob). E [20], where (1²) Lt.. ~ Pi (2(A) (A), j.queryy ~ Prr (Hy. Vi.prob = softmax(1g] A j.query) are the corresponding probability distributions of the ground truth label Ujjqueery- For each j. we i.id.
generate 100 testing prompts P = (Iz*,1 Myk,1, M Mik, M, iquery) where (Z)a, Vzkk,) ~ ( A), E [20). A E [100], E [M]. We denote a model's output for testing prompts Pyks as Tzk- We calculate its inference error with
20x 100 keloo]) maXre(e] (Ubz)) (Us.prob)) which serves an approximation of the expected total variation distance defined in (3).
For the '3-layer' model, we used the x-transformers library and defined it as an encoder-only transformer with 64 embedding sizes, 3 layers, 2 heads and without positional encoding.
For experiments in Figure 1, we set the size of the training dataset to B = 100, 000 and set the batch size to 50. We train the 'layyer' using Adam with learning rate 0.0005 for 10 epochs, and train the '3-layer' using Adam with learning rate 0.0001 for 5 epochs. Each experiment is repeated 3 times with different random seeds. For experiments in Figure 2, we also set the size of the training dataset B 100, 000 and set the batch size 50. We train the 'I-layer' using Adam with learning rate 0.001 for 5 epochs, and train the '3-layer' using Adam with learning rate 0.0001 for 5 epochs. In 'sae norm' and 'same covariance' settings, pre-training data are sampled according to Assumption 4.2 with a fixed A that
A = diag ..., Aa), where X; = |Â.] and Å, i.1.d. ~ N(3, 1). In "different norms' setting, for each r E [BB]. with probability P j) = 1/10, f ~ N(k, ja). = 0, ..., 9, then each Gaussian component is sampled according to N(Pr,, A). In (different covariances) setting, we randomly generate 11 02 My E IRd that half of their elements are 0.1 and the other half elements are 100. Then, we define A, diag(v:), i 2, 3 and generate pre-training data according to Assumption 4.2 with A, A1, A2, Ag. Each experiment is repeated 3 times with different random seeds. For experiments in Figure 3, the structure of the transformer with full parameters 'I-layer, full' is defined as
F(E:WY, WKQ) = E E+ E ET
(59)
P
where WV, WKQ (d++) are the parameters for optimization. For all three transformer models, we set the size of the training dataset to B 400, 000 and set the batch size to 50. We train the "lllayer, sparse' and '1-layer, full' using Adam with learning rate 0.001 for 5 epochs, and train the 3-layer transformer model with softmax attention using Adam with learning rate 0.0001 for 5 epochs. Each experiment is repeated 3 times with different random seeds. For experiments in Figure 4, we train the single-layer transformers with the sparse-form parameters and structures defined in Section 4. We set the size of the training dataset to B- = 10, 000 and set the batch size to 50. We train the transformers using SGD with learning rate (0.1, 0.5, 1} for 10 epochs, and get the best model on each training dataset. Then, we test these trained models on different testing datasets. Each experiment is repeated 10 times with different random seeds.
40


Reviews:
[{"review_id": "gUZC5OOmRC", "reviewer": "Reviewer_yaQy", "summary": "In this paper, the authors provide a theoretical analysis of in-context learning of linear classification tasks on the Gaussian mixture models. By assuming a simplified linear self-attention structure and fixing some parameters during the whole training, the authors prove that linear attention can converge to the global minimum in linear convergence rate when minimizing the population loss with gradient descent. Additionally, the authors prove that the global minimum can achieve a small total variance between the output of the trained model and the true label when applied to a new testing prompt. Finally, the authors conduct simple numerical experiments, supporting their conclusions.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">By assuming an over-simplified attention structure...\n\nSetting some parameters to fixed values and considering spare form parameters is commonly used in ICL theory papers [7,8,9,10], and we adopt a similar parameterization as in [7,8,10]. Even for this simplified structure, our analysis for the convergence of $W^*$ is nontrivial. For example, in order to get a tighter bound for $W^*$, in Lemma D3, we first show the loss function $L(W)$ converges point-wisely to $\\widetilde{L}(W)$, which implies that the global minimizer $W^*$ converges to $2\\Lambda^{-1}$. Then, given the property we prove in Lemma D3, we can derive a tighter convergence rate for $W^*$ in Lemma D4 than the bound we got in Lemma D3.\n\n>total variation cannot be used to evaluate test performance. \n\nWe thank the reviewer for this comment, but we believe that total variation distance can be used to evaluate test performance in our setting. We will illustrate this point using a binary case example: Suppose for data $x_q$, the conditional distribution of the ground truth label $y$ is $P(y=1|x_q)=p$, $P(y=-1|x_q)=1-p$. Suppose the prediction of the model is $z$, which depends on $P=(x_1, y_1, ...,x_M, y_M, x_q)$. Suppose the conditional distribution of $z$ is $P(z=1|P)=q$, $P(z=-1|P)=1-q$. Since $(x_i, y_i)$ are i.i.d. sampled, for given $P=(x_1, y_1, x_2, y_2, ...,x_M, y_M, x_q)$, $z|P$ and $y|x_q$ are independent. Thus, the example given by the reviewer is not valid in our setting. If we consider the accuracy, we can calculate it as $Acc(y, z)=pq+(1-p)(1-q)=1+2pq-p-q$. We can notice that, if $p>1/2$, the best $q$ to maximize the accuracy should be $q=1$ and if $p<1/2$, the best $q$ to maximize the accuracy should be $q=0$. Thus, if the output of the model $\\widehat y_{out} | P$  is the same as $P(y=0|x_q)$, we can output $z$ based on the value of $\\widehat y_{out} | P$ to maximize the possible accuracy. The smaller the difference between $\\widehat y_{out} | P$ and $P(y=0|x_q)$, the better the model performs. Actually, the total variation distance in our case measures this difference.\n\n>What concerns me most is their decision to drop all terms and factors that are independent of $N$ during the training process...\n\nWe'd like to clarify that we did not drop all terms and factors that are independent of $N$ during the training process. For example, in Theorem 3.3 Eqn (10), we retained the coefficient of the $1/N$ term and showed that this coefficient is affected by the distribution of distance between two means $\\mu=\\mu_{\\tau, 1}-\\mu_{\\tau, 0}$, the query, and the covariance matrix $\\Lambda$. Intuitively, the classification problem becomes easier when the distance between the two class means $\\mu$ increases and the variance $\\Lambda$ decreases; conversely, it becomes harder when $\\mu$ is small and $\\Lambda$ is large. Eqn (10) reflects this intuition. For instance, when $\\mu$ is large and $\\Lambda$ is small -- corresponding to an easier classification problem -- the parameter $a$ becomes larger, and the derivative $\\sigma'(a)$ decays exponentially with increasing $a$, thereby reducing the overall training error. On the other hand, when $\\mu$ is small and $\\Lambda$ is large -- corresponding to a harder classification problem -- the training error increases accordingly. This theoretical prediction aligns well with our intuition. We did not explicitly present results in terms of the dimensionality $d$, because intuitively, $d$ is not the primary factor determining the classification difficulty. Instead, it is the relationship between $\\mu$ and $\\Lambda$ that plays the central role. Our theory captures this key insight effectively.\n\n>I also do not understand why the query token $q$ can appear on the RHS of the formula of Theorem 3.6 when you are taking the expectation over the query token pair.\n\nWe did not take the expectation over the query token $q$. You can check Theorem 3.6. Our expectation is take over $x_i, y_i; i=1,.., M$, which does include $q$.\n\n>I find that Assumptions 3.2 and 3.5, which state that the mean vectors $\\mu_1$ and $\\mu_2$ have the same $\\Lambda^{-1}$ norm, are somewhat strong.\n\nSee our second reply to reviewer aULa.\n\n>There are several theoretical studies on the optimization of one-layer transformers that consider training $W^v$ and $W^{KQ}$ simultaneously, even in the more challenging softmax attention setting [2, 3, 4, 5]. There are also results on general convergence guarantees of transformers such as [6]. In addition, [7] considers almost the same question as this paper. I suggest that the authors should compare their results with these works.\n\nThank you very much for providing those related works. We will cite and discuss them in the revised paper.\n\n>Is there a typo in line 719? Should it be $\\frac{1}{4N^2}$ instead of $\\frac{1}{4N}$?\n\nNo. This is not a typo. You can verify that the result is $\\frac{1}{4N}$.\n\nReferences can be found in our reply to reviewer M6vU.", "claims_and_evidence": "From my perspective, there exist some major concerns regarding the assumptions and conclusions of this paper. I will list them as follows and discuss the details in the next \"Methods\" and \"Theory\" sections.\n\n1. By assuming an over-simplified attention structure, the optimization problem equation (7) is a convex optimization problem obviously, which is almost totally understood. Additionally, given the loss function (cross-entropy loss function) considered in equation (7) and equation (17), the optimization problem considered in this paper is essentially a logistic regression for the binary case and multi-class case respectively, which is highly understood. I do not feel that there exist any essential challenges in converting the conclusions from binary cases to multi-class cases. Therefore, I do not feel that the optimization problem considered in this paper is as highly non-linear (actually this is a generalized linear model) and challenging as they claimed. \n\n2. Additionally, the author proposes utilizing total variance as the criteria to evaluate the test performance of the new prompt. However, the total variance would only imply the similarity between the distributions, instead of the random variables themselves. Consequently, the Theorem 3.6 implies nothing regarding the testing performance.", "methods_and_evaluation_criteria": "1. As I mentioned in the previous section, the attention layer considered in this paper, adopts a simplified linear attention structure. Additionally, it constrains all the parameters in $W^V$ and $W^{KQ}$ as fixed except the left-top block in $W^{KQ}$. Therefore, all the nonlinearity of this model comes from the cross-entropy loss function. Such oversimplification has rendered this linear attention model equivalent to logistic regression, for which the conclusion that this represents a strongly convex optimization problem within each compact set is well-established. Additionally, the data model of this paper is specified as a Gaussian mixture, which is obviously non-linear separable when given the size of the training set, i.e. $B$ is infinitely large. Therefore, the loss function is coercive and must have a unique minimum given the strongly convex property. Therefore, the optimization problem considered in this paper, from my perspective, is well-studied and trivial to some extent. I have reviewed the proof details and found that most of the proof in this paper focuses on establishing the strong convexity of the cross-entropy loss, which is a well-known fact. Note that even in the original paper studying the in-context learning for training [1], the authors consider a more practical training strategy without fixing any parameters, which makes their loss non-convex and the theoretical analysis highly non-trivial. Besides, there exists multiple theoretical studies that consider training the $W^v$ and $W^{KQ}$ simultaneously, even in a more challenging softmax attention setting [2, 3, 4, 5]. \n\n2. Additionally, as I mentioned in the previous section, total variation cannot be used to evaluate test performance. To illustrate this, consider a simple example where the true label $y$ is a Rademacher random variable, and I choose $\\hat y = -y$ as the prediction of this label $y$. According to the total variation formula in equation (3), we have $\\Delta(y, \n\\hat y)=0$. However, the test accuracy for this prediction $\\hat y$ is $0$. Consequently, the conclusion of Theorem 3.6 does not provide any meaningful insights regarding test performance.\n\n[1]. Zhang, R., Frei, S. and Bartlett, P.L., 2024. Trained transformers learn linear models in-context. JMLR\n\n[2]. Jelassi, S., Sander, M. and Li, Y., 2022. Vision transformers provably learn spatial structure. NeurIPS\n\n[3]. Li, H., Wang, M., Liu, S. and Chen, P.Y., A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity. ICLR\n\n[4]. Wang, Z., Wei, S., Hsu, D. and Lee, J.D., 2024, July. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. ICML\n\n[5]. Zhang, C., Meng, X. and Cao, Y., 2025. Transformer learns optimal variable selection in group-sparse classification. ICLR", "theoretical_claims": "Besides the major concerns I proposed in the previous sections, I still have the following concerns for the proof details.\n\n1. The authors obscure several important conditions in their presentation of theorems. What concerns me most is their decision to drop all terms and factors that are independent of $N$ during the training process. In most theoretical studies related to ICL or transformers, the embedding dimension $d$, is an extremely critical parameter. However, in this paper, the authors directly treat $d$, along with the norms of the mean vectors $\\mu_1$ and $\\mu_2$, and other factors related to the spectral distribution of the variance matrix $\\Lambda$ at the scale of constant order, neglecting these factors entirely. In comparison, in [1], the authors retain all these terms in their conclusions of theorems, therefore the readers can clearly understand how large the sequence length $N$ is required to cancel the effect from the other terms and achieve a good performance.\n\n2. Additionally, I also do not understand why the query token $q$ can appear on the RHS of the formula of Theorem 3.6 when you are taking the expectation over the query token pair.\n\n3. I find that Assumptions 3.2 and 3.5, which state that the mean vectors $\\mu_1$ and $\\mu_2$ have the same $\\Lambda^{-1}$ norm, are somewhat strong. Given these assumptions, I do not see a significant difference from directly assuming an isotropic Gaussian noise.", "experimental_designs_or_analyses": "I have checked the experimental results. There is no major incorrectness.", "relation_to_broader_scientific_literature": "As I said in the previous section, even compared with the original paper [1], the model considered in this paper is much over-simplified. \n\n[1]. Zhang, R., Frei, S. and Bartlett, P.L., 2024. Trained transformers learn linear models in-context. JMLR", "essential_references_not_discussed": "There are several theoretical studies on the optimization of one-layer transformers that consider training $W^v$ and $W^{KQ}$ simultaneously, even in the more challenging softmax attention setting [2, 3, 4, 5]. There are also results on general convergence guarantees of transformers such as [6]. In addition, [7] considers almost the same question as this paper. I suggest that the authors should compare their results with these works.\n\n[2]. Jelassi, S., Sander, M. and Li, Y., 2022. Vision transformers provably learn spatial structure. NeurIPS\n\n[3]. Li, H., Wang, M., Liu, S. and Chen, P.Y., A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity. ICLR\n\n[4]. Wang, Z., Wei, S., Hsu, D. and Lee, J.D., 2024, July. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. ICML\n\n[5]. Zhang, C., Meng, X. and Cao, Y., 2025. Transformer learns optimal variable selection in group-sparse classification. ICLR\n\n[6]. Gao, C., Cao, Y., Li, Z., He, Y., Wang, M., Liu, H., Klusowski, J.M. and Fan, J., Global Convergence in Training Large-Scale Transformers.  NeurIPS \n\n[7]. Frei, S. and Vardi, G., Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context. ICLR", "other_strengths_and_weaknesses": "No.", "other_comments_or_suggestions": "No.", "questions_for_authors": "Is there a typo in line 719? Should it be $\\frac{1}{4N^2}$ instead of $\\frac{1}{4N}$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qu5GzoXROZ", "reviewer": "Reviewer_je9W", "summary": "This paper looks at ICL for classification using Gaussian mixtures (same covariance across classes but different means) by trained transformers. They show that under the condition that the training and test data come from the same covariate covariance distribution, using linear attention can provably work with a single layer. They also show experimental evidence that multiple layers can do better as well as work even when this assumption of the same covariance distribution is broken.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We conducted additional experiments and uploaded code in https://anonymous.4open.science/r/In-Context-Classification-of-Gaussian-Mixtures-2374\n\n>Fundamentally, I felt that the lack of serious comparison to just the linear-regression approach is a big weakness. Treating classification as linear regression isn't optimal, but it can work decently well especially in the kinds of settings here.\nSo, what parts of the analysis are picking up extra nuances of what single-layer linear-attention transformers can do beyond linear regression and what parts are just casting this problem into its shadow linear-regression form in disguise? I can't tell from the discussion. But this is what I really want to know.\n\nIn SNR_LR_Comp.png, we conducted experiments comparing the transformer with linear regression for classification tasks. We can notice that trained transformers, in general, have better performance than linear regression, especially when the prompt length is small. This is not that surprising since the trained transformers somewhat learned $\\Lambda$ during training. \n\n>Given the work in Zhang, Frei, and Bartlett (2023a), this feels like the critical question.\n\nCompared to [1], we use the cross-entropy loss for the classification problem. Thus, a critical distinction to [1] is that our model does not have a closed-form expression of the global minimizer, while [1] does. This distinction adds a significant challenge and leads to a different technical approach compared to [1]. By analyzing the Taylor expansion at the stationary point, we establish the convergence properties of $W^*$. Moreover, [1] considered optimizing the transformer using gradient flow. In contrast, our work proves the convergence of optimizing the transformer with the more practical gradient descent method. We will add more related discussions in our revised paper.\n\n\n>First straight-up least-squares knowing nothing. Second, one that has learned the covariance $\\Lambda$ by magic. Third, one that has learned the covariance $\\Lambda$ using the same training data provided to the transformers?\n\nIn SNR_LR_Comp.png, we compared the trained transformer with straight-up linear regression, knowing nothing. Actually, as the reviewer aULa pointed out, our pre-trained attention transformer approximately implements LDA. If a LDA model learned the covariance $\\Lambda$ by magic or learned the covariance $\\Lambda$ using the same training data provided to the transformers, this LDA model should have better or comparable performance as the trained transformer.\n\n>If I think about how linear regression for classification works in a high-enough-dimensional space, when the means are large, it works quite decently. Anyone who studies mixture-models as toys for classification knows that the signal-to-noise ratio is important. Given that your plots have inference errors dropping to by $5\\%$ or lower, it suggests that the SNR is decent. So how do the results change if we vary the SNR?\n\nIn SNR_LR_Comp.png, by changing the magnitude of the covariance matrix of testing data ($\\Lambda, 4\\Lambda, 16\\Lambda$), we compared the accuracy of trained transformers and linear regression with different rates of SNR. Our results show that smaller SNR leads to worse accuracy for both the transformer model and the linear regression model.\n\nReferences can be found in our reply to reviewer M6vU.", "claims_and_evidence": "Yes, although I must confess that I did not have time to read the full proofs or check them carefully.", "methods_and_evaluation_criteria": "Yes, sort of. There is a key missing comparison in my opinion: to just doing \"least-squares for classification\" --- what is sometimes called the LS-SVM --- where the shared covariate covariance from the training and test data is known to this LS-SVM as well (either directly by magic or from just doing a dumb natural algorithm involving averaging over the training data to extract it.)\n\nGiven the work in Zhang, Frei, and Bartlett (2023a), this feels like the critical question.", "theoretical_claims": "Not really. Sorry.", "experimental_designs_or_analyses": "Yes, see above.", "relation_to_broader_scientific_literature": "The problem of understanding exactly why transformer layers work and what the limitations are is an important problem. The paper does a good job of surveying the literature.", "essential_references_not_discussed": "They made a good choice I feel.", "other_strengths_and_weaknesses": "Fundamentally, I felt that the lack of serious comparison to just the linear-regression approach is a big weakness. Treating classification as linear regression isn't optimal, but it can work decently well especially in the kinds of settings here. \n\nSo, what parts of the analysis are picking up extra nuances of what single-layer linear-attention transformers can do beyond linear regression and what parts are just casting this problem into its shadow linear-regression form in disguise? I can't tell from the discussion. But this is what I really want to know.", "other_comments_or_suggestions": "None. Just answer my questions.", "questions_for_authors": "How does just doing \"least-squares for classification\" --- what is sometimes called the LS-SVM --- do for the experiment in Figure 3? First straight-up least-squares knowing nothing. Second, one that has learned the covariance $\\Lambda$ by magic. Third, one that has learned the covariance $\\Lambda$ using the same training data provided to the transformers?\n\nIf I think about how linear regression for classification works in a high-enough-dimensional space, when the means are large, it works quite decently. Anyone who studies mixture-models as toys for classification knows that the signal-to-noise ratio is important. Given that your plots have inference errors dropping to by 5% or lower, it suggests that the SNR is decent. So how do the results change if we vary the SNR?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6ZI84KGD9o", "reviewer": "Reviewer_aULa", "summary": "This paper studies the in-context learning (ICL) capabilities of the transformer model. In particular, this paper shows that one layer of linear attention mechanism, after pre-training through gradient descent, can implement classification of Gaussian mixture data. The main results of this paper are the convergence guarantees and test-loss bounds of pre-training a (single-layer, sparse, linear) attention model for both binary and multi-class classification.\n\nCompare to previous results such as [Bai et al. 2023], which only provided a construction, this paper presents rigorous guarantees on the GD pre-training convergence dynamics and additionally studies the case of multi-class classification.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">I suggest the author to condense Section 4. Remarks E.1 and G.2 should be included as part of the main body.\n\nThanks for your suggestions. We will modify them in the revised paper.\n\n>It is okay to assume homo-scedasticity (same $\\Lambda$ for both classes), but the assumption on the means having the same $\\Lambda^{-1}$-weighted norms is hard to justify. I guess this assumption is made for the convenience of the proof. If so, the authors should be upfront about this and explain why the problem is intractable without this assumption.\n\nYes, this assumption is made for the convenience of the proof. Actually, in Remark E.1 and G.2, we explained why we need this assumption and showed that if this assumption is not satisfied, the single-layer transformer cannot correctly perform the in-context classification tasks. We also verified this in experiments (Section 5.1, Figure 2) that the single-layer transformer cannot perform well for varying norms.\n\n>It seems that the derivation is about to connect the pre-trained attention model to linear/quadratic discriminant analysis (LDA/QDA) but then abruptly stops. Actually, I think it is very helpful to describe what kind of learning algorithm is implemented by the pre-trained transformer. Setting $W = 2\\Lambda^{-1}$ implements the LDA decision criterion?\n\nYes, we are connecting the pre-trained attention model to LDA and setting $W = 2\\Lambda^{-1}$ implements the LDA decision criterion. We will add more discussions about the connection between the pre-trained attention model and LDA in the revised paper. Thanks for your suggestion.\n\n>What happens when you stack multiple layers?\n\nIn our experiments (Section 5.1, Figure 2), we showed that multi-layer transformers have better robustness for varying covariances and norms. We believe developing a better understanding of multi-layer and more complex transformers is an intriguing direction for future research.\n\n>The well-conditioned property (9)\n\nes, we prove the that there exist $\\alpha>0$ and $l<\\infty$ that satisfy $(9)$. In Lemma D.2, we prove that in a compact domain, the strong convexity parameter $\\alpha$ is larger than 0. In Lemma D.5, we show that $l\\leq  \\frac{1}{4}\\sum_{i\\in[d^2]}E[(p_{t_1(i)}q_{t_2(i)})^2]$, where $p,q$ are some combinations and rotations of Gaussian random variables. Since Gaussian random variables have finite second moments. Thus, $l<\\infty$. We will clarify this in the revised paper.\n\n>I noticed that in Section D.1, the description of D.3 does not match the actual lemma\n\nWhat we exactly did in Lemma D.3 matches what we described in Section D.1. We are analyzing the Taylor expansion of $L(W)$ in Lemma D.3. In line 892-894, we show that as $N\\to\\infty$, our loss function $L(W)$ point-wisely converges to $\\widetilde{L}(W)$. In line 926, we also show that as $N\\to\\infty$, the global minimizer $W^*$ converges to $2\\Lambda^{-1}$. The logic of the proof here is that we first show the loss function $L(W)$ converges point-wisely to $\\widetilde{L}(W)$, which implies that the global minimizer $W^*$ converges to $2\\Lambda^{-1}$. Then, given the property that the global minimizer $W^*$ converges to $2\\Lambda^{-1}$, in Lemma D.4, we can derive a tighter convergence rate for $W^*$. We will clarify this in the revised paper.\n\n>Figure 3: How many classes are there?\n\nWe mentioned in the paper that we compare the classification of three Gaussian mixtures.\n\n>For the \"inference error\" are you referring to the TV distance in equation 3? I don't think this notion is directly applicable to k-nearest neighbor or SVM.\n\nWe used KNeighborsClassifier and SVC from sklearn, and they provided the prediction probability for given test data. We then use the prediction probability to calculate the TV distance. See the API of sklearn for details.\n\n>I find this comparison to be unfair since the transformer models have been pre-trained on a lot (what is the exact number?) of data\n\nThe exact number of training data can be found in H.2. Experiment Details. It is somewhat unfair to compare Transformer models with classical models. However, the goal of our experiments is not to make a fair comparison between Transformers and classical methods, nor to prove that Transformers are superior. Instead, our main purpose is to demonstrate that trained Transformers have in-context learning capabilities and perform no worse than classical methods.\n\n>Why is LDA not part of the classical baseline? You should include your source code\n\nWe upload our code and additional results in the following link. https://anonymous.4open.science/r/In-Context-Classification-of-Gaussian-Mixtures-2374\nWe added LDA in LDA_Comp.png\n\n>Adding a few references on the classical methods for Gaussian mixtures, e.g. linear discriminant analysis, would be very helpful.\n\nThanks for this comment. We will cite and discuss a few references about Gaussian mixtures and linear discriminant analysis in the revised paper.", "claims_and_evidence": "I find the overall writing of this paper to be good and the theoretical claims to be clear. The study of ICL for classification tasks is under-explored in the literature and this paper is a valuable addition towards that direction.\n\nHowever, I find that the paper is lacking in a few areas:\n1. The multi-class case in Section 4 is certainly a worthwhile exercise, but the setting and results are largely direct extension of the binary case in Section 3. I find that there is very little benefit from essentially repeating Section 3. I suggest the author to condense this section, elaborate more on the technical difficulty of the multi-class case, and leave rest of the details to the Appendix.\n2. I find Assumptions 3.1 and 3.5 to be slightly problematic. It is okay to assume homo-scedasticity (same $\\Lambda$ for both classes), but the assumption on the means having the same $\\Lambda^{-1}$-weighted norms is hard to justify. I guess this assumption is made for the convenience of the proof. If so, the authors should be upfront about this and explain why the problem is intractable without this assumption.\n3. The discussion of the prediction (around eq (12)) is rather lacking. It seems that the derivation is about to connect the pre-trained attention model to linear/quadratic discriminant analysis (LDA/QDA) but then abruptly stops. Actually, I think it is very helpful to describe what kind of learning algorithm is implemented by the pre-trained transformer. This type of argument is in fact a major selling point of [von Oswald, 2022] and really helped with its popularity. In fact, if I am not wrong, setting $W = 2\\Lambda^{-1}$ implements the LDA decision criterion?\n4. What happens when you stack multiple layers? We know this helps in the regression setting. Do you have reasons/hypothesis on why more layers does not seem to help much for the classification task?\n\nAnd some minor points: \n1. remarks E.1 and G.2 should be included as part of the main body.\n2. The well-conditioned property (9) is proven in the paper, but currently it sounds like an assumption. Please clarify this.", "methods_and_evaluation_criteria": "n/a", "theoretical_claims": "The theoretical claims and analysis look to be the natural extension of the techniques of [Zhang et. al. 2023], and the authors did a decent job at differentiating their works from the existing results. While I am not a fan of the sparsity assumption (5), almost everyone in the area imposes this assumption. So I will not complain about this too much.\n\nI noticed that in Section D.1, the description of D.3 does not match the actual lemma, and I cannot seem to find a lemma that fits the description. So I would like to see a more elaborate proof sketch that correctly connects different steps of the proof.\n\nOverall, I think the ideas presented in the paper are interesting and valuable. However, I have some reservations regarding 1) the assumptions on the data, 2) connection to classical Gaussian mixture techniques such as LDA, 3) the gap in the proof sketch. I think this paper barely falls short of the standard for publication, but I am happy to upgrade my score if the authors can address my concerns.", "experimental_designs_or_analyses": "I know this is a theory paper, so the few proof-of-concept experiments in Section 5 are fine. However, I find that the discussion around Figure 3 to be extremely lacking.\n1. How many classes are there?\n2. For the \"inference error\" are you referring to the TV distance in equation 3? I don't think this notion is directly applicable to k-nearest neighbor or SVM.\n3. I find this comparison to be unfair since the transformer models have been pre-trained on a lot (what is the exact number?) of data, whereas I assume the classical models are directly applied to the test prompt.\n4. Why is LDA not part of the classical baseline? This is the perfect setting for LDA.\n\nAlso, you should include your source code even if they are short.", "relation_to_broader_scientific_literature": "see above", "essential_references_not_discussed": "I think the coverage of the recent works on ICL is sufficient, but I think adding a few references on the classical methods for Gaussian mixtures, e.g. linear discriminant analysis, would be very helpful.", "other_strengths_and_weaknesses": "see above", "other_comments_or_suggestions": "see above", "questions_for_authors": "see above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6nSdfvmIOM", "reviewer": "Reviewer_M6vU", "summary": "This work studies the convergence and training dynamics of transformers for in-context classification tasks of Gaussian mixture data. The results show that a single-layer transformer trained by gradient descent converges to the global optimal at a linear rate. A quantification of how the training and testing prompt lengths affect the inference is provided. The analysis can also be extended to multi-classification problems.\n\n--------------------------------\n## Update after rebuttal\n\nI appreciate the authors' clarification. I prefer to keep my current rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">The technical contribution beyond [Zhang et al., 2023a] is not clear to me. You consider the classification problem with a different loss function. But why is it challenging enough compared with [Zhang et al., 2023a]. It is better to include related discussions in the paper.\n\nCompared to [1], we use the cross-entropy loss for the classification problem. Thus, a critical distinction to [1] is that our model does not have a closed-form expression of the global minimizer, while [1] does. This distinction adds a significant challenge and leads to a different technical approach compared to [1]. By analyzing the Taylor expansion at the stationary point, we establish the convergence properties of $W^*$. Moreover, [1] considered optimizing the transformer using gradient flow. In contrast, our work proves the convergence of optimizing the transformer with the more practical gradient descent method. We will add more related discussions in our revised paper.\n\n>The practical insight of this work is not clear. For example, how can the theoretical analysis be used to improve ICL in practice? Are the theoretical conclusions aligned with any existing empirical finding? Or can the theory be used to explain any theoretical finding?\n\nYes, our theoretical conclusions align with existing empirical findings. For example, in Figure 1, we conducted experiments of single-layer and multi-layer transformers for in-context classification of Gaussian mixtures. We can notice that both transformer models' ICL inference errors decrease as training prompt length ($N$) and test prompt length ($M$) increase, and increase as the number of Gaussian mixtures ($c$) increases. These behaviors are consistent with our theoretical claims. Moreover, the results also indicate that some of our insights obtained from studying this simplified model may hold for transformers with more complex structures, and studying this simplified model can help us have a better understanding of the ICL abilities of transformers. We hope this paper can provide valuable insights into the theoretical understanding of the ICL mechanisms of transformers. These insights may be helpful for potential architectural design and building safe and explainable AI systems.\n\n>The writing needs some improvement. It is better to include some remark after each Theorem to provide an intuitive explanation. This can help readers understand the key point of each result.\n\nThanks for your suggestion. We will add more marks and intuitive explanations in the revised paper.\n\n**References**:\n\n[1] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023a. \n\n[8] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.\n\n[9] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024.\n\n[10] 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? arXiv preprint arXiv:2310.08391, 2023.", "claims_and_evidence": "Yes, the analysis is solid to me.", "methods_and_evaluation_criteria": "N/A. This paper is mainly theoretical.", "theoretical_claims": "The theoretical analysis looks rigorous, and the conclusions make sense.", "experimental_designs_or_analyses": "N/A. This paper is mainly theoretical.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "I would like to mention some weaknesses here. \n\n1. Theoretical novelty. The technical contribution beyond [Zhang et al., 2023a] is not clear to me. You consider the classification problem with a different loss function. But why is it challenging enough compared with [Zhang et al., 2023a]. It is better to include related discussions in the paper. \n\n2. The practical insight of this work is not clear. For example, how can the theoretical analysis be used to improve ICL in practice? Are the theoretical conclusions aligned with any existing empirical finding? Or can the theory be used to explain any theoretical finding?\n\n3. The writing needs some improvement. It is better to include some remark after each Theorem to provide an intuitive explanation. This can help readers understand the key point of each result.\n\nZhang et al., 2023a. Trained transformers learn linear models in-context.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Can you prove that $\\alpha>0$ in Eqn 18? Otherwise, the convergence analysis is not that strong.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper studies in-context learning of linear classification on Gaussian mixture data. The authors show that a single-layer transformer with linear attention trained via gradient descent converges to an optimal model, and demonstrate that when the lengths of training and testing prompts are sufficiently large, the trained transformer can perform near-optimal prediction. 
Several reviewers have shared the concern that the results are not very surprising given the simplifications considered in the paper: the training is essentially linear logistic regression. However, after several rounds of discussions, the reviewers agree that the theoretical studies and conclusions are sound. Overall, the paper is on the borderline, and the rating slightly leans towards acceptance.
Decision: Accept (poster)