Non-asymptotic Convergence of Training Transformers for Next-token Prediction

Authors: ruiquan  huang
yingbin  liang
jing  yang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:32:09

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.

Main paper: 
 introduction: 
 The transformer architecture (Vaswani et al., 2017) has revolutionized the field of machine learning, establishing itself as a foundation model for numerous applications, including natural language processing (NLP) (Devlin et al., 2018) , computer vision (Dosovitskiy et 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 et 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 et al., 2021; Belkin, 2024) .To overcome these issues, recent advances in transformer theory have investigated the convergence of training transformers under theoretically amenable setting such as linear regression (Mahankali et al., 2023; Zhang et al., 2023; Huang et al., 2023) and binary classification (Tarzanagh et al., 2023b,a; Vasudeva et al., 2024; Li et al., 2023) . Nevertheless, one of the fundamental task in LLMs and other generative models is next-token prediction (NTP), which involves predicting the next word or token in a 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-asymptotic convergence analysis of the training process, posing the following open questions for further investigation:How fast does the training of a 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 on 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 a self-attention layer and a feed-forward layer for NTP. We summarize our contribution as follows.• We develop a 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. Query-dependent partial orders is a set of partial orders where each partial order classifies tokens into three categories: optimal tokens, non-optimal tokens and non-comparable tokens. These structural properties define favorable max-margin problems on both the feed-forward layer and the self-attention layer. • Second, we design a two-stage training algorithm based on normalized gradient descent. In stage 1 of pre-processing, we use the collocation to train the feed-forward layer. In stage 2, we use the entire dataset to train the self-attention layer. We show that the feed-forward layer and the query-key attention matrix converge sublinearly in direction respectively to the max-margin solution for classifying next token from all other tokens in the preprocessing dataset, and to the max-margin 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 1'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-trivial prediction on unseen data. In particular, the transformer is trained to learn an extended querydependent partial order, where the non-comparable 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 a new sentence and further make desirable prediction.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 4, Section 5, and Section 6, 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 in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Inspired by Brown et al. (2020) , who demonstrated that pre-trained transformers can 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 et al., 2022; Bai et al., 2023; Von Oswald et al., 2023; Fu et al., 2023; Giannou et al., 2023; Lin et al., 2023) . 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 a sparse function.Regarding the training dynamics and optimization of transformers under in-context learning, Ahn et al. (2024) ; Mahankali et al. (2023) ; Zhang et al. (2023) ; Huang et al. (2023) studied the dynamics of a single attention layer, single-head transformer for the in-context learning of linear regression tasks. Cui et al. (2024) 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. Kim and Suzuki (2024) studied the nonconvex mean-field dynamics of transformers, and Nichani et al. (2024) established a convergence rate of Õ(1/t) for the training loss in learning a causal graph. Additionally, Chen et al. (2024) investigated the gradient flow in training multi-head attention. Chen and Li (2024) proposed a supervised training algorithm for multi-head transformers.Another line of research focuses on the training dynamics of transformers for binary classification problems. Tarzanagh et al. (2023b,a) demonstrated an equivalence between the optimization dynamics of a single attention layer and a certain SVM problem. While Tarzanagh et al. (2023b,a) only proved an asymptotic convergence result, Vasudeva et al. (2024) improved the convergence rate to t -3/4 . 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 kernel (NTK) regime.For transformers trained on next-token prediction (NTP), Tian et al. (2023a) analyzed the training dynamics of a single-layer transformer, while Tian et al. (2023b) studied the joint training dynamics of multi-layer transformers. Li et al. (2024) demonstrated the asymptotic convergence of transformers trained with a 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 (2024) 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 (2021) ; 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 ReLU/Leaky-ReLU 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.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Our paper has specified the assumptions (Assumption 1,Assumption 2, and Assumption 3) under which our results hold.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
problem setup: 
 Notations. All vectors considered in this paper are column vectors. We use 1{A} to denote the indicator function of A, i.e., 1{A} = 1 if A holds, and 1{A} = 0 otherwise. ∥W ∥ represents the Frobenious norm of the matrix W . For a vector v, we use [v] i to denote the i-th coordinate of v.We use ϕ(v) to denote the softmax function, i.e., [ϕ(v)] i = exp(v i )/ j exp(e ⊤ j v), which can be applied to any vector with arbitrary dimension. We use {e i } i∈ [|V|] to denote the canonical basis of R |V| , i.e., [e i ] j = 1{i = j}. The inner product ⟨A, B⟩ of two matrices A, B equals to Trace(AB ⊤ ).Next-token prediction. We consider the task of next-token prediction, which aims to predict the subsequent token in a token sequence given its preceding tokens. Formally, suppose that there exists a finite vocabulary set V ⊂ R d that consists of all possible tokens, where d is the dimension of the embedding. Each token x ∈ V is associated with a unique index I(x) ∈ {1, 2, . . . , |V|}, where I is the index function. An L-length sentence X = [x 1 , . . . , x L ] ∈ V L ⊂ R d×L is a sequence of L tokens, where L is an integer. We assume that the maximum length of sentences is L max . The subsequent tokens in sentences are generated from a set of ground-truth model {p * L : V L → V} L<Lmax , where p * L generates the next token x L+1 given the sentence X for any 1 ≤ L < L max . The task of next-token prediction requires us to learn all models {p * L } L<Lmax given a training datasetD 0 = {(X, x L+1 )|L < L max , X ∈ V L , x L+1 ∈ V}. Notably, if X = [x 1 , . . . , x L ] ∈ D 0 , then for any ℓ < L, ([x 1 , . . . , x ℓ ], x ℓ+1) is also a training sample, since it follows p * ℓ as well. Decoder-only transformer. A decoder-only transformer is a stack of blocks consisting of a selfattention layer and a feed-forward layer. For simplicity, we consider one-layer transformer, where the self-attention layer is determined by three matrices:W k ∈ R d×d1 , W q ∈ R d1×d and W v ∈ R d2×d ,namely key, query, and value matrices, and the feed-forward layer is determined byW o ∈ R |V|×d2 .Here d 1 , d 2 are hidden dimensions. Mathematically, given the input X = [x 1 . . . , x L ], we write the one-layer transformer asT θ (X) := ϕ(W o W v Xϕ(X ⊤ W k W q x L )) ∈ [0, 1] |V| , where θ := (W o , W v , W k , W q ),and ϕ is the softmax function. We note that the inner softmax function ϕ is 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 a unified matrix W kq , such that W kq = W k W q . Similarly, we reparameterize the product of the feed-forward (W o ) and value (W v ) matrices as a single matrix W ov , defined as W ov = W o W v . Such a 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 those reparameterization is given by Tθ (X) := ϕ(W ov Xϕ(X ⊤ W kq x L )) ∈ [0, 1] |V| .Cross-entropy loss. Given the training dataset D 0 and the transformer model, we seek to learn p * by minimizing (training) the cross-entropy loss L(θ) defined as follows:L(θ) = - 1 |D 0 | (X,x L+1 )∈D0 log e ⊤ I(x L+1 ) T θ (X),where I(x L+1 ) is the index of x L+1 in V.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Learning. There are many potential societal consequences of our 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 or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
realizable training dataset and two-stage algorithm: 
 In this section, we first provide a mathematical framework based on partial order to formally characterize a realizable training dataset for next-token prediction. We will then describe a two-stage algorithm for next-token prediction that we study.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 sentenceX = [x 1 , . . . , x L ] is a legal training sample, ([x 1 ],x 2 ) is also in the training dataset. In addition, the output of a transformer given one single input token only depends on W ov , i.e. the feed-forward layer. Since training loss can be arbitrarily close to 0, there exists a sequence {W t } such that lim t→∞ -x∈D0 log e ι (x) ⊤ ϕ(W t x) = 0, where ι(x) is the index of next token of x, and the summation is over the case when x is the first token. Due to that ϕ(W t x) is a probability distribution, the equality holds only when ι is injective, since otherwise it is an entropy of some distribution which is strictly greater than 0. Therefore, there exists an injective map n : V → V such that every sentence starts with x, must have a unique next token n(x). We call the set of pairs {x, n(x)} x∈V a collocation. We remark that p * 1 = n. Existence of "order". Second, let us consider the output of a transformer T θ given a legal sentence X = [x 1 , . . . , x L ] with the next token x L+1 = p * L (X). The transformer first calculates a convex combination of x 1 , . . . , x L with corresponding weight φ ℓ ∝ exp(x ⊤ ℓ W kq x L ) for each ℓ ≤ L. Then, the transformer outputs ϕ( ℓ W ov x ℓ φ ℓ ). Recall that the collocation forces x ℓ to map to n(x ℓ ), thus ϕ(W ov x ℓ ) has a peak value at the coordinate equal to I(n(x ℓ )) (the index of n(x ℓ )). Hence, T θ (X) can only have peak value at the coordinates within the set {I(n(x ℓ ))} ℓ≤L . If the training loss can be arbitrarily close to 0, it is desirable to have n -1 (x L+1 ) ∈ {x ℓ } ℓ≤L . Therefore, for those x ℓ with n(x ℓ ) = x L+1 , φ ℓ must be larger than φ ℓ ′ with n(x ℓ ′ ) ̸ = x L+1 . Finally, it worth noting that φ ℓ depends on the final token x L . This observation motivates us to define query-dependent partial orders on V.Definition 1 (x q -partial order) Fix a token x q . An x q -partial order assigns an ordering relationship > x q for certain pairs of tokens in V, and is created as follows. Let D x q 0 be the set of all legal sentences in the training dataset that has the final token (query) x q . Then, for any pair of tokens x, x ′ ∈ V, we assign x > x q x ′ if there exists a sentence X = [x 1 , . . . , x L ] ∈ D x q 0 and x, x ′ are tokens in X such that n(x) = x L+1 ̸ = n(x ′ ), where x L+1 is the next token of X.Note that Definition 1 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 0, the aforementioned discussion shows that if x > x q x ′ , then φ ℓ > φ ℓ ′ , where x = x ℓ and x ′ = x ℓ ′ in some sentence. Thus, exp(xW kq x q ) > exp(x ′ W kq x q ), which indeed need to be well-defined. Otherwise, we will have contradictions such as exp(xW kq x q ) > exp(x ′ W kq x q ) < exp(xW kq x q ). Mathematically, a well-defined (strict) partial order > on a set V satisfies two axioms (Yannakakis, 1982) : (i) there is no x > x; (ii) if x > x ′ and x ′ > x ′′ , then x > x ′′ . Thus, x q -partial order created by D 0 is well-defined for every x q ∈ V.Finally, let us discuss the impact of query-dependent partial orders on D 0 . For a given query x q , the partial order > x q divides tokens in V into four disjoint types.• (Strict) optimal tokens. A token x is optimal, if there is no x ′ such that x ′ > x q x 1 . • Confused tokens. A token x is confused, if there exists x ′ , x ′′ such that x ′ > x q x > x q x ′′ . • (Strict) non-optimal tokens. A token x is non-optimal if there is no x ′ such that x > x q x ′2 . • Non-comparable tokens. A token x is non-comparable if there is no x ′ such that x > x q x ′ orx ′ > x q x.In this work, we assume that there are no confused tokens. This assumption simplifies the problem, making it tractable to provide explicit convergence in direction for training a transformer in Section 5. In summary, we make the following structural assumption on the training dataset.Assumption 1 (Realizable training dataset) D 0 admits (i) a collocation {x, n(x)} x∈V ; (ii) welldefined query-dependent partial orders, where every x q -partial order has no confused tokens.We remark that combining the collocation and query-dependent partial orders, we can regenerate the training dataset as follows. For any sentence with only one token X = [x], the next token is n(x).For other sentences X = [x 1 , . . . , x L ], let x ℓ be optimal under the partial order > x L , and then the next token of X is n(x ℓ ). We next provide a simple example that justifies Assumption 1.Example 1 Consider a language system where the vocabulary consists of four tokens {S, V, O, P}, where S,V,O,P respectively stand for subject, verb, object, and punctuation mark. This system admits the commonly adopted word order (Dryer, 1991) : S, V, O, P. Let the training dataset be {SVOP, VOP, OPP, PSV}.Let us create the corresponding collocation and the query-dependent partial orders from the dataset. The collocation is {(S, V), (V, O), (O, P), (P, S)}. That is, if a sentence starts with a subject, then the next token is a verb. Similarly, if a sentence starts with a verb, then the next token is an object, and so on. The query-dependent partial orders are created as follows:Partial order under query S. S> S P. Partial order under query O. O> O S, O> O V.Partial order under query V. V> V S. Partial order under query P. O> P P.Therefore, if a sentence starts with S (subject), the next token is V (verb) according to the collocation. Then, for the sentence SV, since the query is V and V> V S, the next token of the sentence coincides with the next token of 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 1 and aligns with real-world scenarios. An illustration is provided in Figure 1 . Additional notations of training data. It is worth noting that there are 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 D 0 indexed by n ∈ {1, . . . , N }. For each distinct sentenceX (n) , we calculate its frequency π (n) ∈ [0, 1] in dataset D 0 as π (n) = (X,x L+1 )∈D 0 1{X=X (n) } |D0| .Building upon this, with a little abuse of notation, we use n(X (n) ) ∈ V to denote the subsequent token of the sentence X (n) and In(X (n) ) to denote the index of n(X (n) ).We further denote X(n)-1 as the final token of X (n) , and let Tθ (X) = W ov Xϕ(X ⊤ W kq X(n)-1 ). Then, the loss function L(θ) can be rewritten as follows:L(θ) = n π (n) log v exp e ⊤ v Tθ (X (n) ) -e ⊤ In(X (n) ) Tθ (X (n) ) .(1)For the realizable dataset satisfying Assumption 1, we propose a two-stage training algorithm using normalized gradient descent (NGD). The pseudo code of the algorithm is presented in Algorithm 1.In Section 5, we show that the two-stage algorithm decouples the training of the feed-forward and attention layers without losing the optimality. This is because the training in stage 1 is designed to yield a suitable max-margin solution, which will enable the training of stage 2 to solve a 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 W ov .For simplicity, we introduce the following notation for the training loss of the feed-forward layer. Given a collocation {x, n(x)} x∈V , which can be obtained through extracting all length-2 sentences in the training dataset D 03 , we use normalized gradient descent to train W ov . Equivalently, the loss function can be written asL 0 (W ov ) = - x∈V log exp(e ⊤ In(x) W ov x) v≤|V| exp(e ⊤ v W ov x),where the self-attention elements are removed because the attention matrices are not trained here. Based on the above loss function, we initialize W (0) ov = 0 ∈ R |V|×d , and subsequently take an update at each time t by NGD as in line 4 of Algorithm 1.In the second stage, we fix the trained feed-forward layer and train the self-attention layer based on the loss function given in Equation (1) and using the entire dataset D 0 . Specifically, we initialize W kq = 0 ∈ R d×d , and subsequently take an update at each time t by NGD as in line 7 of Algorithm 1.Algorithm 1 Two-stage Normalized Gradient Descent 1: Initialization:W (0) ov = 0 ∈ R |V|×d , W kq = 0 ∈ R d×d . 2: Input: A collocation {x, n(x)} x∈V , and a training dataset D 0 , learning rate η 0 , η. 3: for t ∈ {0, 1, ..., T -1} do 4: Update W (t+1) ov as W (t+1) ov = W (t) ov -η 0 ∇ Wov L0(W (t) ov ) ∥∇ Wov L0(W (t) ov )∥ . 5: end for 6: for t ∈ {0, . . . , T 1 -1} do 7: Update W (t+1) kq as W (t+1) kq = W (t) kq -η ∇ W kq L(θ (t) )∥∇ W kq L(θ (t) )∥ , where θ (t) = (W 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., 2023a) .Assumption 2 The vocabulary set is orthornormal. Namely, the embedding has unit norm, i.e., ∥x∥ = 1, and x ⊤ x ′ = 0 holds for any distinct tokens x and x ′ . 
convergence of training w ov: 
 To characterize the training dynamics of W ov , we observe that the collocation {(x, n(x))} x∈V defines the following hard-margin problem:EQUATIONIt can be shown that lim B→+∞ L 0 (BW * ov ) = 0. Thus, the loss function L 0 trains W ov to be the max-margin solution with W ov x distinguishing the next token n(x) from all other tokens in V. Since L 0 (•) is convex,) = O(exp(-C 0 t)) for some constant C 0 .Recall that after the training stage 1 with T steps, we obtain a trained feed-forward layer W (T ) ov . Then, we fix W (T ) ov and use normalized gradient descent to train W kq . To characterize the training dynamics of the key-query matrix W kq , we note that each query-dependent partial order also defines a hard-margin problem. Let l(n) ⊂ {1, . . . , L (n) } be the set of indices of the optimal tokens of X (n) . Recall that x ℓ is optimal if there is nox ℓ ′ such that x ℓ ′ > X (n) -1 x ℓ and In(x ℓ ) = In(X (n) ). That is, W kq X (n)-1 should correctly classify optimal tokens x ℓ and non-optimal tokens x ℓ ′ . This is formalized in the following problem:EQUATIONWe will show that the loss function in Equation ( 1) given the well trained W (T ) ov will train W kq towards the max-margin solution W * kq 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 inX (n) , we have |l(n)| ≥ ℓ 1{x ℓ = x}.Assumption 3 is 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 2.Theorem 1 Let Assumptions 1-3 hold. Let W * kq be the solution of Equation (3). Let η < O(1) and W (t) kq be updated by Algorithm 1. Then, for any t ≥ 2, we have that tη2∥W * kq ∥ ≤ ∥W (t)kq ∥ ≤ tη. In addition, the following inequality holds.W (t) kq ∥W (t) kq ∥ , W * kq ∥W * kq ∥ ≥ 1 - 54N L 4 max ∥W * kq ∥ 4 log 2 t tη .Theorem 1 states that the key-query matrix W kq converges in direction to the max-margin solution W * kq /∥W * kq ∥ at a convergence rate of O(log 2 t/t). We further show that the norm of W kq also grows linearly in t, i.e., ∥W kq ∥ = Ω(t). Combining these results, we have the following theorem on the convergence of the loss function and the training accuracy.Theorem 2 (Loss Convergence) For any training sentenceX (n) = [x (n) 1 , . . . , x (n) L ], let φ (n,t) ℓ ∝ exp(x (n) ℓ W (t) kq x (n)L ) be the attention weight. Under the conditions in Theorem 1, there is an absolute constant c 0 such that when T ≥ c 0 ∥W * ov ∥ 5 log(|V|) log T /η 0 and t ≥ c 0 N L 4 max ∥W * kq ∥ 6 log 2 t/η, the optimal token weight satisfiesmin n ℓ * ∈l(n) φ (n,t) ℓ * ≥ (1 + L max exp (-tC 1 ))) -1 .In addition, the loss function L converges linearly 4 to its minimal value:EQUATION4 For fixed W (T ) ov , the minimum loss value L * = |V| exp(-T C0). Then Equation (4) implies L(θ (t) )-L * ≤ L * T C0O(e -C 1 t ) for sufficiently large t, which further implies the linear convergence in t.whereC 0 = η0 4∥W * ov ∥ 2 and C 1 = η/(4L max ∥W * kq ∥ 2 ).Theorem 2 shows that the training loss converges to its minimum value at a linear convergence rate. Furtherm, for T = Ω(log(1/ϵ 0 )) t = Ω(log(1/ϵ)), the optimal token weight is given by 1/(1 + ϵ) 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(n) ℓ * ) with probability 1 -O(ϵ 0 ).The proof consists of the following three main steps. The key proof step lies in carefully analyzing the projection of gradient ∇ W kq L(θ (t) ) onto the token-query outer productEQUATION-1 ) ⊤ , max-margin attention weight matrix W * kq , and the trained attention weight matrix W (t) kq .Step 1 (Lemma 5). By analyzingEQUATION-1 ) ⊤ , we characterize the dynamics of attention weights. Using mathematical induction, we show that the lower bound of optimal token weight is 1/L max .Step 2 (Lemma 6). Then, we show that the cosine similarity between the negative gradient and W * kq is strictly larger than the minimum optimal token weight. Utilizing step 1, due to the NGD update, the norm of the key-query matrix W (t) kq can be shown to grow linearly.Step 3 (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 on those projections, we can show the following inequality for some constant C 0 :∇ W kq L(θ (t) ), W (t) kq ≥ 1 + C 0 log ∥W (t) kq ∥ ∥W (t) kq ∥ ∇ W kq L(θ (t) ), W * kq ∥W (t) kq ∥ ∥W * kq ∥ . Utilizing step 2's result that ∥W (t)kq ∥ 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, we prove the generalization ability of the trained transformers. Recall that Theorem 1 shows thatW (t) kq converges to W * kq ∥W (t)kq ∥/∥W * kq ∥. To characterize the generalization ability, it is desirable to use the property of W * kq , which is given in the following result.Proposition 2 Under Assumptions 1-2, fix a query token x q , let O x q , N x q , M x q ⊂ V be the set of optimal tokens, the set of non-optimal tokens, and the set of non-comparable tokens, under x q -partial order, respectively. Then, the solution W * kq of Equation (3) satisfies x ⊤ 0 W * kq x q = 0 for x 0 ∈ M x q , andx ⊤ * W * kq x q = |N x q | |O x q | + |N x q | , x ⊤ W * kq x q = - |O x q | |O x q | + |N x q | , ∀x * ∈ O x q , x ∈ N x q .Recall that non-comparable tokens (see Section 4) under a query x q never appears in any training sentence data with the same query x q . Thus, Proposition 2 implies an interesting generalization capability -each x q -partial order can automatically incorporate more relationships to expand the query-dependent partial orders. Combining Proposition 2 with Theorem 1, we obtain the following theorem on W (t) kq .Theorem 3 Under the conditions and notations in Proposition 2, let T = Ω(log(1/ϵ)), and t = Ω(log(1/ϵ)). Then there exists a constant C 0 such that(x * -x 0 ) ⊤ W (t) kq x q ≥ C 0 t, (x 0 -x) ⊤ W (t) kq x q ≥ C 0 t, ∀x * ∈ O x q , x 0 ∈ M x q , x ∈ N x q .Moreover, if the trained transformer takes input X with query x q that consists of a non-comparable token x 0 and non-optimal tokens, then the prediction made by T θ (t) (X) is n(x 0 ) with high probability.Theorem 3 suggests that a new partial order is created by the trained transformer. Specifically, it 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 1 to illustrate the generalization ability described above.Example 2 (Generalization to unseen data in Example 1) Recall that in Example 1, the training dataset consists of four sentences: SVOP, VOP, OPP, and PSV. Consider the partial order > P under the punctuation mark P. We have that O> P P and O is an optimal token, P is a 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 a new (unseen) input sentence SP. Since S is non-comparable before training, but is "larger" than P under the trained key-query matrix W (t) kq , the next predicted token is n(S) = V. Case 2. Optimal tokens remain optimal over all tokens after training.Consider a new (unseen) input OSP. O is optimal and S 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-punctuation), the output is SPV (subject-P-verb), which follows a logical subject-verb order and is desirable. However, in cases where the input is VP (verb-punctuation), it may be preferable to terminate the sequence after the verb, i.e., VPP, as the verb alone can suffice to convey the intended meaning. 
experiment: 
 In this section, we verify our theoretical findings via an experiment on a synthetic dataset. Specifically, we randomly generate a realizable dataset as described in Assumption 1 with |V| = 20. Then, we train W  
conclusion: 
 In this work, we investigated the training dynamics of a 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 to zero. These properties allow us to define two max-margin solutions for both the feed-forward layer and the self-attention layer. Then, we showed 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 our theoretical findings.  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: 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 a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing eoective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith eoort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: Our paper is primarily a theoretical work, and does not use any assets. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: Our paper is primarily a 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 was obtained from people whose asset is used.  
 
 Appendix: We first provide the general formula for the gradients of both layers.∇ Wov L 0 (W ov ) = x∈V T 0 (x) -e In(x) x ⊤ ,(5)EQUATIONRecall that we use the loss,L 0 (W ov ) = - x∈V log exp e ⊤ In(x) W ov x i∈[|V|] exp e ⊤ i W ov x .The updating rule of W ov is thatEQUATIONWe know that L 0 is convex respect to W ov . Therefore, we haveW (t) ov -W ′ ov , ∇ Wov L 0 (θ (t) ) ≥ L 0 (θ (t) ) -L 0 (θ ′ ).It is clear that the loss function L reaches the minimum 0 when W ov = ∆W * ov , as ∆ → ∞.Lemma 1 Under the initialization W (0) ov and the updating rule Equation (7) with step size η, the following inequality holds.tη 0 + ∥W (0) ov ∥ ≥ ∥W (t) ov ∥ ≥ tη 0 2∥W * ov ∥ -∥W (0) ov ∥.Proof. Using Equation ( 5), we haveEQUATIONwhere (a) is due the constraints that W * ov satisfies. On the other hand,∥∇ Wov L(W (t) ov )∥ = ∇ Wov L(W (t) ov ) ∥∇ Wov L(W (t) ov )∥ , ∇ Wov L 0 (W (t) ov ) = x∈V i [T (t) 0 (x)] i (e i -e In(x) ) ⊤ ∇ Wov L(θ (t) ) ∥∇ Wov L(θ (t) )∥ x (a) ≤ 2 x∈V i̸ =In(x) [T (t) 0 (x)] i ,where (a) follows from ∥AB∥ ≤ ∥A∥∥B∥ for any matrices A and B. Thus, we obtainW * ov , ∇ Wov L(θ (t) ) ∥∇ Wov L(θ (t) )∥ ≤ -1/2To lower bound the norm of W ov , we recall the updating rule (Equation ( 7)).∥W (t) ov ∥ = W (0) ov - t ′ <t η 0 ∇ Wov L 0 (W (t ′ ) ov ) ∥∇ Wov L 0 (W (t ′ ) ov )∥ ≥ W (0) ov - t ′ <t η 0 ∇ Wov L 0 (W (t ′ ) ov ) ∥∇ Wov L 0 (W (t ′ ) ov )∥ , W ov ∥W * ov ∥ ≥ tη 0 2∥W * ov ∥ -∥W (0) ov ∥.For the LHS of the inequality, it suffices to note that at each iteration, the norm ∥W(t)ov ∥ increases most η 0 due to normalized gradient descent.Lemma 2 At each iteration t, the following inequality holds.W (t) ov ∥W (t) ov ∥ , ∇ Wov L 0 (W (t) ov ) ≥ 1 + 2∥W * ov ∥ ∥W (t) ov ∥ log(2|V|) W * ov ∥W * ov ∥ , ∇ Wov L 0 (W (t) ov )Proof. First, we consider the case whenW (t) ov = W * ov ∥W (t) ov ∥ ∥W * ov ∥ . Due to Equation (8) in Lemma 1, we have W * ov ∥W * ov ∥ , ∇ Wov L 0 (W (t) ov ) < 0In this case, the result is trivial.Then, we consider the case whenW (t) ov ̸ = W * ov ∥W (t) ov ∥ ∥W * ov ∥ . Due the optimality of W *ov , which achieves the minimum norm satisfying the constraints in Equation ( 2), we must have that for some x 0 ∈ V, there exists i 0 ̸ = In(x 0 ) such that the following inequality holds(e In(x0) -e i ) ⊤ W (t) ov x 0 < ∥W (t) ov ∥ ∥W * ov ∥.Therefore, the loss onW (t)ov can be lower bounded as follows.L(W (t) ov ) = x∈V log 1 + i exp (e i -e In(x) ) ⊤ W (t) ov x > log 1 + exp -∥W (t) ov ∥/∥W * ov ∥ (a) > 1 2 exp -∥W (t) ov ∥/∥W * ov ∥ ,where (a) is due to the fact that log(1 + x) ≥ x/2 when 0 < x < 1. On the other hand, letW ′ ov = ∥W (t) ov ∥ ∥W * ov ∥ + 2 log(2|V|) W * ov .Then, the loss on W ′ ov has the following upper bound.L(W ′ ov ) = x∈V log 1 + i exp (e i -e In(x) ) ⊤ W (t) ov x ≤ x∈V log 1 + (|V| -1) exp -∥W (t) ov ∥/∥W * ov ∥ -log(2|V|) (a) ≤ x∈V |V| exp -∥W (t) ov ∥/∥W * ov ∥ -2 log(2|V|) ≤ 1 2 exp(-∥W (t) ov ∥/∥W * ov ∥),where (a) is due to the fact that log(1 + x) < x when x > 0. Thus, L(W(t) ov ) > L(W ′ ov ).Due to the convextiy of L 0 , we have0 < W (t) ov -W ′ ov , ∇ Wov L 0 (W (t) ov ) = W (t) ov , ∇ Wov L 0 (W (t) ov ) - ∥W (t) ov ∥ ∥W * ov ∥ + 2 log(2|V|) W * ov , ∇ Wov L 0 (W (t) ov ) ,which finishes the proof.Proposition 3 (Restatement of Proposition 1) Under the zero initialization W (0) ov = 0 and updating rule Equation (7), for any t ≥ 2, the following inequality holds.W (t) ov ∥W (t) ov ∥ , W * ov ∥W * ov ∥ ≥ 1 - 12∥W * ov ∥ 3 log(2|V|) log t tη 0 . Moreover, tη0 2∥W * ov ∥ ≤ ∥W (t) ov ∥ ≤ tη 0 .Proof. The second argument about the norm of W (t) ov follows directly from Lemma 1. We aim to prove the first part as follows.Let α t = 2∥W * ov ∥ ∥W (t)ov ∥ log(2|V|). By Lemma 2 and the updating rule Equation (7), we haveW (t+1) ov -W (t) ov , W * ov ∥W * ov ∥ = -η 0 ∇ Wov L(W (t) ov ), W * ov ∥W * ov ∥ ≥ - η 0 1 + α t ∇ Wov L(W (t) ov ), W (t) ov ∥W (t) ov ∥ = 1 1 + α t W (t+1) ov -W (t) ov , W (t) ov ∥W (t) ov ∥ = 1 - α t 1 + α t W (t+1) ov -W (t) ov , W (t) ov ∥W (t) ov ∥ = 1 2∥W (t) ov ∥ ∥W (t+1) ov ∥ 2 -∥W (t+1) ov -W (t) ov ∥ 2 -∥W (t) ov ∥ 2 - α t 1 + α t W (t+1) ov -W (t) ov , W (t) ov ∥W (t) ov ∥ (a) = ∥W (t+1) ov ∥ 2 -∥W (t) ov ∥ 2 2∥W (t) ov ∥ - η 2 2∥W (t) ov ∥ + η 0 α t 1 + α t ∇ Wov L(θ (t) ) ∥∇ Wov L(θ (t) )∥ , W (t) ov ∥W (t) ov ∥ (b) ≥ ∥W (t+1) ov ∥ -∥W (t) ov ∥ - η 2 0 2∥W (t) ov ∥ - η 0 α t 1 + α t ,where (a) follows from that ∥W(t+1) ov -W (t)ov ∥ = η 0 , and (b) is due to the fact that x 2 -y 2 ≥ 2y(x-y) for any x, y ∈ R.Summing over t starting from 2, we haveW (t) ov -W (2) ov , W * ov ∥W * ov ∥ ≥ ∥W (t) ov ∥ -∥W (2) ov ∥ - t-1 t ′ =2 η 2 0 2∥W (t ′ ) ov ∥ - t-1 t ′ =2 η 0 α t ′ 1 + α t ′ .Furthermore, due to Lemma 1,t-1 t ′ =2 1 ∥W (t ′ ) ov ∥ ≤ t-1 t ′ =2 2∥W * ov ∥/η 0 t ≤ 2∥W * ov ∥ η 0 log t. Similarly, t-1 t ′ =2 α t ′ 1 + α t ′ ≤ t-1 t ′ =2 2∥W * ov ∥ log(2|V|) ∥W (t ′ ) ov ∥ ≤ 4∥W * ov ∥ 2 log(2|V|) η 0 log t Therefore, W (t) ov ∥W (t) ov ∥ , W * ov ∥W * ov ∥ ≥ 1 - ∥W (2) ov ∥ + 2η 0 ∥W * ov ∥ log t + 4∥W * ov ∥ 2 log(2|V|) log t ∥W (t) ov ∥ (a) ≥ 1 - 12∥W * ov ∥ 3 log(2|V|) log t tη 0 ,where (a) follows from Lemma 1 and ∥W satisfies thatEQUATION∥W (t) ov ∥ ≥ tη 0 /(2∥W * ov ∥) C Proof of(e i -e i ′ ) ⊤ W (t) ov x = 0, ∀i, i ′ ̸ = In(x).Proof. The proof follows directly from induction and the fact that(e i -e i ′ ) ⊤ W (t+1) ov x = (e i -e i ′ ) ⊤ W (t) ov x - η 0 ([T (t) 0 ] i -[T (t) 0 ] i ′ ) ∥∇ Wov L 0 (W (t) ov )∥ , ∀i, i ′ ̸ = In(x).Corollary 1 Under the settings in Proposition 1, let T ≥ 384∥W * ov ∥ 5 log(2|V|) log T /η 0 , and∆ = T η 0 /(4∥W * ov ∥ 2 ) (e In(x) -e i ) ⊤ W ov x ∈ (∆, 3∆), ∀i ̸ = In(x) (e i -e i ′ ) ⊤ W ov x = 0, ∀i, i ′ ̸ = In(x)Proof. The second equality follows directly from Lemma 3.To show the first equation, we analyze(e In(x) -e i ) ⊤ W (T ) ov x = (e In(x) -e i ) ⊤ W * ov ∥W (T ) ov ∥ ∥W * ov ∥ x + ∥W (T ) ov ∥(e In(x) -e i ) ⊤ W (T ) ov ∥W (T ) ov ∥ - W * ov ∥W * ov ∥ W (T ) ov x (a) = ∥W (T ) ov ∥ ∥W * ov ∥ -2 √ 2∥W (T ) ov ∥ 12∥W * ov ∥ 3 log(2|V|) log T T η 0 (b) ≥ T η 0 2∥W * ov ∥ 2 -24T η 0 ∥W * ov ∥ log(2|V|) log T ≥ T η 0 4∥W * ov ∥ 2, where (a) follows from Proposition 1, and (b) is due to Lemma 1. On the other hand, we also have(e In(x) -e i ) ⊤ W (T ) ov x ≤ ∥W (T ) ov ∥ ∥W * ov ∥ + 2 √ 2∥W (T ) ov ∥ 12∥W * ov ∥ 3 log(2|V|) log T T η 0 ≤ 3T η 0 4∥W * ov ∥ 2The proof is finished.Thus, for simplicity, we further assume that (e In(x) -e i )W (T ) ov x = ∆ for all x, because (e In(x)e i ) Ŵov x = Θ(∆) for large enough iteration. Next, we provide the general form of the projection of the gradient of Key-Query matrix W kq follows from a notation for the token weight.The token weight φ(n,t) ℓ of the token x (n) ℓ in the sentence X (n) = [x (n) 1 , . . . , x (n) L ] under θ = (W (T ) ov , W kq ) is calculated as φ (n,t) ℓ = exp (x (n) ℓ ) ⊤ W kq X (n) -1 L ℓ ′ =1 exp (x (n) ℓ ′ ) ⊤ W kq X (n) -1 (9) Lemma 4 (Projection of gradient of W kq ) If W ov satisfies that (e In(x) -e i ) ⊤ W ov x = ∆, ∀i ̸ = In(x) (e i -e i ′ ) ⊤ W ov x = 0, ∀i, i ′ ̸ = In(x)we have∇ W kq L(θ), W ′ kq = ∆ n π (n) ([T (n) θ ] In(X (n) ) -1) ℓ * ∈l(n) φ (n,θ) ℓ * x (n) ℓ * - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 + ∆ n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 .Proof. Recall that l(n) is the set of indices of the optimal tokens in the sample X (n) . Thus, for any ℓ * ∈ l(n), In(x n) ). In addition, we denote T θ (X (n) ) by T(n) ℓ * ) = In(X ((n) θ for simplicity. From Equation (6), we have∇ W kq L(θ), W ′ kq = n π (n) ℓ (T (n) θ -p (n) ) ⊤ W ov x (n) ℓ φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 (a) = n π (n) ℓ i [T (n) θ ] i (e i -e In(X (n) ) ) ⊤ W ov x (n) ℓ φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 = n π (n) ℓ∈l(n) i [T (n) θ ] i (e i -e In(X (n) ) ) ⊤ W ov x (n) ℓ φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 + n π (n) ℓ / ∈l(n) i [T (n) θ ] i (e i -e In(X (n) ) ) ⊤ W ov x (n) ℓ φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 = n π (n) ℓ∈l(n) i̸ =In(X (n) ) [T (n) θ ] i (-∆)φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 + n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 = ∆ n π (n) ([T (n) θ ] In(X (n) ) -1) ℓ * ∈l(n) φ (n,θ) ℓ * x (n) ℓ * - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 + ∆ n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 ,where (a) is due to the fact that i∈[|V|] [Tθ ] i = 1 for any θ, n.The proof consists of three main steps. First, we show that the optimal token weight has a 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 1 and Theorem 2.Recall that the updating rule forEQUATION)We first show that the optiaml token weight has a lower bound during the training.Lemma 5 (Lower bound of optimal token weight) Under the zero initialization and updating rule Equation (10), for any iteration t, and any sample X (n) , if l(n) is the set of indices of the optimal token in X (n) , the following inequality holds.φ (n,t) ℓ ≥ φ (n,0) ℓ ≥ 1/L max , ∀ℓ ∈ l(n) Proof.First, we introduce the notation that φ(n,t) + = ℓ * ∈l(n) φ (n,t) ℓ *as the summation of optimal token weights, and φ(n,t) - = 1 -φ (n,t) +as the summation of non-optimal token weights.At t = 0, due to zero initialization, we have φ (n,0) ℓ = 1/L (n) . Moreover, by Assumption 3, for any ℓ / ∈ l(n), we have φ(n,0) + ≥ q n (x ℓ )φ (n,0) ℓ .We perform induction the hypothesis: φ(n,t) + ≥ φ (n,t-1) + and φ (n,t) ℓ * ≥ φ (n,t) ℓfor all ℓ * ∈ l(n) and ℓ / ∈ l(n).Suppose the hypothesis holds for iteration t. Let x (n)ℓ * be the optimal token in the sequence X (n) .Fix a sample X (n ′ ) . we have (x(n ′ ) ℓ * ) ⊤ ∇ W kq L (t) (θ (t) )X (n ′ ) -1 = n π (n) ([T (n) θ ] In(X (n) ) -1)φ (n,t) +   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (x (n) ℓ * -x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ *   X (n) -1 , X (n ′ ) -1 + n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,t) ℓ ℓ ′ φ (n,t) ℓ ′ (x (n) ℓ -x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ * X (n) -1 , X (n ′ ) -1 Because ℓ ′ / ∈l(n) (x (n) ℓ ′ )x (n ′ )ℓ * = 0 due to Assumption 1, and (x(n) ℓ * ) ⊤ x (n ′ )ℓ * ≥ 0, we immediately have(x (n ′ ) ℓ * ) ⊤ ∇ W kq L (t) (θ (t) )X (n ′ ) -1 ≤ 0. Let x (n ′ )ℓ0 be any non-optiaml token in the sequence X (n ′ ) . Then, we have (x(n ′ ) ℓ0 ) ⊤ ∇ W kq L (t) (θ (t) )X (n ′ ) -1 = n π (n) ([T (n) θ ] In(X (n) ) -1)φ (n,t) +   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (x (n) ℓ * -x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ0   X (n) -1 , X (n ′ ) -1 + n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,t) ℓ ℓ ′ φ (n,t) ℓ ′ (x (n) ℓ -x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ0 X (n) -1 , X (n ′ ) -1 = n π (n) ([T (n) θ ] In(X (n) ) -1)φ (n,t) +   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (-x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ0   X (n) -1 , X (n ′ ) -1 + n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,t) ℓ   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (x (n) ℓ -x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ0   X (n) -1 , X (n ′ ) -1 ≥ n π (n)   (1 -[T (n) θ ] In(X (n) ) )φ (n,t) ℓ+ - ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,t) ℓ     ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (x (n) ℓ ′ ) ⊤ x (n ′ ) ℓ0   X (n) -1 , X (n ′ ) -1 ≥ 0,where the last inequality is due to Assumption 3, the induction hypothesis and i∈[|V|] [Tθ ] i = 1 Therefore, for any n, we haveφ (n,t+1) ℓ * = exp (x (n) ℓ * ) ⊤ W (t+1) kq X (n) -1 ℓ exp (x (n) ℓ ) ⊤ W (t+1) kq X (n) -1 = exp (x (n) ℓ * ) ⊤ W (t) kq X (n) -1 -η(x (n) ℓ * ) ⊤ ∇ W kq L (t) (θ (t) )X (n) -1 /∥∇ W kq L (t) (θ (t) )∥ ℓ exp (x (n) ℓ ) ⊤ W (t) kq X (n) -1 -η(x (n) ℓ ) ⊤ ∇ W kq L (t) (θ (t) )X (n) -1 /∥∇ W kq L (t) (θ (t) )∥ = exp (x (n) ℓ * ) ⊤ W (t) kq X (n) -1 ℓ exp (x (n) ℓ ) ⊤ W (t) kq X (n) -1 + η(x (n) ℓ * -x (n) ℓ ) ⊤ ∇ W kq L (t) (θ (t) )X (n) -1 /∥∇ W kq L (t) (θ (t) )∥ ≥ exp (x (n) ℓ * ) ⊤ W (t) kq X (n) -1 ℓ exp (x (n) ℓ ) ⊤ W (t) kq X (n) -1 = φ (n,t) ℓ * , which implies that φ (n,t+1) + ≥ φ (n,t) + .For the second argument in the hypothesis, we examine φ(n,t+1) ℓ * /φ (n,t+1) ℓ for any ℓ / ∈ l(n). We have φ (n,t+1) ℓ * φ (n,t+1) ℓ = exp (x (n) ℓ * -x (n) ℓ ) ⊤ W (t+1) kq X (n) -1 = exp (x (n) ℓ * -x (n) ℓ ) ⊤ W (t) kq X (n) -1 exp - η ∥∇ W kq L (t) (θ (t) )∥ (x (n) ℓ * -x (n) ℓ ) ⊤ ∇ W kq L (t) (θ (t) )X (n) -1 ≥ φ (n,t) ℓ * φ (n,t) ℓ ≥ 1.The proof is finished.The following lemma shows that the norm of the Key-Query Matrix increases linearly with the number of iterations.Lemma 6 Under the initialization W (0) kq and the updating rule Equation (10), for each iteration t, the following inequality holds.tη + ∥W (0) kq ∥ ≥ ∥W (t) kq ∥ ≥ tη 2L max ∥W * kq ∥ -∥W (0) kq ∥.Proof.We examine the gradient ∇ W kq L(θ) projected onto the optimal direction W * kq /∥W * kq ∥.∇ W kq L(θ (t) ), W * kq = n π (n) ∆ ℓ * ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ * ) -1)φ (n,t) ℓ * a (n, * ) ℓ * - ℓ ′ φ (n,t) ℓ ′ a (n, * ) ℓ ′ + n π (n) ∆ ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ a (n, * ) ℓ - ℓ ′ φ (n,t) ℓ ′ a (n, * ) ℓ ′ = n π (n) ∆([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) +   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (a (n, * ) ℓ * -a (n, * ) ℓ ′ )   + n π (n) ∆ ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ   ℓ ′ ∈l(n) φ (n,t) ℓ ′ (a (n, * ) ℓ -a (n, * ) ℓ ′ )   ≤ n π (n) ∆([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) + φ (n,t) - + n π (n) ∆ ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ (-φ (n,t) + ),where φ(n,t) + = ℓ * ∈l(n) φ (n,t) ℓ *is the summation of optimal token weights, and φ(n,t) - = 1 -φ (n,t) +is the summation of non-optimal token weights.On the other hand,∥∇ W kq L(θ (t) )∥ = ∇ W kq L(θ (t) ), ∇ W kq L(θ (t) ) ∥∇ W kq L(θ (t) )∥ = n π (n) ∆ ℓ * ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ * ) -1)φ (n,t) ℓ * x (n) ℓ * - ℓ ′ φ (n,t) ℓ ′ x (n) ℓ ′ ⊤ ∇ W kq L(θ (t) ) ∥∇ W kq L(θ (t) )∥ X (n) -1 + n π (n) ∆ ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ x (n) ℓ - ℓ ′ φ (n,t) ℓ ′ x (n) ℓ ′ ⊤ ∇ W kq L(θ (t) ) ∥∇ W kq L(θ (t) )∥ X (n) -1 ≤ 2 n π (n) ∆(1 -[T (n) θ ] In(X (n) ) )φ (n,t) + φ (n,t) - + 2 n π (n) ∆ ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,t) ℓThus, we have∇ W kq L(θ) ∇ W kq L(θ) , W * kq ∥W * kq ∥ ≤ - min n φ (n,t) + 2∥W * kq ∥ ≤ - 1 2L max ∥W * kq ∥By the updating rule Equation (10), we have∥W (t) kq ∥ = W (0) kq - t-1 t ′ =0 η ∇ W kq L(θ (t ′ ) ) ∇ W kq L(θ (t ′ ) ) ≥ W (0) kq , W * kq ∥W * kq ∥ - t ′ ≤t-1 η ∇ W kq L(θ (t ′ ) ) ∇ W kq L(θ (t ′ ) ) , W * kq ∥W * kq ∥ ≥ t ′ <t η 2L max ∥W * kq ∥ -∥W (0) kq ∥ = tη 2L max ∥W * kq ∥ -∥W (0) kq ∥.In addition, by the triangle inequality,∥W (t) kq ∥ = W (0) kq - t-1 t ′ =0 η ∇ W kq L(θ (t ′ ) ) ∇ W kq L(θ (t ′ ) ) ≤ tη + ∥W (0) kq ∥.The proof is completed.We next show that the gradient ∇ W kq L(θ (t) ) is close to the optimal direction W * kq .Lemma 7 (Gradient aligns with the optimal direction) Let t 0 = ⌈ 8Lmax∥W * kq ∥ 2 η ⌉. Then, for any t ≥ t 0 , we have∇ W kq L(θ (t) ), W (t) kq ≥ (1 + α t ) ∇ W kq L(θ (t) ), W * kq ∥W (t) kq ∥ ∥W * kq ∥ where α t = 4N L 2 max ∥W * kq ∥ 2 ∥W (t) kq ∥ 1 + log 2L max ∥W (t) kq ∥Proof. During the proof, we denote       a (n,t) ℓ = (x (n) ℓ ) ⊤ W (t) kq X (n) -1 a (n, * ) ℓ = (x (n) ℓ ) ⊤ W * kq X (n) -1 ∥W (t) kq ∥ ∥W * kq ∥ β 0 = 2L 2 max ∥W * kq ∥ 2 ∥W (t) kq ∥ (1 + log(2L max ∥W (t) kq ∥)).We point out a few facts that will be frequently used in the proof.If a (n,t) ℓ ≤ a (n,t) ℓ ′ -C 0 , then we have φ (n,t) ℓ = φ (n,t) ℓ ′ exp a (n,t) ℓ -a (n,t) ℓ ′ ≤ exp(-C 0 ) (11)The same result holds if a(n,t) ℓ ′is replaced any convex combination of a set of a(n,t) ℓ ′ 's.We start the proof by noting that W * kq is the minimum unique solution to the problemEQUATIONTherefore, ifW (t) kq ∥W * kq ∥ ∥W (t) kq ∥ = W * kq , the results is trivial since ∇ W kq L(θ (t) ), W * kq ≤ 0.In the following, we focus on the case when W (t) kq∥W * kq ∥ ∥W (t) kq ∥ ̸ = W * kq .Then, there must be at least a sentence X (n) , such that W (t) kq∥W * kq ∥ ∥W (t) kq ∥violates the contraint on X (n) . In other words, we must havea (n,t) ℓ * -a (n,t) ℓ = (x (n) ℓ * -x (n) ℓ )W (t) kq X (n) -1 ≤ ∥W (t) kq ∥ ∥W * kq ∥ .This implies that for those n, we must have φ(n,t) ℓ ≥ exp(- ∥W (t)kq ∥ ∥W * kq ∥ ) Thus, we consider two types of samples in the folloiwng.Type 1. Let us consider X (n) such that φ (n,t) - ≥ exp(-(1 + β 0 /2)∥W (t) kq ∥/∥W * kq ∥).Recall that the inner product between the gradient ∇ W kq L(θ (t) ) and any other Key-Query matrix θ ′ = W ′ kq has the following form (Lemma 4). n) ) be the loss on sample X (n) .∇ W kq L(θ (t) ), W ′ kq = ∆ n π (n) ([T (n) θ ] In(X (n) ) -1) ℓ * ∈l(n) φ (n,θ) ℓ * x (n) ℓ * - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 + ∆ n π (n) ℓ / ∈l(n) [T (n) θ ] In(x (n) ℓ ) φ (n,θ) ℓ x (n) ℓ - ℓ ′ φ (n,θ) ℓ ′ x (n) ℓ ′ ⊤ W ′ kq X (n) -1 Let L n (θ) = -log e ⊤ In(X (n) ) T θ (X (To proceed, we examine the gradient on each sample X (n) with φ(n,t) - ≥ exp(-(1 + β 0 /2)∥W (t) kq ∥/∥W * kq ∥), which can be divided into two parts. ∇ W kq L n (θ (t) ), W kq = ∆(A (n,t) + B (n,t) ), where              A (n,t) = ℓ * ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ * ) -1)φ (n,t) ℓ * a (n,t) ℓ * - ℓ ′ φ (n,t) ℓ ′ a (n,t) ℓ ′ , B (n,t) = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ a (n,t) ℓ - ℓ ′ φ (n,t) ℓ ′ a (n,t) ℓ ′ .We further letEQUATIONB (n, * ) = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ a (n, * ) ℓ - ℓ ′ φ (n,t) ℓ ′ a (n, * ) ℓ ′.Thus, we aim to find the relationship A (n,t) + B (n,t) between A (n, * ) + B (n, * ) .We first provide the upper bounds for A (n, * ) and B (n, * ) .A (n, * ) = ℓ * ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ * ) -1)φ (n,t) ℓ * a (n, * ) ℓ * - ℓ ′ φ (n,t) ℓ ′ a (n, * ) ℓ ′ = ℓ * ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ * ) -1)φ (n,t) ℓ *   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (a (n, * ) ℓ * -a (n, * ) ℓ ′ )   (a) ≤ ([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) + φ (n,t) - ∥W (t) kq ∥ ∥W * kq ∥, where (a) is due to the fact that (x(n) ℓ * -x (n) ℓ )W * kq X (n) -1≥ 1, and a(n, * ) ℓ = (x (n) ℓ ) ⊤ W * kq X (n) -1 ∥W (t)kq ∥/∥W * kq ∥. On the other handA (n,t) = ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ a (n,t) ℓ - ℓ ′ φ (n,t) ℓ ′ a (n,t) ℓ ′ = ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ (a (n,t) ℓ -a (n,t) ℓ ′ )   = max T      ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ    ℓ ′ / ∈l(n) diff<T φ (n,t) ℓ ′ (a (n,t) ℓ -a (n,t) ℓ ′ ) diff    + ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ    ℓ ′ / ∈l(n) diff>T φ (n,t) ℓ ′ (a (n,t) ℓ -a (n,t) ℓ ′ ) diff         . (a) ≥ max T    ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ   ℓ ′ / ∈l(n) φ (n,t) ℓ ′ T   + ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ   2 ℓ ′ / ∈l(n) exp(-T)∥W (t) kq ∥      ≥ max T    ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ φ (n,t) - T + 2L max exp(-T)∥W (t) kq ∥    (b) ≥ ℓ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ ) -1)φ (n,t) ℓ φ (n,t) - 1 + log 2L max ∥W (t) kq ∥ φ (n,t) - ,where (a) is due to Equation ( 11), and (b) is obtained by choosing T = log2Lmax∥W (t) kq ∥ φ (n,t) - . Recall that φ (n,t) - ≥ exp -(1 + β 0 /2)∥W (t) kq ∥/∥W * kq ∥ and β 0 ≥ 2∥W * kq ∥(1+log(2Lmax∥W (t) kq ∥)) ∥W (t) kq ∥. Thus, we further haveA (n,t) ≥ ([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) + φ (n,t) - 1 + log 2L max ∥W (t) kq ∥ φ (n,t) - ≥ ([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) + φ (n,t) - 1 + log(2L max ∥W (t) kq ∥) + (1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ ≥ ([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) + φ (n,t) - β 0 ∥W (t) kq ∥ 2∥W * kq ∥ + (1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ = (1 + β 0 )([T (n) θ (t) ] In(X (n) ) -1)φ (n,t) + φ (n,t) - ∥W (t) kq ∥ ∥W * kq ∥ ≥ (1 + β 0 )A (n, * )Next, we analyze B (n,t) , and further divide B (n,θ) into B (n,θ) + and B (n,θ) - :                B (n,θ) + = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ   φ (n,t) + a (n,θ) ℓ - ℓ ′ ∈l(n) φ (n,t) ℓ ′ a (n,θ) ℓ ′   B (n,θ) - = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ   φ (n,t) - a (n,θ) ℓ - ℓ ′ / ∈l(n) φ (n,t) ℓ ′ a (n,θ) ℓ ′  Due to Proposition 4, we have B(n, * ) -= 0, and thusB (n, * ) = B (n, * ) + ≤ ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ (-φ (n,t) + ) ∥W (t) kq ∥ ∥W * kq ∥ ≤ 0.We then analyze:B (n,t) + -(1 + β 0 )B (n, * ) + = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ φ (n,t) + a (n,t) ℓ -φ (n,t) + a (n,t) ℓ * -(1 + β 0 ) φ (n,t) + a (n, * ) ℓ -φ (n,t) + a (n, * ) ℓ * = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ φ (n,t) +     a (n,t) ℓ -a (n,t) ℓ * b t,ℓ -(1 + β 0 ) a (n, * ) ℓ -a (n, * ) ℓ * b * ,ℓ     ≥ ℓ / ∈l(n) b t,ℓ <b * ,ℓ [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ φ (n,t) + (b t,ℓ -b * ,ℓ ) (a) ≥ ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) + exp -(1 + β 0 ) ∥W (t) kq ∥ ∥W * kq ∥ -2∥W (t) kq ∥ = -2 ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) + exp -(1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ ∥W (t) kq ∥ exp - β 0 2 ∥W (t) kq ∥ ∥W * kq ∥ (b) ≥ -2L max (1 -[T (n) θ (t) ] In(X (n) ) )φ (n,t) + φ (n,t) - ∥W (t) kq ∥ ∥W * kq ∥ exp - β 0 2 ∥W (t) kq ∥ ∥W * kq ∥ ∥W * kq ∥ ≥ 2L max exp - β 0 2 ∥W (t) kq ∥ ∥W * kq ∥ ∥W * kq ∥A (n, * ) (c) ≥ ∥W * kq ∥ ∥W (t) kq ∥ A (n, * ) ≥ β 0 A (n, * ) ,where (a) follows from that b * ,ℓ ≤ -(1 + β 0 )∥W /2 ≥ ∥W * kq ∥ ∥W (t) kq ∥ log(2L max ∥W (t) kq ∥) For the term B (n,t) - , we have B (n,t) - = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ   φ (n,t) - a (n,t) ℓ - ℓ ′ / ∈l(n) φ (n,t) ℓ ′ a (n,t) ℓ ′   = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ φ (n,t) -   a (n,t) ℓ - ℓ ′ / ∈l(n) φ (n,t) ℓ ′ φ (n,t) - a (n,t) ℓ ′   = max T>0            ℓ / ∈l(n) diff>-T [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ φ (n,t) -       a (n,t) ℓ - ℓ ′ / ∈l(n) φ (n,t) ℓ ′ φ (n,t) - a (n,t) ℓ ′ diff       + ℓ / ∈l(n) diff<-T [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ φ (n,t) -       a (n,t) ℓ - ℓ ′ / ∈l(n) φ (n,t) ℓ ′ φ (n,t) - a (n,t) ℓ ′ diff                  (a) ≥ max T>0    ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) - -φ (n,t) ℓ T -2∥W (t) kq ∥ exp(-T)    (b) ≥ -L max (1 -[T (n) θ (t) ] In(x (n) ℓ * ) )φ (n,t) - 1 + log(2∥W (t) kq ∥) ≥ L max ∥W * kq ∥ φ (n,t) + ∥W (t) kq ∥ 1 + log(2∥W (t) kq ∥) A (n, * ) (c) ≥ L 2 max ∥W * kq ∥ ∥W (t) kq ∥ 1 + log(2∥W (t) kq ∥) A (n, * ) (d) ≥ β 0 A (n, * ) ,where (a) follows from Equation ( 11), (b) is optained by choosing T = log(2∥W (t) kq ∥), (c) follows from Lemma 5, and (d) is due to the fact thatβ 0 ≥ L 2 max ∥W * kq ∥ ∥W (t) kq ∥ (1 + log(2∥W (t) kq ∥)).So far, we have shown that for if φ(n,t) - ≥ exp(-(1 + β 0 /2)∥W (t) kq ∥/∥W * kq ∥), then A (n,t) + B (n,t) = A (n,t) + B (n,t) + + B (n,t) - ≥ (1 + β 0 )A (n, * ) + (1 + β 0 )B (n, * ) + β 0 A (n, * ) + β 0 A (n, * ) = (1 + 3β 0 )A (n, * ) + (1 + β 0 )B (n, * ) ≥ (1 + 3β 0 )(A (n, * ) + B (n, * ) ).Type 2. Now consider sentence X (n) such that φ(n,t) - < exp(-(1 + β 0 /2)∥W (t) kq ∥/∥W * kq ∥).Let n 0 be the type 1 sample such that φEQUATIONkq ∥/∥W * kq ∥) Then, we aim to show thatA (n,t) + B (n,t) ≥ β 0 A (n0, * ) Note that A (n,t) ≥ ℓ * ∈l(n) ([T (n) θ (t) ] In(x (n) ℓ * ) -1)φ (n,t) ℓ * φ (n,t) - 1 + log L max ∥W (t) kq ∥ φ (n,t) - ≥ ([T (n) θ (t) ] In(x (n) ℓ * ) -1) exp -(1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ 1 + (1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ + log(L max ∥W (t) kq ∥) ≥ (1 + β 0 )([T (n) θ (t) ] In(x (n) ℓ * ) -1) exp -(1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ ∥W (t) kq ∥ ∥W * kq ∥ and B (n,t) = ℓ / ∈l(n) [T (n) θ (t) ] In(x (n) ℓ ) φ (n,t) ℓ a (n,t) ℓ - ℓ ′ φ (n,t) ℓ ′ a (n,t) ℓ ′ ≥ -2(1 -[T (n) θ (t) ] In(x (n) ℓ * ) ) exp -(1 + β 0 /2) ∥W (t) kq ∥ ∥W * kq ∥ ∥W (t) kq ∥ Since A (n0, * ) ≤ ([T (n0) θ (t) ] In(x (n 0 ) ℓ * ) -1)φ (n0,t) + φ (n0,t) - ∥W (t) kq ∥ ∥W * kq ∥ ≤ ([T (n0) θ (t) ] In(x (n 0 ) ℓ * ) -1)φ (n0,t) + exp - ∥W (t) kq ∥ ∥W * kq ∥ ∥W (t) kq ∥ ∥W * kq ∥We further note that [T (n0)θ (t) ] In(x (n 0 ) ℓ * ) < [T (n) θ (t) ] In(x (n) ℓ * ) due to φ (n0,t) + < φ (n,t) +. Thus,A (n,t) + B (n,t) ≥ exp -β 0 /2 ∥W (t) kq ∥ ∥W * kq ∥ (1 + β 0 + 2∥W kq * ∥) A (n0, * ) φ (n0,t) + (a) ≥ β 0 A (n0, * ) ,where (a) is due to thatEQUATIONβ 0 ≥ 2∥W * kq ∥ ∥W (t) kq ∥ log 1 + β 0 + 2∥W * kq ∥ β 0 .In summary, we have thatn π (n) (A (n,t) + B (n,t) ) = n is type 2 π (n) (A (n,t) + B (n,t) ) + n is type 1 π (n) (A (n,t) + B (n,t) ) ≥ max n0 is type 1 β 0 A (n0, * ) + n is type 1 π (n) ((1 + 3β 0 )A (n, * ) + (1 + β 1 )B (n, * ) ) ≥ n is type 1 N π (n) β 0 (A (n, * ) + B (n, * ) ) + (1 + 3β 0 ) n is type 1 π (n) (A (n, * ) + B (n, * ) ) ≥ (1 + (N + 3)β 0 ) n is type 1 π (n) (A (n, * ) + B (n, * ) ) ≥ (1 + α t ) n is type 2 π (n) (A (n, * ) + B (n, * ) ) + (1 + α t ) n is type 1 π (n) (A (n0, * ) + B (n0, * ) ) = (1 + α t ) n π (n) (A (n, * ) + B (n, * ) ),where α t ≥ (N + 3)β 0 . The proof is finished. Now, we are ready to prove Theorem 1.Proof of Theorem 1.Recall that α t = 4N L 2 max ∥W * kq ∥ 2 ∥W (t) kq ∥ 1 + log 2L max ∥W (t)kq ∥ . By Lemma 7, we haveW (t+1) kq -W (t) kq , W * kq ∥W * kq ∥ = -η ∇ W kq L(θ (t) ), W * kq ∥W * kq ∥ ≥ - η 1 + α t ∇ W kq L(θ (t) ), W (t) kq ∥W (t) kq ∥ = 1 1 + α t W (t+1) kq -W (t) kq , W (t) kq ∥W (t) kq ∥ = 1 2∥W (t) kq ∥ ∥W (t+1) kq ∥ 2 -∥W (t+1) kq -W (t) kq ∥ 2 -∥W (t) kq ∥ 2 - α t 1 + α t W (t+1) kq -W (t) kq , W (t) kq ∥W (t) kq ∥ = ∥W (t) kq ∥ 2 -∥W (t) kq ∥ 2 2∥W (t) kq ∥ - η 2 2∥W (t) kq ∥ + ηα t 1 + α t ∇ W kq L(θ (t) ) ∥∇ W kq L(θ (t) )∥ , W (t) kq ∥W (t) kq ∥ ≥ ∥W (t+1) kq ∥ -∥W (t) kq ∥ - η 2 2∥W (t) kq ∥ - ηα t 1 + α t Let t 0 = ⌈ 8Lmax∥W * kq ∥ 2 η⌉ be defined in Lemma 7. Summing over t from t 0 , we haveW (t) kq -W (t0) kq , W * kq ∥W * kq ∥ ≥ ∥W (t) kq ∥ -∥W (t0) kq ∥ - t-1 t ′ =t0 η 2 2∥W (t ′ ) kq ∥ - t-1 t ′ =t0 ηα t ′ 1 + α t ′ By Lemma 6, we have t-1 t ′ =t0 1 ∥W (t ′ ) kq ∥ ≤ t-1 t ′ =t0 2L max ∥W * kq ∥/η t ′ ≤ 2L max ∥W * kq ∥ η log t. Furthermore, t-1 t ′ =t0 α t ′ 1 + α t ′ ≤ t-1 t ′ =t0 α t ′ = t-1 t ′ =t0 4N L 2 max ∥W * kq ∥ 2 ∥W (t ′ ) kq ∥ 1 + log 2L max ∥W (t ′ ) kq ∥ = t-1 t ′ =t0 4N L 2 max ∥W * kq ∥ 2 ∥W (t ′ ) kq ∥ (1 + log (2L max )) + t-1 t ′ =t0 4N L 2 max ∥W * kq ∥ 2 ∥W (t ′ ) kq ∥ log ∥W (t ′ ) kq ∥ ≤ t-1 t ′ =t0 8N L 3 max ∥W * kq ∥ 3 /η t ′ log(2eL max ) + t-1 t ′ =t0 8N L 3 max ∥W * kq ∥ 3 /η t ′ log t ′ 2L max ∥W * kq ∥/η = t-1 t ′ =t0 8N L 3 max ∥W * kq ∥ 3 /η t ′ log(eη/∥W * kq ∥) + t-1 t ′ =t0 8N L 3 max ∥W * kq ∥ 3 /η t ′ log(t ′ ) (a) ≤ 8N L 3 max ∥W * kq ∥ 3 η log 2 t,where (a) follows from η ≤ ∥W * kq ∥/e. Therefore, we haveW (t) kq -W (t0) kq , W * kq ∥W * kq ∥ ≥ ∥W (t) kq ∥ -∥W (t0) kq ∥ - t-1 t ′ =t0 η 2 2∥W (t ′ ) kq ∥ - t-1 t ′ =t0 ηα t ′ 1 + α t ′ ≥ ∥W (t) kq ∥ -∥W (t0) kq ∥ -L max ∥W * kq ∥ log t -8N L 3 max ∥W * kq ∥ 3 log 2 t.Finally, by Lemma 6, and t 0 ≤ 1 + 8L max ∥W * kq ∥ 2 /η, we have ∥W (t0) kq ∥ ≤ 9L max ∥W * kq ∥ 2 , andW (t) kq ∥W (t) kq ∥ , W * kq ∥W * kq ∥ ≥ 1 - 2∥W (t0) kq ∥ + L max ∥W * kq ∥ log t + 8N L 3 max ∥W * kq ∥ 3 log 2 t ∥W (t) kq ∥ ≥ 1 - 54N L 4 max ∥W * kq ∥ 4 log 2 t tη .The proof is finished.Proof of Theorem 2.Recall that T ≥ 384∥W * ov ∥ 5 log(2|V|) log T /η 0 , and ∆ = T η 0 /(4∥W * ov ∥ 2 ) due to Corollary 1 and(e In(x) -e v ) ⊤ W (T ) ov x = T η 0 4∥W * ov ∥ 2Note that for all ℓ * ∈ l(n) and ℓ / ∈ l(n)(x (n) ℓ -x (n) ℓ * ) ⊤ W (t) kq X (n) -1 = ∥W (t) kq ∥ ∥W * kq ∥ (x (n) ℓ -x (n) ℓ * ) ⊤ W * kq X (n) -1 + (x (n) ℓ -x (n) ℓ * ) ⊤ W (t) kq - ∥W (t) kq ∥ ∥W * kq ∥ W kq * X (n) -1 ≤ - ∥W (t) kq ∥ ∥W * kq ∥ + 2∥W (t) kq ∥ W (t) kq ∥W (t) kq ∥ - W * kq ∥W * kq ∥ (13) ≤ - tη 2L max ∥W * kq ∥ 2 + √ 2tη L max ∥W * kq ∥ 54N L 4 max ∥W * kq ∥ 4 log 2 t tη (a) ≤ - tη 4L max ∥W * kq ∥ 2 , where (a) follows from t ≥ 1696N L 4 max ∥W * kq ∥ 6 log 2 t/η. Therefore, ℓ * ∈l(n) φ (n,t) ℓ * = |l(n)| |l(n)| + ℓ ′ / ∈l(n) exp (x (n) ℓ -x (n) ℓ * ) ⊤ W (t) kq X (n) -1 ≥ |l(n)| |l(n)| + (L (n) -|l(n)|) exp - tη 4Lmax∥W * kq ∥ 2 ≥ 1 1 + L max exp - tη 4Lmax∥W * kq ∥ 2 ≥ 1 1 + ϵ ,where the last inequality follows from t ≥4Lmax∥W * kq ∥ η log Lmax ϵ .Hence, the loss on the sentence X (n) satisfies that-log e ⊤ In(X (n) ) T θ (t) (X (n) ) = -log exp e ⊤ In(X (n) ) W (T ) ov ℓ x (n) ℓ φ (n,t) ℓ v≤|V| exp e ⊤ v W (T ) ov ℓ x (n) ℓ φ (n,t) ℓ = -log 1 1 + v̸ =In(X (n) ) exp (e v -e In(X (n) ) ) ⊤ W (T ) ov ℓ x (n) ℓ φ (n,t) ℓ = log   1 + v̸ =In(X (n) ) exp   (e v -e In(X (n) ) ) ⊤ W (T ) ov ℓ / ∈l(n) x (n) ℓ φ (n,t) ℓ     = log   1 + v̸ =In(X (n) ) exp   -∆ ℓ * ∈l(n) φ (n,t) ℓ * + ∆ ℓ / ∈l(n) φ (n,t) ℓ     ≤ |V| exp -∆(2φ (n,t) + -1) ≤ |V| exp   -∆ + 2∆L max L max + exp tη 4Lmax∥W * kq ∥ 2   .Thus, the average loss has upper bound, which is|V| exp -∆ + 2∆L max L max + exp(C 1 t) , for C 1 = η/(4L max ∥W * kq ∥ 2 ), and ∆ = C 0 T for C 0 = η 0 /(4∥W * ov ∥ 2 ).D Proof of Proposition 2 and Theorem 3 (n) ℓ * -x (n) ℓ ) ⊤ W x ≥ 1, ∀ℓ / ∈ l(n), ∀n.In addition, for each query x q , if there are k optimal tokens under a x q -partial order, m nonoptimal tokens under x q -partial order, then, for any optimal token x * , non-optimal token x, and non-comparable token x 0 , we havex ⊤ * W * kq x q = m k + m , x ⊤ W * kq x q = - k k + m , x ⊤ 0 W * kq x q = 0. A direct result is that (x (n) ℓ -x (n) ℓ ′ ) ⊤ W * kq X (n) -1 = 0, ∀ℓ, ℓ ′ / ∈ l(n).Proof. Let U ∈ R d be the rotation matrix such that U x = e I(x) . Because U preserves Frobenius norm, the optimization problem in Equation (3) can be written asEQUATIONNotably W * = U W * kq U ⊤ . Note that {e I(X (n)-1 ) } n forms a standard basis. It suffices to minimize the norm of each column of W subject to the constraint (e I(x (n) ℓ * ) -e I(x (n) ℓ ) ) ⊤ W e I(X (n) -1 ) ≥ 1. Let us consider any column c of W , denoted as [w 1 , . . . , w d ] ⊤ . Without loss of generality, we assume that, for all X (n) with I(X (n)-1 ) = c, the set of indices of the optimal tokens of those samples are {1, . . . , k}, and the set of indices of the non-optimal tokens are {k + 1, . . . , k + m}. Then, the optimization problem Equation ( 14) reduces to the following problem min w 2 1 + . . . + w 2 d , s.t. w i -w j ≥ 1, ∀i ≤ k, j ∈ A i ⊂ {k + 1, . . . , k + m}, (15) where A i is the set of indices of the non-optimal tokens in some samples whose optimal token has index i.In other words, each column of the solution of Equation ( 14) is the solution of Equation ( 15).Note that Equation ( 15) is a convex problem with linear constraints. The Lagrangian function isL(λ) = k+m i=1 w 2 i + 2 m i=1 j∈Ai λ ij (1 -w i + w j ),where we directly set w j = 0 for all j ∈ {k + m + 1, . . . , d}. That is, non-comparable tokens have value 0.By KKT-condition, we have           w i = j∈Ai λ ij , ∀i ≤ k w j = - k i=1 λ ij 1{j ∈ A i }, ∀k + 1 ≤ j ≤ k + m Thus, min w 2 1 + . . . + w 2 d = max λ      - k i=1   j∈Ai λ ij   2 - k+m j=k+1 k i=1 λ ij 1{j ∈ A i } 2 + 2 m i=1 j∈Ai λ ij      Let L * (λ) = - k i=1   j∈Ai λ ij   2 - k+m j=k+1 k i=1 λ ij 1{j ∈ A i } 2 + 2 m i=1 j∈Ai λ ij ,where λ ≥ 0. The maximum of L * is achieved when ∇ λ L * = 0. This implies thatj∈Ai 0 λ i0j + k i=1 λ ij 1{j 0 ∈ A i } = 1, ∀1 ≤ i 0 ≤ k < j 0 ≤ k + m.Hence, we have w i -w j = 1 for all 1 ≤ i ≤ k < j ≤ k + m, which means the optimum of the original problem is achieved on the boundary. Therefore, we reduce the original problem to minx k(x + 1) 2 + mx 2 ,where x = w k+1 = . . . = w k+m . Hence, the optimal solution is w 1 = . . . = w k = m/(m + k), andw k+1 = . . . = w k+m = -k/(k + m).Therefore, the solution of Equation ( 15) satisfies that the "optimal values" are the same and the "non-optimal values" are the same as well. This fact proves that(x (n) ℓ -x (n) ℓ ′ )W * kq X (n) -1 = 0, ∀ℓ, ℓ ′ / ∈ l(n).And moreover, if there are k optimal tokens under a x q -partial order, m non-optimal tokens under x q -partial order, then, for any optimal token x * and non-optimal token x, we havex ⊤ * W * kq x q = m k + m , x ⊤ W * kq x q = - k k + m .Proof. The proof follows similar logic to Theorem 2. By Equation ( 13), we have for any x, x ′ ∈ V(x -x ′ ) ⊤ W (t) kq x q ≥ ∥W (t) kq ∥ ∥W * kq ∥ (x -x ′ ) ⊤ W * kq x q - √ 2tη L max ∥W * kq ∥ 54N L 4 max ∥W * kq ∥ 4 log 2 t tη (a) ≥ tη 2L max ∥W * kq ∥ 2 (x -x ′ ) ⊤ W * kq x q -108tηN L 2 max ∥W * kq ∥ 2 log 2 t,where (a) follows from Lemma 6. The first part of Theorem 3 follows from Theorem 3.For the second part of Theorem 3, Let X = [x 1 , . . . , x L ] such that for ℓ 0 ∈ l 0 ⊂ {1, . . . , L}, x ℓ0 = x is a non-comparable token, and other tokens are non-optimal under the x L -partial order.Let φ ℓ ∝ exp(x ℓ W (t) kq x L ) for sufficiently large t = Ω(log(1/ϵ)) such that ℓ0∈l0 φ ℓ0 ≥ 1 -ϵ. Then, we havee ⊤ In(x) T θ (t) (X) = exp e ⊤ In(x) W (T ) ov ℓ x ℓ φ ℓ v≤|V| exp e ⊤ v W (T ) ov ℓ x ℓ φ ℓ = 1 1 + v̸ =In(x) exp (e v -e In(x) ) ⊤ W (T ) ov ℓ x ℓ φ ℓ = 1 1 + v̸ =In(x) exp -∆ ℓ0∈l0 φ ℓ0 + ∆ ℓ / ∈l0 φ ℓ ≥ 1 1 + |V| exp (-∆(1 -2ϵ)) ≥ 1 -ϵ 0 ,where the last inequality follows from T = O(log(1/ϵ 0 )). Therefore, the trained transformer will predict n(x), the next token of the non-comparable token.This is also related to the maximal element in a partially ordered set.This is also related to the minimal element in a partially ordered set.A more general way is to use various standard techniques developed in linguistic analysis(Lehecka, 2015). 


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

Reviews: review: paper_summary: paper_summary: This paper conducts a non-asymptotic analysis of training dynamics for a one-layer transformer, focusing on next-token prediction tasks.
It provides a mathematical framework  based on partial order to formally characterize a realizable training dataset for next-token prediction.
It also introduces a two-stage training algorithm that ensures fast convergence, with both layers approaching their max-margin solutions sub-linearly.
strength_weakness: strength_weakness: The approach described in the paper may not align with practical Transformer training.
Please see the questions below for more details.
questions: 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?
Q2: 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?
Q3: 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.
Q4: How commonly is the normalized gradient used in actual Transformer training? In the training of Transfomer, AdamW is typically used.
Q5: 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: limitations: The assumption of the existence of collocations seems far from real traninng dataset.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. Some of the mathematical concepts, particularly the query-dependent partial orders, are complex and could be better explained.
2. While the theoretical results are strong, additional empirical studies, particularly on real-world datasets, could further strengthen the claims.
questions: questions: What are the potential implications of the findings on transformer training efficiency and computational costs in practical settings?
limitations: limitations: NaN
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

The 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…
questions: questions: How would you change the manuscript to make it more easy to understand your definitions and motivation behind the mathematics? 

How would you change the manuscript to better assess/discuss the limitations of your work?
limitations: limitations: See above, no negative societal impact.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. How is the loss function between lines 227 and 288 obtained? Why is its form different from Eq. 1?
2. In Section 6, the authors theoretically prove the generalization ability of the trained Transformer. Can the generalization ability be proved in the experimental part?
3. Why do the authors design a two-stage training algorithm? Is it valuable in practical applications compared to single-stage training?
questions: questions: Please see the weakness
limitations: 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.

Comments: title: A gentle reminder before the discussion period ends
Thank you
Thank you
Thank you
A gentle reminder
A gentle reminder
comment: Thank you for clarification. I raised my score.
However, I am not satisfied with the composition of n. If we use the frequency of two words, we should consider the probability distribution based on the frequency rather than the maximum one. It seems important as a language model that a word expresses multiple meanings by co-occurring with other words.
Thank you for your insightful questions. Below we answer the reviewer's questions.

**Can every word find a next token?** We do not need every word to find a next token. If a token $x$ does not have next token, then $x$ falls into the category of non-optimal tokens or non-comparable tokens (see lines 182-184 of the paper). Then the trained transformer will not attend to such a token. Thus, $x$ won't play a role to predict next tokens in sentences. However, it is still possible that other tokens can predict $x$ as a next token.

**Deterministic function:** Yes, n is assumed to be deterministic. One reason is that in real world settings, n can be prepared to be a deterministic mapping (see below about construction). In recent line of theoretical research on next-token prediction, deterministic function is commonly adopted, such as in (Li et al., 2024, Tarzanagh et al., 2023.), as an informative model to enable tractable analysis for transformers. Such a model is also analogous to the deterministic structure of 'separable data' widely taken to develop deep learning theory in the literature such as in (Soudry et al, 2018, Taheri et al., 2023). 

**Construct n in real data:** We can construct the collocation n by employing various standard techniques developed in linguistic analysis (Lehecka, 2015). To elaborate a simple version (those practical techniques can include more sophisticated tricks), by processing all sentences in the corpus as detailed in (Lehecka, 2015), the frequency of each 'ordered' word pair $(x,y)$ that appears in the same sentence can be calculated. Then, the most frequent $y$ paired with each $x$ is chosen as its collocated word, i.e., let $n(x)=y$. Such method includes our original construction from length-2 sentences as a special case. There are also several off-the-shelf NLP tools to construct collocations such as Natural Language Toolkit (NLTK).

We hope that our responses have resolved your concerns. If so, we kindly ask the reviewer to consider increasing the score accordingly. If possible, we also kindly ask the reviewer to evaluate the paper based on the theoretical contributions and the novel mathematical techniques that we develop to analyze the next-token prediction problems, which can be applied to studying more sophisticated models in the future. Certainly, we are more than happy to answer your further questions.


Lehecka. Collocation and colligation. In Handbook of pragmatics online. Benjamins 2015.

Li et al. "Mechanics of next token prediction with self-attention." AISTATS 2024.  

Soudry et al. "The implicit bias of gradient descent on separable data." JMLR 2018.

Taheri et al. "On generalization of decentralized learning with separable data." AISTATS 2023.

Tarzanagh et al. "Transformers as support vector machines." arXiv:2308.16898, 2023.
Let me confirm about "collocations". 
Please tell me how to construct n in real data. 
It seems like it needs to be defined for every word, not some words and also, n seems to be a deterministic function.
Considering the above, it seems that there exists an n such that every word can definitively find one next token.
Dear Reviewer U8Jg,

As the author-reviewer discussion period will end soon, we would like to check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

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

Best Regards, 
Authors
Many thanks for your increase of the score and strongly support of our paper!
I have carefully read the rebuttal and, especially following the effort to suggest approaches to generalize the result, I will raise my score from 5 to 6.
Many thanks for your continual and strong support of our paper!
Many thanks for your continual and strong support of our paper!
Thanks to the authors for replying to my questions. I have read the relevant comments and will maintain this score.
I have thoroughly reviewed all the comments and the author’s responses, and I will remain positive about this submission.
Dear Reviewer ZtZ3

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

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

Best Regards, Authors
Dear Reviewer U8Jg, 

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

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

Best Regards, 
Authors