Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer
Authors: Yuandong Tian
Yiping Wang
Beidi Chen
Simon Shaolei Du
Venue: NeurIPS 2023 poster (2023)
Keywords: transformer
training dynamics
theoretical analysis
self-attention
interpretability
neural network understanding
TL;DR: We analyze the 1-layer transformer with next token prediction loss, and rigorously prove its training process and reveal how the token is combined via self-attention layer and the nature of its inductive bias.

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 Dynamies and Token Composition in 1-layer Transformer
Yuandong Tian'
Yiping Wang
Beidi Chen
Simon 2u²
Meta AI (FAIR)
Uniiversity of Washington (yuandong, beidic Ceeidiie@mmet...ee @meta. com,
Carnegie Mellon University *Zhejiang University (ypwang61,ssde sasdusdddsssssseeeeesees bocdujocs.wahhinttonn edu
Abstract
Transformer architectures have shown impressive performance in multiple re- search domains and have become the backbone of many neural network models. However, there limited understanding on how Transformer works. particular, with 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 I-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 in- put tokens, and reveal the nature underlying inductive bias. More specifically, with the assumption (a) positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self- attention acts discriminative scanning algorithm: starting from uniform at- tention, gradually attends more to key tokens that are distinct for a specific next token be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, progressively drops atten- tion 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 winner-takes-all, but decelerates due phase transition that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this sean and snap dynamies on synthetic and real-world data (WikiText).
1 Introduction
The Transformer architecture [ has demonstrated wide applications multiple research domains, including natural language processing 2 4. computer vision [7 speech [8] 9 multimodal- ity 100 11, etc. Recently, large language models (LLLMs) based on decoder-only Transformer architecture also demonstrate impressive performance 4
12 13, after fine-tuned with instruction data [4 or reward models [15]. Why pre-trained model, often supervised by simple tasks such
predicting the next word 4 3 13 filling in the blanks 2 [6 [7]. 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 [8) can approximate Turing machines [9] 20 and
can perform diverse set of tasks, e.g., hierarchical parsing of context-free grammar (21. if its
weights are set properly. However,
unclear whether the weights designed achieve specific tasks are at critical point, can be learned by SoTA optimizers (e.g.. SGD, Adam (22), AdaFac- tor (23] AdamW (24) In fact, many existing MI models, such as k--N, Kernel SVM, or MI.P, are also universal approximators, while their empirical performance is often way below Transformer.
To demystify such a behavior,
important understand the training dynamics of Transformer, i... how the learnable parameters change over time during training. In this paper, as first step, we formally characterize the SGD training dynamics of 1-layer position-encoding-te Transformer for
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

next token prediction, a popular training paradigm used in GPT series B 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 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- attentieo layer tends to become more sparse during training, as suggested by Frequency Bias, we discover that 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 WikiText (25).
Concurrent works on Transformer dynamies. Compared to 22 that uses &g 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] 22 30) give similar conclusions that self-attention attends to relevant tokens. In comparison, our work analyzes richer phenomena in I-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 solu- tions 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 [8]
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 growing interest in explaining the capability of in-context learning 4 of Trans-
former, by mapping the gradient descent steps of learning classification/regress
into feedforward 37 (6). Different from our work, the results in these
steps of Transformer layers [42 43 papers are existential and do not take training dynamics into consideration.
44 45
Training Dynamics of Neural Networks. Previous works analyze the training dynamics in multi- layer linear neural networks [47] 48), in the -tachet-teacher setting [49
50
52
[3]
54 [5]]
(56] [57] and infinitewwidth limit [58
59
60
61
62
63
64
65]
66
68
69
70 including
extentions to attention
models 722
73
For self-supervise learning, works exist to analyze
linear networks 77 and understand the role played by nonlinearity 75 Focusing on attention- based 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 I-layer Transformer can learn a constrained topic model, in which any word belongs to one topic, with E2 loss, BERT l2ke architecture and additional assumptions on learned attention patterns. Snell et al. [77] study the dynamics of single- head attention head to approximate the learning of a Seq2Seq architecture. While these papers also study the optimization dynamics of attention- based models, they focus different settings and do not explain the phenomena presented in our paper.
2

---

3 Problem Setting
Notation. Let {ux be d-dimensional embeddings, and { be discrete tokens. For each token, It takes discrete values from to M, denoted as E [M], and a = Ca, E TR is the corresponding one-hot vector, i.e., the rę-th entry of a, is while others are zero. Ugg the token embedding at location in a sequence.
Let U [ uM] E AR be the embedding matrix, in which the k-th row of U is the (T is the data matrix encoding the sequence of length 7 1. XU E R[T-1)xd the sequence of embeddings for a given sequence {-1, -TI....,TT-13 It is clear that X1M = 11-1-
embedding vector of token k. X [ [fI,..-;TT-I] -I-1]"
We use X to denote i-th sample in the sequence dataset. Similarly, ze[i] is the token located at t in i-th sample. Let be the dataset used for training.
the embedding after L-Layer Transformer Architecture. Given a sequence { { Iy, IT, .r+1). J. 1-layer self attention is:
exp(u WoW K Uz u../ /Vd) beT == T- E exp(u [, WQ W kuz
Here ber is the normalized self-attention weights (E==1 2 ber 1). One important detail is that we mask the weight that the query token attends to itself, which also being used in previous works (e.g., QK-shared architecture (7)). See Sec. Sec.7for discussions about residual connection. Let by = by..... E R be an attention vector, then 6, = 1 and üy = U X by.
lg Og-Normalization, We consider adding normalization to the output of the self-attention layer: ur = LN by), where LN(x) := x/|| ||> 2/22-- NormFormer and RMSNorm [79, 80) also leverages this setting (up to global constant). Our analysis
also be extended standard Lay- erNorm [8]], which also subtracts the mean of æ, while 88) shows that mean subtraction may not affect the empirical results much. LLaMA [82) also uses R.M.Norm. Empirically ür (or Wy)r normalized (instead of X by) and here we use an approximation to facilitate analysis: when the token embedding { (um) are approximately orthogonal to each other, then | æ|12 R ||æ||| and thus ür ss (.(
Objective. We maximize the likelihood of predicted (T + 1)-th token using cross entropy loss:
J= ED U2r+1 WvûT - log Σ exp(u1 (u Wvūr)
max Wx.
For simplicity, we consider single-head attention setting, and multiple-head attention can be re- garded as single-head setting with simultancous different initializations (see Sec. 4 We call Ty m as the query token of the sequence, and 7T+1 = as the next token to be predicted. Other tokens Ię S t contextual and query tokens can take values from 1 to M (i.e., m E [M]) and next token takes the value from to K (i.e.. n [K) where K S M. Fig. Taa shows the overall setting. For an overview of the notation used in the paper, please check Tbl. in the Appendix.
1) that are encoded in X are called contextual tokens. Both the
3.1 Reparameterization
Instead of studying the dynamics with respect to the parameters of token embedding U, key, value and query projection matrices WK. Wq and Wy, we study the dynamics of two pairwise token relation matrices Y ;= UT E MM and Z == UW UWQWTTTT /vd e MMX Intuitively, entries of Y and Z store the "logits" of pairs tokens. We regard the empirical parameterization using U. Wx. Wq and Wy as 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 [33 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 ad advantages and leads to useful insights.
3

---

(a)
Decoding& Softmax
(b)
Concestual tokers 5
Query token.,, Nust token Aydae
F(t|n)
Normaliztioo
#()nn))
Sequence casses
Seft attention
#(()s)
11g
m12
F()mn)
11g
Distinct tokens: There cuiss uniuue with P(]jo 7-0
Contextual bokens
Query taken
Next token
Commos toliens: There cests mutipee with P0.
Overall our setting. (a) A sequence with contextual tokens {æ1,..., TT-11 } and query token zy is fed into l-layer transformer (self-attention, normalization and decoding) to predict the next token 27-1- (b) The definition of sequence classes (Sec.3.2) A sequence class specifies the conditional probability P(]mm, n) of the contextual tokens, given the query token 27 the case that the query token determined by the next token: æT (xT++1) 1) (and thus [[[[m, n) = P(([n), while the same query token m may correspond to multiple next tokens (i.e., 10 ( is not unique). We study two kinds of tokens: common tokens (CT) with P([[n) > 0 for multiple sequence class n, and distinet tokens (DT) with P([nn) 0 for single sequence class n only.
Figure I:
and the token ÆT+1 n.. For simplicity, we consider
Lemma (Dynamics of I-layer Transformer). The gradient dynamics of Eqn. with batchsize is:
Px |
Y -yyNN(X (T)
2 = 12z - -
X diag(br) X (3)
Here I- vo- / projects vector into o's orthogonal complementary spuce, T/v and are the learning rates for the decoder layer Y and self-attention layer Z, := [ai...., aM] bT)).
RM and Om == Cepp(Y'LLNN) LN(("br))/11
br))/1 exp(YTLN(C
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 also more realistic than (27 that assumes dominant initialization in diagonal elements. Since (æT+1 -a) O, have Y1 Ż1 0 and summation of rows of ((t) and
Please check Appendix for the proof. We consider Y(0) Z())
O and P T by X diag(br) X1 XX1
(t) remains zero. Since ør is a one-hot column vector, the update of [ #2......... _M]] is done per row:
Px Mibll Y( -a(1)) |X
im 12. [i]diag(b 1
[1]
[i]ll
where m 2r is the query token for sample i, 2. is the m-th row of Z and 2.m' 0 for row m' * m = [i). Note that if ry [t] m. then by is function of 2m only (but not function of any other 2m.). Here we explicitly write down the current sample index since batchsize 1.
3.2 Data Generation
Next specify data generation model (Fig. (b)), named sequence class, for analysis.
Sequence Class. We regard the input data as mixture of multiple sequence classes. Each se- quence class is characterized by triple 8m.. (IP(I) n)) m, nt). To generate a sequence instance from the class, we first set IT m and xT+1 = n, and then generate the contextual tokens with conditional probability P(])m, n1). Let supp(m, n1) be the subset of token with P(I]m, n) > 0.
In this work, we consider the case that given next token TT+1 = n, the corresponding sequence always ends with specific query token Fr = m =: (n). This means that we could index sequence class with next token IT+1 n alone: Sn = (n)) n), (n),,)) P[//, P([ and supp(n) == supp(v(n))
Note that (m)) predicted, regardless of other tokens the sequence, which is trivial case. When (00(mm) 2 2, the same query token m, combined with other token / in the sequence with non-zero probability P(]m,, n) > 0, determine the next token.
means that the occurrence of token m alone decides next token to be
Overlapping sequence class. Two sequence classes on and sw overlap if rupp(p)))sup(() Asupp(n') + 0.

---

[a] Initialization
1b) Common Token Suppression
(c) Winners-take-all
(d) Attention frozen
Eips
Distinct Token
Common
(mL nt)
Seq class (m,n))
Etjna
Figure 2: Overview of the training dynamics of self-attention map. Here Cijme. [( n) exp(2mit is the un un-normalized attention score (Eqn. 5). (a) Initialization stage. 2m)
C [[]m, n). Distinet tokens (See. (.2) shown in blue, common tokens in yellow. (b) Common tokens (CT) are suppressed (&mt < O, Theorem 2. (c) Winners-take-all all stage. Distinct tokens (DT) with large initial value Gtmm, (0) start to dominate
the attention map (Sec. 5 Theorem D (d) Once passing the phase transition, i.e., 2 to O(K In M/mr), )
attention appears (almost) frozen (Sec. and token composition is fixed in the self-attention layer.
(Global) distinct and common tokens. Let S2(1) := {n P([/n) > 0} be the subset of tokens that co-occur with contextual token I. We now can identify two kinds of tokens: the distinct token I which has [S2(Z] exists one common token so that both [P[1 n) and ((n are strictly positive, e.g., common words like 'the', 'this' 'which that appear in many sequence classes. Sec. [ we will see how these two type of contextual tokens behave very differently when self-attention layer is involved in the training: distinet tokens tend to be paid attention while common tokens tend be ignored.
and the common token with [S((Z] 1. Intuitively, this means that there
3.3 Assumptions
To make our analysis easier, we make the following assumptions:
Assumption 1. We consider (a) positional encoding, (b) The input sequence is long (T + too) and (c) The decoder layer learns much faster than the self-attention layer (i.e., Ty > TTz).
Assumption a suggests that the model is (almost) permutation-invarian Given the next token to predict 27+1 and the query token øT acted query, the remaining tokens in the sequence may shuffle. Assumption [b) indicates that the frequency a token in the sequence approaches its conditional probability [[ n) == IP[[ m, #T+1 n1..
Note that the assumptions are comparable with or even weaker than previous works, e.g., 227 an- alyzes positional attention with symmetric initialization, without considering input data and [28 models the data distribution as discriminative/nov cere/moonndicriminattiee patterns, similar to our com- mon/distinct tokens. Empirically, NoPE [84] shows that decoder-only Transformer models without positional encoding still works decently, justifying that Assumption [Ta is reasonable.
Given the event {xT = IT+1 = n.. suppose for token I, the conditional probability that it appears the sequence P(]mm, n). Then for very long sequence - +oc, expectation the number of token appears in sequence of length T' approaches TPP(m,,,)) Therefore the per- token self-attention weight Cjjm, is computed as:
(),,, n) exp(2mt)
P(1m, n) exp(2ma)
Ol/m.n Er (T"(" [m, n) exp(2mar)
Er P(! (m, n) exp(zm!)
Σr mp
(5)
Here 2md is 2m 1-th entry and Clmma = P(Imm, exp(zma) is un-normalized attention score. Lemma 2. Given the event {xT m, mT+T+1 = n), when T + +o0, we have x (llag(br)).) diag( (Cm, where Cm,n [c1 22m -MI. T RM Note that cTT,"'
X bT Cm...
(6)
By the data generation process (Sec. 3.2), given the next token #T+1 = n. the query token øT m is uniquely determined. In the following, we just use c, to represent Cm,, (and similar for C,
Dynamics of
We first study the dynamies of Y. From Assumption [(c). learns much faster and we can treat the lower layer output (i.e., X bT) as constant. From Lemma when the sequence is long, know
5

---

given the next token F1+1 n, X by becomes fixed. Therefore, the dynamics of Y becomes:
exp(Y T f.) Y = MY fnnen- - an)), On 1 eTexpYYTT.
(7)
Here fn = X br
x by
R.M Obviously 115r|12
and fx 0. Define
1c
== f fK]- Since the vocabulary size M typically is a huge number, and different sequence classes cover diverse subset of vocabulary, we study the weak correlation case:
where the eigenvalues of R satisfies Ail < K and {x(EE) m- satisfies F = I T E,
Assumption 2 (Weak Correlations), We assume M >> K²2 and {fn)
VM Vi [K].
Assumption 2 means that fn share some weak correlations and it immediately leads to the fact that FT F is invertible and F is column full full-rank. Note that the critical point Y* of Eqn. should satisfy that for any given 2r+1 = n, we need a = en- But such Y must contain infinity entries due to the property of the exponential function in and Eqn. 7. we leverage a reparameterized version of the dynamics, by setting W = [w]....., wx]"
not achieve y* in finite steps. To analyze
F E R KMM and compute gradient update on top of W instead of Y:
Lemma 3. Given TT-1 n, the dynamics of W (here a; = exp(w,)/1 exp(w;)):
き - =n = =n) - on.)
(8)
While cannot run gradient update W directly, it can be achieved by modifying the gradient of Y to be Y Ty (fn FE' FEEe,)(e, ax,) X is small, the modification is small well.
Please check Appendix for for the proof. Lemma 3 shows that for every fixed n, only the corre- sponding row of W is updated, which makes the analysis much casier. We now calculate the backpropagated gradient used in Eqn.
Theorem If Assumption 2 holds, the initial condition Y(0) o. M > 100. ny satisfies M 0.99 < ny < 1, and each sequence class appears uniformly during training. then after
t > K2 steps of batch size 1 update, given event IT+1[i] n, the backpropagated gradient
g i = ( axll) takes the following form:
g =
infn
Fn' Brne Fn
n'
Here the coefficients L, (t), Ban (t) and (t) are defined in Appendix with the following properties:
(a) En((1) == 7(1) Enpro Brn (1) (z) > 0 for any E [K] and any L;
(b) The speed control coefficient >(t) 0 satisfies 7(t) (Irtt/K) when lu(M)-K when t 2 2(1+6") In(M)-K with 8' = O In Ia In
n
and ry(t) = O K lu(nyt/K)) ny
In the remark Lemma off-diagonal elements of E. and Theorem [ still holds with a smaller effective learning rate.
Appendix, we analyze the original dynamies (Eqn. 7 with identical
The dynamics of Self-attention
Now we analyze the dynamics of self-attention logits Z, given the dynamics of upper layer Y. Lemma (Self-attention dynamies). With Assumption (Tb)
b) (ie. T +00), Eqn. becomes:
im = 72?
diag(fn)
Σ (fr.. 1)fr', n'zn
(10)
nEW (m)
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., [9(I) = with ((1)) {n P(()n) 0}), and common tokens (CT) that appear across multi- ple next tokens (1((1)) > 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 GCT be the set of common tokens (CT), and G (D)
be the set of distinct tokens (DT) that belong to next token n. Then if Assumption 2 holds, under the self-attention dynamics (Eqn. TOP we have:
(a) for any distinct token I GDT (n), 2md > where m (n)):
(b) if |GoT] = and at least one next token n E (m)) has least one distinct token, then for the single common token 1 E GCT, 2mi < O.
Now we know DTs grow and single CT will shrink. For multiple CT' to shrink, the condition can be a bit involved (see Corollary iin 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 and its two distinct tokens I and V, the dynamics of the relative gain re (t) == Pa analytic form (here the query token m = (n) and is uniquely determined by distinct token l):
( 1= (n (n())))))) 1 has the following
[t/ = TL/11n (0)e (mm =: (0)x()))))x(()
(11)
where xi(t) := e2(emi(t)) is the growth factor of distinct token 1. If there exist dominant token lo such that the initial condition satisfies /lo/ (0) > 0 for all its distinct token & lo, and all of its common tokens satisfy 2mt < O. Then both 2mdg (t) and fnta (t) are monotonousty increasing over t, and
2f2 (t) S 2B.(()
(12)
here Bu(()= -0 Ti T1 So (n'd)).. Intuitively, larger B, gives largerTia/44 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) O, which means
1. Therefore, distinct token 1 with low conditional probability P(]m,, n) will have /1o (0) >> O, According Eqn. 11 this leads to quickly growing ratio /to/1/ (!), which means that the corresponding component fnt will be quickly dwarfed by the dominating component Sugg- On the other hand, token with high conditional probability F[[|m, n) will have
that (to/4((00))
P(la mm.)
3(0mm.))
smaller /lp (0), and the ratio [l)/(]((t) grows slower, costing longer time for lo dominate I.
Initial value as prior information. From the theorems, it is clear that the initial value 11/2 (0) ==
(1)mmm)))xxp.... (01) 2
critically determines the fate of the dynamics. Two tokens and l' with comparable conditional probability P(]mm, n) and [I, n) can be suppressed in either way, de- pending their initial logits (0) and 2mi (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 self- attention [1. From this perspective, each head h has its own Zn 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 En (t) decays slowly (e.g., G. 2 I/t), then (t) - +o0 and all contextual tokens except for the dominant one will drop (i.e., Jnl - 0) following the ranking order of their conditional probability P(I)m, n). Eventually, winner-takes-all happens. Conversely, if (t) drops so fast that B, (1) 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 (t) depends on the decay rate of 7t) and 3 (t), which in turns depends on the inner product f.T (t) (t), which is related to the logits of the common tokens that also decays over time.
Here we perform a qualitative estimation when there only a single common token and every next token shares a single token m (i.e., for any next token n, (n) = m). We assume all normalization
7

---

fzz 0.5
600
500
- 0.5
550
100
200
600 700
Figure 3: Growth factor (t) (Thcorem B over time with fixed f/2 0.5 and changing ny. Each solid line is xi(1) (t) and the dotted line with the same color corresponds to the transition time to for a given n/y-
terms in fa are approximately constant, denoted as Po. which means that f fn' = cxp(22m)) / and Bun' & En R f J R exp(22mc)/dt) as well, and 1 - fn Jn Rs due to the fact that common token components small, and will continue shrink during training.
Under these approximations, its dynamics (Eqn. 10) can be written as follows (here Co := Po/K)
2md 12Y
Σ fal
Baw (fri 1)f = -C nin'' Azml (t) = Co ' 42ml (13)
"e
nEW (m)
Surprisingly, now find phase transition by combining the rate change of (t) in Theorem
Theorem 4 (Phase Transition in Training). If the dynamics of the single common token 2nd satisfies
2nd = -C 127 (t)e and (t) Co (t) tzmi, then we have:
I (Co+ 2(M-1) (1) In (Co- +
to KIn
KM
By
(14)
2K 2K(M-1)
az In (Mmrt/K)) "y
t to 2(1+ 1)) M
M2
As a result, there exists a phase transition during training:
Attention scanning. At the beginning of the training, 7()) = O(yyt/K) and Bx(t) so 1 InnK K - 12 O(Int) This means that the growth factor for dominant token
lo is (sub-) )lnnar:: (1) 2 2 (O)B,(()
[K- 2 (0) and the
attention on less co-occurred loken drops gradually.
Attention snapping. When 2 to = 212(1 + K') M/Ty with 0' = 0 10 In T())= M
t/K) and B, = (Innnt)). Therefore, while B,,() still grows to infinite, the growth Xt, = (In t) grows at a much slower logarithmic rate.
O
tyt
See proof in Appendix This gives a few insights about the training process: (a) larger learning rate 7/y the decoder Y leads to shorter phase transition time to 2K M M/nyy (b) scaling up both learning rate (ny and 772) leads larger (t) when t + +oo, and thus sparser attention maps, and (c) given fixed nzz, small learning rate TγY leads to larger (t) when 2 to. and thus sparser attention map. Fig. 3 shows numerical simulation results of the growth rate (t(.). Here set = 10 and M = 1000, and we find smaller Tлy given fixed 72 indeed leads to later transition and larger (t) (and (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 atten- tion (E.g., T5 [7), ALiBi [85) MusicTransformer [86)). More specifically, the added linear bias exp(2m) + 2o) = exp(=m)) exp(=z)) corresponds to prior of the contextual token be learned in the self-attention layer.

---

Common token (seql) Distinct token (seq2)
Common token
Distinct token (seql) token (seq2)
Iter
her 0
her
6
net 6
Itor
her g
Final
Rinal
10
0
ter 0
ter
her
ter
her 6
10
Final
Rnal
Tokon
Tokon
Figure 4: Visualization of C. (n 1, 2) in the training dynamics of I-layer Transformer using SGD on Syn-Small setting. Top row for query token n = 1 and bottom row for query token n = 2. Left: SGD training with fJy = 7 = 1. Attention pattern e, becomes sparse and concentrated on highest P(]]n) (rightmost) for each sequence class (Theorem B. Right: SGD training with 7/y 10 and f/ = 1. With larger f/y, convergence becomes faster but the final attention maps less sparse (Sec. 6
Residue connection. Residue connection can be added in the formulation, i.e., üT
Lī( Uz,,) where ür is defined in Eqn. [ and ür is used instead in the objective (Eqn. 2 In this case, the Brr in Theorem | now is approximately Ann' ~ fn fw' + [[(b(n) = o(n')) )), which much larger for sequence classes n and ní that share the same query token øT than otherwise.
In this case, Theorem now gives g 7 (Infn En 1((())) Ben Sw.) for 2T+1l1 n.
Due to the additional constraint n' E (()()) (i.c., and n' shares the same query token), can define local distinct and common tokens to be within the sequence class subset e (m) and Theorem 2 now applics within each subset. Empirically this makes more sense, since the query to- ken 27 = my or m2 alone can already separate different subsets & (m)) and (m((m)) 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 ex- plain 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 I-layer attention based on their co-occurrence: even 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 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, 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 PP[[m)
for tokens 10-19 is increasing monotonously (the same for 20-29). The I-layer Transformer is parameterized with Y and Z (Sec. 3.1,, is trained with initial condition Y (0) Z() = 0 and SGD (with momentum 0.9) using a batchsize of 128 and sequence length T = 128 until convergence.
Fig. shows the simulation results. The attention indeed becomes sparse during training, and in- creasing γy with fixed fJz leads to faster convergence but less sparse attention. Both are consistent with theoretical predictions (Theorem 3 and Sec. (). Interestingly, if we use Adam optimizer instead, self-attention with different learning rate v)y Jj picks different subsets of distinct tokens to focus 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-Medis um and compute how attention sparsity for distinct tokens (in terms of entropy) changes with the learn- ing rates (Fig. 6). We can see indeed the entropy goes down (i.e., attention becomes sparser) with larger 1/2, and goes up (i.e., attention becomes less sparse) by fixing 7/2 and increasing 7/y passing the threshold 1/y /zz 2, consistent with Sec. Note that the threshold is due to the fact that our theory is built on Assumption Ic), which requires Tγy to be reasonably larger than 7/2

---

(seel)
Common (seq1)
(seql)
Rer
Rer
Ber
Ber
herd 4
ler
Ber
her 6
her
10
10
10
Final
Final
Pinal
Token
Figure 5: Visualization of (part of) C. for sequence class 1 in the training dynamics using Adam 22
Syn-Small setting. From left right: тлv = a2 = 0.1, 0.5, With different learning rate Adam seems to steer self self-attention towards different subset of distinct tokens, showing tune-able inductive bias.
#lasthnxxt tokens
#lasthnext tokens
#lastheett tokens 5/10
#last/next tokens 10/20
2.05
Average Entrpy efs,
2.0
2.0
2,00
1.95
1.8
m0.5
3.90
12
1.8
2.0
1.4
1.7
1.80
Figure 6: Average entropy of (Eqn.) 5 on distinct tokens versus learning rate ratio T/v /viz with more query tokens M/next tokens K. We report mean values over 10 seeds and standard derivation of the mean.
Iur-0
Ree-500
Ber-1000
iter-1000
15
25.
22
20
25
22
25
20
100
20
20
10
22
30
Ite-0
Iter-500
Iter-1000
Ite- Re--1500
15
15-
15
15-
25
25
25
25-
30
20
10
22
Figure 7: Attention patterns in the lowest self-attention layer for 1-layer (top) and 3-layer (hottom) Trans- former trained on WikiText2 using SGD (learning rate is 5). Attention becomes sparse over training.
Real-world Dataset. We also test our finding WikiText 25 using both 1-layer and multi-layer Transformers with regular parameterization that computes and Z with embedding U. In both cases, attentions of the first layer freeze (and become sparse) at some point (Fig. even if the learning rate remains the same throughout training. More results are in Appendix G
Conclusion and Future Work
In this work, we formally characterize SGD training dynamics of I-layer Transformer, and find that the dynamics corresponds
and snap procedure that progressively pays more attention to key tokens that are distinct and frequently co-occur with the query token the training set. To our best knowledge, data input, and potentially open a new door to understand how Transformer works.
the first analyze the attention dynamics and reveal its inductive bias
Many future works follow. According to our theory, large dataset suppresses spurious tokens that are perceived as distinct in small dataset but are actual common ones. Our finding may help suppress such tokens (and spurious correlations) with prior knowledge, without large amount of data.
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 2017.
[2] Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wttXiv://00.045000 2018.
[3] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, 1la Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
[4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language mod- els are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[5] 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 20Xiv:2010.11922) 2020.
[6] 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 Proceed- ings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[7] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, 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.
[8] Ching-Feng Yeh, Jay Mahadeokar, Kaustubh Kalgaonkar, Yongqiang Wang, Due Le, Ma- haveer Jain, Kjell Schubert, Christian Fuegen, and Michael L Seltzer. Transformerr-ransduuerr End-to-end speech recognition with self-attention. arXiv preprint 1rXiv:1910.12977, 2019.
[9] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for sule-vuue learning of speech representations, Advances in neural infor- mation processing systems, 33: 13:12449-13 12440-12400, 2020.
[10] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, 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.
[11] 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.
[12] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmano, et al. Palm: Scaling language modeling with pathways. arXiv preprint 2rX1v:2204.02111,2222 2022.
[13] OpenAL Gpt-4 technical report, 2023.
[14] Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay. William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, al. Scaling instruction-fieetueel language models. arXiv preprint arXiv:2210. 11416, 2022.
[15] 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-300 2020.
[16) Yi Tay, Mostafa Dehghnni, Vinh Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, U12: Unifying language learning paradigms. In The Eleventh International Conference Learning Representations,200 2022.
11

---

[17] 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 21(1):5485-555) 2020.
[18] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of gequence-t--squeene functions? arXiv preprint ar 10077, 2019.
[19] 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 1-12083, 2022.
[20] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 22(1):346333493) 2021.
[21] Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08117 2023.
[22] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980, 2014.
[23] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory
cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018
[24] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.01101, 2017.
[25] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07865 2016.
[26] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: To- wards a mechanistic understanding. arXiv preprint arX1v:2303.002445 2023.
[27] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37836, 2022.
[28] Hongkang Li, Meng Wang. Sijia Liu, and Pin-Ye Yu Chen. A thcoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. ICLR, 2023.
[29] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi. and Christos Thrampoulidis, On the role of attention in prompt-tuning. ICML, 2023.
[30] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. arXiv preprint arXiv:2306. 3(7):47, 2023.
[31] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of trans- formers recognize formal languages. arXiv preprint arXiv:2009. 2004. 2020.
[32] Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transform- ers and its implications in sequence modeling. arXiv preprint arXiv:2006.0922 2020.
[33] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. arXiv preprint 18Xiv:8807.038199 2018.
[34] Benjamin I. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms. International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[35] N Elhage, N Nanda, C Olsson, T Henighan, N Joseph, B Mann, A Askell, Y Bai, A Chen, Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
[36] Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint 10Xiv:2106.033 2021.
12

---

[37] Ekin Akyürek, Dale Schuunmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211. (55,:,, 2022.
[38] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105. 11111,....11.......111111 2021.
[39] Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ra- masesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint arXiv:2207. 0.400, 2022.
[40] 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.
[41] 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 PrXXiv233110002:: 2022.
[42] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 15:30583-30598 2022.
[43] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677, 2022.
[44] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statis- ticians: Provable in-context learning with in-context algorithm selection. arXiv preprint vrXi:22306.04637,200 2023.
[45] 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. 1895, 2022.
[46] 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.11227, 2023.
[47] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradi- ent descent for deep linear neural networks. arXiv preprint arXiv:1810.02281. 2018.
[48] Peter Bartlett, Dave Helmbold, and Philip Long, Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. In Inter- national conference on machine learning, pages 521-530. PMLR, 2018.
[49] Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for convnet with gaussian inputs. In International conference on machine learning, pages 605-614. PMLR, 2017.
[50] 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.
[51] Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural information processing systems, 30, 2017.
[52] Surbhi Goel, Adam Klivans, and Raghu Meka. Learning one convolutional layer with overlap- ping patches. In International Conference Machine Learning, pages 1783-1791. PMLR, 2018.
[53] Simon S Du, Jason Lee, and Yuandong Tian. When is a convolutional filter easy to learn? arXiv preprint ArXiv:1709.06122 2017.
[54] Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns mee-hideen-layrr enn: Don't be afraid of spurious local minima. In International Con- ference on Machine Learning, pages 1339-1348, PMI.R, 2018. 8.
13

---

[55] 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.
[56] Tianyi Liu, Minshuo Chen, Mo Zhou, Simon Du, Enlu Zhou, and Tuo Zhao. Towards understanding the importance of shortcut connections in residual networks, Advances in neural information processing systems, 32, 2019.
[57] Weihang Xu and Simon S Du. pverrparrmettiizzationn exponentially slows down gradient de- scent for learning single neuron. arXiv preprint urXiv:2302./0033 2023.
(58) Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[59] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in neural information processing systems, 32, 2019.
[60] Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably opti- mizes weerpparameterized neural networks, 2018.
[61] Simon Du, Jason Lec, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference machine learning, pages 1675-1685. PMLR, 2019.
[62] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-paramettrrizaion. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[63] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine grainer analysis of optimization and generalization for overparameterized two-layer neural networks. In Interna- tional Conference on Machine Learning, pages 322-332. PMLR, 2019.
[64] Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterizaaton:: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):84-105, 2020.
[65] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine learning, 109:467-492, 2020.
[66] Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems, 31, 2018.
[67] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018
[68] Song Mei, Andrea Montanari, and Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 175633) E7671, 2018.
[69] Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXiv:2001.11444 2020.
[70] Cong Fang. Jason Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean mean-field framework for werrpparaettrrized deep neural networks. In Conference on learning theory, pages 1887-1936. PMLR, 2021.
[71] 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 Machine Learning, pages 6426-6436. PMLR, 2020.
14

---

[72] Jiri Hron. Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nogp and ntk for deep attention networks. In International Conference on Machine Learning, pages 4376 4386. PMLR, 2020.
[73] 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 neu- ral networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:220..03466, 2022.
[74] Yuandong Tian. Understanding the role of nonlinearity in training dynamics of contrastive learning. arXiv preprint wrX1v:2206.01322, 2022.
[75] Yuandong Tian. Understanding the role of nonlinearity in training dynamics of contrastive learning. ICLR, 2023.
[76] 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 Neural Information Processing Systems, 33:1553 13383-15333,, 2020.
[77] Charlic Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint 0ttf::2003.0600... 2021.
[78] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ICLR, 2020.
[79] Sam Shleifer, Jason Weston, and Myle Ott. Normformer: Improved transformer pretraining with extra normalization. arXiv preprint arXiv:2/10.09456, 2021.
[80] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[81] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06455 2016.
[82] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Tim- othée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
[83] 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 trans- former language models. arXiv preprint arXiv:2205.01068, 2022.
[84] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization transformers. arXiv preprint arXiv:2305. 19466, 2023.
[85] Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2 12409, 2021.
[86] Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Ian Simon, Cur- tis Hawthorne, Andrew M Dai, Matthew Hoffman, Monica Dinculescu, and Douglas Eck. Music transformer. arXiv preprint arXiv: 100.00280000488,, 2018.
[87] Zeyuan Allen-Zhu Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparam- eterized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019.
15

---

Basic Notations
M
Vocabulary size
T
Sequence length
One-hot vector (1 at component k))
X E RT-1XX
Input sequence (of length T 1)
by RT
Vector of self-attention weights to predict the token at time T. contextual token (0 S
0e E RM
2) (one-hot)
#T -1 ERM
Last/query token (one-hot)
PT E RM
Next token (class label) to be predicted (one-hot) i-th training sample of token at location in the sequence Number of possible choices the next token could take. Softmax score of the output layer.
2 [1] E RM
K
(t)
Learnable Parameters
YE RMXM
decoder layer parameters
Z E M
self-attention logits
Zm
m-th row of Z (i.e., attention logits for query/query token m)
Hyperparameters Learning rate the decoder layer
JYY
Learning rate of the self-attention layer
Token Types and Distribution
(()) of (m)
Mapping from next token IT = n to its unique last/query token
The subset of next tokens for last/query token #T-1
P(1m,,)
Conditional probability of contextual token /
given query token m and next token to be predicted as n
GeT GDT(N)
Subset of common tokens
Subset of distinct tokens for IT = n2
Attention Score
čn E R
Unnormalized attention score given next token f2 & norormalieed attention score given next token FT 72 fz-normalized attention score given next token ØT n1 Back- bock-propagated gradient for fn
Cn E RM
fn RM
g RM F
Input matrix of the decoder layer. Each column of F is fr.
Self-attention dynamics
(1)
Relative gain between distinct token / and I' for next token
(t)
Growth factor bound of the relative gain
7(t)
Speed control coefficient
Table 1: Overall notation table of the main symbols in the paper.
A Notation Table
Tbl. gives the notation of the main quantities in the paper.
B Detailed comparison with the concurrent works
B.1 Comparison with
Setting, Assumptions and Conclusions. (28) analyzes the SGD convergence of I-layer VIT model (1 layer self-attention + 2 layer FFN with ReLU, with the top layer of FFN fixed random, token embedding fixed). Under a specific binary data model which the data label is determined by counting the number of tokens that belong to positive/negative pattern, [28] gives generalization bound when the number of hidden nodes in FFN is large, and at the same time, shows that the self- attention attends to relevant tokens and becomes sparse (if number of relevant tokens small).
In comparison, our work focuses on language models, assume broader data distribution (e-g., multi- ple classes, arbitrary conditional probability of token given class label) and incorporate LayerNorm
16

---

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 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 compari- son, our framework and analysis is much simpler by leveraging that certain nonlinear continuous dy- namics systems can be integrated out analytically yield clean solutions (e.g., Theorem (Eqn. [T] and Theorem 4 (Eqn. 128 ) avoiding complicated bounds in (28) This allows us to characterize the converging behavior of self-attentions when +.c.
B.2
Comparison with 29
29 focuses on 1-layer attention- baasd prompt-tuning. in which some parameters the models fixed (Wp, p. Wg). The analysis focuses the initial (3x one-step) SGD trajectory, and constructs the dataset model containing specific soalccttrrdvaa///ooee--rreeee indicates the token relevance. As result, [29] shows the attention becomes sparse (i.e., attending to context-relevant tokens) over time, which consistent with ours, and shows that prompt-attention find the relevant tokens and achieve high aceuracy while afffaattntiio/innaraateeeette
data, and the context-vector
can't '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
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 l-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 attention since its norm can be shown to go to infinity over training.
29
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
Lemma (Dynamics of 1-layer Transformer). The gradient dynamics of Eqn.
with batchsize I is:
Px by X diag(br) X (3) |X brll2
- Y LyNNNXX (21+111)))) T)T, Ž -a)) TT
Here P. := I- /l projects vector into v's orthogonal complementary spacee,ttrrsspee,,eeeeeee ny and are the learning rates for the decoder layer and self-attention layer Z, a == ax) E RM and Um := exp(Y L LN(X
br))/1 exp(Y LNN(" bT)).
Proof. With the reparameterization of Y and Z, the loss function is the following:
J(Y, ED [T LN((X br) log(1 T exp(Y by))))
(15)
and
exp(e, T LIN(X bT)) 1 exp CONOYYLLN(XT LN((XbT))
(16)
Therefore, taking matrix differentials, we have:
dJ (#T+1 - az) Td (#T+1 1 aa (AYTIN(( b) +)+ +T
P x db
X bl]
(17)
17

---

since in general we have d(exp(a)/1 exp(a)) Lda with L = diag(b) bb -bb let = XZ 2r and we have:
ddYLNN(X b) + YT Px b x Ld(X aT)
dJ = (er+1- - a)
(18)
||
PX- b X LXd2 xT ||
= -0rr+1 - a0 a)
dd b) +Y
(19)
Finally notice that PxT b L L= follows.
P P x diag(b) due to the fact that p. v 0 and the conclusion
Lemma 2. Given the event { = m, 2T41 = n). when T + +00, we have
x bT + Cm,..
X diag(bT) X - diag((m,,
(6)
where Cm,n [C1[m,m ".
RM Note that c 1 1.
Proof. Let p [exp(2m1) ..... exp(2mM]] )]T (oxp(2mz)).. c.p.)....ppp.p..........
RM, Pa. == exp(2max). and px =
expp2mxx m) then for any T we have
x
x
Σ
beT
T-1
Σ
Pt.
X XTPx PX
(20)
Ev Pay
1 px
t=1
Combining Lemma 18 and the definition of Cppm,n (Eqn.
5), we have that when T - +00,
M
PP m, n) exp(zmi)er
x b
Cm,n
(21)
Σp v P(l' (m, n) exp(=mv)
Similarly:
X diag(px)) X 1 1 px
X
(22)
Let T' +00, then also get
x diag(br) - diag((mm,))
(23)
D Proof of Section 4
D.1 Notation
For convenience, we introduce the following notations for this section:
Denote E' = (I -)
Apply orthogonal diagonalization E and obtain E U DU where U Uu...., ux] E Oкжк. D diag(A1...., Ax) and [Ai ... |Axl 0.
Denote F' [F, F'] E MM where F E (M is some matrix such that rank( F²) = M. This is possible since (fi)se[K] are linear-independent
Denote W (F") Ty F' [WT, N...T.FTTT. [wi. [w...... WK. WK+1.
wx] E MM XM
Denote Çn := M (en 1 E1)ERM
M-1
Denote Q1 C. 1+ M-1- 90 =
where i,j [M], + 1N- - M
Denote a continuous function that satisfies 1<(0) = and h 1 (M 1 +
exp(Mhh)~- Details in Lemma
18

---

Denote to be the constant defined in Lemma 8 that satisfies O
In (ull(Ad)) In(M)
Denote N, == Ei=1 [[rr+1 = v| to be the number of times the event 4r+1 = n happens.
. Denote N = [N/K] to be the average value of N₂ when PP(n) = 1/K and A ==
[VN (n()) []] ]1 to be the radius of confidence interval centered on N with confidence 1 -. 6.
Here A/N VN V (m()) < since K². Details in Lemma 10 and Remark
Denote W'N) == (0) (N), Øx (N), O,..... OJ E IR JRMxM XM where (n.NN = (M S). Vn [K].
D.2 Proof of Lemma
We assume (mE = [K] for convenience, but we claim that our proof can be easily generalized into the case where f * [K] by reordering the subscript of the vectors. First, we prove the dynamics equation of the reparameterized dynamics Y.
Lemma 3. Given TT-1 n, the dynamics of W (here a, = (xp(())/11 exp(w,)))
w -y = (e)
While cannot run gradient update W directly, it can be achieved by modifying the gradient of Y to be Y m/ (f FE'en) -a,) >1 is small, the modification small well.
Proof. We let F' = [F, E RMXM where rank are linear independent. And we further define W == (F') T - = [F, "" = [WT, YY FTFF
M, this is possible since {fn]nE|K]
[wf,..., uK, wk+1..... "MM] RRMMM, When given IT+1 n, the first term of the differen-
tial of loss function J is:
tr (arTT dY
X (IT+1 (00+1-a)7 ( a)7)
= tr(dY -tr(d(W))Te, e (#T+1 - aa
XX br
(24)
So W = en (xT+1 a)).. This nice property will limit W to independently update its n-th row for any 2T+1 [K], and the last M K rows of W are not updated. Similarly for have
eXp(UWVūT)) = exp(Y fn) 1 exp(y fn) 1 (x 1 exp(w,.) (25)
exp(wn)
a = 1 e'ep(UUVvr))
We get Eqn. g by combining the above results.
If we don't 't run gradient update on W directly, we can run modified gradient update on Y:
= (f.. FF' e''ee)) (e, an)
(26)
This will lead (note that F does not change over time due to Assumption (c)):
W
FT F (Sn FE'en) ax,)
(27)
My F fa - F(I (1 (1+ B)-1)en) (en an)
(28)
ny (FT fn - F Fen + en) (en -an)
(29)
eyen - T
(30)
By Lemma we know that if A is small, so does maX (E,) and thus the modification
small well. In Lemma 5 Remark reduces the learning rate, if all off-diagonal elements of E are the same.
will show that the additional term -F effectively
Lemma shows that 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 M
Cn (e, 1) E A we have:
19

---

Lemma' S. Assume Y is initialized to be zero matrix, Z is fixed, and the learning rate of Y is nl.
Then if event 24 n always holds at step (s 2 1) we have
(x) (M - 1))
(31)
exp(Ahh 1))
(M 1) exp(MM" 1)) j=n
any(s)
(32)
j
(M - 1) + exp Mh* (s- 1))
And thus en - an(ss
M M-1+ 1+11p (s- 133 Sn. Here h² (s) satisfies:
M-1
1² 1) (M 1) exp(M 1))
ny
$21 1
106)) =
(33)
Proof. We prove this Lemma by induction.
Step Note that Y is initialized to be zero matrix, then (0) 0, Vi [K]. So we have
(1) =
Vj
E
[K]
(34)
M
1
1
j=n n
(1)
M
(35)
1
j+n n
M'
1
my (1
j=n
(1y(1)
M
(36)
"/
1 *+1
M
It's easy to check that these equations match that of Lemma 5
Step Assume the equations of Lemma hold for step 1. Then at the step, we have
exp((M - 1)h2 - 1))
exp exp(Mh"(s 1))
exp(MMi 1)) (M 1)' j=n
(s) exp((M 1)/²( (s 1)) + (M exp( (s 1))
(37)
1))
exp(Mh+(s- 1)) + (M 1)¹ 1
exp((M ()h²(s- 1)) + (M 1) exp( -h- 1))
M
j=n
exp(Mh* 1)) + (M - 1)
(n)
(38)
exp(MM> 1)) + (M 1) j+n
(M 1) eep((M>>(s 1)) (M- + +ii - 1)) -(M - (s),
j=n
(n)(s)
(39)
exp(MMr - 1)) + (M- - + hs 1))
= (s),,,
j n
And the equations of Lemma alaso hold for step s. So we finish the proof.
Remark 1. If we following the original dynamics (Eqn. 7), then corresponds to the W dynamics
as follows:
W + (I + ELE' a.) = F fulen - -a,)T
(40)
When all dia-onal elements of E are identical, i.e., fo. fr for n + n', then 0 S p S 1 and
we have
wn ny(en -an)
(41)
w myplen an))
$*n
(42)
20

---

So if different sequence classes are sampled uniformly, then by similar induction argument. we will have
Wnn = -M
Sn +
E.ES S.' 22
=1 (1--)MM
(43)
where the last equation is due to the fact that Σn Sn - M-1 This means that En oon Sn' = -G.. Therefore, the effective learning rate is fly == (1 - p)nTy S ny.
M
En (e. - -
-
M (1 - 1) 0.
M-1
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 directly obtain the equation that needs to satisfy by integrating the corresponding differential equation. Lemma 6. If consider the continuous version of h² (s), namely h, as the following ODE:
dh
'γ
(44)
dt
(M 1) + exp(Mh)
and assume h(0)- = 0, then we have
exp(MM(()) + (M - 1))M)( = Mvy +1
(45)
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 E N,O S (s) - h(s) s 2ny Then there exists some constant c = e(1)
M such that for any < In(M)/m> I/, h(s (s) 2 h(s).
* (b) ht (s) h(s) + 0 when - +00.
Proof. (a) First we show that / (s) 2 h(s) for all 8 E N. and the convex packet function of ht can almost control the upper bound of h. Define ho IR - R as follows:
(46) Here [-] and [:] mean ceil function and floor function, respectively. It's clear that he is strietly monotonically increasing function, and for any E N, (s)) = h² (s), while for any t of N, (t, (t)) lies on the line connecting point (Lt), h" ([t])) and point ([t), h" ([t])). To prevent ambi- guity, we let à (t) to be the left limit of h², i.e., (t()) lim h (t').
(t) (t [t])- [h" ([t]) h (Lt))+ ((tt),, WE
We claim /t()) < (t), VI E Rt. We prove it by induction. First when L = 0, we have (0)
(0) h()) 0. Then we assume h(t') < 1('(t)) hold for time L' S L N and prove that h((() < (L') hold for l' E (1,6+ + 1]. If this is not true, then from the continuity of ht" and h. we know it must exist E" E (t,t+ 1] such that h(t") 2 h²(tt) and h((t)) > (t"). The later condition results that 7)y [M + exp(MM(t"))
my [M 1 exp(Mh"(11* [[t" -!. So h(tt") < h" ((t" = 1 (((")) < (1"))))
(47) This contradicts the hypothesis h(t") 2 1e (t"). So h(t') < hPt(') hold for t' E (t, ( + 1] and thus
for all t E R+ Hence for any N, we have h(s) < (s) = (s). Actually, we can use the similar method to prove that h(s) < h (s) for any s E N
Then we show (s))) 1=())= < 2/yy M by proving that for any N+ h(s) must meet at least one of the following two conditions:
(i (s)) E 1s
(ii) (s) h(s) < 1t*1"((
h(s 1).
If (i) doesn't 't hold, then we have for any E [s - 1, s), h(t) S h(s) <h= 18 1). which
results that in(t) > h (t) for all E [s- - s). Therefore, (*())) -ha))-((s (=---))= (s) (1 (s 1) <
h(s)) h(s - 1) and thus h(s) meets condition (ii). It's clear that h(() and h(1) meet (i).
21

---

frr 0.05,M=
7t= 0. 1000
0.155
9350
0.255
0.10.
0.000
0r
1/h
0.10.1
0,-02
0.-02
3.000
0-00
0.075
0-50.5
0.000
0.-02
0-110
110
0-
Or 1.0
2.0
2.5
3.000
3.000
0
1.0
20
30
50
60
70
80
10
20
30
40
50
62
70
Figure 8: Numerical simulation of At and f with changing vy. The stepped folded line represents ht" and the smooth curve represents h. The gap between h" and h is bounded and goes to zero when time grows.
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 lx*(s) h(s) S h h*(ss) h- 1) (1)
(0) = 7/y / M from Eqn. 33 And for any 2 2, every time h(s) transfer from (i) to (ii) exactly at
8. which means that f(ss
1) meets (i) and thus no smaller than h (s - 2), we get h (s) - h(s) S
h" (s) h(s 1) S h- (s) (s 2) S h" (2) h-" (0) S 2ny /M.
Finally from Eqn. 53 in Lemma [ when < M we get h(s) = (nrttMM and thus there exist
"Y some constant (1) such that fiss + cc h(s)
+ 2ny /M 2 (s) h(s).
(b) Assume that there exist E (0,2ng 2ny /M] such that (s) h(s) 2 for all E N. Since is unbounded, then i((1) - when + op from Eqn. 33 so there exist some % E N such that when
2 $6 hhshhh(ss111---- h(s) S (+1n(1/2)) /2/ Also, from Lemma 9 we know that exists s/ = (3+6) (M(M) where 5 > 0, (11) such that when 2 s0 exp(MM(s)
2(M 1). Since consider the case that = [4] 2 so == max(<6) *6). Then denote A == 2(M-1
co, we just 1, we have:
exp(MA(s])
(t) h(tt) =
N
"
M (xp(MM((s))
M 1 exp (xh(()(t))
f/y
ny
<
M - exp((M(((s) +6) M - exp(MM(ss
M/y exp(MM(+)) [exp exp(M))s -)(
M - 1 +ex++xxpMMiie (M)(M()(s) +]]] (xp
Ty exp(MM((s) e
(48)
2[M - 1 + (x((M(((s) + e))] [M - 11 exp(M(((s)) +c)
My (Me)
(s 2 so max(s6, s())
(1 A expl Mh(ss) exp(4n))) C
exp( (Me)
- 4exp(4my )M -
Here C 4xxp(4nx ) > O and for the last inequality, we the fact that t so
exp(Mc)
3 M and thus by
h(s) S h(t) O
In (n( 9) from Lemma
So we get
A
[A (t) h(O) (M)) h(s0)) S fro Cdl
(49)
This contradicts 10 (1) - h(()) 2 0. So the original assumption doesn't hold, which means that le" (s) - h(8) + when + oc.
Remark 2. By some qualitative estimation, we claim that if ny O(), then there exists some constant c O(I M) such that h((s)) < hs))s) < his + c) for all > 81 : 21m(1+an) where w1 = OIn In M/ In M) ) is defined in Lemma g Denote 5h((t) = he h²(t) -t(t), when 6h(t) < h(t),,
ny
22

---

we have 8h(t) = (t) h(t) x My 8h(t) exp( X ft by computing the second-order derivative of öh, and thus h (t)- 20/ (Mdt) = (In M/ Combining this with the fact that ht) (0InMMyyl)MM when & > 81, we prove our claim. The results of
Lemma 7 and Remark are also confirmed by the numerical simulation results as Fig. g
we just assume Tл < 1 and replace h" with h in the latter parts for convenience. Then we further investigate the properties of Eqn. [45]
So from Lemma and Remark 2
Lemma 8. There exists w;, 0 < w; & =, 2,3, such that for h E J1 == M
(1+w1) In(M)) 21
we have (t)) S (M- - (M-1))Mt M))h(tt.. And for h & J1, we have cxp(Mh(t)) > (M 1)))(t))
Here 11 = O In()))) and if M > 100, we have wo S (M In < 0.01.
In(M)
Proof. It's obvious that exp( - (M 1)Mh((t) has two zero points in RT Let h(()
=
M -22-wo), we get
M
WD In M
(In
M1-w0 O( M In (M))
(50)
M
For another zero point, let w1 E (0, 1) to be some constant such that h(t) exp(Mh) (M 1)Mh, then we get
(1+w) In(M)
satisfies
M
(M- AM-1 1) dIIn(M)
(M (M-1) -1-1] 1)
M (1+
M
M
(51)
e In In(M(M), In(M)
where d' (0.5.2) 2) is some universal constant.
Remark 3. From Lemma 8 if we assume M > 100, then wo < 0.01, and if assume тy > Mi--u > Mous, then (1(1)
(M - 1)Mh(t) has only M in oo). For convenience, we just assume M > 100 and ny
> Mi--o and function exp(Mh(t))
M
one zero point (1+w)) In M
Muls and thus focus on the unique zero point (1+w1) In M of h in the latter parts.
(M (M-1)
We can then show the properties of speed control coefficient y(t) = (MMD)+)x((MMN(//) as below.
Lemma 9, We have two stage for h and y;:
When t & Kln(M), we have exp(MM((/KK) S min(M - 1, (M - 1)Mh(t/K)),
O0 t/(K)) and (t) = O(πt//K).
When 2 2(1+w1) In(M) where W1 = O ( is defined in Lemma 8 we have
cxp(Mh(t//K) K)) Z max(M 1, (M 1)Mh(t/K)), h 0 In In(Mnyt/K)) and ((t) = 0( ny!
K Klln(Mytt/KK)
Proof. For convenience, we just let K = 1. And the proof for K # 1 is similar. We denote exp( (h))
(h) exp(Mh) and (h) (M 1)MK"
M-1
Step 1: < In(Md) If In(M-1) 1) from Eqn. 45 we have:
nr
M - 2 (M 1) In(M - 1) 1)
In(M)
(52)
Mny
ny
So when & < In(M)
we have < In(M-1) and thus exp(M)(()) S min(M - 1, (M (M-))))))
1)M)(().
M
ic., A, Ae < 1. Then from Eqn.
45 we get
M -0 (
h =
(53)
(1 + MM 1)
(M- 1)h
Mny++1
7 = 1+ A
(1+ +1 = 0
(54)
23

---

201+)))l(MM where wi O la(d))) ). So now h In(Ad-1] and thus A >
Step 2: >
M
from Eqn. 52
Then if exp(Mh) MM 1)h, i.c. A2 I 1, from Lemma have
Mny
II(M)) Therefore, M
(M
2(1 + w1) In(M)
1
(1+ A2)
M In M
(55)
'/y
M
T/Y
Contradiction! So when
2 2(1+u)) In(Ad)
we have A2 > 1.
Then from Eqn. 45
we get:
11 A. -
1
h = M
=
In In(Myyt)
(56)
M
M (M 11)ln(
Mny++1 I+a,
(In(Mrrl)
= M (1+
Mayt-1
0
m/y
(57)
I+4,
D.4 The dynamics under multiple uniformly sampled sequence classes
We then generalize our analysis of W to the case where 2T+1 can any value [K] rather than fixing .r+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, == E - E = n] to be the number of times the event JT+1 happens, then we have:
Lemma 10. For (0, 1), with probability least - 5 we have
N
2 <
2
Nnx [NP())]] S
In + 1
N
In
(58)
2
Proof. From Hoeffding's inequality, we have
(1 N N
P
P(m) >t) s 2eepp(2222
(59)
Let / = V 2N (n( and we can get the results by direct calculation.
Remark 4. From Lemma 10 if we consider the uniform sampling case where P(n)) = K. then NP(n) N/K > VN. So Nr are all concentrated around NP(n). Recall the definition of N [N/K] and = I In(+)), with probability at least 1 we have:
IN. NISAEN
(60)
We then further investigate the concentration of (N(N.)::
Lemma 11. For E (0, 1), with probability at least 5 we have
A
(M(N b(N) S (NN) N
(61)
M 1 + exp(MM(N.))
1 exp(MM
(62)
o'
M - -1 exxeex(((NN))
where d' > is some constant such that o' I A < In(M). And i
(K(1+
la In
where
wi is defined in Lemma 8 then o' < $ L
24

---

Proof. First, we note that h has a decreasing gradient, so h(()) 2 h(x) x and - -
h(x1)) x 22 for any 21, 22 0. So with probability at least 1 6, we have:
h(N)A
Ik(N.) - M(N) - h(N) h( A) h( A)
h(N)
(63)
N A
N
For the second inequality, without loss of generality, we let N₂ exp(Mh(s))) and note that:
N. Denote g(t) = (M
1
dg
M exp( (Mh(s)
dh
ds
(M - 1+ exp (M)(s)))²
ds
my M exp(Mh(3))
(64)
M 1 exp(MM((s)
(M - 1 exp(Mh(( (M)(
1
M N
M - 1+ exp(Mh(s))
(M - 1) 4
the last equality holds only when h(s) max,E[N., +)))(s) A, we get:
In(M-1) M
So from 19(N + A) g(Nra) S
M 1+ exp( + A)) M exp(MM(N)) S
M + (xpp(M((N)) 377r
(65)
21+ If N
+ A with W1
O ln la defined in Lemma 8 have o' 7) 4/3
BV
My N < In(M). IN
A, we utilize the Eqn 45 and obtain:
1
M - + exp(( + A)) M - 1 - (MN(M)
(exp(MM(N) + A)) - exp(eexpMM))M.e
M - + (XP(M
M - 1 exp(xM(M +])
1
Mny
M 1 exp(MNNN + A]) (Eqn. 45
M + exp(MM(N))
Mny A
2(1 + wi) In(M)
(Lemma N₂ 2
+ A)
M + exp((MMNN)
M + 2 (N
7/
M - 1 (xx((M(NN) N
So 0' < A/Ñ. When N, < N, with probability at least - 8 we have N₂ Z N A, and similar inequalities also hold for such cases, so we finish the proof.
Recall that S.. E RM is defined as C. = M
M-I (en 1). And we have (1 = C.C. = +
M M-T 90 CJ C; (M- for all i, [M] where + j. For convenience, we denote W'NN = (NV, (x)NN,
OJ E MxM, where (N) == (M L)A((N/K1))).
(M - 1),(N)).. So using these concentration inequalities, we get:
Lemma 12. Assume the assumptions in Lemma hoold but we uniformly sample the training data.
Then if the total number of epochs N satisfies N > K2, we have Y = (F")) - (1 + O')W"(N)
where O' = diag(f....... ................ O).... E RMxx and with probability at least we have
10:1 S
In K Vi [K].
VN
Proof. From Lemma 5 and the first inequality of Lemma 11 we know that
(N)
(M - 1)h(NN))..
(66)
(M 1)A(N)sn + (M - 1))((()))))))))) A(N))S
(67)
(1+ + 0w) -MM - 1))((N)..
(68)
(1+ O,)
(69)
25

---

where for any E (0, 1), with probability at least we have |0:| I vN vin( vn e [K].
K
Therefore, (N) = [N), (N),..... Wx(NN,,, oo = (T O') then from W YY")T, we finish the proof.
Then, we can give out the exact solution of Y by pointing out the properties of F and Fx from the observation that each row of Y should be the linear combination of vectors in {fr, ]ne[K]=
holds and Y()) = Furthermore, we assume the training data uniformly sampled and the total number of epochs N satisfies N > K² Then the solution of Eqn. 26 will be:
Theorem 5. If Assumption 2
Y (F
= F(I E + ())(N)
(70)
Here θ == diaag(0........ and for any E (0, 1), with probability at least 1 - we have [8:] K
x
In( 5 ) Vi [K].
VN
Proof. Let q.z [M] be the i-th vector of (F")-1, then have 2, = j]. From Lemma 12 we get Y = (FF) W'().. And from Eqn. 26 we know all the columns of Y the linear combination of fn, 7 [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')-T, i.e., Qi, E [K]. to be the linear combination of fn, n E [K]. which means that qI,..., qK must be the basis of == span(ff; e [K]) and thus fK+1,-... 9M are the basis of E' sfa K < I M). Therefore, we get E L 3', and thus .q...... q can only be (Ft)T, So the proof is done.
Actually, we see that the result of Theorem 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. in the rate of O x))::
Corollary 1. Assume assumptions in Theorem hold, M >> 100 and by satisfies M 0.99 < nl Then Vn E [K], we have
-
(IT+I - -aa) = (M 1) + (M)(M)) S..
(71)
M 1
(M 1) + exp(Mh(NN) (1+ Sn
where a > -1 and for any 8 E (0, 1), with probability at least 1 we have |o| S nr V (n)(),
and when N >>>>> K(WNNnln +
2(1+w)) with wi defined in Lemma 8 l S VN v In(t).
K
"y
Further, to let |T-I anll2 S with probability least 8 for any n E [K] and < we need the total number of training epochs to be at most O O log( ))
Proof. Note that Ø741 Cn. then we just need to combine Lemma and the second inequality of Lemma 11 to get Eqn. 71] Denote S, to the number training epochs that are needed to let ||T++1 On | x e, then we have
M
(SS))
(72)
M
But note that h(t + 1) - -h(t)= 2 M 14expl MA(SS,)
TVy
JyC M-1
Vt E 10, - 1] from Eqn. 71 we have
h(Sw))
1 M
S, ey (M 11)
In
(73)
EπLY
Note that < 1 and we have N > 22 then we have S =,, Sn <
K In M city
26

---

D.5 Proof of Theorem 1
Finally, we turn to prove Theorem [ Obviously, all the diagonal elements of are zero and the off-diagonal elements of E non-negative since Cljm,n 0. Note that E is real symmetric matrix, then can be orthogonal diagonalization by U DU where U [ui,... ukl OK ; D = diag(A (1,, Ax) and [vıl 2 2 |xx 2 0. Then we can get the following properties of E and E":
Lemma 13. J (11) (1,))) - |A1
Proof. We have:
Eigl = =u, Duj S [Ail (alll/llllll Vi,j
(74)
Lemma 14. E RK satisfies [^al < X < 1, then (I + E) is invertible and (I+ E) ----==-- 10-1=111 where E' satisfies E' = =T D'U D'U
E
and D' diag(X; (y ..., Xx) and X; 1+2, vi E [K].
Proof. Since U is orthonormal and | < X < 1, we have E = U D - O. Then from the property of the Neumann series, we get I E is invertible and
20
(1+ = I+ Σ E"
(75)
T+U
( D
(76)
T- U D'U F
(77)
Here we define D' diagg(,,.... Xx) and use the fact that Σ (~X;) )m8
A
1+X
Lemma 15. If|aa S < 1, then maxιejk (A((E)) < - S b-xx
Proof. We have
X 1+
[ Av
max K
(A(EE))
max SE[K]
maX [Ai S - Pal
(78)
Lemma 16. Assume that Assumption holds, then all the diagonal elements of E' are non- i.s. Ei S 0, E [K]. Further, if there exist any k + i E [K] such that Exs > 0, then E/, O.
Proof. Note that E. = ΣK Aucife 0 (here usa is the k-th component of eigenvector u.) and [&l < 1, we have
E
Σ
1+
A
Uik
Σ k=/
Σ
2
Σ
2 A
50
(79)
When Ey O, then X = (Ai,..,, Ax) must don't have overlapping entries with respect to u.
which results that Eig = Σ -=1 dsunuus = 0 holds for any j E [K]. So prove the results.
Lemma 17. A < 1. then EE Enn' -1
27

---

Proof. From Lemma 14 we have:
K
K
Ak
Enn''
Enn'
U
unk Un/kl
A
k=
k=1
U
+
A
k=1
(80)
A11
K
Unnk U
A1
2
111
Mank P))
Man 122
A1
1-1
Finally we can prove our main theorem in Sec.
Theorem If Assumption 2 holds, the initial condition Y(0) = O, M > 100, Ty satisfies M- < nl < and each sequence class appears uniformly during training, then after t >> K2 steps of batch size 1 update, given event TT+1 = n, the backpropagated gradient
g Y ( cxli)) takes the following form:
g = Infn
Σ nint
Brne
Bune fn
(9)
Here the coefficients (t), Ban (t) and 7(t are defined in Appendix with the following properties
(a) (t) == 7(t) En Brn (t) (t) > 0 for any n E [K] and any t;
(b) The speed control coefficient y(t) 0 satisfies 7(t) Ofyy when $
Iu(M)-K
Ny
and 7()) = O K t/() when & 2 2(1+6') In(M)-K with 8' O - In In ).
n
Proof. Note that if Assumption 2 holds, then Ft = (I - E Recall @1 = 1 1 M-1 = 1 and M Qo = (M- Rs Then given xT+1 n. we get:
1
((or+1lil-ap)) gii
(81)
F(I- +"]=EE)(+
- a(i)), (Theorem
(82)
(1+0)7 * F(1 E')(1 + ()(go....................)..)( 40
(Lemma Corollary (83)
Infn Inln
Σ
Bun'
(84)
n' [[[R
where
(M 1)?((t//RT)
(t)
(85)
(M - 1) + exp(MM((t/KK))
in
(1 + a)(qn (1 + 0w) EEnm) 90 Σ
(1+ Ox)
(86)
kefn,Re(K]
(1 + [0 - E (1 +1+00 81) +
(87)
Ban'
(1 0)(0 (1 0n) E' go((1 0w))
Σ
(1 0z.) Fx.)
(88)
k
(1+07) [Em' (1 + 61) + 8z]
(89)
28

---

Here a is defined in Cor. and satisfies
-1 < < M. 161| &
VN
K vin(2))
+
ir <
and
102l I (M-1P x 2(1 + 3)
M
3 M.
Here we use the fact that 10), [0:1 x
vN
K
In })
ΣME[K] AbU = Exn and the fact from Lemma 15
E S max (,,(E S 1/K M 1 K-1
(90)
kE[K]
(a) Now let's prove that En (1) O. First from (T + FT) F!) I we have F. E² F.F O.
Then use the symmetry of E' and E", we get
(EE') mn
Σ 1=1
E
Σ
Enk En
Σ k=1
Enk En
Σ
En
Eun E
(91)
k=1
lighn
Note that FT F = I + E. we have Enny = STT Sni, Vn' + n and Enn = 0. Then
F
(E E' E.E')n Onn =0
Σ
En E
Fy
(92)
k
Note that (A(E)) > 0, Vi E K] in Assumption
2 immpies that Egs > holds for some + E [K].
Then from (1) of Lemma 16 we get Σx Fm f fn O.
From Theorem
we have Brn' (1 [)( (1 51) + 6z). Note that 0 < 1++ < In(A)), we
have:
Σ n' o/
Bun fn fn' = (1 + o)]
Σ
E + +1) + 83) Emn'
n'yn
=1 [1++o]-[0 (1 5,)
5z Σ Ensy]
n'nn
K
= (1+0)][( + 81)
Σ k=1
2
+ 8z
Σ Emarl
(Eqn. 79
1 X
n'stn
1 81
3
3
2 (1 a)
(min (b((E)
KJAIl
(Eqn. 90
|03|
A1
M
M
3
> (1 [1
(min (A)())
K|Ail.. (161 < 1,111
< 1)
M
K
> 0. (Assumption
2
(93)
We directly use Lemma then we finish the proof.
E
Proof of Section 5
Lemma (Self-attention dynamies). With Assumption (1b) (ie., T +00), Eqn. f becomes:
2m ? n (m)
Σ
diag(fn)
Σ Bun (frf. T)fm',
(10)
Proof. Taking long sequence limit (T +0o), and summing over all possible choices of next token 21+1 = n, plugging in the backpropagated gradient (Bqn. 9 into the dynamics of with query token m (Eqn. 4.. we arrive at the following:
Pr. ||c,112
2m
12
Σ
diag(c,)
Y (œrr4ill (il))
(94)
RE# (m)
727
Σ
(..)
Σ
Brui frx
(95)
n (m)
= 727
Σ
diag( (nn))nnntt - T) Σ
Bun Swe
(96)
nEW (m)
n'
29

---

Note here we leverage the property that Pj f = 0 and Pc., = J
Theorem 2 (Fates of contextual tokens). Let Ger be the set of common tokens (CT), and (D)
be the set of distinct tokens (DT) that belong to next token n. Then if ('Assumpptin 2 holds, under the self-allention dynamics (Eqn. (O), we have:
(a) for any distinct token E GDI (n), 2md > O where m = (n));
(b) if |Ger] = and at least one next token n E (m)) has least one distinct token, then for the single common token I GCT, 2ml < O.
Proof. For any token we have:
Emd = 72"
Σ
fai Σ Ban fne Swel
(97)
nEW ((7r))
12' nigeee
Distinct token. For token I distinct to by definition, for any n' * n. P(Imm,))) = 0 and fie((t) & P( rt') exp(zmi) = 0. Therefore, we have:
Emi =1z7sp Σ Bon Su -n 12.2.1. >
(98)
Note that Emd > O is achieved by &n > O from Theorem
Common token. For any query token m, consider n E (m) and n' * if n and n' does diagg(f))... = O. When and n' overlaps, let not overlap then diag(f.) (f. 0} be the subset of common tokens shared between and n', Ger(n, n²) == (1 P( P(I
I)fm'
since |Gerl = 1 and 0 Gcr(n, nn) S Ger == Untn' Gcr(n, n'), we have |Ger [Gr((,, nn) = and E Gnr(n, i.e., the common token l is the unique overlap. Then we have:
Sni [UTT fn -wu] = (J. Sw ST Sn -(1 Pa)
(99)
So we have:
Emi -12
Σ (1- nEv-1((m)
Pui Σ Bun' fn'
72
Σ nEW-1 (m)
Pu
(100)
a'/n
Since (t) >> 0, the only condition that will lead to 2mi = 0 fa = However, since at least one such has another distinct token I', and thus fny > by normalization condition, fat 1 and thus Em! < O.
Note that for multiple common tokens, things can be quite involved. Here we prove a case when the symmetric condition holds.
Corollary (Multiple CTs, symmetric case). If Assumption holds and assume
(f) Single query token my For any next token [K], vpn) mp-
(2) Symmetry. For any two next tokens n', there exists one-to-one mapping that maps token 1 GDT(n (n) to GDD( so that [P[[m) = (()()))n));
(3) Global common tokens with shared conditional probability; i.e., the global common token set GeT satisfies the following condition: for any GeT, P((]n) Pt. which is independent of next token n;
(4) The initial condition 7(0)=
Then for any common token I E GCT, Emod < O.
Proof. Since there is global query token mg, we omit the subscript mo and let z == Smo.l
We want to prove the following induction hypothesis: for any (a) (t) = (t) for n, n' which are next tokens that the distinet token (and I') belongs to, and (b) the normalization term
(t) = E, (t, = o² (t), i.c.. it does not depend on
30

---

We prove by induction on infinitesimal steps öt. First when = O, both conditions hold due to the initial condition 20) Then we assume that both conditions hold for time t, then by symmetry, we know that for any n and any distinct token 11 E (DDT(n)::
3, (1) 1
Σ Tt' fnn
Sm Sw n2rlhals
Bugne Bugn = =12(1) n'zin2
Σ
(101)
where l/ = 6(1)) is the image of the distinct token 41 This is because (1) f fn EieGer pf (x)((2 is independent of nı and n' by inductive hypothesis, therefore, B
/o²(t) is also independent of its subscripts. And (2) fhul == / = /0 = fmals'
Therefore, 21 (1) (1). which means that 21, (t') mgg (t') for t' - öt.
Let Ger(n) m22 == (1: P( )P(1)22) 0} be the subset of common tokens shared between n and ngg then for their associated ni and nlg, obviously Ger(n1, n2) C GeT and we have:
Ong (1') = Σ am (1)
Σ ((((m))xp(((2((())
(102)
Σ
pd (01lm))exp((1) (t')
E
(1())) ((Im)))xp(22)1 (103)
heGor(m)
IEGor(m):n))
Σ
(x0(1))l2))xp((2220 +
Σ
pt exp((2((()) (104)
()G
(EGor(mI:ng)
= Σ
(12ln2))exp(22 (1)) +
Σ
(1 exp))xx(((2) exp ((') (105)
12E0pr(m2)
(EGeT(m):z))
Ong (t')
(106)
So we prove the induction hypothesis holds for t' + öt. Let 8t + 0 and we prove it for all f.
Now we check the dynamics of common token / E GeT. First we have for any n + n', (t)
& (t)/o?(t) pî (xp(22(t))/oo(t) (t) = (t) and thus fni(t) fwi((t)
fi(t)= > O, therefore:
= Jns [(fT [w] fw) <0
(107)
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 we know Ban' > and independent of the subscripts. Therefore, E < O.
O
Theorem 3 (Growth of distinct tokens). For a next token and its two distinct tokens I and i', the dynamics of the relative gain (a == fa (t)/ (t) (t) (t)- (t) has the following analytic form (here the query token m = (n) and is uniquely determined by distinct token l):
(1 F2/rin (0)e ((0)2(2))))))) =t)
(11)
where xt(t) := (t)m (0) is the growth factor of distinct token l. If there exist a dominant
token to such that the initial condition satisfies Tlo (0) > 0 foror all its distinct token I # lo, and all of its common tokens I satisfy 7nt < O. Then both zmt (t) and fnto (t) are monotonously increasing overt t, and
e (0 S X1, (1) S (B.(1))
(12)
here Bn(t): == 1l So En (t')ddt' Intuitively, larger B,, gives larger Tlp/1jn and sparser attention map.
Proof. Let m (n) be the query token associated with next token n. For brievity, omit subscript m in the proof and let 21 := Zmi-
First of all, for tokens and l' that are both distinct for specific next token n, from Eqn. 98 it is
clear that
a ry/pnn + (rifein (0) + (1) (0))
(0=((0)))
(108)
Ey
31

---

This means that 2()()==((0)); 21 (rife)n (0) 1)e 2(2 ((()))r(0)) 21
(109)
Integrate both side over time and we get:
2((((0==(00)
(r1/r1n(0) + 1) ( (1) i]
(110)
From this we can get the close-form relationship between bet (t) and 2x()):
T1/M (t) (1)=
(111)
Now let lo be the dominating distinct token that satisfies fl,/t n (0) for any distinct token 1, then
we have it, > 0 due to Theorem
for any token l' * lo that is distinct to n, we have:
Mo/t' 11 (03,)) 2
(112)
for any common token 1,, since by
O, we have
d
Culo ne
p2a (lo
|1 dt
2 -ev) 2 2v) 0
(113)
pe (2
Therefore, we have:
d
d
dt (f
>0
(114)
dt
M M+ E /o/o
Therefore, Pal, (t) is monotonously increasing, Combined with the fact fai, (t) S due to normal-
ization condition [|5f/|12 1, we have:
(t) 2 1/
Sio Pai (1))(()) 2 fal, (1)
(115)
Integrate over time and we have:
B(t)= ( 24 (t') =2, 21,(t) (t) 2% (0) f (0)B(t)
(116)
where B(t) == 122 fo (n't')df dt', Plugging that into Eqn.
111 and we have:
é (0) (to - 2B())
(117)
F Estimation in Sec. 6
Theorem (Phase Transition in Training). 1he dynamics of the single common token znd satisfies 2nd = Co "122(()ee" 42md and (t) = Cp (tmm,, then we have:
(t) {
In Coo 2(M-1)2
t << to = Kln.M
KM
(14)
In (Co + 2K 1)2 az In (Murl/K)) 12 40 = 2(1 1)) M
M
As a result, there exists a phase transition during training:
Attention scanning. At the beginning of the training, 7t) O/yytKK) and Bn(t) se 2 In (pt + mzt2) = (II This means that the growth factor for dominant token To is (sub- Ylinear: Xto 2 2252, (L.(( (K- 21 10.5 fhip (0) and the
attention on less co-occurred token drops gradually.
32

---

Attention snapping. When 2 to := 2(1 + 8') In M/Ty with 8' = e M ) r(t)
O K Kin(nyt/K) and (1) = (luInn)) Therefore, while B,, still grows to infinite, the growth factor XI, (t) = (In t) grows at a much slower logarithmic rate.
Ty
Proof. Since every next token n shares the same query token m, we omit the subscript m and let
2md-
We start from the two following assumptions:
& -Co exp(Azi)
(118)
(t) = Co' ex)()z)
(119)
Given that, can derive the dynamics of (1) and (n)::
exp(-Az)) = -Cp (2)(1
(120)
d dexp(- -4zz) = nz7()dl
(121)
exp(-4z)) = -47 nz / ((t)dtt' +1 (use 2)( = 0)
(122)
Let r(t) := 1l 1o ((')de', then r(0) = 0 and (t) (()(()) Therefore, we have
7(t)
(t) = (t) = Co + 4T
(123)
and thus Bn(()= = nz So (n''')tt' can be integrated analytically, regardless of the specific form of (()):
(f) = 12 L' (t')dtt
dr Co + AT(t) Lo" Co+ 4r In(Co + (n(C)+4T(tt)
(124)
(MI+)))(MM(()K) (Theorem 1.. Note that (if treated in continuous time step) is strictly monotonically increasing and satisfies h(0) = 0, (t(Kk) = 7/y (M + and Lemma
Recall that >()) (M-1) (M-1))
exp(Mh((/KK)) dt/K (Lemma and
7h, we can let >(h) M-1+exp(MR)
(M-1)²A and get:
T(() = 1z
70')ae'
(125)
t=
Jman ((t/kK)
M + exp(Mh') dh'
12 K
7(h)
(126)
ny
((t/k) R'ah'
12 K(M - 1)2
(127)
"/v
h(0))
n K(M 1)
(2(1/K)
(128)
T/Y 2
Therefore, (1) has close form with respect to h. regardless of the specific form of h(t):
(t) = - In Co + 22 K(M 1))22(1/K))
(129)
I/Y
(1) When / <<6 <=;== K KInMM)/m from Lemma 9 we have h(t/K) = (1 (1+o(1) We
neglect the o(1) term and denote v = V/ / /1z, then we have when L < l6:
2(M - 1)2 2
(t) In Co
(130)
vKM²
And B, In (Co + 2K - - In (n
33

---

== 10.0, M= 1000
v= 1.0, M 10000
28-
14.
2.
1.2
85-
0.6
22- 5.
drn 0.0.yy
8.1
Gr 2 20.00-000 20.0
2.0 20 2.0
40.0
5.0
5.0
2
20
Figure 9: Numerical simulation of (L) with changing 7 and fixed 1i2/ /1/ The dotted line denotes the transition time to, and B,, (to) marked with the solid dot is independent of 122
(2) Similarly, when > to == 2(1 + wi) M/ny with W1 O(In In M/ M) is defined in Lemma 8 from Lemma we have h(t/K) = (1 + ((Mnyt//) We neglect the o(1) term get when to:
1 Bn((1) = 7 In (Co Co Co+ DM²
2K(M - -1)2 In2 (n2(Mr)/K))
(131)
From this we know (to) = In(Co + 2K(M 11) In (2 w1 In M)). It's interesting to find that B,, (fo) just depends on K, M and v, and thus fixing v and changing 7z will
not influence the value of B,, (to), which means that the main difference between B,, is arises at the stage > to. This matches the results in Fig.
(3) Finally, estimate (t) when is large. When is fixed and t > (Myr) exp(1/V20), we have
1 In 11 12) K
(t) (1+0(1))
In In Mnzut,
(132)
θ In In(
(133)
K
Therefore, from Eqn. 133 we get:
(a) Fix larger 7/2 result in larger (t) and sparser attention map.
(b) Fix vzz. larger (i.e., larger v/y)) result in smaller (t) and denser attention map since In In(x)) is much slower than In(x).
These match experimental results in the main paper (Fig. 6
G Experiments
We use WikiText 25 dataset to verify our theoretical findings. This includes two datasets, Wiki- Text2 and W0k We train both on 1-layer transformer with SGD optimizer. Instead of using reparameterization Y and Z (Sec. (Sec.3.T). we choose to keep the original parameterization with token embedding U and train with a unified learning rate 77 until convergence. Fig. I0 shows that the av- eraged entropy of the self-attention map evaluated in the validation set indeed drops with when the learning rate 7 becomes larger.
Note that in the original parameterization, it is not clear how to set ny and n/2 properly and we leave it for future work.
Furthermore, we use the recall-threshole relationship to reshow that smaller "jy /rzz and larger 7z will result in sparser self-attention map as Fig. I1 and Fig. 12 Here we use some thresholds to retag every entry of the attention as a 0-1 value based on 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
34

---

wikitext2
wikitexti
2.0
Averogie entopy
1.0
0.5
Average self-attention map entropy over the validation sets in 1-layer transformer after training. when the learning rate and changes. Note that higher learning rate leads to higher B,, (t) and thus low entropy (i.e., more sparsity). which is consistent with theoretical finding (Sec. 6 All the experiments repeated in random seeds. Error
Figure 10:
1-std shown the ligure.
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 sparser attention map will result in slower descent rate of the recall-treeshol. line since sparser attention corresponds to fewer distinct tokens with higher attention weights, and the results of Fig. and Fig. 12 match that of Fig.
H Technical Lemma
Lemma 18. Let h =
hM] E RM is some M M-dimensional vector, hx ==
(Pay 1 E R- is vector selected by input sequence X, then given event FT
m, 4T+1 = n, there exists some qm,n = 191mm, 42mm. 4 [mdmmm)] and
RRM so that q 2
T-1 X hx Σ
Ummhte = qm. o
(134)
M
1
ddigg diag(hx))
Σ
hmmne = diag mm. h)
(135)
T
1=1
where mm satisfies Eie Ammm = 1. And with probability at least 8 we have
In(2//)
In(2/5)
max 0, 0. m)
V -1)
2(
< mmmm I P(11m, n)
V 2(T - 1)
(136)
And thus q1/m,,n + P(4m, when +0c.
Proof. Given that IT m and IT+1 = n., then we have
1 T
M
T-1
M
T - X hx T-
h
Σ
Σ I|. 1|x=11
(1m==1) =1) - 1)
h = Σ hymn (137)
hic:
T-1
1=1
And similar equations hold for X T x Then consider the case that the previous tokens are generated by conditional probability P((Im,,)) n) as the data generation part, so T[x; I]. 71 E [T- 1 are i.i.d. Bernoulli random variables with probability P(]mm, n)) and T mm satisfies binomial distribution. By Hoeffding inequality, we get
P(a PP(1 m, n)) < 2 2exp( -(T 1)(2)
(138)
Then we get the results by direct calculation.
35

---

01
1.20
0.63
1.40
0.53
1.25
0.93
0.50
2028 0.000 2.02 6:23 6035 0:038
0.628 3232 6602 3404 6:00 6.000
6.023 0.030 0.002 3:04 0.006 1200
6.08
3.000 0.332
0.094
3836 0.000
1.75
6.8
5.70
1.00
0.7
2.
2.50
6.4-
0.05
2012
0021 0605
6035 0.000
6:12
0223
0.004
6915
0.112
2.05
0056
0.00
6.7%
0.70
0.6
6.5
1.50
0.00
0115
1.0
0.8
0.0
Mouthoo
0.-
24
3.063
1000
0.3045
2.0045
Figure 11: Recall value of attention all distinct tokens versus threshold with changing learning rate ratio my izz Smaller Ty nt/nn /nz corresponds a smaller descent rate and thus sparser attention.
36

---

02 0.5
0.7
83
36-
3384
321
3,608 ndss zaba ndss subs 0d3s
6:0 0:332 do2 6:534 dis 6:538
6d08 nete eés n.est ede ó
3808 odaa 3.52 adas 1dbs nd3s
69-
3.-
3.7
277
2.6
66-
55
32-
7.-
1.015
3.-
3.7
2.6
66-
35
3.4
1.0
5.
Moothoo 1003
78-
0.-
2.a
47-
2.6
3.0
66
35
54
03045
Figure 12: Recall value of attention all distinct tokens threshold with changing learning rate 72.
Larger 7/2 corresponds smaller descent rate thus sparser attention.
37


Reviews:
[{"review_id": "Npozm1jwnk", "reviewer": "Reviewer_sxfK", "summary": "The 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.  ", "questions": "Again, 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", "limitations": "I believe assumptions and limitations are discussed. ", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "I 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. ", "weaknesses": "The 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments! Here are the answers:\n\n> Lemma 1: gradient dynamics of batchsize 1 -> Transformer trained with single example only. \n\nFirst of all, in machine learning, training with batchsize = 1 means that in each gradient step, only one sample is used to update the gradient, it does not mean we train on a single example only, since each time we could use a different sample from the dataset. \n\nFurthermore, the fact that Lemma 1 takes the form of batchsize = 1 does not mean it cannot be applied to batchsize > 1: we can simply sum over sample index i (omitted in Eqn. 3) to get the gradient of Y and Z. \n\n> 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.\n\nThis is a common reparameterization technique used when analyzing Transformer and many previous theoretical works, including works [2] suggested by **TXYd**, leverage this technique. In [2], key/query matrices are merged to one (in their Eqn. 4), q is defined as W_kW_q^T P^T (in their Eqn. 4) and its dynamics is computed instead of the dynamics of prompts embedding P. In [3] the key-query weights are merged into one matrix W and its dynamics is studied. In [A], similar to ours, they use X to replace QK^T as the variable to be optimized and study the property of X in optimization instead. \n\n> Why choose to integrate layernorm in your analysis? \n\nLayerNorm plays an important role in our analysis. LayerNorm provides the additional projection operator $P^\\perp_{X^T b_T}$ (as shown in the proof) that will cancel out one term in the derivative of softmax (line 539 in supp material), leading to much simpler analysis (e.g., Eqn.95 in supp material). We will make it clear in the revised version. \n\n**Reference**\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 (Jun. 6)\n\n[3] Tarzanagh et al., \"Max-Margin Token Selection in Attention Mechanism\", arXiv’23 (Jun. 23)\n\n[A] S. Li et al, The Closeness of In-Context Learning and Weight Shifting for Softmax Regression, arXiv\n\n"}, {"review_id": "loJvyW4CGv", "reviewer": "Reviewer_zgiA", "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).\n\nThe reviewer didn't check the proof in mathematics for this manuscript.", "questions": "N/A", "limitations": "addressed", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. 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.\n\n3. 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.\n\n2. 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.\n\n3. In the training and fine-tuning of the Transformer model, Adam/AdamW is used more often than the SGD analyzed in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments!\n\n> 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. \n\nWe emphasize that most of the experiments are focused on verification of the theory. Also we tried a 3-layer transformer on WikiText-2 and showed the behavior of attention sparsity predicted by the theory. Overall our work is mainly theoretical and we focus on checking whether our theoretical prediction is correct or not, so we mainly focus on 1-layer transformer.  \n\n> Adam/AdamW are used more often in practice \n\nNote that previous and concurrent works [1][2], some works published in ICLR'23/ICML'23, also focus on the setting of SGD training in 1-layer transformer. While Adams are used in practice, SGD is a base case that we want to understand well first. \n\nAlso, in this work, we empirically show the behavior of Adam optimizer in our synthetic setting (Fig. 5), which is indeed different from SGD and has very interesting behaviors. In Adam, the frequency bias seems to be controlled not only by the co-occurrence frequency of key and query tokens, but also controlled by Adam’s learning rate. This could explain why we need cosine learning rate for Adam, in order to sweep through possible co-occurrence frequency pairs. Note that this phenomenon is not mentioned by previous works, to our best knowledge, and can trigger interest in the community. \n\nWe leave a rigorous analysis on Adam to future work. \n\n> Assumptions may limit the generalizability\n\nWe want to emphasize that these assumptions are reasonable. People already show decoder only approaches also work reasonably well in no positional encoding [B] and currently both GPT4/Claude and the LLM community is now exploring much longer sequence input [C,D] (e.g., 32k, 65k or longer). We explained the third assumption ($\\eta_Z \\ll \\eta_Y$), which is a technical assumption for the theorems, in the overall rebuttal. \n\n**Reference**\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\n[B] A. Kazemnelad et al, The Impact of Positional Encoding on Length Generalization in Transformers, arXiv'23\n\n[C] S. Chen et al, Extending Context Window of Large Language Models via Positional Interpolation, arXiv'23\n\n[D] J. Ding et al, LONGNET: Scaling Transformers to 1,000,000,000 Tokens, arXiv'23"}, {"review_id": "M5ml2ZNScA", "reviewer": "Reviewer_TXYd", "summary": "This 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.", "questions": "1. 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", "limitations": "There is no potential negative societal impact of their work. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks the reviewer for the insightful comments. Here are the answers. \n\n> I feel the decoder layer is not formally defined. In around line 92, only the self-attention layer is introduced. \n\nWe define the decoder layer $Y$ and self-attention layer $Z$ after reparameterization $Y = UW_V^TU^T$ and $Z = UW_QW^T_K U^T / \\sqrt{d}$. Then we study the dynamics of $Y$ and $Z$ directly. \n\n> Overclaimed sentence \"we are the first to analyze the attention dynamics and reveal its inductive bias on data input\"\n\nWe acknowledge that previous and concurrent works[1][2][3] also show that self-attention attends to relevant tokens, and we will definitely reference the works and tune town our claim! In comparison, our work analyzes a lot more phenomena in 1-layer transformers that belong to frequency and discriminative bias, which has not been brought up before. For example, sparse attention patterns are connected with the co-occurrence frequency of contextual token and query, characterization of such connection over training including two stage behaviors of attention scores, etc. Please check the main rebuttal.\n\n>  A comparison and discussion between the manuscript and these works\n\nPlease check the overall rebuttal for an overall comparison. See below for detailed comparison:\n\n**Comparison with [1]**\n\n**Setting, Assumptions and Conclusions**. [1] 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 pos/neg pattern, [1] gives a generalization bound when the number of hidden nodes in FFN is large, and at the same time, shows that the self-attention attends to relevant tokens and becomes sparse (if #relevant tokens are small). \n\nIn 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. \n\n**Techniques**. The techniques used in [1] are based on feature learning techniques applied to MLP [R4 etc]. 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. 127)), avoiding complicated bounds in [1]. This allows us to characterize the converging behavior of self-attentions when $t->+\\infty$. To our best knowledge, our framework and techniques are novel, which is also acknowledged by the reviewer. \n\n**Comparison with [2] (published on Jun. 6, after the submission deadline)**\n\n[2] focuses on 1-layer attention-based prompt-tuning, in which some parameters of the models are fixed (Wp, Wq). 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, [2] 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.\n\nIn 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 [2] 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.\n\n**Comparison with [3] (published on Jun. 23, after the submission deadline)**\n\nCompared to [2], [3] also analyzes the dynamics of the query-key matrix $W$ and the embedding of a single tunable token $p$ (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 $p$ is not an attention since its norm can be shown to go to infinity over training.  \n\nIn 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.\n\nWe thank the reviewer for the references and we will put the detailed comparison, as well as reference these papers in the next revision. \n\n> Nonlinearity after the decoder $Y$.\n\nAdding nonlinearity right before cross entropy loss will make our analysis a bit more complicated but not impossible. Specifically, the nonlinearity will modify the back-propagated gradient from cross entropy loss, and Theorem 1 will take a different (and maybe more complicated) form. For simplicity, we choose not to add the nonlinearity layer for this paper.  \n\n> Possible typo in equation 15. \n\nWe indeed missed the term $X^T$ within LN() and will fix.\n\n**Reference**\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'23 (Jun. 6)\n\n[3] Tarzanagh et al., \"Max-Margin Token Selection in Attention Mechanism\", arXiv’23 (Jun. 23)\n\n[R4] Z. Allen-Zhu et al, Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers, NeurIPS’19"}, {"review_id": "usib86ibSV", "reviewer": "Reviewer_ZSxf", "summary": "This 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", "questions": "1. 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.", "limitations": "It 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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This 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.", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thanks the reviewer to give insightful comments!\n\n> 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. \n\nWe explain the intuition of the Assumption 2 (Weak correlation) as below:\n\n**Regarding upper bound**.  $|\\lambda_{\\max}(E)| < 1/K$ ensures that $f_i$ are almost orthogonal to each other since $E_{ij} = f_i^\\top f_j < 1/K$ holds for all $i \\neq j \\in [K]$. And this means that for different sequence classes, the conditional probability of each contextual token given the query/last token is quite different from each other. This scenario is common in real cases, in particular when the vocabulary size $M$ becomes large, and each sequence class may only relate to a small portion of the tokens. \n\n**Regarding the lower bound**. $|\\lambda_i(E)| \\ge 6/\\sqrt{M}$ illustrates that all tokens can be common tokens with small conditional probabilities, resembling a uniform distribution. This condition is only used in formula (93), which proves property (a) of Theorem 1. Specifically, $E_{nn}'$ and $u_{nk}$ in (93) are hard to be estimated, and we use this lower bound to prevent the worst case and prove that $\\xi_n > 0$ always holds.\n\nAdditionally, even if Assumption 2 doesn't hold, Theorem 5 in the Appendix shows that formula (9) and property (b) in Theorem 1 still hold. We just need Assumption 2 to ensure $\\xi_n > 0$ and to maintain the favorable properties of distinct tokens as discussed in Sections 5 and 6.\n\n> 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\nAs mentioned in the reply to reviewer **sxfK**, Lemma 1 gives a gradient formula when batchsize = 1 for simplicity, and the gradient of batchsize > 1 can be easily obtained by summing over different samples, and our framework applies to large batch sizes as well. We indeed tested the case where the learning rate of the decoder is larger than self-attention (see Fig. 6) and the trend is consistent with our theoretical findings. \n\n> I don’t find much surprise in the result that Transformers sparsify and focus attention on unique tokens.\n\nNote that as mentioned in reviewer **TXYd**, previous and concurrent works [1][2], published in ICLR/ICML, gave theoretical justification on why Transformers gradually attend more to distinct key tokens, under SGD training of 1-layer transformer. In comparison, our work analyzes more rich phenomena in 1-layer transformers that can be categorized as frequency and discriminative bias, which has not been brought up before. For example, we analyze multi-class settings, we connect sparse attention patterns with the co-occurrence frequency of contextual token and query, characterize relative growth in details for all tokens that are relevant to the class label (rather than only compare relevant with irrelevant tokens [1][2]), characterize the complete training dynamics (rather than initial steps [2]) and summarize two stage behaviors of attention scores toward convergence, etc. \n\nWe want to emphasize that we want to provide a theoretical framework and lay the foundation of a rigorous theoretical understanding of Transformer. Most of the time there should be no surprise and the theory should explain what happens in practice, but in a more fine-grained and quantitative manner. Such quantitative explanations lead to new discoveries that could go beyond empirical findings and intuitions. \n\n> Residual connections \n\nFirst, note that none of the previous/concurrent works [1][2][3] analyze residual connections in their training dynamics. Here [1][2] are accepted in ICLR'23/ICML'23.\n\nIn addition, our framework can already incorporate residual connection as part of the input $f_n$ to the decoder layer $Y$ (Sec. 7). We could further characterize the scale of the coefficients $\\beta_{nn’}$:  when $\\psi(n) \\neq \\psi(n’)$, i.e., the last/query tokens of seq class $n$ and $n’$ are different, $\\beta_{nn’}$ should be much smaller. In this case, conclusions about distinct and common tokens can be further refined. For simplicity, the main paper focuses on cases without residual connection. \n\n**Reference**\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"}, {"review_id": "D2dQBQb2Rv", "reviewer": "Reviewer_nH6i", "summary": "The 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.", "questions": "Extending 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.", "limitations": "The authors have implicitly discussed the limitations. However, the broader impact was not discussed. Please add a broader impact section.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- 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.", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful and encouraging comments! We appreciate the suggestions! Here are the answers to the questions. \n\n> Attention patterns in multi-layer model\n\nWe haven’t systematically analyzed the attention patterns in multi-layer models yet. Initial experiments show that the attention score in the top-layer is not as sparse as the bottom-layer attention patterns. We will leave it for future work. \n\n> Will there be a specific pattern regarding the phase transition across different layers? e.g. slower phase transition in the last layer etc.\n\nWe haven’t done multilayer analysis systematically yet. The intuition is that the sparsity attention patterns in the lower-layer will lead to combinations of tokens that co-occur a lot, which then are further combined in the higher layer. So it is possible that top layer attention will freeze after bottom layer attention (and slower phase transition), since they provide the building blocks for the top.  \n\n> Is there a possibility that each layer may specialize in different roles? e.g. initial layers \"scan\" while latter layers \"snap\"\n\nYeah it is possible that different layers are in different stages, as mentioned in the previous answers. \n\n> (optional) Will these findings be consistent across different domains that use Transformer models? e.g. computer vision models, graph transformers etc.\n\nIt is possible since for different domains, the token combinations hierarchy are different. E.g., in vision, each patch may carry less amount of information, information tends to spread evenly across the image, and a deep and balanced hidden hierarchy is needed; while in language, a few tokens may carry the majority of information, and tokens may have nonlocal relationships, etc. Graphs from different domains may contain very different feature structures. Our goal is not to design specific models for specific domains, but to characterize how features are learned in (multi-layer) Transformer architecture, under different data distribution.   \n\n> Please add a broader impact section\n\nThanks for your suggestions. Here is one tentative version: \n\nOur work gives a framework to analyze the dynamics of a 1-layer transformer using the tools of nonlinear dynamics systems. Our finding that self-attention pays sparse attention to a subset of distinct tokens that co-occur a lot with the query can have many implications. First, it helps understand the inductive bias of self-attention, reduce spurious correlation (by removing tokens that are not likely to be attended to) and improve generalization. Second, it can also help understand the mechanism of hierarchical token combination, and potentially open the door of understanding multi-layer transformers. The two stage training process, namely “scan and snap”, can also be an important tool to understand how the multi-layer transformer works in practice. \n"}]

Author rebuttal:
We thank all reviewers for their insightful feedback.

We are glad to hear that reviewers agree that a rigorous framework/analysis of the training dynamics of Transformer is valuable, interesting, novel and timely[**sxfK**, **TXYd**, **nH6i**], with clear high-level intuitions [**TXYd**, **nH6i**] and experiments on both synthetic and real-world datasets [**zgiA**]. We will address the common questions below and will reply to each reviewer about answers to their detailed questions.  

**[ZSxf, zgiA, nH6i] Regarding the setting of 1-layer transformers**. We observe a mixed point of view around this aspect. Some reviewers think it is too simplistic (**ZSxf**, **zgiA**) while others (nH6i) think it makes the overall picture clear and helps understanding. 

From our point of view, analyzing a 1-layer transformer is a necessary step. As shown in our paper, there are many interesting and non-trivial behaviors even in this apparently simple case. To see this, let’s check previous and concurrent works [1][2][3] listed by reviewer **TXYd** asking for comparison. [1][2][3] are important works, [1][2] published in ICLR'23/ICML'23, and all focus on 1-layer transformers ( = 1 layer self-attention + 1 FFN). A common theme is that they show self-attention in1-layer attends to relevant tokens during training under various settings and data models (e.g., when the model is fixed and we fine-tune on soft-prompts), therefore, even with 1-layer, many phenomena are nontrivial and deserve dedicated work published in top-tier conferences. 

**Our contributions**. Compared to [1][2][3], our work is novel in the following ways:

+ Among relevant/distinct tokens of a sequence class, we characterize their relative growth quantitatively, showing that self-attention attends to *a subset of distinct tokens* that co-occur a lot with the query, leading to attention sparsity even among relevant tokens. In comparison, [1][2] only show high attention scores for relevant tokens and thus attention sparsity relies on number of relevant tokens in the assumption. 

+ We characterize complete training dynamics for multi-class settings, by *integrating nonlinear dynamics analytically*. We also summarize two stage behaviors of attention scores (scan & snap) toward convergence. In comparison, [1] relies on a much more complicated technique and [2] focuses on initial gradient steps and [3] assumes linear separability of tokens generated from two classes and uses max-margin SVM framework. 

We list the detailed comparison in the our rebuttal to reviewer **TXYd**.

Note that even with 1-layer, our work is already criticized as very dense (**sxfK**). In fact, it can be **unrealistic** to expect a rigorous and detailed formulation of a multi-layer transformer in a single 9-page conference paper, and at the same time, with a rich connection with real-world applications. It'd better to focus on simplified settings to make the take-home message clear. Therefore we focus on 1-layer Transformer and leave its multi-layer extension as the future work. We really appreciate reviewer **nH6i** for the understanding! 

**[zgiA, nH6i] Future directions to address multi-layer cases**. For multi layer, an important component is how the input tokens are combined together to form high-level concepts during training. Our work 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 10 contextual tokens are relevant to the query, the self-attention may pick 1-2 token to combine first due to attention sparsity. This can be regarded as a starting point to study how tokens are composed hierarchically. In comparison, showing that attention attends to all relevant tokens [1][2][3] may not suggest a hierarchical / multi-layer architecture, which is used in practice. 

**[ZSxf, sxfK, zgiA] Concerns on the strength of assumptions.** Understanding transformers in a mathematically rigorous manner is a highly nontrivial problem, and the assumptions we made are comparable with or even weaker than previous works. For example, [R1] analyzes positional attention with symmetric initialization, without considering input data. [1] models the data distribution as discriminative/non-discriminative patterns, similar to ours, assume hinge loss (rather than cross entropy loss), and perform SGD training. [2] also models the data distribution as relevant/irrelevant. [R3] models transformer dynamics near the initialization and freeze many parameters at random initialization. In comparison, we characterize the entire training dynamics with a few well-defined assumptions. We will list the detailed comparison in the reply to reviewer **TXYd**.

**[ZSxf, TXYd] Fast training of the decoder Y than the self-attention layer Z ($\eta_Z \ll \eta_Y$)**: This is a technical assumption in order to obtain Theorem 1, which establishes the relationship between the input signal $f_n$ of the decoder Y and the back-propagated gradient $g$, once the decoder $Y$ has been sufficiently trained. We indeed see works that use different learning rates at different layers in the empirical studies, e.g., [R2]. 

As a future work, we are actively looking for better approaches that can get rid of Assumption 2 and separate learning rate in the analysis. 
 
**References**

[R1] S. Jelassi et al. Vision transformers provably learn spatial structure. NeurIPS’22. 

[R2] E. Dinan et al. Effective Theory of Transformers at Initialization

[R3] A. Bietti et al, Birth of a Transformer: A Memory Viewpoint

[1] Li et al., "A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity", ICLR 2023.

[2] Oymak et al., "On the Role of Attention in Prompt-tuning", ICML 2023 (Jun. 6)

[3] Tarzanagh et al., "Max-Margin Token Selection in Attention Mechanism", arXiv’23 (Jun. 23) 

Meta-review: All reviewers agree that this paper studies the important problem of analyzing training dynamics of Transformer models, but under very strong assumptions, such as, 1 layer transformer model without positional encoding, different convergence rates for different parts of the model, no non-linearity in the MLP layers and the results are in the limit of infinite sequence length.

While these are pretty strong assumptions, given the popularity of Transformer models in practice, the paper makes small but important progress towards better understanding of the Transformer dynamics. 

Reviewers also raised some concerns about comparison to some of the recent papers analysing Transformer training dynamics. Authors provided a detailed response on the difference with their work and promised to update the paper with detailed comparison.

Overall the paper is on borderline and I suggest acceptance given the importance of the problem studied and the results in the paper. I suggest authors to update the paper to improve presentation following review suggestions. Further the paper (including abstract) needs to state the assumptions correctly as the results are in the limit of infinite sequence length not just long sequence length.
Decision: Accept (poster)