Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis

Authors: hongru  yang
bhavya  kailkhura
zhangyang  wang
yingbin  liang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 15:12:54

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.

Main paper: 
 introduction: 
 Ever since the invention of self-attention [VSP + 17], transformers have become a dominating backbone architecture in many machine learning applications such as computer vision [DBK + 20, LLC + 21] and natural language processing [DCLT18] . Nowadays, ChatGPT and GPT-4 [Ope23] have demonstrated astonishing abilities in many areas such as language understanding, mathematics and coding, which have sparked artificial general intelligence [BCE + 23]. In the meantime, there has been a burgeoning development of large language models (LLMs) [TLI + 23, MH23, ADF + 23] as well as multi-modal models [Tea23] .Despite the huge empirical success, theoretical understanding of why a pre-trained language model can possess such impressive performance has been significantly lagging behind. Some previous efforts have been made in understanding the capacity and representational power of transformers [EGKZ22, LAG + 22, ZPGA23, SHT23, BCW + 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, a 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 [JSL22, LWLC22, TWCD23, TWZ + 23, LLR23, TLZO23, TLTO23, ZFB24, 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 dynamics will focus on a basic problem of recognizing cooccurrence 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 [MS99, Dam18] predicts the next word based on co-occurrence of multiple words. Consider the following scenario: if the task for the language model is to read a paragraph describing a children and then answer some questions, say, "Is Bob eating a banana?". In order to answer the question correctly, the model must be able to detect the co-occurrence 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-attention module followed by a linear multi-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 to 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 1, 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 to 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.Transformer representational power. Several previous works have studied the expressiveness of transformers. One line of work was from a universal approximation perspective and thus provided the existential results [YBR + 19, WCM22, PBM21, ZPGA23, LAG + 22, BCW + 23]. As a separate view, [EGKZ22] showed that a single attention head can represent a sparse function over the input sequence with sample complexity much smaller than the context length. [ZZYW23] studied the approximation and generalization performance of transformers in in-context learning. [SHT23] proved that transformers can represent certain functions more efficiently than MLPs.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. [JSL22] proved that transformers can learn spatial structures. [LWLC22] studied how a shallow transformer learns a dataset with both label-relevant and label-irrelevant tokens. [TWCD23] studied a next-token prediction problem and showed that self-attention behaves like a discriminating scanning algorithm. [LLR23] analyzed a layer-wise optimization scheme on how transformers learn topic structures. [TLZO23, TLTO23, VDT24] studied a setting where transformers can learn a 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, [ZFB24] provided the training dynamics of linear attention models, [HCL23] characterized the training dynamics of softmax transformers, and [CSWY24] studied a multi-task linear regression problem with a multi-headed softmax transformer. Further, [LWL + 24] focused on nonlinear selfattention and nonlinear MLP over classification tasks in in-context learning. [WLCC23] proved the convergence of transformers via neural tangent kernel. [NDL24] showed that two-layer transformers can learn causal structure via gradient descent. [CL24] developed algorithms for provably learning a multi-head attention layer. [HWCL24] studied how transformers learn feature-position correlation.This paper studies a different problem of detecting co-occurrence of words via transformers. Such a setting has not been considered in the literature. More importantly, the previous studies of training dynamics of transformers have adopted various assumptions/simplifications such as weight reparameterization, 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.2 Problem Setting Notations. For a vector v ∈ R d , we use diag(v) to denote a diagonal matrix with v being the diagonal entries. When we subtract the vector v by a scalar a, we subtract each entry of v by a, i.e., v -a ∈ R d and (v -a) i = v i -a. We use Ω, Θ, O to hide polylogarithmic factors. 
data model: 
 Definition 2.1 (Data distribution). Given a set of orthonormal vectors {µ i } d i=1 as word embedding, let µ 1 , µ 2 ∈ R d be two target signals whose co-occurrence needs to be detected by the model, and let µ 3 ∈ R d be a common token vector. A data entry (X, y) ∈ R d×L × {±1}, where X = [x 1 , x 2 , . . . , x L ] consists of L tokens, is generated by the distribution D as follows:1. Uniformly randomly select an index i 3 ∈ [L] and set x i3 = µ 3 .• With probability 1/2, set y = 1 and uniformly randomly select two indices i 1 ̸ = i 2 ∈ [L] \ {i 3 } and set x i1 = µ 1 , x i2 = µ 2 . For i ∈ [L] \ {i 1 , i 2 , i 3 }, set x i = Uniform({µ i } d i=4 ). • With probability 1/6, set y = -1 and uniformly randomly select one index i 1 ∈ [L] \ {i 3 } and set x i1 = µ 1 . For i ∈ [L] \ {i 3 , i 1 }, we set x i = Uniform({µ i } d i=4 ). • With probability 1/6, set y = -1 and uniformly randomly select one index i 2 ∈ [L] \ {i 3 } and set x i2 = µ 2 . For i ∈ [L] \ {i 3 , i 2 }, we set x i = Uniform({µ i } d i=4 ). • With probability 1/6, set y = -1. For all i ∈ [L] \ {i 3 }, we set x i = Uniform({µ i } d i=4 ).In summary, there are 4 types of data: (1) both µ 1 , µ 2 appear, (2) only µ 1 appears, (3) only µ 2 appears, and (4) neither µ 1 nor µ 2 appears. We denote the set of indices of the above 4 different types of data by I 1 , I 2 , I 3 , I 4 ⊆ [n]. We further define R = {µ i } d i=4 . For simplicity, our data distribution assumes µ 1 , µ 2 , µ 3 appear only once in a data entry. The occurrence probability of each type of data is chosen in the above way to make the distribution label-balanced. We assume there is a fixed set of orthonormal vectors as word embedding, which is analogous to the one-hot embedding of a 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 a logical AND problem. Remark 2.2. 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 a 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". If 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".For simplicity of our analysis, we make the following assumption on our training data set.Assumption 2.3. The training set satisfies: (i) |I1| n = 1 2 and |I2| n = |I3| n = |I4| n = 1 6 ; and (ii) for all i 1 , i 2 ∈ [n], l 1 , l 2 ∈ [L], if X (i1) l1 , X (i2) l2 / ∈ {µ 1 , µ 2 , µ 3 }, then X (i1) l1 ̸ = X (i2)l2 , i.e., all irrelevant words are different.The first assumption can be approximately satisfied with high probability given the total number n of samples is large enough. Such an assumption can be removed by applying the standard concentration theorems. The second assumption implicitly assumes nL ≤ 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.Consider a training set {(X (i) , y i )} n i=1 with n training samples. Each data point X (i) ∈ R d×L contains L tokens, i.e.,EQUATIONL ]. We consider the transformer model with a self-attention module followed by a linear MLP:EQUATIONwhere the query matrix W Q ∈ R m×d , the key matrix W K ∈ R m×d , the value matrix W V ∈ R m×d , the hidden-layer MLP weights W ∈ R m1×m (with w ⊤ j being the j-th row of W ), and the output-layer weights of the MLP a ∈ R m1 . We define the linear MLP function of the transformer to be G(µ) = m1 j=1 a j w ⊤ j W V µ. We now introduce some shorthand notationsK = W K X, Q = W Q X, V = W V X and let k l = W K x l . Notice that K = [k 1 , k 2 , . . . , k L ].We further extend this shorthand to q l and v l . We also define functions k(µ) = W K µ, q(µ) = W Q µ, v(µ) = W V µ.We introduce the shorthand for the score vector s l :=X ⊤ W ⊤ K W Q x l √ mand the attention vector p l := Softmax(s l ). For the attention vector, if µ, ν ∈ X (i) , let l(i, µ), l(i, ν) be the indices such thatEQUATIONl(i,ν) = ν, and we define p(i) q←µ,k←ν := p (i) q←l(i,µ),k←l(i,ν) := Softmax X (i)⊤ W ⊤ K W Q µ √ m l(i,ν).Initialization. We initialize a j i.i.d.∼ Uniform(±1) and the value of a is fixed during training.[W, W V , W K , W Q ]. We initialize [W, W V , W K , W Q ] by W i,j i.i.d. ∼ N (0, σ 2 1 ) and (W V ) i,j , (W K ) i,j , (W Q ) i,j i.i.d.∼ N (0, σ 2 0 ). Training. We adopt the cross-entropy loss l(x) = log(1 + exp(-x)). The gradient of the crossentropy loss is given by l ′ (x) = -1 1+exp(x) and we define g(x) = 1 1+exp(x) . The model is trained by gradient flow to minimize the following empirical loss:EQUATIONSimilarly, we define the generalization loss L := E (X,y)∼D ℓ(yF (X)). We introduce the parameter condition that we take throughout the entire analysis and proofs.Condition 1. We make the following parameter choices in our analysis:• The embedding dimension and network width satisfy m ≥ Ω(max{m 1 , L 2 }) and m 1 ≥ Ω(1).• The network weight initialization variance satisfiesσ 0 = 1 Θ( √ Lm) ; and σ 1 = 1 Θ( √ m1) .• The number of training samples and tokens satisfy n ≥ Ω(L 2 ) and L ≥ Ω(1).• The failure probability satisfies 1/δ ≤ 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 row 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, W V , W K , W Q directly, it is sufficient to analyzing their impact on inputs, i.e.,X ⊤ W ⊤ Q W K X and a ⊤ W W V X, which are sufficient to compute F (X; W, W V , W K , W Q ).Based on this observation, we formulate two differential equations to keep track of w(t)⊤ j W (t) V µ and ν ⊤ W (t)⊤ K W (t)Q µ (with respect to t) for all µ, ν ∈ {µ i } d i=1 (See Equation (6) and Equation (7) in Appendix C). We further include additional equations to keep track of w(t) j1 , w (t) j2 , ν ⊤ W (t)⊤ V W (t) V µ, ν ⊤ W (t)⊤ K W (t) K µ, and ν ⊤ W (t)⊤ Q W (t)Q µ to complete the system. Intuitively, the additional equations keep track of the shape of the neurons and the word embedding after W V , W K , W Q transform. Although the dynamical system does not directly track the softmax, the softmax probability can be calculated via the scores ofν ⊤ W (t)⊤ K W (t)Q µ. The full dynamical system is presented in Appendix C. Then the training dynamics can be characterized by analyzing these differential equations (see a proof outline in Section 4).In the next two theorems, we present our characterization of the training process into two phases.Theorem 3.1 (Phase 1). With probability at least 1 -δ over the randomness of weight initialization, there exists a timeT 1 = O(1/m) such that • The linear MLP functions satisfy: G (T1) (µ 1 ) ≥ Ω(1), G (T1) (µ 2 ) ≥ Ω(1), G (T1) (µ 3 ) ≤ -Ω(1).• All training samples are correctly classified:y i F (T1) i > 0 for all i ∈ [n]. • For t ∈ [0, T 1 ], all dynamical variables w (t) j1 , w (t) j2 , ν ⊤ W (t)⊤ K W (t) Q µ, ν ⊤ W (t)⊤ V W (t) V µ, ν ⊤ W (t)⊤ K W (t) K µ, and ν ⊤ W (t)⊤ Q W (t)Q µ are close to their initialization values.• Training loss is still large: L (T1) = Θ(1).In Theorem 3.1, item 1 implies that in a short time, the linear MLP function G (T1) (•) positively aligns with the two target signals µ 1 and µ 2 , but negatively aligns with the common token µ 3 . 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 to learn any knowledge yet. This results in item 4 of Theorem 3.1, which shows that the training loss is still large. Theorem 3.2 (Phase 2). With probability at least 1 -δ over the randomness of weight initialization, there exists a time range (T 1 , T 2 ) with T 2 = poly(m) such that for all t ∈ (T 1 , T 2 )• µ ⊤ 2 W (t)⊤ K W (t) Q µ 1 and µ ⊤ 1 W (t)⊤ K W (t) Q µ 2 increase, whereas µ ⊤ 3 W (t)⊤ K W (t) Q µ 1 and µ ⊤ 3 W (t)⊤ K W (t) Q µ 2 decrease. • µ ⊤ 1 W (t)⊤ V W (t) V µ 2 increases, whereas µ ⊤ 1 W (t)⊤ V W (t) V µ 3 and µ ⊤ 2 W (t)⊤ V W (t) V µ 3 decrease. • Linear MLP functions satisfy: G (t) (µ 1 ) ≥ Ω(1), G (t) (µ 2 ) ≥ Ω(1), -G (t) (µ 3 ) ≤ Ω(1). • G (t) (µ 1 ) + G (t) (µ 2 ) + G (t) (µ 3 ) ≥ Ω(1), G (t) (µ 1 ) + G (t) (µ 3 ) ≤ -Ω(1) and G (t) (µ 2 ) + G (t) (µ 3 ) ≤ -Ω(1).In Theorem 3.2, item 1 indicates that, during Phase 2, gradient flow drives the self-attention module to weigh more between the two target signals µ 1 and µ 2 , and to weigh less between one of these signals and the common token µ 3 . Item 2 indicates that gradient flow drives the value matrix W V to positively align the two target signals µ 1 and µ 2 , but negatively align one target signal (µ 1 or µ 2 ) with the common token µ 3 . 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 collectively 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 -δ, there exists a time T ⋆ = Θ(poly(m)) such that• The training and generalization losses satisfy L (T ⋆ ) ≤ 1/poly(m) and L (T ⋆ ) ≤ 1/poly(m).• The attention matrices satisfies:EQUATIONwhereC (T ⋆ ) 1,2 , C (T ⋆ ) 2,1 , -C (T ⋆ ) 3,1 , -C (T ⋆ ) 3,2 = Θ σ 2 0 m √ mLσ 2 1 mm1 + σ 2 0 √ m √ mσ 2 1 mm1andC (T ⋆ ) i1,i2 ≤ O σ 2 0 m n √ mLσ 2 1 mm1 + σ 2 0 √ m √ mσ 2encourages strong negative coupling of one target signal µ 1 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.We next verify our theory and the two-phase characterization of the training process via synthetic experiments (see the experiment setup in Appendix A). 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|G (t) (µ 1 )|, |G (t) (µ 2 )|, |G (t) (µ 3 )| becomesufficiently 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 MLP 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 1 -δ over the weight initialization,∀µ ∈ {µ 1 , µ 2 } : ∂aj w (0) j W (0) V µ ∂t = Θ((σ 2 0 + σ 2 1 )m).Further, by the definition of Phase 1 (see Definition E.2 for a formal definition), the gradients of the attention matrices in the dynamical system are much smaller than that of linear MLP given in Lemma 4.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 ∂ ∂t G (t) (µ 1 ) = Θ((σ 2 0 + σ 2 1 )mm 1 ) which implies that it takes only O(1/(σ 2 0 + σ 2 1 )mm 1 ) iterations for G (t) (µ 1 ) to reach a certain constant magnitude. Lemma 4.1 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 G (t) (µ 3 ), 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 a positive value). This is verified empirically by our experiments in Figure 1 (0) j W (0) V µ3 ∂t = O σ 2 1 √ mm 1 + σ 2 0 √ mL + σ 2 1 mF .Notice that the model output F depends on the weight initialization scale and can be made small.We next show in the following lemma that the gradient ∂G (t) (µ3) ∂t of the common token will quickly become negative soon after the training begins, which drives the transformer model to output a 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 1. This is also verified empirically by our experiments in Figure 1 (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.19). There exists a time T 0.5 ≤ T 1 and a constant C such that for all t ∈ [T 0.5 , T 1 ](1 + C) max ∂G (t) (µ1) ∂t , ∂G (t) (µ2) ∂t ≤ -∂G (t) (µ3) ∂t ≤ (1 -C) ∂G (t) (µ1) ∂t + ∂G (t) (µ2) ∂t .Proof Intuition of Lemma 4.3. We first note that the term1 n n i=1 g (0) i y i L l2=1 m1 j2=1 ∥w (0) j2 ∥ 2 2 (X (i) p (0,i) l2 ) ⊤ µmakes the major contribution to the gradient ∂G (t) (µ3). Such a term is small at the initialization due to the cancellation effect from positive and negative y i 's. However, since the linear MLP G will positively align the two target signals at the beginning, for the samples with positive labels, g (t)i will decrease, whereas for samples with only one target signal, g (t) i will increase. Hence, ∂aj w(t) j W (t)V µ3 ∂t will become negative. This trend will continue until the gradient from µ 3 starts to match the gradients from µ 1 , µ 2 , which is what we establish in Lemma 4.3.Using Lemma 4.3, we can show that all training samples are correctly classified at end of Phase 1. 
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 2. Let T 2 denote the end of Phase 2. Recall that at the end of Phase 1, we have G (t) (µ 1 ), G (t) (µ 2 ), -G (t) (µ 3 ) ≥ Ω(1). We will mainly need to show that such a 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:EQUATION+ m 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ.It has been proved that at the end of Phase 1, for µ ∈ {µ 1 , µ 2 , µ 3 }, we have is large. Assume this can hold for long enough (which we can indeed prove later). Then, we only need to focus on the first term in the sum on the right-hand side in Equation (4). On the other hand, from the dynamical system, we can calculateEQUATIONThus, ify i F (t) i > 0 for all i ∈ [n], then m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2is always increasing.Thus, ∂G (t) (µ) ∂t mainly depends on the behavior of1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 p (t,i2)q←l2,k←µ . Further, this is also a key quantity we need to analyze∂ν ⊤ W (t)⊤ V W (t) V µ ∂t and ∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t . Note that 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 p (t,i2) q←l2,k←µ ≈ 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 if p (t,i2)q←l2,k←µ ≈ 1/L which holds at the beginning of Phase 2. Later, we are going to prove convergence of the training loss via the following: (1) the training loss can 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 of1 n i2: µ∈X (i 2 ) g (t)i2 y i2 . 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 i∈I2 g (t)i and i∈I3 g (t) i close to each other during training; and (ii) Lemma 4.5, which shows that the gradient gaps i∈I1 g(t) i - i∈I2 g (t) i and i∈I2∪I3∪I4 g (t) i -i∈I1 g (t)i are not too small compared with i∈[n] g (t)i . Both Lemmas 4.4 and 4.5 establish that the ratio of those important gradients are kept within certain ranges during training. We call such a key property as automatic balancing of gradients, which is further used for proving that the gradient flow can drive the training loss small. Lemma 4.4 (Same as Lemma F.5). For t ∈ [T 1 , T 2 ], there exists a small constant C ≪ 1 such thati∈I2 g (t) i -i∈I3 g (t) i min( i∈I2 g (t) i , i∈I3 g (t) i ) ≤ C.Proof Intuition of Lemma 4.4. The intuition behind the result is as follows. If i∈I2 g (t)i becomes much bigger than i∈I3 g (t)i during the training, then i∈I1 g(t) i -i∈I2 g (t) i is much smaller than i∈I1 g (t) i -i∈I3 g (t) i which makes ∂G (t) (µ1) ∂t < ∂G (t) (µ2) ∂t. It is not hard to show that random tokens make negligible contributions to the gradient. Thus, for i ∈ I 2 , we have∂F (t) i ∂t ≈ ∂G (t) (µ1) ∂t + ∂G (t) (µ3) ∂t. By the chain rule, we have∂g (t) i ∂t = g ′ (y i F (t) i )y i ∂F (t) i ∂t . Since ∂G (t) (µ3) ∂t < 0, if ∂G (t) (µ1) ∂t < ∂G (t) (µ2) ∂t , then i∈I2 g (t)i will drop faster than i∈I3 g(t) i , i.e., -∂ ∂t i∈I2 g (t) i > -∂ ∂t i∈I3 g (t)i . In Appendix, we formally prove Lemma 4.4 by analyzing the ratio i∈I2 g(t) i / i∈I3 g (t)i , and show that this ratio hangs over around 1 during training.Lemma 4.5 (Abbreviated from Lemma F.6). For t ∈ [T 1 , T 2 ], the gradient satisfies thati∈[n] g (t) i i∈I2∪I3∪I4 g (t) i -i∈I1 g (t) i = O(1), i∈[n] g (t) i i∈I1 g (t) i -i∈I2 g (t) i = O(1).Further, for some constant C, we have(1 + C) max ∂G (t) (µ 1 ) ∂t , ∂G (t) (µ 2 ) ∂t ≤ - ∂G (t) (µ 3 ) ∂t ≤ (1 -C) ∂G (t) (µ 1 ) ∂t + ∂G (t) (µ 2 ) ∂t .Proof Sketch of Lemma 4.5. The proof of Lemma 4.5 relies on analyzing how the ratio between∂G (t) (µ1) ∂tand -∂G (t) (µ3) ∂t changes. We show that this ratio will hang over around some range. Recall the relationship thatEQUATIONj2 . It is not hard to show that-∂G (t) (µ3) ∂t ∂G (t) (µ1) ∂t ≈ -i∈I1 g (t) i + i∈I2∪I3∪I4 g (t) i i∈I1 g (t) i -i∈I2 g (t) i . Define R(t) := -i∈I 1 g (t) i + i∈I 2 ∪I 3 ∪I 4 g (t) i i∈I 1 g (t) i -i∈I 2 g (t) i. Solving when ∂ ∂t R(t) ≥ 0 yields a quadratic inequality, and further analysis shows that the root is contractive and is within some specific range.Utilizing the gradient automatic balancing properties, the following corollary characterizes how the attention matrices in the dynamical system change in Phase 2. In particular, we can show that after W V -transform, µ 1 and µ 2 become more positively correlated whereas µ 1 and µ 3 (also µ 2 and µ 3 ) become negatively correlated. This is a direct result following from updates of the dynamical system. Corollary 4.6 (Abbreviated from Corollary F.13).For t ∈ [T 1 , T 2 ], ∂ ∂t µ ⊤ 2 W (t)⊤ V W (t) V µ 1 > 0, ∂ ∂t µ ⊤ 1 W (t)⊤ V W (t) V µ 3 < 0.Since we have analyzed how G (t) (•) 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∂µ ⊤ 1 W (t)⊤ K W (t) Q µ 2 ∂t = 1 n √ m n i=1 g (t) i y i L l=1 µ ⊤ 1 W (t)⊤ K K (t,i) • diag G (t) (X (i) ) -(G (t) (X (i) )) ⊤ p (t,i) l p (t,i) l x (i)⊤ l µ 2 + 1 n √ m n i=1 g (t) i y i L l=1 µ ⊤ 2 W (t)⊤ Q q (t,i) l p (t,i)⊤ l • diag G (t) (X (i) ) -(G (t) (X (i) )) ⊤ p (t,i) l X (i)⊤ µ 1 .The following lemma shows that the attention score between the two target signals µ 1 and µ 2 increases, whereas that between one target signal µ 1 or µ 2 and the common token µ 3 decreases. Lemma 4.7 (Abbreviated from Lemma F.16). For µ, ν ∈ {µ 1 , µ 2 }, µ ̸ = ν, and for t∈ [T 1 , T 2 ], ∂ ∂t ν ⊤ W (t)⊤ K W (t) Q µ = 1 √ m Θ( L (t) σ 2 0 m) 1 L , ∂ ∂t µ ⊤ 3 W (t)⊤ K W (t) Q µ = - 1 √ m Θ( L (t) σ 2 0 m) 1 L .Lemma 4.7 is keeping track of the attention coefficientsC (t)i1,i2 in Theorem 3.3 via gradient flow, which proves the second item of Theorem 3.3. 
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 a two-phase characterization of the training process. In Phase 1, the linear MLP layer is trained to classify samples correctly, with attention weights almost unchanged. In Phase 2, both attention matrices and the linear MLP jointly evolve to 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 multi-headed attention, multi-layer transformer, etc. Further, it is of interest to 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. We conduct synthetic experiment to verify our theoretical results. We create a dataset following our data distribution in Definition 2.1 with 60 training samples: 30 samples have both µ 1 and µ 2 in it, 10 samples have only µ 1 , 10 samples have only µ 2 , and 10 samples have neither µ 1 nor µ 2 . Each data consists of 5 patches and each patch has dimension 64. The embedding dimension m 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.We provide the gradient flow update for each weight matrix as follows.∂w (t) j ∂t = 1 n n i=1 g (t) i y i L l=1 a j V (t,i) p (t,i) l = 1 n n i=1 g (t) i y i L l=1 a j L h=1 v (t,i) h p (t,i) l,h ∂W (t) V ∂t = 1 n n i=1 g (t) i y i L l=1 m1 j=1 a j w (t) j X (i) p (t,i) l ⊤ ∂W (t) K ∂t = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j L h=1 w (t)⊤ j v (t,i) h L h ′ =1 ∂p (t,i) l,h ∂s l,h ′ q (t,i) l x (i)⊤ h ′ = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j L h=1 w (t)⊤ j v (t,i) h L h ′ =1 p (t,i) l,h (I(h = h ′ ) -p (t,i) l,h ′ )q (t,i) l x (i)⊤ h ′ = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j -w (t)⊤ j V (t,i) p (t,i) l q (t,i) l (X (i) p (t,i) l ) ⊤ + q (t,i) l w (t)⊤ j V (t,i) diag(p (t,i) l )X (i)⊤ = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j q (t,i) l -w (t)⊤ j V (t,i) p (t,i) l p (t,i)⊤ l + w (t)⊤ j V (t,i) diag(p (t,i) l ) X (i)⊤ = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j q (t,i) l p (t,i)⊤ l diag w (t)⊤ j V (t,i) -w (t)⊤ j V (t,i) p (t,i) l X (i)⊤ ∂W (t) Q ∂t = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j L h=1 w (t)⊤ j v (t,i) h L h ′ =1 p (t,i) l,h (I(h = h ′ ) -p (t,i) l,h ′ )k (t,i) h ′ x (i)⊤ l = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j -w (t)⊤ j V (t,i) p (t,i) l (K (t,i) p (t,i) l ) + K (t,i) diag(p (t,i) l )V (t,i)⊤ w (t) j x (i)⊤ l = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j K (t,i) -w (t)⊤ j V (t,i) p (t,i) l p (t,i) l + diag(p (t,i) l )V (t,i)⊤ w (t) j x (i)⊤ l = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l p (t,i) l x (i)⊤ lWe first provide our complete dynamical system. The derivation of each equation is provided in Appendix C.1.∂w (t)⊤ j1 W (t) V µ ∂t = 1 n n i2=1 g (t) i2 y i2 L l2=1 m1 j2=1 a j2 w (t) j1 , w (t) j2 X (i2) p (t,i2) l2 ⊤ µ + 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ (6) ∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ K K (t,i) • diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l p (t,i) l x (i)⊤ l µ + 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i)⊤ l • diag w (t)⊤ j V (t,i) -w (t)⊤ j V (t,i) p (t,i) l X (i)⊤ ν (7) ∂ w (t) j1 , w (t) j2 ∂t = 1 n n i=1 g (t) i y i L l=1 a j2 w (t)⊤ j1 V (t,i) p (t,i) l + 1 n n i=1 g (t) i y i L l=1 a j1 w (t)⊤ j2 V (t,i) p (t,i) l ∂ν ⊤ W (t)⊤ V W (t) V µ ∂t = 1 n i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←µ + 1 n i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←ν ∂ν ⊤ W (t)⊤ Q W (t) Q µ ∂t = 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j ν ⊤ W (t)⊤ Q K (t,i) • diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i m1 j=1 a j µ ⊤ W (t)⊤ Q K (t,i) • diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,ν) p (t,i) l(i,ν) ∂ν ⊤ W (t)⊤ K W (t) K µ ∂t = 1 n √ m i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ K q (t,i) l p (t,i) q←l,k←µ • w (t)⊤ j v (t,i) (µ) -w (t)⊤ j V (t,i) p (t,i) l + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ K q (t,i) l p (t,i) q←l,k←ν • w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l C.1 Derivation of the Dynamical System Lemma C.1. Let µ ∈ {µ i } d i=1 . For all j ∈ [m], we have ∂w (t)⊤ j1 W (t) V µ ∂t = 1 n n i2=1 g (t) i2 y i2 L l2=1 m1 j2=1 a j2 w (t) j1 , w (t) j2 X (i2) p (t,i2) l2 ⊤ µ + 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ = 1 n i:µ∈X (i) g (t) i y i L l=1 a j ∥w (t) j ∥ 2 2 + ∥v (t) (µ)∥ 2 2 p (t,i) q←l,k←µ + ε, where ε = 1 n i:µ∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 p (t,i) q←l,k←µ + 1 n n i=1 g (t) i y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ) I(v (t,i) l2 ̸ = v (t) (µ)).Proof. Let l(i, µ) denote the index such that X (i) l(i,µ) = µ. By the gradient flow update, we have∂w (t)⊤ j1 W (t) V µ ∂t = 1 n n i2=1 g (t) i2 y i2 L l2=1 m1 j2=1 a j2 w (t) j1 , w (t) j2 X (i2) p (t,i2) l2 ⊤ µ + 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ = 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 m1 j2=1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ + 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ) = 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1   a j1 ∥w (t) j1 ∥ 2 2 + j2̸ =j1 a j2 w (t) j1 , w (t) j2   p (t,i2) q←l2,k←µ + 1 n i1:µ / ∈X (i 1 ) g (t) i1 y i1 L l1=1 a j1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ) + 1 n i1:µ∈X (i 1 ) g (t) i1 y i1 L l1=1 a j1   ∥v (t) (µ)∥ 2 2 p (t,i1) q←l1,k←µ + l2̸ =l(i1,µ) p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ)   = 1 n i:µ∈X (i) g (t) i y i L l=1 a j1 ∥w (t) j1 ∥ 2 2 + ∥v (t) (µ)∥ 2 2 p (t,i) q←l,k←µ + 1 n i:µ∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 p (t,i) q←l,k←µ + 1 n n i=1 g (t) i y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ) I(v (t,i) l2 ̸ = v (t) (µ)).Lemma C.2. The following equation on the gradient flow holds:∂ w (t) j1 , w (t) j2 ∂t = 1 n n i=1 g (t) i y i L l=1    l ′ :X (i) l ′ ∈U a j2 w (t)⊤ j1 V (t,i) l ′ p (t,i) l,l ′ + a j1 w (t)⊤ j2 V (t,i) l ′ p (t,i) l,l ′    + 1 n n i=1 g (t) i y i L l=1    l ′ :X (i) l ′ / ∈U a j2 w (t)⊤ j1 V (t,i) l ′ p (t,i) l,l ′ + a j1 w (t)⊤ j2 V (t,i) l ′ p (t,i) l,l ′    .Proof. By gradient flow update, we have∂ w (t) j1 , w (t) j2 ∂t = 1 n n i=1 g (t) i y i L l=1 a j2 w (t)⊤ j1 V (t,i) p (t,i) l + 1 n n i=1 g (t) i y i L l=1 a j1 w (t)⊤ j2 V (t,i) p (t,i) l = 1 n n i=1 g (t) i y i L l=1    l ′ :X (i) l ′ ∈U a j2 w (t)⊤ j1 V (t,i) l ′ p (t,i) l,l ′ + a j1 w (t)⊤ j2 V (t,i) l ′ p (t,i) l,l ′    + 1 n n i=1 g (t) i y i L l=1    l ′ :X (i) l ′ / ∈U a j2 w (t)⊤ j1 V (t,i) l ′ p (t,i) l,l ′ + a j1 w (t)⊤ j2 V (t,i) l ′ p (t,i) l,l ′    . Lemma C.3. Let µ, ν ∈ {µ i } d i=1 .Then the following equation on gradient flow holds:∂νW (t)⊤ V W (t) V µ ∂t = 1 n i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←µ + 1 n i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←ν .Proof. The gradient flow update can be derived as follows:∂νW (t)⊤ V W (t) V µ ∂t = 1 n n i=1 g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j X (i) p (t,i) l ⊤ µ + 1 n n i=1 g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j X (i) p (t,i) l ⊤ ν = 1 n i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←µ + 1 n i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←ν . Lemma C.4. Let µ, ν ∈ {µ i } d i=1 .Then the following equations on gradient flow hold.∂ν ⊤ W (t)⊤ Q W (t) Q µ ∂t = 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j ν ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i m1 j=1 a j µ ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,ν) p (t,i) l(i,ν) , ∂ν ⊤ W (t)⊤ K W (t) K µ ∂t = 1 n √ m i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ K q (t,i) l p (t,i) q←l,k←µ w (t)⊤ j v (t,i) (µ) -w (t)⊤ j V (t,i) p (t,i) l + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ K q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l , ∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t = 1 n √ m i:µ,ν∈X (i) g (t) i y i m1 j=1 a j ∥k (t) (ν)∥ 2 2 v (t)⊤ (ν)w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) q←µ,k←ν + 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j L l=1 ν ⊤ W (t)⊤ K K (t,i) l V (t,i)⊤ l w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) q←µ,k←l I(K (t,i) l ̸ = k (t) (ν)) + 1 n √ m i:ν,µ∈X (i) g (t) i y i m1 j=1 a j ∥q (t) (µ)∥ 2 2 p (t,i) q←µ,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l I(q (t,i) l ̸ = q (t) (µ)).Proof. To prove the first result, we have∂ν ⊤ W (t)⊤ Q W (t) Q µ ∂t = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l p (t,i) l x (i)⊤ l µ + 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l p (t,i) l x (i)⊤ l ν = 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j ν ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i m1 j=1 a j µ ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,ν) p (t,i) l(i,ν) .To prove the second result, we have∂ν ⊤ W (t)⊤ K W (t) K µ ∂t = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ K q (t,i) l p (t,i)⊤ l diag w (t)⊤ j V (t,i) -w (t)⊤ j V (t,i) p (t,i) l X (i)⊤ µ + 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ K q (t,i) l p (t,i)⊤ l diag w (t)⊤ j V (t,i) -w (t)⊤ j V (t,i) p (t,i) l X (i)⊤ ν = 1 n √ m i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ K q (t,i) l p (t,i) q←l,k←µ w (t)⊤ j v (t,i) (µ) -w (t)⊤ j V (t,i) p (t,i) l + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ K q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l .To prove the third result, we have∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t = 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ K K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l p (t,i) l x (i)⊤ l µ + 1 n √ m n i=1 g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i)⊤ l diag w (t)⊤ j V (t,i) -w (t)⊤ j V (t,i) p (t,i) l X (i)⊤ ν = 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j ν ⊤ W (t)⊤ K K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l = 1 n √ m i:µ,ν∈X (i) g (t) i y i m1 j=1 a j ∥k (t) (ν)∥ 2 2 v (t)⊤ (ν)w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) q←µ,k←ν + 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j L l=1 ν ⊤ W (t)⊤ K K (t,i) l V (t,i)⊤ l w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) q←µ,k←l I(K (t,i) l ̸ = k (t) (ν)) + 1 n √ m i:ν,µ∈X (i) g (t) i y i m1 j=1 a j ∥q (t) (µ)∥ 2 2 p (t,i) q←µ,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l I(q (t,i) l ̸ = q (t) (µ)).Lemma D.1. With probability at least 1 -δ over the randomness of the initialization of W K and W Q , for any l 1 , l 2 ∈ [d], we haveEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONQ µ l2 ∼ N (0, σ 2 0 I). The rest of proof applies Lemma H.2.Corollary D.2. For all i ∈ [n], l, k ∈ [L], we have p (0,i) l,k = 1 L ± O 1 Lm .Proof. Following from Lemma G.2 and from the first-order Taylor approximation on the softmax function from 0, we havep (0,i) l,k = 1 L ± O σ 2 0 m L √ m 4 m log 2d δ + 4 m log 2d δ .The corollary then follows from Condition 1.Lemma D.3. With probability at least 1 -δ over the randomness of the initialization of W and W V , then forl 1 ̸ = l 2 ∈ [d],we haveW (0) V µ l1 , W (0) V µ l2 ≤ σ 2 0 m 4 m log 2d δ + 4 m log 2d δ , for j 1 ̸ = j 2 ∈ [m 1 ], we have w (0) j1 , w (0) j2 ≤ σ 2 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ ,and for all j∈ [m 1 ], l ∈ [d],we haveEQUATIONV µ l ≤ σ 0 σ 1 m 4 m log 2m 1 d δ + 4 m log 2m 1 d δ ∥w j ∥ 2 2 = σ 2 1 m 1 ± 4 m log 2m 1 δ + 4 m log 2m 1 δ ∥W (0) V µ l ∥ 2 2 = σ 2 0 m 1 ± 4 m log 2d δ + 4 m log 2d δ .Proof. The proof is similar to that for Lemma D.1 and is omitted. V X (i) p(0,i) l ∥ 2 2 = σ 2 0 m L 1 ± 4 m log 2nL δ ± 4 m log 2nL δProof. First of all, by Corollary D.2 and Assumption 2.3,E ∥W (0) V X (i) p (0,i) l ∥ 2 2 = E   m j=1 W (0) V j , X (i) p (0,i) l 2   = σ 2 0 m∥X (i) p (0,i) l ∥ 2 2 = σ 2 0 m L .Finally, applying Bernstein's inequality and taking a union bound over [n] and [L] we finish the proof.Corollary D.5. Conditioned on the success of Lemma D.4, with probability at least 1 -δ over the randomness of W , for all j∈ [m 1 ], i ∈ [n], l ∈ [L], we have w (0)⊤ j V (0,i) p (0,i) l ≤ σ 1 σ 0 m L log m 1 nL δ .Proof. Conditioned on V (0,i) p (0,i) l, we have wEQUATION2 ). Thus, by Gaussian tail bound and a union bound over j∈ [m 1 ], i ∈ [n], l ∈ [L], with probability at least 1 -δ, we havew (0)⊤ j V (0,i) p (0,i) l ≤ σ 1 σ 0 m L log m 1 nL δ .Lemma D.6. With probability at least 1 -δ over the randomness in the initialization of a, we have|S +1 | = m 1   1 2 ± 2 log(4/δ) m 1   , |S -1 | = m 1   1 2 ± 2 log(4/δ) m 1   .Proof. The proof follows by applying Hoeffding's inequality.Lemma D.7 (Initial sub-network output). Assume that the success of the events in Lemma D.3 holds. For all i ∈ [d], with probability at least 1 -δ over the randomness in the weight initialization, we have|G (0) (µ i )| ≤ O(σ 1 σ 0 √ mm 1 ). Proof. Consider a fixed i ∈ [d]. By Lemma D.3, we have ∥W (0) V µ i ∥ 2 2 = Θ(σ 2 0 m). Thus, conditioned on W (0) V µ i , we have m1 j=1 w (0) j W (0) V µ i ∼ N (0, σ 2 1 σ 2 0 mm 1 ). Thus, by Gaussian concentration bound, we have|G (0) (µ i )| ≤ O(σ 1 σ 0 √ mm 1 ).Lemma D.8 (Initial network output). Assume that the success of the events in Lemma D.3 and Lemma D.4 holds. For all i ∈ [n], with probability at least 1 -δ over the randomness in the weight initialization, we have|F (0) (X (i) )| ≤ 2σ 0 σ 1 2Lm 1 m log(2Ln/δ) ≤ 0.01. Proof. For fixed l ∈ L, j ∈ [m], i ∈ [n], by Corollary D.5, we havew (0)⊤ j V (0,i) p (0,i) l ≤ σ 1 σ 0 m L log m 1 nL δ .Thus, this implies that a j w(0)⊤ j V (0,i) p (0,i) lis a sub-Gaussian random variable with variance proxy σ 2 0 σ 2 1 m L log m1nL δ . Therefore, the following inequality holds.P   m1 j=1 a j w (0)⊤ j V (0,i) p (0,i) l ≥ 2σ 0 σ 1 2m 1 m(log m1nL δ ) log(2/δ) L   ≤ δ.i ∈ [n], l ∈ [L], with probability at least 1 -δ, for all i ∈ [n], we have|F (0) (X (i) )| ≤ 2σ 0 σ 1 2Lm 1 m(log m 1 nL δ ) log(2Ln/δ).Finally, by Condition 1, we can make |F (0) (X (i) )| ≤ 0.01.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 to the two signals µ 1 , µ 2 , the common token µ 3 and the random tokens. Definition E.1 (Radius of keys and queries). Define the radius of keys and queries R K , R Q respectively to beR K := max i,j∈[d] µ ⊤ i W (t)⊤ K W (t) K µ j -µ ⊤ i W (0)⊤ K W (0) K µ j , R Q := max i,j∈[d] µ ⊤ i W (t)⊤ Q W (t) Q µ j -µ ⊤ i W (0)⊤ Q W (0) Q µ j .Definition E.2 (Phase 1). Define the range of Phase 1 to be [0, T 1 ], where T 1 = min{t ′ , C T1 /(σ 2 1 mm 1 )} for some sufficiently large constant C T1 and t ′ is defined to be the maximum time such that for all t ≤ t ′ , all of the following hold:1. max j∈[m],µ∈{µi} 3 i=1 w (t) j W (t) V µ -w (0) j W (0) V µ ≤ R where R < O(1/m 1 ); 2. max j∈[m],µ / ∈{µi} 3 i=1 w (t) j W (t) V µ -w (0) j W (0) V µ ≤ O(R/n + R/ √ m); 3. max µ,ν∈{µ} d i=1 µ ⊤ W (t)⊤ Q W (t) K ν -µ ⊤ W (0)⊤ Q W (0) K ν ≤ R S where R S ≤ O(1/(m √ m)); 4. R K , R Q ≤ O(σ 2 0 √ m).Based on this definition, we can further obtain the maximum softmax probability change as follows.Proposition E.3. DefineR P := max i∈[n], µ,ν∈X (i) p (t,i) q←µ,k←ν -p (0,i) q←µ,k←ν .ThenR P = O 1 √ mL + L m .Proof. By Lemma G.2, we haveR P ≤ O(R S /L + R 2 S L) = O 1 √ mL + L m .Initially, the loss for the samples with one signal will increase.Lemma E.4 (Signal updates, same as Lemma 4.1). At t = 0, for µ ∈ {µ 1 , µ 2 }, we have∂a j w (0) j W (0) V µ ∂t = Θ((σ 2 0 + σ 2 1 )m).Proof. Take µ = µ 1 . First of all, by the gradient flow update in Lemma C.1, we have∂w (t)⊤ j1 W (t) V µ ∂t = 1 n n i2=1 g (t) i2 y i2 L l2=1 m1 j2=1 a j2 w (t) j1 , w (t) j2 X (i2) p (t,i2) l2 ⊤ µ + 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ = 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 m1 j2=1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ + 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ) = 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1   a j1 ∥w (t) j1 ∥ 2 2 + j2̸ =j1 a j2 w (t) j1 , w (t) j2   p (t,i2) q←l2,k←µ + 1 n i1:µ / ∈X (i 1 ) g (t) i1 y i1 L l1=1 a j1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ) + 1 n i1:µ∈X (i 1 ) g (t) i1 y i1 L l1=1 a j1   ∥v (t) (µ)∥ 2 2 p (t,i1) q←l1,k←µ + l2̸ =l(i1,µ) p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ)   = 1 n i:µ∈X (i) g (t) i y i L l=1 a j1 ∥w (t) j1 ∥ 2 2 + ∥v (t) (µ)∥ 2 2 p (t,i) q←l,k←µ + 1 n i:µ∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 p (t,i) q←l,k←µ ε1 + 1 n n i=1 g (t) i y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ) I(v (t,i) l2 ̸ = v (t) (µ)) ε2 .Now, by Lemma D.3, Corollary D.2 and Lemma D.8, we havea j1 1 n i:µ1∈X (i) g (0) i y i L l=1 a j1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ)∥ 2 2 p (0,i) q←l,k←µ1 = 1 n i∈I1 g (0) i L l=1 ∥w (0) j ∥ 2 2 + ∥v (0) (µ 1 )∥ 2 2 p (0,i) q←l,k←µ1 - 1 n i∈I2 g (0) i L l=1 ∥w (0) j ∥ 2 2 + ∥v (0) (µ 1 )∥ 2 2 p (0,i) q←l,k←µ1 = 1 3 ± 0.01 L • σ 2 1 m + σ 2 0 m 1 ± 4 m log 2d δ + 4 m log 2d δ 1 L (1 + o(1)).On the other hand, by Proposition E.5, we haveEQUATION(0,i) q←l,k←µ1 ≤ σ 2 1 √ m 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ log m 1 δ ; |ε 2 | = 1 n n i=1 g (0) i y i L l1=1 a j1 L l2=1 p (0,i) l1,l2 v (0,i) l2 , v (0) (µ 1 ) I(v (0,i) l2 ̸ = v (0) (µ 1 )) ≤ σ 2 0 m √ L 4 m log 2nL δ + 4 m log 2nL δ . If m ≥ Cm 1 log 2 m1 δ in Condition 1 for some sufficiently large C, then |ε 1 | ≤ 0.01σ 2 1 m; and if m ≥ C ′ L log 2nLδ for some sufficiently large C ′ , then |ε 2 | ≤ 0.01σ 2 0 m. Proposition E.5. Assume the events in Lemma D.3 and Corollary D.2 succeed. With probability at least 1 -δ over the randomness in the weight initialization, for all j 1 ∈ [m 1 ], we havej2:j2̸ =j1 a j2 w (0) j1 , w (0) j2 ≤ σ 2 1 √ m 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ log 4m 1 δ .Further, for all µ ∈ {µ i } d i=1 , we haveEQUATION1 n n i=1 g (0) i y i L l1=1 a j1 L l2=1 p (0,i) l1,l2 v (0,i) l2 , v (0) (µ) I(v (0,i) l2 ̸ = v (0) (µ)) ≤ σ 2 0 m √ L 4 m log 2nLd δ + 4 m log 2nLd δ .Proof. First, fix i, l, and consider the randomness of a. By Lemma D.3, a j2 w(0) j1 , w (0) j2 is a sub- Gaussian random variable with variance proxy σ 4 1 m 2 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ 2. This implies that with probability at least 1 -δ/2, for all j 1 ∈ [m 1 ], we havej2:j2̸ =j1 a j2 w (0) j1 , w (0) j2 ≤ σ 2 1 √ m 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ log 4m 1 δ .Thus, by Corollary D.2, for all µ ∈ {µ i } d i=1 , we have1 n i:µ1∈X (i) g (0) i y i L l=1 j2:j2̸ =j1 a j2 w (0) j1 , w (0) j2 p (0,i) q←l,k←µ ≤ |i : µ ∈ X (i) | n σ 2 1 √ m 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ log m 1 δ .We next derive the second inequality. Consider the randomness inW (0) V . Note that L l2=1 p (0,i) l1,l2 v (0,i) l2 I(v (0,i) l2 ̸ = v (0) (µ)) ∼ N 0, σ 2 0 L l2=1 (p (0,i) l1,l2 ) 2 I(v (0,i) l2 ̸ = v (0) (µ))I .Thus, by Lemma H.2 and Corollary D.2 and taking a union bound over i ∈[n], l 1 ∈ [L], µ ∈ {µ i } d i=1 , we have with probability at least 1 -δ/2, L l2=1 p (0,i) l1,l2 v (0,i) l2 I(v (0,i) l2 ̸ = v (0) (µ)), v (0) (µ) ≤ σ 2 0 m 1 √ L 4 m log 2nLd δ + 4 m log 2nLd δ .Therefore,1 n n i=1 g (0) i y i L l1=1 a j1 L l2=1 p (0,i) l1,l2 v (0,i) l2 , v (0) (µ) I(v (0,i) l2 ̸ = v (0) (µ)) ≤ σ 2 0 m √ L 4 m log 2nLd δ + 4 m log 2nLd δ . Lemma E.6 (Random token updates). For µ ∈ {µ i } d i=4 , we have ∂a j w (0) j W (0) V µ ∂t = O 1 n (σ 2 0 + σ 2 1 )m + σ 2 0 √ mL .Proof. Following the proof of Lemma E.4, we have∂a j w (0) j W (0) V µ ∂t = 1 n n i2=1 g (0) i2 y i2 L l2=1 m1 j2=1 a j2 w (0) j1 , w (0) j2 X (i2) p (0,i2) l2 ⊤ µ + 1 n n i1=1 g (0) i1 y i1 L l1=1 a j1 p (0,i1)⊤ l1 V (0,i1)⊤ W (0) V µ = 1 n i:µ∈X (i) g (0) i y i L l=1 a j1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ)∥ 2 2 p (0,i) q←l,k←µ + 1 n i:µ∈X (i) g (0) i y i L l=1 j2:j2̸ =j1 a j2 w (0) j1 , w (0) j2 p (0,i) q←l,k←µ ε1 + 1 n n i=1 g (0) i y i L l1=1 a j1 L l2=1 p (0,i) l1,l2 v (0,i) l2 , v (0) (µ) I(v (0,i) l2 ̸ = v (0) (µ)) ε2By Proposition E.5 and the fact that only one X (i) satisfies µ ∈ X (i) , we have1 n i:µ∈X (i) g (0) i y i L l=1 a j1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ)∥ 2 2 p (0,i) q←l,k←µ = Θ 1 n (σ 2 0 + σ 2 1 )m ,and|ε 1 | ≤ 1 n σ 2 1 √ m 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ log 4m 1 δ , |ε 2 | ≤ σ 2 0 m √ L 4 m log 2nLd δ + 4 m log 2nLdδ .Lemma E.7. For all t ≤ T 1 , for all i ∈ [n], l ∈ [L], we havew (t) j W (t) V X (i) p (t,i) l -w (0) j W (0) V X (i) p (0,i) l ≤ O(LR P σ 0 σ 1 √ m) + 4R/L.Proof. By Definition E.2, we havew (t) j V (t,i) p (t,i) l -w (0) j V (0,i) p (0,i) l ≤ L l ′ =1 w (t) j V (t,i) l ′ p (t,i) l,l ′ -w (0) j V (0,i) l ′ p (0,i) l,l ′ ≤ l ′ :V l ′ ∈v(µ1,µ2,µ3) (R P O(σ 0 σ 1 √ m) + R/L) + l ′ :V l ′ / ∈v(µ1,µ2,µ3) (R P O(σ 0 σ 1 √ m) + R/(nL)) ≤ O(R P σ 0 σ 1 √ m + R/L) + O(LR P σ 0 σ 1 √ m + R/n) = O(LR P σ 0 σ 1 √ m) + 4R/L.By our choice of parameters in Condition 1 and Definition E.2, we have O(LRP σ 0 σ 1 √ m)+4R/L < σ 0 σ 1 m/L.Proposition E.8. With probability at least 1 -δ, for all µ ∈ {µ i } d i=1 , we havem1 j=1 a j µ ⊤ W (0)⊤ V w (0) j ≤ O(σ 0 σ 1 √ mm 1 ).Proof. The proof is similar to Proposition E.22, and is omitted.Lemma E.9 (Value correlation change). For all µ, ν∈ {µ i } d i=1 , we have ∂ v (t) (µ), v (t) (ν) ∂t ≤ i : µ ∈ X (i) + i : ν ∈ X (i) n O(1),and thus,v (t) (µ), v (t) (ν) -v (0) (µ), v (0) (ν) ≤ t i : µ ∈ X (i) + i : ν ∈ X (i) n O(1)for all t ≤ T 1 . Thus, for µ ̸ = ν, we havev (t) (µ), v (t) (ν) ≤ O(σ 2 0 √ m) and ∥v (t) (µ)∥ 2 2 = Θ(σ 2 0 m) for t ≤ T 1 .Proof. By Lemma C.3, we have∂νW (t)⊤ V W (t) V µ ∂t = 1 n i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←µ + 1 n i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←ν .Further,m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j - m1 j=1 a j ν ⊤ W (0)⊤ V w (0) j ≤ m 1 R.Thus, by Proposition E.8, we havem1 j=1 a j w (t)⊤ j W (t) V ν ≤ m1 j=1 a j σ(0) i,l,j w (0)⊤ j W (0) V ν + m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j - m1 j=1 a j ν ⊤ W (0)⊤ V w (0) j ≤ O(σ 0 σ 1 √ mm 1 ) + m 1 R, which implies 1 n i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←µ + 1 n i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←ν ≤ 1 n i : µ ∈ X (i) + i : ν ∈ X (i) O(σ 0 σ 1 √ mm 1 ) + m 1 R ≤ i : µ ∈ X (i) + i : ν ∈ X (i) n O(1)where the last inequality applies Lemma E.7.Corollary E.10. For all t ≤ T 1 , we have1 n n i=1 g (t) i y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ 1 ) I(v (t,i) l2 ̸ = v (t) (µ 1 )) ≤ L • O σ 2 0 √ m .Proof. We derive the following bound:1 n n i=1 g (t) i y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ) I(v (t,i) l2 ̸ = v (t) (µ)) ≤ L l1=1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ) I(v (t,i) l2 ̸ = v (t) (µ)) ≤ L • O σ 2 0 √m , where the last inequality follows from Lemma E.9 and Lemma D.3.Lemma E.11. For all t ≤ T 1 , we have1 n i:µ1∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 p (t,i) q←l,k←µ1 ≤ |{i : µ 1 ∈ X (i) }| n O(σ 2 1 √ mm 1 ).Proof. We first derive:EQUATIONEQUATION.For term (1), by Proposition E.5, we have1 n i:µ1∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (0) j1 , w (0) j2 ≤ |{i : µ 1 ∈ X (i) }| n O σ 2 1 √ mm 1 L .For term (2), note that1 n i:µ1∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 -w (0) j1 , w (0) j2 ≤ |{i : µ 1 ∈ X (i) }| n Lm 1 • O 1 m ,where the inequality is by Lemma D.3, Proposition E.12.Finally, since p (t,i) q←l,k←µ1 ≤ 1 L + 1 L 2 + R P , combining the upper bound for both terms (1) and (2), we haveα n i:µ1∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 σ(t) i,l,j2 w (t) j1 , w (t) j2 p (t,i) q←l,k←µ1 ≤ α |{i : µ 1 ∈ X (i) }| n • O σ 2 1 √ mm 1 L .Proposition E.12 (Neuron correlation change, Phase 1). For t ≤ T 1 , for all j 1 , j 2 ∈ [m 1 ], we haveEQUATIONj1 , w(t) j2 -w (0) j1 , w (0) j2 ≤ t2L σ 1 σ 0 m L log m 1 nL δ + O(LR P σ 0 σ 1 √ m) + 4R/L .Proof. By the gradient flow update in Lemma C.2, we have∂ w (t) j1 , w (t) j2 ∂t = 1 n n i ′ =1 g (t) i ′ y i ′ L l ′ =1 a j2 w (t)⊤ j1 V (t,i ′ ) p (t,i ′ ) l ′ + 1 n n i ′ =1 g (t) i ′ y i ′ L l ′ =1 a j1 w (t)⊤ j2 V (t,i ′ ) p (t,i ′ ) l ′ = 1 n n i ′ =1 g (t) i ′ y i ′ L l ′ =1 a j2 w (0)⊤ j1 V (0,i ′ ) p (0,i ′ ) l ′ + (w (t)⊤ j1 V (t,i ′ ) p (t,i ′ ) l ′ -w (0)⊤ j1 V (0,i ′ ) p (0,i ′ ) l ′ ) + 1 n n i ′ =1 g (t) i ′ y i ′ L l ′ =1 a j1 w (0)⊤ j2 V (0,i ′ ) p (0,i ′ ) l ′ + (w (t)⊤ j2 V (t,i ′ ) p (t,i ′ ) l ′ -w (0)⊤ j2 V (0,i ′ ) p (0,i ′ ) l ′) .Now, since by Lemma E.7 we havew (t) j W (t) V X (i) p (t,i) l -w (0) j W (0) V X (i) p (0,i) l ≤ O(LR P σ 0 σ 1 √ m) + 4R/L,by Definition E.2 and Corollary D.5, we have1 n n i ′ =1 g (t) i ′ y i ′ L l ′ =1 a j2 w (0)⊤ j1 V (0,i ′ ) p (0,i ′ ) l ′ + (w (t)⊤ j1 V (t,i ′ ) p (t,i ′ ) l ′ -w (0)⊤ j1 V (0,i ′ ) p (0,i ′ ) l ′ ) ≤ L σ 1 σ 0 m L log m 1 nL δ + O(LR P σ 0 σ 1 √ m) + 4R/L 1 n n i ′ =1 g (t) i ′ y i ′ L l ′ =1 a j1 w (0)⊤ j2 V (0,i ′ ) p (0,i ′ ) l ′ + (w (t)⊤ j2 V (t,i ′ ) p (t,i ′ ) l ′ -w (0)⊤ j2 V (0,i ′ ) p (0,i ′ ) l ′ ) ≤ L σ 1 σ 0 m L log m 1 nL δ + O(LR P σ 0 σ 1 √ m) + 4R/L .Thus,w (t) j1 , w (t) j2 -w (0) j1 , w (0) j2 ≤ t τ =0 ∂ w (τ ) j1 , w (τ ) j2 ∂τ ≤ t2L σ 1 σ 0 m L log m 1 nL δ + O(LR P σ 0 σ 1 √ m) + 4R/L .Corollary E.13 (Neuron Norm Change, Phase 1). For all t ≤ T 1 and all j ∈ [m], we have∥w (t) j ∥ 2 2 -∥w (0) j ∥ 2 2 ≤ t2L(σ 1 σ 0 m L log m 1 nL δ + O(LR P σ 0 σ 1 √ m) + 4R/L).Theorem E.14 (Signal correlation growth, phase 1). For t ≤ T 1 , for µ ∈ {µ 1 , µ 2 },∂a j w (t) j W (t) V µ ∂t = 1 n i:µ∈X (i) , yi=1 g (t) i L l=1 p (t,i) q←l,k←µ - i:µ∈X (i) , yi=-1 g (t) i L l=1 p (t,i) q←l,k←µ Θ((σ 2 0 + σ 2 1 )m) + ε where |ε| ≤ O(Lσ 2 0 √ m + σ 2 1 √ mm 1 ).Proof. We take µ = µ 1 and the proof is similar for µ = µ 2 . By Lemma C.1, we have∂a j w (t) j W (t) V µ ∂t = 1 n i:µ1∈X (i) ,yi=1 g (t) i L l=1 ∥w (t) j ∥ 2 2 + ∥v (t) (µ 1 )∥ 2 2 p (t,i) q←l,k←µ1 - 1 n i:µ1∈X (i) ,yi=-1 g (t) i L l=1 ∥w (t) j ∥ 2 2 + ∥v (t) (µ 1 )∥ 2 2 p (t,i) q←l,k←µ1 + ε = ∥w (t) j ∥ 2 2 + ∥v (t) (µ 1 )∥ 2 2 1 n i:µ1∈X (i) , yi=1 g (t) i L l=1 p (t,i) q←l,k←µ1 - i:µ1∈X (i) , yi=-1 g (t) i L l=1 p (t,i) q←l,k←µ1 + ε where ε = 1 n i:µ1∈X (i) g (t) i y i L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 p (t,i) q←l,k←µ1 + 1 n n i=1 g (t) i y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ 1 ) I(v (t,i) l2 ̸ = v (t) (µ 1 )).We can now bound the magnitude of ε by Corollary E.10, Lemma E.11 and Proposition E.5 and obtain:|ε| ≤ L • O σ 2 0 √ m + O(σ 2 1 √ mm 1 ).Finally, by Lemma D.3 and Corollary E.13, we have ∥w(t) j ∥ 2 2 = Θ(σ 2 1 m) and by Lemma E.9, we have ∥v (t) (µ 1 )∥ 2 2 = Θ(σ 2 0 m). The proof is completed.Theorem E.15 (Random token growth, Phase 1). For t ≤ T 1 , for µ ∈ R, we have∂a j w (t) j W (t) V µ ∂t = O 1 n (σ 2 0 + σ 2 1 )m + O(σ 2 0 L √ m).Proof. Fix a µ ∈ R. By our Assumption 2.3, µ appears at most once in the training data set. Now, assume µ is in the training set and let i ⋆ be the index of the sample containing µ. Applying Lemma C.1 on the random token µ, we have∂a j w (t) j W (t) V µ ∂t = 1 n g (t) i ⋆ y i ⋆ L l=1 a j ∥w (t) j ∥ 2 2 + ∥v (t) (µ)∥ 2 2 p (t,i ⋆ ) q←l,k←µ + ε, where ε = 1 n g (t) i ⋆ y i ⋆ L l=1 j2̸ =j1 a j2 w (t) j1 , w (t) j2 p (t,i ⋆ ) q←l,k←µ + 1 n n i=1 g (t) y i L l1=1 a j1 L l2=1 p (t,i) l1,l2 v (t,i) l2 , v (t) (µ) I(v (t,i) l2 ̸ = v (t) (µ)).We can now bound the magnitude of ε by Corollary E.10 and Lemma E.11 as follows:|ε| ≤ L • O σ 2 0 √ m + 1 n O(σ 2 1 √ mm 1 ).Finally, by Lemma D.3 and Corollary E.13, we have∂a j w (t) j W (t) V µ ∂t = O 1 n (σ 2 0 + σ 2 1 )m + O(σ 2 0 L √ m).Lemma E.16 (Initial Per Neuron Gradient of Common Token, same as Lemma 4.2). Let F = max i F (0)i . We have∂a j w (0) j W (0) V µ 3 ∂t = O σ 2 1 √ mm 1 + σ 2 0 √ mL + (σ 2 0 + σ 2 1 )mF .Proof. Following from the proof of Lemma E.4, we have∂a j1 w (0) j1 W (0) V µ 3 ∂t = 1 n i:µ3∈X (i) g (0) i y i L l=1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ 1 )∥ 2 2 p (0,i) q←l,k←µ3 + a j1 1 n i:µ3∈X (i) g (0) i y i L l=1 j2̸ =j1 a j2 w (0) j1 , w (0) j2 p (0,i) q←l,k←µ3 ε1 + a j1 1 n n i=1 g (0) i y i L l1=1 a j1 L l2=1 p (0,i) l1,l2 v (0,i) l2 , v (0) (µ 3 ) I(v (0,i) l2 ̸ = v (0) (µ 3 )) ε2 .For the first term, we have1 n i:µ3∈X (i) g (0) i y i L l=1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ 3 )∥ 2 2 p (0,i) q←l,k←µ3 = 1 n yi=1 g (0) i L l=1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ 3 )∥ 2 2 p (0,i) q←l,k←µ3 - 1 n yi=-1 g (0) i L l=1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ 3 )∥ 2 2 p (0,i) q←l,k←µ3 = 1 n ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ 3 )∥ 2 2 yi=1 g (0) i L l=1 p (0,i) q←l,k←µ3 - yi=-1 g (0) i L l=1 p (0,i) q←l,k←µ3 .By Lemma D.8, and Corollary D.2, we have1 n yi=1 g (0) i L l=1 p (0,i) q←l,k←µ3 - yi=-1 g (0) i L l=1 p (0,i) q←l,k←µ3 ≤ 1 n yi=1 1 2 + F L l=1 1 L + 1 L 2 - yi=-1 1 2 -F L l=1 1 L - 1 L 2 ≤ 1 2 + F 1 L + 1 Lm L 2 - 1 2 -F 1 L - 1 Lm L 2 ≤ 2F + O(1/m).By Lemma D.3, we have ∥w(0) j1 ∥ 2 2 = Θ(σ 2 1 m) and ∥v (0) (µ 3 )∥ 2 2 = Θ(σ 2 0 m), and thus,1 n i:µ3∈X (i) g (0) i y i L l=1 a j1 ∥w (0) j1 ∥ 2 2 + ∥v (0) (µ 3 )∥ 2 2 p (0,i) q←l,k←µ3 ≤ O(α(σ 2 0 + σ 2 1 )mF ).Further, by Proposition E.5, we have|ε 1 | = 1 n i:µ3∈X (i) g (0) i y i L l=1 j2̸ =j1 a j2 w (0) j1 , w (0) j2 p (0,i) q←l,k←µ3 ≤ σ 2 1 √ m 1 m 4 m log 2m 2 1 δ + 4 m log 2m 2 1 δ log m 1 δ . |ε 2 | = 1 n n i=1 g (0) i y i L l1=1 a j1 L l2=1 p (0,i) l1,l2 v (0,i) l2 , v (0) (µ 3 ) I(v (0,i) l2 ̸ = v (0) (µ 3 )) ≤ σ 2 0 m √ L 4 m log 2nL δ + 4 m log 2nL δ .Finally, we have∂a j1 w (0) j1 W (0) V µ 3 ∂t ≤ O σ 2 1 √ mm 1 + σ 2 0 √ mL + (σ 2 0 + σ 2 1 )mF .Lemma E.17 (Per neuron gradient of common token). For t ≤ T 1 , we have∂a j w (0) j W (0) V µ 3 ∂t ≤ O(σ 2 1 m)Proof. The proof is similar to that for Theorem E.14 and we omit it here.Definition E.18 (sub-network). Define the sub-network structure asG(µ) = m1 j=1 a j w j W V µ.We can compute the gradient of the sub-network as∂G(µ) ∂t = m1 j=1 ∂a j w (t) j W (t) V µ ∂t = m1 j=1 a j w (t) j ∂W (t) V µ ∂t + a j ∂w (t) j ∂t W (t) V µ = 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ + 1 n n i1=1 g (t) i1 y i1 m 1 L l1=1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2, v (t) (µ) .There exists a T 0.5 ≤ T 1 and a constant 0 < C < 1 such that for all T 0.5 ≤ t ≤ T 1 ,(1 + C) max ∂G (t) (µ 1 ) ∂t , ∂G (t) (µ 2 ) ∂t ≤ - ∂G (t) (µ 3 ) ∂t ≤ (1 -C) ∂G (t) (µ 1 ) ∂t + ∂G (t) (µ 2 ) ∂t .Further,∂G (t) (µ 1 ) ∂t , ∂G (t) (µ 2 ) ∂t ≥ Ω((σ 2 0 + σ 2 1 )mm 1 ).Proof. First of all, we have∂F (t) i ∂t = L l1=1 m1 j=1 L l2=1 ∂a j w (t) j W (t) V X l2 ∂t p (t,i) l1,l2 + a j w (t) j W (t) V X l2 ∂p (t,i) l1,l2∂t .By Lemma E.4 and Lemma E.16, we have∂a j w (0) j W (0) V µ 1 ∂t - ∂a j w (0) j W (0) V µ 3 ∂t = Θ((σ 2 0 + σ 2 1 )m), ∂a j w (0) j W (0) V µ 2 ∂t - ∂a j w (0) j W (0) V µ 3 ∂t = Θ((σ 2 0 + σ 2 1 )m).By Theorem E.15 and Corollary E.24, we haveL l1=1 m1 j=1 L l2=1 ∂a j w (t) j W (t) V X l2 ∂t I(X l2 / ∈ {µ i } 3 i=1 )p (t,i) l1,l2 + a j w (t) j W (t) V X l2 ∂p (t,i) l1,l2 ∂t = O L n (σ 2 0 + σ 2 1 )mm 1 .Thus, for all i ∈ I 1 ∪ I 2 ∪ I 3 ,∂F (0) i ∂t > 0, which implies ∂g (0) i ∂t < 0 for i ∈ I 1 and ∂g (0) i ∂t > 0 for i ∈ I 2 ∪ I 3 .Next, we show that there must exist a time such that -∂G (t) (µ3) ∂t = max( ∂G (t) (µ1) ∂t , ∂G (t) (µ2) ∂t ).Without loss of generality, assume∂G (t) (µ1) ∂t > ∂G (t) (µ2). We now analyze the condition when the magnitude of the update of the common token is less than the magnitude of the update of the signals. By Lemma C.1, we havem1 j=1 ∂a j w (t) j W (t) V µ 1 ∂t ≥ - m1 j=1 ∂a j w (t) j W (t) V µ 3 ∂t ⇔ 1 n ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 1 )∥ 2 2 i∈I1 g (t) i L l=1 p (t,i) q←l,k←µ1 - i∈I2 g (t) i L l=1 p (t,i) q←l,k←µ1 + m1 j=1 a j ε (t) (µ 1 ) ≥ 1 n ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 3 )∥ 2 2 i∈I2∪I3∪I4 g (t) i L l=1 p (t,i) q←l,k←µ3 - i∈I1 g (t) i L l=1 p (t,i) q←l,k←µ3 + m1 j=1 a j ε (t) (µ 3 ) ⇔ 1 n i∈I1 g (t) i L l=1 p (t,i) q←l,k←µ1 + 1 n ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 3 )∥ 2 2 ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 1 )∥ 2 2 i∈I1 g (t) i L l=1 p (t,i) q←l,k←µ3 ≥ 1 n ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 3 )∥ 2 2 ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 1 )∥ 2 2 i∈I2∪I3∪I4 g (t) i L l=1 p (t,i) q←l,k←µ3 + 1 n i∈I2 g (t) i L l=1 p (t,i) q←l,k←µ1 + m1 j=1 a j ε (t) (µ 3 ) ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 1 )∥ 2 2 - m1 j=1 a j ε (t) (µ 1 ) ∥W (t) ∥ 2 F + m 1 ∥v (t) (µ 1 )∥ 2 2 ⇔ 1 n (1 + o(1)) i∈I1 g (t) i + 1 n (1 + o(1)) i∈I1 g (t) i ≥ 1 n (1 + o(1)) i∈I2∪I3∪I4 g (t) i + 1 n (1 + o(1)) i∈I2 g (t) i ± O m 1 mwhere the last equality applies Corollary E.24. If m is sufficiently larger than m 1 (by some large constant factor C), then the last term is negligible. Note that the above implies that ifEQUATIONi = Ω(1) for all t ≤ T 1 . Now, by Proposition E.21, if the initialization level is small enough, then i∈I1 g(t) i -i∈I3 g (t) i ≥ Ω(1). Thus, if ∂G (t) (µ1) ∂t ≥ -∂G (t) (µ3) ∂t , then 1 n i∈I1 ∂g (t) i ∂t = -Θ((σ 2 0 + σ 2 1 )mm 1 ). Therefore, there must exist a time t ′ = Θ(1/m) such that ∂G (t ′ ) (µ1) ∂t = -∂G (t ′ ) (µ3) ∂t . Further, for t ≥ Θ(F/m), we have y i F (t) i > 0 for all i ∈ I 4 and i∈I4 g (t) i < min( i∈I2 g (t) i , i∈I3 g (t) i ), where F = max i∈[n] |F (0) i |. Now, consider a time point t ′ when ∂G (t ′ ) (µ1) ∂t = -∂G (t ′ ) (µ3) ∂t . At t ′ , we must have min( i∈I2 g (t ′ ) i , i∈I3 g (t ′ ) i ) -i∈I4 g (t ′ ) i = Ω(1). Thus, 2 i∈I1 g (t ′ ) i -2 i∈I4 g (t ′ ) i = Ω(1), which implies ∂ ∂F 2 i∈I1 g (t ′ ) i - ∂ ∂F i∈I4 g (t ′ ) i= Ω(1).For i ∈ I 1 , we have∂F (t ′ ) i ∂t = ∂G (t ′ ) (µ 1 ) ∂t + ∂G (t ′ ) (µ 2 ) ∂t + ∂G (t ′ ) (µ 3 ) ∂t + O L n σ 2 1 mm 1 = ∂G (t ′ ) (µ 2 ) ∂t + O L n σ 2 1 mm 1 . For i ∈ I 2 , ∂F (t) i ∂t = ∂G (t ′ ) (µ 1 ) ∂t + ∂G (t ′ ) (µ 3 ) ∂t + O L n σ 2 1 mm 1 = O L n σ 2 1 mm 1 .For i ∈ I 3 , by Proposition E.21,∂F (t) i ∂t = ∂G (t ′ ) (µ 2 ) ∂t + ∂G (t ′ ) (µ 3 ) ∂t + O L n σ 2 1 mm 1 = O F σ 2 1 mm 1 ,and for i ∈ I 4 ,∂F (t) i ∂t = ∂G (t ′ ) (µ 3 ) ∂t + O L n σ 2 1 mm 1 .Thus, by chain ruleEQUATIONThis implies that there exists a constant L such that∂ ∂t i∈I1 g (t) i - i∈I2 g (t) i + (1 + L) i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i = -Θ(σ 2 1 mm 1 ),which implies that there must exist a time T 0.5 ≤ O(1/m) such that for all t ≥ T 0.5 ,- ∂G (t) (µ 3 ) ∂t ≥ (1 + L) ∂G (t) (µ 1 ) ∂t .Moreover, by Theorem E.26, if C T1 is sufficiently large, we have T 0.5 < T 1 .Finally, consider the time when∂G (t) (µ1) ∂t + ∂G (t) (µ2) ∂t > -∂G (t) (µ3). During this time, note that we haveEQUATIONi = -Θ(σ 2 1 mm 1 ) for all i ∈ I 2 ∪ I 3 ∪ I 4 . Thus, there must exist a constant U such that- ∂G (t) (µ 3 ) ∂t ≤ (1 -U ) ∂G (t) (µ 1 ) ∂t + ∂G (t) (µ 2 ) ∂t for all t ≤ T 1 .Corollary E.20. There exists a small positive constant C such that if we define T ′ := min{t :min i∈I2∪I3 F (t) i ≥ C}, then T ′ ≤ O(1/m). Proposition E.21. Let F = max i∈[n] |F (0) i |. For t ≤ T 1 , we have 1 n i∈I2 g (t) i - 1 n i∈I3 g (t) i ≤ O (F ) .Proof. First of all, by Lipschitzness of g, we have1 n i∈I2 g (t) i - 1 n i∈I3 g (t) i ≤ |G (t) (µ 1 ) -G (0) (µ 1 ) -(G (t) (µ 2 ) -G (0) (µ 2 ))| + 2 max i∈[n] F (0) i + o(1).Without loss of generality, assume G (t) (µ 1 ) > G (t) (µ 2 ) for all t ≤ T 1 ; otherwise, we can break the interval [0, T 1 ] into sub-intervals by the time points when G (t) (µ 1 ) -G (t) (µ 2 ) changes its sign and then apply the analysis below to each sub-interval. We first derive∂G (t) (µ 1 ) ∂t - ∂G (t) (µ 2 ) ∂t = m1 j1=1 1 n i2∈I1∪I2 g (t) i2 y i2 L l2=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ1 + 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 1 - 1 n i2∈I1∪I3 g (t) i2 y i2 L l2=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ2 - 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 2 = m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 1 n i∈I1 g (t) i y i L l=1 p (t,i) q←l,k←µ1 - 1 n i∈I1 g (t) i y i L l=1 p (t,i) q←l,k←µ2 + m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 1 n i∈I2 g (t) i y i L l=1 p (t,i) q←l,k←µ1 - 1 n i∈I3 g (t) i y i L l=1 p (t,i) q←l,k←µ2 + m 1 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 1 - 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 2 = m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 1 n i∈I3 g (t) i - 1 n i∈I2 g (t) i + o(1) + m 1 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 1 - 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 2 ≤ m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 G (t) (µ 1 ) -G (0) (µ 1 ) -(G (t) (µ 2 ) -G (0) (µ 2 )) + G (0) (µ 1 ) + G (0) (µ 2 ) + O max i∈[n] |F (0) i | + o(1) + m 1 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 1 - 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 2 .By Theorem E.14, we have By Grönwall's inequality, we havem 1 T 0 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 1 - 1 n n i1=1 g (t) i1 y i1 L l1=1 a j1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ 2 dt ≤ T • O(Lσ 2 0 √ mm 1 ), T 0 G (0) (µ 1 ) + G (0) (µ 2 ) + O max i∈[n] |F (0) i | + o(1) m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 dt ≤ O G (0) (µ 1 ) + G (0) (µ 2 ) + max i∈[n] |F (0) i | T σ 2 1 mm 1 ,G (T ) (µ 1 ) -G (0) (µ 1 ) -(G (T ) (µ 2 ) -G (0) (µ 2 )) ≤ T • O(Lσ 2 0 √ mm 1 ) + O G (0) (µ 1 ) + G (0) (µ 2 ) + max i∈[n] |F (0) i | T σ 2 1 mm 1 + T 0 t • O(Lσ 2 0 √ m) + O G (0) (µ 1 ) + G (0) (µ 2 ) + max i∈[n] |F (0) i | tσ 2 1 m • O(1) dt ≤ O G (0) (µ 1 ) + G (0) (µ 2 ) + max i∈[n] |F (0) i | .This implies that1 n i∈I2 g (t) i - 1 n i∈I3 g (t) i ≤ O G (0) (µ 1 ) + G (0) (µ 2 ) + max i∈[n] |F (0) i | .Thus, if the initialization scale is sufficiently small, 1 n i∈I2 gi and 1 n i∈I3 gi are only differed by a small constant.Proposition E.22. Conditioned on the success of Lemma D.1 and Lemma D.3, with probability at least 1 -δ over the randomness of a, for all i ∈ [n] and µ, ν ∈ X (i) , we havem1 j=1 a j ∥k (0) (µ)∥ 2 2 w (0)⊤ j v (0) (ν) ≤ O σ 2 0 mσ 1 σ 0 mm 1 log m 1 d δ log nd δ , m1 j=1 a j ∥q (0) (µ)∥ 2 2 w (0)⊤ j v (0) (ν) ≤ O σ 2 0 mσ 1 σ 0 mm 1 log m 1 d δ log nd δ and for all l, l ′ ∈ [L] with K (0,i) l ̸ = k (0) (ν) and q (0,i) l ̸ = Q (0) (ν), we have m1 j=1 a j ν ⊤ W (0)⊤ K K (0,i) l V (0,i)⊤ l ′ w (0) j ≤ O(σ 2 0 √ mσ 0 σ 1 √ mm 1 ), m1 j=1 a j ν ⊤ W (0)⊤ Q q (0,i) l V (0,i)⊤ l ′ w (0) j ≤ O(σ 2 0 √ mσ 0 σ 1 √ mm 1 ).Proof. Fix i ∈ [n] and µ, ν ∈ X (i) . Consider the randomness of a. By Lemma D.1, Corollary D.2 and Lemma D.3, a j ∥k (0) (µ)∥ 2 2 w (0)⊤ j v (0) (ν) is a sub-Gaussian random variable with variance proxy O((σ 2 0 mσ 1 σ 0 √ m log(dm 1 /δ)) 2 ). Then the following inequality holds.m1 j=1 a j ∥k (0) (µ)∥ 2 2 w (0)⊤ j v (0) (ν) ≤ O σ 2 0 mσ 1 σ 0 mm 1 log m 1 d δ log 2 δ .Finally, take a union bound over i ∈ [n], µ, ν ∈ {µ i } d i=1 . The analysis for the second term is similar.Next, note thata j ν ⊤ W (0)⊤ K K (0,i) l V (0,i)⊤ l ′ w (0) j is a sub-Gaussian random variable with variance proxy O((σ 2 0 √ mσ 0 σ 1 √ m) 2 ). Thus, m1 j=1 a j ν ⊤ W (0)⊤ K K (0,i) l V (0,i)⊤ l ′ w (0) j ≤ O(σ 2 0 √ mσ 0 σ 1 √ mm 1 ).Lemma E.23 (Score change). For all t ≤ T 1 , for all ν, µ ∈ {µ i } d i=1 , we have∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t ≤ 1 √ m i : µ ∈ X (i) + i : ν ∈ X (i) n O 1 L + 1 √ mand thus,ν ⊤ W (t)⊤ K W (t) Q µ -ν ⊤ W (0)⊤ K W (0) Q µ ≤ t 1 √ m i : µ ∈ X (i) + i : ν ∈ X (i) n O 1 L + 1 √ m .Proof. First of all, by Lemma C.4, we expand the per step gradient descent update as follows:EQUATIONEQUATION+ 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l I(q (t,i) l ̸ = q (t) (µ)) (4)Analysis of (1): By triangle inequality, we haveEQUATIONEQUATION.To analyze (a), since(a + ε 1 )(b + ε 2 ) = ab + aε 2 + bε 1 + ε 1 ε 2 , we have m1 j=1 a j ∥k (t) (ν)∥ 2 2 w (t)⊤ j v (t) (ν)p (t,i) q←µ,k←ν - m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν)p (0,i) q←µ,k←ν ≤ m1 j=1 a j ∥k (t) (ν)∥ 2 2 w (t)⊤ j v (t) (ν) - m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν) p (t,i) q←µ,k←ν + m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν) • p (t,i) q←µ,k←ν -p (0,i) q←µ,k←ν .Further, by Lemma D.1, Lemma D.3, Definition E.1 and Definition E.2, we haveEQUATION)Combining with Proposition E.22, this impliesm1 j=1 a j ∥k (t) (ν)∥ 2 2 w (t)⊤ j v (t) (ν)p (t,i) q←µ,k←ν - m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν)p (0,i) q←µ,k←ν ≤ m 1 • max ν,j ∥k (t) (ν)∥ 2 2 w (t)⊤ j v (t) (ν) -∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν) p (t,i) q←µ,k←ν + m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν) • p (t,i) q←µ,k←ν -p (0,i) q←µ,k←ν ≤ m 1 R K O(σ 0 σ 1 √ m) + R O(σ 2 0 m) + R K R 2 L + R P + O R P σ 2 0 mσ 1 σ 0 √ mm 1 .Thus,|(a)| = m1 j=1 a j ∥k (t) (ν)∥ 2 2 w (t)⊤ j v (t) (ν)p (t,i) q←µ,k←ν ≤ m1 j=1 a j ∥k (t) (ν)∥ 2 2 w (t)⊤ j v (t) (ν)p (t,i) q←µ,k←ν - m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν)p (0,i) q←µ,k←ν + m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j v (0) (ν)p (0,i) q←µ,k←ν ≤ m 1 R K O(σ 0 σ 1 √ m) + R O(σ 2 0 m) + R K R 2 L + R P + O R P σ 2 0 mσ 1 σ 0 √ mm 1 + O σ 2 0 mσ 1 σ 0 √ mm 1 1 L .On the other hand, to analyze (b), we haveEQUATIONNote thatEQUATIONwhich impliesEQUATION≤ maxl∈[L] m1 j=1 a j ∥k (t) (ν)∥ 2 2 w (t)⊤ j V (t,i) l - m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j V (0,i) l + LR P max l∈[L] m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j V (0,i) l • p (t,i) q←µ,k←ν + m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j V (0,i) p (0,i) l(i,µ) • p (t,i) q←µ,k←ν -p (0,i) q←µ,k←ν + m1 j=1 a j ∥k (0) (ν)∥ 2 2 w (0)⊤ j V (0,i) p (0,i) l(i,µ) p (0,i) q←µ,k←ν (iii) ≤ m 1 R K O(σ 0 σ 1 √ m) + R O(σ 2 0 m) + R K R + LR P O(σ 2 0 mσ 0 σ 1 √ mm 1 ) • 1 L + R P + O σ 2 0 mσ 0 σ 1 √ mm 1 1 L + R Pwhere (i) follows from Equation ( 10), (ii) follows from Equation ( 11) and (iii) follows from Equation ( 9) and Proposition E.22. Combining the upper bound for both (a) and (b), we obtain|(1)| ≤ 1 √ m i : µ, ν ∈ X (i) n O m 1 R K O(σ 0 σ 1 √ m) + R O(σ 2 0 m) + R K R + LR P O(σ 2 0 mσ 0 σ 1 √ mm 1 ) • 1 L + R P + σ 2 0 mσ 0 σ 1 √ mm 1 1 L + R P = 1 √ m i : µ, ν ∈ X (i) n O(1/L).Analysis of (2): (2) can be analyzed similarly as (1) with ∥k(0) (ν)∥ 2 2 replaced by ν ⊤ W (t) K K (t,i) l. Thus, we only need to replace σ 2 0 m with σ 2 0 √ m and then take a sum over l. We have|(2)| ≤ 1 √ m i : µ ∈ X (i) n O m 1 R K O(σ 0 σ 1 √ m) + R O(σ 2 0 √ m) + R K R + LR P O(σ 2 0 √ mσ 0 σ 1 √ mm 1 ) • (1 + LR P ) + σ 2 0 √ mσ 0 σ 1 √ mm 1 (1 + LR P ) = 1 √ m i : µ ∈ X (i) n O(1/ √ m).Analysis of (3): (3) can be analyzed similarly to (1), and we obtain|(3)| ≤ 1 √ m i : µ, ν ∈ X (i) n O(1/L).Analysis of (4): (4) can be analyzed similarly to (2), and we obtain|(4)| ≤ 1 √ m i : ν ∈ X (i) n O(1/ √ m).Finally, combining the bounds on (1) -(4), we have∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t ≤ 1 √ m i : µ, ν ∈ X (i) n O(1/L) + 1 √ m i : µ ∈ X (i) + i : ν ∈ X (i) n O(1/ √ m) ≤ 1 √ m i : µ ∈ X (i) + i : ν ∈ X (i) n O(1/L + 1/ √ m).Corollary E.24 (Softmax change). For t ≤ T 1 , we have the following:• if both X (i) l1 , X (i) l2 ∈ R, then ∂p (t,i) l1,l2 ∂t ≤ O 1 L(L 2 + n) √ m ,• otherwise,∂p (t,i) l1,l2 ∂t ≤ O 1 L 2 √ m .Thus,• if both X (i) l1 , X (i) l2 ∈ R, then p (t,i) l1,l2 -p (0,i) l1,l2 ≤ O 1 L + 1 n 1 Lm 2 1 L + 1 √ m ,• otherwise,p (t,i) l1,l2 -p (0,i) l1,l2 ≤ O 1 Lm 2 1 L + 1 √ m .Proof. Consider fixed i, l 1 , l 2 . By Lemma G.2, we have∂p (t,i) l1,l2 ∂t ≤ p (t,i) l1,l2 ∂s (t,i) l1,l2 ∂t + p (t,i) l1,l2 p (t,i)⊤ l1 ∂s (t,i) l1 ∂t .By Lemma E.23, we have the following cases:• if both X (i) l1 , X (i) l2 ∈ R, then ∂s (t,i) l1,l2 ∂t ≤ O 1 nm 1 L + 1 √ m ;• otherwise,∂s (t,i) l1,l2 ∂t ≤ O 1 m 1 L + 1 √ m .Next, during Phase 1, we have p(t,i) l1,l2 ≤ 1 L + 1 Lm + R P . Therefore, • if X (i) l1 ∈ R, then p (t,i)⊤ l1 ∂s (t,i) l1 ∂t ≤ O 1 L + 1 n 1 m 1 L + 1 √ m ; • otherwise, p (t,i)⊤ l1 ∂s (t,i) l1 ∂t ≤ O 1 m 1 L + 1 √ m .Thus,• if both X (i) l1 , X (i) l2 ∈ R, then ∂p (t,i) l1,l2 ∂t ≤ O 1 L + 1 n 1 Lm 1 L + 1 √ m ; • otherwise, ∂p (t,i) l1,l2 ∂t ≤ O 1 Lm 1 L + 1 √ m .This implies that• if both X (i) l1 , X (i) l2 ∈ R, then p (t,i) l1,l2 -p (0,i) l1,l2 ≤ O 1 L + 1 n 1 Lm 2 1 L + 1 √ m ;• otherwise,p (t,i) l1,l2 -p (0,i) l1,l2 ≤ O 1 Lm 2 1 L + 1 √ m . Lemma E.25 (K, Q self-correlation change). For t ≤ T 1 , for µ, ν ∈ {µ i } d i=1 , we have µ ⊤ W (t)⊤ K W (t) K ν -µ ⊤ W (0)⊤ K W (0) K ν ≤ t 1 √ m i : µ ∈ X (i) + i : ν ∈ X (i) n O 1 √ m , µ ⊤ W (t)⊤ Q W (t) Q ν -µ ⊤ W (0)⊤ Q W (0) Q ν ≤ t 1 √ m i : µ ∈ X (i) + i : ν ∈ X (i) n O 1 √ m .Proof. By Lemma C.4, we only need to replace O(σ 2 0 m) by O(σ 2 0 √ m) in the proof of Lemma E.23. Thus, we omit the proof here.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 ≤ C T1 /(σ 2 1 mm 1 ) (where the constant C T1 is from the definition of T 1 in Definition E.2), all of the following hold:1. max j∈[m],µ∈{µi} 3 i=1 w (t) j W (t) V µ -w (0) j W (0) V µ ≤ R, where R = O(1/m 1 ); 2. max j∈[m],µ / ∈{µi} 3 i=1 w (t) j W (t) V µ -w (0) j W (0) V µ ≤ O(R/n + R/ √ m); 3. max µ,ν∈{µi} d i=1 µ ⊤ W (t)⊤ Q W (t) K ν -µ ⊤ W (0)⊤ Q W (0) K ν ≤ R S , where R S ≤ O(1/m √ m); 4. max µ,ν∈{µi} d i=1 µ ⊤ W (t)⊤ Q W (t) Q ν -µ ⊤ W (0)⊤ Q W (0) Q ν ≤ R Q ; 5. max µ,ν∈{µi} d i=1 µ ⊤ W (t)⊤ K W (t) K ν -µ ⊤ W (0)⊤ K W (0) K ν ≤ R K . Thus, T 1 = C T1 /(σ 2 1 mm 1 ).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.Theorem E.27 (End of Phase 1). At t = T 1 , we haveEQUATIONG (T1) (µ 1 ) = Θ(1), G (T1) (µ 2 ) = Θ(1), -G (T1) (µ 3 ) = Θ(1), ∀µ ∈ R : G (T1) (µ) = O(σ 0 σ 1 √ mm 1 ). Further, m1 j1=1 m1 j2=1 a j1 w (T1) j1 , a j2 w (T1) j2 = Θ(σ 2 1 mm 1 ).Proof. By Theorem E.26, we have y i F (T1) i= O(1) andG (T1) (µ 1 ) = O(1), G (T1) (µ 2 ) = O(1), -G (T1) (µ 3 ) = O(1).Further, by Theorem E.14 and Theorem E.19, we have y i F (T1) i≥ Ω(1) for i ∈ I 1 andG (T1) (µ 1 ) ≥ Ω(1), G (T1) (µ 2 ) ≥ Ω(1).Finally, by Theorem E.26, we have that T 1 = C T1 /m. And note that if the constant C T1 in Definition E.2 is sufficiently large, then by Corollary E.20, we have T ′ ≤ T 1 . Thus, we havey i F (T1) i≥ Ω(1) for i ∈ I 2 ∪ I 3 ∪ I 4 and -G (T1) (µ 3 ) ≥ Ω(1).By Lemma D.7 and Theorem E.15, we have∀µ ∈ R : G (T1) (µ) = O(σ 0 σ 1 √ mm 1 ).Finally, by Lemma D.3, Proposition E.5 and Proposition E.12, we have = Θ(σ 2 1 mm 1 ).Theorem E.28 (Phase 1, formal restatement of Theorem 3.1). With probability at least 1 -δ over the randomness of weight initialization, there exists a timeT 1 = O(1/m) such that • G (T1) (µ 1 ) ≥ Ω(1), G (T1) (µ 2 ) ≥ Ω(1), G (T1) (µ 3 ) ≤ -Ω(1).• All the training samples are correctly classified:y i F (T1) i = Ω(1) for all i ∈ [n]. • For t ∈ [0, T 1 ], w (t) j1 , w (t) j2 -w (0) j1 , w (0) j2 ≤ O 1 m v (t) (µ), v (t) (ν) -v (0) (µ), v (0) (ν) ≤ O 1 m ν ⊤ W (t)⊤ K W (t) Q µ -ν ⊤ W (0)⊤ K W (0) Q µ ≤ O 1 m 3/2 1 L + 1 √ m µ ⊤ W (t)⊤ K W (t) K ν -µ ⊤ W (0)⊤ K W (0) K ν ≤ O 1 m 2 µ ⊤ W (t)⊤ Q W (t) Q ν -µ ⊤ W (0)⊤ Q W (0) Q ν ≤ O 1 m 2• The training loss satisfies L (T1) = Θ(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 on the training loss is a direct consequence of Definition E.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 Ω(poly(m)) time, which will serve as the end of Phase 2.Definition F.1. We define Phase 2 of the training to be t∈ [T 1 , T 2 ] such that • The change of K, Q self-correlation is small: max µ,ν µ ⊤ W (t)⊤ K W (t) K ν -µ ⊤ W (0)⊤ K W (0) K ν = O(σ 2 0 √ m) max µ,ν µ ⊤ W (t)⊤ Q W (t) Q ν -µ ⊤ W (0)⊤ Q W (0) Q ν = O(σ 2 0 √ m)• The change of softmax probability satisfies:max l1,l2∈[L], i∈[n] p (t,i) l1,l2 -p (0,i) l1,l2 < O(1/L 2 )• The sum of neuron correlation satisfiesm1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2≤ O(σ 2 1 mm 1 ).• The gradient of W, W V satisfiesmax µ∈{µi} d i=1 m1 j=1 a j ∂w (t) j ∂t W (t) V µ m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 ≤ o(1/L) 1 n i∈[n] g (t) i max i∈[n] L l1=1 L l2=1 G (t) (X (i) l2 ) ∂p (t,i) l 1 ,l 2 ∂t m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 ≤ o(1) 1 n i∈[n] g (t) i • The gradients i∈I g (t) i for I ∈ {I 1 , I 2 , I 3 , I 4 } satisfies i∈I4 g i ≤ min i∈I2 g (t) i , i∈I3 g (t) i ≤ max i∈I2 g (t) i , i∈I3 g (t) i ≤ i∈I1 g (t) i ≤ i∈I2∪I3∪I4 g (t) i and 1 2 ≤ i∈I2 g (t) i i∈I3 g (t) i ≤ 2. • y i F (t) i > C for all i ∈ [n] for some fixed constant C. • |G (t) (µ)| ≤ O(log m) for µ ∈ {µ i } 3 i=1 . • L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 )p (t,i) q←l1,k←l2 ≤ O(1). • T 2 ≤ O(poly(m)). Corollary F.2. For t ∈ [T 1 , T 2 ], if i, j ∈ I where I ∈ {I 1 , I 2 , I 3 , I 4 }, then max i,j∈I g (t) i g (t) j ≤ O(1).Proof. Take I = I 1 and the proof is similar for the remaining cases. For fixed i, j ∈ I 1 , we haveg (t) i g (t) j = 1 + exp(y j F (t) j ) 1 + exp(y i F (t) i ) ≤ 2 exp(y j F (t) j ) exp(y i F (t) i ), where the inequality is due toy i F (t) i ≥ C in Definition F.1. Now we consider exp(y j F (t) j ) exp(y i F (t) i ) = exp(y j L l1=1 L l2=1 G (t) (X (j) l2 )p (t,j) q←l1,k←l2 ) exp(y i L l1=1 L l2=1 G (t) (X (i) l2 )p (t,i) q←l1,k←l2 ) = exp     L l1=1 µ∈{µ k } 3 k=1 G (t) (µ)p (t,j) q←l1,k←µ   -   L l1=1 µ∈{µ k } 3 k=1 G (t) (µ)p (t,i) q←l1,k←µ     (1) • exp       L l1=1 l2: X (j) l 2 ∈{µ k } d k=4 G (t) (X (j) l2 )p (t,j) q←l1,k←l2    -    L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 )p (t,i) q←l1,k←l2       (2). By Definition F.1, it is easy to see that |(2)| ≤ O(1). For (1), we have  L l1=1 µ∈{µ k } 3 k=1 G (t) (µ)p (t,j) q←l1,k←µ   -   L l1=1 µ∈{µ k } 3 k=1 G (t) (µ)p (t,i) q←l1,k←µ   ≤ µ∈{µ k } 3 k=1 G (t) (µ)(1 ± o(1)) - µ∈{µ k } 3 k=1 G (t) (µ)(1 ± o(1)) ≤ O(1).Thus, we haveexp(y j F (t) j ) exp(y i F (t) i ) ≤ O(1) for t ∈ [T 1 , T 2 ].Lemma F.3. Phase 2 in Definition F.1 is well-defined.Proof. We need to show that the definition is valid at t = T 1 with conditions satisfied with strict inequality. Then since everything changes continuously, there naturally exists a T 2 > T 1 such that Definition F.1 is well-defined.First of all, by Lemma E.25, for µ, ν ∈ {µ i } d i=1 , we havemax µ,ν µ ⊤ W (T1)⊤ K W (T1) K ν -µ ⊤ W (0)⊤ K W (0) K ν = O(σ 2 0 √ m), max µ,ν µ ⊤ W (T1)⊤ Q W (T1) Q ν -µ ⊤ W (0)⊤ Q W (0) Q ν = O(σ 2 0 √ m).Next, by Proposition E.12, Proposition E.5 and Lemma D.3, we havem1 j1=1 m1 j2=1 a j1 w (T1) j1 , a j2 w (T1) j2 = Θ(σ 2 1 mm 1 ).Recall that 1 n n i=1 g (T1) i = Θ(1). On the other hand, by Corollary E.10, Lemma D.3, we havemax µ∈{µi} d i=1 m1 j=1 a j ∂w (T1) j ∂t W (T1) V µ = O(σ 2 0 mm 1 ) =⇒ max µ∈{µi} d i=1 m1 j=1 a j ∂w (T 1 ) j ∂t W (T1) V µ m1 j1=1 m1 j2=1 a j1 w (T1) j1 , a j2 w (T1) j2 ≤ O σ 2 0 σ 2 1 = O m 1 Lm .Also, by Corollary E.24, we havemax i∈[n] L l1=1 L l2=1 G (T1) (X (i) l2 ) ∂p (T1,i) l1,l2 ∂t = O 1 √ m =⇒ max i∈[n] L l1=1 L l2=1 G (T1) (X (i) l2 ) ∂p (T 1 ,i) l 1 ,l 2 ∂t m1 j1=1 m1 j2=1 a j1 w (T1) j1 , a j2 w (T1) j2 ≤ O 1 m 3/2 ≤ o(1) i∈[n] g (T1) i .Further, by Corollary E.24, we havemax l1,l2∈[L], i∈[n] p (T1,i) l1,l2 -p (0,i) l1,l2 < O(1/L 2 ).Now, by Proposition E.21, if F is small enough (which can be achieved by making the initialization scale small enough), then1 2 < i∈I2 g (T1) i i∈I3 g (T1) i < 2.Lastly, by Theorem E.27, we have y i F (T1) i≥ Ω(1). And it is straightforward to see that|G (T1) (µ)| ≤ O(log m) for µ ∈ {µ i } 3 i=1 .Finally, we proveEQUATION. A simple corollary from Lemma D.7 and Lemma D.8 is thatL l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (0) (X (i) l2 )p (0,i) q←l1,k←l2 ≤ O(1). Next, we have L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (T1) (X (i) l2 )p (T1,i) q←l1,k←l2 - L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (0) (X (i) l2 )p (0,i) q←l1,k←l2 ≤ L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 (G (T1) (X (i) l2 ) -G (0) (X (i) l2 )) p (0,i) q←l1,k←l2 + L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (0) (X (i) l2 ) p (T1,i) q←l1,k←l2 -p (0,i) q←l1,k←l2 + L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (T1) (X (i) l2 ) -G (0) (X (i) l2 ) p (T1,i) q←l1,k←l2 -p (0,i) q←l1,k←l2 ≤ O(1) which proves that L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (T1) (X (i) l2 )p (T1,i) q←l1,k←l2 ≤ O(1).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.Lemma F.4. For t ∈ [T 1 , T 2 ], we have ∂ ∂t m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 = 2m 1 n n i=1 g (t) i y i F (t) i > 0 and thus, m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2= Ω(σ 2 1 mm 1 ).Proof. By Lemma C.2, we obtain∂ ∂t m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 = m1 j1=1 m1 j2=1 1 n n i=1 g (t) i y i L l1=1 a j1 w (t)⊤ j1 V (t,i) p (t,i) l1 + 1 n n i=1 g (t) i y i L l2=1 a j2 w (t)⊤ j2 V (t,i) p (t,i) l2 = 2m 1 n n i=1 g (t) i y i F (t) i .Then, note that the final term is positive since by Definition F.1, we havey i F (t) i > 0 for all i ∈ [n]. Finally, by Theorem E.27, we have for all t ∈ [T 1 , T 2 ], m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2≥ Ω(σ 2 1 mm 1 ).Lemma F.5 (Same as Lemma 4.4). For t ∈ [T 1 , T 2 ], there exists a small constant C ≪ 1 such that i∈I2 g(t) i -i∈I3 g (t) i min( i∈I2 g (t) i , i∈I3 g (t) i ) ≤ C.Proof. Without loss of generality, assume i∈I2 gEQUATIONi . Then, we havei∈I2 g (t) i -i∈I3 g (t) i min( i∈I2 g (t) i , i∈I3 g (t) i ) = i∈I2 g (t) i -i∈I3 g (t) i i∈I3 g (t) i = i∈I2 g (t) i i∈I3 g (t) i -1.Now by the quotient rule, we haveEQUATIONNote that i∈I gEQUATION.By Definition F.1, we have that for i ′ ∈ I 2 ,∂y i ′ F (t) i ′ ∂t = y i ′ L l1=1 m1 j=1 L l2=1   ∂a j w (t) j W (t) V X (i ′ ) l2 ∂t p (t,i ′ ) l1,l2 + a j w (t) j W (t) V X (i ′ ) l2 ∂p (t,i ′ ) l1,l2 ∂t   = y i ′ L l2=1 ∂G (t) (X (i ′ ) l2 ) ∂t (1 + o(1)) + L l1=1 L l2=1 y i ′ G (t) (X (i ′ ) l2 ) ∂p (t,i ′ ) l1,l2 ∂t = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i + i∈I1 g (t) i - i∈I2 g (t) i + o(1) i∈[n] g (t) i   .Similarly, for i ′ ∈ I 3 , we haveEQUATION .t ∈ [T 1 , T 2 ], if i∈I2 g (t) i is sufficiently larger than i∈I3 g (t) i (i.e., i∈I2 g (t) i / i∈I3 g (t) iis bigger than some threshold), then∂G (t) (µ1) ∂t < ∂G (t) (µ2) ∂t and min i∈I2 ∂yiF (t) i ∂t > max i∈I3 ∂yiF (t) i ∂t. Thus, the ratio will decrease. Similarly, if i∈I2 g(t) i / i∈I3 g (t) iis too small, the ratio will increase. Next, we compute a bound onEQUATION,EQUATION,we have1 i∈I 3 g (t) i i∈I3 1 (1+exp(yiF (t) i ))(1+exp(-yiF (t) i )) 1 i∈I 2 g (t) i i∈I2 1 (1+exp(yiF (t) i ))(1+exp(-yiF (t) i )) ∈ 1 + min i∈I2 exp(-y i F (t) i ) 1 + max i∈I3 exp(-y i F (t) i ) , 1 + max i∈I2 exp(-y i F (t) i ) 1 + min i∈I3 exp(-y i F (t) i ). By Definition F.1, we have-i∈I1 g (t) i + i∈I2∪I3∪I4 g (t) i -i∈I1 g (t) i + i∈I2 g (t) i ± o(1) i∈[n] g (t) i i∈I3 g (t) i ≤ 6. Thus, i∈I2 g (t) i i∈I3 g (t) i ∈ 1 ± 6 • 1 + min i∈I2 exp(-y i F (t) i ) 1 + max i∈I3 exp(-y i F (t) i ) , 1 + max i∈I2 exp(-y i F (t) i ) 1 + min i∈I3 exp(-y i F (t) i ) -1 . Lemma F.6 (Complete version of Lemma 4.5). For t ∈ [T 1 , T 2 ], we have i∈[n] g (t) i i∈I2∪I3∪I4 g (t) i -i∈I1 g (t) i = O(1), i∈[n] g (t) i i∈I1 g (t) i -i∈I2 g (t) i = O(1) i∈[n] g (t) i i∈I1 g (t) i -i∈I3 g (t) i = O(1), ∂G (t) (µ1) ∂t ∂G (t) (µ2) ∂t = Θ(1).Further, there exists a constant C such that(1 + C) max ∂G (t) (µ 1 ) ∂t , ∂G (t) (µ 2 ) ∂t ≤ - ∂G (t) (µ 3 ) ∂t ≤ (1 -C) ∂G (t) (µ 1 ) ∂t + ∂G (t) (µ 2 ) ∂t .Proof. Without loss of generality, assume∂G (t) (µ1) ∂t > ∂G (t) (µ2) ∂t. First of all, by Definition F.1, we have∂G (t) (µ 1 ) ∂t = m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2 g (t) i ± o(1) i∈[n] g (t) i   , ∂G (t) (µ 2 ) ∂t = m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I3 g (t) i ± o(1) i∈[n] g (t) i   , ∂G (t) (µ 3 ) ∂t = m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i ± o(1) i∈[n] g (t) i   .This implies that∂G (t) (µ1) ∂t ∂G (t) (µ2) ∂t = i∈I1 g (t) i -i∈I2 g (t) i ± o(1) i∈[n] g (t) i i∈I1 g (t) i -i∈I3 g (t) i ± o(1) i∈[n] g (t) i , -∂G (t) (µ3) ∂t ∂G (t) (µ1) ∂t = -i∈I1 g (t) i + i∈I2∪I3∪I4 g (t) i ± o(1) i∈[n] g (t) i i∈I1 g (t) i -i∈I2 g (t) i ± o(1) i∈[n] g (t) i . We next analyze ∂ ∂t - ∂G (t) (µ 3 ) ∂t ∂G (t) (µ 1 ) ∂t . We first define the ratio R(t) = i∈I 2 ∪I 3 ∪I 4 g (t) i -i∈I 1 g (t) i i∈I 1 g (t) i -i∈I 2 g (t) i. Since the dependence of R on t is clear and for the ease of notation, we omit this dependence below.Rearranging this definition, we obtainEQUATIONWe consider the range of R ∈ [1, 3]. By Definition F.1, we haveEQUATIONwhere the second equation follows from Lemma F.5. This implies that∂G (t) (µ1) ∂t ∂G (t) (µ2) ∂t = Θ(1), i∈[n] g (t) i i∈I1 g (t) i -i∈I3 g (t) i = O(1),which proves the first result andEQUATIONRecall that by Definition F.1, we have for i 1 ∈ I 1 ,∂y i1 F (t) i1 ∂t = m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   3 i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i - i∈I2∪I3 g (t) i + o(1) i∈[n] g (t) i   ; for i 2 ∈ I 2 , ∂y i2 F (t) i2 ∂t = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i + i∈I1 g (t) i - i∈I2 g (t) i + o(1) i∈[n] g (t) i   ; for i 3 ∈ I 3 , ∂y i3 F (t) i3 ∂t = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i + i∈I1 g (t) i - i∈I3 g (t) i + o(1) i∈[n] g (t) i   ;and for i 4 ∈ I 4 ,∂y i4 F (t) i4 ∂t = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i + o(1) i∈[n] g (t) i   .The above implies that for i 1 ∈ I 1 ,∂yi 1 F (t) i 1 ∂t i∈I1 g (t) i -i∈I2 g (t) i = m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n 1 + i∈I1 g (t) i -i∈I3 g (t) i i∈I1 g (t) i -i∈I2 g (t) i -R + o(1) ; for i 2 ∈ I 2 , ∂yi 2 F (t) i 2 ∂t i∈I1 g (t) i -i∈I2 g (t) i = - m1 j1=1 m1 j2=1 a j1 w (t)j1 , a j2 w (t) j2 1 n (1 -R + o(1)) ; for i 3 ∈ I 3 , ∂yi 3 F (t) i 3 ∂t i∈I1 g (t) i -i∈I2 g (t) i = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n i∈I1 g (t) i -i∈I3 g (t) i i∈I1 g (t) i -i∈I2 g (t) i -R + o(1) ; and for i 4 ∈ I 4 , ∂yi 1 F (t) i 1 ∂t i∈I1 g (t) i -i∈I2 g (t) i = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n (-R + o(1)EQUATIONBy Lemma F.5, we obtaini∈I1 g (t) i -i∈I3 g (t) i i∈I1 g (t) i -i∈I2 g (t) i -1 ≤ ε,where we use ε to denote the small deviation. Note that Equation ( 16) is a quadratic inequality in R and can be rearranged as aR 2 + bR + c ≥ 0, wherea = i1∈I1 g ′ (F (t) i1 ) + i2∈I2 g ′ (F (t) i2 ), b = i3∈I3 g ′ (F (t) i3 ) + i4∈I4 g ′ (F (t) i4 ) - i1∈I1 g ′ (F (t) i1 ), c = (1 + o(1)) - i2∈I2 g ′ (F (t) i2 ) -(1 ± ε) i3∈I3 g ′ (F (t) i3 ) -(2 ± ε) i1∈I1 g ′ (F (t)i1 ) .Now we analyze the equality condition in Equation ( 16) where we calculate the root of the equationEQUATION.By Equation ( 14), we have-b 2a ≥ -1/4 and b a ≤ (1 + O(max i∈[n] exp(-y i F (t) i ))) max i1∈I1 g(F (t) i1 ) -i3∈I3 g(F (t) i3 ) i1∈I1 g(F (t) i1 ) + i2∈I2 g(F (t) i2 ) , i4∈I4 g(F (t) i4 ) i1∈I1 g(F (t) i1 ) + i2∈I2 g(F (t) i2 ) ≤ C < 1, and c a = (1 ± O(ε) ± O(max i∈[n] exp(-y i F (t) i ))) • 2.Recall that we are considering the case of R ∈ [1, 3]. Thus, we only need to consider the root that is positive and we can calculate the rootR ⋆ = -b 2a + b 2 4a 2 -c a ∈ (1, 2) if ε and max i∈[n] exp(-y i F (t) i ) are both suffiently small. Next, since g ′ (F (t) i ) < 0, 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 1 implies that R(T 1 ) ∈ [1, 3], which completes the proof.Corollary F.7. For t ∈ [T 1 , T 2 ], we haveG (t) (µ 1 ) G (t) (µ 2 ) , G (t) (µ 1 ) -G (t) (µ 3 ) , G (t) (µ 2 ) -G (t) (µ 3 ) = Θ(1).Proof. We first prove G (t) (µ1) G (t) (µ2) = Θ(1). Following from Lemma F.6, we have∂G (t) (µ 1 ) ∂t ∂G (t) (µ 2 ) ∂t = Θ(1). By Theorem E.27, we have G (T 1 ) (µ1) G (T 1 ) (µ2) = Θ(1). Thus, for t ∈ [T 1 , T 2 ], we obtainG (t) (µ 1 ) G (t) (µ 2 ) = G (T1) (µ 1 ) + t T1 ∂G (τ ) (µ1) ∂τ dτ G (T1) (µ 2 ) + t T1 ∂G (τ ) (µ2) ∂τ dτ = Θ(1).Note that Lemma F.6 and Definition F.1 imply that∂G (t) (µ 1 ) ∂t - ∂G (t) (µ 3 ) ∂t = Θ(1). Since G (T 1 ) (µ1) -G (T 1 ) (µ3) = Θ(1), similarly, we have G (t) (µ1) -G (t) (µ3) = Θ(1). Finally, G (t) (µ1) G (t) (µ2) = Θ(1) and G (t) (µ1) -G (t) (µ3) = Θ(1) imply that G (t) (µ2) -G (t) (µ3) = Θ(1).F.2 How Fast the Loss Decreases Theorem F.8. For t ∈ [T 1 , T 2 ], we haveL(t) = 1 Θ(σ 2 1 mm 1 )(t -T 1 ) + (1/ L(T 1 )).Proof. First of all, the gradient flow update for the empirical loss is given by∂ L ∂t = n i=1 ℓ ′ (y i F (t) i ) ∂yiF (t) i ∂t. By Definition F.1, we have for i 1 ∈ I 1 ,EQUATIONi + o(1)i∈[n] g (t) i   ; for i 3 ∈ I 3 , ∂y i3 F (t) i3 ∂t = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i + i∈I1 g (t) i - i∈I3 g (t) i + o(1) i∈[n] g (t) i   ;and for i 4 ∈ I 4 ,∂y i4 F (t) i4 ∂t = - m1 j1=1 m1 j2=1 a j1 w (t)j1 , a j2 w(t) j2 1 n   i∈I1 g (t) i - i∈I2∪I3∪I4 g (t) i + o(1) i∈[n] g (t) i   .By Lemma F.6, we have∂y i F (t) i ∂t = m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 Θ   1 n i∈[n] g (t) i   for all i ∈ [n]. Therefore, ∂ L ∂t = 1 n n i=1 ℓ ′ (y i F (t) i ) ∂y i F (t) i ∂t = 1 n n i=1 -g (t) i m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 Θ   1 n i ′ ∈[n] g (t) i ′   = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 Θ      1 n i∈[n] g (t) i   2    = - m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 Θ L 2 ,where the last equality follows from the property of binary cross-entropy loss that ℓ(x) = Θ(-ℓ ′ (x)) for x > 0. By Definition F.1 and Lemma F. = Θ(σ 2 1 mm 1 ).Thus, we have∂ L ∂t = -Θ(σ 2 1 mm 1 ) L 2 .Now, consider the differential equation dL dt = -C 1 L 2 . Note that this is a separable differential equation in t and we can solve it by1 L 2 dL dt + C 1 = 0 ⇒ d dt C 1 t -L -1 + C 2 = 0 ⇒ L(t) = 1 C 1 t + C 2 .This implies for t ∈ [T 1 , T 2 ], we haveL(t) = 1 Θ(σ 2 1 mm 1 )(t -T 1 ) + (1/ L(T 1 ))Corollary F.9. The following bound holds: Proof. This is a direct consequence of Lemma F.4 and Theorem F.8. Proof. By Lemma F.4, we have∂ ∂t m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 = 2m 1 n n i=1 g (t) i y i F (t) i .By Theorem F.8, we have = Θ(σ 2 1 mm 1 ).1 n n i=1 g (t) i = O( L (t) ) = O 1 (σ 2 1 mm1)(t-i y i F (τ ) i dτ ≤ O(m 1 log m) t T1 O 1 (σ 2 1 mm 1 )(τ -T 1 ) + 1/ L(T 1 ) dτ = O m 1 log m log t σ 2 1 mm 1 = O m 1 log 2 m σ 2 1 mmWe now analyze the correlation term with value-transformed data. Lemma F.11 (Growth of correlation of value-transformed data). For µ, ν ∈ {µ i } d i=1 , we have∂ ∂t µ ⊤ W (t)⊤ V W (t) V ν = G (t) (ν) 1 n i:µ∈X (i) g (t) i y i L l=1 p (t,i) q←l,k←µ + G (t) (µ) 1 n i:ν∈X (i) g (t) i y i L l=1 p (t,i) q←l,k←ν .Thus, for t ∈ [T 1 , T 2 ], we havemax µ,ν µ ⊤ W (t)⊤ V W (t) V ν -µ ⊤ W (T1)⊤ V W (T1) V ν ≤ O log m log t σ 2 1 mm 1 .Proof. By the gradient flow update, we have∂µW (t)⊤ V W (t) V ν ∂t = 1 n i:µ∈X (i) g (t) i y i L l=1 m1 j=1 a j ν ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←µ + 1 n i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ V w (t) j p (t,i) q←l,k←ν = G (t) (ν) 1 n i:µ∈X (i) g (t) i y i L l=1 p (t,i) q←l,k←µ + G (t) (µ) 1 n i:ν∈X (i) g (t) i y i L l=1 p (t,i) q←l,k←ν .Thus, we obtainµW (τ )⊤ V W (τ ) V ν -µW (T1)⊤ V W (T1) V ν ≤ τ T1 |G (t) (ν)| 1 n i:µ∈X (i) g (t) i y i L l=1 p (t,i) q←l,k←µ + |G (t) (µ)| 1 n i:ν∈X (i) g (t) i y i L l=1 p (t,i)q←l,k←ν dt.By Definition F.1, for t ∈ [T 1 , T 2 ], we have |G (t) (µ)| ≤ O(log m) and L l=1 p (t,i)q←l,k←µ ≤ O(1). Further, by the property of the cross-entropy loss, we have1 n i:µ∈X (i) g (t) i y i = O( L (t) ).Therefore, by Theorem F.8, we obtainµW (τ )⊤ V W (τ ) V ν -µW (T1)⊤ V W (T1) V ν ≤ O(log m) τ T1 L (t) dt ≤ O(log m) • O log τ σ 2 1 mm 1 .Corollary F.12. For t ∈ [T 1 , T 2 ], we havemax µ∈{µi} d i=1 m1 j=1 a j ∂w (t) j ∂t W (t) V µ m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 ≤ o(1/L) 1 n i∈[n] g (t)i .Proof. By Lemma E.9 and and Lemma F.11, for all µ̸ = ν ∈ {µ i } d i=1 , we have |µ ⊤ W (t)⊤ V W (t) V ν| ≤ O(σ 2 0 √ m + 1/m) and ∥W (t)V ν∥ 2 2 = O(σ 2 0 m + 1/m). Thus, by Lemma F.11, we havemax µ,ν µ ⊤ W (t)⊤ V W (t) V ν m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 ≤ O σ 2 0 √ m + 1/m σ 2 1 mm 1 , max ν ∥W (t) V ν∥ 2 2 m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 ≤ O σ 2 0 m + 1/m σ 2 1 mm 1 . Recall that m1 j=1 a j ∂w (t) j ∂t W (t) V µ = m 1 n n i1=1 g (t) i1 y i1 L l1=1 p (t,i1)⊤ l1 V (t,i1)⊤ W (t) V µ.This implies thatEQUATIONi .Corollary F.13 (Complete version of Corollary 4.6). For t ∈ [T 1 , T 2 ], we have∂ ∂t µ ⊤ 2 W (t)⊤ V W (t) V µ 1 > 0, ∂ ∂t µ ⊤ 1 W (t)⊤ V W (t) V µ 1 > 0, ∂ ∂t µ ⊤ 2 W (t)⊤ V W (t) V µ 2 > 0 ∂ ∂t µ ⊤ 1 W (t)⊤ V W (t) V µ 3 < 0, ∂ ∂t µ ⊤ 2 W (t)⊤ V W (t) V µ 3 < 0.Proof. This is a direct consequence of Lemma F.11, Lemma F.6 and Definition F.1.F.5 Change of Random-Token Sub-Network Lemma F.14.For t ∈ [T 1 , T 2 ] and µ ∈ {µ i } d i=4 , we have ∂G (t) (µ) ∂t ≤ O 1 n L (t) σ 2 1 mm 1 + O( L (t) m 1 (L(σ 2 0 √ m + 1/m) + σ 2 0 m)).Thus,G (t) (µ) -G (T1) (µ) ≤ O 1 n + m 1 (L(σ 2 0 √ m + 1/m) + σ 2 0 m) σ 2 1 mm 1 .Proof. By Lemma C.1 and Definition F.1, we have∂G (t) (µ) ∂t = m1 j=1 a j w (t) j ∂W (t) V µ ∂t + a j ∂w (t) j ∂t W (t) V µ ≤ 1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ + 1 n n i1=1 g (t) i1 y i1 m 1 L l1=1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ) .By Definition F.1 and Corollary F.2,1 n i2: µ∈X (i 2 ) g (t) i2 y i2 L l2=1 m1 j1=1 m1 j2=1 a j1 a j2 w (t) j1 , w (t) j2 p (t,i2) q←l2,k←µ ≤ O 1 n L (t) σ 2 1 mm 1 .By Lemma E.9 and and Lemma F.11, we have|µ ⊤ W (t)⊤ V W (t) V ν| ≤ O(σ 2 0 √ m + 1/m) for µ ̸ = ν and ∥W (t) V µ∥ 2 2 = O(σ 2 0 m). Thus, 1 n n i1=1 g (t) i1 y i1 m 1 L l1=1 L l2=1 p (t,i1) l1,l2 v (t,i1) l2 , v (t) (µ) ≤ O( L (t) m 1 (L(σ 2 0 √ m + 1/m) + σ 2 0 m)).Thus, by Theorem F.8, for t ∈ [T 1 , T 2 ], we haveG (t) (µ) -G (T1) (µ) = t T1 ∂G (τ ) (µ) ∂τ dτ ≤ O σ 2 1 mm 1 n + O(m 1 L(σ 2 0 √ m + 1/m)) t T1 L (τ ) dτ ≤ O 1 n + m 1 (L(σ 2 0 √ m + 1/m) + σ 2 0 m) σ 2 1 mm 1 . Corollary F.15. For t ∈ [T 1 , T 2 ], we have L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 )p (t,i) q←l1,k←l2 ≤ O(1).Proof. By Lemma F.3, we haveL l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (T1) (X (i) l2 )p (T1,i) q←l1,k←l2 ≤ O(1) By Lemma F.14, for t ∈ [T 1 , T 2 ], G (t) (µ) = O(1) for µ ∈ {µ i } d i=4 .On the other hand, by the triangle inequality, we haveL l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (T1) (X (i) l2 )p (T1,i) q←l1,k←l2 - L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 )p (t,i) q←l1,k←l2 ≤ L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 (G (T1) (X (i) l2 ) -G (t) (X (i) l2 )) p (t,i) q←l1,k←l2 + L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 ) p (T1,i) q←l1,k←l2 -p (t,i) q←l1,k←l2 + L l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (T1) (X (i) l2 ) -G (t) (X (i) l2 ) p (T1,i) q←l1,k←l2 -p (t,i) q←l1,k←l2 ≤ O(1),where the last inequality applies Definition F.1. This implies that, for t ∈ [T 1 , T 2 ], by Lemma F.14, we haveL l1=1 l2: X (i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 )p (t,i) q←l1,k←l2 ≤ O(1).F.6 Change of Score and Softmax Probability Lemma F.16 (Change of score, complete version of Lemma 4.7). For t ∈ [T 1 , T 2 ], the attention scores are changing in the following way:• For µ, ν ∈ {µ 1 , µ 2 }, µ ̸ = ν, the query-key-correlation score between the two target signals increases, while the query-key-correlation score between one target signal and the common token decreases, i.e.,∂ ∂t ν ⊤ W (t)⊤ K W (t) Q µ = 1 √ m Θ( L (t) σ 2 0 m) 1 L , ∂ ∂t µ ⊤ 3 W (t)⊤ K W (t) Q µ = - 1 √ m Θ( L (t) σ 2 0 m) 1 L .• The change of score satisfies:max µ,ν∈{µi} 3 i=1 ν ⊤ W (t) K W (t) Q µ -ν ⊤ W (T1) K W (T1) Q µ = Θ σ 2 0 m √ mLσ 2 1 mm 1 + σ 2 0 √ m √ mσ 2 1 mm 1 . • For all µ ∈ {µ 1 , µ 2 , µ 3 }, γ ∈ {µ i } d i=4, the query-key-correlation score changes as follows:EQUATIONµ ⊤ W (t)⊤ K W (t) Q γ -µ ⊤ W (T1)⊤ K W (T1) Q γ ≤ O σ 2 0 m n √ mLσ 2 1 mm 1 + σ 2 0 √ m √ mσ 2 1 mm 1 , γ ⊤ W (t)⊤ K W (t) Q µ -γ ⊤ W (T1)⊤ K W (T1) Q µ ≤ O σ 2 0 m n √ mLσ 2 1 mm 1 + σ 2 0 √ m √ mσ 2 1 mm 1 . Proof. By Lemma C.4, we have ∂ν ⊤ W (t)⊤ K W (t) Q µ ∂t = 1 n √ m i:µ,ν∈X (i) g (t) i y i m1 j=1 a j ∥k (t) (ν)∥ 2 2 v (t)⊤ (ν)w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) q←µ,k←ν + 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j L l=1 ν ⊤ W (t)⊤ K K (t,i) l V (t,i)⊤ l w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) q←µ,k←l I(K (t,i) l ̸ = k (t) (ν)) + 1 n √ m i:ν,µ∈X (i) g (t) i y i m1 j=1 a j ∥q (t) (µ)∥ 2 2 p (t,i) q←µ,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i L l=1 m1 j=1 a j µ ⊤ W (t)⊤ Q q (t,i) l p (t,i) q←l,k←ν w (t)⊤ j v (t,i) (ν) -w (t)⊤ j V (t,i) p (t,i) l I(q (t,i) l ̸ = q (t) (µ)).Now, we take µ = µ 1 , ν = µ 2 . By Theorem E.27, we have G (T1) (µ 2 ) ≥ Ω(1). A consequence of Definition F.1 and Lemma F.6 is that ∂G(t) (µ2) ∂t > 0 for t ∈ [T 1 , T 2 ]. Thus, we have G (t) (µ 2 ) ≥ Ω(1).Further by Theorem E.27, we have G (T1) (µ) = O(σ 0 σ 1 √ mm 1 ) for µ ∈ R and then by Lemma F.14, we haveG (t) (µ) = O(σ 0 σ 1 √ mm 1 ) + O 1 n + m1(L(σ 2 0 √ m+1/m)+σ 2 0 m) σ 2 1 mm1 for t ∈ [T 1 , T 2 ]. Also, Definition F.1 implies that G (t) (µ 2 ) - m1 j=1 w (t) j V (t,i) p (t,i) l(i,µ) ≥ Ω(1). Now, this yields ∂µ ⊤ 2 W (t)⊤ K W (t) Q µ 1 ∂t = 1 √ m Θ( L (t) σ 2 0 m) 1 L + O 1 √ m L (t) σ 2 0 √ m .On the other hand, by the analysis similar to the above, we obtain∂µ ⊤ 3 W (t)⊤ K W (t) Q µ 1 ∂t = - 1 √ m Θ( L (t) σ 2 0 m) 1 L + O 1 √ m L (t) σ 2 0 √ m .Next, to prove the maximum change of the score, we havemax µ,ν ∂ν ⊤ W (t) K W (t) Q µ ∂t ≤ 1 √ m Θ( L (t) σ 2 0 m) 1 L + O 1 √ m L (t) σ 2 0 √ m .By Theorem F.8, we haveν ⊤ W (t)⊤ K W (t) Q µ -ν ⊤ W (T1)⊤ K W (T1) Q µ ≤ t T1 1 √ m Θ( L (τ ) σ 2 0 m) 1 L + O 1 √ m L (τ ) σ 2 0 √ m dτ ≤ O σ 2 0 m √ mLσ 2 1 mm 1 + σ 2 0 √ m √ mσ 2 1 mm 1 . Finally, for γ ∈ {µ i } d i=4 , µ ∈ {µ i } d i=1 , we have ∂ ∂t µ ⊤ W (t)⊤ K W (t) Q γ = 1 n √ m Θ( L (t) σ 2 0 m) 1 L + O 1 √ m L (t) σ 2 0 √ m , which implies that µ ⊤ W (t)⊤ K W (t) Q γ -µ ⊤ W (T1)⊤ K W (T1) Q γ ≤ O σ 2 0 m n √ mLσ 2 1 mm 1 + σ 2 0 √ m √ mσ 2 1 mm 1 .Corollary F.17 (Change of softmax). For t ∈ [T 1 , T 2 ], the softmax probability is changing in the following way:• For µ, ν ∈ {µ 1 , µ 2 }, µ ̸ = ν, the softmax probability between the two target signals increases, whereas the softmax probability between one target signal and the common token decreases, i.e., ∂ ∂t p(t,i) q←µ,k←ν = 1 √ m Θ( L (t) σ 2 0 m) 1 L 2 , ∂ ∂t p (t,i) q←µ,k←µ3 = - 1 √ m Θ( L (t) σ 2 0 m) 1 L 2 .• For all µ ∈ {µ 1 , µ 2 }, γ ∈ {µ i } d i=4 , the softmax probability between one target signal and a random token changes as follows:∂ ∂t p (t,i) q←µ,k←γ ≤ 1 n √ m Θ( L (t) σ 2 0 m) 1 L 2 + O 1 L √ m L (t) σ 2 0 √ m .Furthermore, we havemax i∈[n],l1,l2∈[L] p (t,i) l1,l2 -p (0,i) l1,l2 ≤ O(1/L 2 ).Proof. Take µ = µ 1 , ν = µ 2 . First of all, by Definition F.1 and Lemma F.16, we havep (t,i)⊤ q←µ1 ∂X (i)⊤ W (t)⊤ K W (t) Q µ 1 ∂t ≤ 1 √ m O( L (t) σ 2 0 m) 1 L 2 + 1 n √ m Θ( L (t) σ 2 0 m) 1 L + O 1 √ m L (t) σ 2 0 √ m .Thus, by Lemma G.1 and Lemma F.16, for i ∈ I 1 , we obtain∂ ∂t p (t,i) q←µ1,k←µ2 = 1 m Θ( L (t) σ 2 0 m) 1 L 2 , ∂ ∂t p (t,i) q←µ1,k←µ3 = - 1 m Θ( L (t) σ 2 0 m) 1 L 2 .On the other hand, for γ ∈ {µ i } d i=4 , we have ∂ ∂t p(t,i) q←µ,k←γ ≤ 1 nm Θ( L (t) σ 2 0 m) 1 L 2 + O 1 Lm L (t) σ 2 0 √ m .Finally, by Lemma F.16, we havemax µ,ν ν ⊤ W (t) K W (t) Q µ -ν ⊤ W (T1) K W (T1) Q µ ≤ O σ 2 0 m √ mLσ 2 1 mm 1 + σ 2 0 √ m √ mσ 2 1 mm 1 .Thus, by Lemma G.2, we obtainmax i,l1,l2 p (t,i) l1,l2 -p (T1,i) l1,l2 ≤ O σ 2 0 m mL 2 σ 2 1 mm 1 + σ 2 0 √ m mσ 2 1 mm 1 L + L σ 2 0 m mLσ 2 1 mm 1 + σ 2 0 √ m mσ 2 1 mm 1 2 = O(1/L 2 ).Corollary F.18. For t ∈ [T 1 , T 2 ], we havemax i∈[n] L l1=1 L l2=1 G (t) (X (i) l2 ) ∂p (t,i) l 1 ,l 2 ∂t m1 j1=1 m1 j2=1 a j1 w (t) j1 , a j2 w (t) j2 ≤ o(1) 1 n i∈[n] g (t) i .Proof. This is a direct consequence of Definition F.1, Lemma F.10 and Corollary F.17.F.7 Change of Self-Correlation of Key/Query-Transformed data Lemma F.19 (Change of self-correlation). For µ, ν ∈ {µ i } d i=1 , we haveν ⊤ W (t)⊤ Q W (t) Q µ -ν ⊤ W (T1)⊤ Q W (T1) Q µ ≤ O σ 2 0 √ m √ mσ 2 1 mm 1 , ν ⊤ W (t)⊤ K W (t) K µ -ν ⊤ W (T1)⊤ K W (T1) K µ ≤ O σ 2 0 √ m √ mσ 2 1 mm 1 . Proof. We prove the result for ν ⊤ W (t)⊤ Q W (t) Q µ and the proof for ν ⊤ W (t)⊤ K W (t) K µ is similar. By Lemma C.4, we have ∂ν ⊤ W (t)⊤ Q W (t) Q µ ∂t = 1 n √ m i:µ∈X (i) g (t) i y i m1 j=1 a j ν ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,µ) p (t,i) l(i,µ) + 1 n √ m i:ν∈X (i) g (t) i y i m1 j=1 a j µ ⊤ W (t)⊤ Q K (t,i) diag V (t,i)⊤ w (t) j -w (t)⊤ j V (t,i) p (t,i) l(i,ν) p (t,i) l(i,ν) . A simple result from Lemma F.16 is that |ν ⊤ W (t)⊤ K W (t) Q µ| ≤ O(σ 2 0 √ m) for t ∈ [T 1 , T 2 ] and µ, ν ∈ {µ i } d i=1 . Therefore, by Definition F.1, ∂ν ⊤ W (t)⊤ Q W (t) Q µ ∂t ≤ O 1 √ m L (t) σ 2 0 √ m log m , which implies ν ⊤ W (t)⊤ Q W (t) Q µ -ν ⊤ W (T1)⊤ Q W (T1) Q µ ≤ t T1 O 1 √ m L (t) σ 2 0 √ m log m dτ ≤ O σ 2 0 √ m √ mσ 2 1 mm 1 ,where the inequality follows from Theorem F.8 and Definition F.1.Theorem F.20. Define T ⋆ = min t {t :L (t) = Θ(1/poly(m))}. Then, T ⋆ ∈ [T 1 , T 2 ].Proof. The following results altogether show that Phase 2 can last for at least Θ(poly(m)) time:• Lemma F.19 proves that the change of K, Q self-correlation as follows:max µ,ν µ ⊤ W (t)⊤ K W (t) K ν -µ ⊤ W (0)⊤ K W (0) K ν = O(σ 2 0 √ m), max µ,ν µ ⊤ W (t)⊤ Q W (t) Q ν -µ ⊤ W (0)⊤ Q W (0) Q ν = O(σ 2 0 √ m).• Corollary F.9 proves that ≤ O(σ 2 1 mm 1 ).• Corollary F.12 proves thatEQUATIONi .• Corollary F.15 proves that L l1=1 l2: X(i) l 2 ∈{µ k } d k=4 G (t) (X (i) l2 )p (t,i) q←l1,k←l2 ≤ O(1).• Corollary F.17 proves thatmax i∈[n],l1,l2∈[L] p (t,i) l1,l2 -p (0,i) l1,l2 ≤ O(1/L 2 ).• Corollary F.18 proves thatEQUATIONi .• Lemma F.5 implies that1 2 ≤ i∈I2 g (t) i i∈I3 g (t) i ≤ 2.• Lemma F.6 implies that the gradients i∈I g (t)i for I ∈ {I 1 , I 2 , I 3 , I 4 } satisfies i∈I4 g i ≤ min i∈I2 g (t) i , i∈I3 g (t) i ≤ max i∈I2 g (t) i , i∈I3 g (t) i ≤ i∈I1 g (t) i ≤ i∈I2∪I3∪I4 g (t)i .• Lemma F.6 and Corollary F.17 proves that y i F(t) i ≥ y i F (T1) i ≥ C.The above shows that all the requirements needed to satisfy the definition of Phase 2 (Definition F.1) can hold for at least Ω(poly(m)) time. Thus, T 2 -T 1 ≥ Ω(poly(m)). Further, by Theorem F.8, we have that L (t) ≤ O(1/poly(m)) implies t = Θ(poly(m)).F.9 Proof of Theorem 3.2Proof. This is proved in Corollary F.13 and Lemma F.16.Proof. This is proved as a direct consequence of Lemma F.6 and Theorem F.20.For the generalization loss, since the training loss satisfies L (T ⋆ ) ≤ 1/poly(m), for each class I ∈ {I 1 , I 2 , I 3 , I 4 } there exists a sample X ⋆ I such that ℓ(X ⋆ I ) ≤ 1/poly(m). Note that by Definition F.1 the random tokens only contributes to O(1) in F (T ⋆ ) (X). Thus, given a fixed new sample X ∼ D, we have |F (T ⋆ ) (X) -F (T ⋆ ) (X ⋆ I )| ≤ O(1) which implies ℓ(yF (T ⋆ ) (X)) ≤ 1/poly(m). Since this holds for all X ∼ D, we have L (T ⋆ ) ≤ 1/poly(m).The gradient of p(x) i = Softmax(x) i for x ∈ R n : • 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 a 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 a 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 to 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 to replicate the results, access to 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 NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Question: Does the paper provide open access to the data and code, with sufficient instructions 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 NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 not include experiments.• The experimental setting should be presented in the core of the paper to a 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, or as supplemental material. 7. • Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.Experiment12. Licenses for existing assets Question: Are the creators or 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?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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are 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 creators. 
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 was 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. 
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 as details about compensation (if any)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we 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 as the institution conducting the review.mm1 if one of i 1 , i∈ [d] \ [3]. Theorem3.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 C (T ⋆ ) 1,2 and C (T ⋆ ) 2,1 capture strong coupling of the two target signals µ 1 and µ 2 , and the large negative coefficients C (T ⋆ ) 3,1 and C (T ⋆ ) 3,2 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: - In line 151 and 152, it is confusing why concentration theorems lead to the specific probability in (i) of Assumption 2.3.
- Lack of explanation for $\langle w_{j_1}^{(t)},w_{j_2}^{(t)} \rangle$ in line 194.
- It is not obvious why "the samples with only one target signal may be classified incorrectly as co-occurence" in line 282.
- The notation in line 169 is somewhat misleading.
limitations: 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.
paper_summary: paper_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).

Overall, 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.
strength_weakness: strength_weakness: - This word cooccurrence task is very simple, and thus it is not surprising that a single transformer layer can easily learn it.
- Only full gradients are considered, whereas transformers are typically trained with mini-batch Adam(W).
questions: questions: - What other tasks (beyond word cooccurrence) do you think could be analyzed with this methodology?
- What are the implications of this result to more complex/realistic tasks, like next token prediction?
- If mini-batch Adam is used during training, do the two phases still occur?
- Can you add more details about the experimental setup to the main paper?
- 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: limitations: Yes, limitations have been discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.  

The 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.
questions: 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?

2. 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?

3. Regarding the proof sketch:

    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?

    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: limitations: Yes

Comments: title: A gentle reminder
Thank you
Please respond to rebuttal
A gentle reminder
A gentle reminder
A gentle reminder
comment: Thank you very much for your feedback. 

As we mentioned in our rebuttal, we expect the training phases of Adam will exhibit some differences from gradient descent considered in our work. 
This is due to the coordinate-wise normalization operation uniquely occurring in Adam but not in GD.
Our experiments confirmed our thoughts. 
We provide the experiment results on AdamW with default parameter setup below, where we show the changes in $G(\mu_1)$, which is the MLP alignment with the signal $\mu_1$, and the attention score when the query is $\mu_1$ and the key is $\mu_2$. 
All the values are rounded to a tenth of decimal. 
As you can see from the experiment results below, both the MLP and score change dramatically in the first 50 steps and then the changes slow down.
The behavior thus is very different from gradient descent, where only MLP changes rapidly in the initial training and then both MLP and score jointly change substantially.

| Step | 0 | 50 | 100 | 150 | 200 | 250 |
| --- | --- | --- | --- | --- | --- | --- | 
| $G(\mu_1)$ | 0.2 | 50.8 | 56.8 | 59.7 | 62.0 | 64.0 |
| Score $(\times 10^{-1})$ | -0.8 | 3.0 | 3.5 | 3.8 | 4.1 | 4.3 |

We further emphasize that the main contribution of this paper lies in the development of the **new techniques** for analyzing the training dynamics of transformers, especially the joint optimization analysis of all transformer parameters, as Reviewer cFaT pointed out. We expect that these mathematical techniques can be generalized to more complicated transformer architectures in the future.

We thank the reviewer for the time and efforts.
We hope that our response answers your question.
Thank you very much for your response. I keep my score unchanged, due to what I perceive to be the limited impact/scope of the work.

Regarding my question: "If mini-batch Adam is used during training, do the two phases still occur?" ---> Could you add an experiment to check this?
Dear Reviewer G3j5,

It occurs to us that we uploaded our response earlier and OpenReview didn't send out notification. 
This message serves as a gentle reminder for the reviewer to look at our most recent response. 

We once again appreciate the reviewer's time and effort. 
If our response resolved your concerns, could you please kindly consider to increase your current score accordingly?
We are happy to hear your response. 

Best,

Authors
We thank the reviewer very much for the highly inspiring question. 

To clarify the reviewer's concern, we note that the attention quantities that we study in one-layer case are the projections of attention to the input signals. In multi-layer transformers, the inputs to the middle layers are changing from iterations to iterations, and hence it is hard to find a fixed subspace to project on. Consequently, it is hard to find a natural and meaningful projection of the softmax attention in the middle layers to illustrate the phase change. Similarly, the inputs to the middle MLP layers are also changing, which makes it hard to find a fixed space for projection.

This said, we still expect that the phase-wise learning (first MLP then attention) can occur for broader class of architectures. One feasible setting to demonstrate this is one-layer **multi-headed** transformer, which generalizes our original architecture and still maintains the meaningful projection of softmax attention to the fixed input space. We conducted additional experiments for one-layer 3-headed transformer. In our table below, we provide the values $G(\mu_1)$ of the MLP alignment with $\mu_1$ and the attention scores when the query is $\mu_1$ and key is $\mu_2$. The experiment results are all rounded to a tenth of decimal for better illustration. Take the first head as an example. The result clearly indicates a two-phase training process: (i) in phase 1 (during the first 250 steps), the MLP layer $G(\mu_1)$ grows quickly whereas the change in the attention score is very small compared to its later changes; and (ii) in phase 2 (after 250 steps), both MLP and attention score change substantially.  

| | Step | 0 | 250 | 500 | 750 | 1000 | 1250 | 1500 | 1750 | 
|---|---|---|---|---|---|---|---|---|---|
| Head No.1 | score ($\times 10^{-2}$) |7.7 | 7.8 | 8.2 | 8.6 | 9.0 | 9.3 | 9.5 | 9.8 |
| | $G(\mu_1)$ | 0.8 | 4.8 | 8.0 | 10.7 | 13.0 | 14.8 |16.4 | 17.6 |
| Head No.2 | score ($\times 10^{-2}$) |2.1 | 2.2 | 2.6 | 3.0 | 3.3 | 3.6 | 3.9 | 4.1 |
| | $G(\mu_1)$ | -0.1 | 2.5 | 4.6 | 6.6 | 8.3 | 9.7 |10.9 | 11.9 |
| Head No.3 | score ($\times 10^{-2}$) |-12.4 | -12.2 | -11.7 | -11.1 | -10.6 | -10.2 | -9.9 | -9.6 |
| | $G(\mu_1)$ | 0.7 | 4.6 | 7.6 | 10.2 | 12.4 | 14.2 | 15.6 | 16.8 |

We further emphasize that the main contribution of this paper lies in the development of the **new techniques** for analyzing the training dynamics of transformers, especially the joint optimization analysis of all transformer parameters, as Reviewer cFaT pointed out. We expect that these mathematical techniques can be generalized to more complicated transformer architectures in the future.

We thank the reviewer again for your time and efforts. If our response resolved your concerns, could you please kindly consider to increase your score. We are also more than happy to answer your further questions.
Thank you for your clarifications. Specifically, I have noticed the new experiment. However, these results remain preliminary and do not verify that the training dynamics of weights in this work can be generalized to broader cases. For example, the dynamics "first learning MLP and then learning Attention" cannot be adequately captured by the loss or accuracy presented.

I will maintain my score.
We thank the reviewer very much for your further feedback. We will follow the reviewer's suggestion and enhance the presentation of the proof outline in the revision. We also highly appreciate that you increased the rating of the paper.
Dear reviewers,

The author-reviewer discussion period will end soon. Please make sure you read the authors' rebuttal and respond to it. If you have additional questions after reading the rebuttal please discuss with the authors.  For those who have done so, thank you!


AC
Thanks for your response.

Given that the data setup is simple, the findings on the efficient learnability of this task are not particularly surprising, and might not provide much insight on practical tasks. 

However, the minimal assumptions on the model weights and the joint optimization analysis of all transformer parameters could make the techniques valuable for future transformer analyses. Thus, I will increase my score to weak accept. Since the main contribution of the paper lies in the proof technique, I suggest enhancing the presentation of the proof outline in the final manuscript.
Dear Reviewer cFaT,

We've taken your initial feedback into careful consideration in our response. Could you please check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards, 
Authors
Dear Reviewer ritc,

We've taken your initial feedback into careful consideration in our response. Could you please check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards, 
Authors
Dear Reviewer G3j5,

We've taken your initial feedback into careful consideration in our response. Could you please check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards, 
Authors