Title: Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformers

gradient flow dynamics

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

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

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

1 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.

1.1 Related Work

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.

2.1 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:

Uniformly randomly select an index i 3 ∈ [ L ] and set x i 3 = µ 3 .

Then, one of the following cases occurs:

With probability 1 / 2 , set y = 1 and uniformly randomly select two indices i 1 = i 2 ∈ [ L ] \ { i 3 } and set x i 1 = µ 1 , x i 2 = µ 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 i 1 = µ 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 i 2 = µ 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) | I 1 | n = 1 2 and | I 2 | n = | I 3 | n = | I 4 | n = 1 6 ; and (ii) for all i 1 , i 2 ∈ [ n ] , l 1 , l 2 ∈ [ L ] , if X ( i 1 ) l 1 , X ( i 2 ) l 2 / ∈ { µ 1 , µ 2 , µ 3 } , then X ( i 1 ) l 1 = X ( i 2 ) l 2 , 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.

2.2 Transformer Architecture and Training

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., X ( i ) = [ x ( i ) 1 , x ( i ) 2 , . . . , x ( i ) L ] . We consider the transformer model with a self-attention module followed by a linear MLP:

where 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 m 1 × m (with w ⊤ j being the j -th row of W ), and the output-layer weights of the MLP a ∈ R m 1 . We define the linear MLP function of the transformer to be G ( µ ) = ∑ m 1 j =1 a j w ⊤ j W V µ . We now introduce some shorthand notations K = 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 √ m and the attention vector p l := Softmax ( s l ) . For the attention vector, if µ, ν ∈ X ( i ) , let l ( i, µ ) , l ( i, ν ) be the indices such that X ( i ) l ( i,µ ) = µ, X ( i ) l ( 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. The trainable parameters are [ 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:

Similarly, 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 ˜ Θ( √ m 1 ) .

The number of training samples and tokens satisfy n ≥ ˜ Ω( L 2 ) and L ≥ ˜ Ω(1) .

The failure probability satisfies 1 /δ ≤ poly ( m ) .

3 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 ⊤ WW 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 ) j 1 , w ( t ) j 2 〉 , ν ⊤ 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 time T 1 = ˜ O (1 /m ) such that

The linear MLP functions satisfy: G ( T 1 ) ( µ 1 ) ≥ Ω(1) , G ( T 1 ) ( µ 2 ) ≥ Ω(1) , G ( T 1 ) ( µ 3 ) ≤ -Ω(1) .

All training samples are correctly classified: y i F ( T 1 ) i > 0 for all i ∈ [ n ] .

For t ∈ [0 , T 1 ] , all dynamical variables 〈 w ( t ) j 1 , w ( t ) j 2 〉 , ν ⊤ 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 ( T 1 ) = Θ(1) .

In Theorem 3.1, item 1 implies that in a short time, the linear MLP function G ( T 1 ) ( · ) 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:

Theorem 3.3 indicates that both training and test losses converge nearly to zero as long as the embedding dimension m is sufficiently large, because both the attention and MLP matrices are trained towards enlarging the classification margin in Phase 2. Theorem 3.3 also provides the explicit form of the attention matrix in Equation (3), in which the second term captures the learned information of the self-attention module. It can be seen that the large coefficients 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 encourages 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.

Synthetic Experiment: We next verify our theory and the two-phase characterization of the training process via synthetic experiments (see the experiment setup in Appendix A).

GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=20,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=20,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=19,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>GLYPH<c=17,font=/AAAAAB+DejaVuSans>GLYPH<c=18,font=/AAAAAB+DejaVuSans>

GLYPH<c=16,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=4,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

Phase 1

Phase 2

(a) Attention score correlation

GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=15,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=10,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=11,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=14,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>

GLYPH<c=17,font=/AAAAAB+DejaVuSans>

GLYPH<c=25,font=/AAAAAB+DejaVuSans>GLYPH<c=19,font=/AAAAAB+DejaVuSans>GLYPH<c=21,font=/AAAAAB+DejaVuSans>GLYPH<c=23,font=/AAAAAB+DejaVuSans>GLYPH<c=21,font=/AAAAAB+DejaVuSans>GLYPH<c=23,font=/AAAAAB+DejaVuSans>GLYPH<c=20,font=/AAAAAB+DejaVuSans>GLYPH<c=3,font=/AAAAAB+DejaVuSans>GLYPH<c=16,font=/AAAAAB+DejaVuSans>

GLYPH<c=24,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>

GLYPH<c=22,font=/AAAAAB+DejaVuSans>GLYPH<c=24,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>

GLYPH<c=22,font=/AAAAAB+DejaVuSans>GLYPH<c=24,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=18,font=/AAAAAB+DejaVuSans>GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

GLYPH<c=22,font=/AAAAAB+DejaVuSans>GLYPH<c=24,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=18,font=/AAAAAB+DejaVuSans>GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

GLYPH<c=22,font=/AAAAAB+DejaVuSans>GLYPH<c=24,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=18,font=/AAAAAB+DejaVuSans>GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

GLYPH<c=22,font=/AAAAAB+DejaVuSans>GLYPH<c=24,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=18,font=/AAAAAB+DejaVuSans>GLYPH<c=4,font=/AAAAAB+DejaVuSans>

GLYPH<c=5,font=/AAAAAC+DejaVuSans-Oblique>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>

Phase 1

Phase 2

Figure 1 (a) shows how the attention score correlation µ ⊤ i W ( t ) ⊤ K W ( t ) Q µ j evolves during the training. It is clear that these scores do not change significantly in Phase 1, verifying Theorem 3.1. In Phase 2, the score correlation between two target signals µ 1 and µ 2 increases, and the score between one target signal and the common token decreases, verifying Theorem 3.2.

Figure 1 (b) plots how the training loss changes during the two phases of training. The blue curve (indexed by 'loss') represents the overall training loss of all samples. The other curves correspond to the training loss of four types of samples as indicated in the legend. In Phase 1, the training loss for samples with both target signals (i.e., orange curve) decreases because the linear MLP layer aligns with the target signals (verifying Lemma 4.1 in Section 4.1). The training loss for samples with one target signal and the common token (i.e., green or red curves) first increases because the linear MLP layer initially has not aligned negatively enough with the common token yet (as captured by Lemma 4.2 in Section 4.1), and then decreases in the later stage of Phase 1 when the MLP layer aligns negatively with the common token (as captured by Lemma 4.3 in Section 4.1). All loss functions decrease in Phase 2 because all attention matrices and linear MLP jointly enlarge the classification margin, verifying Theorem 3.2.

4 Proof Outline: Two-phase Gradient Flow Analysis

4.1 Phase 1: Alignment of Linear MLP for Correct Classification

In Phase 1, the linear MLP quickly aligns with the two target word tokens while all attention matrices stay roughly unchanged from their initialization values. We show that the linear MLP functions | G ( t ) ( µ 1 ) | , | G ( t ) ( µ 2 ) | , | G ( t ) ( µ 3 ) | become sufficiently large (larger than some constant threshold) so that all training samples are correctly classified at the end of Phase 1.

We first analyze the dynamical system at the initialization. In particular, the following lemma shows that at the initialization, the linear 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,

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 (b), where the loss function corresponding to only one target signal and the common token first increases in Phase 1.

Lemma 4.2 (Same as Lemma E.16) . Let F = max i | F (0) i | . With probability at least 1 -δ over the weight initialization,

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 ]

Proof Intuition of Lemma 4.3. We first note that the term

makes the major contribution to the gradient ∂G ( t ) ( µ 3 ) ∂t . 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, ∂a j 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.

4.2 Phase 2: Evolution of Attention and MLP for Large Classification Margin

In Phase 2, both attention and MLP matrices evolve towards enlarging the classification margin, thus driving the loss value small.

We now analyze what happens in Phase 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:

It has been proved that at the end of Phase 1, for µ ∈ { µ 1 , µ 2 , µ 3 } , we have ∣ ∣ ∣ ∣ ∑ j 1 ∂w ( t ) ⊤ j 1 ∂t W ( t ) V µ ∣ ∣ ∣ ∣ ≪

∣ ∣ ∣ ∣ ∑ j 1 w ( t ) ⊤ j 1 ∂W ( t ) V ∂t µ ∣ ∣ ∣ ∣ since the magnitude of ∑ m 1 j 1 =1 ∑ m 1 j 2 =1 〈 a j 1 w ( t ) j 1 , a j 2 w ( t ) j 2 〉 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 calculate

Thus, if y i F ( t ) i > 0 for all i ∈ [ n ] , then ∑ m 1 j 1 =1 ∑ m 1 j 2 =1 〈 a j 1 w ( t ) j 1 , a j 2 w ( t ) j 2 〉 is always increasing. Thus, ∂G ( t ) ( µ ) ∂t mainly depends on the behavior of 1 n ∑ i 2 : µ ∈ X ( i 2 ) g ( t ) i 2 y i 2 ∑ L l 2 =1 p ( t,i 2 ) q ← l 2 ,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 ∑ i 2 : µ ∈ X ( i 2 ) g ( t ) i 2 y i 2 ∑ L l 2 =1 p ( t,i 2 ) q ← l 2 ,k ← µ ≈ 1 n ∑ i 2 : µ ∈ X ( i 2 ) g ( t ) i 2 y i 2 if p ( t,i 2 ) q ← l 2 ,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 of 1 n ∑ i 2 : µ ∈ X ( i 2 ) g ( t ) i 2 y i 2 . 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 ∈ I 2 g ( t ) i and ∑ i ∈ I 3 g ( t ) i close to each other during training; and (ii) Lemma 4.5, which shows that the gradient gaps ∑ i ∈ I 1 g ( t ) i -∑ i ∈ I 2 g ( t ) i and ∑ i ∈ I 2 ∪ I 3 ∪ I 4 g ( t ) i -∑ i ∈ I 1 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 that

Proof Intuition of Lemma 4.4. The intuition behind the result is as follows. If ∑ i ∈ I 2 g ( t ) i becomes much bigger than ∑ i ∈ I 3 g ( t ) i during the training, then ∑ i ∈ I 1 g ( t ) i -∑ i ∈ I 2 g ( t ) i is much smaller than ∑ i ∈ I 1 g ( t ) i -∑ i ∈ I 3 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 ∈ I 2 g ( t ) i will drop faster than ∑ i ∈ I 3 g ( t ) i , i.e., -∂ ∂t ∑ i ∈ I 2 g ( t ) i > -∂ ∂t ∑ i ∈ I 3 g ( t ) i . In Appendix, we formally prove Lemma 4.4 by analyzing the ratio ∑ i ∈ I 2 g ( t ) i / ∑ i ∈ I 3 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 that

Further, for some constant C , we have

Proof Sketch of Lemma 4.5. The proof of Lemma 4.5 relies on analyzing how the ratio between ∂G ( t ) ( µ 1 ) ∂t and -∂G ( t ) ( µ 3 ) ∂t changes. We show that this ratio will hang over around some range. Recall the relationship that ∂G ( t ) ( µ ) ∂t ≈ 1 n ∑ i 2 : µ ∈ X ( i 2 ) g ( t ) i 2 y i 2 · ∑ m 1 j 1 =1 ∑ m 1 j 2 =1 a j 1 a j 2 〈 w ( t ) j 1 , w ( t ) j 2 〉 . It is not hard to show that

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 ] ,

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

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 ] ,

Lemma 4.7 is keeping track of the attention coefficients C ( t ) i 1 ,i 2 in Theorem 3.3 via gradient flow, which proves the second item of Theorem 3.3.

5 Discussion and Future Directions

In this work, we developed a novel gradient flow based framework for analyzing the training dynamics of a one-layer transformer to recognize co-occurring tokens. We provided 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.

References

[ACDS24] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems , 36, 2024.

[ADF + 23] Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403 , 2023.

[BCE + 23] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712 , 2023.

[BCW + 23] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637 , 2023.

[CL24] Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. arXiv preprint arXiv:2402.04084 , 2024.

[CSWY24] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442 , 2024.

[Dam18] Friederick J Damerau. Markov models and linguistic theory , volume 95. Walter de Gruyter GmbH & Co KG, 2018.

[DBK + 20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2020.

[DCLT18] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pretraining of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

[EGKZ22] Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

[HCL23] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 , 2023.

[HWCL24] Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn feature-position correlations in masked image modeling. arXiv preprint arXiv:2403.02233 , 2024.

[JSL22] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:3782237836, 2022.

[LAG + 22] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations , 2022.

[LLC + 21] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

[LLR23] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. arXiv preprint arXiv:2303.04245 , 2023.

[LWL + 24] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607 , 2024.

[LWLC22] Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference on Learning Representations , 2022.

[LWM + 23] Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, Zaixi Zhang, and Pin-Yu Chen. What improves the generalization of graph transformer? a theoretical dive into self-attention and positional encoding. NeurIPS 2023 Workshop: New Frontiers in Graph Learning, 2023 , 2023.

[MH23] James Manyika and Sissie Hsiao. An overview of bard: an early experiment with generative ai. AI. Google Static Documents , 2, 2023.

[MS99] Christopher Manning and Hinrich Schutze. Foundations of statistical natural language processing . MIT press, 1999.

[NDL24] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 , 2024.

[Ope23] OpenAI. Gpt-4 technical report, 2023.

[PBM21] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research , 22(1):3463-3497, 2021.

[SHT23] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. arXiv preprint arXiv:2306.02896 , 2023.

[Tea23] Gemini Team. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 , 2023.

[Thr24] Christos Thrampoulidis. Implicit bias of next-token prediction. arXiv:2402.18551 , 2024.

[TLI + 23] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

[TLTO23] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023.

[TLZO23] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Maxmargin token selection in attention mechanism. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

[TWCD23] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv preprint arXiv:2305.16380 , 2023.

[TWZ + 23] Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. In The Twelfth International Conference on Learning Representations , 2023.

[VDT24] Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738 , 2024.

[VONR + 23] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

[VSP + 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

[WCM22] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems , 35:12071-12083, 2022.

[WLCC23] Yongtao Wu, Fanghui Liu, Grigorios Chrysos, and Volkan Cevher. On the convergence of encoder-only shallow transformers. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

[WXM21] Colin Wei, Sang Michael Xie, and Tengyu Ma. Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning. Advances in Neural Information Processing Systems , 34:16158-16170, 2021.

[YBR + 19] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations , 2019.

[ZFB24] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research , 25(49):1-55, 2024.

[ZPGA23] Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08117 , 2023.

[ZZYW23] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305.19420 , 2023.

Contents

A Setup of Synthetic Experiment

We conduct synthetic experiment to verify our theoretical results. We create a dataset following our data distribution in Definition 2.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.

B Gradient Flow Update for Weight Matrices

We provide the gradient flow update for each weight matrix as follows.

C Gradient Flow Dynamical System

We first provide our complete dynamical system. The derivation of each equation is provided in Appendix C.1.

C.1 Derivation of the Dynamical System

Lemma C.1. Let µ ∈ { µ i } d i =1 . For all j ∈ [ m ] , we have

where

̸

Proof. Let l ( i, µ ) denote the index such that X ( i ) l ( i,µ ) = µ . By the gradient flow update, we have

̸

̸

Lemma C.2. The following equation on the gradient flow holds:

̸

̸

̸

Proof. By gradient flow update, we have

Lemma C.3. Let µ, ν ∈ { µ i } d i =1 . Then the following equation on gradient flow holds:

Proof. The gradient flow update can be derived as follows:

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 m 1 ∑ 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 m 1 ∑ 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 )

̸

̸

Proof. To prove the first result, we have

To prove the second result, we have

To prove the third result, we have

̸

̸

̸

D Initialization

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 have

̸

and for any l ∈ [ d ] ,

Proof. Note that W (0) K µ l 1 , W (0) Q µ l 2 ∼ 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

Proof. Following from Lemma G.2 and from the first-order Taylor approximation on the softmax function from 0 , we have

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 for l 1 = l 2 ∈ [ d ] , we have

̸

for j 1 = j 2 ∈ [ m 1 ] , we have

and for all j ∈ [ m 1 ] , l ∈ [ d ] , we have

Proof. The proof is similar to that for Lemma D.1 and is omitted.

Lemma D.4. Conditioned on the success of the event in Corollary D.2, for all i ∈ [ n ] , l ∈ [ L ] , with probability at least 1 -δ over the randomness in the initialization of W V ,

Proof. First of all, by Corollary D.2 and Assumption 2.3,

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

Proof. Conditioned on V (0 ,i ) p (0 ,i ) l , we have w (0) ⊤ j V (0 ,i ) p (0 ,i ) l ∼ N (0 , σ 2 1 ∥ V (0 ,i ) p (0 ,i ) l ∥ 2 2 ) . Thus, by Gaussian tail bound and a union bound over j ∈ [ m 1 ] , i ∈ [ n ] , l ∈ [ L ] , with probability at least 1 -δ , we have

Lemma D.6. With probability at least 1 -δ over the randomness in the initialization of a , we have

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

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 ∑ m 1 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

Proof. For fixed l ∈ L, j ∈ [ m ] , i ∈ [ n ] , by Corollary D.5, we have

Thus, this implies that a j w (0) ⊤ j V (0 ,i ) p (0 ,i ) l is a sub-Gaussian random variable with variance proxy σ 2 0 σ 2 1 m L log m 1 nL δ . Therefore, the following inequality holds.

Taking a union bound over i ∈ [ n ] , l ∈ [ L ] , with probability at least 1 -δ , for all i ∈ [ n ] , we have

Finally, by Condition 1, we can make | F (0) ( X ( i ) ) | ≤ 0 . 01 .

E Training Dynamics: Phase 1

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 be

Definition E.2 (Phase 1) . Define the range of Phase 1 to be [0 , T 1 ] , where T 1 = min { t ′ , C T 1 / ( σ 2 1 mm 1 ) } for some sufficiently large constant C T 1 and t ′ is defined to be the maximum time such that for all t ≤ t ′ , all of the following hold:

Based on this definition, we can further obtain the maximum softmax probability change as follows. Proposition E.3. Define

Initially, the loss for the samples with one signal will increase.

E.1 Initial Gradients

Lemma E.4 (Signal updates, same as Lemma 4.1) . At t = 0 , for µ ∈ { µ 1 , µ 2 } , we have

Proof. Take µ = µ 1 . First of all, by the gradient flow update in Lemma C.1, we have

̸

Then

̸

and

̸

Now, by Lemma D.3, Corollary D.2 and Lemma D.8, we have

On the other hand, by Proposition E.5, we have

̸

̸

If m ≥ Cm 1 log 2 m 1 δ in Condition 1 for some sufficiently large C , then | ε 1 | ≤ 0 . 01 σ 2 1 m ; and if m ≥ C ′ L log 2 nL δ 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 have

̸

Further, for all µ ∈ { µ i } d i =1 , we have

̸

̸

̸

Proof. First, fix i, l , and consider the randomness of a . By Lemma D.3, a j 2 〈 w (0) j 1 , w (0) j 2 〉 is a subGaussian random variable with variance proxy σ 4 1 m 2 (√ 4 m log 2 m 2 1 δ + 4 m log 2 m 2 1 δ ) 2 . This implies that with probability at least 1 -δ/ 2 , for all j 1 ∈ [ m 1 ] , we have

̸

Thus, by Corollary D.2, for all µ ∈ { µ i } d i =1 , we have

̸

We next derive the second inequality. Consider the randomness in W (0) V . Note that

̸

̸

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 ,

̸

Therefore,

̸

Lemma E.6 (Random token updates) . For µ ∈ { µ i } d i =4 , we have

Proof. Following the proof of Lemma E.4, we have

̸

̸

By Proposition E.5 and the fact that only one X ( i ) satisfies µ ∈ X ( i ) , we have

and

E.2 Maximum Perturbation of Neuron Outputs

Lemma E.7. For all t ≤ T 1 , for all i ∈ [ n ] , l ∈ [ L ] , we have

Proof. By Definition E.2, we have

By our choice of parameters in Condition 1 and Definition E.2, we have ˜ O ( LR P σ 0 σ 1 √ m )+4 R/L < σ 0 σ 1 √ m/L .

E.3 Perturbation Term Involving Correlation of Value-transformed Data

Proposition E.8. With probability at least 1 -δ , for all µ ∈ { µ i } d i =1 , we have

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

and thus,

̸

Further,

Thus, by Proposition E.8, we have

which implies

where the last inequality applies Lemma E.7.

Proof. We derive the following bound:

̸

where the last inequality follows from Lemma E.9 and Lemma D.3.

̸

̸

E.4 Perturbation Term Involving Correlation of Neurons

Lemma E.11. For all t ≤ T 1 , we have

̸

Proof. We first derive:

̸

For term (1) , by Proposition E.5, we have

̸

For term (2) , note that

̸

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

̸

Proposition E.12 (Neuron correlation change, Phase 1) . For t ≤ T 1 , for all j 1 , j 2 ∈ [ m 1 ] , we have

and thus,

̸

̸

̸

Proof. By the gradient flow update in Lemma C.2, we have

Now, since by Lemma E.7 we have

by Definition E.2 and Corollary D.5, we have

Thus,

Corollary E.13 (Neuron Norm Change, Phase 1) . For all t ≤ T 1 and all j ∈ [ m ] , we have

E.5 Neuron Weights Align with Signal Value

Theorem E.14 (Signal correlation growth, phase 1) . For t ≤ T 1 , for µ ∈ { µ 1 , µ 2 } ,

where

Proof. We take µ = µ 1 and the proof is similar for µ = µ 2 . By Lemma C.1, we have

where

̸

We can now bound the magnitude of ε by Corollary E.10, Lemma E.11 and Proposition E.5 and obtain:

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

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

where

̸

We can now bound the magnitude of ε by Corollary E.10 and Lemma E.11 as follows:

Finally, by Lemma D.3 and Corollary E.13, we have

̸

̸

E.6 Alignment of Common Token

Lemma E.16 (Initial Per Neuron Gradient of Common Token, same as Lemma 4.2) . Let F = max i F (0) i . We have

Proof. Following from the proof of Lemma E.4, we have

̸

For the first term, we have

By Lemma D.3, we have ∥ w (0) j 1 ∥ 2 2 = Θ( σ 2 1 m ) and ∥ v (0) ( µ 3 ) ∥ 2 2 = Θ( σ 2 0 m ) , and thus,

Further, by Proposition E.5, we have

̸

̸

Proof. First of all, we have

By Lemma E.4 and Lemma E.16, we have

̸

Finally, we have

Lemma E.17 (Per neuron gradient of common token) . For t ≤ T 1 , we have

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 as

We can compute the gradient of the sub-network as

Theorem E.19 (Complete version of Lemma 4.3) . There exists a T 0 . 5 ≤ T 1 and a constant 0 < C < 1 such that for all T 0 . 5 ≤ t ≤ T 1 ,

Further,

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 ) ∂t . 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 have

where 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 if ∂G ( t ) ( µ 1 ) ∂t ≥ ∂G ( t ) ( µ 3 ) ∂t then ∑ i ∈ I 1 g ( t ) i -∑ i ∈ I 2 g ( t ) i = Ω(1) since ∑ i ∈ I 3 ∪ I 4 g ( t ) i = Ω(1) for all t ≤ T 1 . Now, by Proposition E.21, if the initialization level is small enough, then ∑ i ∈ I 1 g ( t ) i -∑ i ∈ I 3 g ( t ) i ≥ Ω(1) . Thus, if ∂G ( t ) ( µ 1 ) ∂t ≥ -∂G ( t ) ( µ 3 ) ∂t , then 1 n ∑ i ∈ I 1 ∂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 ∈ I 4 g ( t ) i < min( ∑ i ∈ I 2 g ( t ) i , ∑ i ∈ I 3 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 ∈ I 2 g ( t ′ ) i , ∑ i ∈ I 3 g ( t ′ ) i ) -∑ i ∈ I 4 g ( t ′ ) i = Ω(1) . Thus,

2 ∑ i ∈ I 1 g ( t ′ ) i -2 ∑ i ∈ I 4 g ( t ′ ) i = Ω(1) ,

which implies

For i ∈ I 1 , we have

For i ∈ I 2 ,

For i ∈ I 3 , by Proposition E.21,

and for i ∈ I 4 , for all t ≤ T 1 .

Thus, by chain rule ∂g ( t ) i ∂t = ∂g i ∂F i ∂F ( t ) i ∂t , we have

This implies that there exists a constant L such that

which implies that there must exist a time T 0 . 5 ≤ O (1 /m ) such that for all t ≥ T 0 . 5 ,

Moreover, by Theorem E.26, if C T 1 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 ) ∂t . During this time, note that we have

Further, if ∂G ( t ) ( µ 1 ) ∂t + ∂G ( t ) ( µ 2 ) ∂t = -∂G ( t ) ( µ 3 ) ∂t , then ∂ ∂t g ( t ) i = -Θ( σ 2 1 mm 1 ) for all i ∈ I 2 ∪ I 3 ∪ I 4 . Thus, there must exist a constant U such that

Corollary E.20. There exists a small positive constant C such that if we define T ′ := min { t : min i ∈ I 2 ∪ I 3 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

Proof. First of all, by Lipschitzness of g , we have

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

By Theorem E.14, we have

and

By Grönwall's inequality, we have

This implies that

Thus, if the initialization scale is sufficiently small, 1 n ∑ i ∈ I 2 g ( t ) i and 1 n ∑ i ∈ I 3 g ( t ) i are only differed by a small constant.

E.7 Small Score Movement in Phase 1

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 have

̸

and for all l, l ′ ∈ [ L ] with K (0 ,i ) l = k (0) ( ν ) and q (0 ,i ) l = Q (0) ( ν ) , we have

̸

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.

Finally, take a union bound over i ∈ [ n ] , µ, ν ∈ { µ i } d i =1 . The analysis for the second term is similar. Next, note that a 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,

Lemma E.23 (Score change) . For all t ≤ T 1 , for all ν, µ ∈ { µ i } d i =1 , we have

and thus,

Proof. First of all, by Lemma C.4, we expand the per step gradient descent update as follows:

̸

Analysis of (1) : By triangle inequality, we have

To analyze ( a ) , since ( a + ε 1 )( b + ε 2 ) = ab + aε 2 + bε 1 + ε 1 ε 2 , we have

̸

Further, by Lemma D.1, Lemma D.3, Definition E.1 and Definition E.2, we have

Combining with Proposition E.22, this implies

Thus,

On the other hand, to analyze ( b ) , we have

Note that

which implies

where ( 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

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

Analysis of (3) : (3) can be analyzed similarly to (1) , and we obtain

Analysis of (4) : (4) can be analyzed similarly to (2) , and we obtain

Finally, combining the bounds on (1) -(4) , we have

Corollary E.24 (Softmax change) . For t ≤ T 1 , we have the following:

Thus,

Thus,

if both X ( i ) l 1 , X ( i ) l 2 ∈ R , then

, then

)

∈ R

Proof. Consider fixed i, l 1 , l 2 . By Lemma G.2, we have

By Lemma E.23, we have the following cases:

if both X ( i ) l 1 , X ( i ) l 2 ∈ R , then

Next, during Phase 1, we have p ( t,i ) l 1 ,l 2 ≤ 1 L + 1 Lm + R P . Therefore,

otherwise,

otherwise,

•

if both

X

(

i

l

1

otherwise,

otherwise,

)

, X

(

i

l

2

otherwise,

Further,

This implies that

if both X ( i ) l 1 , X ( i ) l 2 ∈ R , then

otherwise,

Lemma E.25 ( K,Q self-correlation change) . For t ≤ T , for µ, ν ∈ { µ } d , we have

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.

E.8 All Variables are within Range in Definition of Phase 1

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 T 1 / ( σ 2 1 mm 1 ) (where the constant C T 1 is from the definition of T 1 in Definition E.2), all of the following hold:

Thus, T 1 = C T 1 / ( σ 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 have y i F ( T 1 ) = Θ(1) for all i ∈ [ n ] , and

Proof. By Theorem E.26, we have y i F ( T 1 ) i = O (1) and

Further, by Theorem E.14 and Theorem E.19, we have y i F ( T 1 ) i ≥ Ω(1) for i ∈ I 1 and

Finally, by Theorem E.26, we have that T 1 = C T 1 /m . And note that if the constant C T 1 in Definition E.2 is sufficiently large, then by Corollary E.20, we have T ′ ≤ T 1 . Thus, we have y i F ( T 1 ) i ≥ Ω(1) for i ∈ I 2 ∪ I 3 ∪ I 4 and -G ( T 1 ) ( µ 3 ) ≥ Ω(1) .

Finally, by Lemma D.3, Proposition E.5 and Proposition E.12, we have

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 time T 1 = ˜ O (1 /m ) such that

All the training samples are correctly classified: y i F ( T 1 ) i = Ω(1) for all i ∈ [ n ] .

For t ∈ [0 , T 1 ] ,

The training loss satisfies ̂ L ( T 1 ) = Θ(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.

F Training Dynamics: Phase 2

The idea of the proof is to first define conditions for Phase 2, which guarantees that the small training loss can be achieved. Then we will show that those conditions can be satisfied starting from the end of Phase 1 and up to at least Ω( 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:

The change of softmax probability satisfies:

The sum of neuron correlation satisfies

The gradient of W,W V satisfies

The gradients ∑ i ∈ I g ( t ) i for I ∈ { I 1 , I 2 , I 3 , I 4 } satisfies

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 .

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

Proof. Take I = I 1 and the proof is similar for the remaining cases. For fixed i, j ∈ I 1 , we have

where the inequality is due to y i F ( t ) ≥ C in Definition F.1. Now we consider

i

By Definition F.1, it is easy to see that | (2) | ≤ O (1) . For (1) , we have

Thus, we have 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 have

Next, by Proposition E.12, Proposition E.5 and Lemma D.3, we have

Recall that 1 n ∑ n i =1 g ( T 1 ) i = Θ(1) . On the other hand, by Corollary E.10, Lemma D.3, we have

Now, by Proposition E.21, if F is small enough (which can be achieved by making the initialization scale small enough), then

Lastly, by Theorem E.27, we have y i F ( T 1 ) i ≥ Ω(1) . And it is straightforward to see that | G ( T 1 ) ( µ ) | ≤ O (log m ) for µ ∈ { µ i } 3 i =1 .

Finally, we prove ∑ L l 1 =1 ∑ l 2 : X ( i ) l 2 ∈{ µ k } d k =4 G ( T 1 ) ( X ( i ) l 2 ) p ( T 1 ,i ) q ← l 1 ,k ← l 2 ≤ O (1) . A simple corollary from Lemma D.7 and Lemma D.8 is that

Next, we have

which proves that

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

and thus,

Proof. By Lemma C.2, we obtain

Then, note that the final term is positive since by Definition F.1, we have y i F ( t ) i > 0 for all i ∈ [ n ] . Finally, by Theorem E.27, we have for all t ∈ [ T 1 , T 2 ] ,

F.1 Automatic Balancing of Gradients

Lemma F.5 (Same as Lemma 4.4) . For t ∈ [ T 1 , T 2 ] , there exists a small constant C ≪ 1 such that

Proof. Without loss of generality, assume ∑ i ∈ I 2 g ( t ) i > ∑ i ∈ I 3 g ( t ) i . Then, we have

Now by the quotient rule, we have ∂ ∂t ( ∑ i ∈ I 2 g ( t ) i ∑ i ∈ I 3 g ( t ) i ) ≤ 0 if and only if

Note that

Similarly, for i ′ ∈ I 3 , we have

Now, we analyze Equation (12). Note that for t ∈ [ T 1 , T 2 ] , if ∑ i ∈ I 2 g ( t ) i is sufficiently larger than ∑ i ∈ I 3 g ( t ) i (i.e., ∑ i ∈ I 2 g ( t ) i / ∑ i ∈ I 3 g ( t ) i is bigger than some threshold), then ∂G ( t ) ( µ 1 ) ∂t < ∂G ( t ) ( µ 2 ) ∂t and min i ∈ I 2 ∂y i F ( t ) i ∂t > max i ∈ I 3 ∂y i F ( t ) i ∂t . Thus, the ratio will decrease. Similarly, if ∑ i ∈ I 2 g ( t ) i / ∑ i ∈ I 3 g ( t ) i is too small, the ratio will increase. Next, we compute a bound on ∑ i ∈ I 2 g ( t ) i / ∑ i ∈ I 3 g ( t ) i so that ∂ ∂t ( ∑ i ∈ I 2 g ( t ) i ∑ i ∈ I 3 g ( t ) i ) = 0 . Substituting ∂y i F ( t ) i ∂t in Equation (12), we have

Since we have

By Definition F.1, we have

Thus,

Lemma F.6 (Complete version of Lemma 4.5) . For t ∈ [ T 1 , T 2 ] , we have

Further, there exists a constant C such that

Proof. Without loss of generality, assume ∂G ( t ) ( µ 1 ) ∂t > ∂G ( t ) ( µ 2 ) ∂t . First of all, by Definition F.1, we have

This implies that

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 obtain

We consider the range of R ∈ [1 , 3] . By Definition F.1, we have

where the second equation follows from Lemma F.5. This implies that

which proves the first result and

Thus, to analyze ∂ ∂t -∂G ( t ) ( µ 3 ) ∂t ∂G ( t ) ( µ 1 ) ∂t , we can instead analyze

Recall that by Definition F.1, we have for i 1 ∈ I 1 ,

The above implies that for i 1 ∈ I 1 ,

for i 3 ∈ I 3 ,

Substituting the above into Equation (15) and divide both sides by 1 n ∑ m 1 j 1 =1 ∑ m 1 j 2 =1 〈 a j 1 w ( t ) j 1 , a j 2 w ( t ) j 2 〉 , we have

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 , where

Now we analyze the equality condition in Equation (16) where we calculate the root of the equation aR 2 + bR + c = 0 . We have

Note that

By Equation (14), we have -b 2 a ≥ -1 / 4 and

and

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 root R ⋆ = -b 2 a + √ b 2 4 a 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 have

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 obtain

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 have

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 ) ∂y i F ( t ) i ∂t . By Definition F.1, we have for i 1 ∈ I 1 ,

and for i 4 ∈ I 4 ,

for all i ∈ [ n ] . Therefore,

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.4, we have for all t ∈ [ T 1 , T 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 by

This implies for t ∈ [ T 1 , T 2 ] , we have

Corollary F.9. The following bound holds:

Proof. This is a direct consequence of Lemma F.4 and Theorem F.8.

F.3 Growth of Neuron Correlation

Lemma F.10. For t ∈ [ T 1 , T 2 ] , we have

and thus,

By Theorem F.8, we have 1 n ∑ n i =1 g ( t ) i = O ( ̂ L ( t ) ) = O ( 1 ( σ 2 1 mm 1 )( t -T 1 )+1 / ̂ L ( T 1 ) ) . Further, by Definition F.1, we have | F ( t ) i | ≤ O (log m ) . Thus, for t ∈ [ T 1 , T 2 ] , we obtain

Thus, we have

where the last line follows because T 2 ≤ O ( poly ( m )) in Definition F.1. Finally, by Theorem E.27 we have

F.4 Growth of Correlation of Value-Transformed Data

We 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

Thus, for t ∈ [ T 1 , T 2 ] , we have

Proof. By the gradient flow update, we have

Thus, we obtain

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 have

Therefore, by Theorem F.8, we obtain

Corollary F.12. For t ∈ [ T 1 , T 2 ] , we have

̸

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 have

Recall that

This implies that

Corollary F.13 (Complete version of Corollary 4.6) . For t ∈ [ T 1 , T 2 ] , we have

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

Thus,

Proof. By Lemma C.1 and Definition F.1, we have

̸

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,

Thus, by Theorem F.8, for t ∈ [ T 1 , T 2 ] , we have

Corollary F.15. For t ∈ [ T 1 , T 2 ] , we have

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 have

where the last inequality applies Definition F.1. This implies that, for t ∈ [ T 1 , T 2 ] , by Lemma F.14, we have

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.,

The change of score satisfies:

For all µ ∈ { µ 1 , µ 2 , µ 3 } , γ ∈ { µ i } d i =4 , the query-key-correlation score changes as follows:

and

̸

Now, we take µ = µ 1 , ν = µ 2 . By Theorem E.27, we have G ( T 1 ) ( µ 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 ( T 1 ) ( µ ) = ˜ O ( σ 0 σ 1 √ mm 1 ) for µ ∈ R and then by Lemma F.14, we have G ( t ) ( µ ) = ˜ O ( σ 0 σ 1 √ mm 1 ) + ˜ O ( 1 n + m 1 ( L ( σ 2 0 √ m +1 /m )+ σ 2 0 m ) σ 2 1 mm 1 ) for t ∈ [ T 1 , T 2 ] . Also, Definition F.1 implies that G ( t ) ( µ 2 ) -∑ m 1 j =1 w ( t ) j V ( t,i ) p ( t,i ) l ( i,µ ) ≥ Ω(1) . Now, this yields

On the other hand, by the analysis similar to the above, we obtain

Next, to prove the maximum change of the score, we have

Finally, for γ ∈ { µ i } d i =4 , µ ∈ { µ i } d i =1 , we have

which implies that

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.,

For all µ ∈ { µ 1 , µ 2 } , γ ∈ { µ i } d i =4 , the softmax probability between one target signal and a random token changes as follows:

Furthermore, we have

Proof. Take µ = µ 1 , ν = µ 2 . First of all, by Definition F.1 and Lemma F.16, we have

Thus, by Lemma G.1 and Lemma F.16, for i ∈ I 1 , we obtain

On the other hand, for γ ∈ { µ i } d i =4 , we have

Corollary F.18. For t ∈ [ T 1 , T 2 ] , we have

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

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

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,

which implies

where the inequality follows from Theorem F.8 and Definition F.1.

F.8 Small Loss is Achieved

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:

Corollary F.9 proves that

Corollary F.12 proves that

Corollary F.15 proves that

Corollary F.17 proves that

Corollary F.18 proves that

Lemma F.5 implies that

Lemma F.6 implies that the gradients ∑ i ∈ I g ( t ) i for I ∈ { I 1 , I 2 , I 3 , I 4 } satisfies

Lemma F.6 and Corollary F.17 proves that y i F ( t ) i ≥ y i F ( T 1 ) 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.2

Proof. This is proved in Corollary F.13 and Lemma F.16.

F.10 Proof of Theorem 3.3

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 ) .

G Auxiliary Results

The gradient of p ( x ) i = Softmax ( x ) i for x ∈ R n :

̸

Lemma G.1 (Gradient of softmax) . Let s ( t ) ∈ R l be differentiable in t and p ( s ) = Softmax ( s ) . Denote p i ( s ) = Softmax ( s ) i . Then

Proof. By the chain rule and the gradient of softmax in Equation (17), we obtain

Lemma G.2 (Perturbation of softmax) . Let s ∈ R l and p ( s ) = Softmax ( s ) . Denote p i ( s ) = Softmax ( s ) i . Consider a small perturbation ε ∈ R l to s . Then where ∥ ξ ∥ ∞ = O ( ∥ ε ∥ 2 )

Proof. By Taylor's expansion theorem on softmax and the gradient of softmax in Equation (17), we have where ∥ ξ ∥ ∞ = O ( ∥ ε ∥ 2 ) .

H Probability

Lemma H.1 (Bernstein's inequality for bounded random variables) . Assume Z 1 , . . . , Z n are n i.i.d. random variables with E [ Z i ] = 0 and | Z i | ≤ M for all i ∈ [ n ] almost surely. Let Z = ∑ n i =1 Z i . Then, for all t > 0 ,

which implies with probability at least 1 -δ ,

Lemma H.2. For w 1 , w 2 ∈ R m with w 1 , w 2 i.i.d. ∼ N (0 , I m /m ) , we have

Note that w 2 1 ,i is a sub-Gamma random variable with parameters ( 4 m 2 , 4 m ) . Thus, by Bernstein's inequality,

Next,

By Bernstein's inequality, we obtain

Proof. We first have

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not include experiments.

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. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., 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.

5. Open access to data and code

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.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. 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. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: The experiments are done with synthetic data of small scale and the behavior of the experiment results are pretty consistent.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

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.

12. 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.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent 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.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well 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.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe paper makes no restricting assumptions on the weights of the transformer model and performs the analysis on the joint optimization of all the parameters.\n\nAlthough the paper and its proof are notation-heavy, the authors have broken down the complexity of the proof and notation in the main body to clarify the steps needed to prove the results.\n\nWEAKNESSES:\nThere are some restrictive assumptions on the synthetic data model: The vocabulary set $d$ is considered to be larger than the number of training tokens, which is not the case in realistic setups. Thus, some tokens are not visited at training time. Also, they assume, apart from the two target tokens, the remaining tokens appear at maximum once in the training set.  \n\nThe proof outline in the main body helps in understanding the high-level steps involved. However, it could still benefit from additional clarifications on some intermediate steps. For instance, in phases 1 and 2, it's mentioned how the alignment of the MLP with the target tokens $G^{(t)}(\\mu_{1,2})$ behaves during training. However, it's not clear how this connects to the evolution of the attention scores in phase 2 as stated in Lemma 4.7.\n\nQUESTIONS:\n1. As far as I understand, in your synthetic task,  in the first phase of training, effectively only the MLP weights are learning the task. That is, the model can achieve 100% accuracy only by aligning the MLP weights with the relevant tokens $\\mu_1,\\mu_2$. So, the attention layer is not needed for identifying the co-occurrence of the tokens in this setup?\n\n2. Can you also report validation and accuracy plots in your synthetic experiments? Does the validation loss decay at the same rate as the training loss as stated in Thm 3.3?\n\n3. Regarding the proof sketch:\n\n    a) The alignment of parameters with the target tokens is discussed in the main body. Can you also clarify how the gradients related to irrelevant tokens evolve? In particular, regarding the tokens that do not appear in the training set (since $nL\\leq d$), does the model learn not to attend to them at test time?\n\n    b) I find it confusing that the softmax output remains close to $1/L$ long in the training (line 320) and assigns uniform attention to all tokens in the sequence. Does this statement hold for all training samples? If yes, then how does the model learn to attend to the target tokens?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the training dynamics of a single hidden layer transformer network (self-attention + linear MLP) trained on a binary word cooccurrence task. Specifically, given a data matrix $X \\in R^{d \\times L}$ representing L \"words\" (each column of X is a word vector of dimension d), the model must output +1 if words 1 and 2 both occur in X, and -1 otherwise. The paper shows that a transformer layer is able to learn this task, and that the training occurs in two stages: First, the linear MLP layer learns to classify data points correctly by positively aligning with the embeddings for words 1 and 2 (but without making large changes to attention matrices). Second, it drives the loss down further by using the attention matrices to positively correlate q,k,v for words 1 and 2, and anti-correlate the q,k,v for a common word (denoted word \"3\" in the paper) relative to words 1,2.  After these phases, both the training and generalization losses go to zero (as long as embedding dimension is large enough).\n\nOverall, I found the results interesting and insightful, though not very surprising, and the practical implications of these results were not very clear to me. Thus, I currently recommend weak accept. Importantly, my primary research area is not learning theory, so my knowledge of the related work is relatively limited, and thus my review confidence is relatively low.\n\nSTRENGTHS:\n- It is interesting to see that the training dynamics for this word cooccurrence task can be analyzed rigorously, with relatively few assumptions.\n- The theoretical results are validated with a nice synthetic experiments, that demonstrates that the two phases predicted by the theory do occur in practice.\n\nWEAKNESSES:\n- This word cooccurrence task is very simple, and thus it is not surprising that a single transformer layer can easily learn it.\n- Only full gradients are considered, whereas transformers are typically trained with mini-batch Adam(W).\n\nQUESTIONS:\n- What other tasks (beyond word cooccurrence) do you think could be analyzed with this methodology?\n- What are the implications of this result to more complex/realistic tasks, like next token prediction?\n- If mini-batch Adam is used during training, do the two phases still occur?\n- Can you add more details about the experimental setup to the main paper?\n- Can you add more discussion about the automatic balancing of gradients, and its significance, in a more central part of the text (e.g., section 3, not 4)?\n\nLIMITATIONS:\nYes, limitations have been discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

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

STRENGTHS:
- This article noticed an interesting phase transition during training in this special setting and demonstrates it with solid calculation and experiments.
- A new property of gradient flow is noticed and contributes to prove near minimum training loss together with the analysis of softmax.

WEAKNESSES:
The setting of empirical experiments is also simple and ideal and readers may have no idea if this is a general phenomenon during training for detecting word co-occurrence.

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:
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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13