Task Descriptors Help Transformers Learn Linear Models In-Context
Authors: Ruomin Huang
Rong Ge
Venue: ICLR 2025 Poster (2025)
Keywords: Transformer
in-context learning
linear regression
task descriptor
Primary area: learning theory

Abstract: Large language models (LLMs) exhibit strong in-context learning (ICL) ability, which allows the model to make predictions on new examples based on the given prompt. Recently, a line of research (Von Oswald et al., 2023; Aky¨urek et al., 2023; Ahn et al., 2023; Mahankali et al., 2023; Zhang et al., 2024) considered ICL for a simple linear regression setting and showed that the forward pass of Transformers is simulating some variants of gradient descent (GD) algorithms on the in-context examples. In practice, the input prompt usually contains a task descriptor in addition to in-context examples. We investigate how the task description helps ICL in the linear regression setting. Consider a simple setting where the task descriptor describes the mean of input in linear regression. Our results show that gradient flow converges to a global minimum for a linear Transformer. At the global minimum, the Transformer learns to use the task descriptor effectively to improve its performance. Empirically, we verify our results by showing that the weights converge to the predicted global minimum and Transformers indeed perform better with task descriptors.

Published as a conference paper at ICLR 2025
TASK DESCRIPTORS HELP TRANSFORMERS LEARN LINEAR MODELS IN-CONTEXT
Ruomin Huang
Rong Ge
Duke University
Duke University
ruomin. hhwngug edu
rongge duke edu
ABSTRACT
Large language models (LLMs) exhibit strong in-context learning (ICL) ability, which allows the model to make predictions on new examples based on the given prompt. Recently, line of research (Von Oswald al. 2023 Akyürek al.)
2023
Ahn et al. 2023 Mahankali et al. 2023; ZZhng et al. 2024a Vladymyrov
et al. 2024) considered ICL for a simple linear regression setting and showed that
the forward pass of Transformers is simulating some variants of gradient descent (GD) algorithms on the in-context examples. In practice, the input prompt usu- ally contains a task descriptor in addition to in-context examples. We investigate how the task description helps ICL in the linear regression setting. Consider a simple setting where the task descriptor describes the mean of input in linear re- gression. Our results show that gradient flow converges to a global minimum for linear Transformer. At the global minimum, the Transformer learns to use the task descriptor effectively to improve its performance. Empirically, we verify our results by showing that the weights converge to the predicted global minimum and Transformers indeed perform better with task descriptors.
INTRODUCTION
Transformer-baseee large language models (LLMs) have exhibited surprising abilities. One of their most remarkable abilities is to perform well even on tasks that they are not explicitly trained on. This is partially attributed in-context learning (ICL) mechanism, where in-context examples provided to significantly improve the prediction of LLM on new query input (Brown et al. 2020). For langauge models, classical example of ICL adapted from [rown et al. (2020) is illustrated as in Figure [ where we instruct the language model to translate English to French by providing examples.
Please translate English ish into French:
description
hello => bonjour
in-context axampie #1
many => beaucoup
In-context axample #2
apple m>
query example
Figure 1: An input with both task descriptions and in-context examples.
To understand ICL mechanism without getting into the difficulty of language modeling, Garg et al. (022) investigated the simpler problem of learning function class H in-context. the Transformer is given a sequence S = (x1, (zx))..... Ln, h((x,) ), ..query. Here (x,, h(x;;) (i 1,2,...,n) n) are n in-context examples where h is a function in class H. and &query is new input (the query) that want to evaluate h on. The goal of the Transformer to output a prediction that is equal to h(Z.uery). ). Garg et al. (2022) found that Transformers pretrained to perform this task can successfully in-context learn many concept classes. Later works (Von Oswald et al. 2023 Akkürek
this setup,

---

Published as a conference paper at ICLR 2025
al. 2023 Ahn et al. 2023 Mahankali et al. 2023 Zhang et al. 2024a bj Wu et al. 2024) focused 2g for an unknown vector w and the goal of in-context learning to predict Mquery w Zquery.
more on in mn-onnext learning linear models, where ((x)) = w
These existing theoretical works focuses entirely on leveraging in-context examples. However, as in Figure in practice often provide additional task descriptor ("Translate English to French."). Such descriptors improve the ICL. performance for language models (Brown et al.. 2020). While the prefix before in-context examples could contain various information, this work focuses on descriptions containing distributional information about inputs. In this paper, we investigate whether such task descriptors can help in-context learning for linear models.
1.1 OUR RESULTS
We consider new family of mean-varying linear regression problems. For these problems, the in-context examples (I1, 11), ()) (22, y2).... (In, Yn) are generated based on two parameters μ and w. Here w is the underlying linear relation and we have Y/i w H3. The vector μ specifies the mean of rs's, that is, xt's are sampled from Gaussian (u, I) with mean μ.. The mean will be given the task descriptor in the main paper [ Given μ and in context examples (2,, yx)' and query Equary, the goal of the Transformer is to estimate Vquery = w Equery.
We first show that with a specific embedding of the imput (see Section 2.1 p. I-layer linear self- attention (LSA) network can indeed leverage the task descriptor. We can formally construct the optimal parameters for 1-layer LSA network for any number of samples n. When n is large (going to infinity), the linear Transformer uses the mean 14 straight- forwarllyyrdly to remove the mean of x; 's, leading to a better solution than previous constructions without task descriptor. When the optimal solution is more complicated (see Theorem 4. In this case, we show that the training loss function can be decomposed into two terms and -layer linear Transformers leverage its full capacity to minimize both of them.
small,
For all values of n, we also prove that gradient flow training from a reasonable initialization is guar- anteed to converge to the global optimal solution. Even though the training loss is nonconvex, give a detailed characterization for the set of all global optimal solutions, and show that the training dynamics maintain useful invariants that help the algorithm avoid saddle points and converge to unique global minimum.
Finally, we empirically verify findings in Section 5 For 1-layer Transformers, show that the weights indeed converge to the global optimal solution that we have constructed. We also consider different settings where the Transformer may have multiple layers, or use a different embedding. We empirically observe that Transformers with task descriptor always outperform Transformers without task descriptors all these settings.
1.2 RELATED WORKS
In-context learning for linear regression Garg al. (2022 investigated the function classes that
Transformers can learn in-context, finding that they can learn various function classes, including lin-
functions. Later works Von Oswald et al. 2023) Ahn et al. 2023); Mahankali et al. (2023): ) Bai
et (2023) proposed that, under certain parameters, one forward pass of Transformer is equivalent single step of some variant of gradient descent on linear models. Specifically, Ahnn et al. (2023)
showed that Transformers learn to simulate preconditioned GD, which is optimal for one-layer lin- ear Transformers. Vladymyrov et al. (2024) and Fu al. (2023) demonstrated that preconditioned GD can serve as a second-order optimization algorithm. Zhhang et al. (2024b) found that with an ap- pended linear MLP layer, Transformers can learn an initialization for GD. Wu et al. (2024) explored the task complexity bounds for in-context learning in linear regression. Zhhng et al. (2024a) inves- tigated how Transformers can be trained for ICL by proving that one-layer linear Transformers with appropriate initialization will converge to the global minimum under gradient flow dynamics, Huang ct al. (2023) showed that Transformers with softmax non-linearity can be trained to learn linear re- gression distinct features and analyzed the training dynamics. Chen et al. (2024a) investigated the scenario of multiple-head Transformers and multiple-task linear regressions.
GD
Experiments in Appendix B.5 explore one-hot encoding an altemative form of task descriptors.
2

---

Published as a conference paper at ICLR 2025
Task descriptor To the best of our knowledge, there is lack of theoretical studies how trans- formers make use of task descriptors. Many empirical studies demonstrate their effectiveness. For example, adding token that indicates the domain from which the data originates helps LLMs learn from context more efficiently (Allen-Zhu & Li 2024). BBronn et al. (2020); Honovich et al. (2023) found that incorporating natural language task descriptions helps GPT models improve in-context learning.
Other in-context learning works There is line of works on in-context learning with latent variable models within a Bayesian framework (Xie et al. 2022; Zhang et al. 2023 Wang et al.)
2023
Tiang 2023). Another line of works focus on mechanisms of induction heads. Olsson et al. 2022) proposed that two-layer Transformers can be induction heads, which complete the prompt A, B, A. B. A] + B. Some generalized versions induction heads were later stud- ied: Nichani et al. (2024) explored the mechanism and training dynamics of Trrasformmrs learning causal graphs in-context, and Chen et al. (2024b) analyzed the setting of in-context learning 7t- n-gram Markov chains. There are other works showing Transformers can in-context learn various algorithms and tasks, such as sparse token selection (Wang et al. 2024), CFGs (Yao et al. [021]] Allen- -hu &
Li 2023 Zhao al. 2023), discriminative scanning (Tian et al. 2023) and Transformers (Panigrahi
et al. 2024).
PRELIMINARIES
In this section we first introduce basic notations, then we describe the mean-varying linear regression problem in Section 2.1 Details about linear self-attention architecture and training are given in Section 2.2 Finally Section 2.3 we briefly review how previous work understood ICL for linear regression as doing (preconditioned) gradient descent.
Notations We lowercase letters denote variables and vectors and uppercase letters for matrices. For matrix A, we use A-1,-1 to denote the bottom-right entry (the last row and the last column). denote l2 norm of vector and and || |[ denotes the Frobenius norm of matrix. We use La to denote the x d identity matrix and O,, Oaxd to denote the zero vector and the zero matrix of size d and d d respectively. We omit the subscripts the size can be inferred from the context. We denote Sym(A) = (A + A the symmetrized version for any real square matrix A.
We use ||
2.1 MEAN-VARYING LINEAR REGRESSION
Following the previous line of work Von Oswald et al. 2023 Akyürek et al., 2023 Ahn et al., 2023 Mahankali et al. 2023 Zhang al. (024a). we introduce the mean-varying linear regressions problem.
In this problem there are different linear regression tasks. For each linear regression task т, we first sample a mean μ4, ~ (V Ia) and the linear weight (0, Ia). We then independently sample in-context examples ær,; and the query example Tr.qury from Gaussian distribution A N(Jr, A). The input with task descriptors is S, (u+: I7,1, w F7.1.-... Fr,n w Fr,ns FT.query). The goal of the Transformer is to compute Vrrqquery = w., Tr.qu
Embedding matrix E,. There are many ways encode the input as series of tokens for the transformer. For most parts of the paper, we consider the following embedding matrix E, which duplicates the task descriptor before each stack of (z. y)" T. That is,
Prt
[rr
Pr
Pr
E,
Fr.1
Fr
Fr.n
Fr, query
Vr.1
V7.2
yr
Here set the last query stack to be (My queay, 0) and the entry remains to be filled with the prediction of the model. This particular embedding is chosen to simplify the optimal solution and optimization process. One can of course think of alternative embeddings. For example, also
3

---

Published as a conference paper at ICLR 2025
consider the following preftx embedding Eyre in experiments:
0
0
0
0
1
1
E
M1
1r,1
2rr2
Dr,, query
Vr,1
Vr,2
Vr,n
In this embedding the task descriptor frr is just represented in the first token. The first rows serve as a simplified version of positional encoding that distinguishes the task descriptor, in-context examples and the query example.
2.2 TRAINING DETAILS
To do training on the mean-varying linear regressions problem, we first describe the linear self- attention layer, and then give the training loss and initialization details.
Model Architecture. A standard self-attention Transformer layer with one head computes the fol- lowing update to the embedding E:
(W (iNKE\TO W
S(E; () = E W EM softmax
Here p normalizing factor, E' is the input embedding matrix and M is masking matrix. Similar to previous results(Ahn et al. 2023 Von Oswald et al. 2023 Zhang et al., 2024a)., we consider a simplified version of one-layer linear self-attention (LSA) Transformer. Specifically, the projection matrix and the value matrix are merged into projection-value matrix WPV E RRdxxde and the key matrix and query matrix are merged into a key-query matrix WKQ E ggdexde, Here de the embedding dimension. The attention is also restricted to the first tokens that represent in-context examples (and excludes the query token):
E WWQE n
(
In 0 P
On
fISA(E;W) = E WWW EM
M
(3)
Here W = -WWKQ, WPV) and the normalizing factor is set to be the number of in-context examples n. Note that the masking matrix M excludes attentions to the query token.
Model Prediction. For input embedding E = Er, the prediction is read out from the bottom-right entry of the output
Vr.query fLSA (E,; W)
Training Loss. Let E(W, r) be the expected least-square error for task r. That is,
e(W, = E 2 aaary
(FLAA (Er: W)- w Fn )errqurr))
(5)
Note that different tasks may have different expected loss as they have different H4. training we take expectation over all tasks:
L(W) = EE ( [[(
This represents the population loss for training. In practice, we can generate m sequences S. S.y. S.......... and the empirical loss just the mean-squared error for all the sequences
L(W) E (Ei,
W)-1,-1 W) - W7, queeyy)
Initialization. We make the following assumption on the initialization. The assumption is motivated by the initialization in Zhang et al. (2024)). Assumption 2.1 (Initialization). We assume the initialization of the Transformer satisfies
Σn
Σ12
0g
Oxxd Oaxd
0axd
0g
W kk)
Σg
Egg
0a
0axd
0a
0
0
0
0
where En1, Moz Dig, 221
PSD matrices and a satisfies the equation:
En
|E
|E + 122)) 0.
4

---

Published as a conference paper at ICLR 2025
A simple way to satisfy the requirement is to take Σ11 E12 Σ21 E22 La and 2vd. As we will later see in Section 3.2 setting a this way ensures that thoughout training W and WPV matrices
always balanced they always have the same Frobenius norm. This invariant is important for our analysis since, as demonstrated solutions, need only work with those that are balanced.
Section 3.2 among the infinite global optimal
Training procedure. We run gradient flow on the population loss L(W) from the initialization above.
dW
di -
is possible to use standard techniques to discretize the process and run gradient descent, and is also possible to use a polynomial number of samples to estimate the gradient. For simplicity in this paper we only work with the population loss and gradient flow.
2.3
1-LAYER LSA PERFORMING PRECONDITIONED GRADIENT DESCENT
As Ahn et al.
2023 observed, in the standard linear regression setting where μ4, is always set to O,
then the optimal one layer LSA network computes a preconditioned gradient descent step. That is, they construct weights for I-layer LSA network such that the predicted Squery can be computed by
the input data just consists of E y1
I2
In
Equery
1/2
I/n 0
Iquery = (Iiqurry
Σ y:..)-
Here the vector Σ b= Ue is just multiple of the gradient of least squares objective f(w) = > 1 (y/ de ) The matrix A-1 is related to the covariance matrix of æ;'s and serves preconditioner. As we will see later, with task descriptors I-layer LSA networks can discover more complicated strategies for the mean-varying linear regression problem.
3
WARM UP: LARGE SAMPLE SIZE
The global optimal solution and analysis for the general case are complicated. To highlight our main ideas first describe our results in the limit of infinitely many samples oo.
3.1
MAIN RESULTS
Our main results for infinitely many samples are summarized in the following theorem. As we will see, trained Transformers learn use task descriptors to "standardize" keys by removing the mean.
Theorem 3.1 (Main result). Using initialization as in Assumption 2.1 if the number of samples + oc and satisfies 0 < a for some constant then the gradient flow (8) will converge the global minimizer W. (W w.P) and the corresponding loss lim L(W.) 0. Here n-000
have
Oaxd
-A
0a
Odxd Oaxd 0
Odxd
0g
W.Q
Oaxd 0g
A-1
0a
and W.PV
Oaxd
0a
(9)
u²
0
0
0
u²*
where = (2]|A- 122)
Keys standardization. To understand what the Transformer is doing in this case, we notice that the k matrix can be decomposed the product of two matrices. Let C' and W be
Oaxd
Oaxa
0g
Oaxa
Oaxd
0,
C
Ta
La
0g
A
0g
(10)
0
0
1
Please Lemma A.4 in the appendix for the value of a.
'erre the gradient flow becomes dW
lim L(W).
dt 12

---

Published as a conference paper at ICLR 2025
Then we have W C w. Here W.K is in fact the optimal solution for 1-layer LSA without task descriptors (see Zhang et al. (2024a) and C' is the standardization operator. The effect of matrix C' suggests that when computing the key for the attention, the optimal solution will first standardize the 3y by removing its mean component to get C,,,, which has form
0
CE -
21
22
Zp
Equery
(1I)
y/
1/2
yn
Here we call 2, = I2 H, the standardized versions of z₁. Doing this helps remove the spurious correlations introduced by the nonzero mean May More precisely, we compute the output Iquery as:
Uquery = queery A
i=1
1
(12)
Equerr
A-
iml
query w n 00.
The last step uses the fact that #/2 Ei= ب converges to A as n goes to infinity, so we see that the transformer outputs the correct estimate and achieves 0 loss. On the other hand, if we did not have the C matrix (and hence did not use the task descriptor). the corresponding matrix would be 1
nt >>= _ which only converges to A + Mr This matrix is different for different pr and it is impossible to invert it by changing the weight matrix (the weight matrices need to work simultaneously for all tasks).
3.2 PROOF SKETCH FOR THEOREM
Now we discuss how to prove the convergence result in Theorem 3.11 Several steps of our proof are similar the proof in Zhang al. 2024a) for the case without task descriptors, but as we shall see having task descriptors introduces additional challenges to the proof.
We first characterize the set of global minimum and give a lowerbound of the gradient when the current solution is not globally optimal. As we shall see, there are actually infinitely many global optimal solutions and it would be difficult to lowerbound the norm of the gradient by the distance to the particular optimal solution in Theorem 3.1 We get around this issue by showing that the training dynamics maintain several invariant properties, and the only global optimal solution that satisfies all these invariant properties is the solution in Theorem 3.1
Characterizing global minima and gradient lowerbound Our main lemma below gives a strong characterization of global minima and show that the gradient can be lowerbounded if we are away from the set of global minima.
Lemma 3.2. !f our initialization satisfies Assumption 2.7 and n - oc, then we have
VLLWII 20
-WKQ
+ w.k
+ W.Q 2
+
WKQ
+ K.
A-1 2
u
+lWKQ
A-1
12
WWQ
A
1B
(13)
for some constant c 0.
All four terms in the lemma above needs to be equal to 0 in order to achieve global optimal solution. After examining the four equations, we can show that there are two types of symmetry for the loss function. First, if WKQ is scaled by factor x and WPV is scaled by factor 1/k, then the function computed by the LSA layer does not change. Second, we add skew- symmeriii matrices U TU -U) to the 11 block of WKQ it also doesn't 't change the solution. The second type of symmetry is unique to the setting with task descriptors. Every global minimum is equivalent to the minimum constructed in Theorem 3.1 up these transformations.

---

Published as a conference paper at ICLR 2025
To make the training process easier to analyze and allow use to focus the particular global min-
imum in Theorem 3.1 we compliment Lemma 3.2 with the following invariant result of training dynamics:
Lemma 3.3. If the initialization follows Assumption 2.1 then throughou gradient flow training, the following invariants are maintained:
1. Balancing condition:
x² = =|| 1|
||
11:
(14)
2. W KQ is symmetric matrix.
3. u, WKR W w.k and WhO the only non-zero weights.
4. If is smaller than some positive constant initialization, then 3 for another positive constant B throughout training.
This shows that with the initialization in Assumption 2.1 never need worry about the blocks except for u, W W.K KQ w and W The first symmetry of scaling is now fixed because w.ke W2 are always balanced. The second symmetry also cannot happen because W 11 KQQ remains symmetric. By combining invariants from training trajectory (Lemma 3.3
u and w.Ke w.ke 22
and the landscape result of lowerbounding the gradient (Lemma 3.2
we can get Theorem 3.1
OPTIMAL SOLUTION FOR FINITE NUMBER OF SAMPLES
Surprisingly a 1-layer LSA network comes up with much more complicated strategy when the sample size n is finite. In fact the optimal solutions differ significantly from that of infinitely many samples when is not much larger than d. In this section for simplicity assume that the covari- ance matrix for 2;'s are identity (A = Ia), and summarize the results in theorem below:
Theorem 4.1 (Main theorem, finite sample). Under Assumption [. the number of samples and a satisfies O < a < Or fofor some constant o, then the gradient flow ( will converge to global minimizer W- = (WKQ, W. If the covariance matrix A Ia. then the global minimizer
W. (W.KQ KQ W.P) satisfies
011 01
a12la Ja aggld
Og
Oaxd Oaxd 0
0g Oaxa Oaxd
w.ko_
a21la
Os w.
0g
(15)
02
O
0
b
Here b, (11, (22), 012, 022 are all numbers depending only on n. d. For z E (b, a11, a21 412, (22) we denote = (n, d), the exact formulas for J& are given Theorem A.8 in Appendix.
As see, the optimal solution here makes of two new blocks in W corresponding to an Ja and d:1 Ja. Intuitively, these two blocks leverage the descriptors μ4, in queries, which equal E Exqun query] and thus can be help reduce the bias introduced by Zquery. To see this, we need careful way to decompose the training loss function.
Decomposition of the training loss bias and variance terms, we have
Taking expectation over Aquery and decomposing the loss into
L(W)= = Er. Exr.,tr..rr 2
1
(Basuey [ [ur.query BT.qquery))
bins:L2
(16)
F Ex per
Ezqunry (Br.query
Vr.query
[Urqueery
Br.query 1)²
variance:
We will show that the L, and L2 terms can achieve their individual optimal value independently. To do that we rely on reparametrization trick to separate the variables used in Ta and L.2.

---

n Ei=1 2,. Substitute the Published as a conference paper at ICLR 2025
Simplify the loss in the new parametrization. First note that W31 Wa32 w31 uge Oa is the critical point and they are zero at initialization. We show that this is maintained throughout training. This allows us to simplify the prediction Dr.query- Denote jú, -i= Iy and Ei= m.n! We can then simplify the prediction to be
KQ
KQ
Ur.query = bw, [nnllla11/u? fr + am Ap 012/" py + azz Ar
(17)
Intuitively, for this to predict w, 2r.query. it makes sense to use 012 and a22 this is exactly what happens in the infinite sample regime because there Á ur + and p pr so it is possible to get guequery just by combining these two terms. However this is no longer true for the finite sample case the empirical mean f and second moment Å are different from their expectations, and the first two terms provide useful information about them that helps reducing the loss.
Recall 23 Zr,i H4, is the standardized version of the i-th sample. We can in turn write the empirical variance matrix 2 == n Σ 2:2; and standardized mean == 1 simplified prediction Qr.query back to (16) and we reach
[Motan + a1 + a21 + a22)llu> it + (a21 + a22) (1rST + 2) pr 147) 1=t)l"
I(W)- = Bea.t
2
||oca12 + a2)) ar + bazz (1T²" + 2) La F
+ 2 Exiuu
(18)
Therefore we can write the loss L(W) in the new parametrization A:
b(a11 + 012 + a21 + a22)
02 b(a21 + a22)
(19)
A b(a12 + Q22)
f4 bazz-
Note that for any values of 01,02, 03, 04, there are always values of 6, a11 Q12, 421 that can achieve these 0-values (indeed, the parameters b011, ba12, ba21, bozz). Therefore for the optimal solution it suffices to consider the parametriza- tion of f1, 0g. 0s, f4. Since L1 only depends on 0g and 0₂ and L2 only depends on 01 and 62, we can simply consider the optimal solution for the following two functions separately:
just a full rank linear transformation of
E. [|0akinnt
Li (4) = Ee
04 (Mr + 22) I1
L2(0,, (z) Ben. [| || Oz (pr + B) Mr uT 11]
No need of non-zero 011 and 021 if n oc. the infinite sample regime, notice that f -> pr. à la and & + 02 as n + oc, the decomposition becomes even simpler
1
[1((:1) 11- pr 02/4 --))-11)111 Ezsr, [[0spt-1 + +(0s (04-1) Tal Tall
lim I(O)=
2 Ezc.e
1/2
in
(20)
To minimize both L1 and L2 here, it suffices to choose 02 04 and 0, 0g 0, which 021 =
corresponds to the optimal solution in Theorem 3.1 where 411
Non-zero an and a21 to minimize bias for finite n. When n finite, 2 and deviate from their limits hence just setting 0₂ = 04 1 and 0, 0g 0 are not sufficient to minimize I, and L
8

---

Published as a conference paper at ICLR 2025
at the same time. That is, there will be tradeoff between bias and variance if only using ba12 and bazz. To avoid the tradeoff and minimize both, Transformers need to learn a more delicate structure where b and b021 also used to better estimate μ4,. These requirements give the closed form formulas for 0,'s and then we can solve for the optimal solutions for b, d11; 012, 421, 022- The final result is in Section A.2 Appendix.
The optimization of the finite sample case follows from similar ideas as the infinite sample case, except that the calculations are more complicated. We leave the details in Section A.2 in Appendix.
EXPERIMENTS
In this section show that Transformers can indeed leverage the task descriptor in the mean- varying linear regression problem. For 1-layer Transformers with descriptor embedding, the trained Transformer converges to the optimal solution our theory predicts. Deeper Transformers can further improve the performance both with or without task descriptors, but the Transformers with task descriptors can always leverage that information to achieve lower loss.
5.1 EXPERIMENT SETUP
Model architecture. We train I-layer LSA Transformers parametrized below:
H. W KQ He
He+1 He W MH₃
for l = L, 1.
Here Ho the input matrix, set to be the embedding matrix and the prediction is read out from the bottom-right entry of the output Ūr.query [ff.)-1.-11 Recall M the masking for attention that restricts attention to in-context tokens (excluding the query). We have also done experiments with separate W Wº, W., WV matrices in Section B.1 at Appendix and results are similar.
Embedding matrix. Our weight matrices W wrv are both dg x dg matrices. We use three types of input embedding matrices: embedding with task descriptor E, @.. prefix embedding E in (2) and embedding without task descriptor
E=
Fr.1
Fy.2
Zr,n
Fr. query
(21)
Vr,1
Vr.2
Fr.
0
The first row of E and first two rows of ET serve as a simplified version of positional encoding.
Data hyperparameters. We generate 4096 i..dd input sequences for each episode of training. For all experiments this section, the data dimension = 5 and the covariance matrix A Ta- We leave experiments for different d and A to Section B.2 at Appendix.
Training algorithms. For all experiments, we use Adam optimizer (Kingma & Ba 2015) to train Transformers. We also la gradient clipping to stabilize training.
5.2 EXPERIMENT RESULTS
We first show that for different number of layers and embeddings, Transformers can find ways to leverage the additional information task descriptor.
We plot the ICL training loss curves in Figure 2 First we note that embedding with task descriptors have smaller loss values than embedding without descriptors for 1,2, 3-layer Transformers.
Single-layer Transformers trained on prefix embeddings and embeddings without descriptors exhibit comparable performance. However, as the depth of the Transformer increases, models trained on prefix embeddings achieve lower loss. Surprisingly, 3-layer Transformers outperform those trained on embeddings without descriptors and even surpass those trained on embeddings with duplicated descriptors. This suggests that with prefix embeddings, deeper Transformers may be able to use the descriptor more efficiently. We explore some possibilities by observing the attention matrices in Section B.3 Appendix.

---

Published as a conference paper at ICLR 2025
9J Svarger anbec 200 bridde 3kppe
coctpar creecding 210
ka Susppor ordek Hip petcenbiddeg
Figure 2: The MSE loss curves for Transformers in different depths. We display mean and std of random seeds. The number of samples in each sequence = 50 and data dimension d = 5.
MC MUPNE-L-APW. Tieebmaae
24
1+00
red0
Figure 3: The heat map of W W for well-trained Transformer with task descriptors μrr the training sequences. The number of in-context examples n vary from 100 to 600.
Next we show that 1-layer LSA network indeed finds the optimal solution as our theory predicts To do that, we plot heat maps of the product W WkQ of trained Transformers Figure As Theorem 4.1 suggests, in the heat maps there are four blocks in forms of multiples of Lg. The two left blocks fade while the two right blocks remain prominent as n grows large, which consistent with the infinite sample regime in Theorem 3.1 Detailed formulas for parameters in Theorem 4.1 (appearing in Theorem A.88 in appendix) suggest that ba I will converge to 0 from above as n + 0o and baz1 will converge to 0 from below as oo. The other values bai2 and baz2 will be approximately equal to -1 and respectively as long as is much larger than d. These trends are all observed in Figure 3
CONCLUSIONS
In this work, we investigate how Transformers leverage task descriptions for in-context learning. Specifically, we consider the mean-varying linear regression problem where the task descriptors be set to be the mean μ4 for each task т. We give a global convergence result for Transformers trained with task descriptors. Our theoretical result shows that even 1-layer linear Transformers discover interesting ways to leverage the task deseriptor. We empirically show that Transformers can achieve much lower loss for ICL when task descriptors are provided. We also find clear pattern in the parameters of well-trained Transformers, which verifies our theoretical result. Immediate open problems include understanding how Transformers can make use of task descriptors when the embedding does not duplicate the descriptor, well as what happens for multi-layer and/or nonlinear versions of Transformers. As future work, would also be interesting to explore how Transformers leverage more general task descriptors across broader range of ICL tasks.
7 ACKNOWLEDGEMENTS
Rong Ge and Ruomin Huang are supported by NSF Award DMS-2031849,
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenRe- ntt///fftfde http:/ppdddd./00000000/011111110 Dx
view.net, 2023. URL https: http://pppeeeveieeee htt..tttsss//oppeeviieeee
Zeyuan Allen-Zhu and Yuanzhi Li. Physies of language models: Part 1, context-free grammar. CoRR, abs/2305. 136733 2023. doi: 10.4855/ARXIV.335..1560 URL picttss://oo orgg 10. 4w55888500iii.... 2305. 13673
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws, 2024.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisti- cians: Provable in-context learning with in-context algorithm selection. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NewrIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL
nips wttpppppcccpppppppeeeeeeeeeeeeee files/ palles//aaper/222//aaaa 122688807711300155562222333--553 Cesttrcct-Coffreenne html
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vooee Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. CoRR, abs/2005. 14165, 2020. URL http:: //tppi. org abs
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality (extended abstract). In Shipra Agrawal and Aaron Roth (cds.), The Thirty Seventh Annual Conference on Learning Theory, June 30 July 3, 2023, Edmonton, Canada, volume 247 of Proceedings of Machine Learning Research, pp. 4573. PMLR, 2024a. URL https: hatttss:/prooeediingg mlrr html
kress/v247/siyee
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409. 10559, 2024b.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimiza- tion methods for in-context learning: study with linear models. CoRR, abs/2310.170866 2023. doi: 10.48550/ARXIV.3110.11 17086. URL httpt::tt/sddt 425500/00//0.45//2111/ 17086
Shivam Garg. Dimitris Tsipras. Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Or Honovich, Uri Shaham, Samuel R. Bowman, and Omer Levy. Instruction induction: From few examples to natural language task descriptions. In Anna Rogers, Jordan L. Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-14, doi: 10.18653/V1/003.. acl long 108
2023, pp. 1935-1952. Association for Computational Linguistics, 2023,
ACL-LONG.10S htRh:tps::0 /tattpas://oot http:///..oo//.///////// 10.00.8653333 (28653/v1/220
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers CoRR, abs/23 1b.0223,, 2023. doi: 10.4550/ARXXW.23100622,, URL https: :ttps://doo //i.i 10.
4 48550/ 2310. 05249
11

---

Published as a conference paper at ICLR 2025
Hui Jiang. A latent space theory for emergent abilities in large language models. CoRR, 0bs/2304.09960, 2023. doi: I044550/ARXX.220000990. URL https: http://doo..oq/...... 48550/arXiv 2t150//xxiv.234...... 399600900)
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. URLhttp: /
aba/ 1412. 6980
Zhi-Quan Luo and Paul Tseng. Error bounds and convergence analysis of feasible descent methods:
general approach. Annals of Operations Research, 46(1):157-178, 1993.
Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One step of gradient descent is prov- ably the optimal in-context learner with one layer of linear self-attention. CoRR, abs/2307.03576, URL http:////oo.orr/1...8855///////0
2023. doi: 00448550/RXXI..337..000 3307.035777
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers leam causal structure with gradient descent. In Forty-first International Conference on Machine Learning, ICML 2024, Vi- enna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL
net/
[Eorum?id 1MM4jmMMM1m100
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. CoRR, abs/ 11895, 2022. doi: 10. arXiv 2209. 11895
/tt.tpp:///oi..o//11... 48550/ARXIV 11895, URL httpttt////dtt org/ 0101:/10.48550001 48550/
Abhishek Panigrahi, Sadhika Malladi, Mengzhou Xia, and Sanjeev Arora. Trainable transformer in transformer. In Forty-first International Conference Machine Learning, ICML 2024, Vi- enna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL
Eorum? --mmmidd--ooe FeGG
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S. Du. Scan and snap: Understand- ing training dynamics and token composition in I-layer transformer. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http::/ttt//...to /ttpt:::/ppppr...pee n nips. http:///pppppe.aeeeaa/////////3/ 6b6 hrg.eeeeeeea1b01eeee3e3933366633
html
Max Vladymyrov, Johannes von Oswald, Mark Sandler, and Rong Ge. Linear transformers are versatile in-context learners. CoRR, abs/2402. 14180, 2024. doi: 10.4455/ARXX1.2422.411... 120...1..1110
URL https::/ http:///oooorgg//.... 005.44/555rii/
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174, PMLR, 2023.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang- Large language models are latent variable models: Explaining and finding good demon- strations for in-context learning. In Alice Oh, Tristan Naumann, Amir Glober- son, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural In- formation Processing Systems 36: Annual Conference cessing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http::// /tttp://papers. www.ii..cpppaaeeee cc/paper abst html]
Neural Information Pro-
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D. Lee. Transformers provably learn sparse token selection while fully connecco nets cannot. In Forty-first International Conference Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https:://ppnrevieee net wet./oorumiid=grgill IWDcTd
12

---

Published as a conference paper at ICLR 2025
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https:/ /ttp::///ppeeeviiieee ew. net. forum? ? VSh5ePa0ph
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL,
iew. ew net. www.:nee/forrree forum?id-RdJJFFC JUMT
Shunyu Yao, Binghui Peng, Christos H. Papadimitriou, and Karthik Narasimhan. Self-attention net- works can process bounded hierarchical languages. Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computa- fional Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/JJCLPP 2021, (Volume 1: Long Papers), Virtual Event, August 1-6, 2021, pp. 3770-3785. URL
Association for Computationall Lingatiiinnnnnnnn Linguistics, 2021. doi: 10. 10.11853/V1/20000
https::/aa...or//0...18.8 1ttp://0..16633/11//// aate:22.ccc---1 lng 292
Ruiqi Zhang, Spencer Frei, and Peter I. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024a.
Ruiqi Zhang, Jingfeng Wu, and Peter L. Bartlett. In-context learning of linear transformer block: Benefits the MI. component and one-step GD initialization. CoRR, abs/2402. 14951, 2024b. doi: 10.48550/ARXI1224222111 19551. URL hhttps:///o/ www.orrg//00 htt.a//r///....////e... 4w0.444555055a00 .2w... 2402.
14951
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang. and Zhaoran Wang. context learning learn? bayesian model averaging, parameterization, and generalization. CoRR, abs/2305. 19420, 2023. doi: 10-48550/ARXIV.2355.1122 URL httpp:://d/ /ttt::::/doi..r//1... 48550/arXive 2305. 19420
What and how does in-
Haoyu Zhao, Abhishek Panigrahi, Rong Ge. and Sanjeev Arora. Do transformers parse while pre- dicting the masked word? In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pp. 16513-16542. Association for Computational Linguistics, 2023. doi: N..18653/V1/20233 EMNLE- 1029, URL htthhtss://ooo http///i...rr11..11111118811 18653/ v1 v1/2023. emn lp main. 1029
13

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
1.1
Our Results
Related Works
2
Preliminaries
2.1
Mean-varying Linear Regression
2.2
Training Details
2.3
l-layer LSA Performing Preconditioned Gradient Descent
Warm Up: Large Sample Size
Main results
3.2
Proof sketch for Theorem 3.1
Optimal Solution for Finite Number of Samples
5
Experiments
5.1
Experiment Setup
5.2
Experiment Results
6
Conclusions
10
Acknowledgements
10
Omitted proofs
15
Proof for Theorem 3.]
15
A I. Proof Sketch
15
A. 1.2 Proofs of Lemma
A.3 Lemma A.4 and Lemma A.5
17
A.2
Finite Sample Proof
27
A.2.1
Proof of Theorem A.7
27
A.2.
Proof of Lemma A. and Lemma A.12
30
2.3
Proof of LemmalA.10
35
3
Proof of Theorem A.8
37
Additional experiments
38
B.
Experiments for separate weights matrices.
38
B.2
Experiments for higher dimension and non-spherical covariance matrix A.
38
B.3
Experiments for two-layer Transformers trained on prefix embedding.
39
B.4
Weights patterns for 3-layer Transformers trained E,
40
B.5
Experiments for one-hot task descriptors
40
14

---

Published as a conference paper at ICLR 2025
A OMITTED PROOFS
In this section, we assume E[: is taking expectation over all randomness without specifying. We denote Amin (A) and Amax (A) the smallest and largest eigenvalue for any matrix A. For any PSD matrix A, denote VA the unique PSD matrix such that (VA)² A. Sometimes we also write VA as A4, We denote [k] (1,2......4
We first prove the main result for infinite number of samples in Section A.1 then we extend the result to the more complicated finite sample setting Section A.2
A.1 PROOF FOR THEOREM 3.1
In this section proof Theorem which restate here.
Theorem A.1. Using initialization as in Assumption 2.1 if the number of samples nt - oc and a
satisfies < a < a for some constant Q then the gradient flow (8) will converge to the global
minimizer W. = (W.k WWV) and the corresponding loss lim L(W.) = 0. Here we have
a-000
Oaxd -A
0a
0axd
Oaxd
0g
W.Q
Odxd
-
0a
and W.PV. =
0axd
Odxd
0g
(22)
"²
0g
0,
0
0,
0
u*
where u (2 - 122)
A.1.1 PROOF SKETCH
Here we give the sketch of our proof to Theorem
which follows the proof framework in Zhang et al. (2024a). Before we start, let's write W PV and WKQ into blocks:
WPV 11 W.P
WPV W
W PV
WPPV
Wag
(23)
(w
PV
)T
(wz2 YT
AV
PV
W33
and
WKQ
W-QQ
KQ wia
wkQ=
w.ko
wkk RQ
Way
(24)
(w (K TT
(wx
W33
where all the W11, W12, W21. W22 E Raxd Wis, W2z3, W31, W32 E Rd and W33 E R. By expand- ing the prediction Ur.query = fLSA W2d441144 we know the prediction only depends on the weight blocks w W.KQ wgy. W:KQ w wPy was way
KQ KQ and w Therefore will
only consider the training dynamics of these relevant blocks. To simplify notation we gather all the relevant parameters the following block matrix U.
21
Un
U12
U13
WkQ wkk (w KQ
W W
Ws PV
U =
U2i
Uge
Uus
w
(25)
ug1
Ma2
u-1
T
(w32
w33
As we mentioned in Section the proof of Theorem Lemma A.2. If the initialization follows Assumption 2.1 then throughout gradient flow training, the following invariants are maintained:
3.1 relies on Lemma
which we restate here
1. Balancing condition:
2 = || LIWKQ12
||
1W +
WW +
||
KQ 22
(26)
2. WKQ is symmetric matrix.
KQ 3. W.QQ KQ w.kQ and W 22 are the only non-zero weights.
*Please Lemma A.4 for the value of
15

---

Published as a conference paper at ICLR 2025
4. If is smaller than some positive constant at initialization, then 3 for another positive constant / throughout training.
We split Lemma 3.3 into Lemma A.3 and Lemma A.4 We start with the dynamics of u13, U23, M3i and u32, which shows that these parameters stick to 0 during the training phase so the dynamics of U could be simplified. Then show that there is balance between u_1 and U11 U12 U21 U22- Specifically, we have the following lemma.
Lemma A.3. If our initialization satisfies Assumption 2.7 then we have both
u13(f) = u23(t) u3i (t) uz2(t) =
(27)
and
= for all 0. Addiiioon,l,yy we have U11 remains symmetrie throughout the training.
(1) + (t)l (122( [[) =--(t)==[U[n((([]
(28)
Given the balanced condition, we can prove u.1 could be lower bounded by some positive constant during the training phase in Lemma A.4 which suggests the trajectory of u-1 is away from the saddle point u-1 = O.
Lemma A.4. ofurr initialization satisfies Assumption 2.7 - oc and satisfies a where is equal to
2+2
(29)
2||A F (1 + 2tr (A) + 3d²) + 28drr (A) + 60d3
then have (t) 2 B for all 2 0. Here
(d + 2)
(30)
(1+ (IA + 21r(A) + d2 + 2d)
Combining Lemma
A.3 and Lemma A.4 we have the Lemma 3.3 in the main paper.
With the lower bound of u.1,, we are finally able to give an error bound (Luo & Tseng 1993) of our loss L(U) in Lemma A.5 which is the Lemma 3.2 in the main paper.
Lemma A.5. If our initialization satisfies Assumption 2.1 and + oc, then we have
[VL(
A
20
[sym (U11 + U12 + U22 + U2) +
U22 U21
u-
(31)
-
A-1 1122
+1|U12 +
||
1
u-1
u-
where
Amin (A)²
Amia (A)a
c c= 82 min
270d3
270d3*
90d2
'0dd
With Lemma A.3 in hand, we can finally prove Theorem 3.1
Given these supporting lemmas,
now ready prove the main theorem.
Proof of Theorem 3.1 Since L(U)* > 0 is bounded below, we know L(U:) the loss over gradient flow will converge. Any stationary point U* of the gradient flow must satisfy VL(U*) There-
fore, combining with the error bound (3) we have |1 + U - 4- |SSm + Uja +
U2 12 = |U + A- = U1
A A= 11: = 0, which implies that Uh2 "", Ui2
, Ug 0ax and (U)) Oaxa- By lemma A.3 know this implies Ui, Oaxd-
Finally by direct computation we know the corresponding loss is L(U*) 0, which implies that U is global minimizer. Combining 28). we have w², (2]]A 11: Translating U back to W according to (25), we obtain Theorem
The rest of this section will give detailed proofs for the supporting lemmas.
16

---

Published as a conference paper at ICLR 2025
A.1.2 PROOFS OF LEMMA A.3 LEMMA A.4 AND LEMMA A.5
mostly relies on calculations of gradients. We will show that (i gradients are always O for certain blocks and (ii) the growths of W and WPV are balanced in training dynamics.
The proof Lemma A.3
Proof of Lemma A.3 We first prove 27 by showing is the critical points for certain blocks.
Part I: critical points at zero.
The gradient of the loss is
aE(U, r) E((ur.query - wi Truqqery)
Qur.query 0U
au
To give the detailed gradient formulation, we need to expand 9 Dr.query in terms of U first. Denote Then have
= Ei= sii, and for = Ei= Tr,s-
MTH jr
Mr .. Pr.
w. Aw, )
Un
U12
gqeey = (ul3 1223 u-1)
Åw-
U21
U22
Me FTr.query
,1 - ir
w A
U31
U32
U1s (MTH U11 + MT U2i + w RT U3i) Pr
"/s (MT U12 + M4 pi U22 + w P py U32)) Qr.query (unu, Un + AU21 + Awruja) Mr
U23
423 ( U12 + AU22 + Aw,uya Tq.query
u- (n, w jrr 11111 w. AU21 w Aw, Awrus1m))
+ u-1 (ut w ft, U12 + w, AU22 + w Aw, Urz
(32)
If letting U13 = = = U32 = 0, then we have
graq u- u-1 (ي, , ü. Un1 + w. AU21P + pr_ w pr 14,U1112r.qur + w. AU
(33)
The gradient on 1/3 is
@f(U,T)
(ir.query - w Iy query) (HT U11 + u U21) pr]
Oui3
(Ur.query
W xr guery)) (1+ U12 + M=A U22) Ir..ueryy
Note that
Dr.query W. quecy = -.
" U124 + AU22 r.qerry Zr.query
W-1
degree-1 with respect to Wr. Also note that (K- Un1 uT U21) Mr and (uru) U12 u U22) gr.query do not contain wr. Since [[w,] and wy is independent with all other random variables, we have 84(U.) = 0.
Jusa
Similarly, we have aE(U.T) given that 413 U23 = = 0. Owzs
Let = (iiru, Unp? + AUZ1 + f U12 + A AU22 qeay |T, Then the gradient
on is
al(U,T)
w Aw (2 - w. x xI..query)
Ougi
122 w.. Amoy =F
17

---

Published as a conference paper at ICLR 2025
Similarly, we have
0(UUr)
Bugg
given that 1/1 Ug3
Ug1 U32
0. Taking expectation
(r,, have OLUU)
Ouia
8L(U)
Buga
OL(U)
0031
BL(U)
thigg
0 given that M13
M22
Ug1
U32 O.
Now we prove 28 by showing the dynamics of w² and 1Un l? are balanced.
|| + ||U21|1 + ||222112
Part 2: Balancing condition. We can simplify the prediction Or guary by letting v13 = u23
U32 = 0 in (32), which gives the prediction
U31
query 1w- (1i-H, Un + AUZ1) prr + (a-N) U12 + AU22) xT.qpers)
This implies that
1 u-1 ( (pruT Un + AU2i) Mr + [2rp,U12 + AU22 - La zr.quory):
Graq Vriquery
u-1
(34)
Now
we
can
compute
the dynamics of U by the chain rule
a((Ur)
00
8(B auery 1/, (.sy)
E
(Ur.query Vr.query)
80
Therefore, we have the dynamics of Un. U1 U21, U22 and u as follows:
aE(U,T)
E [(rr.quory Vr Vrr.query Ju- 1 w P Me
(35)
a0n1
BE(U,T) T)
E [Ur.quary Mrrrqurry 1
Aww -
(36)
0021
TE(U,T) [(ir.quory Yr Vr..query )u of for i qqquery)
(37)
0012
(E E (Urrquerry - Vr.qqeery) yu )u
Aw a qquery
(38)
0U12
OE(U,T) T) =E (0r.quary - )ury w T M2ftr + (Ms
Dr.query ].
M1 + u-
Ia
(39)
tu_1
Here M1 = 11 U12 + AU22 - 1 Id and M2 = PT U1 + AU21. Therefore we have
dl (U,T))
On-1 u- ir (UJ de(U,T)
de (U,
+ U Ol (U,
+U 0R(U,T)
+ Uj aU12
(40)
oUn
ØU21
ØU2
Taking expectation over Pr., we have the same thing holds for L(U)
aL (U).
Ju u- = tr U aUn + +UT ØL
aL(U)
OL (U)
+UT OL(U)
+ UT
(41)
0U12
0021
OUn
This implies that du (t),()) tr 00000000000000000000000000000100
(42)
di di
Therefore if we set 20) 110::0000 (122
at initialization,
have
2 = (100000 ( (01
(43)
for all t t 0.
Finally, the gradient of Un as in 35 is obviously symmetric. Given that we also initialize Un to be symmetric, Un will remain symmetric throughout the training.
18

---

Published as a conference paper at ICLR 2025
Now prove Lemma A.4 which shows that u-1 is actually away from the saddle point at origin with appropriate initialization. This proof relies on a careful analysis on loss L as a function of u-1- We can rewrite the loss (,(u1) = a(t)u²
b(t)u_1 + as a time-dependen 2 function of u-1- Note that a and b depend on Un, U1 U2,, U2 and are in turn time-dependent, while e is constant independent of time. If we set u (0) small enough (smaller than the b/2a but greater than 0) at initialization, Lo Lo(u_1 (0)) will be dominated by the degree-1 1 term, which has a negative coefficient at time O. Specifically, we will have an upper bound Lo(u-1()) for some positive number < c. Another important observation is that the gradient flow of U guarantees that L will be no larger than (u(w=1(0) throughout the training, which actually implies that La (u (1)) < c-é for all time I. For any time I, the set of ut that gives loss value at most (f--1: : (u-1) < c e}) will always have a support that is always positive and bounded away from O. We take a lower bound of (w_1) over all time / to obtain the desired time-independent lower bound of u-1- It is worth noting that to obtain time-independent lower and upper bounds of (u-1), we need the balancing condition between u and U11 U12, U21 U22 to eliminate the dependence of L on Un, U12 U21, U22
Lo(0) - c
Proof of Lemma We first rewrite loss function ((U,)) as a function of u-1 By taking expec- tation over pe,, we can obtain L(u-1) as a function of u-1
Step 1: Rewrite loss degree-2 function of u-1- We will decompose the loss E(U, r) into ((U, r) = (U, + (u, and rewrite & and lg separately. Recall M1 = P U12 + AU22 1 la and M2 = far fe- Un + AU21. Then we have
1
((U,T)
E Riuey [Urr.qusyy Vrr.quer)
(44)
2
22
(Rzz, [tr(M] M2NTH17) + Ez, [tr (M, Mi (A + (+P7))
+2E,, [tr (Mr M2H+HHH))
(45)
2
Ez, [tr (MT []]] _2 Eg, tr ((M2 + Mi) (M2 + M))
(46)
(V..))
As we can see in (45) when taking expectation over Equery given pr., the variance term is
2 ". [tr (MT M] which is l1 This implies that E[A,] = L1 and E(fzl = L2 where L2 and Li the bias and variance decomposition of the loss function. Now we compute the expec- tation in E and f2. Define a positive value 7 = || + , (r((A and a positive definite matrix
T= m+1 (++M+P7) + (tr (A) + 1| Ia. By direct computation we have
TU,,T) = 2², - tr (TUTN--N U12A + U (A + mT117) U22A + TT 1--2u-h (T- n
2
[1t!? U22
1
-u U1AA (A + pTr+) U22A)) + z (A) (A)
= -- + +
1 tr tA).
2
(47)
Here C1.1 = ir (1-11, U12A + (A + /T11,) A) is the coefficient of degree-1 term - and
C12 ] tr (TU |ب U12 + U (A (T U2 + 2U (r- "M-HT) U22A) is the
coefficient of degree-2 term u² in E1
19

---

Published as a conference paper at ICLR 2025
Similarly we have
(2( 2 " tr (U(1 Un)) PTb, (U12 + Un) + (U22 + Un1) T T (A + [(b1e) (U22 U21) MT14, T
+2 2U12 + Uni) pTR (r (r- 2
n -1 (U22 + U21) HrH
1 -1-1tr (u+H1 (U12 U11) |+ + (A + (TMiTT) (U22 + UnI) + 2 (m-T1
1
-2,1 2 2 (1+N+)
(48)
Here C2 = Ir (u- (U12 U11) uTH + (A (z1+) (U22 - U21) p=u+) is the coefficient of degree-1 term -th 1- and
02.2 2 tr (7 (U12 + Un)) prll (U12 + Un)) + (U22 + U21) T (A + mTH+) (U22 + Un) PT17
+U12 pri (r [1 (U22 + U21) Mott.
the coefficient of 2 term u², in lz.
Note that coefficients C1,1, C1,2; C2,1, C2 are all time denendenndnnt positive functions of fr. Taking expectation over My, have
L.(u-1)= = Epg, [61 (U, + +2
1 = B,, [C1,2 + c22)221 - (02,1 + c1,1) u-1] + T [ll 12] + 2 (A).
(49)
Step bound the coefficients. Now we give several useful lower and upper bounds on C + C2.2 and 22 + C1,1- Denote (t) as the corresponding coefficient at time t under the gradient flow. The first two bounds yield a upper bound of Lo (0)) < (0) - for some e >> 0. The third bound yields lower bound for (u-)). Specificcl, we have the following claim.
Claim 1. We have the following three bounds:
1.
E() + c2, 2 (d 2))-)(0),
(50)
2.
0,2 + 022 S 11, + / 11 (211 + I[Ale)),
(51)
3.
C1, + C221 S (2+ 22uu111A + 1u 1|2.
(52)
Step 3: Upper bound (o)) Now we can upper bound Lo(u-1(0)).
(o(w = E + (U(O)
1
[(ci,2 + 0222) (0)2 (c2,1 c1,1) [0]] + - [[||-1]]] + 2 (A)
$ (0)
A + u- 1/= ([AI) + +2[1 ( (d + +2)
(53)
+ B [1[++]11) 11 (r 50 and (1)
< 2 (d+ 20) + E [1|er [1a] 1 tr (A).
-
The last inequality comes from that (0) <a
d+2
at initializa-
2B 2E|IAA+++
1| (IIA - +2 [P]]
tion.
20

---

Published as a conference paper at ICLR 2025
Step 4: Show u-1 is always positive. Note that when u-1 0, the loss L,( 2 || +
ftr (A). Therefore, u-1 non-zero whenever (u(u1)) < 1 p. 11 + 1tr Since
Le(u-1(0 (()) is non-increasing we have (t)))((t) < Eu [||a~[?]] + tr (A) for all 2 O, which implics that u is non-zero
and by
(3) we know (u-u (0) gBp [||и~1|]] + 1/r (A),
for all t 2 0. Further since we have u.1 (0) 0 and u (t) is continuous on L, we have u
for all 2 0.
Step 5: Lower bound [((u1((t) for all t. Now we lower bound Le(u-1(t).
L((ur((t) E [(1,2 + 02,2)u (02,1 + (11) +
E
||| + tr (A)
2
2 -u-1 + c1.1
EB
||| ||
+ 2 tr (A)
(A)
2
(2 22 2 2 E
+ |1:
+ B [1], |12] + (r (u-1 > and (52)
(54)
Step 6: Combine lower and upper bounds together. Since Le(u-1(t)) S Lo(u-1 combin
ing (53) and (54) we have
(d (d+2)u22 2) (0) (4+2/2) IIA + M
u-1 2
=B> 0.
(55)
Now we prove Claim
Proof.
Recall that
61 =t tr (I-MuTU22) + (A + prP1T) U224))
(56)
and
22 tr (R+H) (U12 + Un) PTH7 + (A + p=H+t) (U22 + U21) p.)))
(57)
Computing the expectation, we have
E|cr,] = tr (U12A + U22 + U2213)
(58)
and
E[ (d 2)tr(U12 + Un + U22 U21) + tr((222 U21)
(59)
By are
Assumption 2.1 time
0 we have Uia(0), (1) (0). (2) and
U21 (0)
PSD matrices. 2)tr (U12(0) + (0) (2) + (0)), which implies that
Therefore we have (0)] 2 0 and [[)21
2
(d +
(0) + (0))
2(d+2)² (I + II
(0) II + (Z) + +1 (U))))))
+ 2(d 2)² (IU02((0)1 + 31
2)
+ (0))
(60)
(submultiplicatiiity)
=(d + (0)=
(Assumption 2.1
Therefore we have [E[11,1 (0) + (0)) 2 (d + (0).
Recall that
C = tr (NUTHHTM UnAA + U.I (A + (T14T U22A + 2U (r rr- (111 U22d
2
(61)
21

---

Published as a conference paper at ICLR 2025
and
222 = 2 tr (y + Un) Mep (U12 + U1i) M+M
2 (r-=urrt u (U22 + U21) mzMh
+2 1U + Un1) M+p-
(62)
n
+ tr (U22 + Un) _ (A + /-H+ (Uz2 U21) p+H
Note that r + + /- and y || 12 if n + xc. Therefore we have
| 11 11 + Un 11 14 T |1
|U22 U21l1/11A n-i1 11 l|r
022
2
+ |112 + U11 F
U12 + U21lF |M+ 11 A + 14 MT ||
"rauchy-Schwarle incquality)
1
||A + M= T
| T
(1lU12 Un r + ||U2 U2)l1=)²
2
<2 || A + /4 - 11
||-|1
||
(11 + |1n + 1|22211 + 1U21 l??)
(Triangle inequality)
=2 A + | 11 F1
|| 11?
(Lemma A.3
(63)
Here the second last inequality comes from |/-|12 = || |1 S IIA PTP 1r- Similarly, for 21 we have
1
C122 tr (I [U UIA U.e (A + )+M?)² U22A + 21 (A + uTUTT U22A)
2
1
|1 11-112111112
+ U22 |/2 | + /4 |
2
2
+
|| 11 U U12l1
(Cauchy-Schwwee Schwartz incquality)
1
||A + uT T
F || (11U12ll (| [[
2
1
$ 2 |IA + M4 =
II (1 + |1 || + 10 1 + 1U21 11:)2
< || A + u T 11 || (111121|3 |1 11221 + |1 (Cauchy-Schwartz inequality)
|| A + /- T 1|= ||ll u (Lemma A.3
(64)
Here the second inequality comes from || < M| + Mr || and ||A]| S Il + prue IF.
Adding
113) and (122) up, we have
0,2 + 2 I 12 + pT17 (211/=112 (All)).
(65)
Recall that
c1,1 tr ((I
U12A + (A + pT U22A))
(66)
and
tr (1-11 (U12 + Un) PTMT + (A + /=1+)) (U22 + U2) p+1+))
(67)
We have
C1,1 = tr ((MHT U12A + (A + (TPT) U22A))
||
I|A prM || | (Cauchy-Schwarce inequality)
(11112l + (1U221 I|A + -- 11 ||
[1UTellE + 1122) || + /
(Cauchy-Schree Schwartz inequality)
2 [1U12 122 + 1122211 |1 + (1U2)112)) || + PTM 12
2 A + #
(Lemma A.3
(68)
22

---

Published as a conference paper at ICLR 2025
Here the second inequality comes from || || I ||A + prr.fle and ||A|| I l| + MTM
Similarly we have
2,.1 = tr (pTRe (U12 + Un1) p + (A + p= (U22 + U21) pr(+t
s s11p,1 |1 + Unlle + ||A / T ll r Mu 11ol U2il2
SccwaaySSchwattz inequality)
111 + U1 FF + 112 + U21 ;) || + 1r l|r
(111221 F + |1 | |122 f U21 (1)) A + |- l|
(Triangle inequality)
2 [|U12 f + ||Uzel1 + |1n 11 + |121 F || + M I
(Cauchy-Schwart) (nne--aaaattteieuuuitttt
2 A + |r ll (Lemma A.3
(69)
Here the second inequality comes from ||T T || S ||A + M.
Adding
68 and (69) up, we have
C1,1 + + S S2+
+ Mr
(70)
Finally we prove the gradient lower bound of |V The high-level idea of this proof to expand L and L2 under new parametrizations U = (u-1U12, u-1U22) and Ú = (u-1Sym (U11 + U12 + U21 U22) (U12 + U22)) respectively, where the two terms ad- mit gradient lower bounds respectively. We can merge these two gradient lower bounds by merging these two parametrizations into one parametrization è = (Ü,Ū) to obtain a lower bound on VI(@)]lr- We then write out the Jacobian of the parameters transformation (u-1U72 1-1 1-1 21 1=1U + θ and find its spectral norm is bounded by some con- stant. Therefore, by combining the lower bound of u- we obtained in Lemma A.4 we can translate the gradient lower bound of [(L)]] back to |/ up some constant. Specifically, we use the following lemma to relate the gradient norm bounds after reparametrization.
u_1
Lemma A.6 (Gradient lowerbound with reparametrizationn) Suppose f : IP - IP and g: EP R are differentiable functions. Denote y = f(x) where If is a transformation between two parametriza- tions y and x. Suppose gly) has gradient lower bound [V
8(y). Denote Jy(x) Rpxr the Jacobian matrix of f at x. If (z) is invertible and the spectral norm of the inverse (x))l2 S for some constant a, then we have [va({(<))] 2 -o)5 foy)for all y = f(x).
Proof of Lemma Easy to see through chain rule that V is invertible, we obtain ō(y) < Vvg(v))l = completes the proof.
( (f(x)
< al[Vag(((x))) which
With the above lemma in hand, we can prove Lemma A.5
Proof of Lemma A.5 We first develope a gradient lower bound on L1-
Step 1: gradient lower bound on L1. We take a new parameterization U1 = u-1012 and
U22 = 122- Denote (U) = [U, and (2) = E((2(U, Then we can simply
(U) in the new parameterization U (U12, U22) as (U). Specifically, have
L1(V) (0)
= -/ ir (1u+ ( uu U12 + U (A + p-11,) U2A +
(71)
+2 T2 (A + ( ) U22A 2 T Ú12A (A + -THTT Ü12A)]
First we want to show that
[T)
2 1/100 Xmii (A) min (Amin 1110 -..lF
(72)
F 10
23

---

Published as a conference paper at ICLR 2025
for all 2 O, where Ü= (-A-1,
By direct computation, we have the gradients
aL,(0)
(dd+2)U12 + (d+ 2)02 + AU - -1) A;
8012
(L, 80
(dd+2)Uzz + (A3 + 2A) U2 + (d+ + 2)U12 + AU - I 1)
-11) A.
Therefore we have
(0)(l) -1.ll
(Li(U) (Li(U)
(Üis + A- -,00 - A 1
OU12
oU22
lr A (2 (Tio + Um)
(U12 + U2e) + (Un - 18m2-111 1-1) A (U/2 + (TTT+1-1) +--1)
+ (A + 21) (Üm -A-1) A (ūZ - -A-1) + (ū1z A-1) A ÜūS - A 1))
=Ir (A ( A VVT + (012+A-1
(73)
2 11²(1== +2
2A + (A²=2A+
21) -1 U122 A11+
+ - ( 2 21)
21 -
21 ^² (A2 + 2A +
Ps
In the last equation the matrix V is defined as
V = (U12 (i1++a-1) +1-1) 21
1 _ (A² +2A
+ 21) 2A+21) 4)
2
(74)
(ū22 ^ 1 (A 2T) 21
1 1² (1² 2
2
is easy to see P1 is a PSD matrix. Actually P3 is also PSD. To see this, for any eigenvalue a of A,
the corresponding eigenvalue in P2 is a² 2 2( 2a ())2 > _a²² 2 O. Furthermore
30 48
have Amin (P2) 2 Amin (A)2. Similarly, we have Amin (P1) 2 1o (iinnDmnn (A)3, 1}. Removing
the term containing in (73 we have
VLL (U) IIύ - -Ü.llr
2 tr U12++A + +-1) Pi (Ū12 + A-1) + A (ÜZ- - -111) P2 (UZ 1022-17 A-1)
Amin (A) min (A)3, 11|| A - + 1 Amin (A)?1IU22 - A A
(75)
Amin (A) minf2mii (A)*, 1} (11012 + 11 110, A -
10
Amin (A) mingAmii (A)3, HH U.
10
Therefore we have
( 2 Amia (A) minf
1
Amin(A)) minddmin 21) - -Ü.llF.
(76)
10
24

---

Published as a conference paper at ICLR 2025
Step 2: gradient lower bound Ly. Now we derive a gradient lower bound for L2. We define
a new parameterization U1 := u- Sym 1111 + U1 + U21 + U22) and U2 := u-1(U21 + U22). ) By checking the dynamics of Un in 35 and 37,, we can find that Un1 keeps symmetric for all t 12 O. Therefore U1 = (U + Sym (1
+ U2 + U22)). Note that for any d x matrix A, it
holds that 11 Au+ = (y.(ym((A)) Therefore we can write L2 under the new parameterization as
(2))
1 E [I(A22 - I tr p-11"]
(77) T + 1r11 U1 is an affine function of (U1, U2). f(X) = || is convex function and the expectation reserves convexity.
L2(U)= = U2),
2
Therefore we know T,2 is convex in terms of U1 and U2 since AU2 -
By setting U2 = A and UI Oaxd in (77. we have L2(UT, U2) = O. Denote U = (Ui, U2),
U. (UH, U.).
By convexity, we have
(Vla(O),ū- I2(U1, -VLT())))>V((()))
(VI2(0),,--C) 2L2(O).
(78)
Therefore expanding the expectation in
77 we have
L2(0) = ||AU2 - 11 + (d 4) (ir (U1) + ir (UF) + (U/11)
+ Ir ((AU2- - 1) (U))) + (((UA- (1 + + (AU2 1.5--1p 1) (U)
d
2 Vd++ (AU2 + UUS 27) + ((1)
||A2 - -11 + (d + 4) |U
+ d+4
+ (d + 4) ir (U;)² (r)(()))
Vd+4 (AU2 + UU2+ A 22 +
F
+ d+4 d+5) |AU - 11 +
1
tr (AU2 - 1)²
d+5
- 2tr + 4|1U111 + (d+ + 5) tr 2 (AU2- 1) tr (U1)
Ir (AU2 - 1) + vd+ 5
d
d
||AU2 - 71 + 4 |U 1/2
d 44
d+5)
1
(IAU2 - 11 + +[1U51??)
30d
min (Auii 1)
(I15/|/ + || - A 111
30d
min (Auin (A), 1)²
10 1...2.
30d
(79)
Here the first equation comes from that ||^|12- 2 tr(A2) by taking A = AU2 T. The second equation comes from that U, is symmetric hence tr (UF) 11. The first inequality comes from that || 2 1tr(( (rr(A) for any d x d real matrix by taking = AU2 - 7 and A U,. The last inequality comes from that
-
UU2 11
- A-1 (440)
))
( AUU2 U1
1
1|
(AUS -
111
v min (Amin 1A),
U1
Combining (78) by
inequality we have
Amia (A)2
[VLz 10 -alle 2 (VLz(O), t.) 2 L2(
10 -, 112,
(80)
30d
25

---

Published as a conference paper at ICLR 2025
which yields that
min (Amin (A), 1)2
(V2(( 2
110 - -Ü=1Ie
(81)
30d
Step 3: merge two lower bounds and parametrizations. Combine two types of parameteriza- tions to get θ (U1z, Ug. U2 UT) Let Vec(A) be the vectorization operator in row-wise
1 2 (1, 4) We now give gradient lower bound of
order. For example, Vec
3
[|V((ec(())]]l by adding two gradient lower bounds 76 and 81)
[V((Vc(())III
(82)
= =12)2(00)1 + (
2d (IU: U| -
- + + A- + 1122 A A-111²).
where e' = min
AA)
Amim(A)
30d
30a
10
I0
Step 4: Translate the gradient lower bound to U-parametrizztion. We define W
(u-1U u-1 u-1022 1 u-1U,)) Then we have
La
Vec(O)
Vec(W) =: J Vec(W),
(83)
Le [T +T)
1e (Te + T)
(T +T)
L
Here T E Fdd xdd the transpose operator. That T Vec(A) Voc(AT) for any d matrix A. It is easy to see T has exactly 1 and d2 1 of zero in each row. So we have [e + TT p S 4d We note that invertible and
1
Laa
Lx (I TT)
Le -
(I +T)
+
T)
I
Therefore we have || < 9d2 which yields a upper bound of spectral norm || s 3d.
Now we can invoke Lemma A.6 and obtain
IVL())
= [VV(Vec(W))
(84)
d
2d0 (Iwill + |12
1|12-1-11 - A- + |10 - A-1 + Ul A-112t)
Finally translate the gradient lower bound back to U:
LV))
Vo. LUU + +IVo., + IV,, L(U) IV,,
+ L(,,,LLU)I L0
= (W)
(85)
20 (IU ||U - --- + + - + 11 A 11
Here the second inequality comes from the lower bound w 2 8 in Lemma A.4
26

---

Published as a conference paper at ICLR 2025
A.2 FINITE SAMPLE PROOF
Now we consider the general setting where is finite. We first give global convergence result, of which the proof follows the same framework infinite samples. There are differences in the proof: (i) in the infinite-sample case, our converged solution has a population loss of O, which is automatically the optimal, while the converged loss positive when n is finite. Hence we need a new Lemma (Lemma A. which shows that fixing V13, U23, U31, U32 to be zero does not lose optimality; (ii) as we will see later in Theorem A.8 the optimal solution have a much more complicated form when n is finite. Therefore we use the form of gradient lower bound which does not explicitly involve the form of the optimal solution as in Lemma A.II
Theorem A.7 (Finite-sampp sample convergence). Under Assumption 2.7 if the number of samples 2 2 and a satisfies < a < O for some constant a, then the gradient flow 88) will converge to global minimizer W. = (W.KC KQ W.P).
To give characterization of the global minimizer when A = Ia. we have the following theorem. Theorem A.8 (Global minimizer). If the covariance matrix Ia. then the global minimizer W. (W &Q W.P) in Theorem A.7 satisfies
01 la
a12la
0g
Oaxd Odxd 0
Oaxd
0a
WKQ
021 la
0g2
Os
.W.V_
Odxd
0d
(86)
0
0,
0
0d
b
where
2
nd + 2n n²
(n n² (d+1) 1)
n nd?
5dn + 4n +
+ d
n2(d + 1) + 2d 6dn + 2d + -
n(n= 1)
n² + n - 1)
112 nd 5dn + An d2 dl -
n2(d
2d + 6dn + 2dl2 -
(n n² (d 1)
n²dd + ni - 1)
2
+
+
(2 +
2d + 6dn + 2d + d -
2d + 6dn + 2d + d
(87)
and
nd nd. 2n
(n-n² (d++1)
Fnd 5dn In
ntdn 2d
Q21 = 1
n(n
1((++)+
12 224nd 5dn 4n
([d+1)+2244 2222nnnnddnn-
(88)
(d(11 nn- (d+1)
012 1
+d-
n ((d-1)+ndd-)
022 12
Remark A.9 (Non-zero 011 and 0:11 to reduce bias under finite samples). As we see (8). the Transformer uses some and 022 to minimize the variance term L1. In the infiitt-sampll case, these 012 and az2 also minimize the bias term L2 to O. However, when v. is finite, these 12 and am fail to minimize the bias term as 7 deviates from Ta and deviates from Oa- Thus the Transformer uses some non-zero 011 and 021 to better estimate P4r which in turn reduces the bias term Loz
Theorem A.8 can be simply proved by
1. Computing the critical point of the reparametrized optimization problem:
2. Translate the critical point condition back to the original parametrization along with the balancing condition.
A.2.1 PROOF OF THEOREM A.7
Denote 2; Zy Mry 2query query My, - 2,, and -= = 2े, the deviation vectors, the empirical deviation vector and the empirical covariance matrix of the in-context samples.
Now sketch the proof of Theorem A.7
27

---

Published as a conference paper at ICLR 2025
Decompose the loss. variance terms, we have
Taking expectation over Equery and decomposing the result into bias and
L(U)- E :1
[E Zanry Orr.qeryy
Vr, surrryyyy 2
(89)
E Ma 1er (Bzyeng [Urqquery Vr.query )
(E Zqurry (T.qqeery
Vr,query 1) ))²]]
Simplify the loss in the new parametrization. Noting that Lemma A.3 tells us that u13 = W23 U32 Og through the training regardless of how large n is, we can simplify the prediction to
U31
be
Ur.query = M4 W nr Un/4, + wi AU21 + 14 w jar (1-U1277.quu + w. AU222
(90)
where Ū = =u-U. Therefore, by algebra have
(11e+ + 2z (Ün + U12 + U2i + Ü22) 1r + (urΣT + 22) (ū2r + U22 pe vr:) )1²
(2) =
2 Es..pp
(91)
and
1 Bucunr ||le + 2) wT (U12 U22) VA (urET + U22VA -oa---Val vAll?]
Li(
(92)
Note that for any -ym-myymettic matrix A and vector æ we have T Ar = 0. Therefore by defining (Am(A) / (A ), can further simplify the loss function to be
L(Ü)
||ue- + 2) pT Sym (Un1 + Ú12 + U2i + U22) Mr + (1+2T + +2) (Üsr + U22) pr pr
Esilor
2
1.2
+ Ee |l(uu+))) +2) (012 + U22) VA+ + (urΣT + +2) Ü22VA VA 1
(93)
Here we give a lemma ensuring that restricting u13, U23, U3i and u32 to be 0₂ does not affect the global minimum of the loss.
Lemma A.10. Denote S = (U ((24+1)x(2d-1)) U13 U23 Ugi Ug2 02} the set
of simplified parameters. We have the global minimum of loss L in S equal to the original global minimum:
min L(U)= UES
min UERR2 ()x2d111
L(U).
Now we introduce the new parametrization θ:
θ Sym(Un + 012 + Ün + U2))
O₂ U21 + Un2
(94)
θ₂ Ü12 + Ü22
θ Ü22-
Therefore L, is a function of Oy, O₂ and Ly is a function of θ₁, Og. Denote the new parameter space P {(1, O2, O2 θ₂) θ Iddd for 1 < 4 and θ symmetric). The following lemma shows that the loss L( is actually strongly-convex as a function the convex parameter space P.
28

---

Published as a conference paper at ICLR 2025
Lemma A.11. If the number of samples 2, there exists constant 0' E P, we have
0, such that for any
|V [l( 2 c||@' 6
(95)
We can prove u-1 could be lower bounded by some positive constant during the training phase in
Lemma A.12 which suggests the trajectory of u-1 is away from the saddle point u- 0.
Lemma A.12. If the initialization of one-layer LSA Transformer satisfies Assumption 2.7 and 0 <
where a² equal
d+2
+1+ MMA + 44+ d-4-tr(A) ||l (2d + 3d46 (tr(A)
|| - 4(d (r(A)
(d 2)
n
2n
(96)
we have
a
u- 20
(97)
(2+ V2) (MAII + 2tr (A) d3 2d)
for all I 0.
With Lemma A.12 in hand, we can finally prove Theorem A.7
Proof of Theorem A.7 we have
By setting and O" to be the critical point O* L(θ) in Lemma
A.II
|] 2 c|l@ - @r
(98)
for any θ P.
To translate the above gradient lowerbound back to U parametrization, we first note that (35 implies that Un keeps symmetric for all 2 0. Therefore θ Un + Sym(U12 U21 U22).
Rearrange the parameter matrices in column get θ (0) 0 0 θ and ú (ÜZZ: Ü,T, Ug ÜHD.. Let Veo(A) be the vectorization operator in row-wise order. For example,
12:
2
Vec
3 (1,2,3,4) 4) It is easy to check that
Le
La La
Le
Vec(e) =
Veo(ū) =: (ec(U).
(99)
1a= [I TT) I
[Te
TT)
2ee +T)
Here the transformation matrix is invertible and
I
L
L
(100)
I
Le -T T) 1
2
(T + +T)
-1 +T)
is easy to see bound (98) obtain |
12 S || 1| < 3d. Now we can invoke Lemma A.6 on the gradient lower
2 _ 0*
(101)
Finally we translate the gradient lower bound on Ü back to U:
|VV
V LOU + L())II + IV,,, L(U) + VV, L(U)
(102)
=w2, VL
Therefore using the lower bound of u-1 in Lemma A.12 have
CO
a
|| -00
(103)
3d
(2 + V2) (MAI + 2tr (A) d2 2d)
29

---

Published as a conference paper at ICLR 2025
Since L(U)= 2 is bounded below, we know (t)) the loss over gradient flow will converge. Any stationary point U of the gradient flow must satisfy VL(U*) = O, which implies the corresponding ((U*) critical point of L(@). Since L()) is strongly-convex on θ we have that the critical
point ((U*) is a global minimizer I. under the restriction that M13 W23 = U31 = U32 0.
Finally through Lemma A.10 we have the converged U is global minimizer of T..
A.2.2 PROOF OF LEMMA A. AND LEMMA A. 12
The high-level idea of proof for Lemma A.11 is lower bound the gradients of I, and L2 respec- tively and then combine the obtained bounds.
Proof of Lemma A.17
Given any O', O" P, we want to show that
[V - VL("") 2 cl(e' @"I|F.
(104)
Only need to show
(L(6)) - VI(0"), 6' - 6") 2 cl 0
(105)
With slight abuse of notation, we denote θ = 6' O". It is equivalent to show (VI(e')- VI.(0"), θ) 2 c
(106)
for some > 0. Recall that we have decomposed the loss function into two parts under -- parametrization L())= = L2(01. + (1 θ1). Therefore it sufficient to show that
(L, 0²2) (L, O),, (O,, 94)) 2 c. (1/a|1 + |+11)))
and
(L2 0²) VL2(0" (91, θ2)) (111 1 + (
Step 1. Getting the gradient lower bound for L, (O,, θ1). By computation we have
0L,
0h, (O₃A AOAA) + d++++ tr(A) aeg
n U3 Ao,
and
0L, 0L,
AO,A A20,A A tr(A) +
1 A0,A -AOAA. n
0 ae"
n
n
Therefore we have the gradient lower bound
(VL 0²) - (L, O1), (Os, 9.))
19
1
Torroo + tr(A)
AO.AOT + d [d+2+ tr(A)
100 (1+ 4
2t
θgΛ0o
100
81
19
1
Amin AA)) + tr(A) + Amin (A)²) ||01|
+
100
n
tr(A)
100 (1 + 1)²
d 2+
Amin (A)]|0s11
81
Note that + tr(A)
100(1+))2 (1 04) n
> lif 2 2. We have the gradient lower bound for
81
(O', Oa) - (L, 07), (O, θ.))
19
Amin (A) min
Amin (A)2 d 1 (1( + 110+)
(107)
100
30

---

Published as a conference paper at ICLR 2025
Step 2. Getting the gradient lower bound for L2(01, O2) difference
Again first compute the gradient
2L2
0L2
tr(A) d (d+++
1 1
, +2 +1+2
Sym(AA2)
0
n
n
((1+
AO2 + dd+ tr(()
+
0 la
n1
"
and
2L2 00, 00
0L2
[++ 71
A (20) (20)+ A02)
(1+ (r() A+
d+2 2 tr(A) AO2.
n
Therefore we have the gradient lower bound
VL2(0", (VL2(0),,6)
(0),
2, tr (0,)² tt tr
1
3 (d++ tr(a))
T1 +2 (1+ (1+ ") (Sym(A0))07 + 10,0%) 1+ AO
2n
ir (A())) + tr ((1 tr (0)) + ((++++ ) ir (0i) + [+ _ "
d 2+ tr(A) 10,07A1 ndmax (A)
0)
The inequality comes from d 2 tr(02)2 and (r(A0207)) 2 Xma 10213
We rewrite the * term as
1+ 2 A
#1)
+ d( d+4+E
+ (1+ d+4
1 4
(1+
d (d+ 4)m
^020AA
d+4+
d+4
where >> 0 will be determined later.
Substituting tr
d(1+1)²
d+4+E 1029AA)
(1+2)²
2+4+$ tr (A02)² and tr ((1+ 1)) (0 2
(1 + +) 1r into ir(*) and substituting the rewritten ir(*) term back, we obtain
(VI2(0/, L 0%),, (01, θ2))
2 (0t) + 1 1+1 1
(10,))
+
|
1 1+
1
n
02 + (A0+ +
b
tr(A) tr (01)) + ( d +8+++++0
tr(A)
1
1
(d (d+4) 4)
||01|
2n
((4+2ttrr -(10)))+))
(1+1) n) 1))
d
(1 + 1)²
d (1 (14H))
d
)) ||
+
n (A)
(d 4)m
d d+4
where a : v(1 + (d/dd 44 and b = vd+ After subtracting the first three positive square terms at RHS and substituting ||A@2|1e 2 Amin (A)[(A)))2)l we reach
(VI2(0',
(e", 0%), (01, θ2))
3 tr(A)
++++
1
(d + 4) ||0:12
2
2n
(1++2+0rAA)
(1+ 1
d
d [1 11)
d d(1 + 1)2
d 4 - 3 ))
+
+
Amin Amin 1A)2110211
((mmax())
n
(d + 4)m
d 44
31

---

Published as a conference paper at ICLR 2025
Since 2, see the coefficient of ||O || is positive and bounded below by 2. It remains to pick some & such that the coefficient of ||A@2] F is positive. It is easy to see
d 2 tr()) n m. (A)
(
d
Amaxx (A)d
(d + 4)m 1+ n
nAmax (A) +
Now we consider two complementary cases.
Case 1: If Amass (A) I rd, then we have ** 2 0. So we can simply pick 4.5 and we have
(e), 02)) 2 min(2, Amin (A) (110 + 112211?).
(VLz(01, - LL 62), (+1,
55d
Case 2: If (xmaa((A) nd, we can then pick { tr(A) and by computation we have
dc - Amax + ndmax (A)(d +
d+4 (1
d [ 4)
#)
dc
d+4+6
n(d + Amax (A)
when 2.
Therefore with direct computation with the assumption
(A) > nd we reach
d3 Amii (A)2
(VL2 - VL2(0 (%), (91, (2)) 2 min (d A)x (A)) (10111 +02211)))))
Combining both cases, we obtain the gradient lower bound of L2
d3 Xmin (A)2
Aman (A)2
(1000,,,00)-VV2200000)))..00 min(2,
(|(0 |
||0, 11:
(d + 4) 4) (A)2*
55d
(108)
Step 3. Combine both gradient lower bounds (107) and (108). We finish the proof with
(VL(e') - VL(O"), θ) 2 c2)0011)
(109)
Now we prove Lemma A.12 The high-level idea of the proof is same as Lemma A.4 write the loss (u_,) = a(t) a(t)u²; b(t)u-1 + time-dependent degree-2 function of u-1. Note that ((t) is independent of so we can re-use the previous bound on it. Since a()) dependent on we re-calculate the upper bound of the coefficient for degree-2 term in this proof.
We again
Proof. Rewrite loss as degree-2 function of u-1- We first write loss function as degree-2 polynomial of u.1. Recall that
L( Ez..or [[(++++)) 11 + (p-2T + 2) 02fl7 H4) m²)
1
(p=2T + 2)) θ.VA VA vA ||w- + 1) A TT
2
Since θ is degree-1 w.r.t u-,, we can write L2 and L1 as degree-2 polynomials
I =: 01,242 - + constant,
L 22 C2,1 constant.
32

---

Published as a conference paper at ICLR 2025
Here constant refers to the term that is constant w.r.. u-1- By computation, coefficients c1,1, C,,1, C1,2; C2.2. After writing them out, as we did in proof of Lemma A.3 we will give lower and upper bounds on C1,1 C2 and an upper bound on C122 C222. Now we denote ①
write out the
O. We have
2. = tr (Axx) + +2tr 2r (01) + d'r (01) and C1,1 = tr (AGa1) + tr (@a1).
Note that these two coefficients are independent of n, which implies that previous bounds (50 and (22 still holds when n finite. We restate these bounds here as
(0) + (0) 2 (d + 20) and C1.1 + C2,1 S v2)u (IAII + 2tr(A) + d2 + +dd.
The formulas for 212 and C,22 are more complicated:
201,2 =dd+ + 2) tr (6a.13] tr(A), (0a10T)
d - tr (ΛB,AÂT) + n
n
+ tr (((1 ((11+))
1 A +
Ar)) A) @LΛÔT)
+ 2tr (θgΛθT
A) + 2 tr (Λθr^ÂT). (110)
n
n
d + 4+ tr(A))
(tr (
+ tr (6t) + tr ( (8(б1)
202
n
+ d 2
tr (Λ
+ [+ n tr (A9390%)
n
+ tr(A) tr ( +2 (1+ h) (ur (δ1Δ6T) + tr (θT + tr (a6x) tr(Ō1)
n
(111)
Re-calculate the upper bounds for c122 + c2 We first upper bound 62.2-
202,2 1 d 4+
tr(())
d
9 F
+
22|001
+ d
d+2 2=21|A
|| FF
2 F
+
1+
1
||A|/
|
1
n
n
n1
F
+
Ir(A)
||A|| [02]
2 F.
+2 1+ n 1 (2//A1 F |
+ +d
|0 F
θ Fr
n
d 4+++ tr(A)
1)) (d +1+2)1Iaaa)
|6g 1
(d 2) + (1+ n
n
+ 2 tr(A)
1 / || + (1+ 1
|62 12
+
||A|| + [1
(2 + d) |AAlF) (+d1/All/)
n
n
F
(d+++ tr(A))
(d+2)+ + (1+ 1)
d 2+2+ + tr(A)
4
(2 (2 (2+d)|AAl1
||//
n
n
1
1+ 1 (2 ++JaAAALE)) .
+
1+ n1
||AlE +
d (d+4+ ur(A)
(d (d+2) + (1+ 1
(2 + d]] + 2+2+ +
|||| r
n
n
1 n
1+ n 1) (2 + d||A|||) -2,
+
1+ ||A| +
(112)
The first inequality is because Cauchy-Schwartz inequality and tr(A@z) s vall0g/lr S
| and tr(θ,) < V. The second inequality is because || Ór
[ 13- + ||@z
The third inequality is because for E [2] have 10,113 S
([]Un || + ||aal| + |Uai|l + [[U2211))² < 1| as u. The last equality is because u.1 is balanced with U matrices.
33

---

Published as a conference paper at ICLR 2025
We next upper bound 6122.
201,2 S
(++2+
tr(A)
|||| F
|
|
+ =
d
6 F +
1+ 11
||
||0.
(r())
||A1 |6a 12
+
1+
|
[60]
[
n
d + 2 + tr(a)
(1+ n
=
||A| ||00||2 2
||||| +
(113)
Alll + (1+
1
tr(())
1 + ) I1A1
16.1
||
| 1 +
n
2 d+ 2 + tr(A)
(1+
+
|
||A||F
d + 1 tt((A)
+
(1+ ||
1
|
The first inequality comes from Cauchy-Schwand inequality and submultiplicatiiity of Frobenius
norm. The second inequality is true because 2||@s||=|[@,]
| [|0a|l++||0, The last inequality
because |[>1li - 121 + [U22|13) - 1| and ||6,1| = 11222 1 - U 1...
Now combining
[13) and 112) we have
212 C2.2
$
d + 4+ (r(A)
(d (d+2)+ +
(1=- n
1 (2 2d + d+ 2+ (T(A)
|AlL
n
d tr(A) + (1+ [||]s) |IAlIE
+
1
(2 + A|]+
n
+++1+tr(1 (1+ #
1
| 1+
||AlF)) 11A1E) u2
n
n
=:c(d, n, A)u
(114)
Upper bound (0)) Now we can first upper bound Lo(w-1(0) by
(o(w = (0) (0)) 22, (0) (0), + (0)) ==1 + constant
S c(d, n, A) (0) (d 2) (0) + constant
(115)
2+2 22 (0) + constant.
< 2c(d. n, A)
The first inequality simply combing the upper and lower bounds we just obtained. The second d+2
inequality is because by assumption we have -0) < 2c(d,n,)))
Lower bound (t))) Then we lower bound ()( by
Lt(w- = (c1,2 + 62.2) u - (c1,1 + c2,1) u-1 + constant
2 (c1,1 + 02,1) u-1 + constant
(116)
2 -2 + v2) (IA 2 tr((A) + di + 2d) + constant.
The first inequality is because C122 + 2, > 0. The second inequality is because we have the upper
bound on C1 + C221 and the fact that u-1 > 0. To see this, just noting that (115) implies that (O)) < constant and hence L(U) < constant throughout training. When u-1 = 0, the loss is equal to the constant. This implies that u_1 (t) 4 0 for all t > 0. Since (0) > 0 and u-1
is continuous w.r.t. t, we have (t) > for all t; otherwise there will be some to such that u- (to) = 0, which is impossible.
Combine both upper and lower bounds. Combining have
115 and 116 and using L(U) < L(U())).
d+2
2
V 2c(d, n, A) (2+ 22) (1A + 2tr(A) + d3 2d) u (0),
(117)
which finished the proof.
34

---

Published as a conference paper at ICLR 2025
A.2.3 PROOF OF A.MMAAA.10
remains to prove Lemma A This proof inspired by Ahn et al. (2023). The high-level idea of this proof is to show the loss functions Lı and L2 are convex under the product the last row of WVV and the first two d columns of W This implies the global minimum is achieved at the critical point under the
parametrization, which is
parametrization. Note the value of the new parameters also depend on M13, M23, M32 and 8311- We then show that there exists a critical point such that M13 W23 t32 U31 Oa. which completes the proof. When the covariance A is non-diagonal, we use the SVD decomposition of to reparametrize I, and further simplify the calculations on L1.
Proof. Reparametrize the loss function to be convex. We will show that setting W13 = W23 =
U32 1g11 Oa does not lose optimality. Denote (u 423 u-1) IR2d+1 Then we
write prediction as
EME (W1, W2) n
Qrar
TTrqqerry
Here E the input embedding matrix and with slight abuse of notation Wi, W2 E (R(2++)xxd denote the first d and the second d columns of WKQ respectively. Note that M32 and U31 the last row of W and W2 respectively.
Now again decompose the loss into variance and bias
2 Ex; Reuary Nt+
(Dr.Aqusry
w. Qur.qu)))
1
T EMEETT WIU, LEMEEWW L -TEME Wo4, - w ur ) - [tr (AA)]
{
Ezi,err
2
71
n
(118)
1
E
Big
(tr (Buru))] + Bar, [tr (AA)
2
L2
L1
where
Ww. EME Wrr (H LWWEMEE--- a-
A =
and
1 B = 1 WT EMET a - W. EME we) (H EMET a
1 W. EME a W7
Denote Wi E Rdd and Wz E Rd the i-th column of W and Wz respectively. Denote Uri the i-th coordinate of wr. Next we will see that L is convex w.r.t. to new parametrization (aW.T.... ... .W ld: aW, 21, ... al
11: We first claim that loss is convex w.r.t. the parametrization (aWT. ... aW Id: aW W ... .WWT) )
With slight abuse of notation denote this parametrization aW. To this, first note that L1 (rA) is convex w.r.t. to the d-dimensional vector WW, EME" matrix. Additionally, we have [ EME
w, since A is PSD is an affine function of (aWT, 21, ,WW,) E R ( since the j-th coordinate [ W, EMET wr], ( E aW.T) wri
an affine function of aW,J for [d]. The desired convexity comes from the fact that the com- position of a convex function and an affine map a convex function. In the same way we can see L2 tr (T(B),)))) is convex w.r.t. the parametrization aW. Therefore, L(aW) is convex function.
Further simplify Lg for general covariance matrix. We write := VΣVT the SVD decomposi- tion of covariance matrix A. Denote diag(o), .aa the diagonal matrix Σ. Let a1, 31, 02 and 8 be d-dimensional vectors to be determined later. Denote diag(a) the x d diagonal matrix whose diagonals the vector a.
35

---

Published as a conference paper at ICLR 2025
VT
Odxd V
Od 0a
VT
Oaxd
Od
Define w = VT Wy, W₂ =
Oaxa 0g
W2V and à =
Oaxd
V
0a
a.
1
0d
0g
For each in-context example i E [n]. define
we define Tz,
V Ir, and in, = V
Ur.
We can in turn
VT
Odxd
0d
juy
Hr
E =
Odxd
V T
0g
E
TT,1
.
Fra
Ir, query 0
Fr.query
(119)
0g
0
1
Vr,1
Vr.2
Vr.
We can simplify L1 by
W. EME - wT) )r
1
Σ B
V
Oaxd
0.
VT Odxd 0
Oaxd V 0
Oa 0a
27
B
V
1
1
W2
Odxd
V
Oa
EMET
Wy
n
0
07
1
1
2
F
W.J EMEE
(120)
We can then write L1 as a convex function of new parameters āW₂ == (T ...
aW 2d
To this
end, we have written Ln as convex function of aW2 and L2 as a convex function of aW.
Show the existence of critical point such that M13 M23 U32 U31 Og- Since the critical point for convex function is a global minimizer, it is sufficient to show that there is set of a, Wi and W2 satisfying
V32 = U31 = Oa, such that the corresponding aW is critical point of L, (@W) and the corresponding aW is a critical point of Lz(a)). It is equivalent to set = U23 U32 U31 = Og and prove the existence of solutions the equations V Vaw, L1 Oaxd and V Vaw, L = Oaxa for all i E [d].
Now we construct such solution in the form of
V
Oaxd
( 0g
0a
diag(ag) (12)
diag((a) diag(31)
e2d+1, W2
0dxd 0
V
diag(S2) 32)
V VT,W=
Wi
W₂
OT
02a41
02441
where ey E Rd is the vector such that the j-th coordinate and others are 0. Note that this construction is consistent with the restriction that W13 made the first 2d elements of and the last row of W2 and W, zero.
u31 = 0g since we have
We compute the corresponding gradient
aa L L1
((LEME) aW²s - -
ÊME
By direct computation we have
0
0
E
EMET
0
0
(10)))ej
e
(1
+
0
and
0
ffej
E
ËMĒT
.WW5
ΣMËT
0
geej
n
pjc
0
where np; = -2jj (d+22 tx(2)
and nqj = Baja5(1+o))
02
3 (1 + ) + tr())
36

---

Published as a conference paper at ICLR 2025
We decide the values of ag and B2 by setting Pj and gj =1+ it is a full-rank linear system the solution exists [ and the corresponding V Vaw,, 41 0. Hence the corresponding dg and B2 is global minimizer for L1. Therefore, the corresponding and W2 is a global minimizer for L1.
o, for all E [d]. Since
Now compute the gradient for L2-
V VaW2,44 Lu
-Va L2
E Mrillri ((HEMET a (Wi: + W2:) > -wri EME
Σ 2n
1
(L EMET a (Wi5 + W22) -W:) EMET
((=
+
n
We have the sum minus terms as
0
0
E
" +
1 Σ
MrillijWTi
EMET
0
OjCj
d+1
T
Egh
2
and the sum of plus terms as
LEMET (Wii + +W2) EME
0
0
8y
Pij + 2
Σ
Mrilisy
t
85e]
[ee]
0
Here
Onj
tr(E)
+ B2j (tr(E) + 20,)
8j =
3d + 12+
and
t, = 0ij n
((d+++ [ [1+ 1
tr(E)
d 2d 8+2 [1+ 2++21++1 oj + 2tr(E)
n1
n
n
+ Bij (tr(E) + 20, + o3))
A sufficient condition for critical point is 8, = d+1 and ty = o,, from which we can solve out and B1. Again it is a full-rank linear system so the solution exists.
A.3 PROOF OF THEOREM A.8
Theorem A.8 give a characterization of the optimal solution when A Lr, which can simply be proved by two steps: (i) Computing the critical point of the @-parametrized loss function; (ii) Trans- late the critical point condition back to the original parametrization along with the balancing condi- tion.
Proof. Since A Id, we have
L Ln(Os, Er,,er [ (n+2T + +2) θs Tal
1 (2 θ2) Exn ( Orpr + (HreT + 2) Ozjer - |r )
Lemma A.1l implies that L()) is strongly-conver Hence the global minima is its critical point.
We write out the critical point condition Ve, T.2 = 0, Ve, L22 = 0, Ve, Ta = and Ve. Tn =
Siine only proving the existence of the solution, do need solve a2 and $2 explicitly.
37

---

Published as a conference paper at ICLR 2025
respectively:
011+)) 2443 n (1(+
2 (00++)00-0000000000000000000+0)
(1+ 02) d) Ia
2
02+ 0 (-1+ + (-++ (1+)) 4=0
(1-1 (1+2+ 00
--- -
(121)
Combining T21 with reparametrization formula and balancing condition in 122 we can solve out
the optimal u Un, U12 U21, U22 in Theorem A.8
+ 1 = w 1-m + U12 U21 + U22)
2 u 121 + U22
3 = u U/2 + U22
(122)
O4 u 1U2
u² |0n + 11222 + 11221 11 + 11222
B ADDITIONAL EXPERIMENTS
In this section, we discuss some additional experiments.
B. EXPERIMENTS FOR SEPARATE WEIGHTS MATRICES.
In this section, we report the results for separate W WQ W W matrices.
We plot the loss curves for Transformers with separate weights matrices in Figure 4 Results are similar to those of merged-weights We can see embedding with task descriptors have smaller loss values than embedding without descriptors for 1, 2, 3-layer Transformers.
ondrring 14per
urgas Enbeking Sayer graft
besdin 24ayer
peficetticlyy 3aayer
10"
Figure 4: The MSE loss curves for Transformers with separate WK W W.P. WV matrices in different depths. We display mean and std of random seeds. The number of samples in each sequence = 50, data dimension d = 5 and A Ia-
B.2 EXPERIMENTS FOR HIGHER DIMENSION d AND NON-SPIIERICAL COVARIANCE MATRIX A.
In this section, we report some results on d = 10 and A = diag(0.2, 0.5, 1, 2, 5) respectively.
In Figure 5 we plot loss curves for Transformers trained on different emheddings with data dimen- sion 10. In Figure we plot loss curves for Transformers trained on different embeddings with covariance matrix diag(0.2, 0.5, 1, 2, 5). In both figures we can see that embedding with task deseriptors have smaller loss values than embedding without descriptors for 1, 3-layer Transform- ers.
To how this performance gap will change across different dimension d, further compare sults on mean- varying linear regressions with different input dimensions 5, 10, 15, 20. In all ex- periments, Transformers trained with descriptors achieve lower MSE loss at convergence compared
38

---

Published as a conference paper at ICLR 2025
N/ Svarger anbec ny 24 bridde 3kkee
coctaar creecdigg 200
No Susrpor ordek Hop petcerbidleg
121
Figure 5: The MSE loss curves for Transformers in different depths. We display mean and std of 5 random seeds. The number of samples in each sequence n = 50, data dimension d = 10 and ld-
dcurppor my. b-kgar pebk biodirg 2kayer
Vo wwdeciingg.no 2/gger perfic andedenn mg kay
deurlpir enbediigg 3ayr potcerbicdig dereeiiinnn..... Haper
Figure 6: The MSE loss curves for Transformers in different depths. We display mean and std of 5 random seeds. The number of samples in each sequence n = 50, data dimension d = 5 and A diag (i.2 0.5, 1,).
with those trained without descriptors, where we denote Lesscriptors and Ino descriptors respectively. We report the performance gap Lno descriptoes Lasseriptors for LSA Transformers with different depths at Table [ As we can see, as the ICL linear regression has higher dimension d, the performance gap widens between transformers trained with and without descriptors.
d d= 5
d 10
d = 15
d = 20
1-layer
2.725410.266
5.43944.361
7.857
10.690_0.742
2-yayer
0.473+40119
1.168+0.095
3.548
4.429
3 l-yayrr
0.04510.034
0.128
0.18440.068
0.622
Table 1: The performance gap Leno descriptors Laeseriptoos between Transformers trained with and without task descriptors. We report the mean and std of random seeds.
B.3 EXPERIMENTS FOR TWO-LAYER TRANSFORMERS TRAINED ON PREFIX EMBEDDING
In this section, we discuss some results on two-layer Transformers trained on prefix embedding.
We plot the attention matrices for a 2-layer Transformer trained on prefix embedding in Figure One incomplete explanation for the results that when descriptor added as a prefix, the Transformer initially uses one layer to attend to the descriptor and remove it from each in-context example. Subsequent layers can then simulate variants of GD to make predictions. As observed in the attention matrix, during the first layer, each in-context example attends to the descriptor in the prefix. However, there are additional attention patterns, such as the prefix attending to the in-context examples in the first layer, and the query token paying attention to the prefix token in the second layer during prediction. These attention dynamics may contribute to improved predictions in other ways. Notably, in Figure 2 the 2-layer Transformer trained on prefix embeddings outperform the 1-layer Transformer tramndd on embeddings with duplicated deseriptors, suggesting that these additional attention patterns assist the Transformer in making more accurate predictions.
2lyyer Transformers trained on prefix embedding converge to loss lower than 10 in the center figure while 1-layer Transfommrs trained on F,, converge to loss above 10 in the left figure.
39

---

Published as a conference paper at ICLR 2025
sampled sequence actentor ccores of bayer 0 heed
sampled sequence ottention scores the 1hth loyer, 0.th head
00
50
00
-44
02
40
-33
40
32
50
30
50
Figure 7: The attention matrices for 2-layer Transformer trained on prefix embedding. The left one is the attention matrix of the 1-st layer and the right one is the attention matrix of the 2-nd layer.
B.4 WEIGHTS PATTERNS FOR 3-LAYER TRANSFORMERS TRAINED ON E,
In this section we report some patterns found in 3-layer Transformers trained on embedding with descriptors. We plot the heat maps of converged weights in Figure ß We can roughly observe that they use descriptors to remove the mean from in-context examples and simulate GD++ (Von Oswald ct al. 2023) as there are diagonals WVV matrices at the first two layers.
2018
Myyt LW
Myp 40
Lov
Figure 8: The heat maps of WKQ and WVP for a 3-layer Transformer trained on embedding with descriptors.
B.5 EXPERIMENTS FOR ONE-HOT TASK DESCRIPTORS
We extend our experiments by introducing one-hot vectors descriptors to indicate different distri- butions of æ. We consider another mean-varying linear regression setting where there are l spherical Gaussian distributions with different means μ11 M1...... 14.- Each ICL linear regression task indepen- dently samples (P,, La) from them as the input distribution of z. Instead directly using the mean μ4; the task descriptor, assign one-hot "task-identig vector e; to each of the Gaussian distributions. N ((4,, Ia) and use this vector as the task descriptor.
As summary, there two key modifications:
1. The mean lr for each task is uniformly drawn from fixed set of k vectors {Mi, μ22 ........ ,
2. For each task t, if H, μy for some j S k. we use the one-hot vector e, (1 at position j. 0 elsewhere) the task descriptor. The corresponding input embedding matrix with task descriptors becomes
e,
Ey =
Fr,1
F7.2
Fr,a
Tr. query
Vr,
Fr,2
40

---

Published as a conference paper at ICLR 2025
We can from Figure that one-hot descriptors also improve the ICL performance of trained Transformers.
15
Roorple mg pect dosill Nape
deetiorr bctile 3kmee
Secrrpor andedding and fodking kays
121
360
Figure 9: The MSE loss curves for Transformers in different depths with one-hot descriptors. We display mean and std random seeds. The number of samples in each sequence n 50, data dimension d = 5 and A = la.
41


Reviews:
[{"review_id": "9Y9a5ZyWQ6", "reviewer": "Reviewer_13Fd", "summary": "The paper explores the role of task descriptors in enhancing the in-context learning (ICL) capabilities of Transformers. The authors show that incorporating task descriptors into the input prompt leads to improved ICL performance for linear models. They present theoretical findings indicating that gradient flow converges to a global minimum for simple linear models when task descriptors are used, and they empirically verify these results by demonstrating weight convergence to the predicted global minimum. The paper includes a theoretical analysis of how task descriptors facilitate in-context learning in linear regression settings, illustrates how Transformers can leverage these descriptors to enhance performance, and provides empirical validation of the theoretical findings through experiments with various Transformer configurations.", "questions": "- What are the specific mechanisms through which task descriptors improve ICL in Transformers? Further analysis or experiments that shed light on this process would be valuable.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper offers a novel perspective on how task descriptors can improve the ICL performance of Transformers.\n- The paper presents a theoretical analysis of how task descriptors enhance in-context learning in linear regression settings, showing how Transformers can utilize these descriptors to boost performance. \n- It also offers empirical validation of the theoretical findings through experiments conducted with different Transformer configurations.", "weaknesses": "- While the paper provides a theoretical analysis, it may lack a deeper exploration into the underlying mechanisms of how Transformers integrate task descriptors during the learning process.\n- The experiments, although comprehensive, are limited to a specific setup. More diverse datasets and problem settings could provide a more robust evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B5TKkZlSXt", "reviewer": "Reviewer_WNXG", "summary": "This paper proposes a new mathematical setting for theoretical in-context learning analysis. It propose a setting which models the task description in in-context learning. The author discussed the optimal soluations for a 1-layer linear self-attention network in two situations, where samples are sufficient or limited, and prove that such optimal soluations can be achieved under gradient flow training with a reasonable initialization. Also, empirical verification for 1-layer transformers are also provided.", "questions": "1.\tMight the author provide some insights/comparisons about what real difference the task description brings for ICL? It would be very helpful.\n2.\tProofs of more complex situation than gradient flow / 1-layer LSA would be very helpful, especially for cases where global optimal can not achieved (which I guess is more often in reality?)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\tThe author models the task description part in in-context learning and provide comprehensive analysis.\n2.\tBoth limited-sample and infinite-sample situations are considered. \n3.\tExperiments further strengthened their analysis.", "weaknesses": "1.\tThe conclusions based on a 1-layer linear self-attention model and gradient-flow might be a little weak.\n2.\tThough introducing task description is very interesting, the techniques used and the theorems got might be a little incremental.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3nTmvrj8ro", "reviewer": "Reviewer_yqsE", "summary": "This paper presents a theory analysis of the task descriptor usage in in-context learning (ICL)  for Transformer models.\nThe analysis is conducted based on a simple linear regression problem (i.e., given $x$ drawn from Gaussian $(\\mu, \\mathcal{I})$ predict $y$  = $w^Tx$ where $w$ is the underlying generative weights), and the task descriptors here are the $\\mu$ for the input distribution.\n\nTheoretically, the analysis results show that the provided task descriptors could help the Transformer directly use the $\\mu$ provided in specialized input embedding given infinite samples $n$. Furthermore, the authors show that the model guaranteed to converge to the global optima from a a reasonable initialization. Empirical results show that the task descriptor could help a 1-layer Transformer achieve significantly lower loss, and the findings generalize to multi-layer Transformer model as well.", "questions": "- Would the loss curve comparison change with larger models and complex data samples (higher dimension than 5)? 12-Layer Transformer such as GPT-2?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Overall solid analysis (theory and empirical ) for the task descriptor usage in ICL of Transformer models.", "weaknesses": "- The analysis and results are overly narrow, focusing solely on linear regression. This limited scope makes it difficult to generalize the findings to fundamental tasks like sentiment classification with in-context learning (ICL).\n\n- The paper's restricted scope can be partially attributed to defining the task descriptor as the mean $\\mu$. This is a non-standard approach, as task descriptions in practical ICL applications typically convey semantic instructions (e.g., `Classify the review as positive or negative.`) rather than characterizing the input distribution $x$ (e.g., the lexical and syntactic patterns of review text).\n\n- (Minor) Figure 2 reveals that the performance gap between models with and without task descriptors narrows as model complexity increases (from 1-layer to 2-layer architectures). This raises questions about the utility of task descriptors in modern large-scale models. The diminishing effect of task descriptors with increasing model size may limit the long-term relevance of this analysis, particularly given the trend toward ever-larger models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HJ6JKUtQAH", "reviewer": "Reviewer_UsGc", "summary": "This paper studies in-context learning with task descriptions. It uses a synthetic setup of linear regression, where the means from which the inputs are samples serve as the task description. Theory shows that a single-layer linear attention model can reach the optimal solution and the optimal solutions are characterized. Experiments show that a single-layer can indeed reach the optimal solution, but deeper models still perform better.", "questions": "1. It's quite puzzling to understanding the embedding matrix formulation of the problem. I'm used to seeing in-context examples as running examples in text, which are made of words that are embedded. But here, the embedding matrix itself is the object to learn. More specifically, given this embedding matrix, the objective is to learn the bottom-right entry, which contains the prediction y_query. What are the trainable parameters? Is E itself updated during training? How does this formulation relate to the usual ICL setup of running text? \n\n2. What are multiple tasks mentioned in equation 6 and does it make sense to take the expectation over all of them?  What is a task specifically?\n\n3. What are the different sequences near equation 7? Are there multiple sequences per task? One sequence per task? \n\n4. Can you motivate the initialization in page 4? any clearer motivation besides having two matrices have the same norm?\n\n5. I am possibly missing some background to understand this, but what do you mean by \"We run gradient flow on the population loss\"? How is the optimization done exactly? \n\n6. What is the significance of characterizing the optimal solutions in section 4? How should they be interpreted and what does it tell us about ICL with descriptions more broadly? \n\n7. If a single linear layer transformer can reach the optimal solution, then why do the experimental results show deeper models to perform better? Is it because of training difficulties with the single layer case? Is it just a sample complexity issue, with the single-layer model not having enough samples?\n\n8. Since there's no pre-training and fine-tuning going on here, it seems like all instances of \"pretraining\" could just be changed to \"training\".", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Important research question on the effect of task descriptions in ICL\n- Simple synthetic setup to study \n- Theory showing optimality of a specific model class", "weaknesses": "- The biggest issue I have is whether the synthetic setup is a good proxy to study ICL, and specifically task descriptors. I understand similar synthetic setups were used to study ICL. But why is giving the mean a good task description? How does it simulate the natural language case, for example the one cited from Brown et al.? \n- It would be helpful to explicitly highlight the insights drawn from the various lemmas and theorems throughout section 4. Unfortunately I had a hard time following this part and so I indicate this is a low-confidence review. \n- The experiments are a nice edition to the theory, but I'm confused about the single-layer being inferior to the deeper models. \n- See other comments and questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)