Non-asymptotic Convergence of Training Transformers for Next-token Prediction
Authors: Ruiquan Huang
Yingbin Liang
Jing Yang
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformer
Convergence Rate
Next-token Prediction
Self-attention
Implicit Bias
Primary area: learning_theory
TL;DR: We show fast convergence rate of training transformers for next-token prediction.

Abstract: Transformers have achieved extraordinary success in modern machine learning due to their excellent ability to handle sequential data, especially in next-token prediction (NTP) tasks. However, the theoretical understanding of their performance in NTP is limited, with existing studies focusing mainly on asymptotic performance. This paper provides a fine-grained non-asymptotic analysis of the training dynamics of a one-layer transformer consisting of a self-attention module followed by a feed-forward layer. We first characterize the essential structural properties of training datasets for NTP using a mathematical framework based on partial orders. 
Then, we design a two-stage training algorithm, where the pre-processing stage for training the feed-forward layer and the main stage for training the attention layer exhibit fast convergence performance. Specifically, both layers converge sub-linearly to the direction of their corresponding max-margin solutions. We also show that the cross-entropy loss enjoys a linear convergence rate. Furthermore, we show that the trained transformer presents non-trivial prediction ability with dataset shift, which sheds light on the remarkable generalization performance of transformers. Our analysis technique involves the development of novel properties on the attention gradient and further in-depth analysis of how these properties contribute to the convergence of the training process. Our experiments further validate our theoretical findings.

Non-asymptotic Convergence of Training Transformers for Next-token Prediction
Ruiquan Huang Penn State University State College, PA, 16801 rzh5514@psu. edu
Yingbin Liang Ohio State University Columbus, OH, 43210 liang. 8898osu. edu
Jing Yang Penn State Univeristy State College, PA, 16801 yangj ing@psu edu
Abstract
Transformers have achieved extraordinary success in modern machine learning due to their excellent ability to handle sequential data, especially in next-token predic- tion (NTP) tasks. However, the theoretical understanding of their performance in NTP is limited, with existing studies focusing mainly on asymptotic performance. This paper provides fine-grained non-asymptotic analysis of the training dynam- ics of a one-layer transformer consisting of self-attention module followed by feed-forward layer. We first characterize the essential structural properties of training datasets for NTP using mathematical framework based on partial orders. Then, we design two-stage training algorithm, where the pre-process pre-processing stage for training the feed-forward layer and the main stage for training the attention layer exhibit fast convergence performance. Specifically, both layers converge sub-linearly to the direction of their corresponding max-margin solutions. We also show that the cross-entropy loss enjoys linear convergence rate. Furthermore, show that the trained transformer presents non-trivial prediction ability with dataset shift, which sheds light on the remarkable generalization performance transformers. Our analysis technique involves the development of novel properties on the attention gradient and further in-depth analysis of how these properties contribute the convergence of the training process. Our experiments further validate our theoretical findings.
Introduction
The transformer architecture (Vaswani et al.) 2017) has revolutionized the field of machine learning, establishing itself as foundation model for numerous applications. including natural language processing (NLP) (Devlin al. (2018), computer vision (Dosovitskiy al. 2020), and multi-modal
signal processing (Tsai et al.] 2019). In particular, transformers achieve tremendous empirical success in large language models (LLMs) such as GPT-3 (Brown al. 2020) Despite the empirical success, limited theoretical understanding of transformers have caused a series of critical concerns about their
robustness, interpretability, and bias issues (Bommasani al. 202 Belkin 2024).
To overcome these issues, recent advances in transformer theory have investigated the convergence of training transformers under theoretically amenable setting such linear regression (Mahankali et al.) 2023, Zhang et al.) 2023 Huang et al. 2023) and binary classification (Tarzanagh et al.) 2023b a Vasudeva al. 2024 Li et al. 2023). Nevertheless, one of the fundamental task LLMS and other generative models next-token prediction (NTP), which involves predicting the next word or token in sequence, given the previous tokens. In NTP, a few recent theoretical studies have started to investigate the training dynamics of transformers (Tian et al. 2023a Li et al. (2024). However, those works lack of fine-grained non-asympioiic convergence analysis of the training process, posing the following open questions for further investigation:
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

How fast does the training of transformer converge in NTP?
In addition, a pre-trained transformer empirically exhibits non-trivial generalization ability. A follow-up question from a theoretical point of view is that
Can we show the generalization capability of a trained transformer unseen data?
In this paper, we take a first step towards addressing the aforementioned questions by studying the training dynamics of a single layer transformer consisting of self-attention layer and feed-forward layer for NTP. We summarize our contribution as follows.
We develop mathematical framework based on partial order to formally characterize the essential structural properties of the training dataset for next-token prediction. In particular, we introduce a realizable setting for training datasets where the loss can be minimized to near zero, which admits a collocation and query-dependent partial orders. A collocation is a set of token pairs where each token is directly paired with its subsequent token. dheerddt partial orders is a set of partial orders where each partial order classifies tokens into three categories: optimal tokens, non-optimal tokens and compoaaae tokens. These structural properties define favorable max-margin problems on both the feed-forward layer and the self-attention layer.
Second, we design two-stage training algorithm based on normalized gradient descent. In stage of pre-processing, we use the collocation train the feed-forward layer. In stage 2, we use the entire dataset to train the self- self-attention layer. We show that the feed-forward layer and the query-key attention matrix converge sublinearly in direction respectively the max-margin solution for classifying next token from all other tokens in the preprocessing dataset, and to the max- mar-in solution for classifying the optimal from non-optimal tokens. In addition, the norm of the transformer parameters grows linearly, which further yields a linear convergence rate of the cross-entropy loss. Our two-stage algorithm decouples the training of the feed-forward and attention layers without losing optimality, as stage I's max-margin solution is judiciously designed to facilitate stage 2's fine-grained classification for optimal token prediction.
Third, we show that the trained transformer has generalization ability for making non- non-trivial prediction on unseen data. In particular, the transformer is trained to learn an extended query- dependent partial order, where the compoaaaaa tokens are inserted in between the optimal tokens and non-optimal tokens. Thus, the trained transformer will attend to non-comparable tokens if optimal tokens are not in new sentence and further make desirable prediction.
2 Related Work
Inspired by Brown et al. 22200), who demonstrated that pre-trained transformers learn in-context i.e., learn new tasks during inference with only a few samples a series of works focus on the expressiveness power of transformers (Akyürek al. [2022] Bai et [003] Von Oswald et al., 2023] Fu et al.) 2023 Giannou et al., [023 Lin et al. (023). These studies have shown that there exist parameter configurations such that transformers can perform various algorithms such as gradient descent. Additionally, Edelman et al. (2022) showed that transformers can represent sparse function.
Regarding the training dynamics and optimization of transformers under in-context learning, Ahn ct (0024) Mahankali et al. (203)); Zhang et al. (2023); [uunn et al. (023)) studied the dynamics of single attention layer, single-head transformer for the in-context learning of linear regression tasks. Cui et al. (024) proved that multi-head attention outperforms single-head attention. Cheng et al. (2023) showed that local optimal solutions in transformers can perform gradient descent in-context for non-linear functions. Kimm and Suzuki (2024) studied the nonconvex mean-field dynamics of transformers, and Nichani et al. (2024) established a convergence rate of Ö(1/L) for the training loss in learning causal graph. Additionally, Chen al. (2024) investigated the gradient flow in training multi-head attention. Chen and Li] (022)) proposed a supervised training algorithm for multi-head transformers.
(2024);
Another line of research focuses on the training dynamics of transformers for binary classification problems. Tarzanagh et al. (2023b al demonstrated equivalence between the optimization dynamics of a single attention layer and a certain SVM problem. While Tarzanagh et al. (202b a) only proved an asymptotic convergence result, Vasudeva et al. (024) improved the convergence rate -/1/ Li et al. (2023) studied the training dynamics of vision transformers and showed that the generalization error

---

can approach zero given sufficient training samples. Additionally, Deora et al. (2023) investigated the training and generalization error under the neural tangent kemne (NTK) regime.
For transformers trained on next-token prediction (NTP). Tian et al. (2023a) analyzed the training dynamics of single-layer transformer, while (Tian et al. (2023b) studied the joint training dynamics of multi-layer transformers. et al. (2024) demonstrated the asymptotic convergence of transformers trained with logarithmic loss function for NTP. Although these works provided valuable insights into the training dynamics of transformers for NTP, they did not provide the finite-time convergence
analysis, which is the focus of this paper. We remark that Thrampoulidis (022) studied NTP without transformer structure.
Our work is also related to the classical implicit bias framework for training neural networks (NNs).
In particular (Soudry et al.] (2018); Nacson et al. (2019): Ji and Telgarsky| (202T): Ji et al. (2021)
established convergence rate of gradient descent-based optimization. |Phuong and Lampert (2020)
Frei et al. (2022); Kou et al. (2024) studied the implicit bias of Rel.U/L.eak-RRl.. networks on
orthogonal data. A comprehensive survey is provided in Vardi (2023). However, these works focused
on classical neural networks, whereas we investigate the implicit bias of transformers for NTP.
3 Problem Setup
Notations. All vectors considered in this paper column vectors. We use 1{A}} to denote the indicator function of A, i... 1 {{A) = 1 if A holds, and [[A] = 0 otherwise. ||W] represents the Frobenious norm of the matrix W. For vector u, use [e]: denote the i-th coordinate of v. We use ((v) to denote the softmax function, i.e., [d(v]]: = exp(x:)/ E, exp(e] v), which can be applied to any vector with arbitrary dimension. We use {ei)ennni! denote the canonical basis of RRV,, i.e. [et]y = i = j]. The inner product (A, B3 of two matrices A, equals to Trace(AB
Next-token prediction. We consider the task of next-token prediction, which aims to predict the subsequent token token sequence given its preceding tokens. Formally, suppose that there exists a finite vocabulary set V C Idd that consists of all possible tokens, where d is the dimension of the embedding. Each token x E V is associated with a unique index 1(x) E 21.2..... MV] where is the index function. An L-length sentence X ( #2] E VL C Rddx. is sequence of L tokens, where L is an integer. We assume that the maximum length of sentences is Lmax. The subsequent tokens in sentences are generated from where Pi. generates the next token F1++1 given the sentence X for any < L < Limax The task of next-token prediction requires us to learn all models (pī, L given training dataset
of ground-truth model {pī V/ - V) 1.
Do (X)))L+ < Immax X FL,+1 E V}. Notably, if X = - 2z, Do. then for any l < L, ([ 22), 12)+) ) is also a training sample, since it follows P as well.
Decoder-only transformer. A decoder-only transformer is stack of blocks consisting of self- attention layer and feed-forward layer. For simplicity, we consider one-layer transformer, where the self-attention layer is determined by three matrices: W& E Rdxdi W E Rd xd and W, Rd namely key, query, and value matrices, and the ford-forward layer is determined by W. E ( Here d1, d2 are hidden dimensions. Mathematically, given the input X [2] ..), we write the one-layer transformer as (X) == @WW, W.X X (XX W,WWa E [0, 1]]], where := (W., W., Wi, Wq). and & is the softmax function. We note that the inner softmax function ái part of the attention model, and the outer softmax function & is the decoder that generates a probability distribution over V for token prediction.
Reparameterization. We reparameterize the transformer architecture by consolidating the key and query matrices into unified matrix Wka. such that Wk., = W&W, Similarly, we reparameterize the product of the feed-forward (W.) and value (Ww) matrices as a single matrix W... defined as W W.W. W.. v. Such reparameterization is commonly adopted in transformer theory works (Huang et al. 2023 Tian et al.) 2023a Li et al. 2024| Nichani et al.] 2024). Thus, the transformer under WIq 1]
those reparameterizatioon is given by Tx(X) @(W..
Cross-entropy loss. Given the training dataset DD and the transformer model, we seek to learn p. by minimizing (training) the cross-entropy loss C()) defined as follows:
[(A)= = iDol
Σ
log T T
(XDD
3

---

where I(xL+1) is the index of x1+1 in V.
4 Realizable Training Dataset and Two-Stage Algorithm
In this section, we first provide a mathematical framework based on partial order to formally characterize realizable training dataset for next-token prediction. We will then describe two-stage algorithm for next-token prediction that we study.
4.1 Realizable Training Dataset
We characterize a realizable training dataset via two structural properties, where the training loss can be made arbitrarily close to zero. We first provide some intuitions about those two properties.
Existence of "collocation". First, we note that if a sentence X = [_i,...,, Ir.] is a legal training sample, ([m1], is also in the training dataset. In addition, the output of a transformer given one single input token only depends on Wor. i.e. the feed-forward layer. Since training loss can be arbitrarily close to O, there exists sequence (Wi} such that lime- EtED, log e, (x) <(W(z) where ((x) is the index of next token of xr and the summation is over the case when x is the first token. Due to that o Wxx is a probability distribution, the equality holds only when is injective, since otherwise it an entropy of some distribution which is strictly greater than Therefore, there exists an injective map n V + V such that every sentence starts with must have unique next token n(z). We call the set of pairs {z, m(x)) ZEE a collocation. We remark that pi n.
Existence of "order". Second, let us consider the output of a transformer To given a legal sentence X = [x1,.... , , with the next token # p1 (X). The transformer first caleulates convex combination of 21....,II IL with corresponding weight ş exp(x= LL) for each S L. Then, the transformer outputs (EE, Wov.e))). Recall that the collocation forces zy to map to n n( thus (WW,a3r) has peak value at the coordinate equal to ((((2))) (the index of n(xr)). Hence, (X) can only have peak value at the coordinates within the set (I(n(x=))}e<. If the training loss can be arbitrarily close to 0, it is desirable have n (x E { Therefore, for those Fy with n(xx)) = TL+1. Ppt must be larger than by with n(re) # IL+11 Finally, it worth noting that fo depends on the final token FL. This observation motivates us to define query-dependent partial orders on V.
Definition (x% -prrtall order) Fix a token x². An x. ---ppaiiall order assigns an ordering relationship >as for certain pairs of tokens in V, and is created as follows. Let D be the set of all legal sentences in the training dataset that has the final token (query) x9, Then, for any pair of tokens x, x' E V,, we assign >x9 x' if there exists a sentence X [21,..., v1. v.)= DD and x. x' are tokens in X such that n(x) = 11+11 # n(x'), where F141 is the next loken of X.
Note that Definition is a "constructive definition" which might not be well-defined. However, as we are under the setting when the training loss can be arbitrarily close to O, the aforementioned discussion shows that if x >an I', then fo > (gg, where x = De and x'' = By in some sentence. Thus, exp(x: =9) > exp(x' Wkg 2),, which indeed need to be well-defined. Otherwise, have contradictions such as exp(=W 29) > exp(x' ()) < xxp(=Wka..)) Mathematically, a well-defined (strict) partial order on V satisfies two axioms (Yannakakis) 1982): (i) there is no x > x;; (ii) if x > x' and x' > x", then x > x". Thus, -arrtiaa order created by Do is well-defined for every x E V.
will
Finally, let us discuss the impact of query-dependeee partial orders on Do. For a given query x², the partial order >2 divides tokens in > into four disjoint types.
(Strict) optimal tokens. A token x is optimal, if there is no a' such that x' >29
Confused tokens. A token & is confused, if there exists x', x" such that x' >20 x > x". (Strict) non-optimal tokens. A token is non-optimal if there is no x' such that a 29 x Non-comparable tokens. A token is non-comparable if there is no x' such that x >za x' or
x' >>9 x.
'This also related to the maximal element in a partially ordered set. This is also related to the minimal element in partially ordered set.

---

In this work, we assume that there are no confused tokens. This assumption simplifies the problem. making tractable to provide explicit convergence in direction for training a transformer in Section In summary, we make the following structural assumption on the training dataset.
Assumption (Realizable training dataset) Do admits (i) a collocation {a, n(x))vev: (ii) well- defined query-dependent partial orders, where every x 299-partial order has no confused tokens.
We remark that combining the collocation and query- deeerddeeenndntt partial orders, training dataset as follows. For any sentence with only one token X = [xr]. the next token is n(r). For other sentences X [I1,.... .L),, let be optimal under the partial order >21, and then the next token of is (xz)). We next provide a simple example that justifies Assumption 1
regenerate the
Example Consider language system where the vocabulary consists of four tokens {S. O, P). where S, S. respectively stand for subject, verb, object, and punctuation mark. This system admits the commonly adopted word order (Dryer 1997): S. V. O. P. Let the training dataset be {SVOP, VOP, OPP, PSV}.
Let create the corresponding collocation and the query-dependent partial orders from the dataset. The collocation is {(S. V), (V. O), (O. (P. S)}. That is, if a sentence starts with a subject, then the next token is verb. Similarly, if sentence starts with verb, then the next token is object, and so on. The query-dependee partial orders are created as follows:
Partial order under query S. S>sP. Purtial order under query O. 0>0S, O>oV.
Partial order under query V. V>
Partial order under query P. O>pP.
Therefore, if sentence starts with (subject), the next token is V (verb) according to the collocation. Then, for the sentence SV, since the query is V and V>vS, the next token of the sentence coincides with the next token V,, V, which is exactly O (object). Finally, for the sentence SVO, following similar argument, the next token is P (punctuation mark). This example satisfies Assumption 7 and aligns with real- real-world scenarios. An illustration is provided in Figure]
decennedpenn partial order
Collocation
Dogs
Chase
Cats
Dogs (5)
Chase
Cats (0)
(P)
Chase
(5)
(V)
[O]
(P)
[V]
(o)
[S)
Figure 1: The left plot shows the mapping from sentence to the next token. The red rectangle indicates the optimal token in the corresponding sentence. The right plot shows the collocation relationship.
Additional notations of training data. It is worth noting that there only finite number of distinct sentences. For ease of presentation, we introduce the following notations. Suppose there are N distinct sentences in the training dataset Do indexed by {1,...; ; N). For each distinct sentence
(x), we calculate its frequency n (1)) E [0, 1] in dataset Do as (n)) DIE, (X=x("))
DD
Building upon this, with a little abuse notation, we use ((x(")) E V to denote the subsequent token of the sentence (") and (x((m)) to denote the index of n(X(")).
We further denote (n) as the final token of X(),, and let Te(X) = Wor XorX Xo X Then, the loss function [(0) can be rewritten as follows:
(()) = n (m)
Σ
log Σ exp (c. [exp(e.To(xre (cTTo(x(m))) - T - To(x(n))).
(I)
4.2 Training Algorithm
For the realizable dataset satisfying Assumption D we propose two-stage training algorithm using normalized gradient descent (NGD). The pseudo code of the algorithm is presented in Algorithm In Section 5 we show that the two-stage algorithm decouples the training of the feed-forward and
5

---

attention layers without losing the optimality. This is because the training in stage 1 is designed to yield suitable max mar-in solution, which will enable the training stage solve fine-grained classifcation problem and identify the optimal token for prediction.
In the first stage of pre-processing, we use the collocation set to train the feed-forward layer Wav For simplicity, we introduce the following notation for the training loss of the feed-forward layer. Given collocation {z, n(x)) JEV, which can be obtained through extracting all length-2 sentences in the training dataset Do we use normalized gradient descent to train Wov. Equivalently, the loss function can be written as
exp(e In(2x) W Wov.:)
Co --
Σ log WEV
Eng exp(e, W
where the self-attention elements removed because the attention matrices not trained here. Based on the above loss function, we initialize Wov = 0 E IR V|xd and subsequently take an update
at each time by NGD as in line 4 of Algorithm [
In the second stage, we fix trained feed-forward layer and train the self-attention layer based on the loss function given in Equation ①1 and using the entire dataset Da. Specifically, we initialize kgg E Rdxd and subsequently take an update at each time by NGD as line 7 of Algorithm
Algorithm staes Normalized Gradient Descent
1: Initialization: W:O 0 E RVV xd Wky 0 E Igdxd
2: Input: A collocation {ar, (()))=2)) and training dataset Dg, learning rate Too 77-
3: for E 10,1, 1} do
4: Update W((+1) as W wi(t)
Ww., (W
V w.. Co (W
5: end for
6: for E {0, (i 1 do
7: Update w-(++1) as w((+1) w(())
"[[ C( where ((1) (W:!? w
Vwn, C(0
kg. kg
8: end for
5 Training Dynamics of the Transformer
In this section, we present the convergence result for Algorithm 1 Before we proceed, we first introduce the following technical assumption, which has been commonly adopted in the previous
theoretical studies of transformers (Huang et al. 2023 Li et al. 2024 Tian et al. (023a)).
Assumption The vocabulary set orthornormal. Namely, the embedding has unit norm, i.e. ||>||| = L, and x' = holds for any distinct tokens x and x'.
5.1 Convergence of Training W..
To characterize the training dynamics of Wov, we observe that the collocation ((x) defines the following hard-margin problem:
/h arg min ||W]],
s.t. (e.- - e.) 21, Vv* In(x), In(z).
(2)
can be shown that g- Lo (BW*)) 0. Thus, the loss function Lo trains Wav to be the max-me maa-ina solution with W a distinguishing the next token nzx) from all other tokens in V.
Since Lo(-) is convex, have the following convergence result the training of (!(!)
Proposition 1 Let W be defined in Equation 2
Under Assumptions 1 let w!( be updated by
Algorithm 1
Then, for any Z 2, we have
I || tro and the following bound holds:
2| I
AA more general way various standard techniques developed in linguistic analysis
Lehecka 2015

---

Wo
Wo.
5 W | log (2))))
log/
[|
[
tno
Moreover, the loss function Lo satisfies that LW w I O(exp((not/(411)) ())).
Proposition states that during the training stage 1, the feed-forward layer W converges in direction to / ov at rate of ((log//t) which classifies the next token from all other tokens. In addition, since the norm of W increases linearly, the loss Lo converges linearly to zero, i.e., Co (W O(exp(-Col) for some constant Co.
5.2 Convergence of Training Wkg
Recall that after the training stage with T steps, obtain a trained feed-forward layer W
Then, we fix w:(?) and use normalized gradient descent to train Wkq- To characterize the training dynamics of the key-query matrix Wkq. note that each query-dependent partial order also defines hard-margin problem. Let I(n) C (1,..., I be the set of indices of the optimal tokens of X(). Recall that ze is optimal if there is Ie such that ze > x(1) 24 and In(x:)) In(X(()). That is,
(m) Wkq-X X -I should correctly classify optimal tokens zz and non-optimal tokens z This is formalized in the following problem:
Wh kg arg min [W]], s.t.
(or) De (1) W 2
JK E I(n), & I(n), Vn.
(3)
i given the well trained w:!?) will train Wkq towards
We will show that the loss function in Equation the max- margin solution Wi,, in direction for classifying between the optimal and non-optimal token. We further make the following technical assumption.
Assumption 3 For any sample X ((n), the number of optimal tokens is not less than the number of non-optimal tokens. Formally, for any non-optimal token x X(m), we have [[())] 2 Σ (xz = x}.
Assumption 3 consistent with practical and empirical observations, where optimal tokens often demonstrate higher relevance, making them more frequent in subsequent outcomes.
We now present the convergence result for the training of the key-query matrix in stage
Theorem
Let Assumptions hold. Let Wkg be the solution of Equation
Let < 0(1) and
Who be updated by Algorithm Then, for any t 2 have that W
|| kg || s tn. In
addition, the following inequality holds.
w(t)
kg
54NTA m max
W kg
log2
21
[
tn
Theorem I states that the key-query matrix Wky converges in direction to the max-margin solution W /tt/llWW. kg at convergence rate of O(log" t/t). We further show that the norm of Wky also grows linearly in i.e., [Walll l(t). Combining these results, have the following theorem the convergence of the loss function and the training accuracy.
Theorem (Loss Convergence) For any training sentence X()) = [m]] ,...,. (m)) let yp (n.t) x exp(x (n) w:(!)
(1) (n)) be the attention weight. Under the conditions in Theorem 1 there is absolute constant on such that when T 2 co [W* ( log((V)) log T/10 and L 2 Co D. Wiggl l log² I/n, the optimal loken weight satisfies min, Ee.Eu()) Pi.
(n.1) 2 (1 + Lmax exp (-(C1))) In addition,
the loss function converges linearty its minimal value:
((0 = L(W(T) w-(t)) ov kg & [V]ex]
-TCo (1
2 2T
(4)
-'=xxxx + exp C1t)
*For tixed W (T) the minimum loss value L C'TCoO(e-CR -Ctt for sufficiently large f, which further implies linear convergence
Then Equation
(4 implies (()()) )(0())-C*
-C'
7

---

where Co
no 2l 2 and C1 = 7/ (4Lmax [ 113).
Theorem shows that the training loss converges to its minimum value at a linear convergence rate. (L(lo)(/)e)) the optimal token weight is given by l/ + e) for any 0, which is close to 1. This implies that the trained transformer attends to the optimal token and thus outputs the correct next token n(x(?") with probability 1 O(e)).
Furtherm, for T ([llo(((/e))))))
5.3 Proof Sketch of Theorem
The proof consists of the following three main steps. The key proof step lies carefully analyzing the projection of gradient V Wie, ((((t) onto the token-query outer product x (x(*)) max max-margin attention weight matrix Wiq. and the trained attention weight matrix W
By analyzing (( (x) (XLT) we characterize the dynamics of
Step 1 (Lemma
attention weights. Using mathematical induction, we show that the lower bound of optimal token
weight is 1/ mmaxx
Step 2 (Lemma
6 Then, we show that the cosine similarity between the negative gradient and Wig is strictly larger than the minimum optimal token weight. Utilizing step 1, due the NGD update, the norm of the key-query matrix w/( can be shown to grow linearly. kq
Step (Lemma 7 Finally, we carefully compare the difference between the projections from gradient to the trained attention matrix and max-margin attention matrix. By separately evaluating the impact of the optimal and non-optimal tokens those projections, can show the following inequality for some constant Co:
1+ Co log [|W kg ||
( |
( Wii (Wr,o (
(Vw., ((0()). ((w0, WiG)
kg
W kq
Utilizing step 2's result that (+) grows linearly, the dynamics of the attention layer can be shown to converge in direction to the max-margin solution in Equation (3).
Generalization Ability
In this section, prove the generalization ability of the trained transformers. Recall that Theorem D shows that W:() converges to W Wig kq
[ W Il- To characterize the generalization ability, is desirable to use the property of Wap which given in the following result.
Proposition 2 Under Assumptions 2 fix query token x9, let Nan, Mze C V be the set of optimal lokens, the set of non-optimal tokens, and the set of non-comparable tokens, under x -partial order, respectively. Then, the solution W* of Equation (3) satisfies To W x 0 ffr vo &Mas and
kgq
War
10
J., W kg
W Wgg
Vr. E Oge, E No..
Ozel
W
OO + W
Recall that non-comparable tokens (see Section 4 under query xº never appears in any training sentence data with the same query x². Thus, Proposition 2 implies an interesting generalization capability each xd =a-paiia order can automatically incorporate more relationships to expand the query-dependent partial orders. Combining Proposition 2 with Theorem theorem on ((())
we obtain the following
Theorem Under the conditions and notations in Proposition ((log(1/e)) Then there exists constant Co such that
let T = (((log(1/e)) and t.
kg (x, ( To ( 2 2 Cut, (zo (zo-x) w(t) x 2 Cot, Yz. E Own, E MM x E N
Moreover, if the trained transformer takes imput X with query 29 that consists of a non-comparable token To and non-optimal tokens, then the prediction made by Tate )) is n(æo) with high probability.

---

Theorem suggests that a new partial order is created by the trained transformer. Specifically, inserts the non-comparable tokens between the optimal and non-optimal tokens. The trained transformer can generalize the token prediction to such new sentences as given in Theorem 3
We use Example to illustrate the generalization ability described above.
Example (Generalization to unseen data in Example 1) Recall that Example 7 the truining dataset consists of four sentences: SVOP, VOP, OPP, and PSV. Consider the partial order >p under the punctuation mark We have that O>pp and O is an optimal token, P is non-optimal token, and S.V are non-comparable tokens. We then have the following non-trivial prediction by the trained transformer.
Case 1. Non-comparable tokens are learned to be "larger" than non-optimal tokens.
Consider (unseen) input sentence SP. Since S is non-comparable before training, but is "larger"
than P under the trained key-query matrix W the next predicted token is n(S) = V. 'qa'
Case 2. Optimal tokens remain optimal over all tokens after training.
Consider a new (unseen) input OSP. optimal and is still "smaller" than O under the trained P-partial order. The trained transformer will consistently predict P.
In both of the above cases, the trained transformer provides desirable prediction for the unseen sentences. We further note that the effectiveness of both cases can vary during the inference time of the trained transformer. For instance, if the input sequence is SP (subject-uunctuation)) the output is SPV (subject-P--ee where the input is VP verb--unctuation), it may be preferable terminate the sequence after the verb, i.c., VPP, as the verb alone can suffice to convey the intended meaning.
which follows a logical subject-vens order and is desirable. However, in cases
7 Experiment
In this section, verify our theoretical findings via experiment synthetic dataset. Specifically, we randomly generate a realizable dataset as described Assumption I with [ = 20. Then, we train Wav and Wkq by Algorithm] 1 each with 900 iterations. The parameters are chosen as d |V]
No = 0.2./v and 1j = 0.05/ Vd. Figure 2 the first three plots show the dynamics of the training stage 1, which indicates the convergence of the loss Lo (W:()) to its minimum value, the convergence of W direction to W and the linear increase of the norm [| ||. respectively. These results verify Proposition] 1 The last three plots show the dynamics of the training stage 2, which indicates
the convergence of the loss C(B()), the convergence of Wr(() in direction to Wiq and the linear
increase of the norm ||W|(]]] kg These results verify Theorem and Theorem 2
All experiments are
conducted on PC equipped with an 12-0F processor and 16GB of memory.
1.0
10
0
300 800
300
300
300
300
Figure 2: Training dynamics of single-layer transformer for NTP.
8
Conclusion
In this work, we investigated the training dynamics of single-layer transformer for NTP. We first characterized two structural properties of the training dataset under the realizable setting where the training loss can be made arbitrarily close zero. These properties allow us to define two max-mare max-margin solutions for both the feed-forwer that both layers converge in direction to their corresponding max-margin solutions sub-linearly, which further yields a linear convergence of the training loss for NTP. We further showed that the well trained transformer can have non-trivial prediction ability on unseen data, which sheds light on the generalization capability of transformers, Our experiments verify theoretical findings.
layer and the self-attention layer. Then, we showed

---

Acknowledgments and Disclosure of Funding
The work of Huang and J. Yang was supported in part by the U.S. National Science Foundation under grants NSF CNS-1956276 and ECCS-2133170 The work of Y. Liang was supported in part by the U.S. National Science Foundation under grants ECCS- 2413528 and DMS-2134144
References
Ahn, K., Cheng, X., Daneshmand, H., and Sra, (2024). Transformers learn to implement pre- conditioned gradient descent for in-context learning, Advances in Neural Information Processing Systems, 36.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2022). What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:221.11661.
Bai, Y., Chen, F., Wang. H., Xiong. C., and Mei, (2023). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637.
Belkin, M. (2024). The necessity of machine learning theory in mitigating ai risk. ACM/JMS Journal of Data Science.
Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosselut, A., Brunskill, E., (2021). On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901.
Chen, S. and Li, Y. (2024). Provably learning a multi-head attention layer. arXiv preprint 2rXiv:2402.04084
Chen, S., Sheen, H., Wang, T., and Yang. Z. (2024). Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.1944.
Cheng, X., Chen, Y., and Sra, S. (2023). Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.065228
Cui, Y., Ren, J., He, P., Tang, J., and Xing, Y. (2024). Superiority of multi-head attention in in-context linear regression. arXiv preprint 47X1v:2401.17426.
Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. (2023). On the optimization and generaliza- tion of multi-head attention. arXiv preprini 20.
Devlin, J., Chang. M.-W., Lee, K., and Toutanova, (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04.005.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2020). An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929.
Dryer, M. S. (1991). Svo languages and the ov: Vo typology1. Journal of linguisties, 27(2):443 482.
Edelman, B. L., Goel, S., Kakade, S., and Zhang, C. (2022). Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, pages 5793-5831. PMLR.
Frei, S., Vardi, G., Bartlett, P. L., Srebro, N., and Hu, W. (2022). Implicit bias in leaky relu networks trained on high-dimensional data. arXiv preprint arXiv:2210.07082.
Fu, D., Chen, T.-Q., Jia, R., and Sharan, (2023). Transformers learn higher-order optimization methods for in in-context learning: A study with linear models. arXiv preprint 17000
10

---

Giannou, A., Rajput, S., Sohn, J.-y, Lee, K., Lee, J. D., and Papailiopoulos, D. (2023). Looped transformers as programmable computers. International Conference on Machine Learning. pages 1 11398-112 100--11442 PMLR.
Huang. Y., Cheng. Y., and Liang, Y. (2023). In-context convergence of transformers. arXiv preprint arXiv:2310.05249.
Ji, Z., Srebro, N., and Telgarsky, M. (2021). Fast margin maximization via dual acceleration. In International Conference on Machine Learning, pages 4860-4869. PMLR.
Ji, Z. and Telgarsky, M. (2021). Characterizing the implicit bias via a primal-dual analysis. In Algorithmic Learning Theory, pages 772-804. PMLR.
Kim, J. and Suzuki, T. (2024). Transformers learn nonlinear features in context: Nonconvex mean- field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258.
Kou, Y., Chen, Z., and Gu, Q. (2024). Implicit bias of gradient descent for two-layer relu and leaky relu networks nearly-orthogonal data. Advances in Neural Information Processing Systems, 36.
Lehecka, T. (2015). Collocation and colligation. In Handbook of pragmatics online. Benjamins.
Li, H., Wang, M., Liu, S., and Chen, P.-Y. (2023). A thcoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015
Li, Y., Huang, Y., Ildiz, M. E., Rawat, A. S., and Oymak, S. (2024). Mechanies of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pages 685 685-693. PMLR.
Lin, L., Bai, Y., and Mei. (2023). Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining. arXiv preprint urXiv:2310.08550
Mahankali, A., Hashimoto, T. B., and Ma, T. (2023). One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint #xXiv22307003577
Nacson, M. S., Lee, J., Gunasekar, S., Savarese, P. H. P., Srebro, N., and Soudry, D. (2019). Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR.
Nichani, E., Damian, A., and Lee, J. D. (2024). How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.1473: 14735.
Phuong, M. and Lampert, C. (2020). The inductive bias of relu networks on orthogonally separable data. In International Conference Learning Representations.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. (2018). The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57.
Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. (2023a). Transformers as support vector machines. arXiv preprint arXiv:2308.1 16898.
Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. (2023b). Max-margin token selection in attention mechanism. Thirty-seventh Conference on Neural Information Processing Systems.
Thrampoulidis, C. (2024). Implicit bias of next-token prediction.
Tian, Y., Wang, Y., Chen, B., and Du, S. S. (2023a). Scan and snap: Understanding training dynamics and token composition in 1 l-layer transformer. Advances in Neural Information Processing Systems, 36:71911-71941
Tian, Y., Wang, Y., Zhang, Z.. Chen, B., and Du S. (2023b). Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535.
Tsai, Y.-H. H., Bai, S., Liang, P. P., Kolter, Z., Morency, L.-P. and Salakhutdinov, R. (2019). Multimodal transformer for unaligned multimodal language sequences. In Proceedings of the conference. Association for computational linguistics. Meeting, volume 2019, page 6558. NIH Public Access.
11

---

Vardi, G. (2023). On the implicit bias in deep-leaming algorithms. Communications of the ACM, 66(6):86-93.
Vasudeva, B., Deora, P., and Thrampoulidis, C. (2024). Implicit bias and fast convergence rates for self-attention. arXiv preprini arXiv:2402.05738.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, A. N., Kaiser, E.., and Polosukhin, I. (2017). Attention all you need. Advances in neural information processing systems, 30.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR.
Yannakakis, M. (1982). The complexity of the partial order dimension problem. SIAM Journal on Algebraic Discrete Methods, 3(3):351-358.
Zhang, R., Frei, S., and Bartlett, P. L. (2023). Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927.
12

---

Contents
Introduction
2
Related Work
Problem Setup
4
Realizable Training Dataset and Two-Stage Algorithm
4.1
Realizable Training Dataset
4.2
Training Algorithm
5
Training Dynamics of the Transformer
5.1
Convergence of Training W
5.2
Convergence Training W kg
5.3
Proof Sketch of Theorem
8
Generalization Ability
Experiment
8
Conclusion
Expression of Gradients
14
Proof of Proposition]
14
Proof of Theorem I lnd Theorem 2
17
Supporting Lemmas
17
C.2
Step
19
Step 2
21
C.4 Step 3
23
Proof of Theorem
28
6 Proof of Theorem
30
D Proof of Proposition and Theorem 3
31
D.I Proof of Proposition
D.2 Proof of Theorem
33
13

---

A Expression of Gradients
We first provide the general formula for the gradients of both layers.
V Wow Lu(Wov) = Σ
(To(æ) eIn(a)) x
(5)
JEV
C(0)
Σ (1) x(n) (diag((o(X(m)) -00((()))))()) (2r(0) Tw. (To(x!())
( - -p(r))
(xT
(6)
B Proof of Proposition 1
Recall that we use the loss,
exp (it))) Wovz)
oo) = -
- log
EiElio) exp (eg W Worz))
NEV
The updating rule of Wov is that
Ww Co(W:
W
W(()) 110
(7)
w., Co(Wor)
We know that Lo is convex respect to W.v. Therefore, we have
(W - W' Vw., Vw.,, Co(o(t)) (600)))-o(00
is clear that the loss function C reaches the minimum when Wor AW" ov: as A + 00.
Lemma 1 Under the initialization w! and the updating rule Equation following inequality holds.
with step size n, the
81J0
tno
|W
2 | 2
|| ||-
2 W
Proof. Using Equation
5 have
W o Vw.. Lo(!!!))
=E (T() -ein(a)) Wiye
NEV
Σ
T (x)) Ein(e)) Wa
< - Σ
-
Σ
[o (z))..
(8)
izA(())
where (a) is due the constraints that W* satisfies. On the other hand,
W (WW()))
Ww... C(W(())]
V... Lo
V...,. L(W
14

---

W ((0())
Σ
Σ (x)) - Ein()))
V L 2(0())
WEV
(e)
2
Σ
T (z)l..
xEV
(2)
where (a) follows from || & ||||||| B|| for any matrices A and B. Thus, we obtain
W (((())
W ov
-1/2
VW... ((00))
To lower bound the norm of Wov. recall the updating rule (Equation
Wor Co (W:(()
|
W
ov
W(0) ov
710
Wor Lo (W(("))
Eet
V Wo Co
W
W:Oo)
Σ pct
T70
V W.. Co(W (1)
|
f7/
||
W
For the LHS of the inequality, it suffices note that each iteration, the norm |[Wd9 (0) increases most ?/ due to normalized gradient descent.
Lemma At each iteration t, the following inequality holds.
Vw. Co
1+
2|lW
10g((111))
W [
Vw. Fow
[
|W
Proof.
First, consider the case when w-(()
W
W W..
Due to Equation 8 in Lemma
have
W ov ||W
Vw.. Lo(W
<0
In this case, the result is trivial.
Then, we consider the case when W:O # W
ov W. Due the optimality of W which achieves
W
the minimum norm satisfying the constraints in Equation (2), we must have that for some zo E V, there exists in of In(xx) such that the following inequality holds
|
(CIn(wo) ex) W=( 20
|
Therefore, the loss on w-(!) can be lower bounded as follows.
((W(()) Σ log 1+ Σ exp - -(i(-))))))) ( wo(t)))
log (1 exp
(a) 1
exp (-1IWW011/////
2
where (a) is due to the fact that log(1 + x) 2 z/2 when < x < 1. On the other hand, let
W (e
+ log(2(11)) W Then, the loss on W has the following upper bound.
[
15

---

L(W
Σ log
1 Σ exp (es -
(m(() (m((w) W x0)x)
NEV
log
11 (12)
exp
exp (-11 11 W ov 1og(2(11))
.EV
(0)
S
Σ
| exp
(-1I log(2(V)))
NEV
exp(==1li
where (a) is due to the fact that log(1 + 2) < when > 0. Thus, L(W()) > L(W... Due to the convextiy of Lo. we have
<< (W100 W V... Co(w(t)
[ 22og((111)) (wā. Vwa LoWdd))) ||&lll
= =(wb V.o. Co(
which finishes the proof.
Proposition 3 (Restatement of Proposition [ Under the zero initialization W: 0 and updat- ing rule Equation ( for any 1 2, the following inequality holds.
W:l
W
12 W 1 log (2
logt
2 1
|| (0)
||Wa>]
t770
Moreover,
trio W
<
||
S
tro-
Proof. The second argument about the norm of W!! follows directly from Lemma prove the first part as follows.
We aim to
ar
2|IW
VI). By Lemma and the updating rule Equation 7 |
we have
w(++1)
W
W
VW.... L(W))))
W
170
ov
[ nvl
I
(
170
V... L(W()))
1 + ag
||
w
1
W((+1)
W(())
+ att
ov
||
(
Cx
(((())
+ oi
||
(1W(0413112_ |W(++1) - w ||
2
at
w((+1)
w ov
#+ az
m
w-(f)
16

---

[w(++)
2
11
2 W II
2 ||
Wo [(O(t))
W
1 + ax
IV [(0(0)]]
|
(6)
[W(++1)1
||
2
(2)
+ az
where (a) follows from that [| wol for any 2, E IR.
1oo and (b) is due to the fact that x² x²=²²² z 2y(20-1
Summing over starting from 2, we have
WW w ov:
W
|
|
Σ
Y00k'
[W
21
1 ap
P=
d==2
Furthermore, due Lemma
W 11/10
t
2'=2 1
2 W ||
logt. t.
70
Similarly,
Σ
O
2 log(2[V])
1 1 ay
|| I1
t=2
W
1| log(2 VI)
logt
770
Therefore,
W:ll
W [Wall
2no W ||
log t 4|| (2g()))) logt
21 1
||
|| ||
12 W 1 log(2 VI)
logt
tmo
where (a) follows from Lemma and
|| (2) s 2no - W 1l. and ||W!
2 (2 (W
Proof of Theorem and Theorem 2
C.1 Supporting Lemmas
Lemma 3 With zero initialization, under the updating rule Equation 7 satisfies that
for any iteration i, w
(x)) (e Tw() = O, Vi, In(x). Proof. The proof follows directly from induction and the fact that
(e. - (e)-er) >rrrrTTee (w((+)) (es - -er)) <. Tw To( [Tr6
Vi, + In(x).
|/V W... Co (o |||
17

---

Corollary Under the settings in Proposition 1 let T > A (4
3 logg(()))))) log T/no. and
(eIn(s) es) W WorS E (A, 34), Vi In(x)
ex) T 0, Vi, In(z)
Proof. The second equality follows directly from Lemma To show the first equation, we analyze
(eIn(z) es) T TW
W
W
W
(eln(x))
ei)
+
IVT'll(eaats ez)
W
[ ovll
|| (7)
[W
W llog((2 12)(2 log' log
2 2
[ |I
ov
Tmo
(6)
TTo
24 W ( log T
2|| W 11
Tnoo
4||W> ||?
where (a) follows from Proposition
1 and (b) is due to Lemma
On the other hand, we also have
12 W ov
3 log (og()(1))
log
(elu(z) - T T.)TTx
+ 2
|W
Tíoo
37'110
4 W
The proof is finished.
Thus, for simplicity, we further assume that (4In()) e w. A for all x, because (ein(z) e:) = O(A) for large enough iteration. Next, we provide the general form of the projection of the gradient of Key-Query matrix Wkq follows from notation for the token weight.
The token weight (.)))) of the token xx in the sentence ((a) (Wof), WKg) is calculated
[a]" (n) ..... (3)) under =
exp (aafil) Wkg.x(m)
(n.t) Pe
(9)
El' exp (xx) (m)))) Wkgx x (m)
Lemma (Projection of gradient of Wkq) If satisfies that
(etn(z) - c,) T Wor A, vi In(x)
(e: 1a--0 (()-er TT.TT.Wn 0, Vi, 1' * In(x)
we have
(Vwn, [(0), W(a)
A Σ
( (n) [m(x(()) (n 1)
Σ
(n,0)
(n)
Σ
(m
(m)))
W Wingx(n)
Pe.
x
pp
e.e(())
Σ
n (n)
(
(m,0)
(m)
Σ
(n.0) (1))
WKg x(m)
x
pp
n
(
18

---

Proof. Recall that ((n) is the set of indices of the optimal tokens in the sample X(1), Thus, for any La [(n), In(x{")) = In(X()).. In addition, we denote (X(())) by T for simplicity. From
Equation (6, we have
(Viin [(0), Whng)
T
Σ
π (x)
(TT)) - p Wore (1)
(
pl p/(00)
Σ
Fp
xt
(n)
W Wagxxin)
n
T
(a)
π (1)
(TT en(x(en) W Worlh (())
pp
(mi0))
x
Σ
Fp x
(n,0) ((()))
Wkg x(m)
T
Σ
π im) Σ
TTt li(e: (mn(x(m))) T
Worey (n) (n0)
"pe
(m))
Σ
Pe
(n,0) (m)
I'
Wkg x(1)
(El(n)
|: ((n(x()))) T W Worz?
(n) (m/))
(n)
Σ
(n.0) (n)
(m)
Wigx'"" (n)
Pe
I'
12
(g((n)
Σ
E
(TTh) ):(- ((--))pp (n))
1z (1/1))
(n)
Σ
(n.0)
Wky x(m)
π
Pp
tif
7
(E)( (n))
Σ
(T
(nn pl())) (n(1))
Σ
(m.0)
Who x(m))
Pp
(g)(n)
A
Σ
π
(n) ((T ( lnn(X(()) 1) Σ (m)))
tta
(n)
Σ
(
(n (x)
Itt'
(n)
W
E.El(n)
(n)
[T
(n))
Σ
(1.0)
(m)
W X
π
pe
(gf((m))
where (a) is due to the fact that (T
1 for any 0, n.
Main Steps
The proof consists of three main steps. First, we show that the optimal token weight has lower bound. Then, we show that the gradient aligns with the optimal direction. Third, we show that the norm of Key-Query matrix grows linearly. Combining these three steps, we can prove the Theorem and Theorem 2
Recall that the updating rule for w:(f)
(f) is
W 2(0()) |V W., 2(B())
(t+1) kg
(((t)
(10)
C.2 Step 1
We first show that the optiaml token weight has a lower bound during the training.
Lemma (Lower bound of optimal token weight) Under the zero initialization and updating rule Equation 10 for any iteration t, and any sample X()), il)) is the set of indices of the optimal
token in X the following inequality holds.
(n,e) 2 st
(2,00 2 1/Lmax
E ((el())
Proof.
19

---

(n.t) First, we introduce the notation that (n.0) = Et. El()) Pt. as the summation of optimal token weights, and
and (p_h.)) =1 1 = the summation of non-optimal token weights.
At / = 0, due to zero initialization, we have pe = 1/L(). Moreover, by Assumption for any $ I(n), we have 4 (n.00) 2 qn (
We perform induction the hypothesis: (n.1) 2 p+ of (m.t-1) and (n.t) 2 pp (n.t) for all la E ((n) and
é $ I(n).
(n) Suppose the hypothesis holds for iteration i. Let ve. be the optimal token in the sequence ((n).
Fix a sample x("), we have
(ob ( ) T VW Wa (0()))
Pe (x - =x())) )T()) ) (x x!mr)
Σ m(()) (1T (m) l[n(x(m) 1)p(( (m.t)
Σ
(m. (m)
((
(m) ) To ) (x (
+ Σ (n)
(1)
(mit)
Σ
(1r)
(m.t)
(m( sPt
Pp
FG((()
Because Ex(((n) (2p (2)))))) yo
1 and (x/ (a(
Tx (n') x 2 0, we immediately have
0 due to Assumption
(o((")) Vwing L (B(t)) x SO.
(n') Let ito be any non-optiaml token in the sequence (1x)), Then, we have
(a ( VWe, ( (0(0))
Σ 7(m0 lin(x())) 1)p(
Σ
(n.t)
(a ((())) Ta(a') (x
De
(
(n)
[T
(n.1)
Pp'
(m,t) (n)
(xz
) 2o (x
(gl(m)
x xT ((an)) ) (xx))xxa)) x!a))
Σ
n (n) ((T(a)) Imnx)0oy
([T - 1) (n.))
ge (- (1
(n.1)
((n))
((f((n))
=x(()) Tx((o)) ) (xery. xx(1))
+ Σ (n)
(mi))
Pe (
(m,t)
(xl
12
(g((n)
((f(in))
2 Σ (n)
(1 (1- (T(())
((((x(m)))(
(mx1)
Σ
Σ
(nie)
(
(2)
(x
x(()))
n
(gl(n)
(gl(())
2 0,
where the last inequality is due to Assumption Therefore, for any n, we have
the induction hypothesis and Eie[nn]] (T =
(n) (eo))TWoie Tw (1+1) (n)) exp (cr
(n,4+1)
Ee exp (wx()) (1)) T (1+1)) x(1))
kg
20

---

exp
( ( x - (1) | W/n - (0(0))
- (1)0)))))
kq
Σ, exp
(x w
kg
(m) 17(26 ) V Wkn (((00)
Wie ( (((1)))
exp (z(².)
(2)) ) T W x(n)
Σ. exp
( ((c""))T"'" Wa (m)) + 1 (m)
_t ) V VW ( (0(0)) / (((E)))
exp
(1))
T kn
((())
Σ2 exp
(x) (
(t) x x(x))
kg
(n,t)
= Pq.
which implies that +pt
((
2 80
(n.1)
For the second argument in the hypothesis, we examine Pi. /48 /
(n.+1) (a,+1) for any & & I(n). We have
(n,++1)
Pt.
[n,++1] = exp (1r It
(n)
( x(m)
= exp
(s/
xt² (n)) (T x
kg -11
exp (-1 VV, L0 (000))) (
(n) - xx (1)) ( VWi., ((000)
(n.t)
Pe.
16 (n.1)
1.
The proof is finished.
C.3
Step 2
The following lemma shows that the norm of the Key-Query Matrix increases linearly with the number of iterations.
Lemma 6 Under the initialization (0) the following inequality holds.
and the updating rule Equation
10 for each iteration t,
kq
bn
+
[| kg
[[(()) kg
2
bn
|| (9) kg
||-
21max
W
kg
Proof.
We examine the gradient We,, [(0) projected onto the optimal direction W W ll.
V Wi. [(B(()), Wiia)
Σ 7()) A
((T pioeoli 1)p(1.))
(n.t)
(m,a)
Σ
(mit) (aie) ax
di
El())
8r
nn)
[T(()) pce
(nit)
Σ
(antt (1,a) ael
In In(or
(gl(n)
(x) A([T(()) (( 1)p (m,))
Σ
(a
(m.t) (n.e)
(a
ag (m)e))
(gl(n)
21

---

Σ
A
A
Σ (Treed
In(fo)) (n.e)
Σ
Pe (al
(nid) (a ((n))) ag ol..tt)
(gl(n)
(EE(N)
A ouslim(ximm)
1) (m)1) p
(T 0cs 1 (mi))
(-p((ii)))
+
2gl(n)
(2.0) where 4 (n.t) Ee, El(n) Pe. is the summation of optimal token weights, and (m.t) = 1 - 1 (2)1))
is the summation of non-optimal token weights. On the other hand,
VV Wko 2(
( V Who L(B(1))
V Wi, ((O()) Win, (((())]
Wi 2(0(0)
(n) A
gc (o(o 1)p(.
(n.1)
It,
(n) -
-E
Σ
Pp
(nid) (n)
Vw,, C(A(O)]] x(a)
2, E((((
V ((0()))
+ Σ
π (m)
(T())
ols Inaaftn 1n0 sp(
(mi))
(n)
Σ
(ni) (n)
x(())
spp
Jif!
| Wing [(A(t)]]
Eg((m)
< 2
A (T Inn(X (mn ) 6p_ (mit))
(m.t)
42 Σ
5((m) (m) A
(Tb
1 (n1))
pgf(()
Thus, we have
Vw2. C())
Wia
min, 40 (n.t)
VWL. [(O) |Wa]]
2 W kg
2I W kg
By the updating rule Equation
1T0
we have
V We, (((()))
W kgq
w. (0)
'q
Win ((0()))]
Wkg
V Wi., C()(e))
W
Wo(
W (0)
Σ
[ W
V Wi., ((O(e))]
||W:(01 ||
Lmax Wiig
tn
kg
(0)
||-
Lmmax
kkg
In addition, by the triangle inequality,
V Wie, C(B(0))
Wie (((
P=O
S tn + (0) kg ||.
The proof is completed.
22

---

C.4 Step 3
We next show that the gradient V Wa. ((0()) ) is close the optimal direction Wkq-
|| W 11
Lemma 7 (Gradient aligns with the optimal direction) Let to t to, we have
Then, for any
[
WWn, C(B10), W(()
(1 an) (Vnn ((B10)), Wika)
[ kq
where
4N L2 2
W Wigll 112
(1 log (2 ww
| kg
Proof. During the proof, denote
(n.e)
=xx/ (x(()))Tpr(()) x(m))
kq
(1)
(1,*) (ar( T W X(n)
kq
W kg
212 max
[ kg 11?
80
(1 +
(w(t)
kq ()).
[[W((0] kg
||
We point out few facts that will be frequently used the proof.
az (n,t) (n.t) S ap - Co, then we have
(n.)) Pp
(n,1) exp
a1
(ine)
- dgr.0) exp(-Co)
(11)
(n.t) is replaced any convex combination of of ag The same result holds if ag
(n.t)s.
We start the proof by noting that W is the minimum unique solution to the problem
W kq arg min [[W]] s.t.
(o - xt (")) 21, vl. E ((m)) , e ((m) Vn.
(12)
Therefore, if w [ Wig the results is trivial since
(Vw.., L(O(t)), Win) 5o. 0.
kg p
In the following, we focus on the case when wilt) |
* W Then, there must be at least
kg |
sentence X(m), such that W [WEll violates the contraint on ((), In other words, we must have
kg pw
W (
(n.t)
Oz (x xx ( x(())
(n.))
(m) (n)
kg
ag
W kkg
I |1 T
(n.t) exp(
This implies that for those n, must have sp
Thus, we consider two types of samples in the folloiwng.
Type 1. Let us consider (n) such that (m.t) 2 exp(-(11 + 30/ (6/2) |
Recall that the inner product between the gradient Vwn. ((119( and any other Key-Query matrix
0' = W has the following form (Lemma
V Wh. (()()), Wha)
23

---

A
E
xl (n) (IT (
In(x(m) -1
Pe.
(n,0)
(x)
Σ
pe
(no (n)
It''
W W (())
(.((())
n (n)
(T
(m))))
(()
Σ
(/100) (n)
Wko x
In(a
De
n
((f(())))
Let (0) log c (o) ()((( be the loss on sample X(m).
To proceed, we examine the gradient on each sample ( with p! 2 exp(-(1
[1/1/// kql). II), which can be divided into two parts. Wen Ln (O())), Wka) - A(A(m.)) +
B(n))), where
A(m,)) Σ
((T
1)
(0.1)
(n.1)
Σ
Fp
(ma,1) (n.1)
ag
E.El(m)
Blnn = Σ
(TT
Inlel (nly (p
ap (m,))
Σ
pp
(n.t) (n.t) ae
gct)
egl(m)
We further let
A(m,)) Σ
((T ((oo
-1)p(1000
(0,0)
(m,))
(nn)) (m.)) Σx
Σ
pe
ag
)
El(()
and
B(n,)=
Σ (T()) (gl(n)
(o (nD))
as (nn
Σ
(n,e) (teit)
agitt)
Thus, aim to find the relationship ((n,)) + B(n()) between A(100) + B(nr))
We first provide the upper bounds for (m,+) and B(n.))
(z,00
Σ
((T Imi - 1)p((1))
(n.t)
(m,o)
Σ
(niit) (m)))
ai.
por ag
l. El(m))
Σ
(IT ((oo))m((a?")) (m))
(1.e)
Σ
Pe
(i (112) ag 1(11)))
(n.)
(al
(l(m))
(rl((n)
(a) ((T())
I ( li(xt)))) - 1)) (m.1)
||
kg
[
where (a) is due to the fact that (o (of"") W / ll.
xt
(n)
1. and (nx)
ag
On the other hand
(m.()_ = Σ
((T(())
800 Im(agn - 11)p( (m,t)
(n.t)
Σ
(nit) (mit)
aj
De
ag
eel(n)
= Σ
((T(R)
- -1)yp( (m.t)
Σ
(n.t) (a (mi.))
a ((me)))
(El(n)
(ed((())
24

---

(IT (Coolliace (1)))
1)sp(m)))
(n,))
Σ
gp (ap
(n
(a.m..)
ag (n.
max
(e)(
efl(())
diffc
+ Σ (1T ((teed -))pp (md)
Σ
pl (af
ae )
REl(n)
diff
Σ
(1T aooe
1) (mie)
Σ
max
RE(()
(gl((n)
+ (1T stoo)n( 1)sp! (m)))
2
Σ
exp(- exp(-T)/lW
(EI(n)
(gg((n)
Σ (Tpeell In( 1)so(m)))
(m.))
(so!m.o)T 2mmax exp( 1ep(-TTp (T)) ((
2 max
T
(El(n)
2 W (n.1)
(Cr(a)
Lac=!") 1)sp(mm0) p! [+ log
(m,e) (n.0)
(El(m)
2hmaax 50
where (a) is due to Equation (11 and (b) is obtained by choosing log
Recall that x! (m.)) 2 exp (//)) (1/1 and 30
2|1W (lggl [W: ID)
[|
Thus, we further have
2Lmax (1)
A 2 ((T)2) (ooltt(x))) 1)) (1.e))
1+ log
50 (nn.t)
[ |I
2 (1r ()pp(m (m)1)
1 loloolgLnggg [max W1 )) + (1 + B2/2)
|
/Bollw kg
|
2 ([T (x)
0co (mnxin)
1)pp (n.2) p
(m.e)
+ (1 + 22/2)
2 kg
W kg
[w:(1)
(1 + (6)(((())))) (ro - 1m)0 (((t)))))))))
[
(1 (m..)
Next, we analyze B(n.), and further divide B(n,0) into B( and B(n.0),
((
Σ
(T(N) aco) Intaf (m.))
pl
((mil)) (n.0)
(m.8)
Σ
so (0,0)
ag
(g((m))
('l(l)
B (m,00)
(T(()) 0(5)
(mid)
(mid) (m.0)
Σ
(n,e) (n,0)
In(a?
snn sp
ai
De
ag
(()(())
eeEl(n)
we have (m,+)) 0. and thus
Due to Proposition
|
(a,+) = (m,*)
pco Jr, (a (p (-
(m.t))
< 0.
W
(
25

---

We then analyze:
B B(n.t) (1 + (o) B(
Σ (((R) (a(a) (n)) Pt pp+ of ((nt))
(n)e)
((m.tt (n,1)
-1 + 30) (vf af (2.00 ar (44:0aggm)) (mu)))
di
(g((m)
[T (2) 0t0 hntaa foly Pe s
(mn)) (ffi))
(n,t)
(m.e) dz.
(1 + 300 (frint) -o (m::))
ai
Oz.
(gl(n)
bi,
ba,t
( Tocoolana (n(a Pi ( (m.1) (n,e) (be,e by,e)
(gllm)
|
OcO En(o exp
(m.t)
-(1 + (b0
| ))
(-21111(21)
(-
Rel(m)
|
( 3o l| ||
- Σ Thee (wn exp
(m.1)
-(1 Bo/2)
||
kg
|
exp
W
egl(m)
||
(
2 -2 (T((r) ( (n,1)
D _p_
(n.1t)
exp
Bo
||agil
|
2
[ kg
( Bo | )
2 2Lmax exp
(-2 2 [|
| [W |1
(0) |
(m,,))
|| kg
3o A (18,*)
where (a) follows from that ba,e x -(1 + (o) [ (1)
kg 1/ and Equation
11 and (b) is due the
fact that (1 = for any 0, n. and (c) follows from 3 2
||W log Lmxx
kg |])
For the term (m.t) we have
B(n.))
(m) (ppp sp
(( a ((nt))
(n.1) (ra))
Ap ag
((f(())
p(l(())
(n.))
Σ (T(i)
Ma (0) Pi (nie) (m (nnt))
(n.t)
Σ
gp
(n,t)
at
(nie) ag
(xl((n)) 60
(gf(())
(n.e)
( (nceera (m.0)
(n1)) (m,0)
(n.t)
Σ
so (m,a) a (n.1)
max
dg
T>0
ag Ja
('gl(())
dinl
6 (n.t)
Σ
IT ote oo sf (m.t))
(n,t)
(n.))
a
(m,t) Ogt
(()
diff
26

---

(Toeed Infal (n) sp!_.)) (-1ffu00 - 2/ exp( (1)
max T>0
(((((1))
Lina (1
gor grcs (m)) )اا! (1+ log(2)111
Lmax Wig =1
(1+ log (2) A(m.))
(1.t)
9+
12
Wig
max
1 + log (2) A (m))))
[|
kg
(1)
||
(n.) IV3
where (a) follows from Equation 11 (b) is optained by choosing T = 12 W kg
II), (c) follows
from Lemma 5 and (d) is due the fact that Bo 2
L Wey Il
(1 + log w/ ID).
1
So far, we have shown that for if (m=,)) 2 exp(-(11 + (6/2))l1 11/l II). then
(m.t) + B(n,t) (m.tt) + B(m.t) + (((.t)
2 (1 + Bo) A (m,-) + (1 + (1))))(())))) + Bod((ar)) + A.
(1 + 3,6 (()))) + (1 (()
2 (1 + 3.6)(( (m,t) + B(a.)))
Type 2. Now consider sentence X (n) such that m! exp(( (11 (1+ + 3,/2) |/| kg l).
Let ng be the type sample such that yp (no.e)
exp(-1| | I)
Then, we aim to show that
(m.)) + B(rr) 2 AoA (mo.*)
Note that
Lmax [ kg
(lr,t)
Σ ((T()) ((Tollin(o??")
1)pp( (1.) (2 (++
(o(1.)))
(.((()
2 ((T()) pce In(a..)) - 1) cxp ( B2/22)
ky I|
)
1+ (1+
Bo/ 2)/
| ||
+ logg mmax W II)
||
[[
(-11+ ((1++00
||
2 (1 (()
010) ( - -1)exp exp
2) (0/2)
[ kg
[
and
B(aa) = Σ (T((T)
(mit)
(n.t)
Σ
(m,t) (n.))
In(o spt
(g!(n)
||
> 2(1 - (T())
@ (n(e exp
-(1 + -(1+
3)/2)
| II
27

---

Since
|
A (mm=+)
< ((T (
1)4+ (ma.t) (m).e)
[
1|
[ ||
S ((T
1)++ exp
(mo.t)
kg
(
||
We further note that [T In(a!ho))
(T gce) ince (,o) due to pt ofmort)
< (m.))) Thus,
A (m.0) + B(m.2)
|
A
2 exp
-Bo/2
(1 8o + 2||W []
WWigl
(n)
B (no.r)
where (a) is due to that Bg 2
2 2/ 21W2u1l(++2
log(1 + 2
[[
||
2 2(e -
1)
kg 1l
and
2 kgl kg
1 8o + 2 W kgg
Bo 2
log
[W((] kg
B0
In summary, we have that
Σ n (n) (A(m0)) + B(ml))
Σ (A Blml)) +
Σ
(()) (A(n.t) + Bla.t)
n is type
is type 1
type 1 3od (mo:))
(1) ((1 (11+3,)) 3,6) (1,+) ( + (1+ B1 B()))
max
is type
Σ (m( Ba + B (m)+ + +11+ 330)
Σ
T'n) (A + B(n-))
type
type
(1 + (N 3)))
Σ x() + +B(()))+ B(nr)))
isstyps
2 (1 + at) Σ
(A(x.)) + B(m,t)) + (1+
Σ a type
( + B(oo.+))
no type
(1 + ar) Σ
(A(i))) + B(r,t)),
where a; 2 (N 3).3. The proof is finished.
Now, we ready prove Theorem I
C.5
Proof of Theorem i
Proof of Theorem
12
I 112 [ log (( By Lemma 7 we have
Recall that Cq
[
(-(1+1) kg
w.(t)
W | ||
kq
28

---

Wkq
Wian C(B())),
17
|
(t)
2
7f
V Winn ((0()))
1 + ax
|
(
II
kq
W
w
kn
+ ar
log
|
||
kq
1
|
W 1 (11
a2
w-(++]
w
kg
11
+ az
| ||
|
[
1|
27
7/c
V Win ((B(1)
W
2/
kg
(t)
2 W (0)
1 + at
V ((00))]
[ kg
1)2
2 ||
kgg
kg
2 2/(w:((1 kg
1 + Oxe
Let to
be defined in Lemma
7 Summing over from to. have
7
Wig kg
wolfo)
| (9
(fo)
7ep
kq
||
2
(t)
+ au
t'=
By Lemma 6
have
2mmax
[ kg 11/1
[I ||
kg
1f to
Lmax [
log t.
Furthermore,
t-1
Σ
Qr
Og
+ o
t'=lo
4NT2
W 112
kg
(1 log (2lmax (((()l))
4NT2 max
W kg 11?
4
W kg 112
(1 + log (2
mmx
log (t)
[[ (e')
LNL3 max
[ kq
113/11
L max
[| |///m
t'
log 22 Lmaxx) +
log
2 L
Wig 11/1)
Mmto
SNL3
[ 11 113/m
SNL3 max | 11//
log e// / (gg 1)) +
log(t')
1'
1'
P'to
P=
3 max
Wtal! 1|3
log21 1
1
where (a) follows from 7 S || Therefore, have
W* kg
7/
wr(0 (1)
(ta)
|w/600) (60
Σ
nop
|[
2 (
1 1 ay
r=
29

---

| || maa W
log
8NL
max
W kg (l log" L.
Finally, by Lemma and to S 1 + 81.aax || kell 112//,, we have [W
97
W 11 and
kq
kg |
Wkg |[
2/| W/ Temax [ kg
log /
8NT.
max
Wk || log" /
||
(4)
54N
W kg
log
1
max
bn
The iroof is finished.
C.6 Proof of Theorem 2
Proof of Theorem
Recall that T' 2 384 W (2)(2 log T//0o, and A = T'no(44] (4]] due to Corollary
|
and
(Cin(s) ex) - W!''
Tmoo
41 W
Note that for all lu ((n) and l & I(n)
(x! (n)
x(( w:(0 x(x)
|
| (0)
kgg
(ate
(m)
(m)) T W (n)
(1z
)T
kg
(t)
I kg
Wkq
x(1)
[ II
((tt)
kg
kg
(13)
[
|| (0)
kg
tn
V2tn
54NL4 max
W llog"
W kg 2 2mma
1
Lmax W
bn
by
4imax
W 11
kg
where (a) follows from /. 1 1696NT. max 1 [og'" 1/n.
Therefore,
[((n)
Σ
(n.1) Pi.
[[(n)] + Ees(on) exp (xl (m)
(
(
(t) kg
((1)
El))
[(n)]
[(m)] (I(m)
[((n)]) exp
TLmaa W
1 Timax exp
4L.mox I|
1+
44mas [W log Linax
where the last inequality follows from
Hence, the loss on the sentence (1) satisfies that
-log (C Torol (X()))())))
30

---

exp (eT I (m(()(o)) Wor
Ee 2 (mm) pp(1)))))
log
Σ Ensi exp e Wil!' Ee 266") 4i
(aat)
=-
1 E(HI(((()) exp ((ew - ((((X(()) TW
Ee 2e Di
(n.t)
=log
+
E
exp
(E, - (m(x(w)) W(T) W(T))
Σ
Di
(n) (nit)
4p8
FE!(m)
llog
1+
Σ
exp
-
Σ
(n,e)
+ A
PE
(m.e)
vxln(X(e))
E.el(m)
(El(n)
[V] exp
A222p
(m,e)
11)
2AL
$ exp
- + Lmax + exp
4T.max [
Thus, the average loss has upper bound, which is
2A 2 Imax
exp - + Linax exp(Cit))
for 1 17/(444 Wxgll 11)) and A CoT for Co 1o/44| 1))
D Proof of Proposition 2 and Theorem
D.1 Proof of Proposition] 2
Proposition 4 (Restatement of Proposition 2 Under Assumption 2 if WWa satisfies Equation
ie.,
W arg min [W]],
(o (n) - (m) - ( VE ((n). Vn.
In addition, for each query x9, if there i optimal tokens under x -aaarrraall order, optimal tokens under x² -arptalall order, then, for any optimal token Iw, non-optimal token x, and non-comparable token To. we have
Tu W W x9
m
x Wa
b
To W&g x9 = 0.
+ m
+ m
A direct result that
(x) (m)
a(( T x (m) 0, ve,l & I(n).
Proof. Let E Rd be the rotation matrix such that Ux = ((()) Because U preserves Frobenius norm, the optimization problem in Equation (3) can be written
W arg min [W]] s.t(()(())) - ((as")))) TWe 2
(14)
Notably W UW kg UT.
Note that (((())))) Jnn forms a standard basis. It suffices to minimize the norm of each column of W subject the constraint ((((=(")) (((a()))) T Weeeye 1[x ((x(m)) (m)) 2
Let us consider any column of W. denoted as [wi,..., wa] T. Without loss of generality, we assume that, for all X (n) with 1(x(*) the set of indices of the optimal tokens of those samples
31

---

k). and the set of indices of the non- non-optimal tokens are { + 1......... . + mm}. Then, the optimization problem Equation (1) reduces the following problem
min +... w s.t. Wy 2 Vi S k,j E A, C {A +l,...... m)) (15)
where A; the set of indices the non-optimal tokens in some samples whose optimal token has index i.
In other words, each column of the solution of Equation 14 the solution of Equation (1)
Note that Equation
[3} is a convex problem with linear constraints. The Lagrangian function
L(A)
Σ
Wi
>i1(1 we + wyy,
SEA: where directly set wy for all j E { + ,,..., d}. That is, non-comparable tokens have value 0.
i=1
By KKT-condition, we have
Xijt
Wisk
JEAL
Wj
E
Ak
++-sjjsttn +m
1=1
Thus,
min w +... + wi
k+m
k
max
Σ
Xij
- Σ
Σ
+ +
X
1=1
jEA,
Ø=
1=1
EEIEEA
Let
2
m
(*(A) = Σ
Σ
A
- Σ
Exal(EAA))
+2 +2 Σ Aig
jEA:
j=h+1
ii=1
i=1jEA,
where A 2 0. The maximum of L is achieved when VxL* = This implies that
Σ jEΛi
Xno +
Ag e A}=1 11
V1 s i0 s < jo S k m.
Hence, we have wi w for all 1 Si A: < S k + m, which means the optimum of the original problem is achieved on the boundary. Therefore, we reduce the original problem to mink(2 + 1)² + ma²
where Wk+1
Wk+mm Hence, the optimal solution is w1 Wk+1 = =... Wk+m - k/(k m).
Wi m/ m + k), and
Therefore, the solution of Equation
15 satisfies that the "optimal values" are the same and the
"non-optimal values" are the same well. This fact proves that
(at - x W W x 0, ve, El & l(n).
And moreover, if there are k optimal tokens under 2 --rrtila order, non-optimal tokens under x² -prrtial order, then, for any optimal token v. and non-optimal token v, we have
3r, W
6 kg
m
w T
W
1"
1: +
1: +
32

---

D.2 Proof of Theorem 3
Proof. The proof follows similar logic Theorem 2
By Equation
13
we have for any 2, x' EV
(x x') (t) x2
|
2tn [Wig] lqq
54N
4
[
llog t
kg
x²) Wigz"
max
kq
| kg
Liax
trj
In
(xc x!) W
kg
N0. T.2 max [W lolo" 1
2
W
kg
where (a) follows from Lemma 6 The first part of Theorem follows from Theorem 3
For the second part of Theorem Let X (..,..., 2r such that for lo lo C {l...., Teg a non-comparable token, and other tokens are non-optimal under the z #u-partial order.
kg Let fp x exp(x W ( xL) for sufficiently large = ((log(1/c)) such that Etosto Pto 2 1 - c.
Then, have
exp (in(s) W Σ, x2pp
(i((a) Taco (X)
E Ensly exp
C W.T Σe Depe
1 E Evfl((2) exp (ee -In(()) ( (W! Ee xesp))
14 Σ (z(((2)) exp (-A -LEEo Pto + A A 9P))
2 + |/ exp (-A(1 2c))
2 1 Co:
where the last inequality follows from T 0(log(1/s)) Therefore, the trained transformer will predict nxx), the next token of the non-comparable token.
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide details of our three main claims made in the abstract in Section Section 5 and Section respectively.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Assumption
Answer: [Yes]
Justification: Our paper has specified the assumptions (Assumption I Assumption 3) under which our results hold.
2
and
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
34

---

Answer: [Yes]
Justification: We provide the full proofs the appendix and title them as the "proof of lemmas, propositions, and theorems". Each proof can be easily found from the table of contents.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it acects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide our experiment details in Section including the model dimension, vocabulary size, the learning rate, and the iteration number. The experiment follows exactly of Algorithm
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
35

---

reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We provide code in the supplemental. We do not open source data.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
public/guides
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7/nips cc/ wtc/publicc/ hoolllccggiiess. gwt.///ieee//eee// Codeeoomiiieeeee /eso/euuuide/ SanSSuubissss iebbbbissssioi ssisssoicccyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Our experiment details are provided in Section 7
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: As our synthetic dataset and training process is fixed, we do not foresee significant errors.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
36

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Our experiment is on synthetic data and computationally lightweight, which does not require any GPU. The rumning time is also within a hour.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics atttps:/neeurps httpp::/aaarrrpp..ccpubl///////
Answer: [Yes]
Justification:
We follow exactly the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field Machine Learning. There many potential societal consequences of work, none which we feel must be specifically highlighted here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
37

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper is primarily a theoretical work, and does not have such risk.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing coective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith eoort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: Our paper is primarily theoretical work, and does not use any assets.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
38

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper primarily theoretical work, and does not release any assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper is primarily a theoretical work, and does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper is primarily theoretical work, and does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
39

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
40


Reviews:
[{"review_id": "niYjrnNlm8", "reviewer": "Reviewer_U8Jg", "summary": "This paper conducts a non-asymptotic analysis of training dynamics for a one-layer transformer, focusing on next-token prediction tasks.\nIt provides a mathematical framework  based on partial order to formally characterize a realizable training dataset for next-token prediction.\nIt also introduces a two-stage training algorithm that ensures fast convergence, with both layers approaching their max-margin solutions sub-linearly.", "questions": "Q1: Although the ground truth model is deterministic, it seems uncommon in machine learning to assume a statistical model because the next word is not deterministically determined from a deterministic context. Besides assuming n, where is the existence of p*L theoretically necessary?\nQ2: The assumption of the existence of collocations seems too far removed from reality. What are some realistic data where such assumptions hold? The setting in Dryer, 1991 seems artificial. Especially, the existence of n() is crucial for the two-stage training and the proof, and it forms a major assumption at the core of this paper. Is it possible to actually prepare n() in real datasets?\nQ3: I am confused about the definition of partial order. In Definition 1, it is stated that if there is at least one sentence where x >{x_q} x' and n(x) = L_{L+1} != n(x'), then it is also permissible for there to be sentences where n(x) != L_{L+1} = n(x'). I could understand it if it were defined to hold for all sentences.\nQ4: How commonly is the normalized gradient used in actual Transformer training? In the training of Transfomer, AdamW is typically used.\nQ5: Assumption 3 may not hold in situations where attention is sparse. It is important that attention is placed on the optimal token, but it is known that the attention map in actual Transformers is generally sparse. Doesn’t this assumption contradict the sparsity of attention?", "limitations": "The assumption of the existence of collocations seems far from real traninng dataset.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper develops a detailed theoretical framework that both analyzes  non-asymptotic convergence of the training and generalization capabilities of transformers in the next-token prediction task.", "weaknesses": "The approach described in the paper may not align with practical Transformer training.\nPlease see the questions below for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments. \n\n**Q1:** Deterministic ground truth model seems uncommon. Besides assuming n, where is the existence of $p_L^*$ theoretically necessary? \n\n**A1:** In recent line of theoretical research on transformers, deterministic models are often adopted such as in Li et al., (2024) on next-token prediction and Tarzanagh et al., (2023) on binary classification, mainly because deterministic models can offer \nvaluable theoretical insights due to their simplicity and clean structure. In particular, it provides a clear cause-and-effect relationship between inputs and outputs, which makes it appealing to facilitate the interpretation of the model. Note that deterministic model of ``separable data'' is also commonly studied in theoretical works on neural networks such as in Soudry et al., (2018). \n \nThe ground truth model $p^*_L$ serves as a theoretical baseline for analyzing an algorithm so that the performance of a training algorithm is well defined. As a notation, $p^*_L$ can be replaced by $\\mathrm{n}(X)$, where $X$ contains $L$ tokens. As a realizable model, it is captured by collocation and partial order. With the existence of such a model $p^*_L$, the training algorithm can have provable performance. \n \nTarzanagh et al. Transformers as support vector machines. arXiv:2308.16898, 2023. \n\nLi et al. Mechanics of next token prediction with self-attention. AISTATS 2024.  \n\nSoudry et al. \"The implicit bias of gradient descent on separable data.\" JMLR 2018.\n\n**Q2:** Existence of collocations seems far from reality; realistic data where such assumptions hold? How to prepare n() in real datasets? \n\n**A2:** Thanks for the question. The existence of collocations can be observed in many natural languages, where certain words frequently appear together (Lehecka 2015). Taking such an assumption enables us to derive interesting theoretical results, while remaining arguably reasonable in practice. Note that the assumption of separable data, as a special case of this, has formed an important line of theoretical research in deep learning to explore the implicit bias and training dynamics of neural networks (Soudry et al. 2018, Chizat and Francis 2020). \n\nGiven a dataset, $\\mathrm{n}$ can be prepared by the counting method described in Lehecka, T. (2015). More generally, it can be extracted by simply including all two-token sequences from the dataset, and each second token $y$ is the next token of the first token $x$, i.e., $n(x)=y$. Furthermore, leveraging domain-specific knowledge, such as medical dictionaries, can also enrich collocations. These approaches allow for the practical extraction of \\( \\mathrm{n}() \\) in a scalable and efficient manner.\n\nLehecka. Collocation and colligation. In Handbook of pragmatics online. Benjamins 2015.\n\nChizat and Francis Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. COLT 2020.\n\n**Q3:** How to understand definition of partial order. \n\n**A3:** Thanks for the question. We clarify that the partial order is **query-dependent**. Therefore, if $x>\\_{x_q} x'$ and $n(x)=x_{L+1}\\neq n(x')$ in some sentences ending with query $x_q$, it is **not** possible that $n(x)\\neq x_{L+1} = n(x')$ in another sentence ending with the same query $x_q$. \n\nThis is because of the structure of the transformer where the attention weights are proportional to $\\exp(x^\\top W_{kq} x_q)$. If $n(x)=x_{L+1}\\neq n(x')$ under query $x_q$,  then $x^\\top W_{kq} x_q > (x')^\\top W_{kq}x_q$ as long as the sentence ends with $x_q$. Thus, $x$ is always more important than $x'$ in $x_q$-ended sentences. However, $n(x)\\neq x_{L+1} = n(x')$ **is** still possible in other sentences ending with **different queries** than $x_q$.  These facts motivates us to introduce *query-dependent* partial orders.\n\n**Q4:** How commonly is the normalized gradient used? AdamW is typically used. \n\n**A4:** Thanks for the question. Although normalized gradient descent (NGD) is often favored in theoretical studies, there are also empirical works showing that the achieved accuracy of NGD is slightly better than Adam in training transformers (Cutkosky and Mehta 2020), where momentum is employed to improve the performance of NGD. Cutkosky and Mehta (2020) also theoretically proved that NGD achieves a fast convergence rate for generic optimization, which aligns with our theoretical findings that NGD achieves a linear convergence rate for NTP with one-layer transformers.\n\nWe acknowledge that AdamW is indeed more commonly used in practice. However, it is much more challenging to analyze the training dynamics under AdamW theoretically. So far, all theoretical analysis on transformers have been focused on GD and its variants. It is somewhat necessary to first develop techniques for GD and its variants, before these tools can be further advanced to study AdamW. Our work takes a first step to understanding the NGD dynamics of Transformers for NTP, and we are actively exploring the analysis for Adam-based Transformer training.\n\nCutkosky and Mehta, Momentum Improves Normalized SGD, ICML 2020.\n\n**Q5:** Assumption 3 may not hold in situations where attention is sparse. \n\n**A5:** Thanks for the question. We apologize that the statement of Assumption 3 may cause confusion. It actually requires that the number of optimal tokens is greater than or equal to the number of **each individual** non-optimal tokens, not the summation of all non-optimal tokens. Hence, such an assumption is not in conflict with sparse attention. One simple example satisfying Assumption 3 is a sentence where every token is distinct. The transformer's attention needs to focus only on the single optimal token, which is clearly sparse. \n\n---\nWe thank the reviewer again for the highly inspiring comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sGrTbd0UGV", "reviewer": "Reviewer_72WK", "summary": "The paper presents a non-asymptotic analysis of training dynamics for a single-layer transformer used in next-token prediction tasks. It introduces a two-stage training algorithm leveraging structural properties of the training dataset, defined via collocations and query-dependent partial orders. The findings include sub-linear convergence of both the feed-forward and self-attention layers to their respective max-margin solutions, and a linear convergence rate for the cross-entropy loss, supporting non-trivial prediction capabilities on unseen data. The approach is validated through theoretical analysis and empirical results, enhancing understanding of transformers' training and generalization behaviors.", "questions": "What are the potential implications of the findings on transformer training efficiency and computational costs in practical settings?", "limitations": "NaN", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper introduces novel theoretical frameworks for analyzing transformer training, focusing on non-asymptotic convergence.\n2. The research is technically robust, with sound mathematical derivations and empirical validation.\n3. The concepts and results are communicated effectively, albeit with room for improvement in some technical descriptions.", "weaknesses": "1. Some of the mathematical concepts, particularly the query-dependent partial orders, are complex and could be better explained.\n2. While the theoretical results are strong, additional empirical studies, particularly on real-world datasets, could further strengthen the claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments, and would like to provide the following responses. \n\n\n**Q1:** Some of the mathematical concepts, particularly the query-dependent partial orders, are complex and could be better explained.\n\n**A1:** We appreciate the feedback on the complexity of the mathematical concepts. To address this, we will provide more intuitive explanations and examples in the revised manuscript. \n\nSpecifically, query-dependent partial order formalizes the concept of importance or relevance of tokens in a given sentence. If one token is more important than others or more relevant to the next token, it is considered as an optimal token and thus is ''greater''  than other tokens.  For example,  assume that the sentence *Machine learning is a popular* has next word *area*, and *popular* is the most important word. Then, under the query *popular*, *popular* itself is the optimal token, and hence it is ''greater''   than other tokens.\n\n\nIn addition, such a partial order is **query-dependent** so that the order can be different given different queries. For example, consider a new sentence *A popular area is machine* with next word *learning*. Then, under the query *machine*, *machine* itself is more important or ''greater'' than *popular*. \n\nFinally, the motivation for defining these concepts is detailed in Response A1 to Reviewer ZtZ3.\n\n**Q2:** While the theoretical results are strong, additional empirical studies, particularly on real-world datasets, could further strengthen the claims.\n\n**A2:** Thanks for the suggestion. We are actively seeking suitable real-world datasets and experiments to further illustrate our results.\n\n**Q3:** What are the potential implications of the findings on transformer training efficiency and computational costs in practical settings?\n\n**A3:** Thanks for the question. Our results have the following implications. (1) The collocation training can be separated from the attention training to reduce the complexity of joint training, and at the same time achieve stable and fast convergence rate for each separate training. (2) Our result shows that the loss function converges fast (with potentially linear convergence rate), and hence the training does not require many iterations. The main computational costs lie in the gradient computations in each iteration, which scales with the number of training parameters. Thus, employing techniques such as LoRA to reduce the number of training parameters is crucial to achieve scalable training.\n\n---\n\nWe thank the reviewer again for the highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VE134z3q00", "reviewer": "Reviewer_ZtZ3", "summary": "In this work, the authors mathematically examine the learning dynamics of simple transformers for next token prediction. To allow a mathematical analysis, they consider a  highly simplified setting: the transformer consists of a single attention layer followed by a single feed-forward layer, the layers are trained one after the other (i.e., in a custom, decoupled training procedure), and the dataset adheres to a set of mathematical properties that make sure the loss can be arbitrarily close to zero. They show that both layers converge in direction to their corresponding max-margin solutions sub-linearly.", "questions": "How would you change the manuscript to make it more easy to understand your definitions and motivation behind the mathematics? \n\nHow would you change the manuscript to better assess/discuss the limitations of your work?", "limitations": "See above, no negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper adresses an important problem, namely how we could/should understand the learning dynamics of transformers. The paper seems a valid contribution to this problem, and provides a step in the direction of a better theoretical understanding. \n\nThe paper defines a set of mathematical properties a dataset for next-token-prediction should adhere to, to obtain a training error arbitrarily close to zero. This might be useful for future work in the community. \n\nThe paper is well-written, well-structured, the related works section is extensive and the work is put into the context of previous works. \n\nGreat care has been taken in the mathematical formulations and proofs; I wasn’t able to check the details, but the mathematical derivations seem correct and they are extensive.", "weaknesses": "The paper completely lacks a discussion of how general the obtained insights are; i.e., what the paper contributes beyond the setting of the simple network architecture, custom dataset and custom training method used to derive the results. It is understandable that assumptions and simplifications have to be made to make a problem amenable to a mathematical analysis, but the paper contains no discussion of these limitations nor any experiments to further examine this. This also makes it hard to judge the actual contribution of the work. \n\nThe current version of the paper is only easy to read for people working on this exact topic. To also be of interest to a somewhat broader audience (e.g., machine learning experts with a good grasp of mathematics), some concepts would need to be better motivated and/or explained in a more intuitive way. E.g., although the text mentions, ‘We first provide some intuitions about those two properties’ (L145), there seems to be no real intuitive explanation of why collocations and partial orders are introduced, and how this should be interpreted. The definition and arguments are purely mathematical. The example dataset provides some additional insights, but does not clarify why these properties are important, how this leads to zero training loss in the limit, etc…", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments, and would like to provide the following responses. \n\n**Q1:** How would you change the manuscript to make it more easy to understand your definitions and motivation behind the mathematics? For example, the example dataset provides some additional insights, but does not clarify why these properties are important, how this leads to zero training loss in the limit, etc. \n\n**A1:** Thanks for the suggestion. We will provide the following high-level description of our dataset and their properties. \n\nThe goal of introducing the collocation and query-dependent partial order is to make training transformer provable and interpretable. We breakdown the insights as follows.\n\n(i) Since the model should apply to all subsequences of sentences in the training set including length-two sentences, then collocation between two words naturally becomes a property of the dataset.\n\n(ii) It can be observed that the input of our loss function is a convex combination of $L$ points each being a projection of one token embedding in the sentence, where the weights in the convex combination are the attention scores of each token with the query token. In this way, the loss values of these tokens determine an order among them, and the token that achieves the minimum loss value serves as the optimal token. In other words, some tokens has larger attention scores than others under a specific query. These facts motivate us to introduce query-dependent partial order. \n\nCombining (ii) with (i), the collocated token of the optimal token becomes the predicted next token. Finally, due to the injection property of collocation and the structure of cross-entropy, the loss function can be trained to approach zero.\n\n**Q2:** How would you change the manuscript to better assess/discuss the limitations of your work? For example, provide discussion of how general the obtained insights are; i.e., what the paper contributes beyond the setting of the simple network architecture, custom dataset and custom training method used to derive the results.\n\n**A2:** Thanks for the suggestion. We provide the following discussion on the generality of our findings and will add it to our revision.\n\nWe remark that our work establishes the basic framework to analyze the convergence performance of transformers for NTP task. Several insights can be extended to more complex settings. \n\nRegarding multi-layer transformers, for example, one possible idea is to view the single layer in our work as the last layer of a multi-layer transformer. This is promising if we consider the last layer training regime (Kirichenko, et al., 2022). In this regime, if the output of other layers satisfies certain conditions such as near-orthogonality, which is often satisfied in high dimensional setting, our convergence and generalization results can be directly applied to multi-layer transformers.\n\nTo generalize dataset, for example, our technique can be extended to study a noisy version of collocation. We expect that most of our proof steps can still be applied by incorporating appropriate noise bound. \n\nTo generalize training method, an alternative natural method is to update the linear layer and the attention layer simultaneously instead of two-stage training. A popular way to implement such an algorithm is via two timescale stepsizes so that the attention update can be in a faster scale and the final convergence can be determined by the feedforward layer by incorporating the suboptimality error of the attention layer. Such suboptimality error can be bounded by our current techniques for training stage 2.\n\nKirichenko et al. \"Last layer re-training is sufficient for robustness to spurious correlations.\" arXiv:2204.02937, 2022.\n\n---\n\nWe thank the reviewer again for the highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bTUloaaYzc", "reviewer": "Reviewer_jsSk", "summary": "This manuscript focuses on the next-token prediction (NTP) task, and provides a fine-grained non-asymptotic analysis of the training dynamics of a one-layer transformer. Specifically, the authors first characterize the essential structural properties of training datasets for NTP using a mathematical framework based on partial orders. Then, they design a two-stage training algorithm, where the pre-processing stage for training the feed-forward layer and the main stage for training the attention layer exhibit fast convergence performance. Finally, they show that the well-trained transformer can have non-trivial prediction ability on unseen data, which sheds light on the generalization capability of transformers.", "questions": "Please see the weakness", "limitations": "This manuscript simplifies the Transformer into a single layer for analysis. While I understand the need for this simplification, it would be nice to include a brief paragraph outlining the views on generalizing to multi-layer Transformers.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This manuscript conducts a theoretical analysis on the convergence speed and generalization of the important NTP task, with giving an example to illustrate. Although not comprehensive, it can serve as a stepping stone for convergence analysis of NTP tasks.\nThis manuscript is well-written and easy to read.", "weaknesses": "1. How is the loss function between lines 227 and 288 obtained? Why is its form different from Eq. 1?\n2. In Section 6, the authors theoretically prove the generalization ability of the trained Transformer. Can the generalization ability be proved in the experimental part?\n3. Why do the authors design a two-stage training algorithm? Is it valuable in practical applications compared to single-stage training?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments, and would like to provide the following responses. \n\n**Q1:** How is the loss function between lines 227 and 228 obtained? Why is its form different from Eq. 1?\n\n**A1:** The loss in Eq.\\ (1) is for the general-length case. The loss between lines 227 and 228 is for the training stage 1 (i.e., training of feed-forward layer), where the inputs are one-token sentences. We note that there was a typo in Eq.\\ (1), i.e., missing $\\exp$ in the first bracket. The correct form should be\n\\\\[ \\log \\left(\\sum_v \\exp( e_v^\\top \\bar{T}(X) ) \\right) - e_i^\\top \\bar{T}(X) ,\\\\]\nwhere $i$ is the index of the next token of $X$. Such a form can specialize to the loss for stage 1 as follows. Let $X = [x_1]$. Then we have $\\bar{T}(X) = W_{ov}x_1$. Thus, the loss becomes\n\\\\[\\log\\left( \\sum_v \\exp(e_v^\\top W_{ov}x_1) \\right) - \\log\\exp(e_i^\\top W_{ov} x_1) = -\\log\\frac{ \\exp(e_i^\\top W_{ov} x_1) }{ \\sum_v \\exp(e_v^\\top W_{ov}x_1) }.\\\\]\nThank you for your question and we will fix the typo in the revision.\n\n**Q2:** In Section 6, the authors theoretically prove the generalization ability of the trained Transformer. Can the generalization ability be proved in the experimental part?\n\n**A2:** Yes. Please see the attached PDF file in the global response for the test loss result. For the test dataset, we randomly construct test data as described in Theorem 3, and plot the average test loss for 20 trials. As shown in the figure, the test loss also converges almost to 0. In addition, since the generalization ability is determined by the parameter $W_{kq}$ and $W_{ov}$,  alignment of these parameters with their corresponding max-margin solutions as shown in Figure 2 of the paper implies the generalization ability as well.\n\n**Q3:** Why do the authors design a two-stage training algorithm? Is it valuable in practical applications compared to single-stage training?\n\n**A3:** Thanks for the question. Our characterization of the realizable dataset is via two steps: (i) identifying collocation and (ii) using collocation to define partial ordering and optimal token. This naturally motivates the two training stages to learn the corresponding information: (i) stage 1 learns collocation, i.e., next-token prediction for length-one sentences; and (ii) stage 2 learns partial ordering to identify optimal token. Thus, it is also meaningful to apply such a two-stage algorithm in practice. We also expect that stage 1 of pre-processing value matrix using one-token sentences could help to stablize the entire training process of transformers. Single-stage training may still yield good solution simply from optimization perspective, and we are currently working on analytically characterizing the dynamics of such a training process.\n\n**Q4:** This manuscript simplifies the Transformer into a single layer for analysis. While I understand the need for this simplification, it would be nice to include a brief paragraph outlining the views on generalizing to multi-layer Transformers.\n\n**A4:** Thanks for the suggestion. A possible approach to generalize it into multi-layer transformers is to view our single layer as the last layer of a multi-layer transformer. Therefore, the input $X$ in the paper can be viewed as the output of a multi-layer network. Then, we can combine our technique for the last layer and the technique of neural tangent kernel (NTK) for other layers (Emami et al., 2021), and analyze the training dynamics of such a multi-layer transformer. We expect that the last layer evolves similarly to our findings given that other layers output $X$ satisfies certain conditions, which may be potentially shown by exploiting the realizability of the dataset and the properties of NTK. \n\nEmami et al. \"Implicit bias of linear RNNs.\" ICML 2021.\n\n---\n\nWe thank the reviewer again for the highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their feedback, which will greatly improve our paper. 

In the attached PDF file, we provide a figure for an additional experiment verifying the generalization ability described in Theorem 3, as a response to the second question Q2 of Reviewer jsSk. In the experiment, we construct a test dataset described in Theorem 3 and plot the average test loss of 20 times against the training loss. As shown in the figure in the attached PDF, the test loss also converges almost to 0.

Meta-review: In this submission, the authors present interesting theoretical insights into transformer training dynamics for next-token prediction, supported by detailed mathematical analysis. Although the setting is simplified, the reviewers concur that it contributes to understanding transformers. The authors have also addressed reviewers' concerns  in their rebuttal. While further empirical validation on real-world datasets would enhance the work's impact, the theoretical contributions in their current form provide a solid foundation for advancing the understanding of transformer training dynamics. To sum up, the decision is to accept the paper.
Decision: Accept (poster)