Transformers Provably Learn Two-Mixture of Linear Classification via Gradient Flow
Authors: Hongru Yang
Zhangyang Wang
Jason D. Lee
Yingbin Liang
Venue: ICLR 2025 Poster (2025)
Keywords: transformer
training dynamics
gradient flow
mixture of linear classification
Primary area: learning theory
TL;DR: We study the gradient flow dynamics of training transformer on two-mixture of linear classification.

Abstract: Understanding how transformers learn and utilize hidden connections between tokens is crucial to understand the behavior of large language models.
To understand this mechanism, we consider the task of two-mixture of linear classification which possesses a hidden correspondence structure among tokens, and study the training dynamics of a symmetric two-headed transformer with ReLU neurons.
Motivated by the stage-wise learning phenomenon in our experiments, we design and theoretically analyze a three-stage training algorithm, which can effectively characterize the actual gradient descent dynamics when we simultaneously train the neuron weights and the softmax attention.
The first stage is a neuron learning stage, where the neurons align with the underlying signals. 
The second stage is a attention feature learning stage, where we analyze the feature learning process of how the attention learns to utilize the relationship between the tokens to solve certain hard samples.
In the meantime, the attention features evolve from a nearly non-separable state (at the initialization) to a well-separated state.
The third stage is a convergence stage, where the population loss is driven towards zero.
The key technique in our analysis of softmax attention is to identify a critical sub-system inside a large dynamical system and bound the growth of the non-linear sub-system by a linear system. 
Finally, we discuss the setting with more than two mixtures. 
We empirically show the difficulty of generalizing our analysis of the gradient flow dynamics to the case even when the number of mixtures equals three, although the transformer can still successfully learn such distribution. 
On the other hand, we show by construction that there exists a transformer that can solve mixture of linear classification given any arbitrary number of mixtures.

Published as a conference paper at ICLR 2025
TRANSFORMERS PROVABLY LEARN TWO-MIXTURE OF LINEAR CLASSIFICATION VIA GRADIENT FLOW
Hongru Yang
Zhangyang Wang
The University of Texas at Austin & Princeton University hy by6385@utexas.ed
The University of Texas at Austin at 1 stlasanngeutexa edu
Jason D. Lee
Yingbin Liang
Princeton University
The Ohio State University liang. B89@osu. edu
jasonlee@princeeoon.cc
ABSTRACT
Understanding how transformers learn and utilize hidden connections between words is crucial to understand the behavior of large language models. To study this mechanism, we consider the task of two-mixture linear classification which features a hidden correspondence structure between words, and study the training dynamics of symmetric two-headed transformer with ReLU neurons. Motivated by the stage-wise learning phenomenon observed in our experiments, we design and theoretically analyze three-stage training algorithm, which can effectively characterize the actual gradient descent dynamics when simultaneously train the neuron weights and the softmax attention. The first stage is a neuron learning stage, where the neurons align with the underlying signals. The second stage is an attention feature learning stage, where we analyze the feature learning process of how the attention learns to utilize the relationship between the tokens to solve certain hard samples. In the meantime, the attention features evolve from a nearly non-separable state (at the initialization) to well-separated state. The third stage is a convergence stage, where the population loss is driven towards zero, The key technique in our analysis of softmax attention to identify a critical sub- system inside large dynamical system and bound the growth of the non- lon-linear sub-system by a linear system. Along the way, we utilize novel structure called mean-field infinite-width transformer. Finally, we discuss the setting with more than two mixtures. We empirically show the difficulty of generalizing our analysis of the gradient flow dynamics the case even when the number of mixtures equals three, although the transformer can still successfully learn such distribution. On the other hand, show by construction that there exists transformer that solve mixture of linear classification given any arbitrary number of mixtures.
INTRODUCTION
Since the invention of self-attention (Vaswani al. 2017), transformers have become the dominat- ing backbone architecture in many machine learning applications such as computer vision (Doso-)
vitskiy et al.) 2020 Liu et al. 202I), natural language processing (Devlin et al. 2018 and protein
structure prediction (Jumper et al. [202T). Within the past two years, ChatGPT and GPT4 (OpenAl
2023 Bubeck et al. 2023) along with other large language models (LLMs)) (Tuvrron et al. 2023
Manyika & Hsiao 2023; Gemini 2023 Anthropic (202)) have demonstrated astonishing abilities
language understanding. math Olympiad (AlphaProof & AlphaGeometry 2024) and coding, etc.
Despite the wide range of success of transformers, how those models can achieve such impressive performance still remains largely unknown. One mystery lies how transformers are trained to utilize the connections between words to solve various tasks. For example, consider the prompt
Work done while visiting Princeton University

---

Published as a conference paper at ICLR 2025
"Bob is watching television in the living room. Where is Bob?". To answer the question, the model must utilize the (hidden) correspondence between "Where" and "living room".
There are many previous works studying the mechanism of self-attention under different settings through training dynamics such as (Li ct al., 2024b; Huang et al. 2024, Li et al. 2022 Tarzanagh
et al. 2023b a Vasudeva et al. 2024). However, many of those previous works were either focusing on attentions that are position- based or studying a setting where the words in the dictionary don't t have explicit relationships between each other. The goal of this work to study some simple yet informative setting such that we can get a clear understanding how the softmax attention learns and utilizes certain connections between words to solve tasks through training. To achieve this goal, we consider the task of mixtures of linear classification. In this task. we have dictionary consisting of group identifiers and classification features, and there hidden correspondence structure: each group identifier is associated with one classification feature. Each input contains a group identifier and a corresponding signed classification feature. The rest of tokens are sampled from signed classification features from other groups. The label of the sequence is determined by the sign of the classification feature corresponding to the group identifier. Thus, a model needs to capture the correspondence structure between group identifiers and classification features to solve this task. To see how attention can solve this task, later in our training dynamics analysis, we are going to show that the softmax attention builds such correspondence structure by allocating more weights on the group identifier when it sees classification feature queries.
Although the intuition suggests that unknown that how the transformer models can be trained to learn such hidden correspondence structure. This inspires the following intriguing question:
construct transformer solve this task, it is still
How do transformers learn and utilize the hidden correspondence structure solve mixture of linear classification via gradient descent?
Our contributions. this work, we study the training dynamics of two-headed transformer given two mixtures. Our contributions are summarized as follows:
1. As a guidance for our theory, we first conduct experiments to observe the training dynamics of the transformer where we train all the weights simultancously. Our experiment results show clear stage-wise learning phenomenon where the neuron weights learn before the attention modules.
2. Motivated by our experimental observations, we design a three-stage layer-wise training al- gorithm and further analyze its gradient flow dynamics. We characterize the feature learn- ing process where the attention features evolve from nearly non-separable state a well- separated state. In particular, our analysis captures how the self-attention associates the group signals and classification signals to solve this task. Our analysis closely reflects the actual behavior of the gradient descent dynamics when we train all weights simultaneously.
3. In order to analyze the change in softmax attention, we formulate a set of relevant variables whose evolution can be characterized by large non-linear dynamical system of ordinary differential equations. To make the analysis of the dynamical system tractable, we reduce the dimension of the system by identifying some key variables and forming a smaller sub- system. We are able to relate the evolution of the smaller system with a linear dynamical system to track its dynamics. Our proof techniques can be of independent interest.
4. We explore the training dynamics of the transformers when the number of mixtures goes beyond two. We empirically show the difficulty of analyzing training dynamics even when the number of mixture equals three, although the two-headed transformer model can still learn this distribution. On the other hand, we give general two-headed transformer struction that can solve mixture of linear classification given arbitrary number of mixtures.
1.1 RELATED WORK
Mixture of linear regrrssionnChasificttion Mixture of linear regressiio/cllasiffcation is a classi- cal model in statistics and machine learning (De Veaux 1989; Jordan & Jacobs 1994) which was applied in areas such as object recognition (Quattoni et al. 2004) and machine translation (Liang et al. 2006). This problem can be solved by tensor-hased methods (Anandkumar al. 2014 2012
2

---

Published as a conference paper at ICLR 2025
Hsu et al. 2012; Chaganty & Liang 2013) and the peteottaioonmaximmzziiii algorithm (Khalili & Chen 2007; Yi et al., 2014, Balakrishnan et al. 2017;; Wang et al.) 2013). Previously, mixture of linear classification has been used in studying mixtures of experts (Chen al. 2022).
Training dynamics of transformers. The training dynamics of transformers have been studied under various settings. One particular category is in-context learning. For example, in-context linear regression has been studied under linear attention (Zhang et al.] (024). softmax attention (Huang
al. 2023), non-linear embedding (Yang al.), and multi-task with multiple heads (Chen al.
2024 Other in-context learning settings include classification (Li et al. 2024b), causal structure
NNichani et al., 2024), Markov chain or n- g-garm (Edelman et al., 2024 Makkuva al., 2024 Chen
et ). nearest neighbor (Li et al., 2024c) and chain -f-houht
(Li et al.
2024a).
For non-in-context settings, has been shown that trained transformer can learn spatial structures (Jelassi et al. 2022), topic models (Li et al., 2023b). and feature-posiin position correlation (Huang et al. 2024). In addition, (Tian et al.] 2023) showed that self-attentio behaves like a discriminative sean algorithm. There is one line of works studying setting where the attention weight can converge to a SVM solution (Tarzanagh et al. 2023b a Vasudeva et al. 2024). In addition, (Li et al. 2022) studied classification task with lelevatelevant and label-irrelevant tokens. (Li et al.) (023a)) pro- vided analysis training graph transformers for node classification tasks. <Wang et al.) showed transformers can learn sparse token selection task which lead to optimization--iinn based separation between transformers and MLPs. There are also works trying to prove convergence of transformer training via NTK (Wu et al. 2023 Deora et al.) and mean-field (Kim & Suzuki 2024
Gao et al.).
Our work also study non-in-context learning setting. However, many the previous works were either considering attentions that are position-based studying setting where the words don't have explicit relationships with each other. Our study complements the above work by considering a setting where the dictionary possesses a correspondence structure.
SETTINGS
Notations. For a vector u E Rd, we use diag(v) to denote diagonal matrix with v being the diagonal entries. When we subtract the vector v by scalar a, we subtract each entry of by a, i.e. a E Rd and (v a); e a. For set S with elements in IR and |S] we can vectorize the set S in any arbitrary but fixed order and denote the vector vec(S) RR. We use 22 O, Ō to hide polylogarithnic factors.
2.1 DATA DISTRIBUTION
Definition 2.1 (Mixture of linear classification). Define the K-mixture of linear classification data distribution as follows: Assume d 2 2K. Let C = {cx )K C. IRd denote the group signals and IR denote the classification signals, where C UUV. is orthonormal. Let L be the number of tokens per sample. Each data entry (X,y) E d {+1} is created as follows:
k=1 )==
Sample y ~ Uni form((+1)).
2. Sample k Uniform([]] and lo Uni fni(o]((L)) Set X1, Ck.
3. Sample 1 Uniform([L) (lu3) and set X1, = yuk--
4. For each l2 E [Z]] (1o, h1, sample k iid. Uniform([K]] \ (k)) and set X1, Obk, where i.idd (informm(=+1))
Based on our data distribution, the model needs to utilize the correspondence between the group signals and classification signals to solve this task. We further elaborate our data model as follows. Take the sequence = [c1, U/. 1/1 U2 yy with y1, //2 Ys E {+1} as an example. Based on the form of X, the token e₁ indicates that this instance in group 1. Then the label of this sequence depends on the sign associated with 2j (that shares the same index with cj), and hence is given by V1.
3

---

Published as a conference paper at ICLR 2025
For the training dynamics part of this work, consider K = and L = 3, and discuss the setting with arbitrary number of mixtures in Section 6 Since the samples are invariant of permu- tation between tokens, for k * k' E 82).8.8 E (+1). We refer to the samples Xk,s, as consistent samples and samples Xk,s,-s conflicting samples.
ignore the order of the tokens and define X (ck, Stg, s' Utr)
2.2
TRANSFORMER ARCHITECTURE AND LOSS FUNCTION
Given inputs X E Rdxt, we first define single softmax attention head H Rdx + R as
L
H(X) = w Xsoftmax (x Woz)) Σ E=1
where E Rd, be I. WK, We Rmxd and xι denotes the 1-th token in X. Then, we define our one-layer symmetric two-headed transformer with ReLU neurons f RdxL IR as
f(X) = (H+ -a (X)
where b E R and a is the ReLU activation function. We denote the weights in H4 as w4. WK+, Wo+ and similarly for H.. We introduce shorthand notation K+ W Wx+A. X and Q+ WQ+X. For n2u X, we define the output of attention, or softmax probabil-
ity, as (+)
where I(v) denotes the index such that Pot (X) = softmax ( (XT WKT WoHh) (())
Xi(p)= = =.
Loss Function. We train W WK+ Wo+ to minimize the population loss E (e(ff((((())) where ((x)) log (1 + exp (-+)) is the binary cross-entropy loss. The gradient of the cross-entropy loss given (x(
and we denote g(x) 1/(1 + e*).
GUIDING EXPERIMENTS AND ALGORITHM
3.1 EXPERIMENT RESULTS
To understand the dynamics of the training process, we train the two-headed transformer in Equa- tion ( on the data distribution Definition 2.1 our experiments, all the weights in the trans- former are trained simultaneously via gradient descent with learning rate 0.1.
4
07
02
6150
+.08
482
3008
6.28
4.00
8:28
1080
1300
360
18800 5450
(a) Loss
(b) Neuron alignments
(c) Attention scores
Figure 1: Experiment results. The attention scores and neuron alignment are of the positive head.
We make the following observations from experimental results in Figure I which will guide us to further develop our three-stage training algorithm and our analysis in Section
Learning for different types of samples follows different orders. From Figure la the loss for the samples with consistent signs drops in the early stage of training, while the loss for the samples with conflicting signs only start drop after sufficient amount of training.
There is also learning order of the weights. From Figure Tb the neuron weights start to align with the underlying signals since the beginning of the training. On the other hand, by Figure Ic only after sufficient amount of training, the attention scores start to exhibit no- ticeable change. We further conduct experiments by fixing the neuron weights as Gaussian and only training the attention. Our results show that the model is not trainable in this case.
4

---

large positive Published as a conference paper at ICLR 2025
The time when the attention scores exhibit noticeable change coincides with the time when the loss for samples with conflicting signs drops. This suggests that attentions play a central role in learning those samples.
The attention weights are positively associating the group signals and classification signals from the same group and negatively associating them from different groups. This can be seen from Figure lc that the score cj W W Q++" "1 is increasing
magnitude and the score cf W w 02 decreasing to a large negative magnitude.
In order to understand the mechanism of neurons and attentions separately, the second observation suggests to decompose the training process into stages: neurons are trained first, followed by training the attention. The process can be finished by training neurons again to drive the loss towards zero.
3.2 THREE--TAGE TRAINING ALGORITIM
Motivated by our experiment results, we introduce our three-stage layer-wise training algorithm in Algorithm [ In our layer-wise training algorithm, in the first stage we train the neuron weights W+, w- In the second stage, we train the attention weights WK+, WQ+, WK-, Wq- on the sam- ples with conflicting signs. In the third stage, we train the neuron weights again. Our algorithm design closely reflects our experimental observation in Section
Algorithm Three-stage Training
Initialize _0
0, (: Kt:
W ~ NN w2 ) and be a sufficiently small positive constant
such as 1/2.
while < T do
Stage I: training the neuron weights with all samples
(xw) V Vus (x)f(
de
[
W (t)
@+1
[W w w!l)
end while
while E [Ti, TT] do
> Stage 2: training the softmax attention with the conflicting samples
d(w w
(x,) We (X)) -x.
de (2:)
end while
while E [T2, T3] do
Stage training the neuron weights to minimize the population loss
= -((x))) V.. (x)))
de
(W
+= = (b)) W.(T))
k=
K Q+
end while
4 MAIN RESULTS
Before we state the main results, we first introduce the parameter conditions used in our analysis. Condition
We use the following parameter conditions:
The attention initialization scale satisfies < C < for some small constant C.
-The attention weight embedding dimension m poly(1/w).
We now state the guarantee for Algorithm I under an idcalized infinite-width transformer model which will introduce later in Definition 5.1 Our experiment results show that this idealized model is indeed an accurate model that captures the real training dynamics.
Theorem 4.1 (Main theorem). Under Condition T the infinite-wide width transformer model in Def- inition 5.7 there exists Ti O(1), T3 Ti (log(I/w)), T3 T2 O(1/e) such that after running Algorithm 1 we have E[[(y] (Ts) S
The neuron weights of the transformer satisfies
5

---

Published as a conference paper at ICLR 2025
Positive alignment with the signed classification signals w (Ta) w (T&) ( (-1))=((1)
(1) and
Uk
Positive alignment with the group signals w
(Ta) Ck = (1)-
Symmetry between heads and groups:
w+ (Ts)
2k
w (Ta]T
(-Ux);
w (Ta)
11
w+ (T&) (Ta) (T (Ta)
02, W+ Ck ===== Ck2 W'+ c
w
(Ta)
Cz
In particular, the output of the softmax attention for sample X1,+.- satisfiex the following properties:
. Connenttattion on the group signal key:
oat (X1,+.-) //- (-)Th)
(+,7s)
vt (x1, 0(1)
(+)Ts)
(X (X1,,,-) pg p 12, (X), (X1,-,-) O(1)
TT)
q u2.t
P (X1,ti)) Pot-c ken (x),
(+.Ts)
1s)
Near symmetry within a single attention head:
Po P (X (X1,t--) -P---99- kt (X1,t,-) S SO(w)
(+TTg)
(+. Tz)
Pa--12kk- (X1,t.-) - Pq ---- k (X1,ti-)
(+.Tg)
(t.Ta)
$ sŌ(w)
(+.Ta) [ Po (X1,t.-)
Pytes kk -t (x1,t,-) $ Ō(w).
We now describe how the transformer can correctly classify the sample X1,+ Based on Theo-
rem 4.1 the output of the transformer
can be analyzed as follows: by the attention's
concentration the group signal key, we have
(Ta)T (+.Ts) w= kgt (X1,t.-) + Pgt kt, (X1,-.-) + Pgttakkkee (X1,t.-))
(t)
(+.Ta)
(-TTs) (Ts) (oo-o1,ke-a (X),+.-) Pae- (X1,t,-) + p (X (X1,t.))) 2 9(1).
(sTTs)
On the other hand, due to attention's near symmetry property, regardless of queries, we have w T
(Ts) (-02) approximately cancelling each other:
U1 and w.
w (T)) (-wz)
(+.Ts) P Pgen (X1,t.-) -o ke- (X1,t--) + P P ph va (X1,t.-))
(+.Ts)
+ ww
(T&) 0i P Poe (Xi,+,-) + Pqt- (X1,t,-) + Pge kt (X1,t.-))
(+.73)
Therefore, the positive head can output value larger than the negative head and thus the sample can be correctly classified. Our results on the behavior of the softmax attention in Theorem 4.1 agrees with our experiment results in Figure Ic By symmetry, the trained transformer will behave similarly on other samples with conflicting signs.
We now introduce another interpretation of Theorem 4.1 that relates to logistic linear regression. Notice that the transformer architecture in Equation
can be rewritten as
a (X) + 6) El Xsoftmax (X K+ Wo+
f(x) [WI w II
(o(H +) L
Xsoftmax (TT K WQ-
6(H bib 5(H (X) + byb
F(X)
We define the vector F(X) E 2g to be the attention feature of input X. Thus, if the transformer model can successfully learn the data distribution D. we must have the set of attention features (F(X)) linearly separable. We show in Lemma 5.4 the attention features under uniform attention is not linearly separable and we describe how the attention features evolve and become well-separater after training in Section 5.2
Finally, for samples with consistent signs, although their attention features also change after stage 2, are able show that those changes do not affect the separability of their attention features.

---

Published as a conference paper at ICLR 2025
5 PROOF OUTLINE
In order to analyze the training dynamics of the softmax attention, first of all, we need to keep track of the dynamics of the score variables ji" W Wov for μr, E CCV. V. However, upon writing down the gradient flow updates for J W Woz, it can be seen that their updates are governed by the variables in the form of J W W and J W Wou which we call key and query self-score variables. By keeping track of the gradient flow updates for those three types of variables all at once, we now have a complete dynamical system where all the terms appearing in the updates are tracked. The gradient flow dynamical system can be found in Appendix B.2
Before we analyze the training dynamics, we observe one complication brought by random ini- tialization: from our data distribution in Definition 2.1 the two mixtures are symmetric to each other and thus, ideally, the transformer's behavior should be the same for each class, e.g., W
4 WOT wr(t)
K+ Q+ Ui = 2 W K Q++
K+ @ 01 of U2. However, due to random initialization, c. (0) W
K+ w U2. In order to solve the above complication, we consider the infinite-width trans- former where the embedding dimension m of the softmax attention goes to infinity. Within this model, we can derive nice symmetry properties between two groups and between two heads (Ap- pendix that are convenient for our analysis.
22 1
Definition 5.1 (Infinite-width transformer). The infinite-width transformer is defined the archi- tecture in Equation
I with m + oc.
Our experiment results show that as long as the initialization scale w is small enough, and the embedding width is large enough, the inftnite-width transformer can indeed serve as an accurate model reflecting the true dynamics. For clarity, we present all of our results and analysis under the wndth transformer model. The finite-width transformer is discretized version of the wndthite-width transformer. Thus, the actual gradient flow dynamics can be seen as the infinite-width dynamics with perturbation introduced by finite-width discretization. We discuss how to take the discretization into account in Appendix G
5.1 STAGE 1: NEURONS ALIGN WITH THE SIGNALS
In stage of Algorithm T we train the neuron weights and keep the attention weights fixed their random initialization. Since the neuron weights are zero-initialized, the outputs of the two attention heads are both zero initialization. Recall that we fix the bias term b in our transformer (Fqua- tion (1)) to be some small positive constant. Thus, there exists period of time from initialization that the outputs of ReLU will be positive for all samples. We first characterize how the neuron weights align with the signals in such case.
Lemma 5.2 (Neuron alignment, stage 1.1). Under Definition
5.1 for
[0,Ti), if
H (X), (X),H( dw d T
for all X E DD for l E {1,2).
dw (T)
Up 1
Σ
(0 (Xx...)))
Uk 1
000) (X....))) [E(
Σ
dt 8
8
At the same time, for k E (1,2).
dw T
du " (t)
+ 0, dt
O.
dt
From the above lemma, we can see that when the output of ReLUs are positive for all samples, the neuron weights will align with the signed classification signals while not aligning with the group signals at all. A direct consequence is that the attention head output H1 (X,,.-) will keep de- creasing until it becomes smaller than -b, as long as b is not too large. Then the neuron weight w++ is no longer able to receive gradient from sample X1.-.-. We can show that when this happens, the neurons start to positively align with the group signals.
Lemma 5.3 (Neuron alignment with group signals, stage 1.2). Under Definition 5.1 for t [0, Ti],
if H (1) (Xx..-) < -br k E (1, then
du (0) Ch
1 g9
(Xx.+.+)))))
dl
7

---

Published as a conference paper at ICLR 2025
At the same time, the neurons still keep aligning with the signed classification signals. Thus, if T = (1), then we have " 2k = (1) and w (0) C (1) after stage
Although the neuron weights can learn the signals in stage 1, it is not hard to show that the trans-
former will output zero for samples with conflicting signs like XI,+,- and thus those samples are not correctly classified. In fact, we can show that under the averaged initial value model, our data distribution D is not learnable by merely training the neuron weights.
Lemma 5.4 Non--epaaabblliy of uniform attention feature). There exist no neuron weights wi such that yf (X) > 0 for all X E [Xe).e-e)) when the softmax attention is uniform.
Recall our definition of attention feature in Equation 2 Lemmaa 5.4 can be interpreted as that the
uniform attention feature is not linearly separable. Due randomness, the actual attention feature at initialization is not exactly uniform. However, by simple continuity argument, Lemma 5.4 can be translated as the maximum margin of the initial attention feature is at most O /vm)
5.2 STAGE 2: ATTENTION LEARNS FROM THE SAMPLES WITH CONFLICTING SIGNS
Since by Lemma 5.4 the samples with conflicting signs are hard to learn by merely training the neuron weights, in stage 2, we directly train the softmax attention weights WK+, Wo+ on those hard samples. As we mentioned the beginning of Section [ in order to keep track how the softmax attention changes during training, need to analyze dynamical system with all the score variables J w
K+ We+" V, key self-score variables μa
w y and query self-score variables "w(oT K+
Q v. From the gradient flow update in Appendix B.2 observe that we can write the dynamical system the following form:
j w W
wi(t)
Q+
vec(f W ( w-(t) K+ a+
vec(fu TW will) (iweeuv)) Kt Q4
d
vec(f w(()
w / m.vECUv) R+
-(00
vec(fn W wilh) Kt K v) (/
dt vec(fn
K+
T w Q+
(
vec((u T w (0]
W@+ w/O v) }INECUv)
where (1)) a square matrix with dimension Hs w-(e) W(2 1WTTTWTe K+
TWOT
T K+ w 22.l
Hu T with wi( "mall 16 + 10 + 10 36. Thus, directly analyzing the dynamical system can Q+ easily become intractable. We need come up with ways to utilize the structure of the problem to reduce the dimension we need to analyze.
One observation from Figure Ic is that the score variables cf W ( W en and cf W w oz exhibit the maximum magnitude change compared with all other score variables. In fact, these two quantities play a key role in making H (X1,t.)) > H- -x), Thus, we are able to prove that cf W K+ w 01 will become sufficiently large at the end of stage 2, then we can make good progress proving that the attention features become separable.
(
To make our presentation simpler, for a moment, assume we have the near symmetry properties:
cf T w(t)
K+ 4+ U1 w( (0O
K+ @+ 0/
OT w
K4 -c (())
K+C²
1 W ((t))
Q+ Q441 U1 RJ -u w
( (()
2/ U2
which can indeed be made rigorous in Lemma
E.4 Applying this near symmetry property to the
gradient of eJ ((TT W at l = Ti (see Appendix E.5.
we have
K+
de ((t (( K+ Q+
U1
3
gg
1 (T)4
11
W(T)
(WWT) Wet v1)
(2 (T)T
C1
dt
c=
(c (TwW
(T) K+ ())
(2 (11)
Recall that the maximum margin of set S is defined y max w|2=1 minges w
8

---

Q+ Published as a conference paper at ICLR 2025
Due symmetry and cancellation at initialization, it only depends 2 self-score variables: c W w
W++0 and Q++ 01 at = = Ti also depends only on of W w U1 Thus, this simpler depen- system analyze. Unfortunately, this is not true. Take the gradient of C W ( W U1 as an example. Training the K+
21. At the same time, the gradient of ej w w K+ K-c C1 and on @+ @++
" W w
K+ W
dence can hold through the rest of stage 2, then we get a much smaller
attention weights can break the symmetry that holds at Tı, and the gradient of cf T ( start to depend on other self-score variables. The way to overcome this issue to realize that as long as the maximum change in the output of the softmax not large, we can bound the effects from other self-score variables and the gradient of ef W w v1 is still mainly contributed by K+ q+
K+ Q+ v1 will
o W (OT WOT K+
w e1 and u ( w Q+ V1.
Q+
Now, we have successfully reduce the dimension of the system that we need to analyze to 3 x 3. However, there is still problem need to handle: the system has matrix with non-linear depen- dence on those 3 variables. Generally, it impossible to derive a close-form solution. Fortunately, can show that (1) those variables are increasing throughout the training and (2) when the change of the softmax attention (from initialization) is within a certain range, the growth of those 3 variables can be lower bounded by linear dynamical system.
Lemma 5.5 (Informal version of Lemma B2
Fort 2 Ti, when all the score variables bounded such that
within some range, there exists a constant
W
T K+
With t1
W
w(t) K+
(TT
O
17
dt
C
T
K+
w
K+ C
2
b
1
C
W
K
W
K+q Q+
" Q+
w () w:
Q+
0
0
W
Q+
W
With Lemma 5.5 we can prove that the score of W w Un will become sufficiently large at the end of stage 2.
K+
On the other hand, when the query token is the group signal, we can prove
of w wilt) c1 < O, 2y W wilh) > 0. Q+
By symmetry, this implies Pate Atce (X1,t.) (-+)
(+.0)
ket (X1,+.)). This finishes the proof of
the trained attention's property on the concentration on the group signal token in Theorem 4.1
Finally, when U1 and U2 are the keys, we can prove the near symmetry property between scores:
WTT K+
Wo+ ~ -wJOOOT Wo+
0
K+
(T) K4
W
Wo+up U2 = - woT Wo+
R+
w (TT
Wo+01 == -2 -WT
Wo+01
K+
K+
which can be translated to the near symmetry property within a single head in Theorem 4.1
5.3 STAGE CONVERGENCE
In stage 3, we train the neuron weights again. the end of stage 2, the attention features for samples with conflicting signs becomes well-separated. addition, we can show that the attention features for the samples with consistent signs are still well-separated. Thus, the training in stage 3 behaves like logistic linear regression on linearly separable data and training the neuron weight can drive the population loss towards zero. This proves the part in Theorem 4. that the population loss can become sufficiently small after training.
6 GENERAL NUMBER OF MIXTURES
In this section, we study the case when the mixtures have more than two groups Our experiment results can be found in Appendix A.4. We summarize experimental findings below.

---

Published as a conference paper at ICLR 2025
Observation 1. The transformer architecture in Equation 1 learn the data distribution for K 2 3. However, the near symmetry property in the attention no longer holds. addition, the training takes longer with more groups.
On the other hand, we observe that for K 3, the training dynamics is still stage- stage-wise and there are score variables exhibiting major growth compared with others. Thus, the techniques developed in our work can be used study the training dynamics for K 3.
Although the training dynamics for K 2 3 has some differences from the K = case, are able to provide a construction of transformers that can solve any arbitrary K mixtures.
Lemma 6.1. There exists a transformer architecture solving K-mixture of linear classification.
Proof. The construction given as follows. Let
K
K + C.
K Σ
W Wo+ =
Σ
UkCR
Σ
C&Ck +
k=
k=l
km
k
K
WR_ W -20
Ck
C.
CHCR
+
U/4Uk
k=1
k=1
k=1
The transformer is given as
K
S(X) =
2 Xssotmax(X" W Wo+ X1,) +
ug X softmax WWW WV--xi)
k=1
Let C co. First of all, consider classification signal query Uh. Then,
K
Uk Xsoftmax(X W WQ+ Uk)
oy Xsofmaxx
W W Wk+
22 Uk
Ug
Uk
=
Similar analysis for query -2k- On the other hand, for group signal query cg. suppose the corre- sponding classification signal in X is UR. we have
kk softmax W W@+c))+ +
oy Xsoftmax(A (X WTT Wo+cx)
Ug Uk
2 Ck
Similar analysis if the corresponding classification signal is Ukk Thus, we have yf()) > 0 for all (X.y) E D.
7 DISCUSSION
In this work, we studied the training dynamics of a two-headed transformer solving two-mixture of linear classification. We characterized the feature learning process of how the softmax attention utilizes the relationship between the class signals and their corresponding classification signals in this task. One open problem is that the current proof techniques critically utilize the fact that both the transformer architecture and the data distribution are symmetric. Can we discover some hidden structures of the problem such that can analyze the training dynamics for K > 3? More interest- ingly, we demonstrated that during training, the attention features evolved from a near non-separable state well-separated state, i.e, the maximum margin of the attention features increases during training. What properties does the data distribution need to have such that the maximum margin can increase in training? We leave those interesting questions future works.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
HY would like to thank Alex Damian for extensive discussion and suggestions. The work of ZW was in part supported by NSF Scale-MoDL grant (award number: 2133861) and the CAREER Award (award number: 2145346). JDL acknowledges support of Open Philanthropy, NSF IIS 2107304, NSF CCF 2212262, NSF CAREER Award 2144994, and NSF CCF 2019844. The work of YL was supported in part by the U.S. National Science Foundation with the grants ECCS-21 211660 and DMS-2134145.
REFERENCES
AlphaProof and teams AlphaGeometry Ai achieves silver-medal standard solving international
mathematical olympiad problems. 2024. URL httppttts::///eem.nnnd http:///eepppiiooooooooooee//a/a www.gooooolggggooooo cisover//icccvr
blog/ai-solves imo -app-iomooproooo wwwpppollsaaaaaeeeep www-sat----e ver meter--edal---
Animashree Anandkumar, Daniel Hsu, and Sham M Kakade. A method moments for mixture models and hidden markov models. In Conference on learning theory, pp. 33-1.JMLA Workshop and Conference Proceedings, 2012.
Animashree Anandkumar, Rong Ge, Daniel J Hsu, Sham M Kakade, Matus Telgarsky, al. Tensor decompositions for learning latent variable models. J. Mach. Learn. Res., 15(1):2773-288 2014.
Anthropic. The claude 3 model family: Opus, sonnet, haiku. 2024.
Sivaraman Balakrishnan, Martin J Wainwright, and Bin Yu. Statistical guarantees for the em algo- rithm: From population to sample-based analysis. 2017.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712 2023.
Arun Tejasvi Chaganty and Percy Liang, Spectral experts for estimating mixtures of linear regres- sions. International Conference on Machine Learning, pp. (040-1048. PMLR, 2013.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. In ICML 2024 Workshop on Theoretical Foundations of Foundation Models.
Siyu Chen, Heejune Sheen, Tianhao Wang. and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442, 2024.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-experts layer in deep learning. Advances in neural information processing systems, 35:23049-23062. 2022.
Richard D De Veaux. Mixtures of linear regressions. Computational Statistics & Data Analysis, 8 (3):227-245, 1989.
Puncesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention. Transactions Machine Learning Research.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1 04805, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition at scale. International Confer- ence on Learning Representations, 2020.
Benjamin L Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution statistical induction heads: In-context learning markov chains. arXiv preprint arXiv:2402.1 1004, 2024.
11

---

Published as a conference paper at ICLR 2025
Cheng Gao, Yuan Cao, Zihao Li, Yihan He, Mengdi Wang, Han Liu, Jason M Klusowski, and Jianqing Fan. Global convergence in training large-scale transformers.
Team Gemini. Gemini: family of highly capable multimodal models. arXiv:2312.11805, 2023.
arXiv preprint
Daniel J Hsu, Sham M Kakade, and Percy Liang. Identifiability and unmixing of latent parse trees. Advances in neural information processing systems, 25, 2012.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn feature-position correlations in masked image modeling. arXiv preprint arXiv:2403.02233. 2024.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structur. Advances in Neural Information Processing Systems, 35:37822-37836, 2022.
Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm.
Neural computation, 6(2):181-214, 1994.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-558) 2021.
Abbas Khalili and Jiahua Chen. Variable selection finite mixture of regression models. Journal of the american Statistical association, 102(479):1025-103 2007.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01255, 2024.
Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning. generalization. and sample complexity. In The Eleventh Interna- fional Conference on Learning Representations, 2022.
Hongkang Li, Meng Wang, Tengfei Ma. Sijia Liu, Zaixi Zhang, and Pin-Yu Chen. What improves the generalization of graph transformer? theoretical dive into self-attention and positional en- coding. NeurlPS 2023 Workshop: New Frontiers in Graph Learning, 2023, 2023a.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yo Chen. How do nonlinear trans- formers acquire generalizatioonguurateeed cot ability? In ICML 2024 Workshop on Theoretical Foundations of Foundation Models, 2024a. URL htw.s:ttssscc..... Ttt.p:://ppezeeiiie net Forum?? iid=
HizDrDvxLP
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear trans- formers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402. 15607, 2024b.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729. PMLR, 2023b.
Zihao Li, Yuan Cao. Cheng Gao, Yihan He, Han Liu, Jason M Klusowski, Jianqing Fan, and Mengdi Wang. One-layer transformer provably learns one-nearest neighbor in context. 2024c.
Perey Liang, Alexandre Bouchard-Côté, Dan Klein, and Ben Taskar. An end-to-end diseriminative approach to machine translation. In Proceedings of the 21st International Conference on Compu- tational Linguistics and 44th Annual Meeting of the Association for Computational Linguisties, pp. 761-768, 2006.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IFFF/CVF international conference computer vision, pp. 10012-10022, 2021.
12

---

Published as a conference paper at ICLR 2025
Ashok Vardhan Makkuva, Marco Bondaschi, Chanakya Ekbote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local to global: Learning dynamics and effect of initialization for transformers. arXiv preprint arXiv:2406. 03072, 2024.
James Manyika and Sissie Hsiao. An overview of bard: an early experiment with generative ai. Al. Google Static Documents, 2, 2023.
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 14735, 2024.
OpenAl. Gpt-4 technical report, 2023.
Ariadna Quattoni, Michael Collins, and Trevor Darrell. Conditional random fields for object recog- nition. Advances neural information processing systems, 17, 2004.
Terence Tao. Nonlinear dispersive equations: local and global analysis. Number 106. American Mathematical Soc., 2006.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023a.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Thirth-sevnnth Conference on Neural Information Processing Systems, 2023b.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du Sean and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305. /6380, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. /3971, 2023.
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738. 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- lion processing systems, 30, 2017.
Zhaoran Wang, Quanquan Gu, Yang Ning, and Han Liu. High dimensional algorithm: Statistical optimization and asymptotic normality. Advances in neural information processing systems, 28, 2015.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In Forty-first International Conference on Machine Learning.
Yongtao Wu, Fanghui Liu, Grigorios Chrysos, and Volkan Cevher. On the convergence of encoder- only shallow transformers. In Thirth- Conference Neural Information Processing Sys- tems, 2023.
Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. In ICML 2024 Workshop on Theoretical Foun- dations of Foundation Models.
Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi. Alternating minimization for mixed linear regression. In International Conference on Machine Learning, pp. 613-621. PMLR, 2014.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55 2024.
13

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
Related Work
2
Settings
2.1
Data distribution
2.2
Transformer architecture and loss function
3
Guiding Experiments and Algorithm
3.1
Experiment results
3.2
Three-stage Training Algorithm
Main Results
5
Proof Outline
5.1
Stage I: Neurons align with the signals
5.2
Stage 2: Attention learns from the samples with conflicting signs
5.3
Stage 3: Convergence
General Number of Mixtures
Discussion
10
Additional Experiment Results
16
A.1
Softmax attention
16
A.2
Only train the attention
16
A.3 Full attention scores of two mixtures
17
A.4
Experiments on Multi-mixture Models
17
A.5
Real World Experiment
18
Gradient Flow
19
B.
Gradients of each weight matrices
19
B.2
Gradient Flow Dynamical System
20
Symmetry
20
Stage 1: Training neurons to learn the signals
23
D.I Stage 1.1: Growth of neuron alignment with classification signals
23
D.2 Stage 1.2: Growth of neuron alignment with group signals
24
D.3 Properties the end of stage
26
Stage 2: Self-attention learns from the conflicting (hard) samples
26
14

---

Published as a conference paper at ICLR 2025
Growth of scores
26
Growth asymmetry
30
E.3
When the query is the group signal
32
Properties at the end of stage 2
34
E.S
The updates of the softmax attention at / 7i
35
F.5.1
The scores
35
E.5.2
The key self-scores
37
E.5.3
The query scores
39
E,6
Sum of partial gradients
41
E.7
Expanding the updates
42
E.7.1
The key self-score
42
E.7.2
The query self-score
45
E.7.3
The scores
48
Stage
53
Handling the finite-width discretization
56
Probability
58
15

---

Published as a conference paper at ICLR 2025
A ADDITIONAL EXPERIMENT RESULTS
A.1 SOFTMAX ATTENTION
We provide the softmax probability of the positive head on different samples.
0.7
07
0.6
05
05
04
04
03
22
-2
07
01
01
(a) X1,++
(b) X1 X
0.40
06
0.5
0.4
0.5
03
-22
02
V2-
(C) X
(d) X1,-i
Figure 2: Softmax attention the positive head on different samples. The vertical axis denotes the query tokens and the horizontal axis denotes the key tokens.
A.2 ONLY TRAIN THE ATTENTION
Recall that in Section 3 we observed there learning order when all the weights are trained simultaneously: the attention weights learn after the neuron weights. To further study whether this learning order is necessary, we conduct another experiment where we initialize the neurons as Gaussian and only train the attention. Our result suggests that in this case, the transformer is longer trainable.
Training
Only training
0.700
0.094
0.602
2000
4000
6020
8000
10000
12000
14030
(a) Loss
Figure 3: The loss for training only the attention while keeping the neurons fixed at initialization.
16

---

Published as a conference paper at ICLR 2025
A.3 FULL ATTENTION SCORES OF TWO MIXTURES
4dkkd
k1ke
92kd
k22k
422kvt
-1
-3
492.es2
2000
4000
5900
3000
10000
12000
14002
(a) Score
Figure 4: Full attention scores for K = 2.
A.4 EXPERIMENTS MULTI-MIXTURE MODELS
In this section, provide experiment results synthetic data for multi-mixture models. Specifi- cally, Figure 3] Figure and Figure 7 respectively indicate that the training dynamics of 3-mixture, 4-mixture and 8-mixture models are all showing similarities that for 2-mixture model, demonstrating the broad applicability of our developed theoretical results.
mixtures
alignment signals
0.7
1.0
0.6
2.5
2.5
0.3
1.5
1.0
0.5
0.0
0.0
2000
4000
6000
8000
10000
:2000
30000
0
2000
4000
5000
0000
10000
11000
14000
(a) Loss
(b) Neuron alignment
Scores mixtures
scotes
gvkkcl
qv1
k22
2:2
-2
-2
-2
2020
4000
6000
8000
10000
120000
16000
5
7000
6000
1050
5000
10600
17600
14000
(c) Scores with width 65
Scores with width 260
Figure 5: The experiment results for the case with K = 3 mixtures.
17

---

Published as a conference paper at ICLR 2025
Loss for mixtures
0.7
222k
212k
0.4
6.2
0.1
-7
10000
(a) Loss
(b) Scores
Figure 6: The experiment results for the case with mixturrs.
Loss mixtures
lons
0.6
0.5
0.2
0.1
0.0
2000
4000
6000
8000
10000
:2000
24000
0
2000
4000
5000
5000
10000
11000
14000
(a) Loss
(b) Scores
Figure 7: The experiment results for the case with K 8 mixtures.
A.5 REAL WORLD EXPERIMENT
In this section, present our real world experiment results. In experiments, dataset and extract the images with label 1 and label 2 to play the role of classification signals. We create random vectors play the role of group signals. We plot the training loss curves for both simultaneously training all the weights and separately training the neuron weights and attention weights in Figure In separate weight training, the three stages
MNIST
100 steps each.
18

---

Published as a conference paper at ICLR 2025
Training Loss
0.7
0.7
0.6
0.5
0.4
0.4
6.2
0.2
0.1
0.2
0.0
0.0
1000
200
(a) Loss for simultaneous training
(b) Loss for separate weight training
Figure 8: The experiment results for simultaneous weight training and separate weight training.
B GRADIENT FLOW
B.1 GRADIENTS OF EACH WEIGHT MATRICES
dw
E
(1)
Xp
(+))
dt
1=1
dw (2)
E
gl)) (X)) jyo
(1)
Xpi
(+.0)
dt
dw(())
(t)
Opih
(+.t)
K+
E
( (X)()) Nyo
w X
W=1 0syt1 941 *1
dt
1=1 h=1
E
(x) (1)
w Xn.
Pi.h
(+.1) (I(h- h²)
Pi,k' 1941 TM
(+.1) (t)
E
X 1o
w
Xpi
(+.0)
(Xp) )
T + 44/ w X ddiagr) -100 (x
E
(X)
9+1
(t) 1p
w. (oiigg(((-))) X
E
(X) y0
4tiPi
(1) (+.0)
diag [ we T._xpp( (p((tt)) X
(w
Q+
E
X
Pith
(++1) (T(h = h1)) - Pi,n'' )k
(+,
39 ( (T)
dt
h=1
h=
E
(X)
(1)
w
T (+.0) [ ( k diag(r) (1))x T "w+) 1i
E
(X)
(1)
(-w. Xp) (+()) (ttp)t)p(())))) (+1) + diag(p) (+++)))xx) (+t)) w+++ xi
(X)yo (1)
K x! diag (x
(xTw_ Tw - w. (xxttt)) pf (ttt)
E
19

---

Published as a conference paper at ICLR 2025
B.2 GRADIENT FLOW DYNAMICAL SYSTEM
Let n, & VUC. We have
dw (E)T 1
(0) (X)yo (2)
Σ
I
Xpi
)
(1))))
11
dt
dv W wh(t)
K+
Q+H
( (X)vo (1) E
1 T (0) (0) (+.0T
a+ Qqili
diag
w+
(0)
X w Xpi spftt()) XT
(1)
dl
+ E
(0 (X)oo
v
(+) K+
Kt (1) diag X [Tw++ u4 w x xxp(+(0)) pi p/tt002zT P
1=1
T ((t (
dv
K4
R+M
E
gith (X) (t)
wllh ( (1) (+,0)
K+ l+iPi
diag (wIX w w_ ((())))) X
+ E
(00 (X)OO (1)
Σ
p T (++)) diag I x - w (oo[xpttt) xT
K+
dv To w
E @
a+
(X)yo (1) Σ
-.--w[[ppj T w _ diag (x=== w Xp, xppp((1))) pd 1(t100
# T wiih xc_diag (x w. xpp xpp++))) ((tt) (x
+ E
g 10 (X)yo
Q+
1=1
C SYMMETRY
In this section, derive symmetry results for the infinite-wide transformer model. We first char- acterize the initial values of the softmax attention weights:
Lemma C.1.
The initial values of the attention weight in the infinite-width transformer satisfy:
V
( K=
Q+#
Ewx.Wo
W
K
WQu
W
= 0
v 1 W0
(0)
Exx
v T W
KM
# μ = = ==
0,
K+ K+#
Bwx
" W W KK
Bia
V WT W WOH
= 0, V μ
T w (O) w0)
Q= Q=M=
Ewa
1 NWWO W WOH
= w2 ==
where WK, WQ E R and (Wk)ej, (WQ)ss (o(0,w21) ).
Proof. Let (WK+)i. (Wo+); denote the i-th row of WK+, WQ-, respectively. Then, Notice that (Wil)) K WeLl has the same distribution as 1 Em abb where ai by w2). By the Law of Large Number,
we have v O11 Q W2), (W()) a ~ 21). Thus, v (
I T (w) T p.
(0)
itd.
lim O W m0000
lim Σ a.b = [|abb]] = E[a = O.
K= Q+/4=
m=+00
The rest can be proved similarly with an additional fact that for w ~ NO,, TT is independent of we v if ps I.
20

---

Published as a conference paper at ICLR 2025
The symmetry results contain two parts: symmetry between two groups and symmetry between two attention heads.
Lemma C.2 (Symmetry between two groups), Under the infinite-width transformer model Def- inition 5.7 let pi, & {c1,0) and μ12, 12 E {(2, V2 such that pi, P2 are the sume type signals (i.e., either P1, or Corpi 422 E V) and v1, V2 are the same type signals. For all t,
w [OT
(Q)
M1
=w+ P2
(3)
R+ 9+/1 = 1/2 WTTT 1:(
1/2 WT w
K+ 0+14 = = T (T (TTW((
w(t)
K+ Q+/2:
R- 0412
(4)
of W w
K+ R+/1 1/2 W w(O
K+ K+12,
"1 TWO W
K+ K+122 = W w
K+ K+M
(5)
in Twit w:ll)
Q+ Q+//1 = 12 w with)
" W wit)
Q+ Q+ = =v2 W w(t)
Qt Q+/2:
Q+ Q4111
(6)
The symmetry between the two groups also holds for the negative head. Corollary C.3. Let 8,88 E {+13. For all t,
p (X1,0(11 1( (x(t:0((2..))) (X = - H (X1,s,') H (X2,,
H (X1,s,)) H (X2,,,0)
(X1ao))) 5 1 (X2,a.)))
all the equalities hold at t = 0. Thus, we only need to show the gradient flow updates satisfy all the equalities if all the equalities hold at time t.
Proof of Lemma C.2
First of all, by Definition 5.1
= p(+,))((2.s.d)) (X2.s.²) (recall that since our model doesn't 't have any positional encoding, we can consider ordered inputs in our analysis), which further implies that H
Assume all the equalities hold at time 1. This implies that p(+.0) (X1.s.)) (itt.((X....))
(2) (X.,.,")) = H (X2e,)), H (X1,,,)) = H (X2,,))
and f(() (X1,,,s)) (X2,,)) for all s, E {=1}.
First, we prove Equation
dw (t)
I
(xp (10))))) T
11
E
g
yo
yo
11
dt
1=1
I
(X1 (er (X1 (1,0)))) E (X (x) pi
Eg,s'
2
+ 59
1
(2) (X2,, (4) (X2,.
X2 (22x0))))(((20s)) (X2ssr) T
11
= Eg,o'
g/ (X2,,or)) (t) (X
X2 (tid)) (X2
U22
1=1
1
X1 Xx...orp/tin (t.1) (XI. T
+ 29
(0
(X1,,,,)) ( (X1,
222
(by ((t,)( = (X2,.))
= E
g
L yo!
+
(xp(th
((p(tte))
T
12
dw (O)
dt
42
Next, we prove the first equality in Equation (4)
du
W T
K 11
(((t)
dt
L
Es,s'
g
(t)
so
(t)T (t)T (t) (+,)
Q+ S+iPi
diag (w X1,0,'
(T)
w X1 ((.0))
(t)
2
1=1
21

---

Published as a conference paper at ICLR 2025
+
2
B
g
80+
+
M T (+ diag (oT X2,000 w x2sosepftit) Xahaore
Q+
9tili
1=1
1
B 8,6'
g(0) 00)0
80
(t)
T (TT K diag [T - w. (1. (((t)))) pl pft: (XLLdt)) un 1 K+
2
W x diag (2 (2 - wI N2..2anopfttt vfft)t pf (ft10 (X2,0)))) K+
+
9
80
V
2
1=1
4
1
Es,s'
g (1)
H2
W (8) (+,4)T
Q+ 4+1Pi
(2.ssofftr)) Xgleoier diag (w (1) X2,0,0 Baa wwt (1)
2
1m1
B
g(t) (0)
M2 T wil (1) ofippt
9tepi
(1)T (+ diag wit X1,0,0 w (1) x(sispftit) X XTonares
2
+
1
Be,0'
9 ( 80
II W k K+
d diag (2 - w x2.2 ((t1)) pi (X2.ee)))))) u2
2
1=1
1
g00) sd
Σ 2/2
I W Ki2diag K+
diag [X (T - w. (Ix (100)) pl (X
2
1=1
du
W K+
(c)
W
The rest equalities in Equation 4 Equation (5) and Equation (6 can be proved similarly.
Symmetry between two heads: notice that for the negative head, we can view it as training the model -X() on the dataset where the label is yy and the set of classification signals {-11, )-
Lemma C.4 (Symmetry of two heads). Suppose Definition 5.7 holds. Let 11 E VUC. For all t,
(0) 11 = (1) M4 WT w(t)
111 (-1)
(7)
K+ Q-111 M1 pb2 W ((t))
Q-111 (11):1--1100ae (-1)S(uzE)) (-1))taE
(8)
P1 WTT (())
K+ K= = 112 W:!'T W(() Pz (-1))((1)))))))) (-1))(GEVV(11))(EE
(9)
ug W (0)
Q+ Q=111= 12 w w(() 11 (-1)
(-1)
(10)
Corollary C.5. Let s E {+1). For all t.
pl+ (Xk.s,')) = 1-00 (X -x,
H (Xkie,r) = 11 (1) (Kk,-
of (Xm..w)) -8 -8f (X (Xn,-s,-s')
Proof of Lemma C.4
First, all the equalities are satisfied when t = 0 by Definition 5.1 Assume all the equalities hold at time i and we are going to show that the gradients with respect to also satisfy the equalities. Equation (8) implies that (Xx,
pl pp) (xx._s,--s))) which further implies H4 (Xa...r) = ( (XR.-a,-r) and (XK...)s))) = - (10 (1000(((x0))) (Xn.a,-')).
Define vl = P (-1))(MEV))) for E VUC and X = Ic Uk] [Ck, Uk: Uke]]
Xk.-a--a..
We first prove Equation
dw P1 Ek..,''
(Q)
L (0))
(X ((ot)) 7
111
dt
+
1=1
22

---

Published as a conference paper at ICLR 2025
_w (f)T 14
X Nhioror/1100 7
glt) s0!!
14
dt
1=1
g( sct
X (k pp0-))
H4
which implies
dw _ (t)
141
dt
dt
We next prove Equation
8
du
(t)T
K+
w
@+
dt
-B
g(t)
b04 1=1
(0)
w (1) (+,9) diag (wt (0) KK.W - w (1) Xka pftt)t X 9tiPt
Ek gl0) go
TT [ diag [I ( - xka pi p(ti0) (XMraorr)
w QQ1"" 1 d(w)) ((o)TTkk_
(1)
K-
dt
g (1) 3o (0)
(n'TTWl (0)
( diag (w (1) XN..,e -w (t) (t.a Xhrsowr
9-1Pi
L
E
(w((() T k
(Xmrww) Tu diag (K - w I. p(-10) pi ((-))
K
g(t) 100
(0)
writ d (-,0) diag (1)
(O)T X (hooopp--t)) X
1=1
TW T k-th diag (X w w XK (11) Pi (X) Dit Σ
which implies
de T wilT w(t)) K+ Q+ d(v))wkk will
K-
dt
The rest of equalities Equation and Equation
10 can be proved similarly.
D STAGE I: TRAINING NEURONS TO LEARN THE SIGNALS
D. STAGE 1.1: GROWTH OF NEURON ALIGNMENT WITH CLASSIFICATION SIGNALS
Due to aero-initialization the neuron weights, both attention heads are outputting zero. Since we add a small positive bias in the ReLU, for a short amount of time from the random initialization, can ignore the ReL.U activation.
Lemma D.I (Neuron alignment with the classification signal, stage 1.1), Under Definition 5.1 for t [0, Tı), (1( (X), HT() (X) for all X E D, then for k (1,2).
dw T
(t)
Uk
Σ
g0t) (XM
dt 8
23

---

Published as a conference paper at ICLR 2025
dw T
Uk
gd (K
dt
Proof. Take k = 1. By the gradient flow update, have
dw dw [
21
E
9
yo
(1)
Xpl (xp(+10) (e(tt.))
11
dt
Xpp ppp(t)0)) T
E
1=1
Σ
g
(1)
(N1,s.a)
1,bt (XL.xr) )+
(+.t)
+
Σ
(x (X
Σ
pot (X1
(t.))
8
8=88
1=1
1=1
Σ (X2..
po (X2
(t.))
Σ (X2,,
(+.1) a (X
8
$=0
1=1
(X
(by Corollary C.3
8
Lemma D.2 (No neuron alignment of group signals, stage 1.1). Under Definition 5.1 long H (X), ( (x) > for all X, and for le E {1,2].
dw (t)T (T) + Ck 0,
!t Ck 0.
dt
dt
Proof. By the gradient flow update, we have
(1) dw
Ck
E
yo
(c)
Xpl
Ck
E
9
(1)
Xpi Spp(ti0)
Ck
1=1
L
Σ (0) [K
(t.t)
Σ (Xx,
(X
(++t)
P
8
1m1
0
(by and Corollary C.5
Corollary D.3 (Duration of stage 1.1), There exists a time To.5
@()) such that
(Tos)) (X (Xa,o,-o) fbror E (+1).
Proof. Lemma D.I and Lemma D.2 imply that (f0) (X) 2(1) for S O(1). Since is chosen to be a sufficiently small constant, there exists a time To.5 @ such that H! (To.s) (Xx, S -b for (#1).
D.2 STAGE 1.2: GROWTH OF NEURON ALIGNMENT WITH GROUP SIGNALS
In this sub-stage, going to show that (1) the neuron weight will keep positively aligning with the signed classification signal and (2) the neuron weights start to positively align with the group signal.
24

---

Published as a conference paper at ICLR 2025
Lemma D.4 (Growth of classification signals, stage 1.2). Under Definition 5.1 for [0,Ti]. if 14 11 (Xa...- < for all (2). then
dw (f)
2k
Σ WEP
(x
g00h
db
8
Similarly, if H1 H (Xk,+i+) < -b for all l E (2), then
dw (() Uk
1
Σ WE(22
(0
(Xx:--))
dt
8
Proof.
Take k = 1. If H (XM, <<>))) - for all k E [2], then
dw _w [t)
V1
(t) (1)
(xpl (1))0)) )
9
yo
11
dl
[X (X1
P (Xititt
(t+t)
+ 8
Σ
sg
(X1
Σ
Pottoo (X1,)
(+.t)
8
1=1
L
( (X2
qe (Xa.t.t)
(+.0)
1
Σ
(t) (x 22,0,'
P Poe- (X2
(+.1)
8
- 8
1
((0) (X..ttt))))
(by Corollary C.3
8
WE/2
Lemma D.5 (Growth of group signals, stage 1.2), For t E [0,Ti]. if (X < -b for l E (1,2), then
dw ( Ck
1
dt
Similarly, if (X.,t.t) < bor k E (1,2). then
dw (T) (t)
1
(Xe,,))
8
dt
Proof. If
( (Xt.,,)) < -b, then
dwo (t)
Ck
= E
9
yo
yo
(1)
Σ
Xpp (10tt)))
Ck
df
1
X
(+.0) l
(X.,t.r) - (00) (X
p let (Xx--+))
(+.t)
8
1=1
g 11) (Kr+
(by Corollary C.5
Similarly, if (1)) (X < -b for 1: E [2], then
dw (!T (1)
L
Ch = -B (0) yo
(t)
Σ
Xpp -:0) Ck
dt
1=1
25

---

Published as a conference paper at ICLR 2025
I Σ
8 - Σ (X
-
p 1=1
pe (Nx,-1o) + g (X
(-1)
Σ
Pelkken (Xxiti-))
1=1
(by Corollary C.5
D.3 PROPERTIES AT THE END OF STAGE
Corollary D.6 (Activation of different types of samples). Fort E [Tb.5, Ti] and {=1}, & e [K],
only d (x)) = and all the rest of cases equal 1.
Proof. First of all, by the neuron signal alignment lemma stage 1.2 (Lemma D.4 have
and and Lemma D.5
dw 4 (+) (vi + U2)
(w (t)T
C1
dt
dt
Since the attention is uniform during stage 1, by Corollary in stage 1.2.
D.3 we continue to have à (Xk--
Corollary D.7 (Neuron alignments at the end of stage 1). There exists a time Ti = (1)) such that
W. (T.)T ((1) for all μ E VCC.
Proof. This is direct consequence of Lemma D.4 and Lemma
D.5
As result of stage 1, we can show that the transformer model now can produce positive margin on the consistent samples but it fails to classify the conflicting samples.
Theorem D.8 (Stage 1). Under Definition
5.1 we have f(Ti) = s 22(1) and
((T)) (X..a.))) =.
Proof. Since the attention during stage uniform, this direct consequence Lemma D.I Lemma D.4 Lemma the symmetry between two heads in Lemma C.4
E STAGE 2: SELF-ATTENTION LEARNS FROM THE CONFLICTING (HARD) SAMPLES
E. GROWTH OF SCORES
The following result states that there exists a range where (1) the growth of (OT ((!)
c K+ Q+ and -c₁ ( (t))
K+ U dominates other score variables; (2) the growth of Q+ Q- uy dominates other query self-score variables; and (3) the K+ K+ ck. -ck W w(t) c dominates other key self-score variables.
T W(T wi(t)
Q+ Q+ Uk. - (C) w-tt)
growth of of ((T ((
K+
Lemma E.1. For (1) witll ( Wo C, then there exist constant 0 < such that
> Ti. there exists constant C > 0 such that as long as
maxp,renve 1u W
c. min KAKEE[K]
(oT (BT (1)
K+ Q+ Uk. - -cT
(T) w-10 UK K+
2 max
U (I.EOUV
{uuT w K+ w- w
U
fooTwhe w:11
((T K+
w
oN (11)
C. min k*kE[[]
is w
w-(l) Q+ Uk.
-
w w
Q+ UKr a+
26

---

Published as a conference paper at ICLR 2025
w |
2 max
U
{ wot w
@+ Q+ ct . W
wof c . - (t) W
(12)
a+
KARE
min kykee([]
ck we W
K++ Ck. -
W w
k+ WK+cH cx)
1 w w vn 11)
2 max
U
{| 1
w a 01 WK+ W Kktto
(13)
KAMEJK
and
(0) Q+ ux1
1 I cl
|
(14)
Proof. The proof is by inspecting the gradient flow updates in Appendix tion. At ( = -T, all the scores equal 0 and further, we have
E.5 and proving by induc-
deg W K+ Q+ tk
deg W ( K+ @
(D) w
dt
dt
Uk
3
(T))
(Ti)
2m
w
(T)T (Ti)
c1
w
Xpi (t:Ti))
t=
g
(Ti)
w a1
(Ti)
On the other hand, we have
(NU fau w 1 (
U (U KAMEE(K)
d arg max
{Ict w K+
cl
w ux
w (T) w |c]
dt I-
MNNECu
dup (OT
(OT K+ Q+0 c1
dt
and
duj (1) ( w-(t) K+ Q+
3 ( 2w (Ti) (t.t))
4
X
- ai
(i)
3
(T))
3 zut
(TiT
dl
T=
Thus, Equation is satisfied at t = = Ti.
Next, have [|W]T (Ti) 1x w² Uk W (1) W uk
o. U/ (T:)
W!TT Ck 0.
Q+
uk WT(
@+ wilh)
(Tn)) ex 0 and of T
w (2))
Q- cw 0. Further,
du w(t) (-(t)
Jup ((T w-(l)
Q+ @
dup ((( w((t)
Q+ We-CH CK'
det w-(t) w
Q+ W
Q+ Q+ U dt
=0
dt
dt
l=
l=
1=
1-T;
Hence, Equation
12 is satisfied at Tı.
Similarly, we have ||
(T)) K+
cxl1 w? of (W((T T (T)
K+ K+ ct. = O, oh (T,)
wT(T))
K+
O, U R+ K+ ov = 0 and uk K+ K+ = 0. In addition.
T (T)T (T))
(T) (T)
(T)
( (-(t) do K+ K4cK
dup w ( w-(t)
d w (t) K+ K+
W
du W (t)
w-(1) K+ K+UW UK!
K+ K+Ck
Ck
0
dt
dt
dt
df
t=
t=
t=
t=
Thus, Equation
13 is satisfied at t = Ti. Finally, Equation 14) is also satisfied at t = Ti.
Now, assume Equation (T]). Equation (12) and Equation (13 hold at some t Ti. We are going to show that the gradients at f also satisfy those inequalities. We are going to prove the inequality be- tween cg (t) ((()) 2g and v, w/O0T W c1. The proof for the rest of cases similar by checking the gradient flow dynamical system in Appendix E.5 as they all follow similar structures.
27

---

Published as a conference paper at ICLR 2025
First of all, notice that if we replace p(+ with p(+-25) in holds that there exists a constant 0 < d' 1 such that
dej W W
" th and
doj W w
de 44
dt
c' 1
3
w w
Q+ 01 u/ W ( (
w
( c1 w (T)
(OT
Xpi (((t)7))
a+ Q+122
+ 39
1
2
c
W
K+
K+G a W (T) (t)
K+ x++0
(w{)T C1 -wf ( Xpp (11)
(
w ( Q+
@+ C1 cf wit w 02)
(-w(fT ( Xpi (111
ci
3
1 + 1
T
W
(
W
K- 01 1j WK (
wi(l) K+
(-w Xpl+ T- s
3
1i
K+
02
since _ c TOT Xp|+
01 Next, we calculate = 3 w ( and w Xpl (+.Ti) = f (T)
(T
the deviation of replacing p (+1) with y (,.)) in the gradients
de_ W
(T w
K+ dt
4 v1
1.
E
yo
2 w(OT (0) (0) (+IIIT
Q+ 9iiPi q.PPi diag [w+ (B)
x - -
-w+ (()
Xp| (11))) x
1=1
(t)
diag c W ( h diag (XTw+ 14 - w xxp xttt) pft-ri)227 21
+ EE
yo
L
Σ
w (t) (pf (+ diag (wo x - -w+ Xpi xp(t)())
(()
(t)
yo
u
a+
9tl Pi
1=1
(T.Ii
(t)
x - (t) Xpi ( x C1
Pi
diag
w
I. (t)
+ E
g yo
Σ
ci
W x (diag (x w Xp! p(ttt)) pi (())))
K+
1=1
diag (XTw+
w+ - T x.__xp(t100 (T)i)) pl p(t:: xo 01
Thus, there exists a constant C > 0 not too large such that max IN TWTTT and thus [|p((t0)) - p(t,Tillloo = O11 and then
K4 W S C.
dcT w((t) K+
w-(1)
Q+ 21
E
n0 yo
(t)
e T wiOT (+.21)
94it) (40e/+ diag ( (1)
x - w xp(+ x XTa
(()
dt
+ E
9 yo (t)
cJ
Tw(T K+
xt diag (x w+ w. XpI (e_xxp(+TTg (it) pi
((t 2oT u1
S c"
3 01) 1/ w
( Wit - of W.TWOOT//// WorTTTW!ooo w((h13) mz)
(w (T) ci - -w+ ( Xp1 (1)
1 +
g
(2)
c
W T w(() K+
K ej
WT K+ K4 02)
w
(w a --- wit Xpi t1771))
for some sufficiently small constant d" 0. This implies that if satisfies d' + (1) < c < 1 then dof WLOT with) v1 do. w w(O
Based on Lemma E.I we can show that there exists a time range where the dynamical system exhibits exponential growth.
28

---

Published as a conference paper at ICLR 2025
Lemma E.2 (Score growth lemma). Fort 2 T1, there exists constant C > 0 such that as long
paxuee Ia T W
W
v S C, then there exists a constant b such that
Q4
ci W
T
W(t
Q+ 1n
T
17
Ci W
Q+ v1 W
W
d
dt
Ci W
T
W
K4
(c)
2 b
0
0
ci
W
K+
(
K+a
0
W
Q+
q+
Q+
Proof. First, as long as lpp lpp(+
(X)
(X) (1) for all X and I, we have gl)) 2 to the update of cj w/ w
+ T1)
Q(1). Now, applying Lemma
E.
K+ Q+ Un in Appendix E.5.1 we can
derive that there exists a constant by
such that
de W W (:(t) K+ Q+
K4 1 wil w v1 bi (cTWRT w
dt
and similarly, there exists b2, by such that
de W (t)
w
k+1
2 . T T
W
dt duf W ((( Q+ W
01 by cT k+ W
dt
Take 6 = min(b1, b2, by) we finish the proof.
Lemma E.3. There exists T2 such that T2 - Ti = Olog + and
cf TwWp
(TA)
W
(T2)
Q+ U1, -c] (T2)
w-(T))
022 = (1). Q+
K+
K+
Proof. We can define a process xtt) by
0
1 17
dz
b
x00)
2
Ax,
A
0 0
2
m
1
0
0
and by Lemma E.2
we have
c T. (T (T++))
w
x(1) S
c T.
(T++TT
W
(1)+)) ci
K+
K+
W +T+
(T)
Q+
Q+
01
Now, we analyze (t). We can compute the eigen-decomposition of A as
det (A AT) det. (
1 -
+2 = ==
1
-
0=
0
-
The eigenvectors can be found by finding the null space of A
XI:
A 1
0
I
i
x
x2
0
1
0
X
_33
Thus, the eigenvector for 0s
0
C
and the eigenvectors for the non-zero eigenvalues
X
C
1
29

---

Published as a conference paper at ICLR 2025
Therefore, plugging in the initial value z=(),, have
V2
-22
w² exp
s01)
(22b1)
1
exp (-V26l)
111
1
1
Hence, there exists T2- -i 2 (log # such that o W W
(Ta)
(1 01 2 x(I2 Ti) 2 Q(1).
E.2 GROWTH OF ASYMMETRY
The following result shows that there exists a near-symmetry between the score and self-score vari- ables.
Lemma E.4. For 2 Th. there exists constant C
0 such that as long as
maxpe Ia TWO K+ W ul S C, if we define
(0)
{ |cc ( W (01 + 12) 12)T W
K+ W/ipl. 1 ww K+ w (ar + to2) |a w kt w (or + 02)
max REVUE
1 (us u2) W w (a T a+ W (01 + 12) TW w }
(a +2) k+
then
(A
dt
S (w)
dec W
w (2)) m1
dt
and
(1) < Ō (w)
Proof. First of all, by the symmetry between two groups in Lemma C.2 we have
(ci + (2) WT
K+ Q+ mn = c. W w(()
K+ Q+ (m) + 02) K+ Q+ = cf WOOT w(())
(ci + 22) TW (t(
(ci + +2) K+ Q+
(N) + v12) T (w(tTT(()))
K+ Q+ "1 = in W w(t)
K+ Q- (m) + 02) (T(T
(r) + 02) Tw w(t))
Q4 CI = of K+ Q-
(er + 02)
Thus,
A (0) = max
(a + 02) ( K- w (m) + 02) T (T( w(0
K+ Q+/4
wevue
(c1 + o2)
w(o W K+
W. (mi + v2)
T K+ K+ (cr 02) w/OT WotH, (v) + u2) T
witt
At = Tı, we have (T) = w and by the updates in Appendix E.5
we have
dA (t)
0
dt
t=
Define Ap(0 = dax Pi d2
(X)
3 J. For 1. 2 Ti, from the gradients of the score
variables in Appendix H.55 and Appendix
E.6 we have
(1)
0 (
(0 Ap
(15)
dt
30

---

Published as a conference paper at ICLR 2025
This can be seen from the following. Take dej W1OT w an example: Define dt
dej Wx K+
a- 01
p(t.t)=p(t.m)
E
yo
U W (1) (0) (+.I6)
9+iPi
diag
(9T w X w+ Xpi (p)+t x Xk,
(t)
I. ci W
+
(4)
yo
K o diag x 04 w
Xpp (t))t pi
((t.Ti) x7 vs Xx.a,-s
Then, following similar steps in deriving Appendix E.5 it is not hard to derive that we replace p(t+t) by p(+.Ti) in de; w with m
have
de T W W
dt
p(+.t)=p(1.7)
= 3 000
(1(t)
Q+ - 2 w W W
2 (TuT a
Q+
+ 3 0
w + c W w
(Ti)T - -o 11
K+
kt xx
and thus,
de
deTwoo
@ (01 + U2)
K+
= 0.
dt
1p(1.)=pp 10000(111))))) 1.31)=
Further, by the chain rule, have
wit de (TT K+
(vi + v2)
dw-(t) Q- (e1 + U2)
de (()
Q+
oJ W(T
(0) + +2) w
di
K+
dl
dl
Q+
aw10 (wi+
Therefore, expanding
via Appendix
E.6 we have
dc WTT K+
@+ (vi + 02) =
de
(-(t) ( K+
Q+ (1) + v2)
de_
WTT
K+
Q
(vi + 02)
dt
dt
dt
=pt +))
I 0 5
Applying similar analysis to all other asymmetry variables, we get Equation (5)
Now, we analyze the growth of (t), Consider specific time 7" such that oj W (T") T
Q+ in w
K-
(w) for the first time and by Lemma E.1
we have A < (w)) for all < T". Further, by
Lemma E.2 and similar analysis in Lemma E.3
have T s ((lo(((/w) (1)w). Thus.
da (0)
A A
<
A
(3))
dt < w² + +(w3).
dt
Ti
This implies that
A
( (w))
< <((w)
(T eT
(T"
K4
Q+
U1
Finally, comparing the gradient of ( with the gradient of c. w-(T (
K+ W in for t E [T', [T] (fol-
lowing similar analysis in Lemma E.
we have
O
I (w))
de W (T W W K+ Q+
01
dt
31

---

Published as a conference paper at ICLR 2025
where the Ö is independent of t. Notice that this inequality the ratio of gradients also holds for < T since Ap(t) O(w) for t I T'. This implies that [A()) I A +
(w))e (CT) ( K+ Q+ 21 oj (T'T w K+ Q+ v1l S ((w).
E.3 WHEN THE QUERY IS THE GROUP SIGNAL
The main goal of this section is to analyze the group signal queries. To do this, we need to know the sign of the group signal queries. This can be analyzed by induction-like argument over all the scores and self-scores below.
Lemma E.5 (Sign of all the score and
variables). Fort E [Ti,T2).
a W w:(l)
aw w:(()
K+ Q+ v1 2 0,
K+ @+ 02 0
"n W (:(t)
in W w 02 22
K+ Q+ 00
K+
c w ( @+ 5
e K+ @+ 2 0,
K+
" w (oOTw())
"1 T w 22
K+ Q+ S 0,
K+
w w(t) 2 0,
ei W
a Q+
Q+ "2 S0
a Q+ witt) 2 0,
ci wiit w(t)) 02 50
Q+
in we w(t)
tn w ((t) (t)
Q+ Q+ 01 20,
Q+ Q+ 5
wit w a K+ K4 50
K- w K4 22
a K+ K+ 2 0,
QWOT 1(0
K+ 55
21 W w
w K+ K+ 02 0
W (1)
K+ K+ 01 22 0,
Proof. First of all, at f = Ti all the variables are zero and thus the results are satisfied. Now, assume all the inequalities hold for some Ti, we show that the gradients of those score variables also satisfy those inequalities. We first prove the inequalities in the first column. First of all, by Lemma we have
de,W (alli (L) wi K+ o+ dej ToT wlo du w w Q+ dt
(1) 01
K+
20.
dt
dt
W ((()) K+ @+ e < as an example and the signs of all the other score variables can
We now prove is be analyzed similarly. By Appendix E.7 have
dapp WIOT w(l) Q--1 e
E
g
yo
c Twi (8) (+,)T
941Pi
diag (0) x w
(1)T Xpl spp(t.l) X Xkiti-
dt
(
T wiit (0) (+.0)
diag (O) X u4 Xpj xpft.0)) x Xx-iit
(O)
yo
c
4ptPi
L
got
Σ
((T K+
KC diag (x W4 - w (.]xp((ttt) + (++)) zu Xaaee
yo
1=1
We first analyze the sign of the sum of the first two terms.
(4)
c T W (0) (+.0T
4tiPi (1 diag [ (O) x -w44 with Xpi spf(t00) x Xk.ti-
g
yo
+ E
g (1) yo
ci
Wrih (t) (++) diag (wf (t)T x - -wt (x(= x m XR-
a+ (i
1=1
())
W Q+
(p (t,0)
ko (-w(f) (0) (- Mi.tt,p(t.t)) P
Pa--oykktmy (-w (1) -pt
(t.t)
32

---

Published as a conference paper at ICLR 2025
(T) + + Q+ W ui ii (Potenktor (X1,t,-) - Pg-- (e--o, (X2.t.))
((tt))
+ +a W w
qt @t oot,ktt (-rof ( (1.t.-pptc) (1t11) - oj W witl
a a+ c2l ge (-wr (TT X2,ti-D
w (-1)) (Po-- (-wof (0) X X1, (1tti-po--1)) p ) + P Pgmegat-- (-wl+ ( (2ittipp/t:1
+ + @+ @+
+ + T wilh (t) (-ppt- (tit)
Q+ @ uy
U1
12t- (X1,t.-) + (+,))
k (X2.t.-)
+ +cTwel o+ @+ ciP k (XI,t,) -wot D1 + o wa+ w c2p cep -e (x V1
(+.0)
(I)
(t,1)
(t)
0 (WWS o+ w (-01) (gt-on
(t.0)
(-w4T (OT X1, (1,-, ) + P p9- (-1of (1) X2, (t. (((tt))))
(t,t)
+ + a+
W (e (-wn)w4 (1) (-11)
P (t- (x1,-++) (t.)
(t.1)
kit (X2 (X2)-it))
+ + W w cap @iptaaak
(+.t)
(1) X1, X1,- pp+ (1.1) - oT w/OT wiit)
Q+ Q+ of (-w: (-u(() o2p
(t.))
(t)T X2 X2
-w+
+ + ot 84 (p
will)
(22)
(-w
(1)T X1, - ppt ) + p (+1))
p2-- (-w(3T (t) X2, X2 ))
(t.0)
t w wilh
@+ U2 (-01)
(0)
1gth 11 (X (X1,- - (+.1)
(+.0)
1y--1g.kt (X2,
t TT (+.)
@+ ciP 8 (X (X1, (t)T (-11)
cJ wet w (t.))
(T(
ea (X2, (x2,-++ U1
(t)
Define
A = max
{|ce + [(c1+02)TW:/) ww
k+ w (0i + W k+ w(f
DEvUe
" T Witt k+
w (c1 + +z2)/ 1 T W (O)
will (01 + u2) }
K+
Notice that
kt (X1,t,- - P og (X2,t.-) = (0)
(t.t)
(+.t)
(t)T
X1,t.-1 -1-- (+.1) = kk (X1,t (11 - p 11/2+,)) (x11,
(t.t)
ke W+
T kp w+ X1, X Po-- (X1,-++)(1 Pg pg
(+.0)
(+.1) (t.))
(t.1)
Now consider the function y(œ) (1 x) 22/ Notice that if x1
<
3 and
< 12 < 2 then w(l
y(22) > y(21). Therefore, there exists a constant C > O such that if lc w have
@1 u1 S C, we
(+.0) ket -1 ktt (x1,t.)))/2) be (X1, - Pg- k 1)/2
- (t.1)
(+.1)
(+.)) ktc
Next,
(+.1) kto (x1,,t)) -0- Pyt ktc (x), - P P (X1,,+))(1 - p Pot- ktc (X1,2
(+.t)
(tit)
(++t)
= =O(c] =
(T w:l w(("1) m1) > (1)s(1)
W
where the inequality by Lemma
E.4
Thus,
of W wht w(() ( (ptt.0
(Pyten (-w (1) X1,t.- (ntt-p((t1)) pb( + p p Dt (-w (OT X2+ p(te))
@+
+ t WoOT (((t)
Q+ u U1 (t.tt) (pottt) (X1,ti-) - (tit)
-o--k (X2iti))
+ c at wif (-v1) ((ttt)
(po--ugkkton (-r04) (oT (p ) + kete (-w4t) (t) X2- (t)
+ of Wot w(t (-u1)w (t)T (-11) (p (X1,,,+) - (tttt) k (X2, (X2-
(
33

---

Published as a conference paper at ICLR 2025
0
Similarly,
c Wot w (-02) 02) (Dot-eakto, (-w4t (1)T X1, -1,-,-pu-- PP 12 ) + p (tit))
(4.1)
P
(-w4t (oT X2,+ ppt-12 -2-t.-P9--
Q+ Q+
+ ci W wilt) (2 (H)
@+w uzwt is (-ptt- (eet- (X1,t,- + P (t,d)
(ti))
(ottz,ktttt (X2it.-))
Qt Q+
+ a W w @4 (o--2zkt-1 (-wo X1, (p. ) + -tー (-w+ (w (t)T (O)
(t.0)
(1,)
X2, -
(t.1)
+ + w wi(t)
(1) (-01) (otee (X1,- - (+,t) Q+ Q+ U2
(tit)
Pg kt (X2, (X2,
0
Finally, since of WK @4 s 0, W (T)
w
K+ Q+ 20
w
T W witt) (t.t)
@+ a+ go (X2,t, -) Q (0t (X1 (x1 (0) (-0)
1 ( (X (x1+ (0)
21 + ci w wilt) (+.0)
( (t)
@+ *4
+ e (TT (T(T((()) (+.0)
ci w(th will) c2p (+.))
oe (X2,,+) (t)
a4 Q+
0
Thus,
duf ( w (t) K- Q+ c1 < 0.
di
The inequalities in the second column be proved by noticing the symmetric structure of the gradients of score variables in Appendix E.6 Thus, the gradient of the variables in the second column has the exact opposite sign of the corresponding variables in the first column.
E.4 PROPERTIES AT THE END OF STAGE 2
Lemma E.6 (Activation in stage 2). For t E [2].
Wk E [2) 8 E (+1) of (Ku.s. 0
Proof. This is a direct consequences of Corollary D.7 and Lemma E.4 Theorem E.7 (Attention at the end of stage 2). There exists T3 (log(1/w) such that
For samples X1,+ and X1,-.-.
when the query token is the group signal C1.
P (X1,+++)
(2. Pot ke (X1,-,)) S (w)
Pyter ke (X1+..+) Pyt ht
(+.15)
(2.
(X1,
(X1,,-) (w(w))
Pat (X1,
(+.T2)
(2)
(8 (1,-,-) (w(ww)
Pgt-c.kkt kt
when the query token is the classification signal,
P (X1++.+) 10 Po
(+.Tg)
(b.
(X1,,--) S kw(w)
P (X1,+.+)
(+.T2)
p/ (2.T2)
(x,,,-)
<ō(w)
22.
(+,2b) ( x1,
Pgt ug h--- m (X 11,-,- (ōww
(+.42)
Pg-uykkeo kg v2
(X1 (X1,+.+) - P P k (X (X,,,-))
(+.T2)
(ō(w)
Pa-12 koc
(+,Ta)
(x)
Pg+ ( bt- (X1,-,-)
(+.T2)
5 (ō(w)
Porez bo
Porng (Xi,t.+)
(+.Ta)
(+)T2)
(X1,,--)
<ō(w)
Pg uy
34

---

Published as a conference paper at ICLR 2025
For samples X1,+.- and X1,-,t+
when the query is the group signal 01
P pat (X1,t:-) > g (X (X1,-it)
(2)
(+.T2)
P (X1,..- Pg P kt (X1, (w(w) (((
Pater kt (X1,- P P kt
(2,
(X1 (Xl,- (w(w)
when the query is the classification signal,
kst (X1,t,-) - P -t- (X1, (X1,-it) = 9(1)
(++t)
(t.1)
q o2kk (X1,t--) - P (+.0)
(+))
ktc (X1 (x1,,++) = ((1)
P pg-u (X1, P kt (X1,
(+,0)
(t.0)
(1)+ ) S (w)
Py-- U2 kt (X1,+.)
(+.0)
Pg-- U2 12 (x1,, < (((w)
(+,))
1++00 p kt (X1,t,-)
P kt (X (X1,+-)) s (w(w)
(+.))
Proof.
The proof is direct consequence of Lemma
E.3 Lemma E.4 and Lemma
E.5
E.5 THE UPDATES OF THE SOFTMAX ATTENTION AT = Ti
In this section, we give out the explicit updates for the scores and self-scores at = Ti. The equations in this section are derived by the following:
1. Notice that Pi
(+TT)) for 1, 12 E [L] and thus w (71)
XR, sPi
(+,11)
w (Ti)T
Cx where
8 E {+1).
2. Therefore, by Corollary D.7 w (Ti)
Ch 2(1) and thus 8 (Xi,t.))
04 (Ti) (XI, = of ( (X2,+.-)= of (T.) (X2.-.+) = 1.
3. By symmetry between two groups and two heads in Lemma C.2 and Lemma we
have w (3)) w
(Ti) 02 and (XT,)(X), =
=g (T) =
(X2 Thus, denote ((T)) g
E.5.1 THE SCORES
dej
(
T K+ Q+
U1
dt
T=
L
= E
11 Q+ 4+Pli
W(T (0) (+.0)
diag (w()
(9) x - w xyp (10) x 4
(()
X
1=1
t=
+
g
(1)
yo
c
x diag ( K+
x T w+ w [xIxx(ttt)) pf-
( xi 11 X1.
1=
= 39 Til2
W TT (T))
e+
Q+ 01 - 0j
TT)
Q+
Q- b2)
w
(T))
(Ti) Xp (-TT1)
+ +39
1
(
(T)T
K+
w c1 c.
W
K+
(Ti)
(T(
K+ 22)
w (Ti)T
C1 -
w
(Ti)T
Xpr (+.21))
= 3 1
(T))
1 - -
(1))
(Ti)
(11)
3 1u (11)
ot
Q+
Q-
+ 39
1
(T)
K+
w(Ti)
R+
9
o
(T)
K+
w(Ti)
K+ 02)
(2
2 (T))
-w
35

---

Published as a conference paper at ICLR 2025
dcf w We
Q+ -
dt
3
1. m
1/2
W
@+
(T)) T (T) -
Q+
@
T
a
Wa-
wiThi
(T:)
T
-w ((T) X
TT)TT ( )
1 +3
1
(
-
 T
ai + G
w K+
T
W
(T
to - =m
xix (10)
K+
kt
(
1 =
(a) 1 1.m
( (1) 2i
W
2+
- w
a+
(1))
w (1)) (
2
(2i)T
+ 3
1
0112 ( -c w
K+
(T)) W (T)
c +
W
(T) (
k+
k+ (3
2
(71)
de. dcTW Wkt T w ci
dt
1 T
( (t)
yo
L
@+ qiPPi
( (0 (+.0T
diag( (w
(U) X - -w w (tt1))
Xpi ) x xTa a
Xx.a-
+
E
g
11)
yo
T
ci
kt
+x++**
k- diag x w+ F T
xpi xptt.'?
(t(1))
)
Pi
1
30
1
m2
W a+ ii- - c w W B+
T (T v2))
Xpi
++
+
3
1
(
(c)
K+
w
K+
(T)
01 - q w
K+
T
k
Xpx
(t))))
( (o) 1
T (T)
1
ii= - c w 
w
(T)
2
Q+
Q+
w2) {
la
+ -
1 00 (
W
WK400 01 -
w
T
w (Ti)
( v2)
(1
(T)) T )
K+
xt w 02 dej
*++++
dt
11
- (T)
-
1 200
()) W
(T (T)) W
a+
Q+
022 -
T (
a+
a+
(1)
(T)) 0 - (T)) xpi (.i)
+
1.
(
[]]
kt
W
(Ti)
K+ 12- -
a Wkt
(T))
K+
W (T)
(T)
xxp xpi (t:t1i)
(T)) T w
1 =
wa+ wo+
u2 - c Q+
T
w
(T:)
2
(Ti)T
(2)
NQ+
1
2/i (9) (
( T w(Th) 022 -
T
( (-
(
(T)) (h))
+ g
K+ WKt
WK+
k+
WK+
(
to
doj Wktt S+++0 @
ci
c
dt
(=
=
E
g yo
11)
(t) +
c
(T) Q+ G.1Pi Si diag(( (wt)T x - x-
[ (t) (++)
(() xp(t.() x Vn
XTe] Xn.a--
1=1
+
(t) yo
u
( K+
K 4 diag(1 [x ww - - xp(((t) pf p5tt0022
xo T.T c
Xh,0-- -
=
-
1
(
[c
W
(T)) T w (T
Q+ -0+00 c1 - cJ w
(T)) (TTTTT((Ti)) w cz) ()) (-w (TU)T Xpf (t.ii) (
36

---

Published as a conference paper at ICLR 2025
+ 3 1 w
1 (T) W W w v2)
(-wfT)) Xpl xp(t )
K+
= 1oon (o) w
1 3
( W wit c c- - T w 02)
3 (u (Ti)T
+ - ou
(T W 01 -1-1111 1y T T
W (12112
kt v2)
1 w F (T))T
01)
w K+
w
9+
dt
t=
1 - (-c]
T (TT W
Q+ a w T
@+ 02)
(-W)
(T)T Xpi (T) )
a+
Q+
+ 1
K+ - 12 w
k+ 03
1mm (TITT T.TXPp
=
1 2ine
(-
+
@+ c1 +
T
w (T()
Q+
02)
a1)
+ n
K+
W
(T)
K+
01
- 02
W
K+
T
W 02)
3
1
w
(TiJT
a
dof W ( W w-(t)
K+ Q 01
dt
t=
g E
yo
2 wel (t) (1) (+.0)
Q+ 4+iPi
diag (uf X - wt (t)T xpl x((+ XTo, 11
XA
1=1
T=
+ E
g yo
Σ
0j
x diag (X w+ - x.][xp(-10 p/- x 2o 01 W K+
Xk,5,-8
t=
1
g 222
(Ti)
q- q 0j w
(Ti)
(Ti) Q+
w((7)10) (2)
(T) Xp( (11)))
Φ+
+ 3
1
(T)
kt K+ c. of (w
T (T)
TT
WXTh
kt (2)
TT)T C1 - -wf (T)T xpl (1)
K+
=
( wiTh (Ti)
@+
q= q ej w
a+
(T))
w
Q+
(Ti)
02)
( w+
(Ti)
a
+ 5 0112
1
(T) (((
Wk K+ K+ c T
T
(T))
K+ (2)
2 (Ti)T
K+
3 w
a
duf doTWTOT WTOT w(t)
Kt Q 02
dt
t=
0/ T
w-(Th)
w 1 - - (T)) (TTTT
will) (2)
(-wo (Ti) xp (1(()
@+
+ 5g
1
(
(1)
+ kt
T
T w(ffle2 (2)
(w (D)
( Xp| (11)
kt
1 5 m0) (o5 w wiTh
of W
3 w (T k2)
- (TiT ))
3
+ 1
1
Ti2
(-op W
w Th
T
W T
w
W
2 (T)T
k+
W
02)
wut 1
E.5.2 THE KEY SELF-SCORES
dup
OTT w (t) K+ K+ U2
dt
37

---

Published as a conference paper at ICLR 2025
E
g
( yo
(1)
W (0) (+.0
4t1li (.ipl
diag [wIX- w xp] xp)- x X02 Xx
R+
1=1
1=1b
+ E
g 10 yo+
02 W 41pp( (+.0)
9tiPi
diag - 1---wi - w.Txp/ x x U1 Xk
XTvr
11
39
1
T)2
1
W
(TI)
k+
W (Ti)
Q+
c - w
(T w (1)
K+
Q+
02)
w
(Ti)
XpI (tit))
+ 3
(1)2
-u (T) T
a+ σ1 + u. w
T
Q+
(Ti)
(Ti)
Xpi
(+.21)
1
W w
4 - .. W
T
(Ti)
TTiT
59
1i
kK
k+
Q+
K+
Q+
61
+ 3
1
(T))
(-
W
(T)T
x+
w
(T)
4 of w
K-
T
W 02 3
1 (T))
du. W
w
K+
dt
T=
3
1 (1)4
(o) W kt
w:(T)) (T))
Q+ c - TT) W (
k+
(2)
(-u+
(TiTT Xpi (4.T1)
39
1 m1)4
2/
W.TW( (T (1(
K+ Q+ 4 of
K+
T W 02)
(
w
(T))T
C1
de] W W w
K+
K+ C2
db
(=
1
W (0) (+.0)
diag [ w. xxp(t()) x C2
=
E
yo
4+iPi
1=9)
+ E
9 yo+
W (0) (+ diag [w - w -.x
K+ g:iPi
XTa c
Xh.
I=h
39
1
(2)2
W K+
(Ti)
Q- 02 c.
TT
kt
Q+ v1)
(T))
C2 w
(T)T
Xpi (t.T.)
1
(T( (T(T)
(Ti)
(i)T
(Ti) Xpi (t
[xpp
+
K+
Q+ 11
K+
Q+ D3)
w
C1
3
1
( (T))
k+
Q- 0/2 c
K+
w m1
W
TT)T C2
+ 3
1
c) ( (T)
K+ Q+ 01 oj
K+
Q+ 02)
3
(T))
dc_
((( W
K+ k+c
dt
(=I)
3 m
1
(T
Q- - T w(T (2)
m (T)T 01 - TT)T xpl x (iT))
c
K+
1
(i)
(T))
(T))
-
Tw w v2)
2 (T))
- cisT C1
K+
Q-
K+
( dep
K+
K+
dt
t=
E
(00)
yo+ yy (t)
(t) (+,)T
4t(Pi
( diag (wIX- - w Exp(t(0)) XTer XA.s-S
diag
38

---

Published as a conference paper at ICLR 2025
(0) (0,, (-.0) diag (wIx- - -w, x XTc 0i
+ E
g yo+
941li
1
K+
(T)
Q+ و -...c W w((
02)
(Ti)T p )
+
1
1 2im
k K+
(1))
++ 0 1j
W
w w2n ++
1o - w
(T)T xp/ ) )
K+
Lir roy
[))
T
w (( W
- -
Q+ -
w
(Ti)
o2)
1
w
Q+
3
K+
a+
(T))
01
W
Kt
TT)
w
w m2)
3
2 (T))T Thi
a)
dep W W w (t)
K+ K+
dl
T=
- - - wa 02)
1 1gg T222
[ (T)) (T)) -
(
(111
3
k+
Xpl Xp]
+ 10
1 1
1
wT (T))
K+
o+
02 02 W
K+
(T))
w
@
a+
( (T)) (T)
02 -w+
-ut
(Ti) Xpi )
39
1
(T)2
( - c
T
w 02)
( wh
T
K+
0
1
(o)
(T
K+
W (T))
@+ 02
W
K+
(T)T
v1)
3 w
2 (TiT
E.5.3 THE QUERY SELF-SCORES
d
W02T
will
Q-
dl
T= 7
E
g yo
(1)
W dia diag (XTw+ w - -. x.[[(((+)) (1) pa o/ 2r 12 Xa..-
a+
1=1
1
+ E
g yo
2y wr K( diag
x w+ - -. ((t(t Dr
(Tw++
( ((t)0) 1T 12 Xk
T- -
=
- (T)2
(T
++
k+ ( с + -
a+
W (TT
kt (2)
(m
(T))
ci - w
(T)T (1.11))
xpp
w w (T))
+ 3
(1)2
Q+
kー 6 - 12
w
Q+
(h)
w (Ti) (2)
( 1 - w+
(Ti) Xpi (t.1i)
= (
(
++
W (T))
K+ 0 +
w
a+
T
W
K+ (2)
2
3
(T)) T
+ 3
( TT.
o+
K- 6
02
Q+
TT)
K+ 02)
3
2
(TI)
w w((l
du Q- U1
dl
((T)) ( 1
1 = 1
( (T)
+
K+ 4
a+
+++++++
kt (2)
(
-w --
(T (+ (t:25))
Xpl (T)
xoxxppl Xp (lui
+
io W
( (T))
a+
k= 01 - 0j w
a+
(T)
w
k+ )
(T))
mm) ( (T)T ai - -
1
(
T w(lh)
++
k+
q - nf
w
(
o+
WK+
03)
(2)
2 (T)T T 41)
39

---

Published as a conference a T paper at ICLR 
(CLR200 2025 0
+ - m
1
T
a1----1 - Wk- a
1 T
(3 3 )
(2 (TIT
1i
24
1a
+++
dcl wet W 2 cz
dej 061
dt
g
10
(1)
T
W'++
2+
K x diag (x ww - - "xxpp) x
Miae 02 al
T
+
(0)
x²O did
diag( (x x-w+-- w xp1
(x
x
c w
=
3
K+
0/ -
mi - 3 w'a
o+
kt
1n
(ri)
xpl
W
3 +
1
W
-1m
X (1T)
Bu 3 =
1
(
- W
(
2n
1
10
to
@+
K+
2
-aat
T
1 +
W
v2)
1
K+
x+
u 3
dp
(uu - =
w//
W
(1)
+ T
(T:)
(1-21) 1id
k+
(
3 +
1
(c
c+
(1)
x+
mi q wà
K+
(
=
(
W
k
w
WK+
(1)
1
+
1
( (
WKt mn
- ai
W
at
(T)
T
(
(T:)T
CI
wa22 T T (1) 02
dej
2+
wat
dt
Ti
k diag (xT w-++- - wI (t(()) ) (i aa
10 (1)
R
g
yoyo
yo++
at
kt
(
pi
1=1
1
1 T=
+
(
L
T (T)
K _ d (x TTwx "+ - -w x
(t1))
pi Df
2
9
2/2
E=
1
(2)
(T(T), (T:)
02 W T
W WK+ (1)
(T:)
C1 - wt -
a+ WK+ K
1
1
W
(1)
T
W
w 1 (1:)) x (tttt))))
+
3
1/
kt
- 02
+t
kt*
1
W
(T) 12) 02
+
- c.
(
a+
TM
(T)
kt
2
Ci
1 m
m) -
+
( 1
kt
W
+t
T
kt
(a)
(
1
(T)) e
w in
dej a
dp
40

---

Published as a conference paper at ICLR 2025
1 w
( G -w- - Xpi sp(+))))
(T))T
3
K+
+ 3
1
Q+ K+ - 1j w
o+
[T)]
W
K+
(1)
02)
w (- ( xpptttn))
30 0i
1
(
c cj W T
w 02
2
(uutiii C1
K+
3
+ 39
1
(2)2
(e) W T (T))
@ W
oj
W
a+
TTiT
W 02)
W
(T:T c1 (-
E.6 SUM OF PARTIAL GRADIENTS
aw
(t) Q+
dwitt)
U1
dt
dt
L
9 g000y yo
K diag (xwwl² w Xpp xp(tht)) pi p/t10002 x7 Ter X1,t
L.
+ E
gdd yo
(2)
(1) diag [x (wu w Xpi (10+ pi
T
11t00)20 xi 02 X2.t
1=1
W UK 01
(w U1 - -w
(1)T
(t)T X1,t, -1.t.-p(tt) p p poe
wko C C1
(w4t (t)T
C1 -w X1,+.- P Potei,kec
(0) wt
(+.1)
WK 12)
W+
(A)T (-12) (B)T X1, 1,,, p P pgt
(tit)
w K4 "2
W+ U22 u+ X2, -1, ) (++e)
(
[1)T
+w(l) + C2
(wt C2 - wt X2,+ pp ( (++t)
(t)T
(0)
+ +w!l) +WK+ (-01) (w!) (-vz) - (i) (() -p )
(t,tt)
Po-
dw(()) Q+(c1 +02)
dt
(1)
k (t)
diag (xT "4 - x.Txpptt10 (100) pi 21 (ar + 02) Xatt--
E
g
yo
1=1
(())
W U1
1u+ U1
(t)T
(() X1,t. 15.t.-p/ttt pft-1] p/tt porakeem
+WO
(1)T
(1) C1 -1- X1,t,-p(tt p(ttt) p potcy
((te))
K+
+ +w(00 WK (-U2)
w (1)T (-U2) - w X1,t - (111) (+.))
(b)
+W(()
W (1)T
(t)T X2,+, - y (t+)
K+* U2
U/2 w+
+ +W(O K+ C2
(w) C2 w
(()
X2,t - pp- P
(t+t)
+W(0 1 (-1)
(1(())(( (-11) - = X2,,, -2 (+1)
(w (t)
Mp K+ 02
K4
di
db
41

---

Published as a conference paper at ICLR 2025
g
1
yo
(
diag w w (xIxp(t()) x is X
1=1
E
g(0) yo
4+1Pi
(1) (+.0)
diag (wfx - w _pp p/(t)()) X 22 X2,+
(0
W (t) (-10)
W+ 0s - w't -wt (1)T
(1)
X1,t.- (
Q+ VIP
vipg-m
W1l cIP 95 K1
(+,1
(1)
(1) T X1,t. -p(t-)
W1
+W (-22)
(+.1)
u+ 1s --- -wt (t)T X1, (t.t)
(0)
12
+ +W(l 12p 92 12
(+,1)
18
(f)
02 w X2, P
(t)
w 02l
c2 02l 9t
(+.1)
W
(1)T
v2 w'
(
X2. -P
- Q+
(- (+.1)
1, 12
W't 1/2
(0)
w X - ppt
(1)
dw(() K+ (01 + 02)
dl
L
n yo
yo
(1)
Σ
4+iPi
(8) (+a)T
diag [ w. Xp w.Ix((=.)) -0) x (ci + 02) Xkst.-
1=1
/0
with W
(t.0)
(+.1)
(w(()) (()
Wi+ Ci - w X1,t,- -1(t11
(t)T
+ Q+ GI 9
(+.6)
u4 C1 - -w -1,ti-putt (tit)
(0T
(t)
Po+e
wilh Q+ U2))
w 67 - w [ X1,t, P
(()
(+.1)
KEc
(il!)
W - (t)T X2,t -aot.-pp/-) p/-
Q+ uop
(+.0) W c2p
(O)
C2 -w+ X2,t -Pa-es
(t)
(t,1)
W+
+ (-
(+.0)
(U)()) - w
(1)
E.7 EXPANDING THE UPDATES
To help analyze the sign of the scores and self-scores in Lemma E.5 we rearrange the gradient flow updates for those variables into forms that are convenient for such analysis.
E.7.1 THE KEY SELF-SCORE
do W/ (OT ( K4 K-
((
01
E
g
yo
W(JT (4,0) diag (wIX-wIXe) K+
4+PP)
wI w.xx((t)) x XTe1 Xn.a--
dl
1=1
E
g(
yo
(t)
T ( (0) (+.0) diag [w.x - (.xxp(tt)) XTa Xk
1=1
L
(t)
w-( ( (6) (6) (+.0)
diag [w,.x x w. Xp, (itt)) x Un Xmati-
yo
c
K+ 4+iPi
42

---

Published as a conference paper at ICLR 2025
= = ci (
w (+.
W 0j - -
(D)
(t) Tx x1,+,-Pp p
k+ Qt vip oipgeon
+ o k with ciP ke
(+.0)
(W
(t)T
V1- - (1) X X1,t, -Pp-o- ((t11]
+ @2
(9)
(t)T X1, 902
(+.1)
W+ 11 -1
- -q Q+ U2ll 10
10 (-11) - -w+ X2.t.
(t)
(t)T
(t.0
pgt
- -c 2 K- C2f (+,0)
(t)T
(-11) - -w-
( X2, p
- -q 2 (t)
(1(( w (-01) - -u+ (t) X2, (t.0)
K- Q+
900 9 W k+ o1 u1 kter (1 (1)T S1,,t.P/ttt ) + p ,tt-0 (-w(fT (-w(!T X2,t. ppt
(td))
+ + T with
Q la u 1a (go (X1,t.-) (+.0)
(T)
1g (X2,ti-))
+ + (t)
02) (pg (oo-e (-w++) (1)T
X1 -1,t.-19--0
(t.1) ) + (4.0)
pe- (-w+t (() p
+
(t) Q+
(- uu T
2l Pg--
Pgt- (X1,t.-)
(4,1)
Po-12.kt-1 (X2,
(+)))
+ w aalgoaaat kmm (-w (1)
(+00)
- X1,t,--,p (( ) - -cTW(( w
X
02p 02 (-wo² (!)T X2,t pbt-22
(++1)
K+ Q+
+ +o T 1 (+.0)
_ cj K+ c2ll c2l kg mn (x2 (!] K- Q+ at (X )0 01 +
(0)
whl (++2)
(t)T
V1
E
80 y0
yo
(1)
c W (0) (T.)
9tiPi
diag (uTx- w xxpxttt)) x 12 X
Gi W w (-
-1 (-
(t00)
[t)T (-u1) - (!) X1 1I, -. (t.1)
c W (TT
will)
@+
G
(+.0)
104
(()
(-w)
=/+ X - +P 9t-c)
(()T
(
+ w
@
(t)
u2p
(t+))
(t)T
(-v1) ww! ( X1, + pb-
w w
012 10 yt
(+.))
w U1 w
(t)
Tt)T
X2 X2,- -2, po (4.1)
kt.
- (((t)
Q+ c2p
kt
w+ 11 - w
(0)
(()
X2, P 16tt0))
- c. w (T (1) (+.+)
K- w PPo
w- U1 - - X ppt
(1)
T
= 0 c w T with (-v1) (p ((())))
(p ( (-w( (t)
(1,- (t.1)) ) + (t:1)
pgooo (-w(QT X2. -p(t111
K+
+ c W wilh W un U1 Pyt-on kt -0 (X1,- - p p oy (X2,-it))
(T) (t.t)
+ + W OT wil U2
(ttt)
(-w(!) x1, +p(1)) (t.1) + p 12- ( (wuf) (t)
X2
(t11)) v2n)
+ + W k- @+
(1)
w uz (t)T
(-01) (p (X1,- - (t+1)
1tー (X2
+ + w (1) 1TT W (+.t)
(-w+ x1,-,t (t.1) - c
(1)
T
W (t)
(+,))
(-w47 (t) Xa X2 (t+1)
K-
cipgr kt
Kt @ 02P
43

---

Published as a conference paper at ICLR 2025
+ tcj
T (0T
arw (-01 ( 1- - (X1,- - ej T (OT (t.0)
Vr c
"r (X2-s)))
I.
E
g
11 yo
+
T WkIT (0,1-
K+ 94ili
diag [ - w.xxp/+ (11ttt)) x xTa XM.a-
)(0)
11 W w:(t)
(++))
(t)
w+ X1,,, Patte )
(f)T
K- Q+
u1
W4
+ Foj w wilt) (+.0)
(1)T
61 -w+
(1) -1.t.-ppt-1
Kt q+ gt
+ to TT wilh (- 12 2 pg
(+.)
W+ Ci -w
(0)
w
[t) X1,t. - -P 99-
(+.1)
bi T @+ (- P
(+.0
W C1 wt X1 - -
(0)
(
+ (+.1)
T w-tt)
(+1)
W+ Ci w X1 $1.- ((t.))
(t)T
(1)T
oi K+ c1l kag
bi W ( w
(++)
[ut a (t) X1, -ppttt)) (t+tt
(t)T
Q+
dup
K+ u1
glth y (t)
L
TWOOT (+:00
diag - -=] xpp (((()))) x Xk.-s
dt
2R
V
K+
94iPi
L
wTlt (8) (+.0T
diag ( - -.w- _p xx( (111)))) x oi Xxyto-
g yo
K+
94iHi
(0)
(
(+.0 Q+ 01 kt
(1) (T)
W+ - - X1 X1,t.-
(t)
+ rj w w (1)
K+ Q+
c1P
(+
W4 U1 -+ X X1,-,-P -Pg-es)
(1)
(1)T
(t.1)
+ +.Tw Kt Q+ (-02)p (+.1)
(t)
W+ 1i - --+
(()
-w+ X1. pg 99--12
(t)
(t,2)
-uTwkk (((t)
(11)
W (2) (-21) - -w X2 Pop-11g
(t)
K+ Q+
u2p
-. T K+ Q+
(
(+44
(1)
( (-vi) -w+ X2,t. pp+ )
())
c2l
- -0,T K+ ++
w (1)
(-11)
(+.0)
w (-01) -+ (O)
(0)
- (t.0) pgt pp
0)
(1) T (2 K-
will
Koto (-w13T (2)T X1,
(t.t)
-ptt1) -1,t.--(+t)
+ (ti0)
(-woff (1)T
- PPg--e))
(4.0)
8- 01
+ + Wit (!
Q+ m 11 (Pomniee (X1,t.-) - Pp P9--vs (X2+i-)
(O)T
(+.0)
+ Tw
(1) (-22) (t.0)
Pg--o2 (-w(fT (1)T
(t.t)
Py--pikt- (-w(() (e) (((
(tit)
@+
- + P
T (0 (C)T U1 P
Pom (X1,t,-)
(t.0)
Po-ozkk- (X2+.-)
(t.t)
o. k
(++1) (-w (1)T
x1,---((t-)) )
U. ((() k+ @+ geg
(+.0)
(-w(() (1)T
X2,t (
t
W K+
Q+ 1p
1P km (X1, (x1 uy + oj w(O
(+-0)
(
K4
We4 cop
p2 (X2+,-)w( U11
(+.1)
(T)
44

---

Published as a conference paper at ICLR 2025
L
gd yo
T Whe (000p (+.0)
diag [ x w xpftttt) x 1i -x
K+
9+iPi
-00
11 W T
w (-U1)p
(t-t)
(1) (-v1) - w- (1) - -P (4.1)
K-
+ vi WT w
Kt Q1
(+.0)
1+ (-11) - -w+
(H)
T
X1 + Pata -
tu T
Kt wilh u2
(4:)
(B)T (-u1) - w
T
X
pl p/tt122
- Kt wilh (-12) (4.)
W+ V1 w X2, 22-
(t)T
(
Pfgt (+.0)
-u -T K+ @+ p
(+.0)
(t)
[f)
X2. x +
w+ 21 wt
- w ( (itt (t+t)
(t)T
T
X2, x2,
K+ @+ k
W U1 w
"
1 W2T w w (-11) (p kt-- (-w (t) (t.0)
-1 + (+.1)
pot (-wffT (T) X2, (p
K-
+ +oJ TTTW with viw (9)
((ttt)) Poo- (X1,- - P -ogerkkeu (82 (22)-14))
(t.))
TTT Wo U2 (+.)
(o
(ww+ (t)T 1/110
- 11
+py P k2 (-w(0T X2 X2
(+11)
(1)
o. T w u2
(1) (-01)
(Potoz (X1,-+t)
(ttd)
P-- 22 12 (X2,-
(+.1)
+ v w w W cip
(+.1)
(-w(f) (1)T X1, p
(t-1)
U1 WK+ W @2P (-wff) (t) X22
T
(tit)
+ t T wilt) (()
K+ ( (-e1)p (X1,-,-)- 1/ W c2 or (
(0)
o1 (Xe-it))
E.7.2
THE QUERY SELF-SCORE
T de
W(())
Q+ W1
V1
dt
1.
(t)
Q+ K4 diag (XTw+ ww -. Xpx (10))) ia p1 I Xuand
E
+ E
9 (1) yo
(1)
U1 w
w ( x diag (xTur &Tw+ wI xxxxx((0 (000 pd (0 a Xka-
Q+
g (1)
(t)
c w 2 (x +++ - w Xxxxx(+*0 (01) pi 2o Ten X
yo
())
ci ( w(l U1
W U1 w X1,+, p5
(1)
P pgo
(te)
Q+
w+ C1 - --- ( p/t p p ktt
(O)
+ c Qt e e1
+ + W
Q+ w (-v2)
(w+ (-1z) - (t) X X1, D (tit)
p Pot
- -cT w U2
(t)
(t)T X2,t $2.t.-Pae-ae P pg
(t,t)
w D22 w
- -c. T ((())
[t) (t)
X 1tt pgt kgt
(tit)
at k
- w+
45

---

Published as a conference paper at ICLR 2025
- -c WoT W W (-01) (w+ 1 (-01) ( X2,t, -ett.-Pot Pot ) k
(t.0)
Q+
= - c ( W ii ( Xi X1, ppz pgt pgtmiiktm + (-w(f) X2,t, 12,t.-Pp Potiing ) p -o
p(tie))
a+
+ oj W w mn 01 petor (X1,t,-) - Pe
Lo
(p(t.0)
-P -1 (X2.t.-))
(t.1)
+ +o @
a+ Wx+ (-v2) (( w (1) X1,ti- (t.1)) ((t-1)) P + (-wtt (9) - pg--mkkr
(oT
(tit))
+ + T
a+ w m2 (9T
02 PPttenkk-o (X1,- - Pot-en u, (Xe,ti-))
+ + w
at e w+ - wt
Tt) X1,t. pb ) P Pgtty
- - a+ Wx402 C2 (w c2 - w+ 1p.t.-Pot )
p p ,g
(t,))
I.
E
gdd yo
(2)
Ci T w:l K diag (XT,, W, - w wIxxp(+ (xp(t(o) Pi
pf(t1000t x7 11 Xx-
1=1
9 ( w K+
( (-11) (w()T(- (-11) - - X1 -,, X1,- +P
(t+f)
p(tttt)
Pot-orkkt-r
c w
(TT
(
w
(t)
(O)
(b)T
(t,1)
(+.t)
Q+
CI -w+ X1, Ppp-mm
p
P9t-v,
(t)
(1)T
(B)T
(+.1)
(+,2)
+ Qt
K+ U2
w+
12- -
x1.-,tPpr pyt
P 11
c. w wilh
Q+ k+ (-12)
(wt (-12) u (b)T
(1)
x2 //tt)t
p ((1)))
Patv
c. w W 02
(t)
(1)T
(tatt X2,- + p ) p kttes
C2 - w+
c w ( Q+
w
T
(B) X2. Pp p p
(t.t) )
(t)))
U1
[wt 11 w
00 9 Q+ (
w
(1.1)) (() P Po--v ko- + (-u(QT (-wlIXx X2. xe_ ) v
(1) X1 -1,-,+
(tit)
q w! T wilh
Qt WKtww uw ts
(T) (t.1)
Pot- (X1,-
P ktv Poter (X2.
(t,t)
+ w
Q+ W u2 ((ww (t)
(T X1 P
(+.1) ((1.)))
P + -wi (
(!T x2 (( (((2))) ) (tt)))
Poronkt-ss)
+ +c Q W
will
u2 02 ( Pott (x),, - P p -t (X2,-.t)
(9T
+ + w Q+ a ci w ci - (())T X1, x1 (t.1) ) p p
(()
(t.t)
- q. w o+
wx 02
w+ C2 - w+ X2,- +p/t
(O)
(
(t,))
1
E
g yo
(t)
Σ
2 will
off
(x w+ - w xxxp((t xpp(ttt))
pi pl IJ as Xman-
(tal)
(00)
of W (T
(1TTTT 19 -
- ()) X1, pfttlt)
-i,t,,p(ttt) ) p f
k+ U1
+ vj wT(
w ci (w ci (()T X1,t, 1 ) (tit)
(t)
kto
46

---

Published as a conference paper at ICLR 2025
+ tuj T W (-12)
(w 10) (-02) - -w+
( X1, Potcs) Patcr
(tit) (t.))
- T a-
w k+ (-01)
(wiTT (-11) - -wt
(t)T X1, x1,-, - (t.t)
Pgeci
oi w w
(t)
(t)T X1, - P p ) P p
((t)))
R C1
C1 -+
- -e Q-
(
U2 -w+ X1. x1.--t +P p (+.0)
(t.0)
K+ 12
w+
)0) 19 w
(OT
W w U1 a+
((ww (1)T X1, pp- (t.))
ptt + (-w4+ (1)T X1,- -ftte) ) p p
(+)))
+ 1i T R u u T9T
(Pytea (Poteiitte) (X1,ti-) - (t.t)
-tt (f (Xl,
+ t T wx (-12) ((
(11T X1 (-w(ITx
-P.t. Pit + (-uft (f) X1, ((ttt) (tit)
(t.1) (t.0)
+oj T Q- Kt m2 w
OT (+.0
(p--,, (X1,t.-) - (+i))
dt (X1,
ut
+ t -
( w
c1 c1
w c1 (t)T X1 S1,t.-pftt p( ) p Rto
(t)
(t.))
- a
(t)
T X1,- (tt)) )
(t.))
k+ c1
c
kete
de_ at
w(O) 0t
2F
g (t)
(t) yo
wrOT K (XT:: w+ - wIxpp
pf p((100) x7 a Xk..-s
dt
+
Q+
E
g00
yo yo
(t)
c T wroo K diag (XTux ww - w [xxxtt0) pi af as Xk.0.-
)))
ci T wl0 K+ U1
uu4 U1 -
(1)
-
(t) 10)xx1,+ X1,t, -...t-p(rtt) p(
P f pgrankter
w w
[ (t)T
Ci -w X1,,, T -1,tt-p///t0 ((tttt)
of
K+ ci
+ +c w
Q+ K+ (-w2)
wwt (-12) ( (t)T X 1 pptt)) p po kt
(()
c. w a+
(()
k+ (111)
w+ (-v1) -w(
(1)T
(b) X1,-,+P/t00 (t,))
D pg ktccktt-u
cf W whe
w
(
C1 w
(t)
X1,- + p//
P ktc
cf w at
WK 12
T
U2
w
(t)T
X1 1,-,++1( 1
(t.t)
P p
(t.1)
900 c ( T
W U1
((-w(1
(t)
X1 x1,t.-p/tt
p + (-w4) (D) X1, (p (t.t) ) (t.)
a+
+ o Q+ Kt un (0) 01 ( (X1++--= - p P (X1,-
((th)
(t.t)
+ + wT(
at whf (-v2) (( (--(() (1) X1.t.- p/- ) p p Potc,,kt-0 + (-w() (1) X1- ((tit)
p
+ + Tw (
at kt v2 (1) Kt- (X1,t--) - P (X1,+)))
+ + wilh (TT
Q+ whic wu
(O)
(t)T w+ X1,t. p/ P
(+.1)
- -c w Q+
(f - w+
(t)
X1,- x1,
k
47

---

Published as a conference paper at ICLR 2025
E.7.3 THE SCORES
de. WT K+
1 yo (1)
ci W (2) (+,))
diag ( (t)T
X w Xpl xpf+t.))) xT oi X
g
dl
+ E
( yo yo (2)
c K+
K diag
xxT ui - xpp xpptte)) p p(+ x9 02 Xkia-5
E
9 10 yo
wrlhh
(tt) diag (tt X - (t))T Xpf -ofTTpff=1) ff(tt.1) x 01 Xkio-
Q+
900 c ( W with 0 mn
(+.0)
(woeT (t)
1 61 - X1,t,.pp -pftt11)
+ +cT a+ Q+ af
wil (+,)
(1]T
10
CI
+ +cTwet will a+ (- P ge- 122 k2
(+,1)
(()T
(!) w+
X1 X1, pg (+,2)
W+ C1 -w+
- -cT ccTwet with) (- p (-m
Q+
(+.1)
[w_ G W+
(t)
X1 -,,-, + Pa qe--en
(+.1)
- cTwett wil) (+
w+
(t)T
w
(t)T
X1 1,-, 11,,,PPg-0 +Pgtts
(+.1)
-cT @+ ke
w (+.1)
( (t)T
W+ c
w X1, x
(
L yo
E
W K diag x (xT++ W4 - w xp(+ Pi
(4.t) xi as Xhirs
00
(++ ( U- ww (t)T Xi,t. -p P pgt
(+.0)
K+
K4 U1
+ cl W K- W C1
(t) a-w+ X1,t p/t (t.))
u4
Patce
+ + K- K+
(( (t)
(-U2)
w+ (-1z) - -
(1)T
T X 1, pftte) P
(+,t)
p
-c W 1 K+
(-11)
(t) (-11) - w (2)T X1 - X1.
((11))
(()
(b) X X1, - pg p kat
(t.1)) (t.t)
K+ ci
w
-w
o. W (TT
Ww U2 w X1, x1,-
(1)
(1)
(til) (+.0)
K-
K4 U2
0
1 @T w ((-w(0 (1) X1, (-mx(TXii++
x1,t,-1(+t) p +
(tit)
(-w177 (t) X X1,, ((-0)) ) (tit)
pgra
W K+ un (9T
0i (ptt (potor (X1, - (X1 (X1,-it))
(+.0)
K-
+ + (T (t)
k+ (-22) (( ((-w(1
(1)T X X1,t, - )
(tif)
+ (-w(ft (t)T X X11 p(tit) p kga
pot 1
+c-
(
K+ u
(T)
12 ( (X1,t.-)
p Po-ck,tuu kt (X1,-)))
(+.t)
+ K-
K+ c1
(O)
ei w
(()T
X X1,,, - (1))
p
(+.))
w (1) w(()
(()
(t)T
x1 )
(+1)
K4
c1
(w c w
k
48

---

Published as a conference paper at ICLR 2025
dof
W K+
w
Q-
E
g yo
yo
(1)
Σ
1i T (+2)) diag w x " Xpi xp(+.0)) xTn, i XR
4+iPi
diag
dt
+
1=1
+ E
g00 yo
1 W K+
K diag XTw++ - w xpffttt)) ps (++)) x7 11 XkIoT
I.
E
(t) yob yo (1)
2 WllT (0) (+.0T
4i1li
diag (T X --w++ (epxx((t)))) xT Xkit
uu Wof @ u1p uip gd
(++)) wet 2 ( X1,t. 11...-Pgr-1 p
(t)
Q+
+ o T W o
(+.0
u+ U1 wt X1,t. 11,++-Put -D ((t11)
(1) (()
+ + w w u2)) (11)
(1) (t)
U1 - w
-
(t) X1 +. Pot 112
(+.1)
12
- - - a- at u2p
T w
(+.1
W (1)
(-v1) - w
(t) X2,t 14 2/t-22)
- -uT w Q+ 02p
w+ (-21) - X2.t pgt
(t)
(t)
g2,
- -u.w a- w (- 12
(+.))
(t) (-1z) ( -
-w+ X2, pg p
(t)
= 000 1j wo w U1 (poter
(ttt))
(-w+ (() X1.t.- + (t.1)
u, (-w(f) (t) 12.t.-Putm1
+ 1 W with ui (9) 01 (pgto) (X1+ti-) - Po--01 11 (X2,+.-))
+ F TT wiyh (-v2) (tit)
Pot-uz -w+ (t)T X1,t.- 13,tt-Ppe-- -/+ ) + ket (-wf (1) ToTTTx22t00
X2,t.
+ oi a- with (- (1)
P oogt-n (X1,t,-) - Pot Pot (X2,
(t,t)
+ (uT Q+ @+ Capate ke + WLOT will) (t.0)
(+.0)
[()
Q+
02
+ + wh w(t) (+.0
(-w(f)
(1) X2,t. p(-
a+ cap
- w TT w((t
a- @+
02p kt
(+.0)
(-w(f (1) X2.t. -2.t.-ph-02
( yo (t)
u Wlll (0) (1pp(t (T. diag (0(TTx x- ( Epp(tit) xT 0g Xx-
)0) (TT W w (-v1)p
(+.t)
(t) (10( (-v1) (1) X1, +
Q+
01 (w+
+ o. TT e- wit cip cipgr
(+.0)
(t)
(-ei) - -w/ $1.1++((-12 (t.1)
w
+ t Te will D2p (
(td)
(-11) - w X1 /
Q+
- -,, wTTT w
Q- @ 12)) (+)))
w U1 w X2.
(t)T
kt
-uTT @ o2
(+)
k0
w+ U1 w
[()
(1)
X2. - P90-2
(t.t)
- -wTw( Q- @t u1ggeop
wxlt (+
(+ Ui w X2, X2_
(t)T
49

---

Published as a conference paper at ICLR 2025
= = ( ( ot wot (p (1 (-w(f) X1, -1,,,tPot-0
(111) ) + (1+
demm (-w(f X2, -Dbtt1
+ 1 T With un (T is (Pa--onkt-n (X1,=,-) - Pot dit (X2,
+ Fuj T wa 22 (tt)) (Potesktte) [wo (1)
X1, -1 + ) + P 02 (-w4
(+.1)
(1)
X2, X2 p(t
+ v w with u2wt ( (-01) (Potoz (X1, - Pat Pp (X2,-
(t+)
+ o wor w @+ C
a -1 + Ui W will
po (1)))
Q+ Q+ C2l p 02
w+ (-11)
(1)
+ i Q- w (+.0)
-wt (1) X1 (t.1))
- - - wilh) (+.0)
(1) X2, xt2. (-w
Q- @+ ko
L
(t)
2 W xxddiag x
(Tw++ w+ wpxxpft* xi)+ pf xT us Xx.ti-
9 yo
K+
w welee U1 [w w u X, -p P pororike
K-
+ bi TWOT
K+ w
# 01 - (T) x1,t.-p(t0) ((111) p p kt-
(2)
+ t T
K+ W (-12) (/ (-12) - -w+ X1 X1, Patm
(1)
(tit)
Poterkk-uu
- K+
w U2 w+ X2 p(t-1))) pat p1
(t)
(1)T
(tit)
Kt
02
-uT
u C2 w+ X2,t -p P
(t)
(t)
p
(t.t)
K+
K+ C2
- -wT K+
K+ (-01)
(w4) (-o1) -
-w+ (OT X2 22,+:-P9t-- P
(t.t)
) Pge Pg kt
(tt))
(t.t) 900 0f T Whit 11 ((-u(0 X1,t,- pg potenkker + (-w4f UT (tit)) ((+1)) -2 p
+ vi TWiT w
K+ WK+UI Witorffo (0 01 (Pyto) (X1,t,-) - -pa--, (X2.t.)))
-eTw K+ 02 02 ((-w0)
(1) X1,t pytes
(+.1) (t,t)
p k + (-w4f (OT - ) (O) (til)
p )
(-11)
+ +u w w
K+ Kt u2uf
02 (p(tte) (gt-en (X1 (X1,- - p(t.1) oi (x22tt.)))
+ t W
(104) c w+ (t)T X1,t.- ) (+.0)
K+ C
-o T K+
w
Cx+02
02 [u+ 02 -w+
(t) (tit) ) (tit)
-P kg
I.
(t)
d(t diag (x w+ - w. w_[xp[+ (1()tt)) (tit) aT
Pi 27 us Xx.-it
g yo
2 K+
(0)
of W ( w(())
(-v)) (wi()T (-v)) - w((hT xx, X1,- - ppt (ttt))
p
K-
K+
+ i K+ c1
c1
[t)
ci - - X X1, (+.1) (t.t)
(t)T
Pot-on
+ o w w
K+
02
ww 02
(t)
w+
(t)
X X1,- - v2) P
(t.t) (ti))
pgr-
50

---

Published as a conference paper at ICLR 2025
- - K+
W (-02)
(0) (wt (-02) (u44T(-0)) - -wt X2, X2-- (tit)) (tt)))
(0)
Potorke-ug
- -u - K+
wxic2 c2
u+ c2 w (1)T X2, X2,- p (tit)
- - K+
W U1
U1 - w X2, x2, (+t))) (w
(
0000
19 w (t) Wk+ K+
(-v1)
((-uf"Tx,, X1, -1, ) (titt) -o + (-wff) Tft X2. pbttt) ) ((t))) pg
+ uj WK+ W U1 1
(OT
2s uy (X1, - P Potus kt (X2,-
(t.t)
K+
+ to T W 02 (-w K+
(1IT X1, X1,
(+,0)
p pqg + -w (1) (tit) ) (11)) p
+ +oTw K+ K+ u2
w
(9)
(Pot-eikeu (X1 (1,-(+) - (+,1)
(+.0
Po- (X2,-it)
+ + K+ c+
(()
(t)T X1 -1, (t.1) ) pgt kot
(t.t)
w c1 w
- - w
(t) (t)
X2.- - (4.1) )
(t.t)
K+
C2
w C2 w
kot
dup W w
dt
Q--0 ci
E
yo
+
C w (8) (++)T
9tiP)
diag (w+ (t) x - w
-- Xpp (2(tt))) x xT Xk
1=1
L
E
yo
U1
W T _ diag (x (XT=+ K+
[ili-w]jeo w.Txp/t-9 pi 2x 03 Xx.e-0
E
g0 y0
yo
(1)
c 9tPPi
T (t) (1) (+.0
( diag
ww+ x w xp(+ TTy Xk.t.-
(t)T
G ( W vIP uipgeng )0)
(4.))
w4 U1 wt X 11,
(()
(t)T
c w ( wilh Gr a k00
(+.0
144 11 w X1,+ p(ttt)
(0) (()
+ W w (-u2)p qt
(12))(((())) (t.t)
w U1 (( [t) X1,+, 1,,,- py
(1)T
(t+t)
@+
-2,
Pyt-u1 12
cJ w w u2 gt2
10 (t)
(-21)
w Tt)T
X2,+, pp2
o+
of w wilh
c2p 0 kt
(+.0)
w (()
(-v1) - w+ X2,t -..t.-Po -p 1
(t)
Qt
- -c ( Q+
w (- 11 (+.1)
(1) w+ (-v1) - (t)
P p
(+.0)
= の
4 wToT ( w Q+
(p (-wef (t) X1 ((st)- (t.t) + +p + -g (-w(!) X2,t, -
(t12)
+ cf Qt W un Us ( (X1,t.-) - pgt - - (X2,ti-)
will
(tit)
(t.1)
+ t T
Q+ cip
(-w((T X1.t.-P(t.)) (t10)) - oJ (OT w(t)) c2p t pegkkton (-w( (() - ppt
@= Q+
+ + T at wift (-02) (((ttt)
(pet- (-w0t (()T X1,t.-P (1.t--pg--mn)) + Porozkkt-u (-w4t) (1)
(t)T
(+.t)
Xe,t, pft
+ c w w v2 Un (-p (td)
(t)T
-Ppg-ookkt (X1,t,-) + (1,1)) kt- (X2,t.-))
at
51

---

Published as a conference paper at ICLR 2025
+ + W will
wat oip koor (w 1T is + + q+ gt in (X2,,,- -wo 2/2
will
L
g yo
(1)
c q+ 9tiPi
T (0 (T
diag (
(0) X - w xpft x T Xki
900)
9 W
a+ q+ ( (- ) (tit)
(t) (-11) (-v1) - (T X1,- -1,-,+ppt -pp
(t.1)
+ q at Q+ PP
wr(t)
(1)
(-11) - -w+t"
(t) +
+ W will
Q+ Q+ u2l
(+)))
10
(t)
(-11) - -w+
(t)
X 1,, + D/ ppt
- -q TT a+ @+ (-U2)1 (1,1)
wilh
(t)
T X2, -
(t.0
W+ U1 -
02
kt
- c] w w
(+)))
(r)T
(f)T
X2, - + + (t.1)
@+ 02p
wt
w
- -c] W (
u1 (+))
-- v1 -+ X2.-
(t)T
(f)
(22.-,PP400) -
o+ Q+
- G (w a+
w/4/ (-11) / ( (-w(f) [!) X1, - ) + (t.0)
(p (t,0)
(t.0)
poterkter (-wi²7 (t) X2, -p ))
+ c Qt 2+ (-ur) (1) (-11)
(1)
Pgt Po--or (x1,-1+) - (t.t)
(tit)
keto (X2 (X2)
+ + TeTwe w @4 aaaaaaakto
(+.)
( X1, -1,- X1, +P p(11 oT W will) (+.1)
Q+ 94 C2P (-w (t) - -t
wot
+ + ot 22 (p9-12kt-011
(++1)
(-wffT
(t)T X1, pttt) (t,1)
) + og-
(+.1)
(-w (T X2, X2,
(+.0)
t w (
ot @4 U2 (1) (-01)
w(O
(p (p (x1,-,+) pg-- (X2,-
(+,1)
+ w (T
W (+.1)
Gip app (X1, (X1,
(t)T (-11)
oJ [((t) w/O0 (ti))
Q+ cap (X2, (X2 (t) U/2
(4) yo
(t)
of WTOT K x(" [x (TT++ w++ (xxxx(tt)) pf (() a Xn.0--
K+
()) of w W T w w U1
(wo(T 11- w X -1,+,-phttt D p
(1)
(21))
+ o. TW W # C1 - ((()) X1,t, 1ft00) P p
(()
+ t T
K+ w
(T)T (-b2) - -wf X +.-+ 9
(tit)
P p
k
- -o K+ K+
w(t)
-1)
w (-t1) - -w+ X, X1.-,tPp
(1)T
+
ppt
(t.0)
P Pgtc)
- -uTwRe
u+ C1 - w X X1,- p ( ) p
(
(1)
K+
C+ CI
- -wTwee K+
(T)
U2 w+ x X1.-, / ((tt))) P (tt
K+ 02
( of 1 We 11 ((-uL)
(- X1,t,,-- ) p + (-wff (1) -1,-, pfttt) ) p
WT w(th (O)
K+ w 01 p (Xi,ts-) - (4.t)
P (X1r=st)
+ 1
52

---

Published as a conference paper at ICLR 2025
+ ij wha W (-02)
X p ((t
-t + (-w(f (9) X1, ) (t,))
K+
p (P-c, (X1,,-- - P pot utou ( (X1,-:t)
+ vi W / m
(t)T
k+
+ tu T W
c1 w X1 X1,t, P1 ppteet )
(ti)) p (X1,ti-)
k+
R c1
-wT K+
oi
K+C1
w+
(t)
<i - (t) X1, X1, p (X1,-
F STAGE
In stage 3, we show that training the neuron weights can drive the training loss to be arbitrarily small. This is possible since the attention features now are linearly separable with at least a constant margin. This can be seen by considering the neuron weight w+ ci +22 +u2 and w- = c1 -1 -2. Simple calculation can show that this will make yf(X) 2 22(1) for all (X, v) E
-1
Lemma F1 (Growth of classification token alignment). Let ofr g (Xi,t))) and gDd
gX1 For t [T2, Tz). & E [2].
( dw (1)T
Uk = 0 + Ō (wg=)).
dt
Proof. First of all, by symmetry, we have -t =g ((2) (x (X and - = (X2,,)))) = (X1, - (00) (X2..++) ). In stage 3, we have
dw dw (
V1
dt
E
g
yo
(e)
Σ (xp(t()
L.
xppttet)
T
21
1=1
8 g
1 (1)
(-1) p (Xi,
p(ttt))
p (X1,t,-) + (011
Po (X1,-+t
040 (X)
P p
(td))
k (X1,- ) (t,e)
Poe- (X1,-,-) + (tid)
P (x1,--))
+ Pa kt- (x2++i) + (+.0)
(+.))
p (X2,t + (+1))
pa (X2,ti+)
+ + (1)
(X2_--) (t.)
P (X2,--- + pp- -u kt (X2,-,-) + P Pot--en -2
(t.t)
(+.1)
(X2,--))
+ 89 P P (X1,t )) + - +p (X (x ) + P g+ U2 kt (X1,
(1) (+.f)
(t+)
(++))
P9- (x1,-+t) + p8 9 k Ktー (X1,-+) + / 1 (X1 (X1,-+)
(+.0)
(+,1)
(+.1)
Pgtoz (X2,+.)) - D (t))
(+.t)
kt - (X2,t.)
P (X2+t--)
(t+t)
-pttog.ktyy (X2,,tt - p pgoo (X2,- - (t.t)
(+1t)
(+.t)
og (X2,
By Theorem E.7 we have
kto (X1, + (t.t)
(+.0)
kot ot (X1++.+) + (+.t)
k (x1+t++)
+ 1 (X),-.)
P P (X1,-,-) + P (X1, (X1,-,-) + (tit)
(+1t)
(tit)
Pg (X1,-.))
ky (X2,t.+) + p btt (X + Pgr (X2++.4
(+.1)
(+it)
(++))
+ +0+ (X2,--) ( (X2,---) + p kt (X2,-,-) + Pgt -2 (X2,--))
(t)
(t.)
(t+t)
(tit)
0(1)
53

---

Published as a conference paper at ICLR 2025
and by symmetry between two groups Lemma C.2
(+.)
kto (X1,t,-) + p (X1, ) (tit)
(+.t)
Po-- (X1,t.
+ Po (X1,-,+) + Pop- kt- (X1,-,t) + g0 (X (X1,-+t)
(t.0)
(+10
(+,))
(+.0)
(X2,+ uu--u, (X2,+,-) q- kt- (X2+t--
+100
(+.1)
ke (x2,,,+) - ket (X2,-i+) Pot kt (X2,-.t)
(+,t)
(+,)
(+.0)
= (ww)
Thus,
dw (t) 01
0(04)))))0 + = +0(w""))
dt
Lemma F.2 (Growth of group token alignment). Let g0 = g00) (X and 900h = (x1,t))) Fort E [T2, Tz]. & E [2].
(C)T ck 2 ((g_²) + + (wig(²)
Proof. In stage 3, by Theorem E.7 we have
dw
C1
dt
E
g 10 yo
(t)
Σ
2xxpfttt x
Xpp
T
C1
g 9 Po (Xist keter kto (X1+++) + (t.0)
1 (1) (+.1)
koo (X1,ts+)
0
(1)
(X1,-1) p (X1,-,-) + P p uy (X1,-,-) + P p Pot- (X1,--))
-
(9)
(ptte (X1,t,-) + P p ktc
ktc (X), + P p kg
kt (X1+t--)
-P (X1,-14) - p uy (X1,-,+) - k (X1--1t)
(+.0)
(t.0)
2 = (wg
+ Ō
Theorem F.3 (Convergence). There exists T3 such that T3 Tr = O(1/e) and ElE(y S
Proof. Let 94 (t)
= (x1,t.+) and g 1 = g0 (XL,,-). For consistent samples XR. we have
(X,+,))))))
dw + (T) dt
Pota + (t.t) ktt + Pgt-vakktmy
(ti))
(+.1)
dt
dw T dt
(p + (110) bye Potoo2kte)
(+.0)
(+.0)
(tit)
dw dw (T)
((ttt)) ( + Pyronkke + P pytee
(tif)
(+.t)
dt
dw ())
U1 (!) (-,) (D + ktt, + D Py-bykeo
dl
54

---

Published as a conference paper at ICLR 2025
dw _w (t)
12 _
(-1))
(-)0)
P ( + k + Pot u2- kt
dt
dw (t)
6i (()) (p + ktt + Po-i
dl
By symmetry, we have
dw (
01
P de (Xi, + (-1))
(-.4)
ke (X1,+i+) + (-1) 0g, (X1stit))
dl
(!T dt
12
Po-on (X (x1,+++) Potu k (X1, (x P k4-2 (X (1,t+t)
(-1)
(-1)
(-1)
dw (
C1
(-11) p (X1,tit) P k (X (x)++)+ + P (11)) 1g (X (X1,+++))
dt
dw ( V1 dt
P
Kt (X1,-,-) + +tt ug kt (X1,- + Put U2 kg (X1,-.-)
(t.t)
duf" dt
U2
(gt -x), Pot Pot-ur bt U2 (X1,- + Pgt -t (X1 (x1,--))
(+.t)
+
(+1))
(+.1)
dw (t)T
C1
P pa (X1,-,- +
(+-))
(X1,-,- + P -
(+.1)
vz kt (X1 (x1,-
dt
q
Thus,
(X1,tt))))) =Ω Ω
dw T 01
dw (t)
dw dw (f)
U2
+
dt
dt
dt
dt
000 δ (n
(by Lemma F1
For conflicting samples Xk, we have
(11+
dw T (+.t)
P p + (t1))
patter + (g--makoo)
dt
dt
dw T !TT
12 (+.t)
Pgtui (o-c + -t- + (ti))
(-1t)
p Pet-oakt-uy
dt
dw (t) Ci (tul)
phttl + (4.1)
P + (t.)) + -
dt
dw (
(1)
(e + (-1t) Pooeiikter Poe--u
(-10)
(-1()
dt
dw (2)
(1)
(-11 (P kt P kt -22 + - U2 12
(-1))
P
dt
du (
(
k + P -P kt- + (-,1)
(-,6)
dl
(T
dw Vi
Ω
Ci
+ Ō
w
dt
dt
22g - = (wg(+)
(1)
(by Lemma F2
Thus,
(E(((uf(((
1 (d(F(()
1 (x
d/
2
dl
2
dl
1 -((00000))))))))
(X
+ (
(x dl
dl
55

---

Published as a conference paper at ICLR 2025
- (1o + (
(((0)))
which implies that
(x]]) =o (1-
1-72 +
G HANDLING THE FINITE-WIDTH DISCRETIZATION
In this section, we describe how to incorporate the finite-width discretization. As we mentioned in the beginning of Section 5 the gradient flow for the finite width transformer can be seen as the gradient flow for infinite-width transformer model with perturbation. Thus, we need to calculate the deviation of the finite-width gradient flow from the infinite-width gradient flow. Since we only train the neuron weights in stage and stage 3, we only need to consider how the finite-width discretization will affect stage 2.
First, at initialization, by the standard concentration bound, we have the following:
Lemma G.1. With probability at least 1 over the randomness of the initialization of WK. K, Wq. for any U.V CUV.
( (0) K H, (001))) sw2
4
8
4
8
log
log
8
771
711
(W u, x u) S w²
4
8
4
8
log
+
log
8
m
(W!o) n, w((0)) b) S w²
(V
4
8
4
8
log
+
log
m
8
and any μ E CUV,
(w K 10 = =m
1
log
log
m
8
5
[W 1ll2 w²
log
log
m
8
m
8
Proof. Since W(O)) W ~ NO0, w21) 1), we apply Lemma
H.2 to get the corresponding bound.
Next, recall that from Section 5.2 the score and self-score variables evolve according to the dynam- ical system
vec(f T (OT w:( d
K+ o+ 1) (uneccv))
vec(fu T w-(t] w-10)
K+ Q4 1) (uvecuv))
vec(( (w()) will u) (IIEECv))
= A
vec((n w will)
K4 u) (Ivecuv)
di
vec {u W wit)
K+ Kt
(OT
@+ 11 (u.eev)))
veo([N] w (e)T witt) Q+
R+
Q+ [[NNE)uv)
Q+
stage 2 lasts Olog(( 1/w)) time. Then, we can bound the effect of the perturbation to the dynamical system as follows. Consider the dynamical system
for some matrix (1) with || 0(1). By Lemma B.3
dx((t) A(1)x((), x(0) = 20 dt
56

---

Published as a conference paper at ICLR 2025
and the perturbed dynamical system
dx. (t) - (A(t) + (A())),.)), x,(0) = Do +
dt
where and (A(() are the perturbation introduced by the finite-width discretization. Notice that we
have |s)A(s)] S c](rr(()) for some constant for all To]. Define òr(() x. (1) x(t).
We have
dōx(t) = A(t)6z(t) + (A())x.(t)
dt
This implies that
[1x(())] 1 ||A(s)6z(s)+
+ + (s)
! [
Applying Grönwall's inequality, we have
[[xx(t)]| [|6x(0)]| exp (L. [|A(s)]
poly(1/w)/
since ||A()]| S 0(1) and ||xr(()]| S 0(1) and I O(log(1/w)) Thus, picking m poly(1/w) can make the perturbation to the gradient flow system most (w))
Proposition G.2. For < T2.
(t)l (t)ll 0(1).
Proof. We are going show that in the -nftth case, ||x(0]] 0(1) for S Th. Later, are going to show that [|x((t]] < 1 and by triangle inequality, we have [æx - 0(1. The key observation is to use Lemma E.1 which provides upper bound the score variables after stage However, there is an issue: x will also contain self-score variables whereas Lemma E. only provides upper bound the score variables.
To solve this issue, we refine our analysis for the infinite-width case. We are going to wifd C within (multiplicative) constant factor of K+ Q+ v1- First of all, direct consequence of Lemma E. is that it also provides up- bound of the growth of the 3 dominating variables and thus Lemma E.2 can be modified as
show that of
(-( @+
w:ll2 01 and cf w K+
of (tT w(t))
TT l(( K+
01 W w: e
d
v1
0
I
1
K+
dt
c W
K+
w
k44
=
-11
1
0
0
( (T
K+
w w
k_q
w
10T
(t
0
0
(OT
w(())
Q+
Q+
@Q+"1 tn
Notice that there is a self-balancing property of this dynamical system: there exists constant w((
C. 0 such that if of w
C cf WK+
K+ U1, then d ee G (t) w-(t) K+ @441 v1 >
w-())
d de.. BcTw(() (( c1 Thus, if those variables already within constant factor of each other, they K+ K4C1-
will maintain this relationship.
Then, the base case be established by the existence of time Ti,s such that Ti.s Ti O(1) and of (TT.)T (T.Tss) K+ the proof of i. w-(e) w(t) Q+ 14+ and ef
Q+ ci (T.s)
w (Ti.s)
Ca, of (Ta) (T..s) Q+ Q+ = 2ww2). This finishes
K+
K
OTT (
are within constant factor of c W w(t K4 @
K+ K
during stage 2.
apply Grönwall's inequality as before and get [|āz(t)]| < (oly(1/w)/ ///mm / < as long as m 2 poly(1/w). This
Finally, for the finite-wide width case, long as || < for t S Th,
argument is proved in way similar Section 1.3 in Tao (2006).
n
57

---

Published as a conference paper at ICLR 2025
H PROBABILITY
Lemma H.1 (Bernstein's 's inequality for bounded random variables). Assume Zi.... nn are n i...d random variables with [ and Zii M for all i [n] almost surely. Let Z Dim Zi. Then, for all t > O,
12/2
1
1)
P > 1 S exp
I exp
min
Th 2j=1
Z[
+ Mt/3
2 ====
R
R
2M
which implies with probability least
6,
1 + 2M log
2
E [2] log 8
i.i.d. N(0, Im mm/,
Lemma H.2. For wi, E II with wi, w2
2
|| 121112- - -1/2 P
log +
log
2
8
0
2
P
(101, W2)
2
log 8
+
log
<
6
g
Proof. We have
m
E
W1
E
Wi.6
Notice that wf, is a sub-Gamma random variable with parameters inequality,
Thus, by Bernstein's 's
4
4
I
|| l1
E
log
log
5 6
m
m
Next,
m
B((w w2)l =
E
W2,
22,
=
By Bernstein's inequality,
4
2
4
2
IP
(wn w2)l 2
log
+
log
2
m
8
m
58


Reviews:
[{"review_id": "xmVz013H9V", "reviewer": "Reviewer_fKbv", "summary": "This paper investigates the dynamics of transformers in learning linear classification tasks where input tokens have hidden correspondence structures. The authors examine a two-headed transformer with ReLU neurons on a specific task involving two classes with hidden linear relationships. They observe a three-stage learning process: (1) neuron alignment with classification signals (2) (2) attention feature learning where features evolve into a separable state and (3) training the neurons again to the final convergence. Baaed on the empirical observation, they propose a three-stage training algorithm and derive the main theorem based on the algorithm. \nTheir theoretical analysis, using gradient flow dynamics, shows how attention mechanisms enable the model to leverage hidden correspondences to solve classification tasks effectively. The study also considers settings with more than two mixtures, highlighting challenges in analysis yet demonstrating the model's empirical success.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "**Theoretical Contribution**: The paper provides a rigorous theoretical analysis of transformers’ training dynamics on structured linear classification tasks, specifically using gradient flow dynamics. This theoretical foundation contributes to a deeper understanding of how attention mechanisms in transformers learn hidden relationships, a topic that is not yet fully understood in the community.\n\n**Algorithmic Insights**: The proposed three-stage training algorithm mirrors the empirical learning dynamics, providing a potentially practical framework for initializing and training transformers on tasks with similar hidden structures. This approach could be valuable for applications requiring attention to specific feature-alignment steps.\n\n**Novel Analysis of Attention Dynamics**: The paper effectively addresses the softmax attention mechanism’s behavior within the training dynamics. The authors introduce techniques to simplify and analyze the system, such as reducing the attention's complex nonlinear dynamics to a manageable subsystem, which may serve as a foundation for future analyses in related settings.", "weaknesses": "- This paper only consider learning two-mixture of binary classification, and the data distribution is of size $d\\times 3$. The main theorem is based on Algorithm 1, an algorithm designed from empirical observations on a synthetic dataset but not tested on any real-world datasets. So, I question whether analyzing Algorithm 1 and this specific choice of dataset is really helpful. I suggest the authors to apply Algorithm 1 on some real-world dataset, and shows that it works. (do not need to be better than standard training). \n\n- Although this paper's title is \"two-mixture\", there should be more experiments and discussions on $K\\geq 3$ than current Appendix A.3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WPCz0EJocM", "reviewer": "Reviewer_Kkd8", "summary": "This paper studies the training dynamics of a two-headed transformer using a three-stage, layer-wise training algorithm. It further characterizes the feature learning process and captures the attention mechanism involved in solving the task.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper presents a rigorous theoretical framework for understanding the training dynamics of non-linear Transformer architectures.\n\nThe most interesting finding, in my view, is how the authors analyze the progressive emergence of connections across diverse data types at various stages, offering a comprehensive perspective on the layered learning process within Transformers. Specifically, the study demonstrates that weights in distinct components of the model are updated in different stages, indicating an organized, sequential approach to learning. This stage-wise learning strategy introduces a novel perspective on the interplay between weight adaptation and data distribution, highlighting how the architecture distinguishes and captures distinct data characteristics over time. Additionally, the paper provides a theoretical characterization to substantiate its findings, anchoring the analysis in formal proofs.", "weaknesses": "(1) The description of the experimental setup is missing, but I assume the results are all based on synthetic data. Therefore, justification with real data is necessary.\n\n\n(2) The theoretical insights for guiding the training and deployment of Transformers are unclear in practical applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UAHkkd2whU", "reviewer": "Reviewer_qpZd", "summary": "This paper analyzes training dynamics of transformer estimated via gradient flow under the mixture of linear classification setting. Based on the empirical findings, they construct their own three-stage training procedures, which essentially give the same training dynamics with actual gradient descent simultaneously train all the parameters in Transformers.", "questions": "A.\tI had trouble understanding the setting of mixture of linear classification. Authors only spend one paragraph in the 2nd page discussing about the model. But I think it should be elaborated more in details to motivate the readers. \n\nB.\tI might have missed the part while reading the paper, but I wonder how training dynamics discovered in the paper corresponds to the actual reading procedures of the human beings? For instance, using the example in the paper: “Bob is watching TV in the living room. Where is Bob?” How do people catch the correspondence between Where and Living room and how does this correspond to the training dynamics in the paper? I wonder if authors have any thoughts on this. \n\nC.\tI think this is important question: is the training dynamics in the paper also observable in the vanilla transformer case? The authors consider two-headed symmetric transformers, but in real practice vanilla transformer is commonly used. If this is the case, I think this paper’s contribution is significant. \n\nD.\tDo authors have any insights on the multiple layer cases?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides rigorous theoretical analysis on the training dynamics of transformer which is a highly non-trivial task.", "weaknesses": "Please see the below questions:", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1QKTwDYwM4", "reviewer": "Reviewer_7eXW", "summary": "The paper analyzes training dynamics of a single attention layer two-headed transformer on a specific task. The task -- a two-mixture of linear classification -- possesses a hidden correspondence structure between the feature vectors in the sequence. Each sample contains an indicator vector that dictates, sign of which of the other feature vectors in the sequence is important. Thus, to efficiently solve the task, a communication between the class indicator and the rest of the feature vectors is necessary. The authors first observe the training dynamics empirically, specifying three distinct training phases. Based on these observation they define a training algorithm similar to standard full-batch gradient flow, for which they prove that the transformer not only learns the task with perfect accuracy, but, more importantly, learns to effectively utilize the hidden structure in a human interpretable way. Finally, the authors discuss the extension to more than just two-mixture of linear classification, constructing a transformer architecture that solves the task, but showing that a gradient-based analysis would need to utilize more refined proof techniques.", "questions": "- Q1: Have you tried to measure this phenomenon and the training dynamics for any of the setting extensions I mentioned in W2? \n\n- Q2: In Figure 1, we see that the neuron alignment grows with $c_k$ immediately from the beginning. This seems to be in contradiction with line 367, especially since I believe the weight initialization in experiments was the same as in the theory. What makes this difference? \n\n- Q3: In figure 1c, how would the curves for $c_1$ being query look like? Can you provide a figure with these included? \n\n- Q4: Where do the slight asymmetries in Theorem 4.1 come from when you consider the averaged initial value model? \n\n**Summary:** I consider this paper to be a very instructive, novel, theoretical contribution. It allows us to understand transformer training dynamics in a human-interpretable way that aligns with our natural understanding of what attention should be learning to do. This highly precise description of the role of attention in simple transformers makes the paper a strong contribution. Therefore, I recommend accepting the paper. However, my high score is conditioned on the authors addressing my concerns in the discussion period.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "What I particularly like about the paper is the very well chosen combination of data distribution, architecture and proof assumptions. This combination is perhaps a minimal example on which it is possible to demonstrate this phenomenon of transformers learning to utilize the hidden communication between tokens at this level of complexity. On the other hand, this setting is perhaps a maximal example under which the proof still remains practically doable and comprehensible. Therefore I find the setting of the paper very suitable. \n\nUnlike many other papers studying transformers, the advantage of this work is that the studied task and the provided results are simple enough that they provide human-interpretable understanding of the transformers' behavior, that agrees with our expectations -- the attention layers are crucial to provide the necessary communication between tokens and they can learn to do so. Further, the importance of the use of multiple heads for certain tasks is instructively demonstrated.\n\nThe authors also do a very good job extracting the essence of otherwise long and tedious proof into the main body and make it understandable for a wide audience. I think the presentation improves reader's own understanding of the inner workings of transformers and allows the reader to re-use some of the techniques in their own transformer research. The presented analysis and the core proof idea makes sense and nicely captures the important aspects of the optimization.", "weaknesses": "- W1: While I appreciate the setting of the paper in general, there are a few questionable decisions when it comes to architecture. In particular, in the main part of the paper, it is a bit unusual that the value matrix basically maps to 1D outputs and the representations are averaged even before applying ReLU. I would expect averaging only after the application of the ReLU and also like to see what would happen if more than 1D value matrices would be used. However, bigger issue comes with Lemma 6.1. In particular, the order in which the operations are made in the transformer defined in the proof is so unusual I am not sure it can be interpreted as transformer anymore. The summation over $l_2$ should definitely come as the inner-most, otherwise we cannot even talk about attention. Then should come the summation over $k$, only after that should come the application of $\\sigma$ and the summation over $l_1$ should be the outer-most. I would ask the authors to redo the statement and proof of Lemma 6.1 so that the summation over $l_2$ is the inner-most or otherwise justify their architectural choice. \n\n- W2: Although the presented phenomenon is intuitive, it is unclear *if* and *how* does it generalize to more general settings. In particular to multi-mixture of linear regression, where I would expect (almost) the same thing happen. It is unclear why the symmetry is broken and therefore it is also unclear whether the described phenomenon even holds the same way in multi-class setting. The same question holds for another natural extensions of the considered setting -- multi-class classification, presence of noise, non-orthogonality of features or class imbalance and more natural weight initialization. I wouldn't mind that the paper only studies the simplest setting if there was evidence that the knowledge obtained can be transferred to the more general settings but in this case it is not clear whether the same mechanism even works in extended settings.\n\n- W3: Some aspects of the presentation should be improved. To provide a couple of examples: \n\n- *It could be clearer if authors re-named classes to something else as they are confused with the classification classes.*\n- *The assumption of averaged initial value model is not discussed enough. In particular it is not clear where any asymmetry at all comes from if we have this assumption.*\n- *The proof in appendix should be more verbose and make some more explicit statements as intermediate steps. For instance, the separability of the consistent samples after phase 2 is only proved implicitly and is never explicitly stated in the proof. This makes it hard to make look-ups of specific proof parts mentioned in the main body, because the proof itself is an ocean of infinite formulas.*\n- *The use of $g$ is somewhat peculiar. I suspect the authors to either make a few typos when using $g$ or not to define it properly.*\n\n**Typos:**\n\n- The property 4 in Definition 2.1 is inconsistent with the verbal description in the rest of the paper as this way, you would leave most of the positions actually empty (zeros).\n\n- line 170: You should write $\\nu$ instead of $X^{(i)}$.\n\n- In lemma 5.2, the use of $g$ is somewhat weird. Please double-check what you meant there. The same holds for lines 427, 430\n\n- Line 408: based on my counting the numbers should be 16, 10, 10. Correct me if I am wrong. \n \n- In line 60 you talk about class identifier getting all the signal from classification features, but later in your theory you consider the class identifier's role as the key to be crucial. I think it better to rewrite this sentence so that the roles are exchanged.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)