Title: Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer
Venue: NeurIPS 2023
Accepted: yes
Keywords: transformer

training dynamics

theoretical analysis

self-attention

interpretability

neural network understanding
Abstract: Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss,  how the representation emerges from the gradient \emph{training dynamics} remains a mystery. In this paper, for 1-layer transformer with one self-attention layer plus one decoder layer, we analyze its SGD training dynamics for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a \emph{discriminative scanning algorithm}: 
 starting from uniform attention, it gradually attends more to distinct key tokens for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but stops due to a \emph{phase transition} that is controllable by the learning rate of the decoder layer, leaving (almost) fixed token combination. We verify this \textbf{\emph{scan and snap}} dynamics on synthetic and real-world data (WikiText-103).

Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

Abstract

Transformer architectures have shown impressive performance in multiple research domains and have become the backbone of many neural network models. However, there is limited understanding on how Transformer works. In particular, with a simple predictive loss, how the representation emerges from the gradient training dynamics remains a mystery. In this paper, we analyze the SGD training dynamics for 1-layer transformer with one self-attention plus one decoder layer, for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that selfattention acts as a discriminative scanning algorithm : starting from uniform attention, it gradually attends more to key tokens that are distinct for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a phase transition that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this scan and snap dynamics on synthetic and real-world data (WikiText).

1 Introduction

The Transformer architecture [1] has demonstrated wide applications in multiple research domains, including natural language processing [2, 3, 4], computer vision [5, 6, 7], speech [8, 9], multimodality [10, 11], etc. Recently, large language models (LLMs) based on decoder-only Transformer architecture also demonstrate impressive performance [4, 12, 13], after fine-tuned with instruction data [14] or reward models [15]. Why a pre-trained model, often supervised by simple tasks such as predicting the next word [4, 3, 13] or filling in the blanks [2, 16, 17], can learn highly valuable representations for downstream tasks, remains a mystery.

To understand how Transformer works, many previous works exist. For example, it has been shown that Transformer is a universal approximator [18], can approximate Turing machines [19, 20], and can perform a diverse set of tasks, e.g., hierarchical parsing of context-free grammar [21], if its weights are set properly. However, it is unclear whether the weights designed to achieve specific tasks are at a critical point, or can be learned by SoTA optimizers (e.g., SGD, Adam [22], AdaFactor [23], AdamW [24]). In fact, many existing ML models, such as k -NN, Kernel SVM, or MLP, are also universal approximators, while their empirical performance is often way below Transformer.

To demystify such a behavior, it is important to understand the training dynamics of Transformer, i.e., how the learnable parameters change over time during training. In this paper, as a first step, we formally characterize the SGD training dynamics of 1-layer position-encoding-free Transformer for

next token prediction, a popular training paradigm used in GPT series [3, 4], in a mathematically rigorous manner. The 1-layer Transformer contains one softmax self-attention layer followed by one decoder layer which predicts the next token. Under the assumption that the sequence is long, and the decoder learns faster than the self-attention layer, we prove the following interesting dynamic behaviors of self-attention during training. Frequency Bias : it progressively pays more attention to key tokens that co-occur a lot with the query token, and loses attention to tokens that co-occur less. Discriminative Bias : it pays attention to distinct tokens that appear uniquely given the next token to be predicted, while loses interest to common tokens that appear across multiple next tokens. These two properties suggest that self-attention implicitly runs an algorithm of discriminative scanning , and has an inductive bias to favor unique key tokens that frequently co-occur with the query ones.

Furthermore, while self-attention layer tends to become more sparse during training, as suggested by Frequency Bias, we discover that it will not collapse to one-hot, due to a phase transition in the training dynamics. In the end, the learning does not converge to any stationary points with zero gradient, but ventures into a region where the attention changes slowly (i.e., logarithmically over time), and appears frozen and learned. We further show that the onset of the phase transition are controlled by the learning rates: large learning rate gives sparse attention patterns, and given fixed self-attention learning rate, large decoder learning rate leads to faster phase transition and denser attention patterns. Finally, the SGD dynamics we characterize in this work, named scan and snap , is verified in both synthetic and simple real-world experiments on WikiText [25].

Concurrent works on Transformer dynamics . Compared to [26] that uses ℓ 2 loss, our analysis focuses on cross-entropy, which is more realistic, imposes no prior knowledge on possible attention patterns inaccessible to training, and allows tokens to be shared across topics. Compared to [27] that analyzes 'positional attention' that is independent of input data with symmetric initialization, our analysis focuses on attention on input data without symmetric assumptions. [28, 29, 30] give similar conclusions that self-attention attends to relevant tokens. In comparison, our work analyzes richer phenomena in 1-layer transformers related to frequency and discriminative bias, which has not been brought up by these works. For example, sparse attention patterns are connected with co-occurrence frequency of contextual token and query, characterization of such connection over training with softmax, including two-stage behaviors of attention logits, etc. We also leverage analytical solutions to certain nonlinear continuous dynamics systems that greatly simplifies the analysis. Detailed comparison can be found in Appendix B.

2 Related Works

Expressiveness of Attention-based Models . A line of work studies the expressive power of attention-based models. One direction focuses on the universal approximation power [18, 31, 32, 33, 20]. More recent works present fine-grained characterizations of the expressive power for certain functions in different settings, sometimes with statistical analyses [34, 35, 36, 37, 21, 38, 39, 40]. In particular, there is growing interest in explaining the capability of in-context learning [41] of Transformer, by mapping the gradient descent steps of learning classification/regression into feedforward steps of Transformer layers [42, 43, 44, 45, 37, 46]. Different from our work, the results in these papers are existential and do not take training dynamics into consideration.

Training Dynamics of Neural Networks . Previous works analyze the training dynamics in multilayer linear neural networks [47, 48], in the student-teacher setting [49, 50, 51, 52, 53, 54, 55, 56, 57], and infinite-width limit [58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], including extentions to attention-based models [72, 73]. For self-supervised learning, works exist to analyze linear networks [74] and understand the role played by nonlinearity [75]. Focusing on attentionbased models, Zhang et al. [76] study adaptive optimization methods in attention models. Jelassi et al. [27] propose an idealized setting and show the vision transformer [5] trained by gradient descent can learn spatial structure. Li et al. [26] show that the 1-layer Transformer can learn a constrained topic model, in which any word belongs to one topic, with ℓ 2 loss, BERT [2]-like architecture and additional assumptions on learned attention patterns. Snell et al. [77] study the dynamics of a singlehead attention head to approximate the learning of a Seq2Seq architecture. While these papers also study the optimization dynamics of attention-based models, they focus on different settings and do not explain the phenomena presented in our paper.

3 Problem Setting

Notation. Let { u k } M k =1 be d -dimensional embeddings, and { x t } be discrete tokens. For each token, x t takes discrete values from 1 to M , denoted as x t ∈ [ M ] , and x t := e x t ∈ R M is the corresponding one-hot vector, i.e., the x t -th entry of x t is 1 while others are zero. u x t is the token embedding at location t in a sequence.

Let U = [ u 1 , . . . , u M ] ⊤ ∈ R M × d be the embedding matrix, in which the k -th row of U is the embedding vector of token k . X = [ x 1 , . . . , x T -1 ] ⊤ ∈ R ( T -1) × M is the data matrix encoding the sequence of length T -1 . XU ∈ R ( T -1) × d is the sequence of embeddings for a given sequence { x 1 , . . . , x T -1 } . It is clear that X 1 M = 1 T -1 .

We use X [ i ] to denote i -th sample in the sequence dataset. Similarly, x t [ i ] is the token located at t in i -th sample. Let D be the dataset used for training.

1-Layer Transformer Architecture . Given a sequence { x 1 , . . . , x T , x T +1 } , the embedding after 1 -layer self attention is:

Here b tT is the normalized self-attention weights ( ∑ T -1 t =1 b tT = 1 ). One important detail is that we mask the weight that the query token attends to itself, which is also being used in previous works (e.g., QK-shared architecture [78]). See Sec. 7 for discussions about residual connection. Let b T := [ b 1 T , . . . , b T -1 ,T ] ⊤ ∈ R T -1 be an attention vector, then b ⊤ T 1 = 1 and ˆ u T = U ⊤ X ⊤ b T .

ℓ 2 -Normalization . We consider adding a normalization to the output of the self-attention layer: ˜ u T = U ⊤ LN( X ⊤ b T ) , where LN( x ) := x / ∥ x ∥ 2 . NormFormer and RMSNorm [79, 80] also leverages this setting (up to a global constant). Our analysis can also be extended to standard LayerNorm [81], which also subtracts the mean of x , while [80] shows that mean subtraction may not affect the empirical results much. LLaMA [82] also uses RMSNorm. Empirically ˆ u T (or W V ˆ u T ) is normalized (instead of X ⊤ b T ) and here we use an approximation to facilitate analysis: when the token embedding { u m } are approximately orthogonal to each other, then ∥ U ⊤ x ∥ 2 ≈ ∥ x ∥ 2 and thus ˜ u T ≈ LN(ˆ u T ) .

Objective . We maximize the likelihood of predicted ( T +1) -th token using cross entropy loss:

For simplicity, we consider single-head attention setting, and multiple-head attention can be regarded as single-head setting with simultaneous different initializations (see Sec. 4). We call x T = m as the query token of the sequence, and x T +1 = n as the next token to be predicted. Other tokens x t ( 1 ≤ t ≤ T -1 ) that are encoded in X are called contextual tokens . Both the contextual and query tokens can take values from 1 to M (i.e., m ∈ [ M ] ) and next token takes the value from 1 to K (i.e., n ∈ [ K ] ) where K ≤ M . Fig. 1(a) shows the overall setting. For an overview of the notation used in the paper, please check Tbl. 1 in the Appendix.

3.1 Reparameterization

Instead of studying the dynamics with respect to the parameters of token embedding U , key, value and query projection matrices W K , W Q and W V , we study the dynamics of two pairwise token relation matrices Y := UW ⊤ V U ⊤ ∈ R M × M and Z := UW Q W ⊤ K U ⊤ / √ d ∈ R M × M . Intuitively, entries of Y and Z store the 'logits' of pairs of tokens. We regard the empirical parameterization using U , W K , W Q and W V as a specific way of parametrization of Y and Z , in order to reduce the number of parameters to be estimated. Previous work also leverage similar parameterization for self-attention layers [27, 46].

For real-world applications, the number of tokens M can be huge (e.g., the vocabulary size M = 50 , 272 in OPT-175B [83] and M = 32 , 000 in LLaMA [82]) and directly optimizing Y and Z would be prohibitive. However, as we will show in this work, from the theoretical perspective, treating Y and Z as independent variables has some unique advantages and leads to useful insights.

Contextual tokens

𝑥

!

𝑥

"

𝑥

#$!

𝑥

#

𝑥

#%!

Query token

Next token

(a)

(b)

Self-attention

Normalization

Decoding & Softmax

ℙ(𝑙|𝑛

!

)

𝑚

!

𝑛

!

ℙ(𝑙|𝑛

"

)

𝑛

"

ℙ(𝑙|𝑛

#

)

𝑚

"

𝑛

,

ℙ(𝑙|𝑛

$

)

𝑛

-

Query token

𝑥

%

Next token

𝑥

%&!

Contextual tokens

𝑥

'

(

1 ≤ 𝑡 ≤ 𝑇 - 1)

Sequence

classes

Common tokens:

There exists multiple

𝑛

with

ℙ(𝑙|𝑛) > 0

Distinct tokens:

There exists unique

𝑛

with

ℙ(𝑙|𝑛) > 0

Lemma 1 (Dynamics of 1-layer Transformer) . The gradient dynamics of Eqn. 2 with batchsize 1 is:

Here P ⊥ v := I -vv ⊤ / ∥ v ∥ 2 2 projects a vector into v 's orthogonal complementary space, η Y and η Z are the learning rates for the decoder layer Y and self-attention layer Z , α := [ α 1 , . . . , α M ] ⊤ ∈ R M and α m := exp( Y ⊤ LN( X ⊤ b T )) / 1 ⊤ exp( Y ⊤ LN( X ⊤ b T )) .

Please check Appendix C for the proof. We consider Y (0) = Z (0) = 0 as initial condition. This is reasonable since empirically Y and Z are initialized by inner product of d -dimensional vectors whose components are independently drawn by i.i.d Gaussian. This initial condition is also more realistic than [27] that assumes dominant initialization in diagonal elements. Since ( x T +1 -α ) ⊤ 1 = 0 and P ⊥ X ⊤ b T X ⊤ diag( b T ) X 1 = 0 , we have ˙ Y 1 = ˙ Z 1 = 0 and summation of rows of Z ( t ) and Y ( t ) remains zero. Since x T is a one-hot column vector, the update of Z = [ z 1 , z 2 , . . . , z M ] ⊤ is done per row:

̸

where m = x T [ i ] is the query token for sample i , z m is the m -th row of Z and ˙ z m ′ = 0 for row m ′ = m = x T [ i ] . Note that if x T [ i ] = m , then b T [ i ] is a function of z m only (but not a function of any other z m ′ ). Here we explicitly write down the current sample index i , since batchsize is 1 .

3.2 Data Generation

Next we specify a data generation model (Fig. 1(b)), named sequence class , for our analysis.

Sequence Class. We regard the input data as a mixture of multiple sequence classes . Each sequence class is characterized by a triple s m,n := ( P ( l | m,n ) , m, n ) . To generate a sequence instance from the class, we first set x T = m and x T +1 = n , and then generate the contextual tokens with conditional probability P ( l | m,n ) . Let supp( m,n ) be the subset of token l with P ( l | m,n ) > 0 .

In this work, we consider the case that given a next token x T +1 = n , the corresponding sequence always ends with a specific query token x T = m =: ψ ( n ) . This means that we could index sequence class with next token x T +1 = n alone: s n := ( P ( l | ψ ( n ) , n ) , ψ ( n ) , n ) , P ( l | m,n ) = P ( l | n ) and supp( n ) := supp( ψ ( n ) , n ) .

Note that | ψ -1 ( m ) | = 1 means that the occurrence of token m alone decides next token n to be predicted, regardless of other tokens in the sequence, which is a trivial case. When | ψ -1 ( m ) | ≥ 2 , the same query token m , combined with other token l in the sequence with non-zero probability P ( l | m,n ) > 0 , determine the next token.

̸

Overlapping sequence class . Two sequence classes s n and s n ′ overlap if supp( n ) ∩ supp( n ′ ) = ∅ .

Seq class

(𝑚,𝑛

!

)

Seq class

(𝑚,𝑛

"

)

Distinct

Token

Common

Token

̃ 𝑐

!|#

#

̃ 𝑐

!|#

$

(a) Initialization

̃ 𝑐

!|#

#

̃ 𝑐

!|#

$

(b) Common Token Suppression

̃ 𝑐

!|#

#

̃ 𝑐

!|#

$

(c) Winners-take-all

̃ 𝑐

!|#

#

̃ 𝑐

!|#

$

(d) Attention frozen

(Global) distinct and common tokens . Let Ω( l ) := { n : P ( l | n ) > 0 } be the subset of next tokens that co-occur with contextual token l . We now can identify two kinds of tokens: the distinct token l which has | Ω( l ) | = 1 and the common token l with | Ω( l ) | > 1 . Intuitively, this means that there exists one common token l so that both P ( l | n ) and P ( l | n ′ ) are strictly positive, e.g., common words like 'the' , 'this' , 'which' that appear in many sequence classes. In Sec. 5, we will see how these two type of contextual tokens behave very differently when self-attention layer is involved in the training: distinct tokens tend to be paid attention while common tokens tend to be ignored.

3.3 Assumptions

To make our analysis easier, we make the following assumptions:

Assumption 1. We consider (a) no positional encoding, (b) The input sequence is long ( T → + ∞ ) and (c) The decoder layer learns much faster than the self-attention layer (i.e., η Y ≫ η Z ).

Assumption 1(a) suggests that the model is (almost) permutation-invariant. Given the next token to predict x T +1 = n and the query token x T = m acted as query, the remaining tokens in the sequence may shuffle. Assumption 1(b) indicates that the frequency of a token l in the sequence approaches its conditional probability P ( l | m,n ) := P ( l | x T = m,x T +1 = n ) .

Note that the assumptions are comparable with or even weaker than previous works, e.g., [27] analyzes positional attention with symmetric initialization, without considering input data and [28] models the data distribution as discriminative/non-discriminative patterns, similar to our common/distinct tokens. Empirically, NoPE [84] shows that decoder-only Transformer models without positional encoding still works decently, justifying that Assumption 1(a) is reasonable.

Given the event { x T = m,x T +1 = n } , suppose for token l , the conditional probability that it appears in the sequence is P ( l | m,n ) . Then for very long sequence T → + ∞ , in expectation the number of token l appears in a sequence of length T approaches T P ( l | m,n ) . Therefore the pertoken self-attention weight c l | m,n is computed as:

Here z ml is z m 's l -th entry and ˜ c l | m,n := P ( l | m,n ) exp( z ml ) is un-normalized attention score.

By the data generation process (Sec. 3.2), given the next token x T +1 = n , the query token x T = m is uniquely determined. In the following, we just use c n to represent c m,n (and similar for ˜ c n ).

4 Dynamics of Y

We first study the dynamics of Y . From Assumption 1(c), Y learns much faster and we can treat the lower layer output (i.e., X ⊤ b T ) as constant. From Lemma 2, when the sequence is long, we know

given the next token x T +1 = n , X ⊤ b T becomes fixed. Therefore, the dynamics of Y becomes:

Here f n := X ⊤ b T ∥ X ⊤ b T ∥ 2 → c n ∥ c n ∥ 2 ∈ R M . Obviously ∥ f n ∥ 2 = 1 and f n ≥ 0 . Define F = [ f 1 , . . . , f K ] . Since the vocabulary size M typically is a huge number, and different sequence classes can cover diverse subset of vocabulary, we study the weak correlation case:

Assumption 2 (Weak Correlations) . We assume M ≫ K 2 and { f n } K n =1 satisfies F ⊤ F = I + E , where the eigenvalues of E ∈ R K × K satisfies | λ 1 | < 1 K and | λ i ( E ) | ≥ 6 √ M , ∀ i ∈ [ K ] .

Assumption 2 means that f n share some weak correlations and it immediately leads to the fact that F ⊤ F is invertible and F is column full-rank. Note that the critical point Y ∗ of Eqn. 7 should satisfy that for any given x T +1 = n , we need α = e n . But such Y ∗ must contain infinity entries due to the property of the exponential function in α and we can not achieve Y ∗ in finite steps. To analyze Eqn. 7, we leverage a reparameterized version of the dynamics, by setting W = [ w 1 , . . . , w K ] ⊤ := F ⊤ Y ∈ R K × M and compute gradient update on top of W instead of Y :

While we cannot run gradient update on W directly, it can be achieved by modifying the gradient of Y to be ˙ Y = η Y ( f n -FE ′ e n )( e n -α n ) ⊤ . If λ 1 is small, the modification is small as well.

Please check Appendix D for the proof. Lemma 3 shows that for every fixed n , only the corresponding row of W is updated, which makes the analysis much easier. We now can calculate the backpropagated gradient used in Eqn. 3.

Theorem 1. If Assumption 2 holds, the initial condition Y (0) = 0 , M ≫ 100 , η Y satisfies M -0 . 99 ≪ η Y < 1 , and each sequence class appears uniformly during training, then after t ≫ K 2 steps of batch size 1 update, given event x T +1 [ i ] = n , the backpropagated gradient g [ i ] := Y ( x T +1 [ i ] -α [ i ]) takes the following form:

̸

Here the coefficients ι n ( t ) , β nn ′ ( t ) and γ ( t ) are defined in Appendix with the following properties:

̸

(a) ξ n ( t ) := γ ( t ) ∑ n = n ′ β nn ′ ( t ) f ⊤ n ( t ) f n ′ ( t ) > 0 for any n ∈ [ K ] and any t ;

In the remark of Lemma 5 in Appendix, we analyze the original dynamics (Eqn. 7) with identical off-diagonal elements of E , and Theorem 1 still holds with a smaller effective learning rate.

5 The dynamics of Self-attention

Now we analyze the dynamics of self-attention logits Z , given the dynamics of upper layer Y

Lemma 4 (Self-attention dynamics) . With Assumption 1(b) (i.e., T → + ∞

̸

. ), Eqn. 4 becomes:

Please check Appendix E for the proof. Now we study the dynamics of two types of contextual tokens (Sec. 3.2), namely distinct tokens (DT) which appear only for a single next token (i.e., | Ω( l ) | = 1 with Ω( l ) := { n : P ( l | n ) > 0 } ), and common tokens (CT) that appear across multiple next tokens ( | Ω( l ) | > 1 ). We show their fates are very different: over training, distinct tokens gain attention but common ones lose it .

Theorem 2 (Fates of contextual tokens) . Let G CT be the set of common tokens (CT), and G DT ( n ) be the set of distinct tokens (DT) that belong to next token n . Then if Assumption 2 holds, under the self-attention dynamics (Eqn. 10), we have:

(a) for any distinct token l ∈ G DT ( n ) , ˙ z ml > 0 where m = ψ ( n ) ;

(b) if | G CT | = 1 and at least one next token n ∈ ψ -1 ( m ) has at least one distinct token, then for the single common token l ∈ G CT , ˙ z ml < 0 .

Now we know DTs grow and a single CT will shrink. For multiple CTs to shrink, the condition can be a bit involved (see Corollary 2 in Appendix E). The following theorem further shows that the growth rates of DTs critically depend on their initial conditions:

Theorem 3 (Growth of distinct tokens) . For a next token n and its two distinct tokens l and l ′ , the dynamics of the relative gain r l/l ′ | n ( t ) := f 2 nl ( t ) /f 2 nl ′ ( t ) -1 = ˜ c 2 l | n ( t ) / ˜ c 2 l ′ | n ( t ) -1 has the following analytic form (here the query token m = ψ ( n ) and is uniquely determined by distinct token l ):

̸

where χ l ( t ) := e 2( z ml ( t ) -z ml (0)) is the growth factor of distinct token l . If there exist a dominant token l 0 such that the initial condition satisfies r l 0 /l | n (0) > 0 for all its distinct token l = l 0 , and all of its common tokens l satisfy ˙ z ml < 0 . Then both z ml 0 ( t ) and f nl 0 ( t ) are monotonously increasing over t , and

here B n ( t ) := η Z ∫ t 0 ξ n ( t ′ )d t ′ . Intuitively, larger B n gives larger r l 0 /l | n and sparser attention map.

Self-attention as an algorithm of token scanning . From Eqn. 11, we could see that self-attention performs token scanning . To see that, consider the simplest initialization that z (0) = 0 , which means that r l 0 /l | n (0) = ( P ( l 0 | m,n ) P ( l | m,n ) ) 2 -1 . Therefore, distinct token l with low conditional probability P ( l | m,n ) will have r l 0 /l | n (0) ≫ 0 , According Eqn. 11, this leads to quickly growing ratio r l 0 /l | n ( t ) , which means that the corresponding component f nl will be quickly dwarfed by the dominating component f nl 0 . On the other hand, token with high conditional probability P ( l | m,n ) will have smaller r l 0 /l | n (0) , and the ratio r l 0 /l | n ( t ) grows slower, costing longer time for l 0 to dominate l .

Initial value as prior information . From the theorems, it is clear that the initial value r l/l ′ | n (0) := ( P ( l | m,n ) exp( z ml (0)) P ( l ′ | m,n ) exp( z ml ′ (0)) ) 2 -1 critically determines the fate of the dynamics. Two tokens l and l ′ with comparable conditional probability P ( l | m,n ) and P ( l ′ | m,n ) can be suppressed in either way, depending on their initial logits z ml (0) and z ml ′ (0) . In the empirical implementation, the initial value of the logits are determined by the inner products of independently initialized high-dimensional vectors, which fluctuate around zero.

The concept of 'initial value as prior' can explain empirical design choices such as multi-head selfattention [1]. From this perspective, each head h has its own Z h and is initialized independently, which could enable more diverse token combination (e.g., a combination of 1st, 3rd, 5th tokens, rather than a combination of 1st, 2nd, 3rd tokens).

6 The Moment of Snapping: When Token Combination is fixed

Theorem 3 suggests two possible fates of the self-attention weights: if ξ n ( t ) decays slowly (e.g., ξ n ( t ) ≥ 1 /t ), then B n ( t ) → + ∞ and all contextual tokens except for the dominant one will drop (i.e., f nl → 0 ) following the ranking order of their conditional probability P ( l | m,n ) . Eventually, winner-takes-all happens. Conversely, if ξ n ( t ) drops so fast that B n ( t ) grows very slowly, or even has an upper limit, then the self-attention patterns are 'snapped' and token combination is learned and fixed.

The conclusion is not obvious, since ξ n ( t ) depends on the decay rate of γ ( t ) and β nn ′ ( t ) , which in turns depends on the inner product f ⊤ n ( t ) f n ′ ( t ) , which is related to the logits of the common tokens that also decays over time.

Here we perform a qualitative estimation when there is only a single common token l and every next token shares a single token m (i.e., for any next token n , ψ ( n ) = m ). We assume all normalization

0

100

200

300

400

500

600

700

t

0

100

200

300

400

500

600

l

0

(

t

)

Z

=0.5

Y

=0.5

Y

=1.0

Y

=2.0

Y

=5.0

terms in f n are approximately constant, denoted as ρ 0 , which means that f ⊤ n f n ′ ≈ exp(2 z ml ) /ρ 2 0 and β nn ′ ≈ E ′ nn ′ ≈ f ⊤ n f n ′ ≈ exp(2 z ml ) /ρ 2 0 as well, and 1 -f ⊤ n f n ′ ≈ 1 due to the fact that common token components are small, and will continue to shrink during training.

Under these approximations, its dynamics (Eqn. 10) can be written as follows (here C 0 := ρ 4 0 /K ):

̸

Surprisingly, we now find a phase transition by combining the rate change of γ ( t ) in Theorem 1:

Theorem 4 (Phase Transition in Training) . If the dynamics of the single common token z ml satisfies ˙ z ml = -C -1 0 η Z γ ( t ) e 4 z ml and ξ n ( t ) = C -1 0 γ ( t ) e 4 z ml , then we have:

As a result, there exists a phase transition during training:

Attention scanning . At the beginning of the training, γ ( t ) = O ( η Y t/K ) and B n ( t ) ≈ 1 4 ln K -1 ( ρ 4 0 +2 η Y η Z t 2 ) = O (ln t ) . This means that the growth factor for dominant token l 0 is (sub-)linear: χ l 0 ( t ) ≥ e 2 f 2 nl 0 (0) B n ( t ) ≈ [ K -1 ( ρ 4 0 + 2 η Y η Z t 2 )] 0 . 5 f 2 nl 0 (0) , and the attention on less co-occurred token drops gradually.

Attention snapping . When t ≥ t 0 := 2(1 + δ ′ ) K ln M/η Y with δ ′ = Θ( ln ln M ln M ) , γ ( t ) = O ( K ln( η Y t/K ) η Y t ) and B n ( t ) = O (ln ln t ) . Therefore, while B n ( t ) still grows to infinite, the growth factor χ l 0 ( t ) = O (ln t ) grows at a much slower logarithmic rate.

See proof in Appendix F. This gives a few insights about the training process: (a) larger learning rate η Y of the decoder Y leads to shorter phase transition time t 0 ≈ 2 K ln M/η Y , (b) scaling up both learning rate ( η Y and η Z ) leads to larger B n ( t ) when t → + ∞ , and thus sparser attention maps, and (c) given fixed η Z , small learning rate η Y leads to larger B n ( t ) when t ≥ t 0 , and thus sparser attention map. Fig. 3 shows numerical simulation results of the growth rate χ l ( t ) . Here we set K = 10 and M = 1000 , and we find smaller η Y given fixed η Z indeed leads to later transition and larger B n ( t ) (and χ l ( t ) ).

7 Discussion and Limitations

Positional encoding . While our main analysis does not touch positional encoding, it can be added easily following the relative encoding schemes that adds a linear bias when computing self attention (E.g., T5 [17], ALiBi [85], MusicTransformer [86]). More specifically, the added linear bias exp( z ml + z 0 ) = exp( z ml ) exp( z 0 ) corresponds to a prior of the contextual token to be learned in the self-attention layer.

0

5

10

15

20

25

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

Common token   Distinct token (seq1) Distinct token (seq2)

0

5

10

15

20

25

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

0

5

10

15

20

25

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

Common token   Distinct token (seq1) Distinct token (seq2)

0

5

10

15

20

25

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

Token #

Token #

̸

Residue connection . Residue connection can be added in the formulation, i.e., ¯ u T = LN(LN(˜ u T )+ u x T ) , where ˜ u T is defined in Eqn. 1, and ¯ u T is used instead in the objective (Eqn. 2). In this case, the β nn ′ in Theorem 1 now is approximately β nn ′ ∼ f ⊤ n f n ′ + I ( ψ ( n ) = ψ ( n ′ )) , which is much larger for sequence classes n and n ′ that share the same query token x T than otherwise. In this case, Theorem 1 now gives g [ i ] = γ ( ι n f n -∑ n = n ′ ∈ ψ -1 ( ψ ( n )) β nn ′ f n ′ ) for x T +1 [ i ] = n . Due to the additional constraint n ′ ∈ ψ -1 ( ψ ( n )) (i.e., n and n ′ shares the same query token), we can define local distinct and common tokens to be within the sequence class subset ψ -1 ( m ) and Theorem 2 now applies within each subset. Empirically this makes more sense, since the query token x T = m 1 or m 2 alone can already separate different subsets ψ -1 ( m 1 ) and ψ -1 ( m 2 ) and there should not be any interactions across the subsets. Here we just present the most straightforward analysis and leave this extension for future work.

Possible future extension to multi-layer cases . For multilayer training, a lasting puzzle is to explain how the input tokens get combined together to form high-level concepts. The analysis above shows that the training leads to sparse attention even among relevant tokens, and demonstrates that there is a priority in token combinations for 1-layer attention based on their co-occurrence: even if there are 10 relevant contextual tokens to the query, the self-attention may only pick 1-2 tokens to combine first due to attention sparsity. This can be regarded as a starting point to study how tokens are composed hierarchically. In comparison, [28, 29, 30] show that attention attends to all relevant tokens, which may not suggest a hierarchical / multi-layer architecture.

8 Experiments

We conduct experiments on both synthetic and real-world dataset to verify our theoretical findings.

Syn-Small . Following Sec. 3.2, we construct K = 2 sequence classes with vocabulary size M = 30 . The first 10 tokens (0-9) are shared between classes, while the second and third 10 tokens (10-19 and 20-29) are distinct for class 1 and class 2, respectively. The conditional probability P ( l | n ) for tokens 10-19 is increasing monotonously (the same for 20-29). The 1-layer Transformer is parameterized with Y and Z (Sec. 3.1), is trained with initial condition Y (0) = Z (0) = 0 and SGD (with momentum 0 . 9 ) using a batchsize of 128 and sequence length T = 128 until convergence.

Fig. 4 shows the simulation results. The attention indeed becomes sparse during training, and increasing η Y with fixed η Z leads to faster convergence but less sparse attention. Both are consistent with our theoretical predictions (Theorem 3 and Sec. 6). Interestingly, if we use Adam optimizer instead, self-attention with different learning rate η Y = η Z picks different subsets of distinct tokens to focus on, showing tune-able inductive bias (Fig. 5). We leave analysis on Adam for future work.

Syn-Medium . To further verify our theoretical finding, we now scale up K to create Syn-Medium and compute how attention sparsity for distinct tokens (in terms of entropy) changes with the learning rates (Fig. 6). We can see indeed the entropy goes down (i.e., attention becomes sparser) with larger η Z , and goes up (i.e., attention becomes less sparse) by fixing η Z and increasing η Y passing the threshold η Y /η Z ≈ 2 , consistent with Sec. 6. Note that the threshold is due to the fact that our theory is built on Assumption 1(c), which requires η Y to be reasonably larger than η Z .

0

5

10

15

Token #

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

Common token  Distinct token (seq1)

0

5

10

15

Token #

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

Common token  Distinct token (seq1)

0

5

10

15

Token #

Iter 0

Iter 2

Iter 4

Iter 6

Iter 8

Iter 10

Final

Common token  Distinct token (seq1)

2

4

6

8

10

Y

/

Z

1.4

1.6

1.8

2.0

Average Entropy of

c

n

#last/next tokens = 1/2

Z

=0.5

Z

=1.0

Z

=2.0

Z

=3.0

2

4

6

8

10

Y

/

Z

1.7

1.8

1.9

2.0

#last/next tokens = 3/6

2

4

6

8

10

Y

/

Z

1.7

1.8

1.9

2.0

#last/next tokens = 5/10

2

4

6

8

10

Y

/

Z

1.80

1.85

1.90

1.95

2.00

2.05

#last/next tokens = 10/20

0

10

20

30

0

5

10

15

20

25

30

iter-0

0

10

20

30

0

5

10

15

20

25

30

iter-500

0

10

20

30

0

5

10

15

20

25

30

iter-1000

0

10

20

30

0

5

10

15

20

25

30

iter-1500

0

10

20

30

0

5

10

15

20

25

30

iter-0

0

10

20

30

0

5

10

15

20

25

30

iter-500

0

10

20

30

0

5

10

15

20

25

30

iter-1000

0

10

20

30

0

5

10

15

20

25

30

iter-1500

■

口

Real-world Dataset . We also test our finding on WikiText [25] using both 1-layer and multi-layer Transformers with regular parameterization that computes Y and Z with embedding U . In both cases, attentions of the first layer freeze (and become sparse) at some point (Fig. 7), even if the learning rate remains the same throughout training. More results are in Appendix G.

9 Conclusion and Future Work

In this work, we formally characterize SGD training dynamics of 1-layer Transformer, and find that the dynamics corresponds to a scan and snap procedure that progressively pays more attention to key tokens that are distinct and frequently co-occur with the query token in the training set. To our best knowledge, we are the first to analyze the attention dynamics and reveal its inductive bias on data input, and potentially open a new door to understand how Transformer works.

Many future works follow. According to our theory, large dataset suppresses spurious tokens that are perceived as distinct in a small dataset but are actual common ones. Our finding may help suppress such tokens (and spurious correlations) with prior knowledge, without a large amount of data.

References

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 2017.

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

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

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. arXiv preprint arXiv:2010.11929 , 2020.

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.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll´ ar, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16000-16009, 2022.

Ching-Feng Yeh, Jay Mahadeokar, Kaustubh Kalgaonkar, Yongqiang Wang, Duc Le, Mahaveer Jain, Kjell Schubert, Christian Fuegen, and Michael L Seltzer. Transformer-transducer: End-to-end speech recognition with self-attention. arXiv preprint arXiv:1910.12977 , 2019.

Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. Advances in neural information processing systems , 33:12449-12460, 2020.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. In International Conference on Machine Learning , pages 1298-1312. PMLR, 2022.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

OpenAI. Gpt-4 technical report, 2023.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416 , 2022.

Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems , 33:3008-3021, 2020.

Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, et al. Ul2: Unifying language learning paradigms. In The Eleventh International Conference on Learning Representations , 2022.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

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.

Jorge P´ erez, Pablo Barcel´ o, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research , 22(1):3463-3497, 2021.

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

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , pages 4596-4604. PMLR, 2018.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843 , 2016.

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

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836, 2022.

Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. ICLR , 2023.

Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. ICML , 2023.

Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. arXiv preprint arXiv:2306.13596 , 3(7):47, 2023.

Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009.11264 , 2020.

Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09286 , 2020.

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819 , 2018.

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.

N Elhage, N Nanda, C Olsson, T Henighan, N Joseph, B Mann, A Askell, Y Bai, A Chen, T Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread , 2021.

Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint arXiv:2106.03764 , 2021.

Ekin Aky¨ urek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115 , 2021.

Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint arXiv:2207.04901 , 2022.

Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems , 35:21750-21764, 2022.

Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey for in-context learning. arXiv preprint arXiv:2301.00234 , 2022.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Jo˜ ao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

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.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895 , 2022.

Shuai Li, Zhao Song, Yu Xia, Tong Yu, and Tianyi Zhou. The closeness of in-context learning and weight shifting for softmax regression. arXiv preprint arXiv:2304.13276 , 2023.

Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. arXiv preprint arXiv:1810.02281 , 2018.

Peter Bartlett, Dave Helmbold, and Philip Long. Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. In International conference on machine learning , pages 521-530. PMLR, 2018.

Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for a convnet with gaussian inputs. In International conference on machine learning , pages 605-614. PMLR, 2017.

Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In International conference on machine learning , pages 3404-3413. PMLR, 2017.

Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural information processing systems , 30, 2017.

Surbhi Goel, Adam Klivans, and Raghu Meka. Learning one convolutional layer with overlapping patches. In International Conference on Machine Learning , pages 1783-1791. PMLR, 2018.

Simon S Du, Jason D Lee, and Yuandong Tian. When is a convolutional filter easy to learn? arXiv preprint arXiv:1709.06129 , 2017.

Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns one-hidden-layer cnn: Don't be afraid of spurious local minima. In International Conference on Machine Learning , pages 1339-1348. PMLR, 2018.

Mo Zhou, Tianyi Liu, Yan Li, Dachao Lin, Enlu Zhou, and Tuo Zhao. Toward understanding the importance of noise in training neural networks. In International Conference on Machine Learning , pages 7594-7602. PMLR, 2019.

Tianyi Liu, Minshuo Chen, Mo Zhou, Simon S Du, Enlu Zhou, and Tuo Zhao. Towards understanding the importance of shortcut connections in residual networks. Advances in neural information processing systems , 32, 2019.

Weihang Xu and Simon S Du. Over-parameterization exponentially slows down gradient descent for learning a single neuron. arXiv preprint arXiv:2302.10034 , 2023.

Arthur Jacot, Franck Gabriel, and Cl´ ement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems , 32, 2019.

Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks, 2018.

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning , pages 1675-1685. PMLR, 2019.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning , pages 322-332. PMLR, 2019.

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , 1(1):84-105, 2020.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep relu networks. Machine learning , 109:467-492, 2020.

Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems , 31, 2018.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665E7671, 2018.

Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXiv:2001.11443 , 2020.

Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean-field framework for over-parameterized deep neural networks. In Conference on learning theory , pages 1887-1936. PMLR, 2021.

Yiping Lu, Chao Ma, Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean field analysis of deep resnet and beyond: Towards provably optimization via overparameterization from depth. In International Conference on Machine Learning , pages 6426-6436. PMLR, 2020.

Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466 , 2022.

Yuandong Tian. Understanding the role of nonlinearity in training dynamics of contrastive learning. arXiv preprint arXiv:2206.01342 , 2022.

Yuandong Tian. Understanding the role of nonlinearity in training dynamics of contrastive learning. ICLR , 2023.

Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems , 33:15383-15393, 2020.

Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ICLR , 2020.

Sam Shleifer, Jason Weston, and Myle Ott. Normformer: Improved transformer pretraining with extra normalization. arXiv preprint arXiv:2110.09456 , 2021.

Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems , 32, 2019.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth´ ee Lacroix, Baptiste Rozi` ere, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. arXiv preprint arXiv:2305.19466 , 2023.

Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409 , 2021.

Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Ian Simon, Curtis Hawthorne, Andrew M Dai, Matthew D Hoffman, Monica Dinculescu, and Douglas Eck. Music transformer. arXiv preprint arXiv:1809.04281 , 2018.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems , 32, 2019.

A Notation Table

Tbl. 1 gives the notation of the main quantities in the paper.

B Detailed comparison with the concurrent works

B.1 Comparison with [28]

Setting, Assumptions and Conclusions . [28] analyzes the SGD convergence of 1-layer ViT model (1 layer self-attention + 2 layer FFN with ReLU, with the top layer of FFN fixed as random, token embedding fixed). Under a specific binary data model in which the data label is determined by counting the number of tokens that belong to positive/negative pattern, [28] gives a generalization bound when the number of hidden nodes in FFN is large, and at the same time, shows that the selfattention attends to relevant tokens and becomes sparse (if number of relevant tokens are small).

In comparison, our work focuses on language models, assume broader data distribution (e.g., multiple classes, arbitrary conditional probability of token given class label) and incorporate LayerNorm

naturally. We propose more detailed quantitative properties, e.g., attention sparsity even among relevant tokens, two-stage evolution of attention scores, with a much simpler analysis.

Techniques . The techniques used in [28] are based on feature learning techniques applied to MLP (e.g., [87]). It identifies lucky neurons if the number of hidden neurons is large enough. In comparison, our framework and analysis is much simpler by leveraging that certain nonlinear continuous dynamics systems can be integrated out analytically to yield clean solutions (e.g., Theorem 3 (Eqn. 11) and Theorem 4 (Eqn. 128)), avoiding complicated bounds in [28]. This allows us to characterize the converging behavior of self-attentions when t → + ∞ .

B.2 Comparison with [29]

[29] focuses on 1-layer attention-based prompt-tuning, in which some parameters of the models are fixed ( W p , W q ). The analysis focuses on the initial (3x one-step) SGD trajectory, and constructs the dataset model containing specific context-relevant/context-irrelevant data, and the context-vector indicates the token relevance. As a result, [29] shows the attention becomes sparse (i.e., attending to context-relevant tokens) over time, which is consistent with ours, and shows that prompt-attention can find the relevant tokens and achieve high accuracy while self-attention/linear-attention can't.

In comparison, our work goes beyond the 2-classes model and further points out that the attention weight will be relevant to the conditional probability of the contextual tokens, which is more detailed than the sparse attention result in [29] that relies on the sparsity assumption of contextual tokens itself. We also focus on the pre-training stage (training from scratch, predicting the next token), characterize the entire trajectory under SGD for the self-attention layer, in particular its converging behavior.

B.3 Comparison with [30]

Compared to [29], [30] also analyzes the dynamics of the query-key matrix and the embedding of a single tunable token (often [CLS] token). It makes connection between the binary classification problem with 1-layer transformer and max-margin SVM formulation, when the tokens are linearly separable. The dynamics is characterized completely, which is nice. Note here is not an attention since its norm can be shown to go to infinity over training.

In comparison, our work does not learn the embedding of an individual token, but focuses on the dynamics of (all-pair) attention scores during training. We also work on multiple-class setup and do not explicitly assume the linear separability among classes.

C Proof of Section 3

Lemma 1 (Dynamics of 1-layer Transformer) . The gradient dynamics of Eqn. 2 with batchsize 1 is:

Here P ⊥ v := I -vv ⊤ / ∥ v ∥ 2 2 projects a vector into v 's orthogonal complementary space, η Y and η Z are the learning rates for the decoder layer Y and self-attention layer Z , α := [ α 1 , . . . , α M ] ⊤ ∈ R M and α m := exp( Y ⊤ LN( X ⊤ b T )) / 1 ⊤ exp( Y ⊤ LN( X ⊤ b T )) .

Proof. With the reparameterization of Y and Z , the loss function is the following:

and

Therefore, taking matrix differentials, we have:

since in general we have d(exp( a ) / 1 ⊤ exp( a )) = L d a with L := diag( b ) -bb ⊤ , let a := XZ ⊤ x T and we have:

Finally notice that P ⊥ X ⊤ b X ⊤ L = P ⊥ X ⊤ b X ⊤ diag( b ) due to the fact that P ⊥ v v = 0 and the conclusion follows.

Lemma 2. Given the event { x T = m,x T +1 = n } , when T → + ∞ , we have

where c m,n = [ c 1 | m,n , c 2 | m,n , . . . , c M | m,n ] ⊤ ∈ R M . Note that c ⊤ m,n 1 = 1 .

Proof. Let p = [exp( z m 1 ) , . . . , exp( z mM )] ⊤ ∈ R M , p x t := exp( z mx t ) , and p X := [exp( z mx 1 ) , . . . , exp( z mx T -1 ) ] ⊤ , then for any T we have

Combining Lemma 18 and the definition of c l | m,n (Eqn. 5), we have that when T → + ∞ ,

Similarly:

Let T → + ∞ , then we also get

D Proof of Section 4

D.1 Notation

For convenience, we introduce the following notations for this section:

Denote E ′ := ( I + E ) -1 -I .

Apply orthogonal diagonalization on E and obtain E = U ⊤ DU where U := [ u 1 , ..., u K ] ∈ O K × K , D = diag ( λ 1 , ..., λ K ) and | λ 1 | ≥ ... ≥ | λ K | ≥ 0 .

Denote F ′ := [ F, F ◦ ] ∈ R M × M where F ◦ ∈ R M × ( M -K ) is some matrix such that rank ( F ′ ) = M . This is possible since { f i } i ∈ [ K ] are linear-independent.

Denote W ′ := ( F ′ ) ⊤ Y = [ F, F ◦ ] ⊤ Y = [ W ⊤ , Y ⊤ F ◦ ] ⊤ = [ w 1 , . . . , w K , w K +1 , . . . , w M ] ⊤ ∈ R M × M .

Denote ζ n := M M -1 ( e n -1 M 1 ) ∈ R M .

Denote q 1 := ζ ⊤ i ζ i = 1 + 1 M -1 , q 0 := ζ ⊤ j ζ i = -M ( M -1) 2 where i, j ∈ [ M ] , i = j .

̸

Denote h to be a continuous function that satisfies h (0) = 0 and ˙ h = η Y · ( M -1 + exp( Mh )) -1 . Details in Lemma 6.

Denote ω 1 to be the constant defined in Lemma 8 that satisfies ω 1 = Θ( ln ln( M ) ln( M ) ) .

Denote N n := ∑ N i =1 I [ x T +1 = n ] to be the number of times the event x T +1 = n happens.

Denote ¯ N := ⌈ N/K ⌉ to be the average value of N n when P ( n ) ≡ 1 /K and ∆ := ⌈ √ N ln( 1 δ ) ⌉ to be the radius of confidence interval centered on ¯ N with confidence 1 -δ . Here ∆ / ¯ N ≍ K √ N √ ln( 1 δ ) ≪ 1 since N ≫ K 2 . Details in Lemma 10 and Remark 4.

Denote ¯ W ′ ( N ) := [ ¯ w 1 ( N ) , ..., ¯ w K ( N ) , 0 , ..., 0 ] ⊤ ∈ R M × M , where ¯ w n ( N ) := ( M -1) h ( ¯ N ) ζ n , ∀ n ∈ [ K ] .

D.2 Proof of Lemma 3

We assume ∪ m ∈ [ M ] ψ -1 ( m ) = [ K ] for convenience, but we claim that our proof can be easily generalized into the case where Ω = [ K ] by reordering the subscript of the vectors. First, we prove the dynamics equation of the reparameterized dynamics of Y .

̸

Lemma 3. Given x T +1 = n , the dynamics of W is (here α j = exp( w j ) / 1 ⊤ exp( w j ) ):

While we cannot run gradient update on W directly, it can be achieved by modifying the gradient of Y to be ˙ Y = η Y ( f n -FE ′ e n )( e n -α n ) ⊤ . If λ 1 is small, the modification is small as well.

Proof. We let F ′ := [ F, F ◦ ] ∈ R M × M where rank ( F ′ ) = M , this is possible since { f n } n ∈ [ K ] are linear-independent. And we further define W ′ := ( F ′ ) ⊤ Y = [ F, F ◦ ] ⊤ Y = [ W ⊤ , Y ⊤ F ◦ ] ⊤ = [ w 1 , . . . , w K , w K +1 , . . . , w M ] ⊤ ∈ R M × M . When given x T +1 = n , the first term of the differential of loss function J is:

So ˙ W ′ = e n ( x T +1 -α ) ⊤ . This nice property will limit W to independently update its n -th row for any x T +1 = n ∈ [ K ] , and the last M -K rows of W ′ are not updated. Similarly for α we have

We get Eqn. 8 by combining the above results.

If we don't run gradient update on W directly, we can run a modified gradient update on Y :

This will lead to (note that F does not change over time due to Assumption 1 (c)):

By Lemma 17, we know that if λ 1 is small, so does max i ∈ [ K ] | λ i ( E ′ ) | and thus the modification is small as well. In Lemma 5 Remark 1, we will show that the additional term -FE ′ e n effectively reduces the learning rate, if all off-diagonal elements of E are the same.

Lemma 3 shows that we can transfer the problem into solving K independent and similar non-linear ODE. And we then show that such a problem can be well solved by following Lemma. Recall that ζ n := M M -1 ( e n -1 M 1 ) ∈ R M , we have:

Lemma 5. Assume Y is initialized to be a zero matrix, Z is fixed, and the learning rate of Y is η Y . Then if event x T +1 = n always holds at s step ( s ≥ 1 ) we have

̸

And thus e n -α n ( s ) = M -1 M -1+exp( Mh ∗ ( s -1)) ζ n . Here h ∗ ( s ) satisfies:

Proof. We prove this Lemma by induction.

Step 1 : Note that Y is initialized to be a zero matrix, then w i (0) = 0 , ∀ i ∈ [ K ] . So we have

̸

̸

It's easy to check that these equations match that of Lemma 5.

Step s : Assume the equations of Lemma 5 hold for step s -1 . Then at the s step, we have

̸

̸

̸

And the equations of Lemma 5 also hold for step s . So we finish the proof.

Remark 1. If we following the original dynamics (Eqn. 7), then it corresponds to the W dynamics as follows:

̸

When all off-diagonal elements of E are identical, i.e., f ⊤ n f n ′ = ρ for n = n ′ , then 0 ≤ ρ ≤ 1 and we have

̸

So if different sequence classes are sampled uniformly, then by similar induction argument, we will have

̸

̸

where the last equation is due to the fact that ∑ n ζ n = M M -1 ∑ n ( e n -1 M 1 ) = M M -1 ( 1 -1 ) = 0 . This means that ∑ n ′ = n ζ n ′ = -ζ n . Therefore, the effective learning rate is η ′ Y := (1 -ρ ) η Y ≤ η Y .

D.3 Property of h ∗ ( s ) and its continuous counterpart.

Before further investigation on Y , we need to get some basic properties of h ∗ , in particular, how fast it grows over time. First, if we consider the continuous version of h ∗ , namely h , then we can directly obtain the equation that h needs to satisfy by integrating the corresponding differential equation.

Lemma 6. If we consider the continuous version of h ∗ ( s ) , namely h , as the following ODE:

and assume h (0) = 0 , then we have

Then we will show that the h is actually almost the same as the original step function h ∗ .

Lemma 7. For h and h ∗ we have:

(a) For any s ∈ N , 0 ≤ h ∗ ( s ) -h ( s ) ≤ 2 η Y M . Then there exists some constant c = Θ(1) such that for any s ≤ ln( M ) /η Y , h ( s + c ) ≥ h ∗ ( s ) ≥ h ( s ) .

(b) h ∗ ( s ) -h ( s ) → 0 when s → + ∞ .

Proof. (a) First we show that h ∗ ( s ) ≥ h ( s ) for all s ∈ N , and the convex packet function of h ∗ can almost control the upper bound of h . Define h ◦ : R + → R + as follows:

Here ⌈·⌉ and ⌊·⌋ mean ceil function and floor function, respectively. It's clear that h ◦ is a strictly monotonically increasing function, and for any s ∈ N , h ◦ ( s ) = h ∗ ( s ) , while for any t / ∈ N , ( t, h ◦ ( t )) lies on the line connecting point ( ⌊ t ⌋ , h ∗ ( ⌊ t ⌋ )) and point ( ⌈ t ⌉ , h ∗ ( ⌈ t ⌉ )) . To prevent ambiguity, we let ˙ h ◦ ( t ) to be the left limit of h ◦ , i.e., ˙ h ◦ ( t ) = lim t ′ → t -˙ h ◦ ( t ′ ) .

We claim h ( t ) ≤ h ◦ ( t ) , ∀ t ∈ R + . We prove it by induction. First when t = 0 , we have h ◦ (0) = h ∗ (0) = h (0) = 0 . Then we assume h ( t ′ ) ≤ h ◦ ( t ′ ) hold for time t ′ ≤ t ∈ N and prove that h ( t ′ ) ≤ h ◦ ( t ′ ) hold for t ′ ∈ ( t, t +1] . If this is not true, then from the continuity of h ◦ and h , we know it must exist t ′′ ∈ ( t, t +1] such that h ( t ′′ ) ≥ h ◦ ( t ′′ ) and ˙ h ( t ′′ ) > ˙ h ◦ ( t ′′ ) . The later condition results that η Y [ M -1 + exp( Mh ( t ′′ ))] -1 > η Y [ M -1 + exp( Mh ∗ ( ⌊ t ′′ ⌋ ))] -1 . So

This contradicts the hypothesis h ( t ′′ ) ≥ h ◦ ( t ′′ ) . So h ( t ′ ) ≤ h ◦ ( t ′ ) hold for t ′ ∈ ( t, t +1] and thus for all t ∈ R + . Hence for any s ∈ N , we have h ( s ) ≤ h ◦ ( s ) = h ∗ ( s ) . Actually, we can use the similar method to prove that h ( s ) < h ∗ ( s ) for any s ∈ N + .

Then we show h ∗ ( s ) -h ( s ) ≤ 2 η Y /M by proving that for any s ∈ N + , h ( s ) must meet at least one of the following two conditions:

(i) h ( s ) ∈ [ h ∗ ( s -1) , h ∗ ( s )] .

If (i) doesn't hold, then we have for any t ∈ [ s -1 , s ) , h ( t ) ≤ h ( s ) < h ∗ ( s -1) = h ◦ ( s -1) , which results that ˙ h ( t ) > ˙ h ◦ ( t ) for all t ∈ [ s -1 , s ) . Therefore, h ∗ ( s ) -h ∗ ( s -1) = h ◦ ( s ) -h ◦ ( s -1) < h ( s ) -h ( s -1) and thus h ( s ) meets condition (ii). It's clear that h (0) and h (1) meet (i).

0

10

20

30

40

50

60

70

80

0.000

0.025

0.050

0.075

0.100

0.125

0.150

0.175

h

/

h

*

Z

=0.05,

M

=50

Y

=0.1

Y

=0.1

Y

=0.2

Y

=0.2

Y

=0.5

Y

=0.5

Y

=1.0

Y

=1.0

Y

=2.0

Y

=2.0

0

10

20

30

40

50

60

70

0.000

0.002

0.004

0.006

0.008

0.010

0.012

h

/

h

*

Z

=0.1,

M

=1000

Y

=0.2

Y

=0.2

Y

=0.5

Y

=0.5

Y

=1.0

Y

=1.0

Y

=2.0

Y

=2.0

Y

=4.0

Y

=4.0

t

t

These two conditions mean that the gap between h ∗ and h will not grow if h ( s ) is smaller than h ∗ ( s -1) . Then for all h ( s ) that meet (i), we have h ∗ ( s ) -h ( s ) ≤ h ∗ ( s ) -h ∗ ( s -1) ≤ h ∗ (1) -h ∗ (0) = η Y /M from Eqn. 33. And for any s ≥ 2 , every time h ( s ) transfer from (i) to (ii) exactly at s , which means that h ( s -1) meets (i) and thus no smaller than h ∗ ( s -2) , we get h ∗ ( s ) -h ( s ) ≤ h ∗ ( s ) -h ( s -1) ≤ h ∗ ( s ) -h ∗ ( s -2) ≤ h ∗ (2) -h ∗ (0) ≤ 2 η Y /M .

Finally from Eqn. 53 in Lemma 9, when s ≤ ln η M Y , we get h ( s ) = Θ( η Y t/M ) and thus there exist some constant c = Θ(1) such that h ( s + c ) ≥ h ( s ) + 2 η Y /M ≥ h ∗ ( s ) ≥ h ( s ) .

(b) Assume that there exist ϵ ∈ (0 , 2 η Y /M ] such that h ∗ ( s ) -h ( s ) ≥ ϵ for all s ∈ N . Since h is unbounded, then ˙ h ( t ) → 0 when t →∞ from Eqn. 33, so there exist some s ′ 0 ∈ N such that when s ≥ s ′ 0 , h ( s +1) -h ( s ) ≤ ϵ +ln(1 / 2) /M . Also, from Lemma 9 we know that exists s ′′ 0 = (3+ δ ) ln( M ) η Y where δ > 0 , δ = Θ(1) such that when s ≥ s ′′ 0 , exp( Mh ( s )) > 2( M -1) . Since s →∞ , we just consider the case that s = ⌊ t ⌋ ≥ s 0 := max( s ′ 0 , s ′′ 0 ) . Then denote ∆ 1 := 2( M -1) exp( Mh ( s )) < 1 , we have:

Here C = exp( Mϵ ) 4 exp(4 η Y ) M > 0 and for the last inequality, we use the fact that t ≥ s ′ 0 > 3 ln M η Y and thus h ( s ) ≤ h ( t ) = O ( ln( Mη Y t ) M ) from Lemma 9. So we get

This contradicts h ◦ ( t ) -h ( t ) ≥ 0 ! So the original assumption doesn't hold, which means that h ∗ ( s ) -h ( s ) → 0 when s →∞ .

Remark 2. By some qualitative estimation, we claim that if η Y = O (1) , then there exists some constant c = O (ln M ) such that h ( s ) ≤ h ∗ ( s ) ≤ h ( s + c ) for all s > s 1 := 2 ln(1+ ω 1 ) η Y where ω 1 = Θ(lnln M/ ln M ) is defined in Lemma 8. Denote δh ( t ) := h ◦ ( t ) -h ( t ) , when δh ( t ) ≪ h ( t ) ,

we have ˙ δh ( t ) = ˙ h ◦ ( t ) -˙ h ( t ) ≍ -η Y M · δh ( t ) · exp( -Mh ( t )) ≍ -δh ( t ) /t by computing the second-order derivative of δh , and thus h ◦ ( t ) -h ( t ) ≍ 2 η Y s 0 / ( Mt ) = O (ln M/ ( Mt )) . Combining this with the fact that h ( t ) = Θ(ln( Mη Y t ) /M ) when t > s 1 , we prove our claim. The results of Lemma 7 and Remark 2 are also confirmed by the numerical simulation results as Fig. 8.

So from Lemma 7 and Remark 2, we just assume η Y < 1 and replace h ∗ with h in the latter parts for convenience. Then we further investigate the properties of Eqn. 45.

Lemma 8. There exists ω i , 0 < ω i ≪ 1 , i = 2 , 3 , such that for h ∈ J 1 := [ 1 M 2 -ω 0 , (1+ ω 1 ) ln( M ) M ] , we have exp( Mh ( t )) ≤ ( M -1) Mh ( t ) . And for h / ∈ J 1 , we have exp( Mh ( t )) > ( M -1) Mh ( t ) . Here ω 1 = Θ( ln ln( M ) ln( M ) ) , and if M ≫ 100 , we have ω 0 ≲ ( 1 M 0 . 99 ln M ) ≪ 0 . 01 .

Proof. It's obvious that exp( Mh ( t )) -( M -1) Mh ( t ) has two zero points in R + . Let h ( t ) = M -(2 -ω 0 ) , we get

For another zero point, let ω 1 ∈ (0 , 1) to be some constant such that h ( t ) = (1+ ω 1 ) ln( M ) M satisfies exp( Mh ) = ( M -1) Mh , then we get

where c ′ ∈ (0 . 5 , 2) is some universal constant.

Remark 3. From Lemma 8, if we assume M ≫ 100 , then ω 0 ≪ 0 . 01 , and if we assume η Y ≫ 1 M 1 -ω 0 > 1 M 0 . 99 , then h (1) ≳ η Y M ≫ 1 M 2 -ω 0 and function exp( Mh ( t )) -( M -1) Mh ( t ) has only one zero point (1+ ω 1 ) ln M M in [1 , ∞ ) . For convenience, we just assume M ≫ 100 and 1 > η Y ≫ 1 M 0 . 99 and thus focus on the unique zero point (1+ ω 1 ) ln M M of h in the latter parts.

Wecan then show the properties of speed control coefficient γ ( t ) := ( M -1) 2 h ( t/K ) ( M -1)+exp( Mh ( t/K ))

Lemma 9. We have two stage for h and γ

as below. :

When t ≤ K ln( M ) η Y , we have exp( Mh ( t/K )) ≤ min( M -1 , ( M -1) Mh ( t/K )) , h = O ( η Y t/ ( MK )) and γ ( t ) = O ( η Y t/K ) .

When t ≥ 2(1+ ω 1 ) K ln( M ) η Y where ω 1 = Θ( ln ln M ln M ) is defined in Lemma 8, we have exp( Mh ( t/K )) ≥ max( M -1 , ( M -1) Mh ( t/K )) , h = O ( 1 M ln( Mη Y t/K )) and γ ( t ) = O ( K ln( Mη Y t/K ) η Y t ) .

̸

Proof. For convenience, we just let K = 1 . And the proof for K = 1 is similar. We denote ∆ 1 ( h ) := exp( Mh ) M -1 and ∆ 2 ( h ) := exp( Mh ) ( M -1) Mh .

Step 1 : t ≤ ln( M ) η Y . If h ≥ ln( M -1) M , from Eqn. 45 we have:

So when t ≤ ln( M ) η Y we have h < ln( M -1) M , and thus exp( Mh ( t )) ≤ min( M -1 , ( M -1) Mh ( t )) , i.e., ∆ 1 , ∆ 2 ≤ 1 . Then from Eqn. 45 we get

Step 2 : t > 2(1+ ω 1 ) ln( M ) η Y where ω 1 = Θ( ln ln( M ) ln( M ) ) . So now h > ln( M -1) M and thus ∆ 1 > 1 from Eqn. 52. Then if exp( Mh ) ≤ M ( M -1) h , i.e. ∆ 2 ≤ 1 , from Lemma 8 we have h = Mη Y t +1 (1+∆ 2 ) M ( M -1) ≤ (1+ ω 1 ) ln( M ) M . Therefore,

Contradiction! So when t ≥ 2(1+ ω 1 ) ln( M ) η Y , we have ∆ 2 > 1 . Then from Eqn. 45 we get:

D.4 The dynamics under multiple uniformly sampled sequence classes

We then generalize our analysis of W to the case where x T +1 can be any value in [ K ] rather than fixing x T +1 = n with the key observation that the row vectors of W ′ can be independently updated. Before formalizing this result, we first conduct the concentration inequality of the sampling number for each next-token case. Let N n := ∑ N i =1 I [ x T +1 = n ] to be the number of times the event x T +1 = n happens, then we have:

Lemma 10. For δ ∈ (0 , 1) , with probability at least 1 -δ we have

Proof. From Hoeffding's inequality, we have

Let t = √ 1 2 N ln( 2 δ ) and we can get the results by direct calculation.

Remark 4. From Lemma 10, if we consider the uniform sampling case where P ( n ) ≡ 1 K , then N P ( n ) = N/K ≫ √ N . So N n are all concentrated around N P ( n ) . Recall the definition of ¯ N = ⌈ N/K ⌉ and ∆ = ⌈ √ N ln( 1 δ ) ⌉ , with probability at least 1 -δ we have:

We then further investigate the concentration of h ( N n ) :

Lemma 11. For δ ∈ (0 , 1) , with probability at least 1 -δ we have

where σ ′ > 0 is some constant such that σ ′ ≤ 1 3 η Y ∆ ≪ ln( M ) . And if N ≥ 2 K (1+ ω 1 ) ln M η Y where ω 1 is defined in Lemma 8, then σ ′ ≲ ∆ ¯ N ≪ 1 .

Proof. First, we note that h has a decreasing gradient, so h ( x ) ≥ ˙ h ( x ) × x and h ( x 1 + x 2 ) -h ( x 1 ) ≤ ˙ h ( x 1 ) × x 2 for any x 1 , x 2 ≥ 0 . So with probability at least 1 -δ , we have:

For the second inequality, without loss of generality, we let N n > ¯ N . Denote g ( s ) := ( M -1 + exp( Mh ( s ))) -1 and note that:

the last equality holds only when h ( s ) = ln( M -1) M . So from | g ( ¯ N + ∆) -g ( N n ) | ≤ max s ∈ [ N n ,N n +∆] ˙ g ( s ) · ∆ , we get:

If ¯ N < 2(1+ ω 1 ) ln( M ) η Y + ∆ with ω 1 = Θ( ln ln M ln M ) defined in Lemma 8, we have σ ′ ≤ η Y ∆ / 3 ≪ η Y ¯ N ≲ ln( M ) . If ¯ N ≥ 2(1+ ω 1 ) ln( M ) η Y +∆ , we utilize the Eqn.45 and obtain:

So σ ′ ≤ ∆ / ¯ N . When N n < ¯ N , with probability at least 1 -δ we have N n ≳ ¯ N -∆ , and similar inequalities also hold for such cases, so we finish the proof.

̸

Recall that ζ n ∈ R M is defined as ζ n = M M -1 ( e n -1 M 1 ) . And we have q 1 := ζ ⊤ i ζ i = 1 + 1 M -1 , q 0 := ζ ⊤ j ζ i = -M ( M -1) 2 for all i, j ∈ [ M ] where i = j . For convenience, we denote ¯ W ′ ( N ) := [ ¯ w 1 ( N ) , ..., ¯ w K ( N ) , 0 , ..., 0 ] ⊤ ∈ R M × M , where ¯ w n ( N ) := ( M -1) h ( ⌈ N/K ⌉ ) ζ n = ( M -1) h ( ¯ N ) ζ n . So using these concentration inequalities, we get:

Lemma 12. Assume the assumptions in Lemma 5 hold but we uniformly sample the training data. Then if the total number of epochs N satisfies N ≫ K 2 , we have Y = ( F ′ ) -⊤ ( I + Θ ′ ) ¯ W ′ ( N ) where Θ ′ := diag ( θ 1 , . . . , θ K , 0 , . . . , 0) ∈ R M × M and with probability at least 1 -δ we have | θ i | ≲ K √ N √ ln( K δ ) , ∀ i ∈ [ K ] .

Proof. From Lemma 5 and the first inequality of Lemma 11, we know that

where for any δ ∈ (0 , 1) , with probability at least 1 -δ we have | θ i | ≲ K √ N √ ln( K δ ) , ∀ n ∈ [ K ] . Therefore, W ′ ( N ) = [ w 1 ( N ) , . . . , w K ( N ) , 0 , . . . , 0 ] ⊤ = ( I + Θ ′ ) ¯ W ′ ( N ) , then from W ′ = ( F ′ ) ⊤ Y , we finish the proof.

Then, we can give out the exact solution of Y by pointing out the properties of F ◦ and F ′ from the observation that each row of Y should be the linear combination of vectors in { f ⊤ n } n ∈ [ K ] :

Theorem 5. If Assumption 2 holds and Y (0) = 0 . Furthermore, we assume the training data is uniformly sampled and the total number of epochs N satisfies N ≫ K 2 . Then the solution of Eqn. 26 will be:

Here Θ := diag ( θ 1 , . . . , θ K ) and for any δ ∈ (0 , 1) , with probability at least 1 -δ we have | θ i | ≲ K √ N √ ln( K δ ) , ∀ i ∈ [ K ] .

Proof. Let q i , i ∈ [ M ] be the i -th row vector of ( F ′ ) -1 , then we have q ⊤ j f i = I [ i = j ] . From Lemma 12 we get Y = ( F ′ ) -⊤ ( I +Θ ′ ) ¯ W ′ ( N ) . And from Eqn. 26, we know all the columns of Y are the linear combination of f n , n ∈ [ K ] . Note that ¯ W ( N ) has only top K rows to be non-zero, so we need to constrain that all the top K columns of ( F ′ ) -⊤ , i.e., q i , i ∈ [ K ] , to be the linear combination of f n , n ∈ [ K ] , which means that q 1 , . . . , q K must be the basis of Ξ := span ( f j ; j ∈ [ K ]) and thus q K +1 , . . . , q M are the basis of Ξ ′ := span ( f j ; K ≤ j ≤ M ) . Therefore, we get Ξ ⊥ Ξ ′ , and thus [ q 1 , . . . , q K ] can only be ( F † ) ⊤ . So the proof is done.

Actually, we see that the result of Theorem 5 matches the modified gradient update on Y (Eqn. 26). And we show that using such reparameterization dynamics, we can still approach the critical point of Eqn. 7 in the rate of O ( 1 N ) :

Corollary 1. Assume assumptions in Theorem 5 hold, M ≫ 100 and η Y satisfies M -0 . 99 ≪ η Y < 1 . Then ∀ n ∈ [ K ] , we have

where σ > -1 and for any δ ∈ (0 , 1) , with probability at least 1 -δ we have | σ | ≲ η Y √ N ln( 1 δ ) , and when N ≫ K ( √ N ln( 1 δ ) + 2(1+ ω 1 ) ln M η Y ) with ω 1 defined in Lemma 8, | σ | ≲ K √ N √ ln( 1 δ ) . Further, to let ∥ x T +1 -α n ∥ 2 ≤ ϵ with probability at least 1 -δ for any n ∈ [ K ] and ϵ ≪ 1 , we need the total number of training epochs to be at most O ( K ϵη Y log( M ϵ )) .

Proof. Note that x T +1 = e n , then we just need to combine Lemma 5 and the second inequality of Lemma 11, to get Eqn. 71. Denote S n to be the number of training epochs that are needed to let ∥ x T +1 -α n ∥ 2 ≍ ϵ , then we have

But note that h ( t +1) -h ( t ) ≥ η Y M -1+exp( Mh ( S n )) ≍ η Y ϵ M -1 , ∀ t ∈ [0 , S -1] from Eqn. 71, we have

Note that ϵ ≪ 1 and we have N ≫ K 2 , then we have S = ∑ n S n ≲ K ϵη Y ln( M ϵ ) .

D.5 Proof of Theorem 1

Finally, we turn to prove Theorem 1. Obviously, all the diagonal elements of E are zero and all the off-diagonal elements of E are non-negative since c l | m,n ≥ 0 . Note that E is a real symmetric matrix, then it can be orthogonal diagonalization by E = U ⊤ DU where U := [ u 1 , ..., u K ] ∈ O K × K , D = diag ( λ 1 , ..., λ K ) and | λ 1 | ≥ ... ≥ | λ K | ≥ 0 . Then we can get the following properties of E and E ′ :

Lemma 13. max i,j ∈ [ K ] ( | E ij | ) ≤ | λ 1 | .

Proof. We have:

Lemma 14. If E ∈ R K satisfies | λ 1 | ≤ λ < 1 , then ( I + E ) is invertible and ( I + E ) -1 = I -E ′ ,where E ′ satisfies E ′ = U ⊤ D ′ U and D ′ = diag ( λ ′ 1 , ..., λ ′ K ) and λ ′ i = λ i 1+ λ i , ∀ i ∈ [ K ] .

Proof. Since U is orthonormal and | λ i | ≤ λ < 1 , we have E n = U ⊤ D n U → O . Then from the property of the Neumann series, we get I + E is invertible and

Here we define D ′ = diag ( λ ′ 1 , ..., λ ′ K ) and use the fact that ∑ ∞ n =1 ( -λ i ) n = -λ i 1+ λ i

Lemma 15. If | λ 1 | ≤ λ < 1 , then max i ∈ [ K ] | λ i ( E ′ ) | ≤ 1 1 -λ | λ 1 | ≤ λ 1 -λ .

Proof. We have

̸

Lemma 16. Assume that Assumption 2 holds, then all the diagonal elements of E ′ are nonpositive,i.e., E ′ ii ≤ 0 , ∀ i ∈ [ K ] . Further, if there exist any k = i ∈ [ K ] such that E ki > 0 , then E ′ ii < 0 .

Proof. Note that E ii = ∑ K k =1 λ k u 2 ik = 0 (here u ik is the k -th component of eigenvector u i ) and | λ k | < 1 , we have

When E ′ ii = 0 , then λ := ( λ 1 , . . . , λ K ) must don't have overlapping entries with respect to u i , which results that E ij := ∑ K k =1 λ k u ik u jk = 0 holds for any j ∈ [ K ] . So we prove the results.

Proof. From Lemma 14 we have:

Finally we can prove our main theorem in Sec. 4.

Theorem 1. If Assumption 2 holds, the initial condition Y (0) = 0 , M ≫ 100 , η Y satisfies M -0 . 99 ≪ η Y < 1 , and each sequence class appears uniformly during training, then after t ≫ K 2 steps of batch size 1 update, given event x T +1 [ i ] = n , the backpropagated gradient g [ i ] := Y ( x T +1 [ i ] -α [ i ]) takes the following form:

̸

Here the coefficients ι n ( t ) , β nn ′ ( t ) and γ ( t ) are defined in Appendix with the following properties:

̸

Proof. Note that if Assumption 2 holds, then F † = ( I -E ′ ) F ⊤ . Recall q 1 := 1 + 1 M -1 ≈ 1 and q 0 := -M ( M -1) 2 ≈ 0 . Then given x T +1 [ i ] = n , we get:

̸

̸

̸

where

Here σ is defined in Cor. 1 and satisfies -1 < σ ≪ ln M . | δ 1 | ≲ K √ N √ ln( 1 δ ) + 1 M ≪ 1 and | δ 2 | , | δ 3 | ≤ M ( M -1) 2 × 2(1 + 3 | δ 1 | ) < 3 M . Here we use the fact that | θ | , | θ i | ≲ K √ N √ ln( 1 δ ) , ∑ k ∈ [ K ] λ k u jk u jn ′ = E kn ′ and the fact from Lemma 15:

(a) Now let's prove that ξ n ( t ) > 0 . First from ( I + E )( I -E ′ ) = I we have E -E ′ -EE ′ = O . Then use the symmetry of E and E ′ , we get

̸

̸

Note that F ⊤ F = I + E , we have E ′ = f ⊤ f ′ , ∀ n ′ = n and E = 0

̸

̸

Note that | λ i ( E ) | > 0 , ∀ i ∈ [ K ] in Assumption 2 implies that E ki > 0 holds for some k = i ∈ [ K ] . Then from (1) of Lemma 16 we get ∑ k = n E ′ nn ′ f ⊤ n f n ′ > 0 .

From Theorem 1 we have β nn ′ = (1 + σ )[ E ′ nn ′ · (1 + δ 1 ) + δ 3 ] . Note that 0 < 1 + σ ≪ ln( M ) , we have:

̸

̸

̸

(b) We directly use Lemma 9, then we finish the proof.

E Proof of Section 5

Lemma 4 (Self-attention dynamics) . With Assumption 1(b) (i.e., T → + ∞ ), Eqn. 4 becomes:

̸

Proof. Taking long sequence limit ( T → + ∞ ), and summing over all possible choices of next token x T +1 = n , plugging in the backpropagated gradient (Eqn. 9) into the dynamics of Z with query token m (Eqn. 4), we arrive at the following:

̸

̸

̸

̸

Theorem 2 (Fates of contextual tokens) . Let G CT be the set of common tokens (CT), and G DT ( n ) be the set of distinct tokens (DT) that belong to next token n . Then if Assumption 2 holds, under the self-attention dynamics (Eqn. 10), we have:

(a) for any distinct token l ∈ G DT ( n ) , ˙ z ml > 0 where m = ψ ( n ) ;

(b) if | G CT | = 1 and at least one next token n ∈ ψ -1 ( m ) has at least one distinct token, then for the single common token l ∈ G CT , ˙ z ml < 0 .

Proof. For any token l , we have:

̸

̸

Distinct token . For a token l distinct to n , by definition, for any n ′ = n , P ( l | m,n ′ ) = 0 and f n ′ l ( t ) ∝ P ( l | m,n ′ ) exp( z ml ) ≡ 0 . Therefore, we have:

̸

Note that ˙ z ml > 0 is achieved by ξ n > 0 from Theorem 1.

̸

Common token . For any query token m , consider n ∈ ψ -1 ( m ) and n ′ = n . if n and n ′ does not overlap then diag( f n )( f n f ⊤ n -I ) f n ′ = -diag( f n ) f n ′ = 0 . When n and n ′ overlaps, let G CT ( n, n ′ ) := { l : P ( l | n ) P ( l | n ′ ) > 0 } be the subset of common tokens shared between n and n ′ , since | G CT | = 1 and ∅ ̸ = G CT ( n, n ′ ) ⊆ G CT := ⋃ n = n ′ G CT ( n, n ′ ) , we have | G CT ( n, n ′ ) | = 1 and l ∈ G CT ( n, n ′ ) , i.e., the common token l is the unique overlap. Then we have:

So we have:

̸

Since ξ n ( t ) > 0 , the only condition that will lead to ˙ z ml = 0 is f 2 nl = 1 . However, since at least one such n has another distinct token l ′ , and thus f nl ′ > 0 , by normalization condition, f nl < 1 and thus ˙ z ml < 0 .

Note that for multiple common tokens, things can be quite involved. Here we prove a case when the symmetric condition holds.

Corollary 2 (Multiple CTs, symmetric case) . If Assumption 2 holds and assume

(1) Single query token m 0 . For any next token n ∈ [ K ] , ψ ( n ) = m 0 .

̸

(2) Symmetry . For any two next tokens n = n ′ , there exists a one-to-one mapping ϕ that maps token l ∈ G DT ( n ) to l ′ ∈ G DT ( n ′ ) so that P ( l | n ) = P ( ϕ ( l ) | n ′ ) ;

(3) Global common tokens with shared conditional probability : i.e., the global common token set G CT satisfies the following condition: for any l ∈ G CT , P ( l | n ) = ρ l , which is independent of next token n ;

(4) The initial condition Z (0) = 0 .

Then for any common token l ∈ G CT , ˙ z m 0 ,l < 0 .

Proof. Since there is a global query token m 0 , we omit the subscript m 0 and let z l := z m 0 ,l .

We want to prove the following induction hypothesis : for any t (a) z l ( t ) = z ϕ m ( l ) ( t ) for n , n ′ which are next tokens that the distinct token l (and l ′ ) belongs to, and (b) the normalization term o 2 n ( t ) := ∑ l ˜ c 2 l | n ( t ) = o 2 ( t ) , i.e., it does not depend on n .

̸

We prove by induction on infinitesimal steps δt . First when t = 0 , both conditions hold due to the initial condition Z (0) = 0 . Then we assume that both conditions hold for time t , then by symmetry, we know that for any n 1 and any distinct token l 1 ∈ G DT ( n 1 ) :

̸

̸

where l 2 = ϕ ( l 1 ) is the image of the distinct token l 1 . This is because (1) f ⊤ n 1 f n ′ = ∑ l ∈ G CT ρ 2 l exp(2 z l ( t )) /o 2 ( t ) is independent of n 1 and n ′ by inductive hypothesis, therefore, β is also independent of its subscripts. And (2) f 2 n 1 l 1 := ˜ c 2 l 1 | n 1 /o 2 ( t ) = ˜ c 2 l 2 | n 2 /o 2 ( t ) = f 2 n 2 l 2 .

Therefore, ˙ z l 1 ( t ) = ˙ z l 2 ( t ) , which means that z l 1 ( t ′ ) = z l 2 ( t ′ ) for t ′ = t + δt .

Let G CT ( n 1 , n 2 ) := { l : P ( l | n 1 ) P ( l | n 2 ) > 0 } be the subset of common tokens shared between n 1 and n 2 , then for their associated n 1 and n 2 , obviously G CT ( n 1 , n 2 ) ⊆ G CT and we have:

So we prove the induction hypothesis holds for t ′ = t + δt . Let δt → 0 and we prove it for all t .

̸

Now we check the dynamics of common token l ∈ G CT . First we have for any n = n ′ , f 2 nl ( t ) = ˜ c 2 l | n ( t ) /o 2 ( t ) = ρ 2 l exp(2 z l ( t )) /o 2 ( t ) = ˜ c 2 l | n ′ ( t ) /o 2 ( t ) = f 2 n ′ l ( t ) and thus f nl ( t ) = f n ′ l ( t ) := f l ( t ) > 0 , therefore:

On the other hand, from the proof on induction hypothesis, we know all off-diagonal elements of E are the same and are positive. Then all all the off-diagonal elements of E ′ are also the same and are positive. Following Theorem 1, we know β nn ′ > 0 and is independent of the subscripts. Therefore, ˙ z l < 0 .

Theorem 3 (Growth of distinct tokens) . For a next token n and its two distinct tokens l and l ′ , the dynamics of the relative gain r l/l ′ | n ( t ) := f 2 nl ( t ) /f 2 nl ′ ( t ) -1 = ˜ c 2 l | n ( t ) / ˜ c 2 l ′ | n ( t ) -1 has the following analytic form (here the query token m = ψ ( n ) and is uniquely determined by distinct token l ):

̸

where χ l ( t ) := e 2( z ml ( t ) -z ml (0)) is the growth factor of distinct token l . If there exist a dominant token l 0 such that the initial condition satisfies r l 0 /l | n (0) > 0 for all its distinct token l = l 0 , and all of its common tokens l satisfy ˙ z ml < 0 . Then both z ml 0 ( t ) and f nl 0 ( t ) are monotonously increasing over t , and

here B n ( t ) := η Z ∫ t 0 ξ n ( t ′ )d t ′ . Intuitively, larger B n gives larger r l 0 /l | n and sparser attention map.

Proof. Let m = ψ ( n ) be the query token associated with next token n . For brievity, we omit subscript m in the proof and let z l := z ml .

First of all, for tokens l and l ′ that are both distinct for a specific next token n , from Eqn. 98, it is clear that

This means that

Integrate both side over time t and we get:

From this we can get the close-form relationship between r l/l ′ | n ( t ) and z l ( t ) :

Now let l 0 be the dominating distinct token that satisfies r l 0 /l | n (0) > 0 for any distinct token l , then

we have ˙ z l 0 > 0 due to Theorem 2.

̸

for any token l ′ = l 0 that is distinct to n , we have:

for any common token l ′ , since ˙ z l ′ < 0 , we have

Therefore, we have:

̸

Therefore, f 2 nl 0 ( t ) is monotonously increasing. Combined with the fact f 2 nl 0 ( t ) ≤ 1 due to normalization condition ∥ f n ∥ 2 = 1 , we have:

Integrate over time and we have:

where B ( t ) := η Z ∫ t 0 ξ n ( t ′ )d t ′ . Plugging that into Eqn. 111, and we have:

F Estimation in Sec. 6

Theorem 4 (Phase Transition in Training) . If the dynamics of the single common token z ml satisfies ˙ z ml = -C -1 0 η Z γ ( t ) e 4 z ml and ξ n ( t ) = C -1 0 γ ( t ) e 4 z ml , then we have:

As a result, there exists a phase transition during training:

Attention scanning . At the beginning of the training, γ ( t ) = O ( η Y t/K ) and B n ( t ) ≈ 1 4 ln K -1 ( ρ 4 0 +2 η Y η Z t 2 ) = O (ln t ) . This means that the growth factor for dominant token l 0 is (sub-)linear: χ l 0 ( t ) ≥ e 2 f 2 nl 0 (0) B n ( t ) ≈ [ K -1 ( ρ 4 0 + 2 η Y η Z t 2 )] 0 . 5 f 2 nl 0 (0) , and the attention on less co-occurred token drops gradually.

Attention snapping . When t ≥ t 0 := 2(1 + δ ′ ) K ln M/η Y with δ ′ = Θ( ln ln M ln M ) , γ ( t ) = O ( K ln( η Y t/K ) η Y t ) and B n ( t ) = O (ln ln t ) . Therefore, while B n ( t ) still grows to infinite, the growth factor χ l 0 ( t ) = O (ln t ) grows at a much slower logarithmic rate.

Proof. Since every next token n shares the same query token m , we omit the subscript m and let z l := z ml .

We start from the two following assumptions:

Given that, we can derive the dynamics of z l ( t ) and ξ n ( t ) :

Let Γ( t ) := η Z ∫ t 0 γ ( t ′ )d t ′ , then Γ(0) = 0 and dΓ( t ) = η Z γ ( t )d t . Therefore, we have

and thus B n ( t ) := η Z ∫ t 0 ξ n ( t ′ )d t ′ can be integrated analytically, regardless of the specific form of γ ( t ) :

Recall that γ ( t ) = ( M -1) 2 h ( t/K ) M -1+exp( Mh ( t/K )) (Theorem 1). Note that h (if treated in continuous time step) is strictly monotonically increasing and satisfies h (0) = 0 , d h ( t/K ) = η Y ( M -1 + exp( Mh ( t/K ))) -1 d t/K (Lemma 6 and Lemma 7), we can let γ ( h ) := ( M -1) 2 h M -1+exp( Mh ) and get:

Therefore, B n ( t ) has a close form with respect to h , regardless of the specific form of h ( t ) :

(1) When t < t ′ 0 := K ln( M ) /η Y , from Lemma 9 we have h ( t/K ) = (1 + o (1)) · η Y t/ ( MK ) . We neglect the o (1) term and denote ν := η Y /η Z , then we have when t ≤ t ′ 0 :

And B n ( t ′ 0 ) = 1 4 ln ( C 0 +2 K ( M -1) 2 M -2 ν -1 ln 2 ( M ) ) .

0

1

2

3

4

5

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

B

(

t

)

=10.0,

M

=100

Z

=0.5,

Y

=5.0

Z

=1.0,

Y

=10.0

Z

=2.0,

Y

=20.0

Z

=4.0,

Y

=40.0

Z

=8.0,

Y

=80.0

0

10

20

30

40

50

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

B

(

t

)

=1.0,

M

=10000

Z

=0.5,

Y

=0.5

Z

=1.0,

Y

=1.0

Z

=2.0,

Y

=2.0

Z

=4.0,

Y

=4.0

Z

=8.0,

Y

=8.0

t

t

(2) Similarly, when t > t 0 := 2(1 + ω 1 ) K ln M/η Y with ω 1 = Θ(lnln M/ ln M ) is defined in Lemma 8, from Lemma 9 we have h ( t/K ) = (1 + o (1)) ln( Mη Y t/K ) /M . We neglect the o (1) term and get when t > t 0 :

From this we know B n ( t 0 ) = 1 4 ln( C 0 + 2 K ( M -1) 2 M -2 ν -1 ln 2 (2(1 + ω 1 ) M ln M )) . It's interesting to find that B n ( t 0 ) just depends on K,M and ν , and thus fixing ν and changing η Z will not influence the value of B n ( t 0 ) , which means that the main difference between B n is arises at the stage t > t 0 . This matches the results in Fig. 9.

Finally, we estimate B n ( t ) when t is large. When ν is fixed and t ≫ ( Mη Y ) -1 exp(1 / √ 2 ν ) , we have

Therefore, from Eqn. 133 we get:

(a) Fix ν , larger η Z result in larger B n ( t ) and sparser attention map.

(b) Fix η Z , larger ν (i.e., larger η Y ) result in smaller B n ( t ) and denser attention map since ln ln( x ) is much slower than ln( x ) .

These match our experimental results in the main paper (Fig. 6).

G Experiments

We use WikiText [25] dataset to verify our theoretical findings. This includes two datasets, WikiText2 and WikiText103. We train both on 1-layer transformer with SGD optimizer. Instead of using reparameterization Y and Z (Sec. 3.1), we choose to keep the original parameterization with token embedding U and train with a unified learning rate η until convergence. Fig. 10 shows that the averaged entropy of the self-attention map evaluated in the validation set indeed drops with when the learning rate η becomes larger.

Note that in the original parameterization, it is not clear how to set η Y and η Z properly and we leave it for future work.

Furthermore, we use the recall-threshold relationship to reshow that smaller η Y /η Z and larger η Z will result in a sparser self-attention map as Fig. 11 and Fig. 12. Here we use some thresholds to retag every entry of the attention as a 0-1 value based on its softmax value for every query token, and then calculate the recall value by the average value of the proportion of the distinct tokens with new labels

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

Average entropy

wikitext2

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

Average entropy

wikitext103

equal to 1 to the total number of distinct tokens. In the figures, the horizontal coordinates denote the threshold values around 1 /M for different last/next tokens setting, and the vertical coordinates denote the recall rate. The dataset is Syn-Medium mentioned in Section 8, and every data point in the figures is the mean value over 10 seeds. It's obvious that a sparser attention map will result in a slower descent rate of the recall-threshold line since sparser attention corresponds to fewer distinct tokens with higher attention weights, and the results of Fig. 11 and Fig. 12 match that of Fig. 6.

H Technical Lemma

Lemma 18. Let h = [ h 1 , h 2 , . . . , h M ] ⊤ ∈ R M is some M -dimensional vector, h X := [ h x 1 , ..., h x T -1 ] ⊤ ∈ R T -1 is a vector selected by input sequence X , then given event x T = m,x T +1 = n , there exists some q m,n = [ q 1 | m,n , q 2 | m,n , . . . , q M | m,n ] ⊤ ∈ R M so that q ≥ 0 and

where q l | m,n satisfies ∑ M l =1 q l | m,n = 1 . And with probability at least 1 -δ we have

And thus q l | m,n → P ( l | m,n ) when T → + ∞ .

Proof. Given that x T = m and x T +1 = n , then we have

And similar equations hold for 1 T -1 X ⊤ diag( h X ) X . Then we consider the case that the previous tokens are generated by conditional probability P ( l | m,n ) as the data generation part, so I [ x t = l ] , ∀ t ∈ [ T -1] are i.i.d. Bernoulli random variables with probability P ( l | m,n ) and Tq l | m,n satisfies binomial distribution. By Hoeffding inequality, we get

Then we get the results by direct calculation.

0.028 0.030 0.032 0.034 0.036 0.038

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

#last/next tokens = 1/2

Y

/

Z

=1

Y

/

Z

=2

Y

/

Z

=5

Y

/

Z

=10

0.028 0.030 0.032 0.034 0.036 0.038

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.028 0.030 0.032 0.034 0.036 0.038

0.30

0.35

0.40

0.45

0.50

0.55

0.60

0.65

0.028 0.030 0.032 0.034 0.036 0.038

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

0.012

0.013

0.014

0.015

0.016

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

#last/next tokens = 3/6

0.012

0.013

0.014

0.015

0.016

0.4

0.5

0.6

0.7

0.8

0.9

0.012

0.013

0.014

0.015

0.016

0.5

0.6

0.7

0.8

0.9

0.012

0.013

0.014

0.015

0.016

0.45

0.50

0.55

0.60

0.65

0.70

0.75

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.4

0.5

0.6

0.7

0.8

0.9

1.0

#last/next tokens = 5/10

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.5

0.6

0.7

0.8

0.9

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

0.90

0.0040

0.0045

0.0050

Z

=0.5

0.4

0.5

0.6

0.7

0.8

0.9

1.0

#last/next tokens = 10/20

0.0040

0.0045

0.0050

Z

=1

0.5

0.6

0.7

0.8

0.9

1.0

0.0040

0.0045

0.0050

Z

=2

0.5

0.6

0.7

0.8

0.9

1.0

0.0040

0.0045

0.0050

Z

=3

0.6

0.7

0.8

0.9

1.0

0.028

0.030

0.032

0.034

0.036

0.038

0.3

0.4

0.5

0.6

0.7

#last/next tokens = 1/2

Z

= 0.5

Z

= 1

Z

= 2

Z

= 3

0.028

0.030

0.032

0.034

0.036

0.038

0.3

0.4

0.5

0.6

0.7

0.8

0.028

0.030

0.032

0.034

0.036

0.038

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.028

0.030

0.032

0.034

0.036

0.038

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

0.012

0.013

0.014

0.015

0.016

0.5

0.6

0.7

0.8

0.9

#last/next tokens = 3/6

0.012

0.013

0.014

0.015

0.016

0.5

0.6

0.7

0.8

0.9

0.012

0.013

0.014

0.015

0.016

0.4

0.5

0.6

0.7

0.8

0.9

0.012

0.013

0.014

0.015

0.016

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.6

0.7

0.8

0.9

#last/next tokens = 5/10

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.5

0.6

0.7

0.8

0.9

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.0075 0.0080 0.0085 0.0090 0.0095 0.0100 0.0105

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.0040

0.0045

0.0050

Y

/

Z

= 1

0.6

0.7

0.8

0.9

1.0

#last/next tokens = 10/20

0.0040

0.0045

0.0050

Y

/

Z

= 2

0.5

0.6

0.7

0.8

0.9

1.0

0.0040

0.0045

0.0050

Y

/

Z

= 5

0.5

0.6

0.7

0.8

0.9

1.0

0.0040

0.0045

0.0050

Y

/

Z

= 10

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe key contribution of this paper is that the training dynamics of a 1-layer Transformer model is demystified theoretically and empirically. The authors have found that the self attention operator performs a discriminative scanning algorithm on the input tokens, attending more on distinct tokens and focusing less on common tokens. They also demonstrate that there is a phase transition to an attention snapping step.\n\nSTRENGTHS:\n- The Transformer's self attention operator has been effective yet a black box in many applications. The authors' attempt in providing a theoretical explanation of the learning process is valuable.\n- A simplified setting helps better understanding.\n- Clear and straightforward analysis that aligns well with the paper's findings.\n\nWEAKNESSES:\nI do not see any critical weakness in this paper. However, more analysis might be needed on multi-layer Transformer settings. Although the authors have provided the attention patterns of a multi-layer model in Figure 7, it only covers a single layer in the model. Considering that most transformer-based models have multiple layers, it would be interesting to see if there are any different learning patterns in different layer depths.\n\nQUESTIONS:\nExtending the discussion on multiple-layer settings from the weakness section,\n1. Will there be a specific pattern regarding the phase transition across different layers? e.g. slower phase transition in the last layer etc.\n2. Is there a possibility that each layer may specialize on different roles? e.g. initial layers \"scan\" while latter layers \"snap\"\n3. (optional) Will these findings be consistent across different domains that use Transformer models? e.g. computer vision models, graph transformers etc.\n\nLIMITATIONS:\nThe authors have implicitly discussed the limitations. However, the broader impact was not discussed. Please add a broader impact section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper analyzes a simple architecture of 1-layer transformer’s SGD training dynamics for the task of next token prediction. The authors prove that self-attention acts as a discriminative scanning algorithm, with an inductive bias to favor unique key tokens that frequently co-occur with the query tokens.\n\nSTRENGTHS:\nThis work aims at understanding how Transformers work, and tries to understand the training dynamics in a simplified setup with 1-layer transformer. The work conducted both theoretical analysis and empirical experiments.\n\nWEAKNESSES:\n1. While it is an important research question to understand how Transformers work, this work only offers some observations regarding how attention weights change over time on two different types of context tokens. The setup is too simplistic to generalize to real-world scenarios. The conclusions depend on a lot of assumptions made in the paper.\n\n2. Several assumptions made in the theoretical analysis are not clearly justified:\na) The decoder layer learns much faster than the self-attention layer. The learning rates are hyper-parameters and in typical settings, they are the same for the decoder layer and the self-attention layer.\nb) The weak correlation assumption in Assumption 2 is not explained at all under what condition it is valid.\n\n3. The empirical experiments don’t share the same conditions as the theoretical analysis: e.g. batch size 1 assumption and much larger learning rate in decoder layer than self-attention.\n\nQUESTIONS:\n1. The assumptions used in the analysis need further justifications regarding the learning rate of Y and Z, and the weak correlation assumption in Assumption 2.\n2. The Residual connections are such a key component of Transformers. Neglecting them in the analysis leaves the conclusions questionable in generalizing to the actual architectures used in Transformer experiments.\n\nLIMITATIONS:\nIt is still unclear how the conclusions and observations from this paper helps to understand how Transformers work. While the research direction is generally good, I don’t find much surprise in the result that Transformers sparsify and focus attention on unique tokens.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper theoretically studies the SGD training dynamics of a 1-layer Transformer. Based on some assumptions, this paper introduces the frequency bias and discriminative bias of Transformers. To be more specific, this work shows that Transformers gradually attend more to distinct key tokens and less to common tokens. Later this procedure ends with a phase transition with a fixed token combination. Experiments are provided to justify the findings. I am going to update my score after the rebuttal if there are any revisions in the manuscript (if possible). \n\n---------------------------------------------------------------------------------\nAfter rebuttal, I increase my score from 4 to 5.\n\nSTRENGTHS:\n1. As far as I know, the proof technique and the framework are generally novel, at least for the Transformer architecture. \n2. Some high-level insights are well presented and intuitively correct.\n\nWEAKNESSES:\n1. Some parts of this paper are not very clear. For example, I feel the decoder layer is not formally defined. In around line 92, only the self-attention layer is introduced. Does the decoder layer refer to $W_V$? I thought $W_V$ usually refers to the matrix for the value embeddings. \n\n2. Although this paper considers a different setup, such as a different Transformer formulation, loss function, and different assumptions, the conclusion that Transformers gradually attend more to distinct key tokens is not a new one ([1] and [2] in the following) in the theoretical works. Especially the sentence \"we are the first to analyze the attention dynamics and reveal its inductive bias on data input, \" in line 320 is an overclaim. Here are several recent works that also study the training dynamics of Transformers. \n\n[1] Li et al., \"A Theoretical Understanding of shallow Vision Transformers: Learning, Generalization, and Sample Complexity\", ICLR 2023.\n\n[2] Oymak et al., \"On the Role of Attention in Prompt-tuning\", ICML 2023.\n\n[3] Tarzanagh et al., \"Max-Margin Token Selection in Attention Mechanism\", https://arxiv.org/pdf/2306.13596.pdf\n\nI know [2] and [3] are concurrent works, while [1] is an older one about training dynamics. I would like to see a comparison and discussion between the manuscript and these works. \n\n3. I suggest presenting a table of notations somewhere in the paper since there are many notations. For example, Section 5 is difficult to follow, although I can understand the conclusion.\n\nQUESTIONS:\n1. Although discussing a different pair of $\\eta_Y$ and $\\eta_Z$ is interesting, how can it happen in practical SGD? Is it an assumption or an algorithm?\n\n2. From my understanding, in this paper, the feed-forward layer after the self-attention layer of the Transformer is a linear layer. I am wondering whether the analysis can be extended to an MLP layer (even only one non-linear activation) with Relu or Gelu. I think feed-forward layers with non-linear activations are more common in practice. \n\n3. I didn't fully check the proof. I found a possible typo in equation 15 in line 536. I think the last term should be $\\log(\\bf{1}^\\top \\exp(Y^\\top LN(X^\\top\\bf{b}_T)))$.\n\nLIMITATIONS:\nThere is no potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
This paper presents a rigorous mathematical analysis of the training dynamics of a 1-layer Transformer architecture without positional encoding for the task of next token prediction. The authors demonstrate that the self-attention mechanism in the Transformer exhibits a discriminative scanning algorithm that gradually focuses on distinct key tokens while paying less attention to common key tokens. The paper also shows that the self-attention layer undergoes a phase transition controlled by the learning rate of the decoder layer, resulting in a stable token combination. The authors verify their findings on synthetic and real-world data (WikiText-103).

The reviewer didn't check the proof in mathematics for this manuscript.

STRENGTHS:
1. The paper provides a formal and mathematically rigorous analysis of the training dynamics of 1-layer Transformer models, contributing to a better understanding of how these models work.

2. The paper demonstrates the impact of the learning rate on the phase transition in the self-attention layer and its influence on the final token combination.

3. The authors verify their findings on both synthetic and real-world data, which strengthens the validity of their conclusions.

WEAKNESSES:
1. The experimental part of the paper is limited in scope, focusing on 1-layer Transformer models without positional encoding and not addressing more complex architectures.

2. The paper's assumptions (no positional encoding, long input sequence, and faster learning in the decoder layer) may limit the generalizability of the findings to a broader range of Transformer models and tasks.

3. In the training and fine-tuning of the Transformer model, Adam/AdamW is used more often than the SGD analyzed in this paper.

QUESTIONS:
N/A

LIMITATIONS:
addressed
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the training dynamics of single layer Transformers. They identify a certain scan and snap procedure of the Transformer that learns a winner-take-all solution given certain data statistics / training dynamics or the self-attention learns to combine tokens.  \nThey accompany their theoretical results and analyses with empirical results.\n\nSTRENGTHS:\nI salute the effort of the authors to mathematically study the training dynamics of Transformers. This is very interesting and timely. I think that the authors put in a lot of effort into this study to work out quite interesting results.\n\nWEAKNESSES:\nThe paper is very dense, hard to parse and therefore difficult to understand. Although I am aware of the difficulty to present theoretical results in a comprehensible manner, I urge the authors to invest time into the presentation of the work. \nIn general, there are various assumptions made along the paper which is fine if justified or discussed. See questions. \nNote that I read the paper a couple of times and still have major difficulties to obtain an overview of the results.\n\nQUESTIONS:\nAgain, in general I found the assumptions very strong and I felt that these shortcomings could be a bit better addressed in the experimental section. \n\n1)  Lemma 1: gradient dynamics of batchsize 1 - you are assuming here training the Transformer on a single example only? This feels very restrictive.  \n2) Lemma 1: Y and Z are assumed to be independent but depend partially on the same parameter matrices. This feels again like a very strong assumption. \n3) Can you elaborate on why you choose to integrate layernorm in your analysis? Could simplify your analyses if not important for your results, or did I miss their importance for your analyses?\n\nLIMITATIONS:\nI believe assumptions and limitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}