Training Dynamics of In-Context Learning in Linear Attention
Authors: Yedi Zhang
Aaditya K Singh
Peter E. Latham
Andrew M Saxe
Venue: ICML 2025 spotlightposter (2025)
Keywords: learning dynamics
in-context learning
linear attention
Primary area: theory->learning_theory
TL;DR: We theoretically characterize how in-context learning abilities evolve during gradient descent training of linear attention, revealing abrupt acquisition or progressive improvements depending on how the key and query are parametrized.

Abstract: While attention-based models have demonstrated the remarkable ability of in-context learning (ICL), the theoretical understanding of how these models acquired this ability through gradient descent training is still preliminary. Towards answering this question, we study the gradient descent dynamics of multi-head linear self-attention trained for in-context linear regression. We examine two parametrizations of linear self-attention: one with the key and query weights merged as a single matrix (common in theoretical studies), and one with separate key and query matrices (closer to practical settings). For the merged parametrization, we show that the training dynamics has two fixed points and the loss trajectory exhibits a single, abrupt drop. We derive an analytical time-course solution for a certain class of datasets and initialization. For the separate parametrization, we show that the training dynamics has exponentially many fixed points and the loss exhibits saddle-to-saddle dynamics, which we reduce to scalar ordinary differential equations. During training, the model implements principal component regression in context with the number of principal components increasing over training time. Overall, we provide a theoretical description of how ICL abilities evolve during gradient descent training of linear attention, revealing abrupt acquisition or progressive improvements depending on how the key and query are parametrized.

Training Dynamics of In-Context Learning in Linear Attention
Yedi Zhang' Aaditya K. Singh"
Peter E. Latham
Andrew Saxe 1
Abstract
can solve unseen tasks based on exemplars in the context of input prompt. In-context learning (ICL) is critical to the flexibility of large language models, allowing them to solve tasks not explicitly included in their training data. How- ever, it remains unclear how architectures like self-attention acquire this ability through gradient descent training.
While attention- based models have demonstrated the remarkable ability of in-context learning (ICL), the theoretical understanding of how these models acquired this ability through gradient de- scent training is still preliminary. Towards an- swering this question, we study the gradient de- scent dynamics of multi-head linear self-attention trained for in-context linear regression. We exam- ine two parametrizations of linear self-attention: one with the key and query weights merged single matrix (common in theoretical studies), and one with separate key and query matrices (closer to practical settings). For the merged parametrization, we show that the training dynam- ies has two fixed points and the loss trajectory exhibits a single, abrupt drop. We derive an an- alytical time-course solution for certain class of datasets and initialization. For the separate parametrization, we show that the training dynam- ies has exponentially many fixed points and the loss exhibits saddle-to-saddls dynamics, which we reduce to scalar ordinary differential equations. During training, the model implements principal component regression in context with the number of principal components increasing over training time. Overall, we provide theoretical description of how ICL abilities evolve during gradient de- scent training of linear attention, revealing abrupt acquisition or progressive improvements depend- ing on how the key and query are parametrized.
Seminal work by Olsson et al. (2022) identified an intrigu- ing trait in the training dynamics of ICI.: the ICL ability often emerges abruptly, coinciding with abrupt drop in loss during training. This abrupt learning phase can reflect the formation of an induction head in the ICL setting (Ols- son et al., 2022; Reddy, 2024; Singh et al., 2024; Edelman et al., 2024), and can also occur more broadly in transformer training dynamies (Nanda et al., 2023: Chen et al., 2024a; Hoffmann et al., 2024; Gopalani et al., 2024). Furthermore, Singh al. (2023) found that ICL may often be transient ability that the transformers acquire and then lose over the course of long training time, phenomenon that has since been reproduced in many settings (He et al., 2024; Anand et al., 2025; Chan et al., 2025; Nguyen & Reddy, 2025; Park et al., 2025; Singh al., 2025). These findings underscore the importance of understanding not only the ICL ability in trained models, but its full training dynamics.
This work aims to provide theoretical description of how the ICL ability evolves in gradient descent training. To do so, we consider the increasingly common setup of linear attention' (Von Oswald et al., 2023) trained on an in-context linear regression task (Garg et al., 2022). The in-context linear regression task, in which the model needs to perform linear regression the data in context, is a canonical in- stantiation of ICL (Garg et al., 2022; Akyürek et al., 2023; Von Oswald et al., 2023; Ahn et al., 2023; Bai et al., 2023). The linear attention model, which has been used in many prior studies (Schlag et al., 2021; Von Oswald et al., 2023; Ahn al., 2023; Zhang et al., 2024a; Wu et al., 2024; Fu et al., 2024; Mahankali al., 2024; Duraisamy. 2024; Li et al., 2024; Yau et al., 2024; Lu et al., 2025; Frei & Vardi, 2025), reproduces key optimization properties of practical transformers (Ahn et al., 2024) and is more amenable to theoretical analysis. Importantly, despite its name, linear attention is a nonlinear model, as it removes the softmax operation but is still a nonlinear function of the input.
1. Introduction
salf-attention-basee models, such transformers (Vaswani et al., 2017), exhibit a remarkable ability known as in- context learning (Brown et al., 2020). That is, these models
'Equal contribution 'aatsby Computational Neuroscience Unit. University College London "Sainsbury Wellcome Centre, University College London. Correspondence to: Yedi Zhang gatt@gaatbyy.ul....
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Wee refer to linear self-attention as linear attention in this paper.

---

Training Dynamics of In-Context Learning in Linear Attention
We study two common parametrization of multi-head linear attention: (i) ATTNM, linear attention where the key and query matrices in each head are merged into single matrix, reparametrization procedure widely used in theoretical studies on transformers (Ahn et al., 2023; Tian et al., 2023; Atace Tarzanagh et al., 2023; Zhang et al., 2024a;b; Chen et al., 2024b; Wu et al., 2024; Kim & Suzuki, 2024; Huang et al., 2024b; Wang et al., 2024b; Ildiz et al., 2024; Ren et al., 2024; Tarzanagh et al., 2024; Yau al., 2024; Julis- tiono et al., 2024; Anwar et al., 2024; Vasudeva et al., 2025; Lu et al., 2025; Chen & Li, 2025; Huang et al., 2025a); (ii) ATTNs, linear attention with separate key and query matrices, which is closer to the implementation of attention in real-world transformers (Vaswani et al., 2017). We spec- ify the fixed points in the loss landscapes, as well as how gradient descent training dynamies traverses the landscape. Our findings summarized as follows.
in gradient descent training. Our results also reveal how parametrization, such as merged versus separate key and query and the rank of the separate key and query weights, influences the loss landscape and training dynamics. This motivates future research to take the parametrization factor into account when studying the landscape and dynamics of attention models.
2. Preliminaries
Notation. Non-bold small and capital symbols are scalars, Bold small symbols are column vectors. Bold capital sym- bols are matrices. || || denotes the 8 norm of vector or the Frobenius norm of matrix. vec(.) represents flatten- ing a matrix to column vector by stacking its columns.
3
For example, vec
2 4 = 4 2 3 4
We use
1, H denote the index of attention head, / = 1, ... P to denote the index of a training sample, and n = 1, ... .N to denote the index of token in sample.
We find two fixed points in the training dynamics of ATTNM. and exponentially many fixed points in that of ATTNs.
2.1. In-Context Linear Regression Task
We show a single, abrupt loss drop in training ATTNM from small initialization and derive an analytical time- course solution when the input token covariance is white. We show saddle-tossaddl training dynamics in training ATTNs from small initialization and reduce the high- dimensional training dynamics to scalar ordinary differ- ential equations through an ansatz. We demonstrate the rank of the separate key and query weights affects the dynamics by shortening the duration of certain plateaus.
We study a standard ICL task of predicting the next token. The input is a sequence {æ11 y1 22, 1/2." 2 N. UN, Tg and the desired output Vg- We refer to By as the query token, {æ1, y1, 22, V2, ... ;
NN UNN } as the context, and N as the context length. By convention (Ahn et al., 2023; Zhang et al., 2024a;b; Chen et al., 2024b; Huang et al., 2024b), the input sequence is presented to the model as a matrix X, defined as
We identify the in-context algorithm of the converged and early stopped models. When ATTNM and ATTN; are trained to convergence, they approximately imple- ment least squares linear regression in context. When the training of ATTNs early stops during the (m + 1)-th loss plateau, it approximately implements principal com- ponent regression in context with the first m principal components.
X
Φ1
22
EN YN
Cq O
E ER(D+ID(N+D) (I)
U2
where "1 NN F. E RD and Vl, VI UN E R
We are given a training dataset {X 1 consisting of P samples. All x tokens are independently sampled from a D-dimensional zero-mean normal distribution with covariance A,
As tool for our analysis, we show that when trained in-context linear regression tasks, ATTNM is equivalent to a two-layer fully-connected linear network with a cubic feature map as input, and ATTNs is equivalent to a sum of three-layer convolutional linear networks with the same cubic feature map as input.
Fp.n Lngq ~
m= , (2)
We consider the in-context linear regression task, where the Yn in context and the target output yy are generated as lin- ear map of the corresponding F, and 2. (Garg et al., 2022). For each sequence fr, we independently sample task vector w. from a D-dimensional standard normal distribu- tion, Wp ~ (0, I), and generate Yu.n = wn). F Up.q = wI æp..: n 1,... N, μ == 1, P. Note that the task vector w is fixed for all tokens in one sample sequence but varies across different samples, and is independent of the tokens xp,1,
We empirically demonstrate that the single and multiple loss drops also occur in softmax ATTNM and ATTNs, respectively.
Comparing the two models, we find that the ICL ability evolves differently in them: ATTNM acquires the in-context linear regression ability through one abrupt loss drop, while ATTNg acquires this ability by progressively improving on in-context principal component regression. This makes a theoretical case for the progressive improvements of ICI.
Dp,,N p..g.
2.2. Multi-Head Self-Attention
A standard multi-hear softmax self-attention layer (Vaswani et al., 2017) takes the matrix X as input and returns a matrix
2

---

Training Dynamics of In-Context Learning in Linear Attention
of the same size,
2.4. Linear Attention with Separate Key and Query
x x W x x
In multi-head attention with separate key and query, we follow the standard practice (Vaswani et al., 2017) of using low-rank key and query matrices where the rank R S. D and D. In practice, usually RH = D. The multi-head linear attention ATTNs with separate rank- R key and query matrices computes
ATTN(X)= = x _ W V X smax
P
where H is the number of heads, P is scaling factor, and W., W.K, W are the trainable value, key, and query matrices in the i-th head. The prediction for By is the bottom right entry of the output matrix:
RH 2
H Σ i=1
9q ATTN(X) D D+I,N+I-
(3)
ATTNs = x
W.
XX
W K W
In this work, we consider multi-head linear self-attention, where we remove the softmax operation and take p N. Specifically, we study two common parametrizations of linear attention: (i) linear attention with merged key and query introduced in Section 2.3 and analyzed in Section 3; (ii) linear attention with separate key and query introduced in Section 2.4 and analyzed in Sections 4 and 5.
We can write the value, key, and query weights in block form,
W.Y 1
W.K
w
k
kis.
The blocks have dimensionalities Ui, ki, E R and 9ir
vi qir E RD (r 1r=11 1.....R,, Similarly to the case with merged key and query, we initialize 0; = 0, kirr = O; they will remain zero throughout training (see Ap- pendix F. With v; = and kür = O, the multi-head linear attention with separate rank-one key and query matri- ces computes
2.3. Linear Attention with Merged Key and Query
The multi-head linear attention ATTNM with the key and query matrices in each head merged as a single matrix W.K W.O =w computes
H
ATTNM (X)= = X Σ
1 W. XX
X
R Σ
ATTNs (X)DNN(
G, ki. K
(S)
where the terms can be written in block form,
where 3 is the input-output correlation in context defined in Equation (4). The expression of Equation (S) already reveals interesting insight. It implies that linear attention with H heads and rank-R key and query differs from linear attention with RH heads and rank-one key and query only in the sharing of certain value weights.
[ T + +N En 2n
Ex anyi Σm=1 Un
XXt
Σm=1 EN 2n
and
W.V
KQ W
U; i
2.5. Gradient Flow Training Dynamics
The blocks have dimensionalities V; un E IR 0g E IR, U; E DxD The blocks denote entries that do not con- tribute to the computation of ATTN(X)) (TTN((XD+1.N411 D+1,N+1- Follow- ing Ahn et al. (2023); Zhang et al. (2024a); Kim & Suzuki (2024); Huang et al. (2024b), we initialize v; u = 0 as they are not required for this model to achieve global minimum loss on the in-context linear regression task. When v, and u, are initialized to zero, they will remain zero throughout training (see Appendix D.I). With the reduction v;, u, O, the prediction for yq. which is the bottom right entry of AX), is
We train the linear attention model using gradient descent on squared loss of the query token", that is L = E(u, )² We analyze the gradient flow dynamics on the loss, given by
dW T dt
Oc aw
00g ow
E (y
94)
(5)
where is the time constant. The gradient flow dynamics captures the behavior of gradient descent in the limit of a small learning rate.
3. Linear Attention with Merged Key and Query
(TTN Σ 3 U.a.y.
(M)
We first study multi-head linear attention with the key and query matrices merged as a single matrix, as described by Equation (M).
where (3 is the correlation between Fr and yn in context,
N
B
Ln.
(4)
N
"We can also handle next token prediction loss (Appendix A.3).
12
3

---

Training Dynamics of In-Context Learning in Lincar Attention
V
AT
1
N Σ
MLP
vec
m=1
U2
UH
Loss
Linear Attention
ve)(U,)
vec(UTg)
c.M.)
201
22
N
2000
4000
6000
1/1
M2
UN 0
ED
Training steps
N1 (X) linear network with cubic feature input, MLP(z). Left: Schematic of the equivalence. Right: Loss trajectories of linear attention and the fully-connected linear network match well. The two models are trained with the same data and initialization. Both exhibit the characteristic abrupt loss drop documented by prior work on the ICL dynamics in linear (Von Oswald et al., 2023) and softmax attention (Singh et al., 2024). Here D = 4, N = 31, H = 8.
Figure/ 1. Multi-head linear attention with merged key and query, ATT
is equivalent to a two-layer fully- cullycoonecctdd
3.1. Connection to A Fully-Connected Linear Network
There are two manifolds of fixed points in this dynamical system: one is the unstable fixed point at zero, denoted Mo. and the other is a manifold of stable fixed points at the global minimum, denoted M..
The H-head linear attention with input sequence X defined in Equation (M) can be viewed as a two-layer width-H fully cullycconeettd linear network with a cubic feature x xX) as input,
Mo = {w2 = 0, W1 = 0}
(9a)
H U,3yy
M, = { W2 W1 W
W2 W1 = E (ygz E
1(2x( (zz
))
)>1}}
(9b)
ATTN X D
Σ R
3.3. Training Dynamics: An Abrupt Drop in the Loss
H
Σ UUV) T vec (Bx,))
We have shown the linear attention defined Equation (M) is equivalent to a fully-connected linear network with cubic feature input. Since this equivalence holds the level of the computation of the model, the equivalence applies to the training dynamics with any initialization and optimizer. Here we discuss the training dynamics from small initial- ization, commonly referred to as the rich learning regime (Woodworth et al., 2020).
2 w MLP(z),
(6)
where
(U1 T vec(U2) (U2)
11
U2
W12 :
W1
xxX) = vec (Bx,)).
vec T
With small initialization, the network initially near the unstable fixed point, MM. at zero. As training progresses, the network escapes from the unstable fixed point, and sub- sequently converges to stable fixed point on the global min- imum manifold, M.. The time it takes to escape from the unstable fixed point is approximately TΛzニ In wail where the initialization scale Wini is the initial 8² norm of layer (see Appendix D.6.1). Because the time to escape from the unstable fixed point starting from small initialization is long, the loss exhibits an initial plateau followed by an abrupt drop, as validated by simulations in Figure 1. In particular, when the input token covariance is white A = I and the initialization is infinitesimally small, we exploit the equiva- lence between linear attention and linear networks derive analytical time-course solution (see Appendix D.5) and obtain
(7)
The feature E DR whose entries are cubic functions of the entries in the original sequence X, is the input to the equivalent two-layer fully-comnected linear network. The stacked value weights correspond to the second-layer weights W2 E R of the fully-connected linear network. The stacked merged key-query weights correspond to the first-layer weights W1 E R xx D² of the fully-connected linear network. A schematic of this equivalence is given in Figure 1.
3.2. Loss Landscape: Two Fixed Points
The gradient flow training dynamics of the linear attention or the equivalent two-layer collecconnected linear network given in Equation (6) is
ATTNM ( (1)))++ o Zqq.
J W we [E ( 2 WiE [xx] ))
(8a)
e2VDz
where o(1) (1+ 1P)
(10)
TW02 Wi (E (3n=T) - W W2E(22)) T))
(e2VT5: 19:511 + w.
YD=
(8b)
iii
4

---

Training Dynamics of In-Context Learning in Lincar Attention
Since (I) is a rescaled and shifted sigmoid function, the weights and the loss trajectories have sigmoidal shapes,
9/
characterized by plateau followed by rapid drop.
q
3.4. ICL Algorithm: Least Squares Regression
When the linear attention model converges to the global minimum manifold MM at the end of training, the model implements
22
24
85
83
24
ATTNM (X) D+1,N+1 = = =VV ") E (zz)
Figure 2. Multi-head linear attention with separate rank-one key and query ATTNs D+) is sum of H (number of heads) three-layer convolutional linear networks with the cubic feature as input. Here we take D 3 to avoid clutter. Entrics in the vectors are denoted as Fy [ 2 m²] ,3 [3 8
= = [ ( -1
(11)
Zgg
N
where the first equality follows directly from Equations (6) and (9b) and the second equality proved in Appendix D.4. Equation (1I) reveals an intriguing duality: the linear regres- sion solution in the cubic feature space of is the in-context linear regression solution in the original space of the n,, U/ token pairs in sequence X. The first line of Equation (11) is the linear regression solution of fitting ya.. with Zμ for all training sequences μ 1, ....... P. The second line of Equation (11) is approximately the in-context linear regres- sion solution, which fits yuun with æn,n (n = I,.. ................................ ,N) for each sequence X,.- When the sequence length N is large, the model recovers the inverse of the true covariance matrix,
rewritten as
H
ATTNs
X
JDIIN-I
::q; K,z,
k 0
0
k
0
where K;
ER DXD*
(13)
0
k
D
lim BT (A+ A 1A+A+EHH ( -1 N00
Ty = BT x...
The matrix K, is convolutional matrix with kernel size D and stride D. A schematic of the three-layer convolutional linear network is given in Figure 2.
N
Here 33r is the En. In correlation in a sequence X, and A the covariance of all a. tokens all training sequences, which approximates the covariance of æn in each individual sequence.
When the number of heads satisfies H 2 D, the linear at- tention with separate rank-ont key and query, ATTNs(X), can express any linear map of 2 (X) and has the same ex- pressivity as linear attention with merged key and query, ATT (X). However, the two models correspond to multi- layer linear networks with different connectivity and depths, resulting in different loss landscape (Kohn et al., 2022; 2024) and training dynamics (Saxe et al., 2014; 2019).
4. Linear Attention with Separate Rank-One Key and Query
We now study multi-head linear attention with separate low- rank key and query matrices. Because the rank-one case captures most of the behaviors of the general rank- R case, we focus on the rank-one case in this section and defer the rank- R. case to Section 5. When = 1, the model definition in Equation (S) simplifies to
4.2. Loss Landscape: Exponentially Many Fixed Points
The gradient flow training dynamics of linear attention with separate rank-one key and query, derived in Appendix E.2, is given by
H Σ
(TTNs(X))4LLAT
TT,Tk.qq k.g. Zg.
(12)
H
k
A
E
Σ
Uy k.q. A
q:
(14a)
4.1. Connection to Convolutional Linear Networks
H1
Tk Di
A E A2
Σ
Up ki q,7 A
qi. (14b)
The H-head linear attention with separate rank-one key and query can be viewed as a sum of H three-layer convolu- tional linear network with the cubic feature z defined in Equation (7) as input. Specifically, Equation (12) can be
V=1
H
T4: vi
A A Σ
Up ky k.g. E (A2)
ki, (14c)
5

---

Training Dynamics of In-Context Learning in Lincar Attention
Simulations
Theory
1.0
(()
1.0
Eigenvectors of A
04
Simulated 4
Theory 4
Loss
((M2)
0.8
0.0
((M2)
0.4
Key weights
(t))
-02
(LMs)
0.21
((M2)
M₁
M2
M2
M4
0.0
0.0
5000
10000
5000
10000
Training steps
Training steps
Training steps
(a) Loss
(b) Loss and value weights
(c) Key weights and eigenvectors
Figure 3. Multi-head linear attention with separate rank-one key and query exhibits sadle--o-saddle dynamics. (a) The loss curve has D abrupt drops, separated by plateaus (six from different random initialization plotted). The loss at each plateau matches theoretical prediction in Equation (19) (dashed gray lines). (b) The value weight v in each head for one the runs in (a) is plotted in solid blue curves. The numerical solutions of from Equation (21) are plotted in dashed blue curves and match the simulations well. The shades of blue distinguish different heads. (c) The key weights during the loss plateau are plotted color. When the model moves from one fixed point the next, the key weight a head, k, aligns with new eigenvector of the input token covariance A. The key weights kis and the cigenvectors 6144 are rows in the heatmaps. A video of the dynamics is provided at URL. Here D = 4, N 31, H = 4, and has eigenvalues 0.4, 0.3, 0.2, 0.1 and eigenvectors as plotted in (c).
Since there are (I) possible ways of choosing m out of indices define Sn in Equation (17), the total num- ber of possible choices summed over m 0,... D is Σ (m) 2D Each choice corresponds to different condition (CI) in Equation (18) and thus different function, ATTNs (X) ( D+1.N++1- Hence, the gradient flow dynamics This dynamical system contains 2n fixed points in the func- tion space of ATTN,(X))DLL,) D+INN++1 We specify the fixed points below and prove their validity in Appendix E.3.
where we denote the in-context covariance of an tokens as nmll C /N and the expectation of Á is
Á Di
2 A A (A) A
(15)
in Equation (14) has 2D fixed points the function space.
The two fixed points of ATTNM (Section 3.2) are contained in the 2n fixed points of ATTNs: the zero fixed point in Equation (9a) corresponds to M(So). i.c., learning no eigen- vector; the global minimum fixed point in Equation (9b) corresponds M(Sp), i.e., learning all D eigenvectors.
Let '..... An be the eigenvalues of the covariance matrix A arranged in descending order, and "1.-.. ep be the corresponding normalized eigenvectors. We use M(S,)) ) to denote a set of fixed points that correspond to learning m
(m 0, 1, ... DD out of the D eigenvectors,
4.3. Training Dynamics: Saddle-to-Saddle Dynamics
S( = {v, k. 4)1. condditins (CI)-(C3)), (16)
Building the exponentially many fixed points have identified, we now analyze which fixed points are actually visited in gradient flow training and in what order. We find that starting from small initialization, the model visits (D + 1) out the 2D fixed points.
where the set Sm specifies the indices of the learned eigen- vectors,
Sn S (1,2, ,D) |Sm.] m.
(17)
With small initialization, the model initially near the unstable zero fixed point, MM M(Ø). As training progresses, the model sequentially visits the fixed points in M2, ,M,, where Mn 2M ,m]). That is, the model trained from small initialization sequen- tially learns to fit the first eigenvector (the eigenvector of with the largest eigenvalue), the second eigenvector, and so on. As shown in Figure 3a, the loss goes through / abrupt drops in training, each corresponding the transition from one fixed point to the next. The abrupt drops of loss separated by plateaus, during which the model lingers near
The three conditions for Equation (16) are:
(C1) The heads sum to fit the eigenvectors with indices in the set Sm
u.kaq
Σxa deS...
1
edey
N
(18)
(C2) For heads with a nonzero value weight, Us # 0, both k; and qi lie in the span of (ea)des...
(3) For heads with zero value weight, v, = 0, at least one of k; or qi lies in the span of {ea]aes...
A fixed point function space corresponds set of fixed points weight space that implement the same input-output map.
6

---

Training Dynamics of In-Context Learning in Linear Attention
an unstable fixed point. Because the time required for head to learn the eigenvector e. from small initialization scales with A (see Appendix E.6), eigenvectors associated with larger eigenvalues are learned faster. This explains why the model learns to fit the eigenvectors sequentially in descend- ing order of the eigenvalues, as well as why we empirically see the later plateaus last longer in Figure 3a.
plateau and the subsequent abrupt drop of loss reduces to 1:
ordinary differential equation about U; (t), m +
(1 1 1+
(r))r/(A////) (A)/Am+t Am
rū; = Xm+1 am-1
ug.
N
(21)
Equation (21) is separable differential equation but does not admit general analytical solution of (m) in terms of & (see Equation (71)). Nonetheless, it greatly simplifies the high-dimensional dynamics in Equation (14) and pro- vides good approximation of the true dynamies: during each plateau and the subsequent abrupt loss drop, weights in one of the heads grow in scale with the key and query weights aligning with the next eigenvector, while the rest of the heads remain approximately unchanged. In Figure 3b, we compare the numerical solution of Equation (21) with the value weights trajectories in the simulation and find excellent agreement.
When the model is at a fixed point in Mms we compute the loss in Appendix E.4 and obtain
1
[(Mm) = tr(A) - Σ Ad 1+
N
d=1
(19)
Equation (19) is highly interpretable in the limit of a large sequence length N. The loss, [(Mm). is the sum of the eigenvalues associated with the remaining unlearned eigen- vectors
In summary, the loss trajectory of linear attention with sepa- rate rank-one key and query trained from small initialization exhibits D abrupt drops, each followed by plateau. The amount of the m-th abrupt loss drop (1 S m S. D) is ap- proximately the eigenvalue Am. during which the key and query weights in an attention head grow in scale and align with the eigenvector C.n-
lim ((Mm)= = tr( N-+00
Aa
Σ d=m+11
Aa
Thus, the loss decreases by approximately Am during the m-th abrupt loss drop. We plot Equation (19) dashed gray lines in Figure 3a and find they match the plateaus of simulated loss trajectories well.
4.4. ICL Algorithm: Principal Component Regression
When the model reaches Mm from small initialization, its weights take on a highly structured form, which is a spe- cific instance of the general definition in Equation (16). As shown in Figure 3c, the key and query weights in a head grow in scale and align with new eigenvector of the input token covariance A during each abrupt loss drop. Based on simulations in Figure 3 and derivations in Appendices E.5 and E.6, we propose an ansatz that during the (m + 1)-th plateau (0 < m < D) and the subsequent abrupt drop of loss, the weights are approximately given by
When the linear attention model is at a fixed point in Mm. based on Equation (18), the model implements
ATTNs( X) D+IIN+1
Σ 17 1 1+tr(A)/Aa d=1
e Bng (22)
N
In the limit of large sequence length N, Equation (22) simplifies and can be interpreted as principal component regression in context with m principal components
k; = = U.e,, U/= = =1,² (1 1+ (rA)/A -
lim ATTN; = w N
ede ag.
N
isism, m, k; = q: = 1z (1)mmt,, i= m +1,
(20a)
Here is the task vector for the sequence X, and En eae, x. is query input T. projected onto the first m principal components. Hence, if training stops during the (m + 1) plateau, the linear attention approximately implements the principal component regression algorithm in context with mn principal components.
(20b)
k; q. 0, 0, m < < H,
(20c)
where v (1) is small during the (m + 1) loss plateau and grows during the (m+ 1)-th abrupt loss drop. Equa- tion (20) implies that the e norms of Us k2 qi in a head are equal, which is consequence of small initialization and the conservation law in Appendix E.8. With this ansatz, the high-dimensional training dynamics during the (m + 1)
After the model has undergone 1) plateaus, converges to the global minimum fixed point, MMD and approximately implements principal component regression in context with all D components, which is least square regression. Thus, the linear attention model with either merged or separate key and query undergoes different training dynamics but converges to the same global minimum solution.
"We permute the heads so that the head aligned with the d-th eigenvector have index d. The signs of any two among vi ki, qi can he flipped with trivial effect on the analysis.
7

---

Training Dynamics of In-Context Learning in Lincar Attention
R Simu ations
R Simulations
R
R
Theory
Theory
Smulltions
Theory
Simulations
Theory
10
1.0
1.0
1,0
0.5
00
0.0
0.0
0.0
20
Training sleps c (xx)
Training sleps I (x10%)
Training sleps (x10')
Training steps i (x100)
Figure 4. Multi-head linear attention with separate low-rank key and query exhibits saddle-to-san to-saddle dynamics, with the duration of platcaus depending on the rank R. Solid black curves are loss trajectories from six random initializations. Dashed gray lines mark the loss values predicted by Equation (19) at nine fixed points, which are [(Mo),, ((Ma),-.. L(Ms) from top to bottom. The four panels differ only in the rank of the key and query weights. Here D 8, N = 31, H = 9, A has trace 1 and eigenvalues da xx d=
5. Linear Attention with Separate Low-Rank Key and Query
6. Related Work
Recent theoretical research on linear attention has investi- gated its expressivity (Vladymyrov et al., 2024; Gatmiry et al., 2024), learnability (Yau et al., 2024), loss landscape (Mahankali et al., 2024; Li et al., 2024), convergence (Zhang et al., 2024a;b; Ren et al., 2024; Fu et al., 2024). and general- ization (Wu et al., 2024; Mahankali et al., 2024; Duraisamy, 2024; Abedsoltan et al., 2024; Lu et al., 2025; Frei & Vardi, 2025). The seminal work by Zhang al. (2024a) analyzed the gradient flow training dynamics of linear attention to prove convergence guarantees, showing what the model con- verges to at the end of training. Our work also analyzes the gradient flow training dynamics but goes beyond cxisting convergence results to describe the entire training dynamics. Moreover, we study multi-head attention with merged or separate key and query weights, while Zhang et al. (2024a) focused on single-head attention with merged key and query.
The linear attention model with separate rank-R key and query shares many behaviors with its rank-one counterpart. For loss landscape, linear attention with rank-R R key and query has the same 2D fixed points in the function space as its rank-one counterpart, corresponding to the model imple- menting in-context principal component regression with subset of all D principal components (see Appendix F.3).
For training dynamics, the loss trajectories differ slightly, depending on the rank R. We plot the loss trajectories with input token dimension
= and different ranks R 1, 2, 4, 8 in Figure 4 (see Figure for R = =,,5,, 6, 7). For R 1, the loss exhibits plateaus at eight values L(MM) (m 0, 0........7 7). plateaus at four values L(Mm) (m 0, 2, 4, 6), and ei- ther brief plateaus or no plateau at the other four values. For R 17=4, 4, the loss exhibits conspicuous plateaus at only two values L(MM) (m = 0, 4). To summarize, with rank- R key and query, the loss trajectory exhibits conspicuous plateaus at value ((Mm) for m that divides R.
For R 2, the loss exhibits
Another line of recent research on the training dynamics of softmax attention models has shown stage-wise dynam- ics. Due to the intractability of softmax attention training dynamics in general, many of these studies made strong assumptions to enable theoretical analyses, including a sim- plified layer-wise training algorithm in place of standard gradient descent (Tian et al., 2023; Nichani et al., 2024; Chen et al., 2024c; Wang et al., 2024a), restricted weights (Boix-Adsera et al., 2023; Chen et al., 2024b; Rende et al., 2024; Edelman et al., 2024), and specifically chosen datasets (Huang et al., 2024b). In comparison, our work leverages the linear attention model without the softmax operation, enabling us to study in fine detail the dynamics of standard gradient descent training without restrictions on weights, Namely, we derive an analytical time-course solution and reduce the high-dimensional dynamics to one-dimensional ordinary differential equations for the two models we study, respectively. Furthermore, we characterize how parametriza- tion (i.c., merged or separate key and query, and rank of the separate key and query weights) affects the loss landscape and training dynamics, an aspect not previously examined.
The difference in the loss trajectories arises from the struc- ture of the model defined in Equation (S). Each atten- tion head has a single value weight &s that is associated with all R pairs of key and query weights in that head, k Irr (r fr=1, 1, RR). During a conspicuous plateau, a new value weight escapes from the unstable zero fixed point and grows in scale. Once the value weight has grown, it leads to larger gradient updates for all the key and query weights in that head, speeding up their escape from the zero fixed point. Hence, in the rank- R case, conspicuous plateau occurs when m divides Rt, corresponding to learning a new head from small initialization. Brief or no plateau occurs when m does not divide R, corresponding to learn- ing new pair of key and query weights in a head whose value weight has already grown, as shown in Figure 11. See Appendix E.4 for further details.
8

---

Training Dynamics of In-Context Learning in Lincar Attention
7. Discussion
Simulations
Theory
1.0
1.0
CLM))
We studied the gradient flow training dynamics of multi- head linear attention and demonstrated how it acquires ICL abilities in training. We begin with a simple setting of linear attention with merged key and query trained for in-context linear regression, following the setting in seminal works (Von Oswald et al., 2023; Ahn et al., 2023; Zhang et al., 2024a). We show an abrupt loss drop in training and give an analytical time-course solution in the case of a white input token covariance and small initialization. However, single abrupt loss drop does not fully capture the evolu- tion of ICL in training practical transformers, where the abilities continue to develop throughout training (Xia et al., 2023; Park et al., 2025). We thus extend our analysis to parametrization closer to the attention in practical transform- ers: attention with separate key and query. In the separate case, we find that the loss exhibits saddle-to-see sa-dde dynam- ies with multiple abrupt drops. The ICL ability evolves progressively, manifesting as implementing principal com- ponent regression in context, with the number of principal components increasing over training time. We thus charac- terize how the linear attention model develops increasingly sophisticated ICI. abilities in gradient descent training.
0.8
0.8
0.6
0.6
(LM)
0.4
((M))
0.0
5000 Training :2pss
10000
0.0
10000
23000
Training stops
(a) Softmax ATTNM
(b) Softmax ATTNs with R =
Figure 5. Loss trajectories of softmax attention with merged or sep- arate key and query. Six runs from different random initialization are plotted. Similar to the lincar attention case, softmax ATT Na exhibits one abrupt loss drop, while sofumax AT TTN exhibits mul- tiple loss drops. The dataset and model setup are the same as Figures 1 and 3 except adding the softmax activation function.
Dynamics of In-Context and In- Weight Learning. Our work studies the training dynamics of in-context learning. Other than in-context learning, attention models can also learn in weight; that is, solving the task by memorizing the map between the query input and the target output without using the information in context. The arbitration between in- context and in-weight learning may depend on properties of training data (Chan et al., 2022). To focus on the dynamics of ICL, we used a purely ICL task, which is in-context linear regression with the task vector sampled from a zero- mean standard normal distribution, w (O, I). Since memorizing any particular task vector does not effectively decrease the loss, linear attention develops only in-context learning ability during training, as shown in Figure 14a. If the task vector w follows a different distribution, the training dynamics involves the development of both in-context and in- in-weight learning abilities, as shown in Figure 14. We provide more details in Appendix G.
Softmax Attention. We empirically find that the different training dynamics of linear ATTNM and linear ATTNs also occur in their softmax counterparts. Figure 5 follows the same setup as Figures and 3 for linear attention, with the only difference being adding the softmax activation func- tion for the attention calculation. We observe that softmax attention with merged key and query exhibits a single abrupt loss drop, whereas softmax attention with separate rank-one key and query undergoes multiple loss drops, separated by phases of conspicuously slower training. This suggests that our findings and theoretical intuition are not unique to linear attention but may also extend to softmax attention.
In our analysis, we draw connections between linear attention and multi-layer lin- ear networks, enabling us to employ the rich theoretical machinery built for linear networks to understand linear attention training dynamics. Beyond training dynamics, many other theoretical results for linear networks can apply to linear attention through the equivalence we draw. For example, the convergence guarantee for multi-head linear attention trained on in-context linear regression tasks can be obtained from the convergence proofs for deep linear networks (Arora et al., 2019; Shamir, 2019). In contrast, without the equivalence, Zhang et al. (2024a) previously ob- tained a convergence guarantee for single-head linear atten- tion, which required highly non- non-trivial derivations. Hence, we believe the connections draw are useful in enabling the applications of theory from one architecture to the other.
Implications for Future Theory.
Effect of Initialization. Having analyzed the small initial- ization case, we now examine how the initialization scale affects training dynamics. For linear ATTNM, Figure 6a shows increasing initialization shortens the plateau before the single abrupt loss drop. For linear ATTNs, Figure 6b shows increasing initialization shortens all platcaus between successive abrupt loss drops. At the largest initialization, both models exhibit an exponential-shaped loss decay - a hallmark of lazy learning (Chizat et al., 2019). In con- trast, rich learning typically exhibits abrupt sigmoid-shaped loss curves as seen in our main result. Theory typically focuses on either the lazy or rich regime, while practical initializations often fall in between. In Figure 6, dynamics from the intermediate initialization seems like mix of the exponential seanential-shaped shaped and higmoid-shaper curves, which are often seen in practice, e.g. in induction head emergence in natural language settings (Olsson et al., 2022, Argument I).
Additionally, we have shown that parametrization signif- icantly affects the loss landscape and training dynamics, motivating future research to examine how their results may or may not be influenced by the parametrization choice.

---

Training Dynamics of In-Context Learning in Lincar Attention
Acknowledgement
httpps:://oeeeeeeee....t.. http://pppneevi..ee//.e/ o orum? orum??1-=-jsmm3405500
lations, 2025. URL https:
We thank Jin Hwa Lee, Sara Dragutinovié, Andrew Lampinen, Basile Confavreux and William Tong for helpful conversations.
Anwar, U., Oswald, J. V., Kirsch, L.., Krueger, D., and Frei, S. Adversarial robustness in-context learning in transformers for linear regression, 2024. URL : / //arxiv. wttp:///r///////1.....o 05189.
We thank the following funding sources: Gatsby Charitable Foundation (GAT3850) to Y7., AKS, PEL, and AS; Well- come Trust (1101 14/Z/15/Z) to PEL; Sainsbury Wellcome Centre Core Grant from Wellcome (219627/2/1//2) to AS; Schmidt Science Polymath Award to AS. AS is a CIFAR Azrieli Global Scholar in the Learning in Machines & Brains program.
Arora, S., Cohen, N., Golowich, N., and Hu, W. A conver- gence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Rep- resentations, 2019. URL https: /openreveew.nnn et et st/fooum?la Siterummmdd-k00gg33555
Ataee Tarzanagh, D., Li, Y., Zhang, X., and Oymak, S. Max-margin token selection in attention mechanism. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 48314- 48362. Curran Associates, Inc., 2023. URL https: new.iiss.nngnnuuusee.....s //proceedings nwwrinsuui.... wwtpppps.cccpppeeeeeee wwtp:pteepppppppeeeeeeeeeee. 1 faapper/2083/1lee///059922222222
Impact Statement
This paper presents work whose goal is to advance the field of machine Icarning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Abedsoltan, A., Radhakrishnan, A., Wu, J., and Belkin, M. Context-scaling versus task-scaling in in-context learning, 2024. URL 2783.
Atanasov, A., Bordelon, B., and Pehlevan, C. Neural net- works as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2022. URL
net forum
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Sys- tems, volume 36, pp. 45614-45650. Curran Associates, Inc., 2023. URL /Rthhttss://pprrceeddg.... neurip wwtp//ppee__illlepppppe////////e
?id= 1Nv-1gddooooom
Bai, Y., Chen, F., Wang. H., Xiong, C., and Mei, S. Trans- formers as statisticians: Provable in-context learning with in-context algorithm selection. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, (eds.), Advances in Neural Information Processing Sys- tems, volume 36, pp. http:/tpsss//ppcceeeeeigeee .nurip fi
Curran Associates, s cc s:cc:pppperr__filsspppp/r/////// 63a36c57ee13300015ecc22222222 onference.
Bed
Inc., 2023. URL ht
3d61 onference pdf.
Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A., and Sra, S. Linear attention is (maybe) all you need (to understand transformer optimization). In The Twelfth International Conference on Learning Representations, hete/fort
Boix-Adsera, E., Littwin, E., Abbe, E., Bengio, S., and Susskind, J. Transformers learn through gradual rank increase. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 24519 24551. Curran Associates, Inc., 2023. URL https:
2024. URL /tttps://openreeeee Pid-OuI5415ry7.
Akyürek, E., Schuurmans, D., Andreas, J.. Ma. T., and Zhou, D. What learning algorithm in-context learn- ing? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https:////.. wpw.revieneeeeeeeeoeeeeeeeoo net/ forum ? ?id=0g0X4H8yN41.
[2aper/0005//iee/666620000282800 Ba7b7laaee331-Paae--oofforee
Boix-Adserà, E., Saremi, O., Abbe, E., Bengio, S., Littwin, E., and Susskind, J. M. When can transformers reason with abstract symbols? In The Twelfth International Conference on Learning Representations, 2024. URL https: http::///ppezeev... net foot/fooumm SooorumiidSSTu STTG FUz8ob.
Anand, S., Lepori, M. A., Merullo, J., and Pavlick, E. Dual process learning: Controlling use of in-context vs. in- weights strategies with weight forgetting. In The Thir- teenth International Conference on Learning Represen-
10

---

Training Dynamics of In-Context Learning in Lincar Attention
Bordelon, B., Chaudhry, H., and Pehlevan, C. Infinite lim- its of multi-head transformer dynamics. In Gloherson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tom- czak, J., and Zhang, C. (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp. 35824- 35878, Curran Associates, Inc., 2024. URL https: /
Chen, S., Sheen, H., Wang, T., and Yang. Z. Training dynamies of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. In Agrawal, S. and Roth, A. (eds.). Proceedings of Thirty Seventh Conference on Learning Theorr, volume 247 of Proceedings of Machine Learning Research, pp. 4573- 4573. PMLR, 30 Jun-03 Jul 2024b. URL https: : /pr oceedings mlr wwwppppeessssseeeeee
html.
Chen, S., Sheen, H., Wang, T., and Yang, Z. Unveiling induction heads: Provable training dynamics and fea- ture learning in transformers. In Globerson, A., Mackey, L.., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang. C. (eds.), Advances in Neural Information Pro- cessing Systems, volume 37, pp. (6479-665677 Curran Associates, Inc., 2024c. URL fitps_/passsseaae/222222 ngs www.inurppp..m wc/paperppppeccecceeeceee.eeee.e 4/ A/file/ hFFFie//aaeeaaeec222244485222222 441 441-Paper-Conference.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D.,
Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G.,
Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H.,, Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020, URL https: http:///ppccceednggeeee neurips.esppppp.pp cc paper files filess/pepeer/2000//llee////ee 0d644500006dcbcbbb 6b4 pdf.
/
Chizat, L.., Oyallon, E.,, and Bach, F. On lazy training in dif- ferentiable programming. In Wallach, H., Larochelle, H., Beygelzimer, A., TAlché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL fitee//pppe//0005f//e/eeee463333
Chan, B., Chen, X., György, A., and Schuurmans, D. To- ward understanding in-context vs. in-weight learning. In The Thirleenth International Conference on Learning Representations, 2025. https:/ /openreveeen .net/ thet/forunte ineimiiddd-eee aKdkkKnNNNN8..
pppper (df5266b29657222555---888
Du, S. S., Hu, W., and Lee, D. Algorithmic regular- ization in learning deep homogeneous models: Layers are automatically balanced. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Gar- nett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 31. Curran Associates, Inc., 2018 URL URLtttps://prrcceednggg..s neurps cctpppsspcppccccoocoooooppo pper_files/paree pper_fflles/pppee//00//////e//// a Duraisamy, K. Finite sample analysis and bounds of gen- eralization error of gradient descent in in-context linear regression, 2024. URL https: wwwpaaxxvvvv..... abs/ 2405.02462.
Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, F. Data distri- butional properties drive emergent in-context learning in transformers. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neurul Information Prucessing Systems, volume 35, pp. 18878-18891 Curran Associates, Inc., 2022. URL http http://cceddins..euurrseeeeec//p es /t/ppppe/002///ia////cccaaaf///2 07 P
Paper pdf.
Chen, A., Shwartz-Ziv, R., Cho, K., Leavitt, M. L., and Saphra, N. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in MLMs. In The Twelfth International Conference Learning Represen- tations, 2024a. URL https: httt:::://pppereee hpwp/eveeeeeeeeeeeee//eeeeeeeeee ew net /forum? Meedd-MMPiiMMM KHELW..
Edelman, E., Tsilivis, N., Edelman, B., Malach, E., and Goel, S. The evolution of statistical induction heads: In-context learning markov chains. Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tom- czak, and Zhang, C. (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp. 64273- 64311. Curran Associates, Inc., 2024. URL https: / //proceedings. wwwpiiinnnnneeuuusspeeee.ee.. wwwpouurriscccppppppeeeee cappaperr_pplplpeeeee
Chen, S. and Li, Y. Provably learning a multi-head attention layer. In Proceedings of the 57th Annual ACM Symposium on Theory of Computing, STOC '25, pp. 1744-1754, New York, NY, USA, 2025. Association for Computing Machinery. ISBN 9798400715105.. doi: 10.1145/3717 823.3718174. URL hatpps:/ / www..oorg.......oo 0717823.371812
4d0e8ff446bc1 1aceee6cccPPpppeeee Pwwer-eoofeeeeneeeeeee
Frei, S. and Vardi, G. Trained transformer classifiers gen- eralize and exhibit benign overfitting in-context. In The

---

Training Dynamics of In-Context Learning in Lincar Attention
Thirteenth International Conference on Learning Repre- sentations, 2025. URL https:/ /opnnevveeweee.ee//o// net. SoorumiddjjsPSSB0
Mackey, L.,, Belgrave, D., Fan, A., Paquet, U., Tom- czak, J., and Zhang, C. (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp. 13244- 13273. Curran Associates, Inc., 2024. URL https: / //proceedi ngs. www.:neeriiiscccm wwt.aarripp..cppppeeeeeeee
Fu, D., Chen, T.qq., Jia, R., and Sharan, V. Transform- ers learn to achieve second-order convergence rates for in-context linear regression. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Pro- cessing Systems, volume 37, pp. 98675-98716. Curran Associates, Inc., 2024. URt.bttpss:/pppcceedde ssterrriss.cc/ppppppeeeee///pppp
b136f1 wwwe9dda-peeeeeeee Paw.s--apee----eea..ee Compeerr-oooffeencee pdf.
Hoffmann, D. T., Schrodi, S., Bratulić, J., Behrmann, N.,
Fischer, V., and Brox, T. Fureka-moment, in transform- ers: Multi-step tasks reveal softmax induced optimization problems. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 18409-18438. PMLR, 21-27 Jul 2024. URL https: /ttpp//rooceeggaeeeee/ ml r. pressss/ / v225/223355/// ho fmann 24 html.
file/b2d4 #ile/2244011100//033277111deeed2 e e-Paper-Con pdf.
Fukumizu, K. Effect of batch learning in multilayer neural networks. Gen, 1044):EE-33, 1998.
Garg, S., Tsipras, D., Liang, S., and Valiant, What can transformers learn in-context? a case study of simple function classes. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 300833333598 Curran Associates, Inc., 2022. URL httpp://prroceednnn...eu........ www.:euerrippp.ppppppppe..eepe / sft.eee//ppeer/2222////l//222222 1aa86ff7.aee99000---ppp---ooooe pdf.
Huang, J., Wang. Z., and Lee, J. D. Transformers learn to implement multi-step gradient descent with chain of thought. In The Thirteenth International Conference on Learning Representations, 2025a. URL https: /op enreview net/ fontt/foumm?? r 3DD
Huang, R., Liang, Y., and Yang. J. Non-asymptotic conver- gence of training transformers for next-token prediction. In Globerson, A., Mackey, L.,, Belgrave, D., Fan, A., Pa- quet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 80634-80673 Curran Associates, Inc., 2024a. URL pttts::/prooceedigg..o htwp:iinnnnnnuuuuiee wwwrips. ca/pppe _ httpa//pappe///02/////e///730000 0dd66 Baf78850448Bccb8882222 pdf.
Gatmiry, K., Saunshi, N., Reddi, S. J.. Jegelka, S., and Kumar, S. On the role of depth and looping for in-context learning with task diversity, 2024. htLps//:: xiv. xiv.org/ abs/ 2410 21698.
Geshkovski, B., Koubbi, H., Polyanskiy, Y., and Rigollet, P. Dynamic metastability in the self-attention model, 2024. URL pattps://arxidn org. abs/ 2410. 06833.
Huang, Y., Cheng, Y., and Liang, Y. In-context convergence of transformers. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 19660-19722. PMIR, 21-27 Jul 2024b. URL htwps: http///oooeeeeenn..ce/ .ml r. r.press/ / /22//2233///// hht.////h/2g22///
Gopalani, P., Lubana, S., and Hu, W. Abrupt learning in transformers: A case study on matrix completion. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U.. Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 55053- 55085. Curran Associates, Inc., 2024. URL https: /
Huang, Y., Wen, Z., Chi, Y., and Liang, Y. A theoreti- cal analysis of self-supervised learning for vision trans- formers. In The Thirteenth International Conference on Learning Representations, 2025b. URL https: 6Uovh.
He, J., Pan, X., Chen, S., and Yang, Z. In-context linear regression demystified: Training dynamics and mechanis- tic interpretability of multi-head softmax attention, 2025. URL https:://// : /ttpxaaxvvv///////////////..aav. org/ abs/ 25w../55.... 12734.
/ //ppppeevieww../ net http://ooummmmrieeie
Ildiz, M. E., Huang, Y., Li, Y., Rawat, A. S., and Oymak, S. From self-attention to Markov models: Unveiling the dynamics of generative transformers. In Salakhutdinov,
He, T., Doshi, D., Das, A., and Gromov, A. Learning to grok: Emergence of in-context learning and skill com- position in modular arithmetic tasks. In Globerson, A.,
R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett,
J., and Berkenkamp, (eds.), Proceedings of the 41st
12

---

Training Dynamics of In-Context Learning in Lincar Attention
International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp- 20955-20982. PMLR, 21-27 Jul 2024. URL https: //proceedings. mlr. wrw.sprsssseeee..... 2trres///235/1ddi2222 12112224. html.
Kohn, K., Merkh, T., Montúfar, G., and Trager, M. Ge-
ometry of linear convolutional networks. SIAM Jour- nal on Applied Algebra and Geometry, 6(3):368-406, URL https:
2022. doi: 10. 1 0 37/21112111 441183.
Jang, U., Lee, D., and Ryu, E. K. LoRA training in the NTK regime has no spurious local minima. In Salakhut- dinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the st International Conference on Machine Learning. volume 235 of Proceedings of Machine Learning Re- search, pp. 21306-21322 PMLR, 21-27 Jul 2024. URL https: : http://pooceedeeieeeee ml wwwps/ 2og22d.htmi.
Kohn, K., Montúfar, G., Shahverdi, V., and Trager, M. Func- tion space and critical points of linear convolutional net- works, SIAM Journal on Applied Algebra and Geometry, 8(2):33 362, 2024. doi: 10. 10.1137/33M1565504. URL https: htt::/dd..rrrr 10. 11 2.M1 565504.
Lee, I., Jiang, N., and Berg-Kirkpatrick, T. Is attention re- quired for ICL? exploring the relationship between model architecture and in-context learning ability. In The Twelfth International Conference on Learning Representations, 2024. URL https: /wwp/rppeevvieeeeeeeeeeeeee net forum ?id QwgAcLLtoX.
Jelassi, S., Sander, M., and Li, Y. Vision transformers prov- ably learn spatial structure. In Koyejo, S.. Mohamed, S.,
Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.),
Li, Y., Rawat, A., and Oymak, Fine-grained analysis of in- context linear estimation: Data, architecture, and beyond. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 133324-138344.. Curran Associates, Inc., 2024. URL https: / https::/procceediiggg neurips cc cc/paper /lLee//paer//02///Eee//444333333 Pwwe-PPaaprr---eeeee
Advances in Neural Information Processing Systems, vol- ume 35, pp. 37822-37836. Curran Associates, Inc., 2022. URL / www.:eeedddnggs.uuuuus...... cc/p opee__filee/apper/222///LLLLLL Siea0000663333321555113333--3-23 nce.
Ji, Z. and Telgarsky, M. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019. URL https: /open review. .et wtt.::eeeeeeeueeeeeeeeeeeee
pdf.
Lu, Y. M., Letey, M., Zavatone- Veth, J. A., Maiti, A., and Pehlevan, C. Asymptotic theory of in-context learning by linear attention. Proceedings of the National Academy of Sciences, 122(28):2202259912) 2025. doi: 10.1073/pn as as.2502599122 URL https: / /wwww pnas. org/d / / /10.000000 107./paas...... 2502599122
Jang, J., Huang, W., Zhang, M.. Suzuki, T., and Nie, L. Un- veil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. In Glober- son, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 135464- 135625. Curran Associates, Inc., 2024. URL https: ierroceedings.nneurippeo cc/ wwtp:percpppppppeeeeeteeteeeeeee files /ppperr/022//ELle444227711111111
Mahankali, A. V., Hashimoto, T., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations, 2024. URL http:s///ppeeeeieeee net. forum ?id= 1td=8p3fuu61kc.
Julistiono, A. A. K., Tarzanagh, D. A., and Azizan, N. Opti- mizing attention with mirror descent: Generalized max- margin token selection, 2024. URL
Makkuva, A. V., Bondaschi, M., Girish, A., Nagle, A., Kim, H., Gastpar, M., and Ekbote, C. Local global: Learning dynamics and effect of initialization for transformers. In Globerson, A., Mackey, L., Belgrave. D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 86243 86308. Curran Associates, Inc., 2024. URL https:
1455.
Kim, and Suzuki, T. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the atten- tion landscape. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 24527-2444 24561. PMLR, 21-27 Jul 2024. URL ht attpps: / http://pccceeenn..ceee ml r. r.press/ v235 //////imm kim24af html.
/ //proceedings.oo neurips wttp::pppp.ccppppr_rieeeele haape//2222//iieeecccafaaftt 2 -oa.eeppeeeoooeeeeeeeeeeeeee
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Stein-
hardt, J. Progress measures for grokking via mechanistic
13

---

Training Dynamics of In-Context Learning in Lincar Attention
interpretability. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: woopenneviee.net/// net forum? (ettoorummiiii----- w0w.d
/ //proceedings....o neurips cc/paper_filees http://cng...noirsopoooooppeeeee / b240bc83fba12 -aapseaaaaaaaaarrrr-eepeaeeapeee Conference pdf.
Nguyen, A. and Reddy, G. Differential learning kinetics gov- ern the transition from memorization to generalization during in-context learning. In The Thirteenth Interna- tional Conference on Learning Representations, 2025. URL https::t//tt.... http///ppeeeevveeeeeee net forum? -d== INyi7qUdjz.
Rende, R., Gerace, F., Laio, A., and Goldt, S. A distribu- tional simplicity bias in the learning dynamics of trans- formers. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, vol- ume 37, pp. 96207-96228, Curran Associates, Inc., 2024. URL /
Nichani, E., Damian, A., and Lee, J. D. How transformers learn causal structure with gradient descent. In Salakhut- dinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the slst International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Re- search, pp. 38018-380777 PMLR, 21-27 Jul 2024. URL https: : http://occeeddnng// ml wttpa//gs........e/////233350 ichani 24mmllll.
07 o79dcbaB004bbee11ccff---eee nce.
Saxe, A. M., McClelland, L.. and Ganguli, S. Exact solu- tions the nonlinear dynamics of learning in deep linear neural networks. In The Second International Conference on Learning Representations, 2014. URL https: wwwevvewwwneeteeeeee
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Heraandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, I.., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads, 2022. URL httph:ttp/// 2202 /wwpp///oo////nccooeee index
Saxe, A. M., McClelland, J. L., and Ganguli, S. A mathe- matical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sci- ences, (16(23):11537-111555 .18202261 16. URL. http:::tt///// //t::://wwwww /www. ppass org/dat / /abs/10. 1073 1 1822022111...
2019. doi: 10. 10.1073/pnas
Schlag. I., Iric, K., and Schmidhuber. Linear transformers are secretly fast weight programmers. In Meila, M. and Zhang. (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pp. 9355-9366. PMLR, 18-24 24 Jul 2021. URL httptttt/pocceeeee ngs mmmmlr... r. pressppeess//eeee https////ss//shsaaaaaaaa html.
xt-
.html.
Otto, F. and Reznikoff, M. G. Slow motion of gradient flows.
Journal of Differential Equations, 237()):372 2007. ISSN 0022-0396. doi:
07.03.007. URL URL.httpe://ww..cceeneeeee /tw.
Shamir, O. Exponential convergence time of gradient de- scent for one-dimensional deep linear neural networks. In Beygelzimer, A. and Hsu, D. (eds.), Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Re- search, pp. 2691-2713. PMLR, 25 28 Jun 2019. URL https: http://oooceeee...e/ee mlr, presssvv99
com/ http://ccceeerrrrrreeeep///o2202 00824.
Park, C. F., Lubana, E. S., and Tanaka, H. Algorithmic phases of in-context learning. In The Thirteenth Inter- national Conference on Learning Representations, 2025. URL https:ttttooooo /pwpr//vieeeee.e.eeeem net htre/formmmm id= XgH1wFHSX8.
Reddy, G.
The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference Learning Represen- tations, 2024. URL https:////// wpwp//ween.eeeeep////p www.wwiwwwwwww.eee.m.....eee or orum? Arummid=aN44JfGCCM
Sheen, H., Chen, S., Wang. T., and Zhou, H. H. Implicit regularization of gradient flow one-layer softmax at- tention, 2024. URL https::/aaran wrg. ab4 03. 08699.
Singh, A. K., Chan, S., Moskovitz, T., Grant, E., Saxe, A.. and Hill, F. The transient nature of emergent in-context learning in transformers. In Oh, A., Naumann, T., Glober- son, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, vol- ume 36, pp. 27801-27819, Curran Associates, Inc., 2023. URL httpttttps... hww.thtts////ooooee http::////ccceeeeeeeeeee www.neiippp.. www.ies.cc......cc
Ren, Y., Wang, Z., and Lee, J. D. Learning and transferring sparse contextual bigrams with linear transformers. In Globerson, A., Mackey, L., Belgrave, D.,, Fan, A., Paquet,
U.. Tomczak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 20304 20357. Curran Associates, Inc., 2024. URL https:
14

---

Training Dynamics of In-Context Learning in Lincar Attention
aper gpe.ffllee/zaarr/002///l////222
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. Atten- tion is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Gar- nett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 30. Curran Associates, Inc., 2017. URL httphttss/pproceednngs..seee aper a0e__file/pppe///77/6/eeefee 4 17eeee11foo0533ccc444444a-222222
nce.
Singh, A. K., Moskovitz, T., Hill, F., Chan, S. C., and Saxe, A. M. What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formation. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarleu, J., and Berkenkamp, F. (eds.), Proceedings of the 4sst International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 45637-45662. PMLR, 21-27 Jul 2024. URL ht.p::::tt//..t. http///oceeddieeeoooo nig.. .ml ml html.
cc/p
Vladymynov, M., von Oswald, J., Sandler, M., and Ge, R. Linear transformers are versatile in-context learners. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, Advances in Neural U., Tomezak, J., and Zhang, C. Advanees Information Processing Systems, volume 37, pp. 48784
r. wtpprsssv23///snnnnn
Singh, A. K., Moskovitz, T., Dragutinovic, S., Hill, F., Chan, S. C. Y., and Saxe, A. M. Strategy coopetition explains the emergence and transience in-context learning, 2025.
48809, Curran Associates, Inc., 2024. URL https:
/ //proceedings neurips cc wc/paper.pppeee.eeeceeeeeeee..ee pttpppreee_eeeeeeeseepeepppeeeee gaaper/2222//illee5/222222222222
files
URL https::////. /atiiviiiivv///ii///aaa//aaaavat http://bb//20000 05631.
Song, B., Han, B., Zhang, S., Ding, J., and Hong, M. Unrav- eling the gradient descent dynamics of transformers. In Globerson, A., Mackey, L.., Belgrave, D., Fan, A., Paquet, U.. Tomezak, J., and Zhang. C. (eds.), Advances Neural Information Processing Systems, volume 37, pp. 92317- 92351. Curran Associates, Inc., 2024. URL https:
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In Krause, A., Brunskill, E.,, Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174, PMLR, 29-29 Jul 2023. URL https: / //proceedings mlr press v202222222//2/ /von dsw aid23a.html
db25cb1 aarobbaaaaabb--o -opeeeeeeeee-ooeeeeeeeeeeee
Tarzanagh, D. A., Li, Y.. Thrampoulidis, C., and Oymak, S. Transformers as support vector machines, 2024. URL https://arxii: org/ abs/ /2008.20... 16888.
Wang. M., Yu, R., E, W., and Wu, L. How transformers implement induction heads: Approximation and optimiza- tion analysis, 2024a. URL.
Tian, Y., Wang, Y., Chen, B., and Du, S. S. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Oh, A., Naumann, T., Glober- son, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, vol- ume 36, pp. 71911-71947, Curran Associates, Inc., 2023. URL htt.t:tt/pooceeeeeee www.:igs.....uiuu...oom...m.. cc aper oper__/leee/paprr//22//////////s 0a3366667449952249960444844 pdf.
abs 2obs/24100 2410.
Wang. Z., Wei, S., Hsu, D.,. and Lee, J. D. Transform- ers provably learn sparse token selection while fully- connected nets cannot. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st Interna- tional Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 51854 51912. PMLR, 21-27 Jul 2024b. URL https: / /proc html.
eedings. ml www.mmrrrrrrrseemm press v22 hwtp//////w/a222
Tong, W. L. and Pehlevan, C. MLPs learn in-context on regression and classification tasks. In The Thirteenth International Conference on Learning Representations, 2025, URL https://///....... /ttps://pppeneee hpwp//weeeeeeeeee////e/ee net/ forum ? fid-MbX0tlrUlp.
Woodworth, B., Gunasekar, S., Lee, J. D., Moroshko, E., Savarese, P., Golan, I., Soudry, D., and Srebro, N. Kernel and rich regimes in overparametrized models. In Aber- nethy, J. and Agarwal, S. (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Pro- ceedings of Machine Learning Research, pp. 3635-3673. PMLR, 09-12 Jul 2020. ngs mlr. press vv25 http://wwwoooth..aaaaoa html.
Vasudeva, B., Deora, P., and Thrampoulidis, C. Implicit bias and fast convergence rates for self-attention. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL https:/ http::://pppreeeeeee //ppnrevieeee...////////p. net forum? id= pKilnj0sbo.
Wu, J.. Zou, D., Chen, Z., Braverman, V., Gu, Q., and Bartlett, P. How many pretraining tasks are needed for
15

---

Training Dynamics of In-Context Learning in Lincar Attention
in-context learning of linear regression? In The Twelfth International Conference Learning Representations, 2024. URL https:://// wpwpreveee..eeeeeeeeeeeee ner/forumm forum ?id=vSh5ePa0ph.
Xia, M., Artetxe, M., Zhou, C., Lin, X. V., Pasunuru, R., Chen, D., Zettlemoyer, I.., and Stoyanov, V. Training trajectories of language models across scales. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceed- ings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp- 1371 1-13388,, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/w1/0223 acI- lcll-oogg 767. URL hwwpt:ttt///oo //al:///acclc anthology. org /2023 acl long. 767.
Yau, M., Akyürek, E., Mao, J., Tenenbaum, J. B., Jegelka, S., and Andreas, Learning linear attention in polynomial time, 2024. URL https: / //tp://aaavv.///aa org abs 2 10.10101.
Zhang, R., Frei, S., and Bartlett, P. L. Trained transform- ers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024a. URL http:
htw.a/..oggpppees/eeee/e .rg. htap:/apprrssseeeee wvt./23sr///////--2// 1022231-22...00 html.
Zhang, R., Wu, J., and Bartlett, P. In-context learning of linear transformer block: Benefits of the mlp compo- and one-step gd initialization. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tom- czak, J., and Zhang, C. (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp. 18310- 18361. Curran Associates, Inc., 2024b. URL https: woarcceeiingggeuurrr.......//// /ppper//00///Laa//bBB8811111311 P pdf.
files
16

---

Training Dynamics of In-Context Learning in Lincar Attention
Table of Contents
Introduction
2
Preliminaries
2
2.1
In-Context Linear Regression Task
2.2
Multi-Head Sclf-Artention
Linear Attention with Merged Key Query
Linear Attention with Separate Key Query
Gradient Flow Training Dynamics..
Linear Attention with Merged Key and Query
3.1
Connection to
Linear Network
3.2
Loss Landscape: Two Fixed Points
3.3
Training Dynamics: An Abrupt Drop in the Loss
3.4
ICI. Algorithm: Least Squares Regression
4
Linear Attention with Separate Rank-One Key and Query
4.1
Connection to Convolutional Linear Networks
4.2
Loss Landscape: Exponentially Many Fixed Points
4.3
Training Dynamics; Saddle-to-Saddle Dynamics
4,4
ICL Algorithm: Principal Component Regression
5
Linear Attention with Separate Low-Rank Key and Query
8
6 Related Work
8
7 Discussion
9
Additional Figures
19
A.I Effectt of Initialization Scale
19
A.2 Higher Dimensions
19
A.3
Varying Context Lengths
19
B Additional Related Work
20
Additional Preliminaries
20
C.I
Data Statistics
20
C.2 Initialization
21
C.3
Kronecker Product
21
D Linear Attention with Merged Key and Query
22
D.I Justification for Zero Blocks Assumption
22
D.2
Gradient Flow Equations
23
D.3 Fixed Points
23
D.4 Duality of the Global Minimum Solution
23
D.5 Analytical Time-Course Solution for White Covariance
24
D.6 Training Dynamies for General Covariance
25
D.7
Conservation Law: All Heads Are Parallel
27
Linear Attention with Separate Rank-One Key and Query Justification for Zero Blocks Assumption
27
27
E.2
Gradient Flow Equations
27
Fixed Points
28
E.4
Loss Value at A Fixed Point
29
E.5
Saddle-to Sadddee Dynamics: From My to Mi
30
E.6
Saddle-to-Sadee Saddee Dynamics: From Mmn to M
32
E.7
Weight Configuration with Minimal L2 Norm
34
E.8
Conservation Law
36
F
Linear Attention with Separate Low-Rank Key and Query Justification for Zero Blocks Assumption
36
F.1
36
F.2
Gradient Flow Equations
37
E3
Fixed Points
37
F4
Saddle-to-Saddde Dynamics
38
F.5
Dynamics with Repeated Eigenvalues
39
E6
Conservation Law
39
Training Dynamies of In- Context and In- Weigtttt Learning
40
17

---

Training Dynamics of In-Context Learning in Lincar Attention
Appendix
1.0
10
0.8
10
0.8
10-3
0.8
0.2
0.6
10
080
0.05
10
10
0.02
0.2
0.2
0.0
1000
2000
5000 Training steps t
10000
Training steps
(a) Linear ATTNM
(b) Linear ATTNs with
Figure 6. Loss trajectories of linear ATTNM and ATTN; with varying initialization scales. The colors indicate the initialization scale. Increasing the initialization scale shortens the plateaus. With small initialization, the models the rich feature learning regime, exhibiting abrupt sigmoid-shapes dynamics. With large initialization, they are in the lazy learning regime, exhibiting exponential-shaper loss decay. The loss curve from intermediate initialization seems like mix of the exponential mixed curves are often seen in practice, such as in induction head emergence in natural language settings (Olsson et al., 2022, Argument 1). dataset and model setup are the same as Figures and 3, except that we vary the initialization scale.
and sigmoid-shaped curves. Such
D= 32, =
D=61,
D= 32, 328
D= D=
1.0
1.0
1.0
0.8
0.8
0.8
0.6
0.8-
06
0.6
0.4
0.2
0.2
02
0.0
0.0
0.0
Thsining stepa
Training staps
Training staps
Training slaps
(a) Linear AT TNM
(b) Linear TNs
Figure 7. Loss trajectories of linear ATTNx and ATTNs with high-dimensional data. Here the sequence length is N 127, A has trace eigenvalues Aa x - Other hyperparameters are labeled at top each panel.
Simulations
Theory
((Mb)
1.0
0.8
0.8
((M))
0.6
Loss
((M))
0.4
0.4
E(M:)
C(Ma)
0.2
0.0
2000
20000
40000
Training steps
Training steps
(a) Lincar ATTNM
(b) Linear ATTNs with R
Figure 8. Loss trajectories of linear AT and ATTNs trained with the next token prediction loss defined Equation (23) with Nmax = 31. In this case, the models are trained on sequences of varying lengths, which they can handle due to the 1/N scaling factor Equations (4), (M) and (S). The dataset and model setup the Figures and 3, except that switch the next token prediction loss.
18

---

Training Dynamics of In-Context Learning in Lincar Attention
A. Additional Figures
A.1. Effect of Initialization Scale
We vary the initialization scale of the linear attention models and plot their loss trajectories Figure 6.
A.2. Higher Dimensions
We train M and ATTNs on a dataset with larger N, D. The loss trajectories are qualitatively similar to those in lower-dimensional cases, despite being noisier. This suggests that our findings do not break in high-dimensional settings.
A.3. Varying Context Lengths
In our main results, we consider a fixed context length N, because our training sequences have the same length and the loss is computed only for the last query token, = E(g - "g)². In practice, however, the training sequences may have varying lengths, and the loss can be computed for every token the sequence, that
Nmax +1
Cmp
(Un - -n))
(23)
N
Nmax
where yN+1 Vy, and Un is the attention model's prediction for yo when given only the first columns of X input. We demonstrate how our results apply to the of varying context lengths. Specifically, the distribution of the context lengths only influences our results through statistic, E(1/N).
For ATTNM, derivations in Appendix D.4 show that the converged model implements
)
ATTNM (X))DLN++ Σ B U,B 3T
E xn
1 N
Azg.
(24)
Substituting Equation (30) into Equation (24), we obtain
+E ( (A (A+=
(TTTNM())DLLNNN) BT A A
Bq.
(25)
The distribution of context lengths only influences Equation (25) through the expectation E(1/N). For fixed context length, E(1/N) L/N, which recovers Equation (11) in the main text. For the next token prediction loss, the distribution of context lengths, p(N), follows a uniform distribution over f1, NNuax) The expectation E(1/N) is the harmonic number divided by Nmaxx which doesn't have a closed-form expression but can be easily computed for a specific finite Nmax-
Similarly, for ATTNs trained with varying context lengths, the fixed point condition (C1) takes the form
Σ k.k
>x
1+E
1 N (1+Er(A)/Aa)
c
deSm..
where the expectation E(I/N) reduces to 1/N in the fixed context length case as in Equation (18). Consequently, when the model at a fixed point Mm. the loss value is
((Mm)=trlee
Σda
1
++E (1+ (+)((AAA
1 N
(26)
d=1
which reduces to Equation (19) when ((1/N)=1//
=
We train ATTN and ATTNs with the next token prediction loss as in Equation (23) and plot the loss trajectories in Figure 8, The loss trajectories are qualitatively similar to those Figures and 3a, modulo the different loss values during the plateaus. We plot the loss values computed from Equation (26) as dashed gray lines and find they match the plateaus the simulated loss trajectories well.
19

---

Training Dynamics of In-Context Learning in Lincar Attention
B. Additional Related Work
concurrent work by Geshkovski et al. (2024) studies saddle-to-saddll--ikk dynamics in softmax attention models following mathematical framework for slow motion of gradient flows (Otto & Reznikoff, 2007). A subsequent work by He et al. (2025) examines the training dynamics of softmax attention trained on the in-context linear regression task; that is, the case we briefly touch on in Figure 5.
A broader body of theoretical literature have explored the transformers training dynamics but addressed different problem from ours, such as the effect of initialization (Makkuva et al., 2024), convergence results (Song et al., 2024; Huang et al., 2024a), sample complexity guarantees (lldiz et al., 2024), scaling limits (Bordelon et al., 2024), and implicit regularization (Ataee Tarzanagh al., 2023; Tarzanagh et al., 2024; Julistiono al., 2024; Vasudeva et al., 2025; Sheen et al., 2024). Other studies considered special training regimes, such as the neural tangent kernel regime (Jang et al., 2024) and the mean-field regime (Kim & Suzuki, 2024). A few works focused on vision transformers (Jelassi al., 2022; Jiang et al., 2024; Huang et al., 2025b). In contrast, our work focuses on the training dynamics and the development of ICL abilities over time.
It is recognized that transformers can perform ICL, whereas it is an open question whether fully-connected networks can perform ICL (Lee et al., 2024; Boix-Adserà et al., 2024; Tong & Pehlevan, 2025). In Section 3.1, we revealed an equivalence between linear attention and a fully-connected linear network with cubic feature input, which is an instance of a fully-connected network performing ICI. Furthermore, we demonstrate that fully-connected networks may perform ICL more comparably to attention models when provided with polynomial features instead of the original sequence. This may explain why Boix-Adserà et al. (2024, Figure 25) observed that fully-connected networks fail learn ICL with the original sequence imput, but succeed when the input is augmented with XXx
C. Additional Preliminaries
C.1. Data Statistics
Recall that we use 3 to denote the in-context correlation between Vn and yn in sequence X, as defined in Equation (4). We additionally denote the in-context covariance of F, a sequence as A
1
(27)
We can thus write Xx T /N as a block matrix
1
XX
N
( 2 T + En x
N Σ Enlii
N
20
+
B
(28)
N
N En x
N U
3
w
Aw
Due to the definition of the in-context linear regression task, we have that
B Aw.
(29)
We will need a statistic, E (A2) obtain:
Let p(N) denote the distribution of context lengths, and recall that 2n, ~ ~N A). We
N
2
2
E
Fn a
N2 N
N +
E
Σ
Σ
x
x
N2
N2
nzn''
N
1
En
Eye (x)
Eg (z)) + +EN
Em
(T x0
N
N
1
1
E
A +
(2A2 2 tr(A)A)
N
1
2
+
E
(A + tt(()
(30)
20

---

Training Dynamics of In-Context Learning in Lincar Attention
For our main results, fixed context length, that p(N) is point mass distribution and B(1/N) 1/N. In this case, Equation (30) simplifies to
E A2
A
A
tr(A)T A.
(31)
N
We note that the cigenvectors of E
(A3)
the same as those of A, which are e ,ep
(z
[[1+
1 N
tE(A) Ad
E
A2
ed
11 N
A ed
Ae
ed.
N
We denote the eigenvalues of (A?)) corresponding to eigenvectors " eD 01, ... apD These eigenvalues are given
by
ad
(
1+
X2
tr(A) +
Ad
23
1+ 1+)
(32)
N
N
N
The matrix B(A2) can be expressed through its eigen-decomposition, which will be useful in later derivations:
E
A2
(33)
d=1
C.2. Initialization
For linear attention with merged key and query, we initialize the entries of the value and the merged key-query weights
Us NOO, / Undd" ~ ~N,0,, winh (HD)).
(34)
At initialization, the following e norms are
H
H U1 O(wmi)
(35)
For linear attention with separate rank-R R key and query, we initialize the entries the value, key, and query weights as
0/2 ~ (O, / H),
Rd 1r
NoO,, w HHRD),
Gir ~ ~O, w HHD).
(36)
At initialization, the following norms are
H
H
R
H
Σ
k
13,
| Owiii))
(37)
1=1
C.3. Kronecker Product
The Kronecker product, denoted as 8,, is defined for two matrices of arbitrary sizes. The Kronecker product of the matrix
A E _2xx and the matrix B E grxs is a block matrix of shape x
an
O1g
a1 B
01g B
A &
B
④ B
ap1
Opq
ap1
B
apg
Based on the definition, it holds for any pair of column vectors a and b
vec(baT)
(38)
21

---

Training Dynamics of In-Context Learning in Lincar Attention
We quote some properties of the Kronecker product to be used in our derivations:
(cA) @ B A &
cBB) = c(A @ B) for any scalar C,
(39a)
(A 8 B) A @ B for any matrices A, B,
(39b)
(A B)- A B-1 for invertible matrices A, B, (A 8 () & D) (AC) & (BD) for compatible matrices A, B, C, D,
(39c)
(39d)
B ()))c(M)) vec(AMB) for compatible matrices A, B, M.
(39e)
D. Linear Attention with Merged Key and Query
D.I. Justification for Zero Blocks Assumption
We prove our claim Section 2.3 that v; and u, remain zero throughout training if their initialization zero.
Proof. The bottom right entry of (TTNM((X) is given by
0g MTTNM X D+1,N+1
H
12
Us]
N
T + En )
N Σ Enlynn
U; iu
Zzg 0
En D.
Σn y
1=1
N
N
T
1
U; v. UU. + 0. Bug + v.w Awny Bn.
v.
N
If we initialize v, u, 0. Ug is
H
H uy
0q Σ 0uB U, UU w A
Σ Uaaq
We now calculate the gradient updates of Us, u, and prove their gradients are zero if their initialization is zero. The gradient update of u, contains E( w)) which is zero. Specifically, we have, from Equation (5),
TŪ, E
(yg 0g) ((A+
1 N
U. + Bu, IQ
(A+ 1
w
Fg
w
A
2y U
N 2 U
H
(A+
Ew (w)
TE
Bg
A
N x Upiiy
O.
(40)
Note that we separated the expectation of w because of the independence between and all tokens.
The gradient update of v; contains En (w kwoT) whose entries are linear combinations of third moments of the zero-mean normal random variable w, and thus zero. Specifically, we have
Tu; = E [0,,
ug0 Au)) (yg 3y )
H
= E wiw T Aw
w xn w
T A
U:o0
im
H1
En (ww "www" T
E
Bg
A
Ui
i=
0.
(41)
22

---

Training Dynamics of In-Context Learning in Lincar Attention
D.2. Gradient Flow Equations
We here derive the gradient flow dynamics for linear attention with merged key and query given in Equation (8).
For linear attention with merged key and query, the prediction for the query output can be written 04 w. Wiz due to Equation (6). Based on the gradient flow training rule in Equation (5), the gradient flow dynamics is
W E w2 (vo- - 2. Wi z) z =1| W2 [E (y0
w2
E (zz )), WiE (zz 1) T
W
2 E
Wi (Ug- 2i Wi w z| W E (ug2]
which was introduced in Equation (8) in the main text.
D.3. Fixed Points
To find the fixed points, we set the gradients in Equation (8) to zero
W E 22 Wi w2 (E (y w2
set
0,
TW0 W1 (E (ug2 w.
W E [zz
sot 0,
which yield the two manifolds of fixed points introduced in Equation (9) in the main text:
0, W1 = 0
Mo {w2 0, W = 0},
EE (H W.2 W,E E (zz 0 M. W2 Wi W Wi E (y,z ) 1xxz²² 1zxz²)=
D.4. Duality of the Global Minimum Solution
We here prove the second equality in Equation (11), that is
A Axt((A)]) -1
E
34
E
zz
")'''==== x == A
Fqq.
N
This equality implies the intriguing duality that the linear regression solution in the cubic feature space of is the in-context linear regression solution in the original space of the Fn Vn. token pairs for each sequence X.
Proof. We first calculate the input and onput-output correlations in the cubic feature space. We denote Ag E (xg27))
While Ag = A, this equality is not needed in this proof.
Due to the property in Equation (38), the cubic feature z can be written as
vec (Bx)) = 2g ? B.
(42)
We substitute in y4 = 2 w, z xy X 3 and use the properties of the Kronecker product in Equation (39) to obtain
IE (yg=
IE [x (w [] E x x ww A
E
x
3 x. A
Evec (Amr
= =vec(AA,)]
(43)
Similarly, we have
E (xx )) E (xx & (xx & [T]
E [ 8 (3,3 )
E( 2 B(Aww Á)
A. @ E (A
(44)
23

---

Training Dynamics of In-Context Learning in Lincar Attention
Using Equation (39c), the inverse of E (zz
(zz )==]T>'== = A, & E &E(
(45)
Multiplying Equations (43) and (45) with
Bg X B3 and using Equation (39e) twice, we obtain
E (vgz)) ) E (zz ) -1 vec(AA,) T A, E A2 ((A)))
-1
(Eg @ (3)
vec B (Λ² AA
(xe @ (3)
3 E (2)) Ade..
(46)
Substituting in E (A²) obtained from Equation (31) finishes the proof
E (ygz) E (zx (zx))= == A2 222 (022
A
A
A)- Ax, =B (A+ A A+tr(AAI
(A+A+IAAA) N
Eg.
N
D.5. Analytical Time-Course Solution for White Covariance
We include a derivation of the time-course solution of two-layer fully-connected linear network with white input covariance and vanishing initialization following (Saxe et al., 2014), and then apply it to linear attention. With vanishing initialization, the conserved quantity given in Equation (55) is exactly zero throughout learning,
WiWi = 0.
Hence, there exists unit norm vector m such that Wi u2mm With the assumption of white covariance, (22T) al (Saxe et al., 2014; Atanasov al., 2022) have shown that the unit norm vector is parallel with the correlation between yy and throughout training, that is
E(ygz)
Wi u2m where
(47)
[E(y,z)]]
We substitute Equation (47) and the white covariance assumption, E (zz = ax pa, into the gradient flow dynamics given in Equation (8) and obtain
Tu = W2m' (B(y,z) aw, w2m) w22 (? anwo w2). where y = [[E(y,))) |||
Notice that the square of the 8 norm of wz follows solvable ordinary differential equation. Let wa w2. The dynamics of s(t) is
7 2wd Teby 2wd wz (r aw. w2) 2s(r ass).
(48)
We solve this differential equation by separating variables and integrating both sides.
2(0) s(y as) ds / 2 dt = in s(t) - as(()) 2 t.
(())
7
s(0) -as(t) 7
The solution of s(t) is given by
2
8(1)
22 - 1) + 7
x(0)
24

---

Training Dynamics of In-Context Learning in Lincar Attention
The time-course of the total weights is given by
E (3
2/2 W s()m s(t)
(49)
[gy
We now apply this solution to linear attention. If the input token covariance is identity, A = ID. we calculate the input and input-output correlations in the cubic feature space according to Equations (43) and (44) and get
[y4 ) = vec(ID) (zz ID 0
(1 (1+-
[D==1++
ID ID=
1+D N
ID
N
The parameters in the dynamies of the equivalent two-layer fully-connected linear network are
1+D
==
N lIvec(AD)] = =V.
(50)
Substituting Equation (50) into Equation (49), we obtain
vec(ID)
2DeVVD4
w wi ((t)
where s(1)
VD
(1 1))) (2VD (1
VD x(0)
Due to the equivalence between linear attention and the two-layer fully-connected linear network given in Equation (6), we obtain
vec(TD)
1 [(t)BBIbagg vD" st) Teq.
ATTNM(X:f) w. W, s(t)
B p== vD Igg
IDet
VD
which is Equation (10) in the main text where we have rewritten o(1) = =(1/vVD.
The time-course of loss can also be expressed in terms of ott)
E(1) (1-20(() + 1
1 1+ 1+D)
D
a(1))) D.
(51)
N
D.6. Training Dynamics for General Covariance
1. EARLY DYNAMICS PREDICTS DURATION OF PLATEAU
For a general input covariance matrix, the full time-course solution to two-layer fully-connected linear networks is currently unavailable. Nonetheless, the training dynamics is well understood and we can analyze the early phase dynamics to estimate the duration of the loss plateau.
In the early phase of training when the loss plateaus, the weights have not moved much away from their small initialization. The training dynamics of Wi is mainly driven by the first term in Equation (8a), and similarly for we in Equation (8b)
W1 wy (E (yg w. W,E zz T)) E (vy 2z )) + Ow (wwi).)
Tw Wi (E (y,z w. W (zz T)) IE (2,=) + (w
Thus, the early training dynamics is well approximated by the linear dynamical system
Wi = w2E (y,
W1E (y₁=).
In the case of nonwhite covariance, the change of variable in Equation (47) is valid in the early phase of training but no longer valid when the loss starts to decrease appreciably (Atanasov et al., 2022). For the early training dynamics, we apply the change of variable in Equation (47) and obtain
TW02 w2m E (y)z) -2
25

---

Training Dynamics of In-Context Learning in Lincar Attention
ATTNM(X)
MLP(2)
Statistics of Training Data
pooc 500 during plateau (multi-head)
(U))
W
(su= )
(ec(U))
Epoch 6000 converged (multi-head)
U,))
vec (A =s (v62 (1=="))(="")
Wi
(ve(U)))
z reshope (Ba+
Epoch 6000 converged (in one head)
reahaape (W?7))
Figure 9. The dynamics of weights in multi-head linear attention with merged key and query can be predicted with statistics of the training dataset. plot the weights at different times training, corresponding to loss trajectories Figure 1 (right). The weights in linear attention (first column) stay close to the weights in the fully-connected linear network (second column) throughout training. During the initial plateau, the vectorized key-query weights in attention vec(U1), ... vec(U:) and the first-layer weight in the fully- connected network Wi are rank-one and align with the correlation between the output and the cubic feature input (24 (top row). At convergence, voc(U1), ... vcc(U.) in attention and Wi the fully-connected lincar network are rank-one and align with the linear regression solution in the cubic feature space (v, E (zz T) (middle row). which is also the in-context linear regression solution in the original token space (bottom row) described by Equation (11). The approximate equality the third column exact when the sequence length N - Co.
Recall that s w2 W2. The early phase dynamics of s(t) approximately
T$ = 2ys.
We solve the differential equation and obtain
(In s(t) - I0s))
2y
Due to small initialization, In s(l) at the end of the plateau is much smaller compared to - IIn s(0) Hence, the duration of the initial plateau of loss, &platemn is
5 In
Iplateon
R
(52)
27
s(0)
Here, the scalar is
? = [g 2
||B (w a ve. ())
E (w 2,gBar,)) )|1
IE' "Aww a x )|
E A End (ww Em, (xr ||-
||
(53)
Substituting Equation (53) into Equation (52), we obtain the approximate duration of the loss plateau
?plateau
In
In
(54)
2
A 11p
s(0)
|
Winit
where we used the definition s(0) = [[]]2 w
26

---

Training Dynamics of In-Context Learning in Lincar Attention
D.6.2. WEIGHTS DYNAMICS
For white input covariance, the training dynamics reduces to scalar ordinary differential equation about s(t) given in Equation (48). For general imput covariance, the vector m in the change of variable defined in Equation (47) rotates during training. As shown in the top row of Figure 9, during the initial loss plateau, the rows of the first-layer weight align with the input-output correlation E (uz²T) but do not change appreciably in scale (Atanasov et al., 2022). Later, when the loss decreases rapidly, the first-layer weight grows in scale and rotates to align with the global minimum solution, E (y₂2 E (zz ( The alignment and rotation behaviors apply to the rows of the first-layer weight in the fully-connected network, corresponding to the merged key-query weights in the different heads in linear attention, shown in Figure 9.
D.7. Conservation Law: All Heads Are Parallel
The weights in fully-connected linear network are known obey conservation law during training (Fukumizu, 1998; Saxe et al., 2014; Du et al., 2018; Ji & Telgarsky,2015 2019)
dt [ W1 = 0,
(55)
which follows directly from the gradient flow dynamics in Equation (8). Under small initialization, the quantity w2 waw2 small at initialization and remains small throughout training. Since the vector w2 rank-one, the conservation law forces W1 to also be approximately rank-one, which means that the rows of W1 are approximately parallel. Since each row of Wi is the vectorized merged key-query matrix of head, vec(U.). rank-one W, implies that the key-query weight matrices of all heads are parallel, differing only in scale. As shown in Figure 9, simulations indeed show that the key-query weights in different heads are parallel.
W1 WT
E. Linear Attention with Separate Rank-One Key and Query
E.I. Justification for Zero Blocks Assumption
This is a special case of linear attention with separate rank-R key and query. The proof for the more general rank-R case can be found in Appendix F.1.
E.2. Gradient Flow Equations
We here derive the gradient flow dynamies for linear attention with separate rank-one key and query introduced in Equation (14).
Based on the gradient flow training rule in Equation (5), the gradient flow dynamics for the value, key, and query weights in the i-th head are
TŪ; k.T E (B(yn Dn) T qi
(56a)
rk2 = ny (3(u,, - (1))or1 Uis
(56b)
TΦ: = a (Ir(U1 (a) k.s
(56c)
We calculate the common term in Equation (56), that is
H
E (B(y (g) 0,)=²))=
B u
Bg
Σ
k.g. Dy
xi T
i=1
h
Awos T
Σ
U/ AKg T.
H
E A
Eup (ww 7)
Bary (xr - E (Nww T A
Σ
u. Ba, (x
1=1
H
^² E (A²
Ug u.kg
A
(57)
27

---

Training Dynamics of In-Context Learning in Lincar Attention
Substituting Equation (57) into Equation (56), arrive the equations as Equation (14) in the main text
H
70, k2
A E (A²)
A
Uy ky q
A.A
A
gi.
1
Tk; = 0/
A E A
Σ
Uy q A
q:.
Tė: =
A 2
A
Uy ki q B (A2)
P=1
where the data statistics EE (A2) is calculated in Equation (31).
E.3. Fixed Points
We prove that the fixed points given in Equation (16) are valid.
Proof. When the model is at fixed point in set M()m), it satisfies Equation (18). Equation (18) be rewritten using as (defined in Equation (32)) as
H
X
Σ
edec
(59)
ad
deS...
Using Equations (33) and (59), we can simplify common term in the gradient descent dynamics in Equation (14) to
D
A E (2
k.qq
A A
Σ
Σ
esered
Σ deSm
Aa
T
A
ad
d=1
Σ
AgeaeT
Σ
Adcae
d=1
deSm..
(60)
dgs...
Substituting Equation (60) into Equation (14), we obtain the dynamics when the model is at a fixed point in M ((
TŪ
k
Σ
AgeaeJ
qi:
(6la)
Tk 0/
qir
(61b)
d
Th: 04
Ageaed
ki-
(61c)
dgs...
For the heads with nonzero value weight, i A O, the key and query weights fixed point satisfy condition (C2) for Equation (16), that is the key and query weights lie in the span of (ea)des.. and thus can be written as
k;
baeas ba RR,
(62a)
desm
caCa,
ca R.
(62b)
dESm..
28

---

Training Dynamics of In-Context Learning in Lincar Attention
Substituting Equation (62) into the gradient flow dynamies given in Equation (61), we obtain
TØE = kjj
Σ
X
Σ
CF - 0,
d'ESm
rk = 0i
Σ
Pheaed
Σ
Cdl = O,
de5m
d' d'ESm
T = 0i
Σ d'ESm
bare O,
where we have used the fact that e] e 0 if + d', because eigenvectors of the covariance matrix A are orthogonal.
(ii) For the heads with zero value weight,
0, the gradients of the key and query weights in Equations (61b) and (61c) contain and are thus zero, k; = 0, q. = O. Further, the key and query weights of a head with zero value weight satisfy condition (C3) for Equation (16). Without loss of generality, suppose that qi lies in the span of (ea)a that is qi satisfies Equation (62b). Substituting Equation (62b) into the gradient of e given in Equation (6la), we obtain
Σ d
Σ d'ESm.
0;
Ca = 0,
where we have again used the fact that eigenvectors of A are orthogonal.
Hence, when the model has weights specified in Equation (16), the gradients of the weights are zero, meaning that the fixed points are valid.
E.4. Loss Value A Fixed Point
We derive the loss when the model is at fixed point in set M(S,),, where the loss is given by
((M(S.)) = tx(A)- - E Aa 1+ [1+1+
(63)
N
des...
Equation (19) in the main text follows directly from Equation (63) when taking S.. = (1,2,..., ,m)
Proof. We substitute Equations (33) and (59) into the mean square loss and obtain
((M(Sa)) = E(y, )z)²
Xd w Reaetar zig
E
T
xiq
ad
desm.
E
g
I
Ad AeseT ad
Ew Ew(ww
I
Σ
dd
Aesed
Ig
ad
deS
deSm
(I-
E
42
T
I
Xd Reaed
Σ
Aa Reded
Bqq
Od
Od
S
dESm..
21
Xd Resee +
Xa Reaed
T.
x
- Σ
Σ
Big
(64)
ad
ad
deSm..
deS
29

---

Training Dynamics of In-Context Learning in Lincar Attention
Since À is independent of T.g. we can calculate the expectation of the purple and teal terms first,
I
Σ
Ad
Reaed
Ad
Aeded
Σ
2
T
ad
ad
dESm ad
de
deSm
27
Σ
Ad Nexed
=
12
Ede A
E
Aade Gd0g
e
ad
dd
de
deSm
ede d
(A
e
+0
deSm
eded Σ
d
ede
ad
dESm
Substituting them back into Equation (64), we get
X
12
((M(Sm))
= E
I-2
Σ
eded
+
Σ
eded
ad
ad
deSm
deSm.
X
E
g
T
I
e.ded
Bq
Od
d
X2
E
(x) Dzg)
E (x) T
agg)
Od
deSm
X
cr(A)-
Σ dESm
c
Aca
ad
=tr(A)
deSm
We plug in the definition of ag in Equation (32) and arrive at the desired result:
1 Artr(A)//a
((M(S.))) = tt(AA
Σ
1+
X
N
dESm..
1
=tr(A)- Σ Ad 1+
N
deS,
E.5. Saddle-to-Saddle Dynamies: From Mo MM
We denote the time which the loss has just undergone the d-th abrupt drop as tu (d I...., D), illustrated in Figure 10.
1. ALIGNMENT DURING THE PLATEAU.
In the initial loss plateau, the weights have not moved much away from their small initialization and thus the training dynamics are mainly driven by the first terms in Equation (14), which are
Tē; = k. A q: + (
(65a)
The A qg + O 5 )
(65b)
- A 2 k: + O (w
(65c)
30

---

Training Dynamics of In-Context Learning in Lincar Attention
Simulation
Theory
1.0
C(M))
0.8
0.6
((Ms)
((M))
0.2
(LMs)
4/2 f
10000
Training steps
Figure 10. Illustration of .., tp. The loss trajectory plotted one the trajectorics of linear attention with separate rank-one key and query in Figure 3a. The time ta (d 1,..., D) denotes the time when the loss has just undergone the d-th abrupt drop.
With small initialization scale Wintt the key and query weights in head evolve approximately
d
k2
0
A 2
k;
(66)
di
qi
=
A
qi
The matrix
0
2
2D2 has eigenvalues {>a.
-22
D
d=1² corresponding eigenvectors
A
0
2
Ted
2
ca
-
ea d=1,...,D
0
ed
ed
A
0
edj
Edj -ed
where recall that Ad, ea(d = 1,... D) are eigenvalues and eigenvectors of A. Hence, the solution to Equation (66) takes the following form
[t)
(
(A) A3 / (((()at))
ce ed
e (k)(0)) (0)) exp
(t)
7 J
d=1
(67)
Σ ej (k:(0) - q: (0)) exp + 2 d=1
1
X L' [ (-
ed
If u/ > 0, the first summation term in Equation (67) grows and the second summation term decays. The key and query weights k;, both grow in size along the directions of the eigenvectors ea. If Us < 0, the first summation term in Equation (67) decays and the second summation term grows. The key and query weights k;; qi grow in opposite directions, ed and Ey respectively. In either case, the multiplication v;k,q, grows along e,e.
E.5.2. REDUCTION TO SCALAR DYNAMICS WITH AN ALIGNMENT ANSATZ.
The dominating term in Equation (67) is the term with the largest positive eigenvalue. In other words, the key and query weights grow the fastest along the first eigenvector e and thus approximately aligned with e Motivated by this insight, we make an ansatz that the key and query weights in a head are exactly aligned with e and the rest of the heads are zero²:
k = qı = U1C1.
(68a)
k; = q: 0, v2 0- H.
(68b)
Note that Equation (68) also assumes that the l norms of ki, q1 are equal, which is true under vanishing initialization due to the conservation law in Equation (83). This ansatz can greatly simplify the training dynamics and provide good approximation of the true dynamics, where weights in one of the heads grow in scale with the key and query weights aligning with e1, while the rest of the heads remain near zero from time 0 to t1.
We substitute the ansatz into the training dynamics in Equation (14) to reduce the high-dimensional dynamics to one- dimensional ordinary differential equation. To do that, first calculate the common expectation term in the training
WWe head aligned with e to have index
31

---

Training Dynamics of In-Context Learning in Lincar Attention
dynamics with the ansatz,
H U/ U A A?
D
2
E
A2
Σ
araae u efee
= A Aaaeee
(69)
1=1
where d is the first eigenvalue of () defined in Equation (32). Substituting Equations (68) and (69) into Equation (14), we find that the training dynamics of the first head simplify and the dynamics of the rest of the heads are zero
TŪE uie, (A ilarely u/) = \iaioi
rk = of (A ilarele oi
oi)e Xtoje Alaroeen
e1 = Xfuee A1aiofe1 T④I = of (A2 A v) vi)er
bi k; 0, q: 0, 2, H.
We further substitute in kn dimensional dynamics about 11 (t)
e q1 ri and find that the high-dimensional training dynamics reduce one-
Tùj = 2prî \aaryy
TÚjE = >ffoee - Aaafee
TŪE = 22oo Alaroy
(70)
Tūje Afofei Alaiee
Equation (70) is separable ordinary differential equation. By separating variables and integrating both sides, we can solve t terms of 1
A
v1 (t)
1
dei Ui
(0)
o
A1
+
A 11 +
2 / al1l +1
In
22233tm tan
constant.
(71)
6
o 2 11 +1
v3
Since Equation (71) does not have straight-forward inverse, we cannot obtain general analytical solution of in (4) in terms of Nonetheless, we can readily generate numerical solutions and obtain approximate analytical solutions when d is near its small initialization to estimate the duration of the first loss plateau.
When U1 is small, the dominating term in Equation (70) is Xfuî and thus the dynamics can be approximated by
TŪ
(0)
v1(1)
At the end of the plateau, Ui (t) has grown to be much larger than 01 (0). Hence, the duration of the first loss plateau, &1, is
~
(72)
E.6. Saddle-to-Saddle Dynamics: From Mm to Mm+11
In Appendix F.5, have analyzed the training dynamics from time 0 (1, during which the model moves from saddle Mo to saddle M, We now analyze the general saddle-to-saddle dynamics from time tm to fm+1 (m = 0,. D 1), during which the model moves from Mm, to Mm+11
1. ALIGNMENT DURING THE PLATEAU.
Based on our dynamics analysis from time to t1 and by induction, the weights during the m-th plateau are approximately described by Equation (20). Namely, there are heads whose key and query weights have grown and become aligned with
32

---

Training Dynamics of In-Context Learning in Lincar Attention
the first eigenvectors while weights in the of the heads have not moved much from their small initialization. Thus, similarly to Equation (61), the heads that are small initialization have the following training dynamics
TŪ k
Σ
9: (
d=m+1
=k =
Σ
q: + ((wm),
d=m+1
T4. U/
Σ
+ (w
d==++1
With a small initialization scale Winit. the key and query weights this head evolve approximately as
D
k;
=
22
2
k;
where Ω =
Σ
Xeaae.-
(73)
dt
qi
qi
d=m+1
The matrix
O
Ω
2D22D has 2m zero eigenvalues and (2D 22mmmm nonzero eigenvalues,, which are (>2, -22) P
2m)
52
0
amm+1"
The nonzero eigenvalues correspond to eigenvectors
0
ea
ea
0
Ed ed
ed -23
0
ed
ed
0
=
-ed
d= m d.
Hence, the solution Equation (73) takes the following the form
(t)
D
23 L vi(()d))
ee eea
Σ
ej (tm(tm) + (tmm))
exp
[q:(t)
d=m+1
A L a(t))a)))
+ 2
Σ ed (ki(tm) q: (t(mm)) exp
ea -ed
(74)
d=m+1
[ed
Σ el (ky(tm) + (tm))
ed
For 02 > 0, the first term grows and the second term decays with time. The third term does not change with respect to time.
E.6.2. REDUCTION TO SCALAR DYNAMICS WITH AN ALIGNMENT ANSATZ.
The dominating term in Equation (74) is the term with the largest positive eigenvalue. In other words, during the (m + 1)-th plateau, the key and query weights that are still near small initialization grow the fastest along the (m 1)-th eigenvector em+1. Based on this insight, we make the ansatz in Equation (20). This ansatz can reduce the high-dimensional training dynamics to one-dimensional ordinary differential equation and provides good approximation of the true dynamics, where weights in one of the heads grow in scale with the key and query weights aligning with em+1, while the rest of the heads do not change much from time tm to tm+1-
To calculate the training dynamics in Equation (14) with the ansatz, we first calculate a common term with the ansatz
H
A² E
A2
uz A = A²
Σ
a
Σ
Ad e + 1
Sm T
A
Od
d=1
A
- Am+11m+4 em+1 em+ Um++
(75)
d=1
By substituting Equations (20) and (75) into Equation (14), we find that the dynamics for the heads with index * m zero
+
à; 0, 0, ği 0, i + m 11
33

---

Training Dynamics of In-Context Learning in Linear Attention
For the head with index + 1, the dynamics reduce one-dimensional dynamics about vz(()
Tà. e e
A
Σ
Ageaed -m++1 a1111mm+11mmmm 1mmm1mmm-mm+mmmm-
em+1
d=1
22 m
Am+1am++0
m
Th, = Thee++1 = v/ A²
ejeed - 1m++10++emm 1mm bm+1
= - 17 em+1
Am+ am+ 10 em+1
T④: = T0:em+1 =
^² Σ
Xgened - Am+ amm+ em+ e
em+1
d=1
= A em+1 A em+1
TŪ; 1 Am+1 am++1
(76)
Equation (76) is the same ordinary differential equation Equation (70) modulo the constant coefficients. Therefore, with the same analysis, can estimate the duration of the (m 1) loss plateau.
When Um++1 is small, the dominating term in Equation (76) is A 2s and thus the dynamics is well approximated by
1
1
Tum+1 X Am+1
1 mm+1
4m
X
m++1
Um++1
Un++ (t)
At the end of the plateau, 9m41(1m-t has grown to be much larger than Um+1 (tm). Hence, the duration of the (m + 1)-th loss plateau is
fm++1 fy
(77)
m+ mm++1 (tm)
We note that the Equation (77) involves u (4m), which depends on the random initialization and the dynamics from time 0 to fm- This explains why we observe the variance tm increases with a larger m, that the timing of later abrupt loss drop varies more across random seeds as shown in Figure 3a.
E.7. Weight Configuration with Minimal L,2 Norm
We prove that Equation (20) with 2 = is the weight configuration with minimal E² norm that satisfies Equation (18). To do this, we find the weight configuration with minimal e norm satisfying a general equality constrain and apply the solution to Equation (18).
Consider the equality constrained optimization problem
minimize
|| 12 + ||
subject
ptkiq
A
where A is positive semi-definite matrix.
Proof. We use Lagrange multiplier to solve this equality constrained optimization problem. First, we construct the Lagrangian function L(M) where the Lagrange multiplier M E IR DxD is symmetric matrix
1
I(M)=
vvec + ||: + 11)= + vec vea
A
-
Σ
vakug
2
i=1
1=1
1
Σ
(uf
|
|| +tr M
A
Σ
u,ki
34

---

Training Dynamics of In-Context Learning in Lincar Attention
Differentiating the Lagrangian with respect all the variables and setting them zero, we get
0L
0c; v; k; Mq: 0
(78a)
ki 0, M:Mq: q:
(78b)
kk:
OL
qi
u. MM
(78c)
0q:
0L
"
ok T
(78d)
0M
Equation (78) suggests that, for each head, the value, key, and query weights are either all zero or satisfy a constraint; i.e., for each either U; k; qi
u. Uē M
(79)
We got Equation (79) by substituting Equation (78c) into Equation (78b). Equation (79) implies that k; an eigenvector of 2 Let us denote the normalized eigenvector of M as &. Substituting Equation (79) into Equation (78a) and rearranging, get
k; q: E
(80)
With Equations (78d) and (80), we obtain
A = Σ UiE
82 = 1/3 &, = = =
(81)
where As; e, are the eigenvalue and eigenvector For the optimization problem, the solution is that there are rank(A) heads with nonzero weights and (H rank(A))
heads with zero weights. The nonzero heads have weights
k; uie; U;
1/3
i 1, rank(A). i=a,.".,raa((A))
(82)
The indices of heads can be trivially permuted. The signs of any two among u, hs, can be flipped without affecting the optimization problem.
We apply the solution in Equation (82) to find a weight configuration with the minimal e norm that satisfies Equation (18). Equation (18) can be rewritten as Equation (59), namely
H
Σ
ugk
Xd Σ e 0d
deSm..
The matrix the right hand side has rank and cigenvectors ea with eigenvalues Aa/as (d E Sm). Hence, the weight configuration with minimal e norm has (H - m) heads with zero weights and m heads with nonzero weights. The nonzero heads have weights
1+ - i= 1=1,- 1,,,,,,0 m.
Ad ad
( 1+
k; = = U;Cis U4
X;
N
This is the same weight configuration as Equation (20) with Um++1 O.
35

---

Training Dynamics of In-Context Learning in Lincar Attention
E.8. Conservation Law
The gradient flow dynamics of linear attention with separate rank-one key and query in Equation (14) implies a conservation law. The value, key, and query weights in head obey
d (k: ki q. q:) =0, dt
d (ke k; - -vf) 0,
(83)
dt
Under small initialization, the quantities k. k; q. qi R and k K; - of Rs 0 are small at initialization and remain small throughout training. Thus, the conservation law enforees the f norms the value, key, and query be approximately the same throughout training, ||& 2| |1 = o?.
We here prove that Equation (83) holds regardless of the choice of the loss function.
Proof. We can use the generic gradient flow equation in Equation (5) to calculate the gradients of k k;s q qi and vf
dk: k;
ck:
dL dg
dC
2k;
2E
k/
2B
dūg 3qg 26
U;ki
dc
dl
dūg dk;
dq.
dq
q dC dgg
dc
24:
T
2E
2E
U:q x kT (3)
dt
dt
dgg dq:
dg
duž
du
dC dũg
dC
20
21E
Ut dig des
2E
U;
k: By
dl
dl
dg
We that the gradients of kj k;; q qí, and of are equal, regardless of the specific choice of the loss function C. Hence, the following conservation law holds for any loss function:
d (k k: qTq;) =0, di (k:T k; vt) O. dl
d
F. Linear Attention with Separate Low-Rank Key and Query F1. Justification for Zero Blocks Assumption
We initialize v; 0, kirr 0 (i 1....
.R), and prove that they will stay zero throughout training.
Proof. The bottom right entry of the output of linear attention with separate rank-It key and query is
Úg ATTNs (X) ())+N,
q
H
2 En 2,
Fnyyn
Ki
kin
D2
Ui]
N Σn Un
Σn sh
ki,
k
Fqq
1=1
H
R
R
Σ
Å
N L
Σ
ki + v.3 2
k T
qir
+ 0. B
ki, qrrr + w,
Aw Σ Ri.
xqq
If we initialize
0, ki 0, Dg is
H R
H R À
9n
3
ki
w
ocki, Zg.
1=1
We now calculate the gradient updates of v. = 0, ki, and prove their gradients are zero if their initialization zero.
36

---

Training Dynamics of In-Context Learning in Lincar Attention
The gradient update of or contains B(w), which is zero. Similarly Equation (40), we have
((*+
n
TØ;
E
(yg
Ún)
N
Σ
kir Uurt
+B²
2
ki, Aur
Lqq
11
R
n Σ
1 (A
E
w
w
oekir D
im
rm
H
(A+
1
R
Ew (w))
BB
B
Zq
A
Ug Ii
N
Σ
ki
2=1
r=1
0.
The gradient update of kisr contains Euw (w Aww! whose entries are linear combinations of third moments the zero-mean normal random variable w, and are thus zero. Similarly to Equation (41), we have
Tki. = E U. 3 + w: (w)
(yg
gg) (n)
H
E
dw
ww Fig
w
A
The ki
gir
1=1
R
Ew
w "ww")
E
21
Bg
A
o q,
0.
F.2. Gradient Flow Equations
Based on the gradient flow training rule in Equation (5), the gradient flow dynamics of linear attention with separate rank- R key and query
R k E (3(y4-
R
H
TÚE Σ
(1))0 qr.r =
Σ
A
A2
w uaker R. A
Is.r
(84a)
r=1
i=1
Ther e,E (B(y (1)=z qirr U/3
A
E
A²
Σ
Uiki, TTT- A
asrro
(84b)
1=1 y'
R
Tqu.
V.k. -ukk, (B(yg 0g)=zg) = Us
A A
Σ
K E
(A2) ker
(84c)
where i = 1, 1,... H, r= 1, ... R. and the data statistics E (A²) is calculated in Equation (31).
F.3. Fixed Points
We use M(S..) to denote a set of fixed points that correspond to learning m (m = 0, 1, D) out of the D eigenvectors,
M(S..) = -(or.. WK 1:H,
QQ 1: conditions (C1)-(C3) are met
(85)
where the set Sn specifies the indices of the learned eigenvectors.
Sn S (1,2, .D. |mml m.
(86)
The three conditions for Equation (85) are:
37

---

Training Dynamics of In-Context Learning in Lincar Attention
(CI) The heads sum up to fit the eigenvectors with indices Sm.
H
R
-1 [+1+tx(A)/A)
U k.k qúr f==== =
Σ
\a 1+
e
(87)
N
deSm
(C2) For heads with a nonzero value weight, 1/2 + 0, k a Zor (r 11=1, R) all lie in the span of (ealdes...
C3 For heads with zero value weight, 0 = 0,
R
Σ E 1=1 ddS
eae qur = 0.
(88)
With the same reasoning as Appendix E.3, one can show the weights satisfying these three conditions have zero gradients and thus are fixed points. Though conditions (3),C3 do not explicitly specify the weights, they are feasible conditions. One possible weight configuration that satisfies all three conditions is to let ki,r qi.r (r * 1) be zero and let V/ k;,1 q:,1 be the same as the fixed point for linear attention with rank-one key query, where the low-rank case falls back into the rank-one case. Therefore, the fixed points described in Equation (85) are valid and feasible. Linear attention with separate rank-R R key and query has the same 2D fixed points in the function space as its rank-one counterpart.
F.4. Saddle-to-Saddle Dynamics
For linear attention with rank- R key and query, the gradient updates of the key and query weights in Equation (84). k.r qi.r include the factor Us which is the shared across ranks r 1,... R but unique to each head. In linear attention with rank-one key and query initialized with small weights, the weights in head, vs ki, qı escape from the unstable zero fixed point to drive the first abrupt drop of loss. Similarly, in the rank- R model, the value weight U; and pair of key and query weights ki.,, qs. in head escape from the zero fixed point to drive the first abrupt drop of loss.
However, the subsequent dynamics differ between the the rank-one and rank- models. In the rank-one model, the loss will undergo a conspicuous plateau until weights a new head, uy ky, qu (i' + i), escape from the zero fixed point to grow. By contrast, in the rank-R model (R 1), the loss will plateau briefly or not plateau because new pair of key and query weights in the same i-th head, k qi.r' (r' # r), can quickly grow to drive the loss drop. A new pair of key and query weights in the i-th head grows faster than the key and query weights in a new head, because the value weight in the i-th head, U. has already grown during the first abrupt loss drop. Since the gradient updates of all key and query weights in the i-th head include the factor es larger value weight leads to larger gradient updates for the associated key and query weights. We plot the value weights with D = 4 and ranks R = 2, 3, 4 in Figures 3b and 11 to show: the loss drop after a conspicuous plateau corresponds to a new value weight escaping from zero, while the loss drop after brief plateau does not.
We plot the loss trajectories with D 8 and different ranks in Figure 12 complement Figure in the main text.
1.0
Loss
((M)))
1.0
Loss
E(MM)
1.0
c(Ms)
%
0.8
0.8
0.8
((M,))
((Ms))
0.6
Loss ((Ms)
((Ms))
((Mb)
0.4
((Ms)
2(M.)
((M))
2(M)
0.0
0.0
2500 Training staps
5000
2500 Training steps (
5000
2500 Training steps è
5000
(a) R = 2
(b) R=
(c) R =4
Figure 11. Loss and value weights trajectories. The setting is the same Figure 3b except different ranks R 2,,3, 4. In the rank-one case in Figure 3b, value weights in four heads grow, each corresponding to an abrupt loss drop from L(Mm) to L(Mm-1) (m = 0, 1, 3). the rank- case, new value weight grows big from small initialization when the loss decrcases from L(M.... to L(Mm+1) for m that divides R. Here D= = 4, N = 31, H = 5, and A has eigenvalues 0.4, 0.3, 0.2, 0.1.
38

---

Training Dynamics of In-Context Learning in Lincar Attention
R
R
R
R
Simulations
Theory
Theary
Theary
Theory
1.0
1.0
1.0
0.5
00
0.0
0.0
0.0
10
2
5
10
10
Training steps (x10*)
Training steps (x:10")
Training steps (xx10r)
Training sleps (x:10%)
Figure 12. Same as Figure 4 but with ranks R 3,5, 7. Here == 8, N= 31, H = 9, A has trace and cigenvalues Ad x d_
R.5. Dynamics with Repeated Eigenvalues
We have demonstrated that linear attention with separate key and query exhibits loss plateaus during training when the cigenvalues of the input token covariance matrix, A, are distinct. When A has repeated cigenvalues, linear attention with separate key and query can also exhibit loss plateaus due to the different random initial weights in each head. In the case with distinct eigenvalues (Figure 3a), plateau duration is determined by both the size of the eigenvalues and the random initialization. the case with repeated eigenvalues (Figure 13, leftmost panel), the plateau duration is determined solely by the random initialization.
R
R
Simulations
Theory
Simulations
Theory
Simulations
Theory
Simulations
Theory
1.0
1.0
1.0
2(M))
0.8
0.8
0.8
0.8
((M))
Loss
0.4
0.4
0.4
0.4
2(M))
(LMs)
((M))
0.0
0.0
0.0
0.0
5000
10000
10000
10000
10000
Training steps
Training steps
Training stops
Training steps
Figure 13. Loss trajectories of multi-head linear attention with separate key and query. The setup is the same as in Figure 3a except that A eigenvalues 0.35, 0.35, 0.15, 0.15. The four panels differ only in rank of key and query weights. Although some eigenvalues are equal, the loss trajectory of ATTNs with R = 1 can still exhibit plateaus when learning them, due the different random initial weights in each head. The plateaus may also be skipped certain random seeds.
F.6. Conservation Law
The gradient flow dynamics of linear attention with separate key and query in Equation (84) implies conservation law, The value, key, and query weights in head obey
d ki, ((r) =0, dt
R
Σ
k
kirr
0.
(89)
dt
We here prove that Equation (89) holds regardless of the choice of the loss function.
Proof. We can use the generic gradient flow equation in Equation (5) to calculate the relevant gradients
dk J kirr
1,p
=2k,
dkir dt
2E
kT
dC dÿg 2E
d.
U.
k.
B.t
(90a)
dt
dØ dk,rr
dØ,
dq., qir = 24 dt
dqa.
dL dÿg 2F
dc
2E
viq T
(90b)
dt
diy dqu.r
dgg
39

---

Training Dynamics of In-Context Learning in Lincar Attention
du du 2E
dC
diuu
dC dg
20;
= 2
E
viB
(90c)
dt
dt
dÿn du;
r=1
Comparing Equations (90a) and (90b), we see that the following holds regardless of the specific choice of the loss function
dk T kirr dq. -i.r
dt
dt
Similarly, comparing Equations (90a) and (90b) with Equation (90c), we obtain
n
dk kir
duž
dt
dl
G. Training Dynamies of In-Context and In- Weight Learning
In this work, we focused on the training dynamics of ICL abilities. Other than ICL, attention models can also learn in weight, that is solving the task by memorizing the map between the query input and the target output without using the information in context. The arbitration between in-context and in-weight learning may depend on the properties of the training data (Chan et al., 2022). To focus on the dynamics of ICL, we considered purely ICL task, which is in-context linear regression with the task vector sampled from zero-mean standard normal distribution, (0, I). Since memorizing any particular task vector does not effectively decrease the loss, linear attention develops only ICL ability during training, shown in Figure 14a.
If the task vector uo follows different distribution, the training dynamics involves the development both in-context and in- weigt learning abilities. In Figure 14, we let the task vector for some of the training sequences be fixed and sample the rest from standard normal distribution to elicit in-weight learning ability. We plot the training loss, in-context learning test loss, and in-weight learning test loss for varying portions of fixed task vectors Figure 14. The larger the portion of fixed task vectors, the lower the loss the model can achieve by memorizing the fixed task vector in weight. We indeed observe the training loss and in-weight learning test loss are lower right after the first abrupt loss drop when the portion is larger.
The technical consequence of fixing some of the task vectors is that Equations (40) and (41) break. In other words, we cannot assume the certain blocks of the value and the merged key-query matrices are zero as in Appendix D.I. Without the zero block assumption, the linear attention model implements
H (A+ N Bo
Σ 0. (AA++
) U: U²+ Mi++,p a.
U, B
Bug + 0t N
Σ uuu n
Ziq-
(91)
Equation (91) include not only a linear map of the cubic feature z vec(Bx but also linear maps of additional features, [A+ + 1 Zy. 1/2 E u. Future work could analyze the gradient descent dynamics of the model described
by Equation (91), building on our results on the dynamics of in-context learning to explore its interactions with in-weight learning.
Reproducibility
Code reproducing our main results is available GitHub: https: httt::///gthabb..cm//ddzzagag http:/yeiizhann/iiinatnn
40

---

Training Dynamics of In-Context Learning in Lincar Attention
train
test
Loss
test IC
Loss
Loss
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
2500
5000
2500
5000
2500
5000
Epoch
Training steps f
Training steps t
(a) 0% fixed task (in-context)
(b) 20% fixed tasks
(c) 40% fixed tasks
Loss
Loss
Loss
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
2500 Training steps /
5000
2500 Training steps f
5000
2500 Training steps /
5000
(d) 60% fixed tasks
(c) 80% fixed tasks
(f) 100% tasks (in-weight)
Figure 14. Dynamics of in-context and in weigttght learning in linear attention with merged key and query. The training set is the same as in-context linear regression task described Section except that a portion of the task vectors u are fixed. portion of fixed task vectors indicates how much training samples can be fitted with the in ne-weighh learning solution, that is memorizing the fixed task vector. The in context learning test loss is evaluated on test sequences whose task vectors are all sampled from (O, I). The in-weight learning test loss is evaluated test sequences whose task vector the same fixed task vector from the training set. Here D 4,N N 31, H 8, A I/D.


Reviews:
[{"review_id": "t8eKLYbpop", "reviewer": "Reviewer_Hj7w", "summary": "The paper investigates the theoretical understanding of gradient descent training dynamics for multi-head linear self-attention models performing in-context linear regression tasks. It analyzes two parametrizations of linear self-attention: one where key and query matrices are merged (ATTNM), and another with separate key and query matrices (ATTNS). For the merged model (ATTNM), the authors identify two fixed points in training dynamics and find an abrupt loss drop, analytically solved under specific conditions. In the separate model (ATTNS), exponentially many fixed points are found, with the loss trajectory showing saddle-to-saddle dynamics where the model progressively learns principal components of the data. Ultimately, both models implement approximate least squares regression in-context upon convergence, providing insights into how parameterization influences the evolution of in-context learning (ICL) abilities during gradient descent training​", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for a detailed review and for raising constructive questions. We're glad to know that you think our claims are backed by clear, precise, and convincing evidence. Below, we present some new experimental results and respond to questions.\n\n- **Softmax Attention**\n\n  We've added Figure 10 (see [URL](https://icml1939.tiiny.site/)) to show that the different training dynamics of linear $\\mathsf{ATTN_M}$ and linear $\\mathsf{ATTN_S}$ also occur in their softmax counterparts. Figure 10 follows the same setup as Figures 1–2 for linear attention, with the only difference being adding the softmax activation function for the attention calculation. We observe that softmax $\\mathsf{ATTN_M}$ exhibits a single abrupt loss drop, whereas softmax $\\mathsf{ATTN_S}$ undergoes multiple loss drops, separated by phases of conspicuously slower training. This demonstrates that our findings and theoretical intuition are not unique to linear attention but can also extend to softmax attention.\n\n- **Effect of Initialization Scale**\n\n  We'd like to clarify that many of our results do not rely on initialization: the fixed points, the connection to MLP, and the ICL algorithm we identified for the converged network. Then, we analyze the training dynamics under small initialization.\n\n  We've added Figure 11 (see [URL](https://icml1939.tiiny.site/)) to demonstrate the effect of initialization scale on dynamics. For $\\mathsf{ATTN_M}$, increasing the initialization scale shortens the plateau before the single abrupt loss drop. For $\\mathsf{ATTN_S}$, increasing the initialization scale shortens all the plateaus between any two successive abrupt loss drops. For both models, the loss trajectories with the largest initialization exhibit an exponential decay shape, which is the hallmark of lazy learning (NTK regime). In existing theoretical literature, the standard choice of initialization is either the large initialization of exponential-shaped lazy learning, or small initialization of rich feature learning as in our paper. The practical initialization scheme is usually somewhere in between. In Figure 11, we see that dynamics from the intermediate initialization seems like a mix of the exponential-shaped and the abrupt sigmoid-shaped curves. Such mixed curves are often seen in practice, e.g. in induction head emergence in natural language settings ([Olsson et al. 2022](https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html#argument-phase-change); Argument 1). Our dynamics analysis focuses on the rich learning regime and provides analytical insight into such phenomena, which we believe is a first step toward understanding dynamics in naturalistic settings.\n\n- **Simulations with Larger** $N, D$\n\n  Thank you for this suggestion. We've added Figure 13 (see [URL](https://icml1939.tiiny.site/)) in which we train $\\mathsf{ATTN_M}$ and $\\mathsf{ATTN_S}$ on a dataset with $N=128,D=64$. The loss trajectories are qualitatively similar to those in lower-dimensional cases, despite being noisier. This suggests that our findings do not break in high-dimensional settings.\n\nWe've added the new figures to our appendix and will use the extra page in final revision to include as many as possible in the main text.", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence. Specifically, the authors provide detailed theoretical analyses, derivations, and simulation results that align closely with their main theoretical claims. Here’s a breakdown of how effectively the claims are supported:\n\n### Claims Supported by Strong Evidence:\n1. **Existence of Fixed Points**:  \n   - **Evidence:** The authors derive precise analytic forms for fixed points in both the merged (ATTNM) and separate (ATTNS) parametrizations. The existence of these fixed points is rigorously demonstrated mathematically (e.g., equations provided explicitly such as (9a)-(9b), (15)-(17)), providing strong support.\n\n2. **Abrupt Loss Drop (ATTNM)**:  \n   - **Evidence:** Both theoretical analysis (with analytical solutions, such as equation (10)) and detailed simulations (shown in Figure 1) clearly show the existence and nature of the single abrupt drop in loss. The analytical solutions are explicitly derived, and the simulations match these theoretical predictions.\n\n3. **Saddle-to-Saddle Dynamics (ATTNS)**:  \n   - **Evidence:** Authors convincingly illustrate through detailed analytic derivations and numerical simulations (Figure 2) the sequence of abrupt loss drops separated by plateaus. Their scalar ordinary differential equation approximation (equation (20)) closely matches simulated trajectories, reinforcing their claims.\n\n4. **ICL Algorithm Interpretation (Principal Component Regression)**:  \n   - **Evidence:** The authors explicitly derive the algorithm implemented at each fixed point, demonstrating how the model progressively learns principal components (equations (18), (21)). Simulations confirm the model aligns with eigenvectors sequentially during training (Figure 2c).\n\n5. **Impact of Parametrization (Rank of Key-Query Matrices)**:  \n   - **Evidence:** Simulation results (Figure 3) convincingly demonstrate the effects of different ranks (R = 1, 2, 4, 8) on the dynamics and duration of loss plateaus, supporting claims about parametrization affecting training dynamics.\n\n### Potential Minor Limitations (though not significantly problematic):\n\n- **Simulation Scope**:  \n  The numerical simulations strongly align with the theory, yet they predominantly focus on low-dimensional cases (e.g., dimension \\(D=4,8\\)). Although suitable for clear visualization and validation, higher-dimensional simulations could strengthen the general applicability of their results.\n\n\nIn summary, all key claims in this paper are backed by clear, precise, and convincing evidence.", "methods_and_evaluation_criteria": "n/a. theory paper, not so many methods and evaluations", "theoretical_claims": "Yes. They all appear reasonable to me.", "experimental_designs_or_analyses": "n/a", "relation_to_broader_scientific_literature": "The paper \"Training Dynamics of In-Context Learning in Linear Attention\" contributes to the broader scientific literature by providing a theoretical analysis of how multi-head linear self-attention models develop in-context learning abilities through gradient descent training. This work builds upon and extends existing research in several key areas:\n\n1. **In-Context Learning in Transformers**: The study aligns with previous findings that transformers can perform in-context learning, adapting to new tasks based on input data without explicit parameter updates. Notably, it complements empirical observations by offering a theoretical framework that explains how these capabilities emerge during training. \n\n2. **Training Dynamics and Fixed Points**: The identification of fixed points in training dynamics and the characterization of abrupt loss drops or saddle-to-saddle dynamics provide a deeper understanding of how self-attention models learn. This insight is valuable for designing more efficient training protocols and architectures. \n\n3. **Parametrization Impact**: By analyzing different parametrizations (merged vs. separate key and query matrices), the paper highlights how architectural choices influence learning dynamics. This perspective is crucial for developing models that balance computational efficiency with learning efficacy.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "n/a", "other_comments_or_suggestions": "n/a", "questions_for_authors": "**Generalizability Beyond Linear Tasks:**\nThe paper thoroughly analyzes linear attention trained on linear regression tasks. Have you investigated whether similar saddle-to-saddle dynamics and progressive improvement phenomena also occur for non-linear tasks or in practical transformer models with softmax attention?\n\n**Effects of Initialization Conditions:**\nThe analytical derivations assume very small (or infinitesimal) initialization scales. How sensitive are your theoretical results (fixed points, saddle-to-saddle dynamics) to larger or more realistic initialization conditions often used in practice?", "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": "3tQdUVODrW", "reviewer": "Reviewer_mXXZ", "summary": "This paper investigates the training dynamics of in-context learning (ICL) in multi-head linear attention models trained on in-context linear regression tasks. The authors examine two parametrizations: one with merged key and query weights (ATTNM) and one with separate key and query weights (ATTNS).\nThe paper provides a theoretical explanation of how ICL abilities evolve during gradient descent training, revealing either abrupt acquisition or progressive improvements depending on the parametrization of the key and query matrices.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your interest in our work and for a thoughtful review. We appreciate how you positioned our contribution at the intersection of three very important and relevant topics, and we’re glad that you found our results provide a clear explanation for progressive ICL acquisition. Below, we present some new experimental results and respond to questions.\n\n- **Softmax Attention**\n\n  We have added Figure 10 (see [URL](https://icml1939.tiiny.site/)) to show that the different training dynamics of linear $\\mathsf{ATTN_M}$ and linear $\\mathsf{ATTN_S}$ also occur in their softmax counterparts. Figure 10 follows the same setup as Figures 1–2 for linear attention, with the only difference being adding the softmax activation function for the attention calculation. We observe that softmax $\\mathsf{ATTN_M}$ exhibits a single abrupt loss drop, whereas softmax $\\mathsf{ATTN_S}$ undergoes multiple loss drops, separated by phases of conspicuously slower training. This demonstrates that our findings and theoretical intuition are not unique to linear attention but can also extend to softmax attention.\n\n- **Effect of Initialization Scale**\n\n  We'd like to clarify that many of our results do not rely on initialization: the fixed points, the connection to MLP, and the ICL algorithm we identified for the converged network. Then, we analyze the training dynamics under small initialization.\n\n  We've added Figure 11 (see [URL](https://icml1939.tiiny.site/)) to demonstrate the effect of initialization scale on dynamics. For $\\mathsf{ATTN_M}$, increasing the initialization scale shortens the plateau before the single abrupt loss drop. For $\\mathsf{ATTN_S}$, increasing the initialization scale shortens all the plateaus between any two successive abrupt loss drops. For both models, the loss trajectories with the largest initialization exhibit an exponential decay shape, which is the hallmark of lazy learning (NTK regime). In existing theoretical literature, the standard choice of initialization is either the large initialization of exponential-shaped lazy learning, or small initialization of rich feature learning as in our paper. The practical initialization scheme is usually somewhere in between. In Figure 11, we see that dynamics from the intermediate initialization seems like a mix of the exponential-shaped and the abrupt sigmoid-shaped curves. Such mixed curves are often seen in practice, e.g. in induction head emergence in natural language settings ([Olsson et al. 2022](https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html#argument-phase-change); Argument 1). Our dynamics analysis focuses on the rich learning regime and provides analytical insight into such phenomena, which we believe is a first step toward understanding dynamics in naturalistic settings.\n\n- **Practical Implication**\n\n  Given that the scope of our paper is theoretical, we believe our findings and intuition can inform practice by providing a deeper understanding of common phenomena observed in practical settings -- such as abrupt learning and the progressive acquisition of ICL. Though not directly resulting in practical recommendation, we hope our work provides a foundation for future work along the spectrum from theory to practice, ultimately leading to more effective architectural and optimization choices.\n\n  Regarding the implication of the equivalence between $\\mathsf{ATTN_M}$ and MLP with cubic features, it can add a perspective to the open question of whether MLPs can learn in-context (Boix-Adsera et al [1]; Tong and Pehlevan [2]). Moreover, we reveal that MLP may perform ICL more comparable to attention models when provided with polynomial features instead of the original sequence. This may explain why Boix-Adsera et al (Figure 25) observed that MLP fails to learn ICL with the original sequence $X$ as input, but succeeds when the input is augmented with $XX^\\top$. We'll add a discussion of this point in revision.\n\n  [1] Boix-Adsera, E., Saremi, O., Abbe, E., Bengio, S., Littwin, E., & Susskind, J. M. When can transformers reason with abstract symbols? ICLR 2024.\n\n  [2] Tong, W. L., & Pehlevan, C. MLPs Learn In-Context on Regression and Classification Tasks. ICLR 2025.\n\nWe've added the new figures to our appendix and will use the extra page in final revision to include as many as possible in the main text.", "claims_and_evidence": "The authors provide:\n-  mathematical analysis of fixed points in the loss landscape\n- precise characterization of gradient descent trajectories\n- analytical solutions for special cases\n- simulations that match their theoretical predictions", "methods_and_evaluation_criteria": "The authors rely on \n- a well-established ICL regression task \n- empirical validation through simulations that match theoretical results", "theoretical_claims": "I did not check the proofs of the theoretical claims", "experimental_designs_or_analyses": "The simulations given in the paper cover\n- training simulation of both ATTNM and ATTNS models with different initializations\n- tracking of the evolution of weights and loss during training\n- comparison of empirical results with theoretical predictions", "relation_to_broader_scientific_literature": "The paper is in the intersection of several research topics including:\n- work showing abrupt emergence of ICL\n- work on the optimization properties of attention\n- work on the stage-wise dynamics of transformers", "essential_references_not_discussed": "The paper covers thoroughly the relevant literature.", "other_strengths_and_weaknesses": "Strengths:\n-The work bridges optimization theory and the emergent capabilities of transformer models\n- The characterization of saddle-to-saddle dynamics in ATTNS offers a clear explanation for progressive ICL acquisition\n\nWeaknesses:\n- The analysis is limited to linear attention, and it's unclear how much of it generalizes to softmax attention used in practical transformers\nThe paper focuses on ICL regression tasks, which are simpler than the complex language tasks where ICL is typically observed", "other_comments_or_suggestions": "It would be valuable to include a discussion on how the theoretical insights might translate to practical  training recommendation", "questions_for_authors": "How sensitive are your findings to the choice of initialization scale? \nThe paper demonstrates that ATTNM is equivalent to a two-layer fully-connected network with cubic features. Could this have a practical implication on transformer architectural choice ?", "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": "0LNbcbFG9T", "reviewer": "Reviewer_etiN", "summary": "This paper provides a theoretical analysis of how linear attention models acquire in-context learning abilities through gradient descent training on an in-context linear regression task. The authors study two parameterizations of multi-head linear attention. In the merged key-query setting (ATTN$_M$), they show that the model has exactly two fixed points—a zero solution and a global minimum manifold—and that small-initialization training follows a single abrupt drop in loss before converging to a least-squares in-context regression solution. The authors also derive an analytical time-course in the special case of an identity covariance matrix and infinitesimal initialization. By contrast, in the separate key-query setting (ATTN$_S$), they show that the model admits exponentially many fixed points corresponding to different subsets of eigenvectors of the token covariance. In this scenario, the training dynamics involve multiple abrupt loss drops separated by plateaus, each plateau corresponding to learning one additional principal component in context. Although the final solution again approximates least-squares regression, early stopping yields principal component regression on a partial subset of components.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for a very detailed and thoughtful review. We're glad to know that you find our findings instructive, rigorous, and well-presented. We'd like to present some new experimental results and respond to questions. All of our added new figures can be found at [URL](https://icml1939.tiiny.site/).\n\n- **Rationale for $1/N$ Scaling**\n\n  We follow the scaling choice in seminal theory works on linear attention trained for in-context linear regression [1,2]. You're correct that the model can be trained to adapt to the sequence length without the $1/N$ scaling. However, including the scaling factor grants the model the flexibility to be trained and tested with varying context lengths $N$. We've added Figure 14 in which the models are trained on sequences of varying lengths -- this would be unfeasible without the scaling. To explain this point, consider that $\\mathsf{ATTN_M}$ computes\n\n  $$\n  \\mathsf{ATTN_M}(X) = X +\\sum_{i=1}^H W^V_i \\frac{XX^\\top}N W^{KQ}_i X\n  $$\n\n  The blocks in $XX^\\top$ contain the term $\\sum_{n=1}^N x_n y_n$, which diverges as $N$ increases. In contrast, the blocks in $\\frac{XX^\\top}N$ contain $\\frac1N \\sum_{n=1}^N x_n y_n$, which remains bounded as $N$ increases. Including the $1/N$ scaling allows the model to process sequences of varying lengths without needing to adapt weights specifically for each $N$.\n\n  > \"we could also instead scale $W_i^{KQ}$ and $W_i^V$ by $1/N$ ...\"\n\n  You're correct that we can either scale the $XX^\\top$ by $1/N$ or scale $W_i^{KQ}, W_i^V$ by $1/\\sqrt N $. However, in both case, we arrive at the same expression in line 246, Equations (18) and (21). In our case, the $1/N$ term is present because of the $1/N$ we build into the model. In your case, the $1/N$ term is also present because $W_i^{KQ}, W_i^V$ are scaled by $1/\\sqrt N$. Thus, line 246, Equations (18) and (21) are not affected by whether the scaling is built in the model or adapted through training.\n\n  Thank you for this constructive feedback. We'll incorporate this more detailed explanation for the scaling in revision.\n\n- **Why Full-Rank $W^{KQ}_i$ in $\\mathsf{ATTN_M}$?**\n\n  We agree that $\\mathsf{ATTN_M}$ diverges from the established multi-head attention structure in practical transformers. Our intention in studying both $\\mathsf{ATTN_M}$ and the low-rank $\\mathsf{ATTN_S}$ is that: $\\mathsf{ATTN_M}$ is widely used in theoretical literature on transformers (see a list of references in lines 60-65) and low-rank $\\mathsf{ATTN_S}$ is closer to the parametrization in practical transformers. Given that a significant portion of the theoretical literature has used $\\mathsf{ATTN_M}$, we believe that analyzing it can be of independent interest to transformer theorists, if not to practitioners.\n\n  Your insight regarding that the multiple heads in $\\mathsf{ATTN_M}$ are somehow \"repetitive\" or redundant is also right. In other words, the multi-head $\\mathsf{ATTN_M}$ is overparameterized. While overparameterization seems unnecessary in terms of expressivity, its effect on training dynamics and convergence properties is a fundamental, long-standing question in deep learning, and is worth studying. Perhaps surprisingly, we find that the $W^{KQ}_i$ weights in different heads of $\\mathsf{ATTN_M}$ stay parallel in training, as shown in Figure 4 and Appendix C.7. This provides a hint at how overparameterized attention models learn and generalize.\n\n- **Effect of Initialization Scale**\n\n  We've added Figure 11 for this. Please refer to our rebuttal to reviewer Hj7w for details (we reached space limit for this rebuttal).\n\n- **Multiplicity of Eigenvalues**\n\n  Thank you for this interesting question. We have added Figure 12 in which $\\Lambda$ has repeated eigenvalues. Six runs from different random initialization are plotted. Linear $\\mathsf{ATTN_S}$ with rank-1 key and query may still exhibit plateaus when learning the repeated eigenvalues due to the different random initial weights in each head. It may also skip the plateau for certain random seeds. \n\n  In the case with distinct eigenvalues, the plateau duration is determined by both the size of the eigenvalue and the random initialization. In the case with equal eigenvalues, the plateau duration is determined by only the random initialization.\n\n- **Simulations with Larger** $N, D$\n\n  We have added Figure 13 in which we train $\\mathsf{ATTN_M}$ and $\\mathsf{ATTN_S}$ on a dataset with $N=128,D=64$. The loss trajectories are qualitatively similar to those in lower-dimensional cases, despite being noisier.\n\n- We'll add the reference “Transformers as Statisticians\" (Bai et al., 2023). It's indeed an omission -- thanks for your note!\n\n- We'll cut back the introduction for softmax attention in Section 2.2 and clarify the subscript $\\mu$ in $X_\\mu$ denotes the index of a training sample in Section 2.1. Thanks for your suggestion.\n\n[1] Ahn, Cheng, Daneshmand & Sra. NeurIPS 2023. arXiv:2306.00297\n\n[2] Zhang, Frei & Bartlett. JMLR 2024. arXiv:2306.09927", "claims_and_evidence": "The paper’s primary claims that multi-head linear attention with merged key–query exhibits an abrupt loss drop while separate key–query parameterization shows stage-wise, saddle-to-saddle dynamics are supported both by mathematical derivations and simulation results. The evidence is convincing within the assumed setting (small initialization).", "methods_and_evaluation_criteria": "The authors’ method is essentially to analyze multi-head linear attention in two parameterizations (merged vs. separate key/query) under gradient-flow training on an in-context linear regression task. The evaluation criteria revolve around loss trajectories under gradient flow, an the limiting in-context learning algorithm (whether the model replicates least-squares regression or principal component regression). Given that the paper’s focus is on theoretical understanding, the chosen tasks (toy in-context linear regression) and the analysis of training trajectories are appropriate.", "theoretical_claims": "The authors’ method is essentially to analyze multi-head linear attention in two parameterizations (merged vs. separate key/query) under gradient-flow training on an in-context linear regression task. The evaluation criteria revolve around loss trajectories under gradient flow, an the limiting in-context learning algorithm (whether the model replicates least-squares regression or principal component regression). Given that the paper’s focus is on theoretical understanding, the chosen tasks (toy in-context linear regression) and the analysis of training trajectories are appropriate.", "experimental_designs_or_analyses": "I think that the experimental design is well executed. The loss trajectories and weight evolution curves match the analytical results and confirms the soundness of the analysis. One potential limitation is that the experiments focus on relatively small $N$ and $D$ values; it would be valuable to see the experiments in larger settings.\n\nOther than that, some additional things that might worth trying are the cases where the covariance matrix $\\Lambda$ has repeated eigenvalues and check whether multiple heads learn in parallel or whether the training still unfolds in discrete stages. Also, the paper’s analysis hinges on small initialization, it would be instructive to demonstrate how differently scaled inits affect the plateau structure (especially in ATTN$_S$).", "relation_to_broader_scientific_literature": "This paper draws on recent work on linear attention (e.g., Von Oswald et al., 2023; Zhang et al., 2024a,b) and connects the phenomenon of abrupt in-context learning during training transformers to a analytically tractable setting. It benefits from the results in deep linear networks (Saxe et al., 2014; Woodworth et al., 2020) to show how parameterization (merged vs. separate key-query) yields distinct training behaviors of a single abrupt loss drop or multiple progressive plateaus. Overall, the work extends in-context linear regression analyses to a multi-head linear attention setting and also connects to known convergence results from linear networks (Arora et al., 2019; Shamir, 2019).", "essential_references_not_discussed": "One specific omission is the “Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection” (Bai et al., 2023), which frames multi-head attention learning as algorithm selection. Even though it's not a convergence proof, it's still relevant to how, in the rank-1 case of this paper, individual heads each specialize to a different principal component and captures separate directions in the data.", "other_strengths_and_weaknesses": "Overall, I think that the paper is well-written and presented. One of the paper’s main strengths is its rigorous theoretical analysis that clarifies how different parameterizations affect the training dynamics of multi-head attention models. Besides, the paper well-integrates the prior results on linear networks. The experiments complements the theoretical discussions. The connection shown in the paper between progressive eigenvector learning and in-context abilities is also instructive.\n\nThat said, I believe the paper needs more justification on why the model they use for multi-head linear self-attention in the case of merged key-query matrices is appropriate. There's frequent connecting discussions which uses the limit of $N \\to \\infty$ and I believe equivalence to the model which doesn't normalize the output by $N$ should be discussed. Also, some discussions or experiments on how deviations from the idealized assumptions (such as small initialization or white input covariance) might change the results could be helpful.", "other_comments_or_suggestions": "- This is something minor but the matrix $X_{\\mu}$ is not defined before being referenced in Section 2.1.\n- Again something minor, the paper’s main focus is on analyzing linear attention, yet Section 2.2 provides a relatively detailed introduction to the multi-head softmax attention model, which is not used in later analysis. Why did the authors choose to present the standard softmax-based model in detail rather than only mentioning it briefly?", "questions_for_authors": "1. I want to ask why the authors chose to scale the model's output by $1/N$. Wouldn't the training procedure naturally allow the weights to adapt to the sequence length without that additional scaling\n2. Although the Section 2.4 has the low rank matrix assumption to match the realistic multi-head architecture, the Section 2.3 omits that assumption. Why the difference? Doesn't this setup than diverge from established multi-head structure of transformers? Because in this case, you're effectively repeating the same input $X$ with different parameters $H$ times. Moreover, I think this also might hurt the equivalence to the two-layer neural net since this forces $W_1$ to be low rank. Whereas, this does not naturally occur with the random initialization.\n1. I'm also a bit confused by the limit with respect to $N$ in line 246. As I said, introducing $1/N$ is artificial and we could also instead scale $W_i^{KQ}$ and $W_i^V$ by $1/\\sqrt{N}$. This change supposed to not affect the input covariance $\\Lambda$, yet, it seems that when we remove $1/N$ from the line 246, it affects the result, why? This also seems to effect the Equations (18) and (21).\n1. In your analysis of ATTN$_S$, you assume eigenvalues are distinct and that each head learns one eigenvector sequentially. How important is this assumption in your proofs? How does the training behavior change if $\\Lambda$ has multiple identical top eigenvalues? Is it possible for several heads to learn these degenerate eigenvectors in parallel, and does that merge or skip some of the plateaus?", "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": "wVzzVe5tzN", "reviewer": "Reviewer_L7nj", "summary": "The authors consider the training dynamics of linear attention using a joint query-key matrix vs. separate query matrix and key matrix, when learning to learn in-context linear regression. It was then shown that in the former case, there are two (classes of) fixed points, while in the latter case, there are 2^D cases where D is the dimension of the x-vector. Thus, during training, the former case has one single abrupt drop, while for the latter, multiple smaller drops are present. The authors also show how in the latter case, the training would experience D drops, corresponding to learning the eigen structure of the covariance matrix. Several experiments are provided, which mostly align with the theoretical analysis. The converging time courses are also analyzed for some simplified cases, which also match the experiments.\n\n## update after rebuttal\n\nThanks for the response. I like the results and increased the score to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your interest in our work and for raising clear, constructive questions. We're glad to hear that you like our results. Below, we present some new experimental results and respond to questions. The numbering follows the reviewer's \"Questions for Authors\" list.\n\n1. **Nonzero** $\\boldsymbol v_i,\\boldsymbol u_i$\n\n   Thank you for this insightful question. The validity of fixed points is not affect by $\\boldsymbol v_i,\\boldsymbol u_i$. In other words, setting $\\boldsymbol v_i=\\boldsymbol u_i = \\boldsymbol 0$ and the rest of the weights to Equation (9) or (15) results in valid fixed points. This follows from our proof that (i) $\\boldsymbol v_i,\\boldsymbol u_i$ have zero gradient when they are zero; (ii) the rest of the weights have zero gradient when they satisfy Equation (9) or (15). Therefore, all weights in the full model have zero gradient.\n\n   As for dynamics, we agree that in practice,  $\\boldsymbol v_i,\\boldsymbol u_i$ can be nonzero and they evolve in training, affecting the dynamics in interesting ways. We actually discussed this issue in our Appendix F \"Training Dynamics of In-Context and In-Weight Learning\". In our setup, nonzero $\\boldsymbol v_i,\\boldsymbol u_i$ can be interpreted as a form of in-weight learning (IWL). In our main setup of sampling the task vector $\\boldsymbol w$ iid from $\\mathcal N (0,I)$, $\\boldsymbol v_i,\\boldsymbol u_i$ remain near zero throughout training, and the model develops only ICL, as shown in Figure 9a. However, when a portion of $\\boldsymbol w$ is fixed, the evolution of $\\boldsymbol v_i,\\boldsymbol u_i$ is non-negligible. In this case, the model first develops IWL, then loses IWL while simultaneously developing ICL, as shown in Figure 9b-e. This competition between ICL and IWL is indeed intriguing. While a full theoretical treatment is beyond our scope, we hope future studies can build on our analysis of the purely ICL dynamics and explore its interaction with IWL dynamics. If the reviewer finds it beneficial, we're happy to move some of Appendix F to main text.\n\n   We also note that setting $\\boldsymbol v_i=\\boldsymbol u_i = \\boldsymbol 0$ is common in theory literature on ICL in linear [1,2] and softmax [3,4] attention.\n\n2. **Cumulative Loss**\n\n   We have added Figure 14 (see [URL](https://icml1939.tiiny.site/)) to demonstrate that using the cumulative loss doesn't change the training dynamics, other than the loss values being different. Consistent with using the last-token loss, training $\\mathsf{ATTN_M}$ with cumulative loss exhibits a single abrupt loss drop, and training $\\mathsf{ATTN_S}$ exhibits multiple loss drops. In our in-context linear regression setup, training with cumulative loss is effectively equivalent to training with the last-token loss on sequences of varying lengths. Because the linear attention model includes a $1/N$ scaling factor, it can be trained and tested with varying sequence lengths $N$. For a detailed discussion on the $1/N$ scaling, please refer to our rebuttal to reviewer etiN.\n\n   It may be worth noting that computing loss for only the last token is also common in the literature, such as in Von Oswald et al. [5] and Zhang et al. [1].\n\n3. **Softmax Attention**\n\n   We have added Figure 10 (see [URL](https://icml1939.tiiny.site/)) to show that the different training dynamics of linear $\\mathsf{ATTN_M}$ and linear $\\mathsf{ATTN_S}$ also occur in their softmax counterparts. Figure 10 follows the same setup as Figures 1–2 for linear attention, with the only difference being adding the softmax activation function for the attention calculation. We observe that softmax $\\mathsf{ATTN_M}$ exhibits a single abrupt loss drop, whereas softmax $\\mathsf{ATTN_S}$ undergoes multiple loss drops, separated by phases of conspicuously slower training. This demonstrates that our findings and theoretical intuition are not unique to linear attention but can also extend to softmax attention.\n\nWe've added the new figures to our appendix and will use the extra page in final revision to include as many as possible in the main text.\n\n[1] Ruiqi Zhang, Spencer Frei, Peter Bartlett. Trained transformers learn linear models in-context. JMLR 2024.\n\n[2] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.\n\n[3] Juno Kim, Taiji Suzuki. Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape. ICML 2024.\n\n[4] Yu Huang, Yuan Cheng, Yingbin Liang. In-context Convergence of Transformers. ICML 2024.\n\n[5] Johannes Von Oswald et al. Transformers learn in-context by gradient descent. ICML 2023.", "claims_and_evidence": "The empirical studies are well designed and appear to support the theory well.", "methods_and_evaluation_criteria": "The authors train the linear attention models for the setting, and compare the observed behaviors with the theoretical analysis.", "theoretical_claims": "Yes, the proofs mostly look correct.", "experimental_designs_or_analyses": "Yes, the experiments are described with sufficient details, and support the theory.", "relation_to_broader_scientific_literature": "N/A.", "essential_references_not_discussed": "Sufficiently discussed.", "other_strengths_and_weaknesses": "While I like the results in general, there are several weaknesses inherent in the study and the formulation of the problem:\n\n1. One main concern is the assignments of 0's to the v_i, u_i to the query and key matrices. Though this might be acceptable for the analysis of the eventual optimal solution, it is less acceptable for the analysis of the training dynamics, because they might impact the dynamics in an unpredictable way. Note that in practice, these matrices are usually assigned random values at initialization, instead of being assigned 0. \n\n2. In Garg's work (and some others), the training loss function is the average of the losses over the whole context window \\sum_{n=1}^{N}, instead of only the last position N+1. I feel the former is a more accurate reflection of the \"in-context learning\" concept, while the latter seems unrealistic.  \n\n3. The linear attention model allows the analysis to be done more precisely, but also too simplified. The softmax function is important in many settings. Similarly, in practice, a feedforward network is used after the attention mechanism, and multiple attention layers are usually used. \n\n4. While the difference in the training dynamics is interesting, it is not clear whether this distinction is important in practice, because they seem to be unique to the linear-attention, and will likely break with a softmax function, the FF network, and multiple layers.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Will non-zero assignments of v_i and u_i change the structure of the fixed points?\n2. Will using the cumulative loss as in Garg change the training dynamics?\n3. Will using softmax change the structure of the training dynamics?\n\nI would consider raising the score if the authors can answer these questions satisfactorily.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper studies the gradient dynamics of multi-head linear self-attention trained for in-context linear regression and finds that two different parameterizations lead to different evolution of ICL ability. The reviewers agreed upon positive ratings. Overall I recommend accept.
Decision: Accept (spotlight poster)