Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis
Authors: Hongru Yang
Bhavya Kailkhura
Zhangyang Wang
Yingbin Liang
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformers
gradient flow dynamics
implicit bias
Primary area: learning_theory

Abstract: Understanding the training dynamics of transformers is important to explain the impressive capabilities behind large language models. 
In this work, we study the dynamics of training a shallow transformer on a task of recognizing co-occurrence of two designated words. In the literature of studying training dynamics of transformers, several simplifications are commonly adopted such as weight reparameterization, attention linearization, special initialization, and lazy regime. In contrast, we analyze the gradient flow dynamics of simultaneously training three attention matrices and a linear MLP layer from random initialization, and provide a framework of analyzing such dynamics via a coupled dynamical system. We establish near minimum loss and characterize the attention model after training. We discover that gradient flow serves as an inherent mechanism that naturally divide the training process into two phases. In Phase 1, the linear MLP quickly aligns with the two target signals for correct classification, whereas the softmax attention remains almost unchanged. In Phase 2, the attention matrices and the MLP evolve jointly to enlarge the classification margin and reduce the loss to a near minimum value. Technically, we prove a novel property of the gradient flow, termed \textit{automatic balancing of gradients}, which enables the loss values of different samples to decrease almost at the same rate and further facilitates the proof of near minimum training loss. We also conduct experiments to verify our theoretical results.

Training Dynamies of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis
Hongru Yang"
Bhavya Kailkhura Lawrence Livermore National Laboratory kailkhura1011n1 gov
The University of Texas at Austin & Princeton University hy6385@utexas. edu
Zhangyang Wang The University of Texas at Austin atlaswang@utexas.cm edu
Yingbin Liang
The Ohio State University 1iang. 8898osu. edu
Abstract
Understanding the training dynamics of transformers is important to explain the impressive capabilities behind large language models. this work, we study the dynamics of training shallow transformer of two designated words. In the literature of studying training dynamics of trans- commonly adopted such weight reparameter-
task of recognizing co-occurrence
formers, several simplifications ization, attention linearization, special initialization, and lazy regime. In contrast, we analyze the gradient flow dynamics of simultaneously training three attention matrices and a linear MLP layer from random initialization, and provide a frame- work of analyzing such dynamics via coupled dynamical system. We establish near minimum loss and characterize the attention model after training. We discover that gradient flow serves an inherent mechanism that naturally divide the training process into two phases. In Phase I, the linear MLP quickly aligns with the two tar- get signals for correct classification, whereas the softmax attention remains almost unchanged. In Phase 2, the attention matrices and the MLP evolve jointly to enlarge the classification margin and reduce the loss to near minimum value. Technically, we prove novel property of the gradient flow, termed automatic balancing of grudients, which enables loss values different samples to decrease almost the same rate and further facilitates the proof of near minimum training loss. We also conduct experiments to verify our theoretical results.
Introduction
transformers have become dominat- ing backbone architecture in many machine learning applications such as computer vision DBK+20 [LC+21 and natural language processing [DCLTT8]. Nowadays, ChatGPT and GPT-4 [Ope23] have demonstrated astonishing abilities in many areas such as language under- standing. mathematics and coding, which have sparked artificial general intelligence (BCE*2] 23). In the meantime, there has been burgeoning development of large language models (LLMs)
Ever since the invention of self-attention VSP+17 17]
TLI 23 MH23 1DF++23 well multi-modal models [Tea23)
Despite the huge empirical success, theoretical understanding of why pre-trained language model can possess such impressive performance has been significantly lagging behind. Some previous
Work done while doing a internship at Lawrence Livermore National Laboratory and visiting Princeton University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

efforts have been made in understanding the capacity and representational power of transformers LAG [EGKZ22]
ZPGA23 SSH22 23] However, most results of this type of works are existential and rely on manual construction of the weights. It is unclear whether the constructed weights are the actual solutions after training transformers. In order to understand the mechanism behind those pre-trained language models, line of studies have aimed to open the black box of optimization via studying training dynamics of transformers and explaining why transformers can be trained to perform well [SL22 LWLC22] [WCD23 TWZ [LR23] TLZ023 [LTO23] [FF224] HCL23] However, those previous works often relied on various simplifications in their analysis such as weight reparameterization, attention linearization, special initialization, lazy regime, etc. One goal of this paper is to take a further step to demystify the training dynamics of transformers and consider more practical training setup, thus better capturing the actual training process.
Our study of transformers' training dynamies will focus on basic problem of recognizing co- occurrence of words under a binary classification setup, which is an important ability of LLMs to perform many tasks correctly in natural language processing (NLP). For example, the classical n-gram model [M99 Dam 8) predicts the next word based co-occurrence of multiple words. Consider the following scenario: if the task for the language model is to read a paragraph describing children and then answer some questions, say, "Is Bob eating banana?". In order to answer the question correctly, the model must be able to detect the co-occurrenc' of the two words "Bob" and "banana" in the paragraph. Motivated by this, we study the problem of detecting co-occurrence of two target words via the model of a one-layer transformer with a self oeffaatennionn module followed by a linear multi-layer layer perceptrons (MLP) layer. Our goal is to characterize the dynamics of the training process via the gradient flow analysis, thus providing a theory to explain how transformers can be trained perform well.
Our contribution is summarized below:
We study the gradient flow dynamics of detecting word co-occurrence. The training starts with random initialization and then simultaneously updates four weight matrices (including key, query, and value matrices and a linear MLP) in the transformer architecture via gradient flow. We show that gradient flow can achieve small loss although the loss function is highly nonconvex. We further characterize the explicit form of attention matrices after training, which captures the strong positive correlation between the two target signals and strong negative correlation between one target signal and the common token, both leading to large classification margin.
We characterize the training process into two phases. In Phase 1 (alignment of MLP for correct classification), we show that the linear MLP of the transformer quickly aligns with the two target word tokens whereas all other variables in the dynamical system stay almost unchanged from their initialization values. All training samples are correctly classified at the end of Phase I, but the loss value is still large due to small classification margin. In Phase 2 (evolution of attention and MLP for large classification margin), along with the continual evolution of MLP, correct classification by MLP also encourages the gradients of attention matrices to learn. Specifically, the softmax probability increases if the key and query tokens correspond to the two target words, and the value transform of the two words becomes more positively correlated, both leading to enlarge the classification margin. Thus, the training and test loss values both are driven down to nearly zero
Technically, our proof techniques do not rely on several commonly used assumptions in the literature such as weight reparameterization, attention linearization, special initialization, lazy regime, etc. Our main idea is treat the problem as a coupled dynamical system with six different types of dynamic variables, for which we provide an articulated analysis on the gradient flow dynamics. In particular, we prove a novel property of the gradient flow, termed automatic balancing of gradients, which shows that the ratio of several important gradients will evolve closely within the same range during training. This enables us to show that the losses of all training samples can decrease almost at the same rate, and is also a key component in proving the near minimum training loss as well as analyzing the changes of softmax.
1.1 Related Work
Transformer representational power. Several previous works have studied the expressiveness of transformers. One line of work was from universal approximation perspective and thus provided the existential results (YBR+ + 19, [CM22 [BM21 ZPGA23 L2 23)) As a separate
view, [EGKZ22] showed that a single attention head can represent sparse function over the input

---

sequence with sample complexity much smaller than the context length. [ZZYW33 studied the approximation and generalization performance of transformers in in-context learning, [SH223
proved that transformers can represent certain functions more efficiently than MI.s.
Training transformers. Various settings of training transformers have been studied recently.
[WXM21] studied the impact of head and prompt tuning of transformer on the downstream learning tasks. USL22] proved that transformers can learn spatial structures. [LWLC22] studied how shallow transformer learns a dataset with both label-relevant and label-irelevant tokens. [TWCD23 studied next-token prediction problem and showed that self--atention behaves like discriminating scanning algorithm. [LLR23] analyzed a layer-wise optimization scheme on how transformers learn
topic structures. [TLZ023 TLTO23] VDT24] studied a setting where transformers can learn SVM solution. [LWM 23 provided analysis of training graph transformers for node classification
tasks. Thr24 studied the implicit bias in the next-token prediction problem. For in-context linear regression, [VONR+23] constructed transformer weights to solve this task and showed empirically that this is similar to what the transformer learned by gradient descent, [ACDS24] proved that the critical points of the training objective of linear transformers implement a pre-conditioned gradient descent, FFB244 provided the training dynamics of linear attention models, [HCL23] characterized the training dynamics of softmax transformers, and [CSWY24] studied multi-task linear regression problem with a multi-headed softmax transformer. Further, [LWL* 24) focused on nonlinear self- attention and nonlinear MLP over classification tasks in in-context learning. [WLCC23] proved the convergence of transformers via neural tangent kernel. [ND224] showed that two-layer transformers can learn causal structure via gradient descent. [CL24] developed algorithms for provably learning multi-head attention layer. (WW studied how transformers learn reature-position correlation.
This paper studies a different problem of detecting co-occurrence of words via transformers. Such setting has not been considered in the literature. More importantly, the previous studies of training dynamics of transformers have adopted various assumptions/ simimpttionssimliiicatttoossssss such weight reparame- terization, special initialization, attention linearization, lazy regime, etc. In contrast, our analysis here based on gradient flow does not rely on those simplifications, which can be of independent interest for studying transformers in other settings.
Problem Setting
Notations. For a vector Rd, diag(v) to denote diagonal matrix with being the diagonal entries. When we subtract the vector v by a scalar a, we subtract each entry of u by i.e., E IR and (v a); t; a. We use 92, θ Ö to hide polylogarithmic factors.
2.1 Data Model
Definition 2.1 (Data distribution). Given a set of orthonormal vectors {ur]a= as word embedding,
let Pl, 22 RRd be two target signals whose co-occurrence needs to be detected by the model, and let μs E IRd be common token vector. A data entry (X, y) E RdxL {+1}, where X [2,, x2, consists of tokens, is generated by the distribution D as follows:
Uniformly randomly select an index ig E [L] and set Dig = p3-
Then, one of the following cases occurs:
With probability 1/2, set 1 and uniformly randomly select two indices in * is E [4] {is}
and set Fi, = ,4, Tig = μ22 For i E [L] {i,, i2, is), set Uniform| (1s )(=)) With probability 1/6, set = -1 and uniformly randomly select one index i E [L] {ia} and
set 24, #1- For E [L] \ {6s, {1}, we set x; = Uniform({ (
With probabiliy 1 /G, set = -1 and uniformly randomly select one index in E [T.] {iz) and
set xi, = u2- For E [L] {ia, i2), we set x = Uniform( (
With probubility 1 /6, set y = -1 For all i E [Z] (ia). we set 2; Uniform((us)d
In summary, there are types of data: (1) both μ4, μا appear, (2) only ly appears, (3) only μa appears, and (4) neither f4 nor μl2 appears. We denote the set of indices of the above 4 different types of data by J1, T2, Ta, 14 S [n]. We further define = {4+){= For simplicity, our data distribution assumes P1, 42, H3 appear only once in data entry. The occurrence probability of each type of
3

---

data is chosen in the above way to make the distribution label-balanced. We assume there is fixed set of orthonormal vectors as word embedding, which is analogous to the one-hot embedding of set of vocabularies. Furthermore, in our daily language, there are some words appearing in almost every sentence such as "a" and "the". Thus, to model those words, we include a common token in every data entry. Finally, notice that if we ignore the common token and random tokens, the data distribution simplifies to logical AND problem.
Recognizing co-occurrence of words is an important ability for language models to perform many NLP tasks correctly. Consider the example of a language model first reading paragraph describing a children and then answering the question "Is Bob eating a banana?" " If the description is "Bob is watching a television while eating a banana", then the model should answer "Yes". the description is " "Bob is playing computer games", " then the model should answer "No". Thus, the model needs to recognize the co-occurrence of "Bob" and "banana".
Remark 2.2.
For simplicity of our analysis, we make the following assumption on our training data set.
Assumption 2.3. The training set satisfies: (i) [i and [Ll Lal LeI :/r and (ii) for
all i1, le E [n], 4.4 l2 E [L], if (((1), x((2) & {11: Ps) then x (és) + x (82) ie., all irrelevant words are different.
The first assumption can be approximately satisfied with high probability given the total number n of samples large enough. Such an assumption can be removed by applying the standard concentration theorems. The second assumption implicitly assumes nL S d. If the irrelevant words are uniformly sampled from a large entire vocabulary, then each irrelevant word appears only very few times in the training set. Thus, letting irrelevant words appear only once in the entire training set is a reasonable way to simplify our analysis.
2.2 Transformer Architecture and Training
Consider a training set ((x(1), 1))) with n training samples. Each data point X E IR contains L tokens, i.e., X = [ (.) ....,,"%). (.). We consider the transformer model with a self-attention module followed by linear MLP:
my
X W W W vm
WX: WK. Wo
ΣΣe aj
10 Wv. Softmax
(I)
where the query matrix We Rmxd the key matrix WK E g the value matrix Wy E gm
the hidden-layer MLP weights W E Im Xm (with w/ being the j-th row of W), and the output-layer weights of the MLP R We define the linear MLP function of the transformer to be G(u) Ene arwy W. We now introduce some shorthand notations K = Xx Q Wq V kz). We further extend this shorthand to qi and vy. We also define functions k(μ)) = WKH, q(H) = WoH, v(p) = Wv We introduce the
WyX and let k WK... Notice that K = [k1, .e.................... .kz].
shorthand for the score vector 8j := X W. and the attention vector pt == Softmax(si). For the v/m
attention vector, if [4, (1),, let (ii, (i, v) be the indices such that x-()) 1(1,1) 14, ((E.))
x(0 =
and we define p Py-ukke- := 10 Pg- == Softmax (X W Woa
1(1.2)
Initialization. We initialize aj iid. Uniform(t!) and the value of a is fixed during training.
The trainable parameters are [W, Wv, WK, Wol. We initialize [W, Wv. WK. Wol by W s.dd (0 of) and (Wy) (Wi).,, (WQ)., s.idd N/O, az).
Training, We adopt the cross-entropy loss ((x) = log(1 + exp( The gradient of the cross- entropy loss is given by ('(x) [+exp(z] and we define g(x) I+exp(z) The model is trained by gradient flow minimize the following empirical loss:
i(, Wv. WK. Wo) = 1 Σ
1 (x,FF(X() W. WyW. WK, WQ)).
(2)
i=1

---

Similarly, we define the generalization loss : B(xxa)>> (X)). We introduce the parameter condition that we take throughout the entire analysis proofs.
Condition 1. We make the the following parameter choices in our analysis:
The embedding dimension and network width satisfy m 2 [[max L2}) and my 2 9(1).
The network weight initialization variance satisfies do @(vLm) and 01 @(vm))
The number of training samples and tokens satisfy n 2 (22) and L 2 ((1).
The failure probability satisfies 1/8 S poly(m).
Main Results
Challenges. The essential goal is to derive the gradient flow update for each weight matrix (see the gradient expressions for all weight matrices in Appendix [B] However, directly analyzing the dynamics of those weight matrices is extremely challenging. because: (i) keeping track of how the column and spaces of each weight matrix change during training is difficult; and (ii) all attention and MLP weight matrices are affecting each other, leading to highly coupled dynamics.
Our General Idea. To overcome the above challenges, we first note that rather than tracking W, Wv, WK, Wi directly, is sufficient to analyzing their impact on inputs, i..., X We W and a' T WWy X, which are sufficient to compute F(X; W, Wv, WK, Wq). Based on this observation, we formulate two differential equations to keep track of w w
V j and TTW .Tw(O w Q (with
respect to f) for all μl, V E {ui (See Equation 6 and Equation in Appendix
C We further
include additional equations to keep track of (1
w v W- W H4, V T W
(:(t) K and
K
T W W ra to complete the system. Intuitively, the additional equations keep track of the shape of the neurons and the word embedding after Wy. WK. Wo transform. Although the dynamical system does not directly track the softmax, the softmax probability can be calculated via the scores of C Then the training dynamics
Tw w j.. The full dynamical system is presented in Appendix can be characterized by analyzing these differential equations (see a proof outline in Section
In the next two theorems, we present our characterization the training process into two phases.
Theorem 3.1 (Phase 1). With probability at least 1 - 8 over the randomness of weight initialization, there exists time Ti = Õ(1 /m) such that
The linear MLP functions satisfy: ((T) 2 2(1), (P2) 2 2(1). (CN) S -Q(())
All training samples are correctly classified: y, ((T)) for all E [m].
For t E [1) all dynamical variables (w / u ww w n. v TW/ 1.
v TT w K 1s, and v Tw( (T) wlo
are close to their initialization values.
Training loss is still large: ((T)) O(1).
In Theorem 3.1 item implies that in short time, the linear MLP function (-)) positively aligns with the two target signals and μ2, but negatively aligns with the common token μra This further guarantees item 2 of Theorem 3.1 that all training samples are classified correctly. Further, item 3 of Theorem 3.1 indicates that the attention matrices are still close to their initialization values, and hence have not started learn any knowledge yet. This results in item of Theorem 3.1 which shows that the training loss is still large.
Theorem 3.2 (Phase 2). With probability least 1 over the randomness of weight initialization. there exists time range (Ti, T2) with T2 = poly(m) such that for all E (Ti,T) T2)
12 WOT wrll) i11 and Ww w(O (0) n2 increase, whereas
w0O Pg WTT (9T
in and
Ng - w-O #2 decrease.
5

---

TWV T pi V
W P2 increases, whereas 1 W W Mi and p1 W/ w M3 decrease. Linear MLP fumctions satisfy: (pz) (11), (12) 2 2(1), (ps) R(1).
(mn) + (12) (us) 2 2(1), ((n) + (ps) -Q(1) and (122)
((sss) -2(1).
In Theorem 3.2 item indicates that, during Phase 2, gradient flow drives the self-attention module to weigh more between the two target signals μl and μ22, and to weigh less between one of these signals and the common token M3. Item 2 indicates that gradient flow drives the value matrix Wy to positively align the two target signals μ1 and μ2, but negatively align one target signal (F1 or (22 with the common token M33. Further, the last two items indicate that the MLP continue to classify correctly and further enlarge the classification margin. Hence, all items in Theorem 3.2 collectivvy indicate that attention and MLP evolve jointly to enlarge the classification margin and hence drive the loss value to decrease in Phase 2.
Theorem 3.3 (Near Minimum Training Loss and Attention). With probability at least 1 8, there exists a time T = ((poly(m)) such that
The training and generalization losses satisfy î((T)) I 1/poly(m) and I(TT*) S 1 1/poly(m).
The attention matrices satisfies:
WK
T "('')
=
(U)T W(O) + EEniisela
(2*)) 11,82 Mig Megi
where ((T) , 211
(T))
(T*)
(T))
= 0
o m Vma mmny
and
c(TT") 11.02
3,1
3,2
(mmlafmm
δ abm mymLoffmm
oyym fmmm
if one of i1, b2 E [d] [3].
Theorem 3.3 indicates that both training and test losses converge nearly to zero as long as the embedding dimension m is sufficiently large, because both the attention and MLP matrices are trained towards enlarging the classification margin in Phase 2. Theorem 3.3 also provides the explicit form of the attention matrix in Equation 3 in which the second term captures the learned information
of the self-attention module. It can be seen that the large coefficients cf,"" and Ch4T") capture strong
((T") and ((T)) encourages strong negative coupling of one target signal μrl or #2 and the common token μ3. All these attention terms contribute to enlarge correct classification margin. Further, the coefficients between all other random tokens are order-level smaller and hence do not corrupt the correct classification.
coupling of the two target signals μ4 and μl, and the large negative coefficients C,11
Synthetic Experiment: We next verify our theory and the two-phase characterization of the training process via synthetic experiments (see the experiment setup in Appendix A
score
Training Loss
lass
0.7
lows_ tow__to1..0.ppp
(3)
0.6
lnws, _13_
0.05
(0 10-p1.A- p3)
0.5
(0-a1
0.4
(g-p2)A-ph (3)
8.00
(0 (0=p2,A= p4)
0.3
0.2
Phose
Phose
0.1
Phase
Phase
-0.30
0.0
250
500
750
1000
1250
1500
1750
2000
250
500
750
1000
1250
1500
1750
2000
(a) Attention score correlation
(b) Training loss
Figure I: Synthetic experiments with illustration of two training phases. The detailed experiment
setup can be found in Appendix A
Figure (a) shows how the attention score correlation n1 w (O)T (+)
My evolves during the training.
It is clear that these scores do not change significantly in Phase 1, verifying Theorem 3.1 In Phase
2, the score correlation between two target signals μ41 and μ22 increases, and the score between one
target signal and the common token decreases, verifying Theorem 3.2

---

Figure I (b) plots how the training loss changes during the two phases of training. The blue curve (indexed by 'loss') represents the overall training loss of all samples. The other curves correspond to the training loss of four types of samples as indicated in the legend. In Phase 1, the training loss for samples with both target signals (i.e., orange curve) decreases because the linear MLP layer aligns with the target signals (verifying Lemma 4.1 in Section 4.1). The training loss for samples with one target signal and the common token (i.c., green or red curves) first increases because the linear MLP layer initially has not aligned negatively enough with the common token yet (as captured by Lemma 4.2 in Section and then decreases in the later stage of Phase I when the MLP layer aligns negatively with the common token (as captured by Lemma 4.3 in Section 4.1). All loss functions decrease in Phase 2 because all attention matrices and linear MLP jointly enlarge the classification margin, verifying Theorem 3.2
4 Proof Outline: Two-phase Gradient Flow Analysis
4.1 Phase 1: Alignment of Linear MLP for Correct Classification
In Phase 1, the linear MLP quickly aligns with the two target word tokens while all attention matrices stay roughly unchanged from their initialization values. We show that the linear MLP functions (() (1))))))) ((()))))))) (GO become sufficiently large (larger than some constant threshold) so that all training samples are correctly classified at the end of Phase 1.
We first analyze the dynamical system at the initialization. In particular, the following lemma shows that at the initialization, the linear MIP layer receives a sufficiently large gradient from the two target signals, and hence samples with the two target signals will be classified correctly as co-occurrence soon afte the training starts.
Lemma 4.1 (Same as Lemma
E.4
With probability at least - 6 over the weight initialization. Ba, w
Vy E {21,
= =(( + +f)m). Ot
Further, by the definition of Phase 1 (see Definition the attention matrices in the dynamical system are much smaller than that of linear MLP given in Lemma 1 This implies that during Phase 1, mainly the linear MLP is performing learning, whereas all the attention matrices are changing slowly from their initialization. Based on this, we have o GOO) (p1) O(( + mf)m) which implies that it takes only O(1/(oī )m))) iterations for G( (pa) reach a certain constant magnitude.
for a formal definition), the gradients of
Lemma 4. indicates that the samples with co-occurrence of the two target signals are classified correctly. The following lemma shows that the initial gradient from the common token, i.e., the gradient of ((t)( is much smaller than the gradient from the two target signals, which implies that the samples with only one target signal may be classified incorrectly as co-occurrence (since the network in this case will output positive value). This is verified empirically by experiments in Figure (bb, where the loss function corresponding to only one target signal and the common token first inereases in Phase 1.
Lemma 4.2 (Same as Lemma E. Let max; 1. With probability at least 1 over the weight initialization,
Bag m
w (0)
us
= δ (of (oivmm of vLL + ofmF).
Notice that the model output F depends the weight initialization scale and can be made small.
We next show in the following lemma that the gradient ( (Pa) of the common token will quickly at become negative soon after the training begins, which drives the transformer model to output negative value when it sees those types of samples. This implies that negative samples (without co-occurrence of two target tokens) will be classified correctly towards the end of Phase This is also verified empirically by our experiments in Figure (b). where the loss function corresponding to only one target signal and the common token descreases towards the end of Phase 1.
Lemma 4.3 (Abbreviated from Theorem E.19p. There exists a time To.s Z Ti and a constant such that for all E [To.s, Ti] Ti]
(8G() (1n) (G())(()) 0t
aGt (147) Gt
(8G(0(M) 0t +
aG0 (p2)
(1 + C') max
S (1-C)
88
db
7

---

Proof Intuition of Lemma 4.3 We first note that the term
( yi Ei m I (0) 1 Plz
makes the major contribution to the gradient aG (am) Such a term is small at the initialization due to the cancellation effect from positive and negative yi s. However, since the linear MLP G will positively align the two target signals at the beginning, for the samples with positive labels, 9, will
decrease, whereas for samples with only one target signal, g; will increase. Hence, (a de will become negative. This trend will continue until the gradient from μg starts match the gradients from M1, P2, which is what we establish in Lemma 4.3
Using Lemma we can show that all training samples are correctly classified at end of Phase 1.
4.2 Phase 2: Evolution of Attention and MLP for Large Classification Margin
In Phase 2, both attention and MLP matrices evolve towards enlarging the classification margin, thus driving the loss value small.
We now analyze what happens in Phase Let T2 denote the end of Phase 2. Recall that at the end of
Phase I, we have G (fz), G (u2), -G (ps) 2 1111). We will mainly need to show that such condition continues to hold in Phase 2, so that attention matrices will evolve with MLP to learn better classifiers. To this end, we exam the following gradient flow in the dynamical system:
aG (w)
E
(1) D02 Yis
Σ
(1.62) P bgt
Σ
Σ 1=1
aj. Ojz
w!!) )2
(4)
Ot
(eX
1==1
m
(t,61]
V(th))
T((
Dig
yis
anPi
4=1
(w) (((t) M
has been proved that at the end of Phase 1, for μ E {11, U2, paj. we have En,
Σn w (1) aw can hold for long enough (which first term in the sum on the right-hand side in Equation (4). On the other hand, from the dynamical system, we can calculate
8t 4 since the magnitude of Σズ= Σm= (o, u/f az (/ is large. Assume this indeed prove later). Then, we only need to focus the
aj U (1
(e)
2m1
Pt (())
(5)
0t
Thus, if v(P(()) > for all i E [n]. then Σm=
mi Enm
ai j azz (0()) (1)) is always increasing.
(a Thus, 81 mainly depends on the behavior of 1 E () Pig Mb, E1 (t.ig)
1g Fur-
w wto
and ov w( w
ther, this is also key quantity we need to analyze that , Eiy (2 Pig Vis Ely=1 Pg- ~ 1 E te 9ig Via if (1.12) pge-z,kkp R 1/L which holds at the beginning of Phase 2. Later, we are going to prove convergence of the training loss via
Note
(4)
the following: (I) the training loss decrease if the softmax probability is uniform; (2) even though the softmax probability will deviate from uniform distribution during training, we can bound such deviation and the loss value can still decrease.
Automatic balancing of gradients. As argued above, our main focus is on analyzing the behavior of (1) Eis (=x(()) Dig Vig. This consists of two parts: (i) Lemma 4.4 which shows that the two groups of samples with only the presence of one target signal have gradients Σier, g.x)) and Eie 9 close
(t)
to each other during training; and (ii) Lemma 4.5 which shows that the gradient gaps Eie, 9
Eier, Di
(0) and Eiero 9.") Eier, Di are not too small compared with Σieno 9 Both
(0)
8

---

Lemmas 4.4 and 4.5 establish that the ratio of those important gradients are kept within certain ranges during trainin. We call such a key property as automatic balancing of gradients, which further used for proving that the gradient flow can drive the training loss small.
and
Lemma 4.4 (Same as Lemma F.5
For L. E [Ti, T2]. there exists a small constant C < such that
Eiets
Pi (1) -iet, 9.¹)
19 min miin,iis, 9i Eiet, 9:0) SC.
Proof Intuition of Lemma 4.4
The intuition behind the result is as follows, If Lier, 9; becomes
much bigger than Eier, 9i during the training, then Eier, Di -iEr,, ///h is much smaller than
Σien 9,") Exer, 9 which makes (G (w2)) It is not hard to show that random
ae
8t
aF (G")(u))
tokens make negligible contributions to the gradient. Thus, for E 12, we have (ps),
St
St (0)) 5
Ht By the chain rule, we have 0g.") g' (vi ((::F"")) ')u: Be Since Eier, 9j
ap
(G)))
< <0.
8t
8t
0t
(G2)) then Eiel, o will drop faster than Eiers 9. i.e. Dt
8
Σ Eels 92
In Appendix, formally prove Lemma 4.4 by analyzing the ratio Ever, Pi / Eiirs 9
(8)
and show
that this ratio hangs over around during Faining.
Lemma 4.5 (Abbreviated from Lemma F.6
For E [2), the gradient satisfies that
Eie(n) K
(4)
Eieln] Di (1)
O(1)
0(1)
Eiergurgur Di Eien 9i
Eier, Di Eier, 9;
(1)
Further, for constant C. we have
(GG))
aG
(G
(2CU)()))) Ot +
OG
(1+C) max
(
(11- ((1-c)
Ōt
Ōt
Ot
Proof Sketch of Lemma 4.5 The proof of Lemma 4.5 relies analyzing how the ratio between changes. We show that this ratio will hang over around some range. Recall
ag (pa) and 8G (Gg)) Øt the relationship that 8G (M) R Eigs p(x( Dig y62
0t
(1)
E EØ=
mh
Uj 0/z
(() (1)
(w w//) >
It is not
hard to show that
(G1) aG (Pa) Øt
-iEr, 9i + Eieenurgur Di
Eie1, gi - 9; (e)
8t Define R(I) = Exer 8 + Eierousgu
Solving when á R(I) 2 0 yields a quadratic inequal-
Eier of") -Eice gi
ity, and further analysis shows that the root contractive and is within some specific range.
Utilizing the gradient automatic balancing properties, the following corollary characterizes how the attention matrices the dynamical system change in Phase In particular, we can show that after Wy-transform, lu and μl become more positively correlated whereas μu and μ3 (also μ22 and μa) become negatively correlated. This is direct result following from updates of the dynamical system.
Corollary 4.6 (Abbreviated from Corollary F.13 Fort E [T),
i at U2 W (T)
W 11 0,
at i W (T)
W
ps 0.
Since we have analyzed how (-)(-) will change in stage 2, we can utilize this information to analyze the change of softmax attention via the following relationship: by Appendix C we can derive
3pT
WTT K
( μ2

---

(9)
1i W (t) k(0.) diag (x()) - (G:0)
pl pi (1:0)) (0
91 yk
RV T1
1=1
Σ 1=1
U2 w (1.1) diag (60) (G (x(0((1))) (x p( X pr-
"t" Pi
9: 96
ny
The following lemma shows that the attention score between the two target signals Yn and Pa increases, whereas that between one target signal ji or #22 and the common token J43 decreases.
Lemma 4.7 (Abbreviated from Lemma F.16 For p.ve {21, P + and for tort E [2],
T Ōt. K
wioh
- 66m)
a 43 W
(T w
-1.0 T.
vm
L
Q
vm
Lemma 4.7 keeping track of the attention coefficients which proves the second item of Theorem 3.3
((
Theorem 3.3 via gradient flow, f1.0
5 Discussion and Future Directions
In this work, we developed a novel gradient flow based framework for analyzing the training dynamics of a one-layer transformer to recognize co-occurring tokens. We provided two two-phase characterization of the training process. In Phase 1, the linear MLP layer trained to classify samples correctly, with attention weights almost unchanged. In Phase 2, both attention matrices and the linear ML.P jointly evolve enlarge the classification margin, thus reducing the loss to near minimum.
As future work, it will be interesting to analyze more general transformer architectures such as muadi-headed attention, multi-layer transformer, etc. Further, it of interest study the dynamics of more advanced gradient descent algorithms such as gradient descent with adaptive learning rate, with momentum, etc., and explore how the hyperparameters will affect the training dynamics. Another direction is to study more practical language sequences where tokens are generated in a correlated fashion. Then the next token prediction becomes an intriguing problem.
Acknowledgement
H. Yang would like to thank Jason D. Lee and Yunwei Ren for insightful discussion and suggestions.
This work was performed under the auspices of the U.S. Department of Energy by the Lawrence Livermore National Laboratory under Contract No. DE-AC22-00NA27344 and supported by the LUNL-LDRD Program under Project No. 22-SI-004 and 24-ERD-010. The work of Y. Liang was supported in part by the U.S. National Science Foundation under the grants ECCS-22 13860 and DMS-2134145. The work of Z. Wang was in part supported by an NSF Scale-MoDL grant (award number: 2133861) and the CAREER Award (award number: 2145346).
References
[ACDS24] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
[ADF*22] 23] Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et Palm technical report. arXiv preprini arXiv:2305. 10403, 2023.
[BCE+23] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, al. Sparks artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.1 12712, 2023.
[BCW+23] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306 246.04777 2023.
10

---

[CL24] Sitan Chen and Yuanzhi Li. Provably learning multi-head attention layer. arXiv preprint Stt::2402.040844, 2024.
[CSWY24] 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.
[Dam18]
Friederick Damerau. Markov models and linguistic theory, volume 95, Walter de Gruyter GmbH & Co KG, 2018.
[DBK+20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Syl- vain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
[DCLT18] Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding, arXiv preprint 1rXiv:1810.04805. 2018.
[EGKZ22] Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang, Inductive biases and variable creation in self-atiention mechanisms. International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[HCL23] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05229, 2023.
[HWCL24] Yu Huang, Zixin Wen, Yuejic Chi, and Yingbin Liang. Transformers provably leam posittio correlations in masked image modeling. arXiv preprint arXiv:2403.02333, 2024.
[JSL22] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822- 37836, 2022.
[LAG 22] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2022.
[LCC+21] 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 IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[LLR23] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. arXiv preprint arXiv:2303.044 044444 2023.
[LWL+24] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402. 15607, 2024.
[LWLC22] Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference Learning Representations, 2022.
[LWM+23] Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, Zaixi Zhang, and Pin-Yr Yu Chen. What improves the generalization of graph transformer? a theoretical dive into self-attention and positional encoding. NeurIPS 2023 Workshop: New Frontiers in Graph Learning, 2023, 2023.
[MH23] James Manyika and Sissie Hsiao. An overview of bard: an early experiment with generative ai. Al. Google Statie Documents, 2, 2023.
[MS99] Christopher Manning and Hinrich Schutze. Foundations of statistical natural language processing. MIT press, 1999.

---

[NDL24] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735, 2024.
[Ope23] OpenAI. Gpt-4 technical report, 2023.
[PBM21] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 22(1): (3(1)):463-34999 2021.
[SHT23] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. arXiv preprint arXiv:2306.02866 2023.
[Tea23] Gemini Team. Gemini: a family of highly capable multimodal models. arXiv preprint urXi:2231..18800 2023.
[Thr24] Christos Thrampoulidis. Implicit bias of next-token prediction. arXiv:2402.18551, 2024.
[TLI+23] 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. 13971, 2023.
[TTO023] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers support vector machines. In Neur!PS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
[TLZO23] Davoud Atace Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max- margin token selection in attention mechanism. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[TWCD23] Yuandong Tian, Yiping Wang. Beidi Chen, and Simon Du Scan and snap: Understand- ing training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305.16330 16380, 2023.
[TWZ+23] Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. In The Twelfth International Conference on Learning Representations, 2023.
[VDT24] Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738. 2024.
[VONR+23] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexan- der Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[VSP+ 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need, Advances in neural information processing systems, 30, 2017.
[WCM22] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35:12071- 2022.
[WLCC23] Yongtao Wu, Fanghui Liu, Grigorios Chrysos, and Volkan Cevher. On the conver- gence of encoder-only shallow transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[WXM21] Colin Wei, Sang Michacl Xie, and Tengyu Ma. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. Advances in Neural Information Processing Systems, 11116618-161777 2021.
[YBR+19] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of loquence-to-sequenee functions? In International Conference on Learning Representations, 2019.
12

---

[ZFB24] 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.
[ZPGA23] Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08/ 2023.
[ZZYW23] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization, arXiv preprint Xrxv:2305.194200 2023.
13

---

Contents
Introduction
1.1
Related Work
2
Problem Setting
Data Model
2.2
Transformer Architecture and Training
3
Main Results
Proof Outline: Two-phase Gradient Flow Analysis
4.1
Phase 1: Alignment Linear MLP for Correct Classification
4.2 Phase 2: Evolution of Attention and MLP for Large Classification Margin
Discussion and Future Directions
10
Setup of Synthetic Experiment
16
Gradient Flow Update for Weight Matrices
16
Gradient Flow Dynamical System
16
Derivation the Dynamical System
17
D
Initialization
21
Training Dynamics: Phase
23
Initial Gradients
24
Maximum Perturbation of Neuron Outputs
27
Perturbation Term Involving Correlation of Value-transformedd Data
27
E,4
Perturbation Term Involving Correlation of Neurons
29
E.5
Neuron Weights Align with Signal Value
30
E,6
Alignment Common Token
32
Small Score Movement in Phase
37
E.8
All Variables within Range in Definition of Phase
43
Training Dynamics: Phase 2
44
Automatic Balancing of Gradients
48
How Fast the Loss Decreases
................................
54
Growth of Neuron Correlation
55
Growth of Correlation of Value- Talue-Transtorme/ Transformed Data
56
F.5
Change of Random- Token Sub-Network
57
F6
Change of Score and Softmax Probability
59
Change of Self-Correlation of KeyQQuery-Transfomee data
61
14

---

F88 Small Loss is Achieved
62
F9
Proof of Theorem 3.2
63
Proof of Theorem
63
Auxiliary Results
63
Probability
64
15

---

A Setup of Synthetic Experiment
We conduct synthetic experiment to verify our theoretical results. We create a dataset following our data distribution in Definition 2.I with 60 training samples: 30 samples have both μr and #2 in it, 10 samples have only Mi 10 samples have only μ2, and 10 samples have neither (4) nor u22 Each data consists of patches and each patch has dimension 64. The embedding dimension mn is set to be 128 and the number of neurons is set to be 256. We use Kaiming initialization to initialize the transformer weights. The transformer is trained by gradient descent with learning rate 0.01 for 30000 epochs.
B Gradient Flow Update for Weight Matrices
We provide the gradient flow update for each weight matrix as follows.
Owyy (1)
L
Ot n2
y
aj
pt v(:0pp
(1,0)
72
1
Σ
yb
a
1/
(t.t)
(8,8) Pi,h
1=1
1=1
1=1
h=1
ow
(1)
L
(8)
(t)
(x(0), (1:0))
R y1
at
i=1
1=1
owl
K
1
(1)
1
( (1,1)
Öpik
(t.3)
(1,5) (B)
9
1
91 2/2
0t
m
m
U8i
=11
1=1
h=1
h
n
1
m
Σ
7
9
(1)
y6
L Σ
44 h=1
j
(
(1,))
L
Pi,h
(1.i)
(I(h = h²) - Pi,k )9
(1.0) G)
I/2
1=1
1=1
h
1
Σ
A
(1)
ΣΣ
m1
aj (-w, (0) v((.)) (-T
(1.s)
(1,0)
(x (( + qi
w0 ( y(t d
(1.0)
96
V1
n/m
1
(2)
azui agul"
(1,1)
(-w (1) v(tti) (1.)) ((,)T
+ w/y (0) vioo X(i
96 Ub
m
1
(2)
(1,1) (
diag m (ne v(()) - w 10 v(t.)))(tt))) X(O)T
96 9/:
agli
pi
n/m
m
1=1
1=1 1=1
ow(0)
1
n
(4)
m2
(1,))
Σ
(t:))
(I(h = h1') - pl )b xi T)
9 96
W
Un
Pi,h
0t
nV
m
i=1
1=1
h=1
1=
1
(1)
Σ>
m2a
(-w)
( v((.)) (1,5) (K (1:00 ((,1)) + ((t,0)) diag(p) ( w/!) af (Ot
9 y6
0j
7/m
1
(1)
(,K(t
(OT v(t.)) p1 Pi + diag(p) (1.0) ( wnt) (T) (-w
(1.0) (.))
9i y6
n/m
L
1
(1)
K (txt) diag (vitO) T - wg (OT y(tu)p p(.))) pi p(:00) TT
9
xi
m/m
E=
C Gradient Flow Dynamical System
We first provide our complete dynamical system. The derivation of each equation is provided in Appendix C.T.
Ow
T
(t) P
Øt
16

---

1
(t)
Σ
1 m
{ (w/00 :
(0)
Σ
(t.11)
v(chh)
9ig 1/2
12
1 +
Ri yi1
anPh
11
12
12=1 ==
in
(6)
dv
W
K
T (
1
L
1
azv WOOT k(t.). diag (vihiiT w -wj
( y(th)) p( (())) (1 xi u
9
pi
mVm
m
i=1
I.
(0)
ΣΣ
agli WhoT (1.0) (1,T) diag ( v((.) - w v(t.) p(.0)) ((i)T v
95 1/
qi Pi
nVm
m
1=1
==
0 (t) (0)
1
1.
n +
91
(1
4/4
0j
w
p(.00)p (t)
pi
(1.6)
Σ
9i y1
(
w
())
v(th)) (1.))
pt
71
1=1
i=1
1=1
δv W W(( 1
V
Ōt
I
L
1 Σ
P (1)
ajv
WhaT 1 1
po-
1 Σ
9i
(0)
Σ
ajl
Tht (1.))
w/ Loti
n2
EOE C
1=1 1=1
(vEX(0
j=1
dv ((())
w (
Q
Q 4
Öt
1
Σ
9i Us
!1)
Σ
a
TwlOT k-(t.)) diag (vil..T (0) - -w,, (0) (1.1) (1.0)
v(t.0))))))) P(tem)) P(((..))
nym
(()
==1
mi
Σ
9; y/
(1)
Σ
agp
T
w k(t.0) diag y( (1) - - To) ((,,)) D(t.)) ) (0.0)
Pi(i.v) P((i.))
n/m
1=1
W T
w-(t) 1
K
Ōt
1
Σ
(1)
m Σ>
WLeJ (2.0) (t.)) ((
di Pgelken"" (wo") (())))())) - TT ((.))))(11))
9: Us
a
n/m
|=1j=1
L m Σ
Σ
(
WhaT di (100 (00pph Potiker (wf( (() - w v(t-i)
9i
y1
asp
nym m
==
1=1
C.1 Derivation of the Dynamical System
Lemma C.1. Let μ E {ui For all j E [m]], we have
2w0 (()T
w
1
n
L
(w mn(") (
(x(r)) ((oba)) T
Vig
Oja
0t
n
1==1
1==1 1==1
1 T1
(8)
(1,51) p((.)t)) ((((t))T 1
+
Vi
an
n
1==1
1=1
-
Σ
9
(t)
y
a
(Iw 112 + ( (t.1)
pg- tE,
n
1=1
17

---

where
1
Σ
(
Σ
Σ
(1) w/// ) pa
(t) (1))
9
yh
71
1=1
L
1
(tif) (tit) Plide "s
(w) mo(( * (n)).
9
a
n
Proof. Let (i, denote the index such that x (((,p) = μ. By the gradient flow update, we have
July
T
w p
1
(1)
mi
(x((x)) p(:tt)) P
922
Viz
ajz
Øt
11
k==1
12=1 j==1
(1)
(f,ii)
withi p
9i1 U/i
aj Ph,
71
4=1
1
(1.22)
Dig yiz
0/2
P kak
12: mEX (12)
1= f2=1
1
(1.51) (2,61)
: (6()(()))
91 Vi
Pl Viz
1111
4=1
1
9ig Viz
(
Σ
ai
ailllu
(t)
112 + Σ aja ( (w0))) w0(²)
Potlakktu
(tria)
n1
ia: (x2
1=
1
Σ
(4)
Σ
(t.i2) (f,5x)
((
(())(()))
Ii Vis
aj
Pla.2 1/2
71
(()
4=1
1==1
Σ
(4)
Σ
(t.ig) ||0 [1000(00)]
(0 gt (x
0000) (x))
Iú Vi
aj
Phida
n
X(1)
4=1
1
Σ
(1u0f) (10 - 11100 (w)) (1.0)
9 y1
aj
T1
Σ
(
ajz (0 w0l² (t.2)
9 yi
PgllKEN
1=1 12
L
(
Piid2 , (00( Migh
(1.1) (1.6)
(1.0 # (1)).
9
01
1=1
4=1
1=11
Lemma C.2. The following equation on the gradient flow holds:
0 w0
(0) (0)
aL
9
U1
Σ
T (t.i)) (t.1)
Ve Pil + aj w Pi.
(T) ((i))
y(ti)) (t.))
T=l
Eu
T.
(t)
Σ
ui ( vili0) (1,5)
Pi."" + + 1 vi
(1) (10) (.0
9 1/4
Pi.
1=1
1=1
Pix! (
ge
18

---

Proof. By gradient flow update, we have
0
ji
(1)
12
(9)
a
(1)
(T)
p(()) (2,3)
((.))
-
Σ
n
(0)
Σ
L
(O)
vll (1.0)
9i
yi
a32
9j
8/s
an w
0t
==1
1=1
:=1
1=1
a
9
(2)
y1
Wj
(
Vp
Pi
(1.0)
(t)T (t.1) (1.0) + az j/
Vp
Pi
P: Ed
+
y1
Σ
Bj w
(t)
vithi) (t.3)
Pi, + +0 "
(t) v/!
Pi,""
72
P:X(
Ell
Lemma C.3. Let μ1, V. E {uifim.
Then the following equation on gradient flow holds:
Ow (!)T
W(O)
V
Σ
95
Vi
Σ
W(IT (1) (1.1) ajv
w L
ipex
9
==
1
Σ
(8)
W (t) (t) (t,5) (t,i)
9;
aju
Wj Pqt-
Proof. The gradient flow update can be derived follows:
SuW
w a
L
me
agv au TT (x(O))(.))
0t
n
mi
as WhOT (1) (x()))(11)
+
9:
71
1=1
Σ
(1)
ajv w (1) (1.1) (1.)
9; U/
Wi Pott
(
m2
(8)
agn T wrO' (6) (t.i) wj Py-l
Di
Lemma C.4. Let p4, v E {uiler Then the following equations on gradient flow hold.
Ov
WLOT
T
w(th) 4
Öt
1
Σ
(1)
my
k(t>) (/) w ((t))) pi(en))
Di
yi
agv
nym
((EX())
j=1
Σ
(
Σ
1 (vici) w (T) vith)) plot)) Pi(aa
9,
agpi
Q
m/m O
j=1
dv
W ( (t(() 4
ØL
L
Σ
TT (1.1) (4.0) (1.))
Poetkkp (ug9T p(t)))()) - w " vit.0) (1.))
g y/
ajv
nym
())
19

---

L m
Σ ΣΣ
9 y/
9i
0g4
W (000 (1)) "
qi Partkeu ((ti)(()) - -wo v(c.0p)
nV m
1=1 1=
iv W T
01
1
(1)
m1
(19T(w)aj)) (0) w0fTT v(c.)) T p((a))))))))
(e.1)
(1.0)
9: yE
0g
Pi(anp)) PotUKttv
nym
(.PEE()))
j=1
m
L
(1)
T
k(haa)) mgg - -w)" T vit. (t:t) K
Pi(cin)) ) (1.0)
i (K((())) # (v)
91
Vi
aj
BV m
j=1
1=1
Σ
(8)
Σ
aglla
(2.i) (1)
(1))
gf (o/"To(::)//) - _ TOT p(t.)))((1)))
(t.5)
Di Vi
P((ep))
BYm m
/=1
1
Σ
(1)
m2
T WllT (1.0) (t.s)
qi Poot (1o)0 c((.)())) - -wo v((.))))p)) u((f(t)) # (w))
9,
yb
ag
nVm
BREX
j1
Proof. To prove the first result, we have
du T T wll 1
Ôt
L
1
(1)
agv T
WLOT ditd (v(t)) an n p(tt))p))))))) ie Llli a
9i
Ub
nym
f=1
L
agll T Kit. (ViLOT w T v((:)) 1(t Too
n/m
m
Q
1=1
1=1
1
Σ
(1)
W K diag (vicoT w v(t P(li))
9:
Vi
a.v
nym m
())
1=1
my
Σ
Σ
(ViST wggh - (0) vu. plies) Pileer k (
W
yb 1=1
To prove the second result, we have
dv
W
K
T
W
(t)
KK 4
Ø1
1
(t)
m
T (TT (1.)) (t,0)
diag (u.) (t) v(t - wof0 v(cc)))((.)) x110 4
9 Ui
azv
qi Pt
ny
m
i=1
Osp T
WTT (1.0 (T.T diag wm
(t) v((.)) T 110r 2
9i V/
qi Pt
nym m
1=1
1=1
Σ
(1)
a₂v wtit (1.5) (1,5)) (t.))
qi Pgelkkt (w ((()))(() - o vitei)p(t))) (111))
9: yb
n m
1
I. me
Σ
ΣΣ
WLOT (1.00 (000) p ((l)il()) - -wm" (v)) v(l0))(1))
9e
yi
asp
K di
j=1
To prove the third result, we have
dv w (
((( 1
K
Öt
20

---

L
m1
azv' W k (vic.)) =w)" (((())pp)))) (1i))) pa (1:)) io (O))
(t)
yi
BVm
m
1=1
==1
1
(1)
0j T W
qi p/'00)p) p diag ( v(aa) w v(ti))p p...)) X
9i
1/4
nVm
j=1
1
Σ
!1)
m1
O k(t. (i w ( vit (((c)) P(lins)
91 y1
ajv
n/m
()
1=1
L Σ
m
Σ O
(1)
T
w (1.0)
(1:0)p (t.1) part (w)") (a( (((:0(v) - w0s v(t:)pp"..) (((0))
9;
1/
OgH'
qi
RV m
1=1 ==
Σ CXPPEX
9
a
(x)
wfth ((1.0pp(i)) (1,1)
Pi(i.)) Potnuktu
nym
j=1
L
Σ
(8)
T K(0.9) (v.l..T w0)) - -w,, w T v(t.0)p(11) (2.5)
(
Pi(i.t) Poo-uke T(K] * (v))
) (1.))
9,
Vi
nym m
BPEX(O
1=1
m2
Σ
(4)
aylla (1)) (2)) (w. (v)0)() - w To y v((.)) Pi(E.))
(t,5)
9
y:
nVT m
==1
L m2
Σ
wll (1:00
(100)pp))) (t.))
potieeu (T,(TT((0000))) ((())(()) - -ws" v(t.i) I( (1.t) + (1)).
9
y/
asp
qi
nV m
1=1
1=1
D Initialization
Lemma D.1. With probability at least 8 over the randomness of the initialization of W and W q. for any 1 [d], we have
K (w Mc, W/ mms)
2d
2d
S mg
log
log
771
m
8
(w Mr, w mss
S mā
log
2d
4
log
2d
l + l2
m
m
(W pr, w:o (s)
S |,)|= o8m
2d
4
2d
log
log
1 + l2
m71
m
for any E [d].
4
2d
4
2d
I mell3 osm
1
log
+
log
m
8
(W (0)
mrl13 om
1
V
4
log
2d
+
4
log
2d 8
8
Proof. Note that W(()) K Ml, W M4g ~ NO of.. The rest of proof applies Lemma
H.2
Corollary D.2. For all i [n], 1, E [L.] we have
(0.3)
+
Pi,k
= L
Lm
21

---

Proof. Following from Lemma G.2 and from the first-order Taylor approximation on the softmax function from O, we have
(0.3) Pi,k
= + 0
oum
log
2d
+
4
log
2d 8 )
T.mm
5
m
The corollary then follows from Condition
!1
Lemma D.3. With probability at least 1 å over the randomness of the initialization of W and Wy, then for + l2 [d], we have
Kwr M4, wi maa)
S ogm
(v
4
log
2d
log
2d
5
5
for j1 + j2 E [m1]. we have
K (w
(0) w (0)) soim (
( log
2m
4
2m
log
m
8
m
for all j E [mi], I L [d]. we have
(
K [ w (0)
v m4) S 000mm
V
4
log
2mjd
log
2m
m
m
4
2m
2m
|00 1/ = ozm
1
log
log
m
m
5
)
(0) ( Mr 11 mr
4
log
2d
4
log
2d
m
m
Proof. The proof is similar to that for Lemma D.I
D.I and is omitted.
Lemma D.4. Conditioned on the success of the event in Corollary D.2 for all i [m], I e [L]. with probability at least ö over the randomness in the initialization of Wv.
agm L
2nI.
2nL
W x (0.1)
1
log
log
8
8
Proof. First of all, by Corollary D.2 and Assumption
2.3
27
0
x
(1)
0.0
Σ
(ww:)),
x (0.5) )) = abml mmml|X(0 (1)) (0.1)
ajm
pt
L
Finally, applying Bernstein's inequality and taking union bound over [n] and [Z] finish the proof.
Corollary D.S. Conditioned on the success of Lemma D.4 with probability at least over the randomness of W, for all [mi], i E [n]], l E [L]. we have
(O) (10.0) pl (0.0) w
m log L
nL
S 1
Proof. Conditioned on (U.i)) (0,1)
we have w/ (T) (00.0) (0.0)
((.00)(00.)
No ( (0.0 Thus, by
p) Gaussian tail bound and a union bound over e [mi]. i e [n]. e [L]. with probability at least 1 5, we have
(0)T m
(0.0) (1T.)
minL
pi
L L
log 8
22

---

Lemma D.6. With probability at least - 8 over the randomness in the initialization of a, we have
log
|++1| = mi
2
log( (4/6)
S-1 = m1
2
Proof. The proof follows by applying Hoeftding's inequality.
Lemma D.7 (Initial sub-network output). Assume that the success of the events in Lemma D.3 holds. For all E [d], with probability at least - 5 over the randomness in the weight initialization, we have
(G s 001100/mm)
Proof. Consider fixed i E [d]. By Lemma D.3 we have (|W 14 O(amm). Thus, conditioned on w Hés we have Σド=1 w (0) w-( (4 ~ NN( ofohmm1). Thus, by Gaussian concentration bound, we have | (14) & (0rioo/mm.))
Lemma D.8 (Initial network output), Assume that the success of the events Lemma D.3 and Lemma D.Adhlss. For all i [m), with probability at least 1 - 8 over the randomness in the weight initialization, we have
[F(0) (x(()))) < 20 v 22mmm (2L & 0.01.
Proof. For fixed I L, [m], [n]. by Corollary D.5
we have
T)T (0.0)) p(00)
m
MinL
w
/1
log
L
8
Thus, this implies that a (0)T (0.1) (0.4)
a sub-Gaussian random variable with variance proxy oB m log mink, Therefore, the following inequality holds.
pj
2mmm(log mink log(2/6)
Σ amu (O)T (100)) (0.4) ((..)) 2 20
I 8.
Taking a union bound over E [n], [L], with probability at least - 6, for all E [m], we have
mint, log(2Ln//))
[x((( < 20 2Lm1m((og
Finally, by Condition
we can make
(F()) (x()) s 0.01.
E Training Dynamies: Phase
During Phase 1 of training, the linear layer quickly aligns with the target signals and all the remaining quantities stay roughly the same. The analysis need to keep track of the evolution of the above quantities with respect the two signals #1, M2, the common token μ3 and the random tokens.
Definition E.1 (Radius of keys and queries). Define the radius of keys and queries Rk. RQ respec- tively be
Rx == max
W K
w:O 14 - 14 T
(O)T w-( 11
ijeid
RQ = max (JE(d
w
T w
- ppr
w. ( w Mj
23

---

Definition E.2 (Phase 1). Define the range of Phase 1 to be 10. T.). where Ti min(t', Cr, (afmmm } for some sufficiently large constant CT, and t' is defined to be the maximum time such that for all S t', all of the following hold:
(4) W(() w)
11 w/ w p R where R < 0(1/m1);
1. (
2. ma
w W// =- W S O(R//++R/mm)
T w v Rs where Rs S (OL//(m/)) W K V- a w
3. max maxp, 1
4. RK; Re S 50(88/m)
Based on this definition, we can further obtain the maximum softmax probability change as follows. Proposition E.3. Define
Rp
max iE il), (.vE(())
(t.3) Po-pukku
(0,5) P Poe-uek-u
Then
Rp = Rp=0
1
m
Proof. By Lemma G.2 we have Rp O(Rs/L++ + RRgg) 0 0 + 1
Initially, the loss for the samples with one signal will increase.
E.1 Initial Gradients
Lemma E.4 (Signal updates, same as Lemma 4.1 Art = for {21, we have
da,wy (0)
(0) μ O((a) + o))
Öt
Proof. Take μ μl1 First of all, by the gradient flow update in Lemma C.I
we have
Owl
T
wl(()
V 1
w/ wff)) (x(rs) ((it)) (1.i2) P
Dig
yi
aja
01
11
1==1
12=
1==
(t)
(t,)) ((t
Di Vii
aj.Pi
71 =1=1
(0)
(1) (1.12)
Diz Viz
ae W W/
p Po
12 ME (t2)
1= 1==1
(tris) (tr52)
((()
9 Uk
jj
Pladz U12
4=1
1==1
(
a a 11 112+ + Σ ajz 20 w0()))
Dig yi2
P
11
12 (12)
4=1
Ja7f3
Σ
Σ
(1,11) (1,61)
I Ui
aj
Plid2 Vi2
4=1
1=
1 Σ
Σ aj
(tas) +
E
(00w))
I yis
Plida
X(4)
1111
24

---

Σ
(1w0)) (1 - (w))) p
9i
U
jj
aj (w000) w0.0)) PoLiKep
yi
1=1
12
10
(1)
Plids 22 (w)) I(v + (1))
(ta)) (t.1))
9
1/s
12
Now, by Lemma
D.3 Corollary D.2 and Lemma D we have
L
1
Σ
(0)
[1 1 + a|0 (0) (2)122 P 1,c-p1
(0.3)
aj
y
13J11 XG)
1=1
1
(0)
|10g
(0)
12 + 110 (11)13)
P
(0.8)
9i
[mw|] (0)
112 + 1 (21)113) (0,1)
P
Im1
iE
1=1
1 - 0.01) L. (oîm + oîm) 1
4
2d
4
2d
1
m
log
8
+
m
log
8
L
(1 o(1))
On the other hand, by Proposition
E.5
we have
L Σ
1a =
Σ
Σ
Oj
W
P P
(0.2)
1=1
jazji
I $ \mmm
m
log
2m2
+
m
log
2m²
8
V
m log
8
|22 =
9i
(0)
yi
01
Pl U22
(0.1) U(uto vl0) )
# (Hi))
1=1
4=1
1=
4 Vm
( log
2nL
2nL log
< ogmmT.
m
m
8
If m 2 Cm log" m2 8 in Condition for some sufficiently large C, then |ea| S 0.01afm; and 2nL for some sufficiently large C², then |22l S 0.01a 5
m 2 CoL
Proposition E.5. Assume the events in Lemma D.3 least 1 - 8 over the randomness in the weight initialization, for all J1 [m1]. we have
D.3and and Corollary D.2 succeed. With probability at
Σ aja
(0) Whj (w
(
I soi
(vi= log
m
2mn²
5
4
log
2m 8
V log
4m
8
Jaihfi
Further, for all μ E (ni) we have
1
(0)
Σ
(0) (0) (0.)
9: 3/4
pg-
Pgtt
n
1=1 jaftn
i: # E (+)
V mim
2
(V log
2m
2
4
log
2m3
log
mi
n
8
m
and
00)
(0.3) (0.1)
(0) 1(1(000) # v(0)
9; Vi
aj.
Ph.d %/
4=1
1=1
25

---

S ogmvL
(
log
2nLd 4
log
2nLd
m
(0)
(0)
Proof. First, fix i, (, and consider the randomness of a. By Lemma D.3
0
"
w
is sub-
Gaussian random variable with variance proxy oim² (V // log o
2mf
+ log
2mf 8
This implies
that with probability at least 1 - 8/2, for all jı E [m11], we have
Σ
(0)
(0)) S ooiv mmm
(V log
2m,
4
log
2mi
log
4m 5
m
6
m
Jo
Thus, by Corollary
D.2 for all {1:) we have
1
0
10)
(0.1)
9
Uc
aj
n2
(0
1=1
ja:jzボ11
p
E X
(V=* log
2m
4
2m9
log m
mmm
log
n
m
m
We next derive the second inequality, Consider the randomness in ( Note that
Σ
(0.0) (0.0) (1))) # (p)) N 0, 00 Ph % I(1
(0.5)
Σ ?
10.0
10.0 # 2((0))
(P,,1
fg
Thus, by Lemma H.2 and Corollary D.2 and taking a union bound over i E [n]. E [4], p Id have with probability least 1 5/2,
(0,0) (0.4) Pi 12 m(
I( (0.4) * v (0) (H)), (W))).
)
S soām VT.
(
4
log
2nLd
log
2nLd
8
8
Therefore,
00)
(0.3) (0.4)
(1) I(115 # (w))
(0.1)
91
Vi
aj
Pl
122
1=1
1111
12=1
4
2nLd
4
2nLd
S L
(
log
log
m
m
Lemma E.6 (Random token updates). For f {ur we have
da, w (0)
w(())
o (00 + of) + oz obVmL)
01
Proof. Following the proof of Lemmal E.4
we have
(0() 0a,w
mi
v
(0)
at w w
(0)
(xl(i) p/ ((.te)) P
giz
Vi1
71
lg
1==1
1
(0)
(T (T.6i)) 1r(0)
91
V1
anPi
71
4=1
Σ
(0)
Σ
(Iw,
13 + 11 (m))11 (..0)
9;
yi
aj
(
1=1
26

---

y Σ Σ
(w] w/
(0.5)
9i
4/2
7
X
(0)
(0.1) (0.1)
(w) I(v (1//)-)) + (i)
Vi
aj
Pl l2
71
1=1
=1
1=
By Proposition
E.5
and the fact that only one X satisfies μ E x(),, we have
PgrlkKt (0))) (0108 + of)m)
Σ BPEX())
110 /0 Oj
(10(3) 1
(0.8)
n
9
j1
and
4 log
2m
2m
4m
|eil I -oi _oi/mim
Vim
log
log
m
m
8
[22 S ogmvT.
( log
2nLd
4
log
2n La
m
E.2
Maximum Perturbation of Neuron Outputs
Lemma E.7. For all S Tı, for all i [n], I [L], we have
m w
(1)
(0) p1
(t.6)
(0) wy w p(p)) (0.9) s 3(I.RroooIym) + 4R/L.
Proof. By Definition E.2 we have
(t)
(1.))
(1.)
(0) y (0.4) (0,6)
Pi
w
- w 100 vj 00.0 (0.4)
w
PLP
Pi.
Σ
(RpŌ(ono m) + RRL)+ +
Σ
(Rpo(ovo)y)) + R/(n))
< v -m R/n)
O(LLRpoOT Vm) + 4R/L.
By our choice of parameters Condition and Definition E.2 we have v(LRpooormm)aB)) <
T0 L.
E.3
Perturbation Term Involving Correlation of Value-transformed Data
Proposition E.8, With probability at least
5, for all E {m we have
azp
(O)T (0) u
< Ō(0 /mmm).
Proof.
The proof is similar to Proposition E.22 and is omitted.
27

---

Lemma E.9 (Value correlation change), For all p.v. E {uili=D we have @ (vll) (u), (v)
|i: E X H + e x 0(1).
0t
n
and thus,
(100))))01)) ( - (1)) S f u x( + :i: x 0(1)
n
for all S. T₁. Thus, for p & us we have (o)) ((0 s Ō and [1(0) we have
((ogm) fort S Ti
Proof. By Lemma
Ow
1
Whi p L.t-
a V
Ōt
()
Σ
!9) y4
agli W (2) (0,5) (1.5)
Wj P
1=1
Further,
agv T
WoOT wj
(0)
Σ
agy
W w
Smil R.
Thus, by Proposition
E.8
we have
m1
m2
j
T W
v
ago ty tl. mj
(0) (T)T w V +
Σ
w a,v
wj
Σ
a, w w
j=1
j=1
010001 + R,
which implies
WiOT (1) (t.))
95
agv
iuEX (2)
1=1 ==1
aju W (1) (t.i)
9 V1
wj Pat
X
E=
1
x + (ilvex (1(0001/mm) + +miR)
n
fi:ue + 0(1)
where the last inequality applies Lemma E.7
Corollary E.10. For all t S Ti, we have
L
Pliia (000)))) Uoll + (H1)) L (0000m) 1=
(1.5) (1.1)
J
4=1
Proof. We derive the following bound:
(
(t.) (tr)
(()) T(ofh:)
(1:)) # (1))
U1
Phods
(1,1)
(t.f)
2000 (w) ((o(1
(1:0) # (1))
Phnla
1/2
,
4=14=1
< L- L (o3mm))
where the last inequality follows from Lemma E.9 and Lemma D.3
28

---

E,4
Perturbation Term Involving Correlation of Neurons
Lemma E.11. For all S IT,
we have
: p1 E x
Σ EPIEX
9i Us
Σ
(
w P
(1) (1,1)
S
1m
J2nin
Proof. We first derive:
Σ
aj
W/) w// Potl
(1) (1) (t,5)
(
i= jantn
Σ
Σ
(0))
max
1,l
6X)
1=1 jュス
Σ
Σ
Σ
1000
10)
(f,i) max P 2 1g
Vi
0j
1=1 1o
(1)
Σ
(6)
(0)
(1,3) max P.
+
yi
Gja
w w
GALEXO
1=1
(2)
For term (1), by Proposition
E.5 we have
1 Σ
00)
{i: un E xx
δ (oi
aj
w
w
5
1:p1 EX
I=1
For term (2), note that
Oj (WwS), 22)
(10))((gh (w w j2
gi Yi
-
jazti
(1: /1 X
Lmu
where the inequality by Lemma D33 Proposition E.12
Finally, since (t.e)
pgt $ 1 + to + Rp, combining the upper bound for both terms (1) and (2).
we have
Σ
a i1id2 (w , (1,5) Poti
P yb
Ja7hi
{i: p1
(
δ (o L).
Proposition E.12 (Neuron correlation change, Phase 1). For S Tı, for all f1, E [mı], we have
0 (0 (0) (9)
J1
< 2L
0100
m
log L
minL
+ O(LRpono + +
Øt
and thus,
(
(1) (4)
J1 j2
w (1))
s t2L
0100
m
L
log
mnnL
8 + +ÖLLRpono + +4R/L)
29

---

Proof.
By the gradient flow update in Lemma C.2
we have
(t)
w
(
(1,5)
1
(0)
Σ
(1)T
p((00)p0 (28)
De
"
Pp
Iv' De
"j
0t
n1 P==1
y==1
V=1
1=1
(t)
Σ
(0) ((0.1') (0,8")
+ (w)) (t) ((.11)) (t,3") (0)T
w/ V= p()10))
96 DF
ajz
wij
Py
Py
71
jj
w
(0)T v(0,2') (0,5) + {w/ ( ( p(t))) (1.))
1p
w To)Tv pp
Now, since by Lemma E.7 we have
(1)
W
X
() w/
w w x('p; (0.1) - ((LRPpogOiVm) +4R/L,
by Definition E.2 and Corollary D.5,
we have
(O)T V (U,1')
pp + (w p(ts))) w)j (O)T
((+)
yv
0j
"
P=1
I L
11001
m
log
mnnL
+ O(LRpono))m) + +4R/L
L
8
(1)
Σ
(0) (0, (0.1)) 1/0.)")((''' + (w
(t) y((ht) (6,1) (0)T pp..'))p(0 p(00)))) pp(())))))))
1/2
Ojs
P=1
P=1
minL
+ ÖLLRpono vm) 4R/L)
log L 10000
8
Thus,
K W w/
(0) (1)
(0) (0)
wijn w
(T) ((T))
w
/oo01 m
minL + + + 4R/LL
I 12L
log
L
5
r=0
Corollary E.13 (Neuron Norm Change, Phase 1). For all t S Tı and all j E [m], we have
I m | 11 |ml || s (22((1000)
log
minL
+ +(LLRpgno + 4R/L).
E.5 Neuron Weights Align with Signal Value
Theorem E.14 (Signal correlation growth, phase 1), For 1 < Tor p (11
Oayw w (1)
Σ
(4)
(t.i) g
Σ
(c)
(t.t)
O((0 + +of)m) +
Øt
71
lktp
9
(1)
1=1
(u Wn=1"
where
| <
Ō(Log of mm1).
Proof. We take μ = μi and the proof similar for μ = 122 By Lemma C.I
we have
da da,w, (1 w/(0) 4
Øt
30

---

Σ
| / 11 + a| 1 (m1))B) (ta))
Palterm
11
t.p1 EXIEX 91
Σ
(1 10 12 + 1 (un)13) (1.1)
Po-littps tE
n
BRIEX -1=-1
= (1w 112 + (x))))
1
Σ
(8)
(1,3)
Σ D
(1,3)
9
+
1=1
1=1
y/ml
y/m-
where
Σ
(t) (1)) (1.))
Di
Oj Wii w p Pat-
I=1
jzz
(1,8) (e.1) Phile 2/2
(11( mu( # (p1))
1/s
==1
We can now bound the magnitude of E by Corollary obtain:
E.10
Lemma
E.11 and Proposition E.5 and
|e| $ L. Ō (0)vm) i /mm1).
Finally, by Lemma D.3 and Corollary E.13 we have [|mo) 1|3 = @(aîm) and by Lemma E.9 we
have (| (p1)ll )12 O))mm. The proof is completed.
Theorem E.15 (Random token growth, Phase 1). Fort I for E R. we have
Da, w w
(08 0
0
of)m) Ö
Proof. Fix a E R. By our Assumption 2.3 a appears at most once in the training data set. Now, assume /1 is in the training and let i" be the index of the sample containing μs. Applying Lemma C.T on the random token J4, we have
wayw) W 0t
(1,2") (10)f + (1!
9
aj
P E,
t=1
where
1 (t)
94t
yk
Σ
Σ
aiz
W w
(1) (1) (2.5)
P Potlikt
n
1=1
J2
(1.2) Phitz V/ ((0(()) Mo + (p)).
0
h=11
1==1
We can now bound the magnitude of & by Corollary E. and Lemma as follows:
|e| L δ (o)vm) m) + (0mm)
Finally, by Lemma D.3 and Corollary
E.13 we have
da, w
(4)
w-(t) (1)
p
0 (1 (o) of)m) O(g.T.m).
Øt
31

---

E.6 Alignment of Common Token
Lemma E.16 (Initial Per Neuron Gradient of Common Token, same as Lemma 4.2) max; 10 We have
Let F
Dag w
(0)
W (0)
v
Ma
Ō (0)
(oivmmi mmn + vovmL mL + (o + oi)mt).
01
Proof. Following from the proof of Lemma
B.4
we have
dag, w/
(0)
( (0)
ns
Σ
(0)
(0) ji 12 +
||u!
(120)))11113))
9,
Ui
P pae-ikkuus
0t
n
1=1
+ dj
Σ
Σ
"i w
(0) (0)
Im1
jz
L
(0.i) (0.0)
(0.1)
_(0) 1(1/0 + o0 (1a))
+ 01
91
Vi
aj
Pla l2
1=1
4=1
12=1
22
For the first term, we have
L
( yi
[| 51 112+ + (ps) (0.0)
10)
p
())
1=1
W (1)
11 + (,))) po-
#1=1
[1
10 1/3
(w,)) )12) pg p
(0.8)
9
n 44
1
01)
1 + (0)
(1
Σ
(0.)
Σ
(0)
0.1
9
P pg
9,
q ke
1=1
1=1
By Lemma D.8 and Corollary D.2 we have
L
Σ
(0,1)
Σ
(0.))
9
pq k
93
Pa-
P
n
31
1=1
V1=-1
(
F
F
11
L
12
T.
1/1
1 F F
1
I +
Lm
2
-
F
L
Lm
2
2F + (L/mm).
By Lemma
D.3 we have
|/ (0) 1 = @ and )1) (2>) Oafmm). and thus,
1
Σ
[ 10 1 - 11 (Ha))13 P Pye < 5 + of)mf).
(0.1)
24
Further, by Proposition B.5
we have
L
| =
Σ
Σ
Ojz
(w))) (0)
(0.))
1=1
ja
32

---

2 of mmm
2ml
4
log
2m3
log
m1
m
m
[22
9i
(0)
y1
aj
Pli,2 12
(0,0) (0.8)
v00) (1s)) ITvb2 # v (1a))
m i=1
1=11
12=1
(V
4
2nL
2nL
log
log
8
Finally, we have
Dag wi
(0)
zō (o (oiymm + +oVmL.+ + (o) + offmF
Öt
Lemma E.17 (Per neuron gradient of common token). For t S T1, we have (0)
Ōaju w
(0)
11 Ōt
I Ō)
Proof. The proof is similar that for Theorem
E. and omit here.
Definition E.18 (sub-network). Define the sub-network structure as
m1 Σ
G(p)
aaw, Wvu.
We compute the gradient of the sub-network as
@G(u)
Dayw (1)
W
ow
Ow! (t)
(c)
V
j
W
0jj
01
al
0L
06
j=1
I
Σ
(t)
(0)) (tria)
Pi
Via
Oi Djz
Pgelzikep
n
(e)
12 ii
j2
L
1
(t,i1) (f,i1) Phidz 1i
(w))
Yi,
1=
1=1
Theorem E.19 (Complete version of Lemma such that for all To.s St t S Tı.
There exists To.s S Ti and constant 0 < C <
OG (m1) (G
(G
(OGU)() (1)0)(()())) +
(G2)
(1+ C)) max
S
Ot S < (1--))
Öt
Ot
Ot
Further,
(G
(G))
2 2((oh + om)m.))
Ot
Øt
Proof. First of all, we have
ar!")
mi
(o X12 (1.)) Phis
a,w (t) w X12
Opi
(1.4)
öt
Ot
Øt
4=1 1=
By Lemma E.4 and Lemma 1. have
da, w μ1
dazu
w ps
O((o + of)m),
Øt
Øt
33

---

Da, Tazwi (0) ôt
da,w;
(0)
(0) Ms W
M2
= O + +f)mm).
0t
By Theorem E.15 and Corollary E.24 have
Dajwy W X
(t)
ppi
(t.0)
1(1 & (p a,u,/ X12
Ot
0t
12=1 =0 n
( + of of)
8F ot > O, which implies
0n < 0 for i E I and
Thus, for all i E J 1 Ia. E U2U I3.
> 0 for
at
at
Next, we show that there must exist a time such that (Gl) max Without loss of generality, assume (G()) >
(a)) (ua)),
dt
dt
Öt
(G())
We now analyze the condition when the
8t magnitude of the update of the common token is less than the magnitude of the update of the signals. By Lemma
at
C.T
we have
W mi da j
mi da (2)
Ms
Øt
Öt
1=1
1
(1n((0)12 mn (11))
Σ
(
(1.
Σ
(1)
(1,1)
+
m
a a
9i
P
93
P
Pat
P
n
NEL,
FeL2
==1
| | + mm
|| (1s))
Σ
(1.5)
(1)
(1,
Σ ())())
9
1
9i
VELJUI,U4
1=1
NEl
1=1
f=1
7
1 1
12 + mi
|10
(Ha)l3
Σ
(1)
(1.8)
Pat 1
n
|
mi
(110
Ji
Lgtt,
1el
1=1
NEL1
1=1
|
12
+ mi
(wa)
Σ
L
(t.1)
Σ
(t)
(t.i)
n [ 12
+ mm
(pz))
Lgtiktps
7
P
iEh
1=1
E
(M1)
|| + mi
(pT)
[
mi (1:0))
3 (1 + a(1))
Σ
9i
1
(1
(()) Σ 0
n
n
iel.
iEL.
1 (1 (1))
Σ NETgUT
9.
(
(1 + o(1))
Σ
9 (1)
+0
(v
m1
n
n
m
iela
where the last equality applies Corollary E.24 If m is sufficiently larger than my (by some large constant factor C)2, then the last term is negligible. Note that the above implies that if aG (11)
Ōt 2 then Σien, 9i ExEL, of 2(1) since Σierur, 9." (+) 22(1) for all t < Ti. Now, by
(ps)
8t
Proposition E.21 if the initialization level is small enough, then Eien, 9: EiEls H 2 2(1).
Thus, if ((s) 0t
(w()) then Eier ag(")
-O((0) + mf)m).). Therefore, there og(l")
dt
de
must exist a time t' (11m) such that ag(')(())
Further, for O(F/m). have y!!)! 0 for all i e 14 and Eien Di < min 9i Eiet, g(")), where F
01
maxie(1) [1)
34

---

Now, consider time point t' when (G aGu (mm))
At tt, we must have
min(Cier Di Eiers gi ) Eieh, gi = 21)). Thus,
2
Σ
9i
Σ ie4
(t))
12(1),
which implies
2
9i
(1)
8
9i
(t))
-((1)
0p
OF
ich
iE4
For E I1, we have
DF!()
aG
(G()))2)
(G())
+0 L
[oimm)
+
0t
0t
0t
Ot
aG (12) +0 ((oommm))
Ot
n1
For i E 12,
aF!!
(G1)
(G())(s))
+0 [ofmmi) =0 -0" ("
T.
O
Ōl
Ol
For E la, by Proposition
E.21
op(()
(42)
(pa))
+ ( (Homm)))
L = 0 (Foîmmı).
àt
Ot
Ot
and for i 4,
aF!")
(G())(())
L ofmmi)
+ +0
Öt
Öt
n
Thus, by chain rule
ag
( 0F:
0g.
we have
ot
0
2
9i
(1))
2
Σ
9:
Σ
0(ommmi)).
(8)
sch
i
iElgUla
This implies that there exists a constant L such that
di )) -ofommm))),
ΣS
Σ
(1)
Σ
9 ( + (1+
Σ
10
Σ
g
9i
Ot
iEf2
iih
iETgUTSUr
which implies that there must exist a time To.5 S ((1/m) such that for all f. To.5.
aG
(1+ O0
Öt
Moreover, by Theorem E.26 if Or, is sufficiently large, we have To.5 < Ti
Finally, consider the time when (g)) + (wz) have
80 (us) at
During this time, note that we
at
ae
aC (M1) aG > @(ofmm1))
Ot
Øt
Further, if (p()) 31 Thus, there must exist constant U such that
(pa)) (ct))(H)) then 9
(c) O(ofmmi) for all E T2U 74.
81
(G(+)(
(8G (G
< (1 -1)
Øt
0t
Ot
for all < Ti.
35

---

Corollary E.20. There exists a small positive constant C' such that if we define T == minst
mingesuls, F!!) 2 C) then T 1 011mm).
Proposition E.21, Let F max
1/.0 For L S Tı, we have
1 Σ n SEL2
Σ
SO(F).
TI VEZs
Proof.
First of all, by Lipschitzness of g, we have
Σ iEls
(c)
S (G (pz) - ((2) - G (12))] + 2 max
[ + (1).
95
9i
n
72
rE[m)
1EL2
Without loss of generality, assume G (m1) > (12) for all S Ti; otherwise, we can break the interval [0. Ti] into sub-intervals by the time points when G (un) GO (p2) changes its sign and then apply the analysis below to each sub-interval. We first derive
OG ((1) (G()((2))
Öt
Ot
my
m
1 n
L.
Σ
Σ
(0)
Σ
ag ajz {u (w)) jz poe
(0)) (lita)
Σ
(0)
Σ
(t,i1)
withTTW pn
Piz yiu
+
Di
Yis
Ph
71
==
1=111==1
11=1
1-1
my Σ
n
Σ
(1)
(6) (1))
(1(t2)
(1)
Σ
( ((t.s))T.. u2
9 V/12
aj ajz
P 1q,122
9i
Vi
Ph
71
b=12==
11=1
h=1
(1) (1)
1
Σ
(2)
(1.))
1
Σ
(1)
(1,6)
ajj Ujz
"1 1 1
Di
yi
9i
yi
Pat-L
12
1=1
iEh
1=1
M1
( (
j/ )010)
(t.s)
Σ
(t,t)
Oi Ojz
Vi
Pa k
9i
1==1
iEL2
1=1
iEls
1=1
12
L
1 n
1.
(t)
(B v((hi))
Σ
(1)
Σ
(t.k:)
( (((1)) P2
+m1
Ii yis
Ph:
11
9i, yi
Pi
n 11=1
n
111
11=1
4=1
ajs j2
Wii w
(t)
(t)
1 Σ
Di
Σ iE12
9 + t((1)
(1)
1=
El
n2 Σ
(t)
Σ
(T v((&))
w:(t u1
1
Σ
(9)
Σ
L
(t,)) vIth P2
+ +m
Di1 Yis
Pti
9i
V/
Ph
n1 11=1
71
1=11
4=1
(1)
((00())) - (()( - (N2)
(WN) (0 (00)
aj 42z
h==
+ (11) + ((2) + +
max
(,,01) + c(11)
iE(n]
Σ
n
L
1
Σ
(1)
Σ
L
T
T(ht) 12
+ +m
Ri
Ui
Ph
14
Dig
yi
Ph
n
$1
4=1
m1
1=1
By Theorem E.14
we have
L
Σ
(1iii) (((.1) ( w 41
1
(1)
(t.iiT
((ii)tt u2 dt
w
mi
yi
aappl
9ig
V/i
aj.Pi
n
=1
T Ö(Log
36

---

J.' G (11) (12) +0
max iE(m)
1=:001)
fToo
((1p
Σ
mn
Σ
(1)
(0)
dt
0 (( + (42) + max iE(m!
(1::411 T Tofmmi)
and
mi
Σ j1 b1=1j2- j2
ai a2 [w/ - ( = O(1)
exp
By Grönwall's inequality, we have (M)) - ( (v)
ST . O m2m))/mmm)))m)) +0 + (((00(1)) - (12) + max sEin
F/ Tofmmas
+ +[, (1. (.01LO8Vm) +0 ((((0(()) + (12) max |2,(00) lofm )) 0(1) di iE(n]
0o (o + G.
(w) (w2)
max (11000
This implies that
Σ 9
(2)
Σ
14) < (M1) + (p)) ((2) + max F((0)
9i
scIn
X
xETs
Thus, if the initialization scale is sufficiently small, 1 Eier, o: and + Eier, 9h (0) by a small constant.
are only differed
E.7 Small Score Movement in Phase 1
Proposition E.22. Conditioned on the success of Lemma D.1 least 1 8 over the randomness of a, for all [n] and p.v X(), we have
2.1aaa and Lemma D.3
with probability at
|| (u) 3
00 .3morooy mm1 log (T) (1)(v)
mid
log
nd
8
8
Σ ag la (0) (n) 1 (O)T (v)
0 (omooooy mim log
mid
log
nd
5
5
and all 1, e [L] with K(O.)) (v) and (0.5)
(v)) we have qe
ayv
T T
K
x (/ (0)
wn
< Ō Vmo /mmi))
10,0 V (0)
Ō(o% (mmooo//mm)
avv
Proof. Fix i E [n] and u.v E x)) Consider the randomness of a. By Lemma D.1
Corollary D.2
and Lemma D.3 aj (0) (w) 12 (O)T (()))) is aub-Gaussian random variable with variance proxy O(ffmmoo my (mm))))))) Then the following inequality holds.
(osmorony mm, log mid V log 2 f
a (0 (2) ()) T (v) <0
37

---

Finally, take union bound over i E [n]. μ, E {pr The analysis for the second term is similar.
Next, note that ajy WK (T) K(0'0((00))
(0, w is sub- Gaussian random variable with variance proxy
O(co)mmooo m/))) Thus,
Σ
ajv
W TTT
K y
(0.0T w
Lemma E.23 (Score change). For all $ Tı, for all v. μ E {14 we have
Ov T W(t
w (t)
Ki: μ E X + (i: WE x o ( m
m
n
and thus,
v TT w -
T
1 Vm
[i: a
X +
+ )
S
K
n
Proof. First of all, by Lemma C.4 expand the per step gradient descent update follows:
dv
W K
w(lt)
Q μ
1
Σ ((vEE())
(1)
m1
(U9T (v) w) (0) ((+.1) v(t.))p(e)))
(e.1) (1,0)
9i
yE
ag
P((E)p)) PytURtEU
n/m
j=l
(1)
Σ
(t)
(t)T K
(x(()))
(v(h:O0 - (0) ((t00))))))))))) p ((K # kx)
9i
v
Pi(in))
BV
iin (
j=1
1=1
(2)
Σ
(t)
(1 (1)
(1) gen. (w (())) - w y(t))) (t.i)
(4.))
(1,5)
(O)
y1
Pi(i.))
(0
j=1
(3)
m2
Σ
(1)
agp T WOOO (1.t) (1.1)
qi Pot (u) (w,oT ((t.)) -wp w v(t.0) (1.0) I(q((h)) + (x))
9
y/
nV m
6
1=1
(4)
Analysis of (1): By triangle inequality, we have
1(1) <
Σ
gi
ag
(2)
nV m
(
j=1
(a)
(9)
11 (1)x00(())) (v)
(i) (iT v(a.)) (1. (1.))
(1.))
9i
aj
n
())
j=1
(b)
To analyze (a), since (a + ab) + =2) = ab ae2 + Pe + E182, we have
a (v) (0)
(1) w10 (v)pp) pp)pge
Σe
|| (0) (v) (1) (0)
(0.) (p))
a
1=1
38

---

m1
ag [1)) B (v) Σ
aj |1 (u) To) (v) (T(000)))
(1,8)
P pouuu
ag [ (v) 1 (T)
pr- P
(0.3) P Potpuktu
Further, by Lemma D.I
Lemma
D.3
Definition
E.1
and Definition E.2, we have
mi
m1 ag a (v) 1 (0T (v) Σ
Σ
ag || 5 O (0)
| (0) (0) ((t)())))) (v) (, u
(0)T
S max nd
S m
[RxO(opa))m) + RŌ(o)m) Rx R
(9)
Combining with Proposition E.22 this implies
m1
ag 11000 (w)lBur) T
(V)
(t.i)
ag |I (1)13,, (O)T 1100) (v) (0.)
(V)
(1(()(())) (v) 11 (1)T (1()())) - -114:(0(v)) (
(O)T (10) (T
(2,1)
< my
max v.j
Pora
+ -
( T (v)
(1.00) (0,1) Po-ukty -Ptttaktt
j=1
I Sml [RxŌ Vm) + RŌ) + RxR) (7. (Rp)
Rp + Ō mo /mm)
Thus.
(()] =
2 (v) (0T,, (v) (t.s)
mi
و a31lb:00 (w) (T 6(00 (v) Pot
Σ 1=1
ag (v)lB3) (T (v)pp)
-=1
mi
aj || 10) (v) , (u)p
+
(2 +RP) Rp
(RxŌ(ano RŌ(aōm) + Rx RxR)
I mm
+ Ō (Rpořmo10e mm)) +ō (ommorooy/mi)))
On the other hand, to analyze (b), we have
( (v) (T) v(..))
(2.0) (1.))
a (1 (0) ((U.)) (0.1)
Pi(i Pa kak
(1) (2w (t) ((..0)
Σ
as (0) (w)1 (T) y (0.0) (0.
(0.1)
(1.))
P((t.p))
PH(ip))
j=l
Σ a, (0) (,>1 T)T 1 vV p(0.0)1000)) (0.1) Pi(ia))
(1:)) . - -p(o.)
Pg-p.ktu)
(10)
39

---

Note that
0g (1 (v) v(c.))
Σ j=1
|1 (v)333 (w) T) v(..0
P((s.))
ag
P(((.())
m1
ag 11 (0) (1)13u) ( (1(+>) Pi(i.u)d
aj || (0) ( (o) v, v,(0,)) (0.0)
P((i.)).)
1=
m2
mi
az ayllx:00 (v) T (i
(e.3)
ag |1: (v) (0)T Tv,10
(0.4)
Pi(i.))
m2
ag || (v) [1) ( v((,.) Tv(
az (5l((00()))) [1) (0) v((0.))
(e.5)
Pi(iaa)
ag1 (1))3100))))3 Tv
(0.4)
Pi(t.aa.
m1
(1 (w))B, (0) v((1)) Σ
ay (v)1 (0)T [((0.)
< max [ELL]
+ LRp max LE(T.
aj (1 (v) (0)T (0.0)
(11)
which implies
[(b]
m11
ag (11 (v) ( v(.))) (e.6)
==1
my
m:
ay 11 (v) ( ((.))) (t.z) (1.1)
Pi(iaa) p p(.i))
E
ag 11. (w)1B,
(0)T v(0.0)
(1000)/100))
P((i p
(0.6) (0.1)
my
a a31lk:00 (v) (T) (0.1) (0,5) (0.0)
Pi(s,a)
P Pa-
a (v)12 ( vitii) v(te)p(())))
(t.0) Pi(t
Σ ag (v)11)) (0) (00 (0.1)
p(00))
(t.))
||| (3 (0) (1.1)) (0,5) P((ia))
(0.0) P Pomm
a 2111: (1) (w) (0) (0. (0.1) (..)
((00)) Pa-
(10
Σ ag |11 (1) (c)T VO.TT (1,0) - Σ
as 11, (w)l3w (w)B v(000)
(T)
max [E[L]
1=1
Σ a, 11:1 (1) ()) v100,0
(1,0)
LRp max
[E
40

---

ag ajllk ( 1o (i. p( (((sp))
(0,5)
(0.0) PoppAt P pet
ay a-ll (v) 1 (0) v(..)) ((..)) (0.3) (0,4)
(ii) S
(RxO(a0OIVi) RŌ(a)m) + Rx
+ I.RpOl0mm000 1olooooyimmm mmm1)
( Rp) + +0 (ojmooo / + Rp
+
where (i) follows from Equation (10 (ii) follows from Equation 11 and (zid) follows from
Equation and Proposition E.22
Combining the upper bound for both (a) and (b), we obtain
1 Ki: ,, X(1]
((1) < v m
δ mi (Rx +R((@m)) + Rx
n
+ LLROOlo mmm)) (t + Rp)
obmooo + / ( Rp) )
+R+)))
Ki: M.E p Ō(/L))
mm
Analysis of (2): (2) can be analyzed similarly as (1) with (v))l replaced by T K(C.). Thus, we only need replace oi with of /m and then take sum over We have
[2)1 < 1 {i: p E H m (RxŌ(opoI RŌ(O)vm) + Rx R
m n
LRpO mood mmoo (1+ LRp) oz m / ymma (1 + LRp)
1 (i: ME X((
Ō(1/m))
Vm
Analysis of (3): (3) can be analyzed similarly to (1), and we obtain
|i: J. E x( Õo(D).
(3) < Vm
Analysis of (4): (4) can be analyzed similarly to (2), and we obtain
[i v E x() Ō(\/m))
[(4) < Vm
Finally, combining the bounds on (1) (4), we have
du T w-(O) ((1))
1 {i: u. E X H (((/))+ 1 (t: M x(x)+ (t: w (0
K
1
Ø1.
m
Vm
1
{i: μ E X x +((i: u X()
0(1/L+1/Vm)
m
Corollary E.24 (Softmax change). For t S. Tú, have the following:
41

---

bottth x x E R, then
(1,5)
pp
1
1(12 + tnvmm))
01
otherwise,
0p/1112
(5,i)
1
Ot
I2 m m
Thus,
if both x x R, then
(1.3) (0.i)
<ō
((t
+
_
1
1
vmm)
Phida Phila
T
ni
Lm
L
otherwise,
(t,8) (0.0)
I δ
(1
+ vmmm)
Pl,,2
Pl
Lm?
L
Proof. Consider fixed 1,11, 12. By Lemma G.2 we have
Öpi.
(t.s)
00h,
(1.i)
(t,i)
(t.i)
(
I Pi
+ Pl
Pl:
Øl
0l
01
By Lemma H.23 we have the following cases:
if both x x(1)
R. then
05 i2
(1,6)
)
1
L +
nm
otherwise,
0 (t.i)
+
Öt
L
m
(t,a)
Next, during Phase 1, we have
Ph, < I + Lm + Rp. Therefore,
if x((( E R, then
(1,)T
0s
((,f)
W δ ((
1
v )
1
Pt.
L +
0t
n2
m
otherwise,
(t.)) (1.))
( 1 +
1 )
2 m
Pl
Ot
L
Thus,
if both x((1) x-() E R., then
(t.i)
J2
1 (1++1)
1
1mm
+
Øt
n
Lm
I
vm)
42

---

otherwise,
(t.2)
-
<
i +
01
Lin
This implies that
if both x (0)
R. then
Pi - P11.22
(t.)) (0.1)
- 5
+
n
Lmn
1 L
+
vm
otherwise,
(1.6)
(0.1)
(
+ +
)
Plad2
Lm?
L
Lemma E.25 (K, Q self-correlation change), For S To. u4 E {ui) we have
[i: E x + fi: v D x
I T WOT w v μ ( w v
Vm
Vm
(i: M x
Ki: x
a WoT w v - μ T w
(
(0)
Q
v
<
mm
n
\m
Proof. By Lemma C.4 we only need to replace Ō(aîm) by Ö(ag /m) in the proof of Lemma E.23 Thus, we omit the proof here.
E.8 All Variables are within Range in Definition of Phase
Finally, we prove that at the end of Phase 1, all the variables in Definition E.2 stay in the range.
Theorem E.26. For t S. CT, / (ofmm)) (where the constant Cr, is from the definition of T in Definition E.2). all of the following hold:
wj (0) w S R. where R = 0(1/m²)
1. (14))
2. (xx
wj (1) (((0) 1 w, (0) : 1
O(R/n R/ /m);
3. max (
Wi W v T W W K
< Rs, where Rs S 0(1/m/m)
WTT ( W(l!) @ v T T (T w (0)
4. max pexp,.e
Roi
5. max ue
"
TW (((t), v p O
(O) W
(0)
Rx.
Thus, Ti CTT/ (ofmmm))
Proof. The first two results are proved by Theorem E.14
Theorem E.15 Lemma E.17
The third
result is proved by Lemma E.23 The fourth and fifth results are proved by Lemma E.25
O
Theorem E.27 (End of Phase 1), At Tn we have y (( = 0(1) for all i [m], and
((T) = ((1), (P2) = O(1,, G(T)) O(1),
Vp Y R: ((T)) (μ) Ö(000
Further,
Σ
aj w
(Ti)
azz w
(T))
@affmmi).
n=
43

---

Proof. By Theorem E.26 we have y (( ((1) and
G ( = O(1),
(22) - O(1), -G(L)((u)) O(1.
Further, by Theorem E.14 and Theorem E.19 we have yi 2 ((1) for i 1i and
((T) 2 2(1), ((T))()2)) 2 2(1).
Finally, by Theorem E.26 we have that 7i CT,/m² And note that if the constant Cr, in Definition E.2 is sufficiently large, then by Corollary E.20 we have T S Ti. Thus, we have y: P.P (1) ((1 for i 12 U T4 and (G())) 2(1).
By Lemma D.7 and Theorem E. we have
Vy E R: ((T)(()) Ö(000) mm1)
Finally, by Lemma D.3 Proposition
B.5 and Proposition E.12 have
aj " ae w
(Ti) (Ti)
(ofmma)
Theorem E.28 (Phase 1, formal restatement of Theorem
3.1. With probability at least ô over
the randomness of weight initialization, there exists time Ti ((1/m) such that
((T))()) 2(1), ((T) 2 2(1), ( -2(1).
All the training samples are correctly classified: y.P!T)) (2(1) for all i [n].
For E [1,
w j1 11
(1) (1) -
w 100
1
m
( (000) (1), v(0 (01)) (1)) (1)
v. W T (1)
w # w (0)
+
33/2
L
vm
Tw w v 1 w
T)T ())
6
1
K v
m
Twi w v- 1 W O)T
1
Q V
m2
The training loss satisfies ((T) = @(1).
Proof. The first two results are proved in Theorem E.27
The third result is proved by Lemma E.9
Proposition E.12 Lemma E.23 and Lemma E.25 The result the training loss is direct conse- quence of Definition
F Training Dynamies: Phase 2
The idea of the proof is to first define conditions for Phase 2, which guarantees that the small training loss can be achieved. Then we will show that those conditions can be satisfied starting from the end of Phase 1 and up to at least (ppol((m) time, which will serve the end of Phase 2.
Definition F.1. We define Phase 2 of the training to be E [Ti, Th] such that
44

---

The change of K, Q self-correlatie is small:
W n W T W v/ V
P.V
W
T @ V Ō(ob/m)
max
The change of softmax probability satisfies:
Pli.2 ,/ < <0111.2)
(1.)) (0.1)
max
(L), sE(n)
'The sum of neuron correlation satisfies
m2
Σ
(1) / , ay s {0ffmm)).
1==1
The gradient of W, Wv satisfies
Ow W 1
maxpe
j= ag 0r
S LL(L))
Σ
gi
=1
11
====
Ujn wi G,j "
(1)
n vE(N)
G(1) (X (8) 8p
Σt Eiz=1
maxig(n]
1
Σ
(1)
E11=1 Em=
a3s /// azz 20/2
11
VE
The gradients Eier Di for E {11, 12, 1 satisfies
Σ 91 S min
Σ Welz
(t)
Σ VE
(1)
(0)
9:
Ji
max
9
gi
gi
9:
iEf and
eTTs
EEL
Eiel, 92
(0)
2 <
2
EiEt, 9,
Fl for all i E for some fixed constant C.
(w)) < O(log m) for μ E {us
Σte Eige E(mo) GIO (t) (00 Pothi S0(1).
(1.))
Th S Oppolymm)).
Corollary F22 For & E [Ti, T2), ifi,j E where I (1,, Iz, Ia, T4), then
max
0()).
(JE 9j
Proof. Take I = 1, and the proof is similar for the remaining cases. For fixed i,j E 4ı, we have
g (t) + exp(y, P())
exp(y, F(1))
( 1 exp(y,
exp(IxF!"))
g
where the inequality is due to ((( 2 C' in Definition
E.1 Now we consider
ex)())F(1))
((()
45

---

exply, E E 1=== G (X((
(1) (t.))
explvi Li,=1 Ei G (x12 por
exp
(
(t,5)
Σ
E
6o (w)pge
(6.))
ME [14kk=m
41=1 WE (EE
(1)
Σ
E
(X
(1-5)
Σ
Σ
(x) (1,5)
exp
Pattl
11
x
11
X x E(x Rima
(2)
By Definition F.1
is easy to see that [22] S O()). For (1), we have
4 Σ
Σ
Σ
Σ
GO (1)
11=1 WE{ P
Σ
E ((I s 00().
5
p 12
ME|NE
Thus, we have
(cp(yFF(0)
exp(ys P(!)) 5 50(1)
fort E T2].
Lemma F3. Phase 2 in Desinition
F.1 is well-defined.
Proof. We need show that the definition is valid Ti with conditions satisfied with strict inequality. Then since everything changes continuously, there naturally exists T3 > T1 such that Definition/E.I]]s well-defined.
First of all, by Lemma
E.25
for H, E {14) we have
max P.N
(T:)T W
(Ti) v w
(O)T
w (0),
K v/
Ō
p w (Ti) (1)) v T (T) w
Q v Ō
max
Next, by Proposition E.12 Proposition
E.5 and Lemma D.3 we have
(Ti) ag wn azz w
(Ti)
@(ommm1).
Recall that 1 Σn (T)) = (1). On the other hand, by Corollary E.10 Lemma
D.3
we have
mi
aw
(Ti)
max u (M:) 1d
aj
W
(Ti)
p
O(oâmmi)
Bio (T))
(T(T)) m
max
1== aj
0 = =
mg Lm
(T1)
(T1)
Cji W) Ug W
46

---

Also, by Corollary E. E.24 we have
(Ti.))
(x
(0)
max im
0t
vm
4=14=11
82 (7..0)
maX
E1==1 -1z=1
GIT (1)
m1
ag with
(Ti)
a222 W2 w (Ti)
ac
m2/2
1
S
Σ
93
(2:)
1==1
<E(n]
Further, by Corollary E.24 we have
(3)
(0,3) P2
O(1/L2).
max 1,6,,, (E(m]
Piid2
Now, by Proposition [.21 if F is small enough (which can be achieved by making the initialization scale small enough), then
Eigr, Ji (Ti)
EiEt, 9i (Ti) 2.
Lastly, by Theorem E.27 we have y (1) Ollog m) for μ E (aterr
2 (11). And is straightforward to see that (G(D)
Finally, we prove Ei E, ly: elm
G Potll,kt12 S 0(). simple corollary
(Ti.))
from Lemma D.7 and Lemma D.8 is that
Σ
Σ
(0.5) ( 1 50(1).
1=1 ly: X Elanlize
Next, we have
I.
Σ
(X)) (1) (Ti.)
Po (p
Σ
Σ
(X ( (0.1)
4= loz x (uu
11 l2: x E(u
Σ
( - -0 (x(])
P
4=1 ly: x E(m))
(X)((x
P - (0.1)
((Tu.))
(Ti.)
11 la: x Ei
+ Σ Σ
-0 -x()))
Pgtliktl2 - (0.1)
-Pg-i, ke
111
la:
X
0()
which proves that
Σ
Σ
CTT) (p 50()).
(F..))
las ((fe)())
In Phase 2, we will analyze the dynamical system in a different way since all the variables now might change dramatically from their values at initialization.
47

---

Lemma F.4. Fort E [Ti, T2]. we have
0 m1
U3 W 232 w/) 11)
(1)
2m n
n
(
Di
> 0
01
J1
J2
and thus,
Σ
On Jn 43 (0)
= 2(ofmm1)).
12=1
Proof. By Lemma C.2
we obtain
0
m8g
m1
(1)
an
ajj
(O)
y(c,0)
(e.t)
1
Σ 99
(0)
(0)
v(e.0) (e,s)
V/
Un "
yi
Plz
1=
i=1
1 1=
i=1
h=1
2m
y
P(1))
Then, note that the final term is positive since by Definition F.1 Finally, by Theorem E.27 we have for E T2].
we have P( > for all i E [n].
(1) ( a32 2 ((oimm1))
(2)
===
L
F.1
Automatic Balancing of Gradients
Lemma F5 (Same as Lemma 4.4 4.4. For 1 E To), there exists a small constant <<> such that
Eier, Hi (1) -iet, 9.
SC.
(1) min (Eier, Di "Eiets li
Proof. Without loss of generality, assume EiEr, (1) > EiEl, 9:- Then, we have
EiEr, o0 EEEE (00 min gi EiEr,
Eiei, - Eiei, Hi
EiEt, Ji
(8)
1.
Eier 96
Eie Ji (0)
Excr,
Now by the quotient rule, we have
8
I 0 if and only if
Br
Eieig
Σ
Σ Sh
Σ
Σ
(
50
Øt
Ōt
iEh2
iET
Oy: ((t)
g' (( Oy: ((())
Σ
Σ
gt
(t)
Σ
9
(
Σ
50
Neh
VET
Ve
Oy:K!!)
Σ g
Σ gh
Σ 9' ((()
50
Eier, 95
Ø1.
Eie Dr
Ø1
vel2
vills
iEl2
iEta
Oy: ((1)
By, Fy,F(())
Σ g' ((:FF.(())
Ed((rF("))
50
Eiet, Di
(1)
Øt
Eieln Di
Øt
iel2
sets
48

---

-1
By. F(1)
Eiel Di iel, (1 + exp((tFF ))( + exp( exp(-y:P(())
dt
Dy.F%
5 0
Eiels Pi (8)
Els (1 + exp((x + expl -y: )(
0t
By: (())
Eien, 9i
iEls (1 (x+xx(uxp()") exp( happ-y:FFi -y(
Ot
0y:F((1)
(12)
Eierg Di iEl2 (1 exp(yr exp(- ((:t1)
0L.
Note that
EiE 9i
1
EiEl 9i (1)
+ minker exp(-y;F,))
(1+ exp(xx exp(-
1 maxies exp(-y;F,))
ieT
By Definition F.1
we have that for i' E Iz.
Oye F
Øt
dagw 10 (O x-(')
Plade + a w x Ophi1a
(1.2)
12 (1,1")
y
Øt
1-
aG (x
L Σ
3p (1,1")
(1 + (())
+
Yr G(t (x((2)
(1))
a
01
1=1
Σ
mi
aj w 2
(t)
(1) 1
Σ
9
(2)
Σ
9,
9,
(
Σ
9
(8)
+ o(1)
Σ
1am1
iEL
sE12
ic(m)
Similarly, for E Iz, we have
Oye F())
Öt
=-
Σ
ai 5h
(t) a3
w
(1)
1
Σ ieh.
Σ
0.0h + Σ Σ sEf
94
(
9
(4)
+ o(1)
Σ 9000
ich
iE(M)
Now, we analyze Equation
12)
Note that for E [T, Eier, D. sufficiently larger
than Eiel,
Di (1) (i.e., EiEro /iEls is bigger than some threshold), then 0t PtP max dt Thus, the ratio will decrease. Similarly, Eiets (/0) E'Er Di is too small, the ratio will increase. Next, we compute bound on Eicr, g.
(G))
(G()())) and minjel, 81
(1)
EiEr, 9i / have
gi (1)
so that dt Eicr, 8. = 0. Substituting Pl.) 2 in Equation (12), we
EiEl, 95 TETs (1 exp(y:F(1) ()(1 + exp(- (pp(-y)((())
)
(a azz (0(0))
Σ 03 (0)
Σ
Hi
(0)
Σ
(1)
Σ
9
!1)
+ o(1)
Σ 960 Σd(²
1==1
iELuU,u
iEh
1e
<c(m)
Σ
Eiel, Di iela (1 exp(v:F(")) (] + Fp((
(1)
49

---

1711
YTA1
2.51 )
Wi ajz u3 !
Σ
9i
Σ
Ri
(1)
Σ
9i
(0)
Σ
9j
11)
- t(
Σ
1==1f==1
iEL1
iEls
6E(oo)
1 minjer, exp( 1 )( maxiely exp( exp(-y;iF y)F(())
Σ
oi
Σ
Σ
9
(1)
Σ
9
(6)
- t((1)
Σ
96
(1)
iEh
VEIJUTgU4
iEh
ic|ml
-
Σ iEh
9i
+
Σ
9i
(1)
Σ
91
9i
(0)
- o(1)
Σ
9i
(1)
iElzUlaU14
ieh
iEIs
kE(m]
Eier, Di
LiEls (1+exp(1: P!!)) () (1+epp( (x F(
Eieig
Eiets Di
(1)
Ei o! Eiel, (1texp(es (()) 1111 expttepp( (xp(-xx111)
Eien Di Eietsuruu 9:0h -iien Di Eicr 8:)) + o(1) Eie(n) gi
Eierg Di
Since
1 (1 exp(ytP, (()))) + exp(-yil
EiEr Di
1
+ max exp(-1y,F F!!) minter exp( 1/1 exp(-3//3 F7
we have
Eiels (I+oxp(y, +lxpppp F((1)))
Eierg ")
( Eiela (1+exp(y) ]](1]]][[[]]]) -exppp((-xx ::
Eicr (6)
+ minier, exp( (1:1((0) 1 1 maxiel, exp( exp(-9:P%(0) (12)((0)
e
+ max exp( 8/ exp(-y:PP (()) 1 min exp( (-y, y.F,(0)
By Definition F.
we have
EiEn, (1)
Σien, Hi
+ +EEEnUAUr 9/
- Eien, Hi EiEt, + o(1 Eieln| Di (1)
S 6.
Eiets Di
Thus,
Eiera (1) Eier, Di
([I+minee exp( exp(-y;F(1)) 1 1 max expl (1p(-)
6
[1+ maxiets exp(- exp(-M;1!10) minger, exp( exp(-y;l") 1y:
-1 -1)
Lemma F.6 (Complete version of Lemma
4.5
For 1 E (Ti, T2), we have
Eienn D.
Eiejnl 9
0(1,,
0(1)
(0)
Σien 9 !
Eien 9. Eiero 9. aCt))
ΣiEM 9 (t)
O(1),
O1))
Esel, 9 (0)
Eiel, 96
ac
Further, there exists constant C such that
(HG))())
aG
ag (ps)
(aG()(H)
OG (H2)
(1 ++ C) max (
(
<
< (1 (I-
I1 +
Ø1.
Øt
I1
0t
50

---

(G)))
(G()) 81
Proof. Without loss of generality, assume have
8t
First of all, by Definition F.1
we
8G (mL)
1
m1
ais Wji
(t)
(t)
Σ
9i
(4)
Σ
9i
(4)
-
o(1)
E
9i
(0)
01
11=1
sEI2
iE(m)
(G()
Ci Wi (8) (2)
azz
Σ
91
(t)
Σ
9
(4)
- o(1)
Σ
91
(0)
Öt
j==1
ieh
sc(mi
O (ps)
(0)
(2))
(t)
Σ
( + o(1)
Σ
(0)
91
Di
9i
Öt
12==
iEL
iE(n]
This implies that
(G
ΣiEn 9i
(1) Eieh 9i
(1)
- o()) Eis(n) (n) 9i
0 (G())
Σier, Di
(1)
-E 9i
(1) o(1) Eie(n] 9i (1)
ät
0G (ps)
(ps)
Eien, 9i (1)
+ +EiEnurgu 9i (0)
+ o(1) Eie(n)
[m] 9i (0)
0t (GL)
Eien gi -EEE Di = =1) P
(4)
(al)
We first define the ratio R(t) Eicrouiann (0) -icn 9f¹
We next analyze
Since
0002 It the dependence of R on is clear and for the ease of notation, we omit this dependence below. Rearranging this definition, we obtain
Eier, 9. - 8."
(1 + R)
Σ
9
Σ
gi
=
Σ
gi
(13)
vih
iE
iEIgUT
We consider the range of R E [1. By Definition
F.1
we have
Eienni yi
Eie(ni 91 (1)
O(1),
O(1,,
(14)
Eistourut Σien
Σ iin 9,
Eier, 9i
where the second equation follows from Lemma E5
This implies that
(G))
Eve(n) 9."
(G)) (1),
(2) 0
EiEn, gi
-EiEI, Di
at
which proves the first result and
(Gs)
Eien (00 + +EEEEuuuuuL Di
8t
+ o(1).
(GO(
Eier, (1) - 9i
(1)
Thus, to analyze & og(()
we can instead analyze
(uy)
8
Σ IEUUUA Di
Eien, 5.0
20
Σien g²"
- Di
Σ
Σ g
(t)
+ 00
Ji
NiE2J1IUU4
ieh
vel
iE
Σ
9
(6)
Σ
9
i
9
(8)
Σ
9
(t)
20 0
Öt
iEh
ieh2
51

---

dy; (0)
01g (1)
E
R Σ
2
(1
R
(15)
ōt
öt
iEl2
1E7.
Recall that by Definition F.1
we have for 1 E I1.
0yi, p(()
m1
M1
aj a
az w/ (0))
1
9
Σ
9
(0)
E
900
+ o(1)
Σ
Öt
n
ich
iETgUls
for d2 E 1z,
Oyin ((t)
12
m1
TR
Us / az2 w
(1)
(1)
1
Σ
E
D1
11)
Σ
94
(1)
Σ
9
(
+
((1)
Σ
9i
Öt
n
ji
1==1
1E11
VEIsU1gU44
ih
iEl2
iE(m]
for d3 I3.
Oyis Ps
MRg
Uz (t)
1
Σ
o
E
96
11)
Σ
9,
(1)
Σ
9
(0)
+ o(
Σ
9,
Ot
31
1==1
IEL:
iETgUTgU4
iEl
iE(nl
and for is E La.
Oyi, p(t) b
Cj :
(1) 1
Σ
Hi
E
Di
(8) + o(1)
+01)
Σ s
0t
n
11=1j==1
ich
iE(n]
The above implies that for i E 11,
Byn p(t)
E Cien 9: (1)
-
-iieh 9.
(1)
R + +0(1)
Σien Di (1)
- -iee, Di
, a3
w
j2
n
1+
Eien Di (t)
-Eiei, Ji
(2)
12=1
for i2 I2.
Biig p(O)
(1) 1 (1 R o
Dt
Σ
Σ
Eien Di
(1)
-iE Ji
(1)
51 1==1
for i3 Is.
Obig p(()
(1) 1 Lier, 1.0 - -ist 9i
(1)
06
aj u a
(1)
- R -R+o(1))
(1) Eien, 9i
Eie 9i
(0)
EiEL, (1) - 91
(1)
12=1
and for E 14.
Byc (())
8t
b
m2
m1
(0) ai 1 a22 "
(1) 1 (-R
EiEn, 95
(1)
Eiera gi
(1)
1==1-=
Substituting
the above
into
Equation
and divide both
sides by
ンー Σド= aj W ajz w (2)
(1)
we have
(R-1 +01)+ + Σ ('(F(()) ( Eren, Ji (1) -iEe,, 9000
(1 + R) Σ
)
5(F 12
R
+ o(1)
Eien, Di -Eeet, H1
igE
+ Σ
((F (0 (R+
kEA
2 (1 + R) Σ )(F!)))
(1+ (1+I Eieh 110 - -iet (1)
Di
(o(())) o(1).
R +
(16)
Σier, Di
(!)
Eiet di
(0)
hth
52

---

By Lemma F.5
we obtain
Eien, 9i
- Eier, 13 Ji
(0) S,
Eien (1) -
Ei 9i
where to denote the small deviation. Note that Equation (6) is quadratic inequality in R and can be rearranged as aR² + bR + c 2 0, where
a Σ ) Σ
Σ g ( + Σ
g'(P - E g'(F.)))
k
1e14
c= (1+
Σ
( (1 E)
5 - Σ (F(()
(2 c2
fgEla
Now analyze the equality condition in Equation aR + bR + c O. We have
where calculate the root of the equation T6
b Eisei, g'(F!"") + EnEI, ) - E
a
Enen, g'(F(")) EEigs, (')
(1+ ((1) (-Eieh d' )(F(!!)) (1 + E) Eisets I' (F! (2 +1 Enen I' o((F!!))
Enen I (1 + Eiger, I (((
Note that
Σ EEI
Di (1)
min
S
Eg Σ
Di
(()
max
1 exp(y:
Tel
1 explyiF
EE!
YE1
By Equation
14 we have
2 -1/4 and 2a
b
Eijr, ((P(11) - -Eish, ((1:!)) Enen ((F!) ) EEEnet,
Ener E ((F(())
b
< (1 + O(maxexp(-xx max
iE(n]
Eizet,
Enen g(F(1)) -Eiee 9(F(!))
1
and
= (1 + 0(e) O(mxx exp( y/ F(!))) 2. (E(m)
Recall that we are considering the case of R E [1,3]. Thus, we only need to consider the
root that is positive and we can calculate the root 1 max4e[n] exp(-u: P(!) are both suffiently small.
La E (1,2) if and
Next, since g' (F( < O. the root R* is contractive (i.e., if R(t) > R then R(t) is decreasing and if R(t) R then R(t) is increasing).
Finally, the result at the end of Phase implies that R(Ti) E [1,3), which completes the proof. n
Corollary F.7. For t [T1, T2l, we have
(p1) ((1) (p2) (12) (Ga) ((t0))))(((( = O(1)
Proof. We first prove ((a))))
(pa) = (1). Following from Lemma F6,, we have
O(1). By
U(2)
Theorem E.27 have (ML)
+(). Thus, for t [Ti, Tz], obtain
((T(()))
(IT)) 1 (G) dr Dr
((1)
O(1).
G (H22
((T)) 1 oGM dr
53

---

21 (()) Since G
Note that Lemma F.6 and Definition F. imply that (ag))
(1),
-G
Bt
similarly, we have
(m)) = (1). Finally, ((2) (G()
(())
= (1) and
(1س
= (1) imply that
(Gw))s))
(p2) = (1).
(GS)
F.2 How Fast the Loss Decreases
Theorem F8. For t E [2], we have
E(t)= ((ofmmi)(t Ti) + (1/T(T:))
Proof. First of all, the gradient flow update for the empirical loss is given by
aï
E:=1 (' p(t)) Py. By Definition Fl we have for E
Oy p(t) is
Σ
my
as w
(1)
w
(1)
1 n
Σ
9i 10
Σ iEguUlgU
9i
(0)
Σ iETgUT
9i
10
+
o((
Σ
h=
ich
[E(n]
for 12 I2,
Oyig p(()
in
w
(2) 1
Σ
(2)
Σ
9
11)
+
95
(1)
Σ
9
+ o(1)
Σ
95
(1)
Ot
31
12
iE1
RETSLTgUr4
ich
iE12
*E(m]
for i3 E la.
dy:s va
(!)
m1
m2
u (t)
(t)
1
Σ
ot
(t)
E
9;
(6)
Σ
91
(1)
Σ
9
+
o(1) Σ
9i (1)
01
j1
12=1
iEL
iELUUUU
ic
ich
iE(m)
and for 14 E 14.
Oyg, (t()
b1
1711
mi
Oj w
1
Σge
Σ
9.th
+o(1)
Σ
Ot
ich
[n]
By Lemma F.6 we have
Fy
Oj Wj (t)
G2z w
θ
Σ
D
(2)
Øt
12
vE(n)
for all i E [n]. Therefore,
01
Ø1
m
1 Σ
an 1 (t)
azz (0 0
9
7i
1=1
12=1
FE(m
aj. ja azz ) θ (
oie
vEiN]
m2
M11
aj w jr
aij (
(1) + (?.2)
$1 1==1
54

---

where the last equality follows from the property of binary cross-entropy loss that ((x) (
for x > O. By Definition F.1 and Lemmal F.4 we have for all [Ti, T2].
(0) aj u 2z w 32 @cffmmi).
11=1
Thus, have
0L 0 (o/mm,))) )
Now, consider the differential equation equation in / and solve by
dI de
- 1.1 Note that this is separable differential
1 dL - + di
d - + C2) =0 L(t) C2
12 di
This implies for E [T1, T2], we have
I(1) = O(oîmmi)(t Ti) + (L/i(Ti))
Corollary F9. The following bound holds:
mi
mi
(Ta) Gja
(Ta)
(a,
(T
(Ti)
(Ti)
<0
ml
Oj
w
az $2
m
12
Jam1
Proof.
This is a direct consequence of Lemma F.4 and Theorem
F.4
F.8
F.3 Growth of Neuron Correlation
Lemma F.10. For & [Ti, T2), have
at w az w2
(8)
(0)
ag w
(
ag, w 20
(11)
J==1
12=1
J1 =1 f==1
mi log logt
mi log m ofmmi
o
= 0
ammm
and thus,
mi
172
as w azz a (0) @(ommm1))
1==11=11
Proof. By Lemma F4
we have
i
aj W 222 (1)
(1)
2m
Si y(1)
(4)
Øt
By Theorem F.8 we have 1 "n gi = o(T(0) = 0 (of ((tesn)(((ettt ((1-2;))1/00 (1-2,)+//(er))
Further, by
Definition F1 we have |m(()) < O(log m). Thus, for E [2), we obtain
my
mi
mi
my
Oj w ajz w
(1)
aj ( (i)
aazz w ( (Ti)
3:
1=
n
g yi p(T) dr
2m
Ti
t=1
55

---

( (o9mmi)((T - Ti) + 1/ ([(TT)
[") 0
Om1 log logm) m)
dr
= =o mi log m log
= o ( (m log m
aimm1
ofmm1
where the last line follows because T2 I O(poly(m)) in Definition F.1 we have
Finally, by Theorem E.27
(Ti) U as i1 , w (())) ((ffmm))
1==
F,4 Growth of Correlation of Value-Transformel Data
We now analyze the correlation term with value-transformed data.
Lemma E.1 (Growth of correlation of value-transformedd data). For E {ue) we have
ot " w ()) w (v) n (
1
Σ
(t,t) P Lek (w) 1=1
Σ
(t)
Ut Σ
(t.)
Ji
yb
Is
(VEX())
i=1
Thus, fort E (TiTTa) we have
T w (f)
w V- - u witi w so
log m logt ofmmi
max
Proof. By the gradient flow update, we have
oul Ouw(DT ( w v
0t
1
Σ
Di
(1)
Σ
m2
agv
WOOT
wy
(1.1)
1
Σ
Ji Vi
(1)
Σ>
mi
aip
WhOT
w (p.)
(0 (1,1)
n
n
(vEX((
E=
(v)
1
Σ
(1.t)
(z)
1
Σ
(1)
p (1,1)
Di yb
Di
y/
P
n
ExExO
Thus, obtain
w/T) v uw (T:T w-(rriy//
L
ol0 Us
Σ
(1,5) Pg
+ ([)(
1
Σ
di
(1)
yb
Σ
p
(1,1))
dt.
By Definition F.1 for t [2], we have (G)) S Ologm) and Σls ((as)
P 2 O(1)
(t.1)
Further, by the property of the eross-entropy loss, we have
Σ b =
EpEx
Therefore, by Theorem F.
we obtain
( wT)) -ー - pW w:( s (llggm) 1' Dt < O(lgmm) 0
log
56

---

Corollary F.12. Fort E [T,,T,] we have
Ow (2)
max (E
mi
0j σt
1
o(L/L
Σ 9
m h-
O5 Wjj ji 0z2 w
(1)
iE m0
Proof. By Lemma E.9 and and Lemma F.11 for all μ + u {ui we have I W w/uus 2/
O + 1/m) and |W w/12 = Olaßm + 1/m). Thus, by Lemma F.11 we have
maxp,u W W w
oovm m 1/m
5Ō
211=1 1===
mi m1
aa w/f azz 10
oimm
max,
| 213 v11
5ō
aamm + 11m
Σ mi
51 j==
m1
Oij Wi
(1)
Gz "
(1)
opmmn
Recall that
Ow (1)
W (1)
7
(tii)) T m.
a
9i V/1
Ph
0L
11=
1==1
This implies that
m. Sw
((()) 11
max pE 3
2== aj
50
oo obymL+ L/m + oBm 1/m
1
Σ 9, 10
n
ΣØ=
(/1 1 azz "/
opm
Corollary F.13 (Complete version of Corollary 4.6
For t [Ti, T2), have
0 U2 W
w:S > 0,
Ot M1
(x(
1r V P1 > 0,
ot M2 W WTT
w >
0
" W w <0
i
12
w
(T)
w((
V A3 <.
Proof. This a direct consequence of Lemma F11 Lemma F.6
and Definition F.
E5 Change Random-Toker Sub-Network
Lemma F.14. For & E [a] and E {ui]d we have
aG (1)
0 ( 1 2 of)m)))))
(m 1/m) oîm)).
0t
Thus,
(n) - s sō < mi (i(loovm 1/m) oîm) o7mm1
Proof. By Lemma C.I
and Definition F.I
we have
aG (p)
Ow
Ow
(e)
w(())
+
a
0t
01
d6
L
mi
(t)
ΣΣ>
Σ
Σ
w (1) (1i2))
Dig yiu
aj aj
Pgela
n
(c0)
b=
57

---

E
Σ
(1.i) (t.it) Pig 1/
((00m)))
+
yog mi
n
1 1==1
By Definition F.T and Corollary
H2
m1
m1
1 T, ( of ) n
Gi2 Vig ΣΣ 1-1h=11
aj j (w))) -0
Patta 50
(1.12)
n
io: (2)
1=
By Lemma and
E.9 and and Lemma
R.11 we have Ia T W w/ s 0 + 1/m) for μ #
| pl2 O(oîm). Thus,
(00()) s (L + oîm)).
Σ
(a.) (1.5) Plida %2
Ui, m1
Thus, by Theorem
E.8 for t E [Ti, T2]. we have
1// (G)
1 (0
ppmmi + ((0vm++0 +1/m)) [ ( dr
G00 (()))())
dr
Ør
1
mn mi(Loo 1/m) + aîm)
0
o,mm²
n
Corollary F.15. Fort E [T], we have
L
Σ
Σ
cIO(X! (p
$ (O()).
la:
X 1-
Proof. By Lemma
we have
Σ
P 00
4=1 l2: x ((m))
By Lemma F.14 for t E [T, T2l. (() = 0(1) for μ {Hildms. On the other hand, by the triangle inequality, we have
Σ
(X)) (p 11 Σ
Σ
(X)(X
11 la: x Elpek_
1= la: x Eime
Σ
[x())(x(()) - ci(t) (x([() ))}
(1,2)
ly x Efeeldim
Σ
Σ
(x)()
(Ti.d) (Tr.))
(2.2)
P 1q -
-
4= ly: Epe))
+ 3
((())(x(0)) , -x0) 1pp00// -
(2.6)
-
4=1 12: x Ee
<0()),
58

---

where the last inequality applies Definition F.1 we have
This implies that, for t E [Ti, T2). by Lemma F.14
(X,
(f,9) (p p S 01)).
la: X Efp 12-4
F.6 Change Score and Softmax Probability
Lemma F.16 (Change of score, complete version of Lemma 4.7 For 1 E [Ti, the attention scores are changing in the following way:
For p. E {μz, Pe), J # v, the query-kyy-corrllatinn score between the two target signals increases, while the query-key-correlation score between one target signal and the common token decreases, i.e.
0
T (TT
W 1r -
1 (T. mm))
1 L
0t
K
vm
0 (T
whl
Q / = 5 (1) 2
/
K
mm
L
The change of score satisfies:
oim (VmLofmmi
oovm m
with w μ v (:( p)
(T:)
=
max [.E((
For all μ E H21 N3}, E { the quere-key-corrreltton score changes as follows:
① T T wilh? (
1 _ o( m/m) T. +
(1) 1 do
m/m
m
and
# WLOT (t)
(w(TT((t), w - T (T) ((TT))
am o V m f mm
ny/mLommm,
om
o m m Vmafmmm
7 W ( W # - w w(Thi) 1
50
n/mLofmm
Proof. By Lemma C.4
we have
Ov w (
K
W Wo
4
Ōt
nym
y1
a
(v)ll3 (
[|
(V)1e()) waaT v(se)p1)0) (aD p (1)0T()))OO))
((()
1=1
(t)
v (f) k(tb0)) (v( w0 - (O) v(t P((end) (o M(K(1))) # (z)
nV m
9i
y:
K
(
j=1
1=1
m1
Σ
(t)
15ll (1() (1) gopk (w))) ((000())))))))))))))))))) - w ( vit.i)
(t.1)
9
Pi(en)
Th m
j=1
m
Σ
(1)
WOT (1.1) (1.2)
di P (uoOT (((.)()) - w -w To v(t.0) ((.0)))(())) M( (ight
(e.t)
+ (n)).
9:
1V m
EWEX (0)
Now, take μ 11,
1/2. By Theorem E.27
have ((T)) 2(1). A consequence of 0 Ti,Th] Thus, we have (12( ?(1).
ag (Pa)
Definition F.I and Lemma F.6 that
59

---

Further by Theorem E.27 we have (1)) = Ō mm) for R and then by Lemma F.14 we have (p) = O ) + Ō mi (L(o /o / 1mm) fm))))m))))
for t E (Ti, T2). Also,
((..)) (1.1) ofmma
Definition F.1 implies that G (/12) Σド= wi P((i.)) 2 2(11. Now, this yields
3p2 W ( w pn 1 O0 mg + o
Vm
On the other hand, by the analysis similar to the above, we obtain
0pg
W
K
w
Q
1
5( of
1 I +0
+0(
(11
1 10 oz
Ōt
Vm
vm
Next, prove the maximum change of the score, have
Ov
T W (t)
K
(+(1)
S 1 2(1
+ 7 o3vm
max p.M
Øt
Vm
mm
By Theorem F88
we have
w WOOT K
T
q 1 v
((TT wT (
<
+ Jn / ( mg))
{ [t) 10 dr
T,
mm
Jp
oi m v /mo m
< 0
VmLofmmi
Finally, for y {ui)/ma p E {ue]i we have
a 0t 1
WTT w
-6 of)))mm) f
1 0( nym
1 [ 0 ) +0 ( Vm
which implies that
1
W
((th) Q
(T)T
w <
hum L mmm
o oovn m m
Corollary E.17 (Change of softmax), For E [2), the softmax probability is changing in the following way:
For p,v (21 + the softmax probability between the two target signals increases, whereas the softmax probability between one target signal and the common token decreases. i.e.
i (t.i)
O
mmm)
Ot q-
Vm
i (t.b)
1 T am om
1
01 P Pat-p
Vm
1
For all μ E (pr,p2) p2); 7 E {uu]d the softmax probability between one target signal and random token changes follows:
i (1.6) 01 Pa
O 2 m)
1 +0
7 oo 3
nym
Ilm
Furthermore, we have
(t.1)
(0.5)
S 50(1113).
max
Pin.2
Pi
60

---

Proof.
Take μ = M1, = Pz. First of all, by Definition F. and Lemma F.16 we have
(1.1)T Pa-m
0(0 m)
I + (1.0)
1
1 + +0 +
T 7.000 20vmm
ōt
Vm
n/m
Thus, by Lemma G.I
and Lemma
F.16 for i 11, we obtain
0 (t.i)
1 O( (g
ogm)
1
Ø1 wh
T?²
① (t.i)
θ [ mdm)
1 2
0t P a
On the other hand, for {us we have
① (t.4)
1
+0
- (t)
og v
0l g
mm
T
Ta
Finally, by Lemma F.16 we have
W w P v W w / <0
(mLLofmm (
agm
oōvm
max
v
mmofmm
Thus, by Lemma G.2 we obtain
(t.i)
(Ti.))
max
Pi
P1
omm
o
ogm 08yn m
)) 0(1/23).
L
-01/L3))
mLLofmmm
mojmm,
(LLoimm mož
Corollary F.18. Fort E Wi,TTh we have
El,=1 Ei (x)()
Opi
maXig[n]
at
So( Σ 9
Σn= Em==1 a/ w0 as2
(2)
(E(N)
Proof. This a direct consequence of Definition
Lemma F.10 and Corollary F.17
Change Self-Correlation of Key/Query- Transfrrrreet data
Lemma F.19 (Change of self-correlation). For n1 E {1: we have
obym
T (-(0) V
W (T)T (T))
Q
Vmo
ooym
T w
( w W -
W (T) W(Th)
T1)
mmafmmm
We prove the result for v (T((()) w-() we have
# and the proof for v (T( w(t) K # is similar. By
Proof.
Lemma
dv W
(()T w P
Øt
1
Σ
9i Vi
(t)
me
agv T WoOT k(t.) (V(..) (0)) -w (OT vla. vlo (1,8) (c.1)
p((i)) Pi(i.)) Pi(in))
n/m m
(v)
j=l
1
Σ
(0)
(T k-(.0) K-(t.0digg diag (w - n v(tt) (1.)
9,
P(((e)) Pi(eer
mmm
(VEX6)
j=1
61

---

A simple result from Lemma F.16 is that |
W ml S O(oîvm) for l E [Ti, T2] and
u.v. (ui)er Therefore, by Definition
F.1
Ov Th.The T ( w-(t)
1
1
1 og @mvmlogm)
Ōt
Vm
which implies
[" δ T1
( i (
d8ym \mmffmm
y Twe @ H- (=
w(T)
(T)) M
o m dr dr
m
where the inequality follows from Theorem F.8 F.8 and Definition
F1
F8 Small Loss Achieved
Theorem E.20. Define T = min L = 0(1/poly(m))) Then, T E [2],
Proof. The following results altogether show that Phase 2 can last for at least ((poly(m)) time:
Lemma
F19
proves that the change of K, Q self-correlation as follows:
W (G)T
K μ w w Ö(oğvm)
max p.r
(T) W v p T o ( w = Ō
max
Q
Corollary proves that
R9
Oj w/ azz (0
(1)
S m(o7mmm)).
Corollary F.12
proves that
(1) 23 aj st W k
max
_ Σ 10 tE(m)
S o((/L)
m1 -=
Gjj "j ae (1)
(t)
Corollary F.15
proves that
Σ
Σ
(x(() (t.i)
( PP ktl2 50()).
4=1 by x Elax
Corollary E.17
proves that
(t.t)
(0,4)
50(1/1²).
max (E
Ph -
Plgd2
Corollary E18
proves that
(t.1)
Eli=1 Eli=1 (x(()) it
[ax;g[1]
S <((1)
Σ 9
(1)
(1)
Σイ
-iz=1
ajj 1 agj w
Lemma F5 implies that
1
E 9i
2.
2 EiELo 9i
(1)
62

---

Lemma
F.6
implies that the gradients Eier Di
(
for I {11, 12, 13, 14] satisfies
Σ iEL.
I min Σ
- max Σ
Σ
Σ
Σ
9i
Di
91
KEl2
iEls
VEL
EEI
tE/201,U4
Lemma
F.6 and Corollary
F.17 proves that yı F(
yı F 2 C.
The above shows that all the requirements needed to satisfy the definition of Phase (Definition can hold for at least (ppol(mm) time. Thus, T2 Ti 2 (2(poly(m)). Further, by Theorem F.8 we have that T O(1 O(1/poly(m) implies (poly(mn))
F.
Proof of Theorem 3.2
Proof. This proved in Corollary F.1 and Lemma F.16
E.10 Proof of Theorem 3.3
Proof. This proved as a direct consequence Lemma F.6 and Theorem F.20
For the generalization loss, since the training loss satisfies Ī()) I 1 1/poly(m), for each class
{I, Iz, I3. J there exists sample X3 such that {(X}) 1 1/polyymm) (m). Note that by Definition F.]
the random tokens only contributes O(1) in F(T) Thus, given fixed new sample X D,
have (x) p(2") (Xf)l I O(1) which implies ((yF(t)) 1/poly(m)) Since this
holds for all X D, we have ((r)) S 1/poly(m).
G Auxiliary Results
The gradient of p(x): = Softmax(r); for x E R":
0p(z); 0 exp(x,)
E exp(et) exp(x.)
expla;) expla,)
Ex exp(xx]
=-<())p((z)
itj
da; drj 2k= exp(Ix))
EE exp(ix) =p(x);(1 - -p(x); i i=j
exp(xx)
Ee exp(ex)
p(z) = -) - -p(x);)
J(p(x)) = diag(p(x)) - p(x))((1))
(17)
Lemma G.1 (Gradient of softmax). Let s(t) E Rd be differentiable and p(s) = Softmax(s)
Denote ps(s) Softmax(s)) (8); Then
Op(s(t)) 0p(s) 0s(t) (diag(p(s)) - p(s)p(s) Os(t)
Ø1
Os Øl
Ōt
Proof. By the chain rule and the gradient of softmax in Equation 11 obtain
Öp(s(t)) 0p(s) @s(t)
(diag(p(s)) p( p(s)p(s) Os(() (p(s)) Os(t)
Ōt Ös Ot
Øt
Lemma G.2 (Perturbation of softmax). Let R¹ and p(s) Softmax(s)) Denote p((s) Softmax Consider small perturbation IR! to s. Then
p(s +18+5) p(s) = (diag(p(s)) - p(s)p(s) + E.
where || O(
Proof. By Taylor's expansion theorem softmax and the gradient of softmax in Equation 17 we
have
p(s + (1++))-(p)) J(p(s)) + & (diag(r(s)) - p(s)p(s) T) )e + E,
where [| O(1)=11).
63

---

H Probability
Lemma H.1 (Bernstein's 's inequality for bounded random variables). Assume Zi.... Zn are i.i.d. random variables with E[Z 0 and [Za & M for all E [r]] almost surely. Let Z Σi=1 Zi. Then, for all
12/2
t
2M })
PIZ S exp
S exp
-min
Σド= RIZ3I + M1/3
2 2
18 [[
which implies with probability at least 1 - 6,
1 + 2M log 5
S 2
E B[Z]] log 8
iid. \O, I /m), we have
Lemma H.2. For w1, W2
IR" with W1, w2
4
2
4
IP
w1 112
11
log
log
8,
172
771
4
2
4
IP
[w),
w2)1
2
log
log
8.
m
m
Proof. We first have
wh
1.
Note that wf,; is sub-Gamma random variable with parameters m= A). Thus, by Bernstein's inequality,
2 4
P
w| 112 -
log
log
S
m
m
Next,
B[ w2)l = E
Σ
w1, W2, = 0
By Bernstein's 's inequality, we obtain
2
Kwn, w222 2
log
log
I 6.
m
64

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: [NA]
65

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
66

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (nttps://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7 /ttppp.cc/pbllc/guudeeee http://oddeuumbiinoooooooooooooe
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
67

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification: The experiments are done with synthetic data of small scale and the behavior of the experiment results are pretty consistent.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2- 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
68

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httppp///euuuip...ccaaaaa///tttt
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
69

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
70

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
71


Reviews:
[{"review_id": "6zSftEZmEB", "reviewer": "Reviewer_G3j5", "summary": "This article delves into the gradient flow dynamics for detecting word co-occurrence, demonstrating that the gradient flow approach can achieve minimal loss. The training process commences with random initialization and can be delineated into two distinct phases.", "questions": "- In line 151 and 152, it is confusing why concentration theorems lead to the specific probability in (i) of Assumption 2.3.\n- Lack of explanation for $\\langle w_{j_1}^{(t)},w_{j_2}^{(t)} \\rangle$ in line 194.\n- It is not obvious why \"the samples with only one target signal may be classified incorrectly as co-occurence\" in line 282.\n- The notation in line 169 is somewhat misleading.", "limitations": "The semantic mechanism this article focuses on is not general enough in NLP and the the results cannot give instruction useful enough to guide the training process of Transformer.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This article noticed an interesting phase transition during training in this special setting and demonstrates it with solid calculation and experiments.\n- A new property of gradient flow is noticed and contributes to prove near minimum training loss together with the analysis of softmax.", "weaknesses": "The setting of empirical experiments is also simple and ideal and readers may have no idea if this is a general phenomenon during training for detecting word co-occurrence.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer very much for your time and efforts on providing helpful review comments.\n\n**Comment:**\nThe setting of empirical experiments is also simple and ideal and readers may have no idea if this is a general phenomenon during training for detecting word co-occurrence.\n\n**Response:** We clarify that the experiment presented in the paper is used to guide and validate our theoretical analysis, and hence it is designed to match our theoretical setting. Our goal here is to understand the training dynamics of transformers under **analytically tractable** setting so that it can help us understand the mechanism behind self-attention and transformers. \n\nRegarding the reviewer's question on the generality of the phenomenon, we attached a pdf file that contains a preliminary experiment that we worked out. Figure 1 shows that for more general settings with multi-layer transformers, our main theoretical findings still hold: (i) the loss converges to the global minimum (i.e., with zero value) despite being highly nonconvex; (ii) the training process achieves correct detection with zero classification error as seen in phase 1 of our characterization, and the loss value continues to decrease to zero due to enlargement of the classification margin as seen in phase 2 of our characterization. We will continue to work on more experiments.\n\nWe also point out that in multi-layer transformers, it is hard to interpret the meaning of the softmax attention in the middle layers since unlike the one layer case, the inputs to the middle layers are changing from iterations to iterations. Thus, it is hard to find a fixed subspace to project on and study its changes. \n\n**Comment:**\nIn line 151 and 152, it is confusing why concentration theorems lead to the specific probability in (i) of Assumption 2.3.\n\n**Response:** Thanks for the question. Such an assumption helps to simplify our notations throughout the paper so that our main results can have simpler version and it is easy for readers to digest the main insight of the result.\n\nThose specific ratios can be satisfied with high probability if the size of the training set $n$ is large enough. This can be achieved by applying Hoeffding's inequality as follows. Take $I_1$ as an example. We have $\\mathbb{P}\\Big[|\\frac{1}{n} \\sum_{i=1}^n \\mathbb{I}(X^{(i)} \\in I_1) - \\frac{1}{2}| \\geq \\sqrt{\\frac{\\log(2/\\delta)}{n}}\\Big] \\leq \\delta$. \n\n**Comment:**\nLack of explanation for $\\langle w_{j_1}^{(t)},w_{j_2}^{(t)} \\rangle$ in line 194.\n\n**Response:** Thanks!\nThis is the inner product of the neuron weights between $j_1$-th and $j_2$-th neuron. \nThis term appears in calculating the update of $\\nu^\\top W_V^{(t)\\top} W_V^{(t)} \\mu$ and is needed to make the dynamical system complete. We will add this explanation to the revision.\n\n**Comment:**\nIt is not obvious why \"the samples with only one target signal may be classified incorrectly as co-occurrence\" in line 282.\n\n**Response:**\nThis is because during the beginning of training, the linear MLP layer will positively align with $\\mu_1, \\mu_2$, i.e., $G^{(t)}(\\mu_1), G^{(t)}(\\mu_2) > 0$, while for the common token $\\mu_3$, we have $G^{(t)}(\\mu_3) \\approx 0$. \nAt the same time, the linear MLP will also output something near zero for random tokens. \nThus, for those samples $X$ with only $\\mu_1$ or $\\mu_2$ (i.e., $X \\in I_2 \\cup I_3$), we have $f^{(t)}(X) > 0$ which is incorrect as the samples in $I_2$ and $I_3$ has label $-1$. \n\n**Comment:**\nThe notation in line 169 is somewhat misleading.\n\n**Response:**\n$p_{q\\leftarrow \\nu, k \\leftarrow \\mu}^{(i)}$ is the output of the softmax attention given input $X^{(i)}$ when the query is $\\nu$ and key is $\\mu$. \nIn addition, $l(i,\\mu)$ denotes the index in $X^{(i)}$ such that $X^{(i)}_{l(i,\\mu)} = \\mu$. We will add this explanation to the revision.\n\nWe thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.", "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": "ZKujkeWAoO", "reviewer": "Reviewer_ritc", "summary": "This paper studies the training dynamics of a single hidden layer transformer network (self-attention + linear MLP) trained on a binary word cooccurrence task. Specifically, given a data matrix $X \\in R^{d \\times L}$ representing L \"words\" (each column of X is a word vector of dimension d), the model must output +1 if words 1 and 2 both occur in X, and -1 otherwise. The paper shows that a transformer layer is able to learn this task, and that the training occurs in two stages: First, the linear MLP layer learns to classify data points correctly by positively aligning with the embeddings for words 1 and 2 (but without making large changes to attention matrices). Second, it drives the loss down further by using the attention matrices to positively correlate q,k,v for words 1 and 2, and anti-correlate the q,k,v for a common word (denoted word \"3\" in the paper) relative to words 1,2.  After these phases, both the training and generalization losses go to zero (as long as embedding dimension is large enough).\n\nOverall, I found the results interesting and insightful, though not very surprising, and the practical implications of these results were not very clear to me. Thus, I currently recommend weak accept. Importantly, my primary research area is not learning theory, so my knowledge of the related work is relatively limited, and thus my review confidence is relatively low.", "questions": "- What other tasks (beyond word cooccurrence) do you think could be analyzed with this methodology?\n- What are the implications of this result to more complex/realistic tasks, like next token prediction?\n- If mini-batch Adam is used during training, do the two phases still occur?\n- Can you add more details about the experimental setup to the main paper?\n- Can you add more discussion about the automatic balancing of gradients, and its significance, in a more central part of the text (e.g., section 3, not 4)?", "limitations": "Yes, limitations have been discussed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is interesting to see that the training dynamics for this word cooccurrence task can be analyzed rigorously, with relatively few assumptions.\n- The theoretical results are validated with a nice synthetic experiments, that demonstrates that the two phases predicted by the theory do occur in practice.", "weaknesses": "- This word cooccurrence task is very simple, and thus it is not surprising that a single transformer layer can easily learn it.\n- Only full gradients are considered, whereas transformers are typically trained with mini-batch Adam(W).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer very much for your time and efforts on providing helpful review comments. \n\n**Comment:**\nThis word co-occurrence task is very simple, and thus it is not surprising that a single transformer layer can easily learn it.\n\n**Response:**\nWe agree that this is a simple task from a representational perspective. \nHowever, our primary goal here is to use this setting as an **analytically tractable** case and develop theoretical techniques to understand the *training dynamics* of transformers, which is a highly non-trivial task considering the complicated nature of self-attention and transformers. \n\n**Comment:**\nOnly full gradients are considered, whereas transformers are typically trained with mini-batch Adam(W).\n\n**Response:**\nRegarding full gradients, since we are using gradient flow, it is more convenient to analyze the full gradients. We can indeed use gradient descent and relax the full gradients to consider stochastic gradients. However, this will create some unnecessary hassle and won't change the essential nature of two-phase training dynamics of transformers. \n\nRegarding Adam or its variant, it is indeed an interesting direction. However, its analysis would be more complicated due to the adaptive momentum terms in the iterative update. The solution that transformers converge to can also have different implicit bias from SGD. Thus, we will leave this as future work.\n\n**Comment:**\nWhat other tasks (beyond word co-occurrence) do you think could be analyzed with this methodology?\n\n**Response:**\nThank you for asking. \nOur proof techniques open up for a great possibility of settings to analyze such as NLP tasks with one hot embedding. For example, one popular way of modeling natural language is by modeling the language sequences by Markov chains. Our techniques can be used to study the training dynamics when we want the transformer model to predict the next input given the previous tokens in the sequence. \n\n**Comment:**\nWhat are the implications of this result to more complex/realistic tasks, like next token prediction?\n\n**Response:** Our work could have some implications applicable to more complex tasks including next-token prediction. (i) The training can experience two or more training phases, where each phase captures one salient change of certain attention scores or MLP weights. (ii) The gradient of well-designed loss function can facilitate classification and enlarge the margin by driving attention scores between query and relevant key tokens to change properly during the training process.\n\n**Comment:**\nIf mini-batch Adam is used during training, do the two phases still occur?\n\n**Response:**\nWe expect the training phases of Adam will exhibit some differences from gradient descent considered in our work. From a generic nonconvex optimization perspective, (Shuo \\& Li) showed that Adam can be viewed as normalized gradient descent in $\\ell_\\infty$ geometry whereas gradient descent operates on $\\ell_2$ geometry. \nThis is due to the coordinate-wise normalization operation uniquely occurring in Adam but not in GD. We expect that the loss with transformers trained by Adam will somewhat also follow such an observation.\n\nXie, Shuo, and Zhiyuan Li. \"Implicit Bias of AdamW: $\\ell_\\infty $-Norm Constrained Optimization.\" Forty-first International Conference on Machine Learning.\n\n**Comment:**\nCan you add more details about the experimental setup to the main paper?\n\n**Response:** Thanks for the suggestion. We will make more space in the main paper and add the experimental setup.\n\n**Comment:**\nCan you add more discussion about the automatic balancing of gradients, and its significance, in a more central part of the text (e.g., section 3, not 4)?\n\n**Response:**\nThank you for the suggestion. We will make the change as the reviewer suggested. \n\nWe thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.", "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": "lnuP3zy3mU", "reviewer": "Reviewer_cFaT", "summary": "This paper investigates the training dynamics of a single-layer transformer followed by a single MLP layer on a synthetic binary classification task, where the objective is to identify the co-occurrence of two specific tokens in the input sequence. They analyze the gradient flow dynamics for the case that all the attention parameters (key, query, value) and the linear layer all trainable and show that the model can achieve low loss despite the non-convexity of the objective. They identify two phases in the training, 1) the MLP aligns with the two target tokens at the start of the training and the model learns to classify all the samples correctly, 2) the attention and MLP parameters update to increase the classification margin and drive the loss to zero. They also run a small scale numerical experiment in their synthetic setup tp confirm their analysis.", "questions": "1. As far as I understand, in your synthetic task,  in the first phase of training, effectively only the MLP weights are learning the task. That is, the model can achieve 100% accuracy only by aligning the MLP weights with the relevant tokens $\\mu_1,\\mu_2$. So, the attention layer is not needed for identifying the co-occurrence of the tokens in this setup?\n\n2. Can you also report validation and accuracy plots in your synthetic experiments? Does the validation loss decay at the same rate as the training loss as stated in Thm 3.3?\n\n3. Regarding the proof sketch:\n\n    a) The alignment of parameters with the target tokens is discussed in the main body. Can you also clarify how the gradients related to irrelevant tokens evolve? In particular, regarding the tokens that do not appear in the training set (since $nL\\leq d$), does the model learn not to attend to them at test time?\n\n    b) I find it confusing that the softmax output remains close to $1/L$ long in the training (line 320) and assigns uniform attention to all tokens in the sequence. Does this statement hold for all training samples? If yes, then how does the model learn to attend to the target tokens?", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper makes no restricting assumptions on the weights of the transformer model and performs the analysis on the joint optimization of all the parameters.\n\nAlthough the paper and its proof are notation-heavy, the authors have broken down the complexity of the proof and notation in the main body to clarify the steps needed to prove the results.", "weaknesses": "There are some restrictive assumptions on the synthetic data model: The vocabulary set $d$ is considered to be larger than the number of training tokens, which is not the case in realistic setups. Thus, some tokens are not visited at training time. Also, they assume, apart from the two target tokens, the remaining tokens appear at maximum once in the training set.  \n\nThe proof outline in the main body helps in understanding the high-level steps involved. However, it could still benefit from additional clarifications on some intermediate steps. For instance, in phases 1 and 2, it's mentioned how the alignment of the MLP with the target tokens $G^{(t)}(\\mu_{1,2})$ behaves during training. However, it's not clear how this connects to the evolution of the attention scores in phase 2 as stated in Lemma 4.7.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer very much for your time and efforts on providing helpful review comments. \n\n**Comment:**\nThere are some restrictive assumptions on the synthetic data model: The vocabulary set $d$ is considered to be larger than the number of training tokens, which is not the case in realistic setups. Thus, some tokens are not visited at training time. Also, they assume, apart from the two target tokens, the remaining tokens appear at maximum once in the training set.\n\n**Response:** Thanks for the question. In fact, we can relax both assumptions by letting the remaining tokens be uniformly randomly sampled from the vocabulary. Our proof will still be valid. More specifically, our proof holds as long as the number of each individual random token is much less than the number of the signals. For example, if we uniformly sample the irrelevant tokens from the vocabulary, then in expectation each irrelevant token will only appear $\\frac{nL}{d}$ times in the training set which is much less than $n$ (the number of times each signal appears in the training set) if $d \\gg L$. \n\n**Comment:**\nIn phases 1 and 2, it's mentioned how the alignment of the MLP with the target tokens $G^{(t)}(\\mu_{1,2})$ behaves during training. However, it's not clear how this connects to the evolution of the attention scores in phase 2 as stated in Lemma 4.7.\n\n**Response:**\nThis can be seen from the update of the score in the gradient flow dynamical system which is in Appendix C. Take the score between $\\mu_1$ and $\\mu_2$ for example, rearranging the terms, we have\n\\begin{align*}\n    &\\frac{\\partial \\mu_1^\\top W_K^{(t) \\top} W_Q^{(t)} \\mu_2}{\\partial t} \\\\\\\\\n    &= \\frac{1}{n\\sqrt{m}} \\sum_{i=1}^n g_i^{(t)} y_i \\sum_{l=1}^L \\mu_1^\\top W_K^{(t) \\top} K^{(t,i)} \\cdot \\textnormal{diag}\\left( G^{(t)}(X^{(i)}) - (G^{(t)}(X^{(i)}))^\\top p_l^{(t,i)} \\right) p_l^{(t,i)} x_l^{(i) \\top} \\mu_2 \\\\\\\\\n    &+ \\frac{1}{n\\sqrt{m}} \\sum_{i=1}^n g_i^{(t)} y_i \\sum_{l=1}^L \\mu_2^\\top W_Q^{(t) \\top} q_l^{(t,i)} p_l^{(t,i) \\top} \\cdot \\textnormal{diag}\\left( G^{(t)}(X^{(i)}) - (G^{(t)}(X^{(i)}))^\\top p_l^{(t,i)} \\right) X^{(i) \\top} \\mu_1 \n\\end{align*}\nSince $G^{(t)}(\\mu_1) = \\Theta(1)$ after phase 1, we have $G^{(t)}(\\mu_1) - (G^{(t)}(X^{(i)}))^\\top p_l^{(t,i)} = \\Theta(1)$. We can use this to show that $\\frac{\\partial \\mu_1^\\top W_K^{(t) \\top} W_Q^{(t)} \\mu_2}{\\partial t}$ is positive in Lemma 4.7 and thus the score between $\\mu_1$ and $\\mu_2$ is increasing. \n\nWe will add these clarifications in our proof outline. We will also try to enrich our proof outline by adding clarifications on some intermediate steps.\n\n**Comment:**\nAs far as I understand, in your synthetic task, in the first phase of training, effectively only the MLP weights are learning the task. That is, the model can achieve 100\\% accuracy only by aligning the MLP weights with the relevant tokens $\\mu_1,\\mu_2$. So, the attention layer is not needed for identifying the co-occurrence of the tokens in this setup?\n\n**Response:** Although MLP can achieve full accuracy, only MLP layer in phase 1 does not achieve a good performance margin, i.e., the loss is not vanishing yet, and such a model can easily make wrong classification when data have noise. The important role that attention layer plays in phase 2 is to enlarge the classification margin together with MLP so that the loss approaches zero and the trained model can generalize well. \n\n**Comment:**\nCan you also report validation and accuracy plots in your synthetic experiments? Does the validation loss decay at the same rate as the training loss as stated in Thm 3.3?\n\n**Response:** In the attached pdf file, Fig. 2(a) provides the accuracy plot, where both the training and test errors become zero after certain number of training steps. Fig. 2(b) provides the validation (i.e., test) and training losses, and they decay at similar rate.\n\n**Comment:**\nThe alignment of parameters with the target tokens is discussed in the main body. Can you also clarify how the gradients related to irrelevant tokens evolve? In particular, regarding the tokens that do not appear in the training set (since $nL\\leq d$), does the model learn not to attend to them at test time?\n\n**Response:** The alignment of the random tokens is much smaller than the signals, and remain to be small and bounded throughout the training process since they don't occur as much in the training set. \nIn our proof, we handle the random tokens all together no matter if they occur in the training set or not. \nTheir effect can be upper bounded during training (Theorem E.15 for phase 1 and Appendix F.5 for phase 2). \nSo the answer to your question is yes -- the model does learn not to attend to those tokens not in the training set at test time. \n\n**Comment:**\nI find it confusing that the softmax output remains close to $1/L$ long in the training (line 320) and assigns uniform attention to all tokens in the sequence. Does this statement hold for all training samples? If yes, then how does the model learn to attend to the target tokens?\n\n**Response:**\nSorry about the confusion. In line 320, we mean to explain how we prove convergence. In fact, there are 2 steps to show convergence: (1) show that the training loss can decrease when the softmax outputs remain to $1/L$, which is what we explain in the proof outline; (2) show that with the deviation in the softmax output from $1/L$, the loss value will still decrease, which is not explained in the proof outline (we will add this part to the proof outline to avoid confusion). The outputs of the softmax attention indeed change as stated in Lemma 4.7.\n\nWe thank the reviewer again for your comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly increase your score. Certainly, we are more than happy to answer your further questions.", "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:
Dear Reviewers,

Please find all the mentioned additional experiment results in rebuttals (to Reviewer G3j5 and Reviewer cFaT) in the attached pdf. 

Thank you.

Best,

Authors

Meta-review: In this paper the authors propose using gradient flow to investigate the training dynamics of a one-layer transform in a case study of word co-occurrence recognition.  The analysis does not impose commonly employed simplifications in literature such as weight reparameterization, linear attention, special initialization and lazy regime.  By analyzing the gradient flow the authors identify a two-phase training phenomena in the investigated transformer in the word co-occurrence recognition task.   The major novelty of the work is that the authors introduce a new proof technique that may benefit the machine learning community for analysis of training dynamics.  The major weakness of the paper which all reviewers agree upon is that the proposed gradient flow analysis is only conducted on the word co-occurrence task which is a very simple task.  Therefore, it is not clear whether the observed two-phase training dynamics can be generalized to a broader scope of applications.
Decision: Accept (poster)