Title: Non-asymptotic Convergence of Training Transformers for Next-token Prediction
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformer

Convergence Rate

Next-token Prediction

Self-attention

Implicit Bias
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

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.

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

Jing Yang

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.

2 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; V on 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 ˜ O (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.

3 Problem Setup

Notations. All vectors considered in this paper are column vectors. We use ✶ { A } to denote the indicator function of A , i.e., ✶ { A } = 1 if A holds, and ✶ { 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 = { 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<L max , 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<L max given a training dataset D 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 × d 1 , W q ∈ R d 1 × d and W v ∈ R d 2 × d , namely key, query, and value matrices, and the feed-forward layer is determined by W o ∈ R |V|× d 2 . Here d 1 , d 2 are hidden dimensions. Mathematically, given the input X = [ x 1 . . . , x L ] , we write the one-layer transformer as T θ ( 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:

where I( x L +1 ) is the index of x L +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 a realizable training dataset for next-token prediction. We will then describe a 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 = [ 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 ∈D 0 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 ′ or x ′ > 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 > SP. Partial order under query O. O > >

OS, O OV.

Partial order under query V. V > VS. Partial order under query P. O > PP.

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 > VS, 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 sentence X ( n ) , we calculate its frequency π ( n ) ∈ [0 , 1] in dataset D 0 as π ( n ) = ∑ ( X,x L +1 ) ∈D 0 { X = X ( n ) } .

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:

4.2 Training Algorithm

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 0 3 , we use normalized gradient descent to train W ov . Equivalently, the loss function can be written as

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 ∇ W ov L 0 ( W ( t ) ov ) ∥∇ W ov L 0 ( W ( t ) ov ) ∥ .

5: end for

6: for t ∈ { 0 , . . . , T 1 -1 } do

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

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

̸

It 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, we have the following convergence result on the training of W ( t ) ov .

Proposition 1 Let W ∗ ov be defined in Equation (2) . Under Assumptions 1-2, let W ( t ) ov be updated by Algorithm 1. Then, for any t ≥ 2 , we have tη 0 2 ∥ W ∗ ov ∥ ≤ ∥ W ( t ) ov ∥ ≤ tη 0 and the following bound holds:

Moreover, the loss function L 0 satisfies that L 0 ( W ( t ) ov ) ≤ O (exp( -η 0 t/ (4 ∥ W ∗ ov ∥ ))) .

Proposition 1 states that during the training stage 1, the feed-forward layer W ( t ) ov converges in direction to W ∗ ov / ∥ W ∗ ov ∥ at a rate of O (log t/t ) , which classifies the next token from all other tokens. In addition, since the norm of W ( t ) ov increases linearly, the loss L 0 ( W ( t ) ov ) converges linearly to zero, i.e., L 0 ( W ( t ) ov ) = O (exp( -C 0 t )) for some constant C 0 .

5.2 Convergence of Training W kq

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 no x ℓ ′ 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:

Wewill 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 in X ( n ) , we have | l ( n ) | ≥ ∑ ℓ { 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.

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 sentence X ( 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 NL 4 max ∥ W ∗ kq ∥ 6 log 2 t/η , the optimal token weight satisfies min 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:

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

5.3 Proof Sketch of Theorem 1

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 product x ( n ) ℓ ( X ( n ) -1 ) ⊤ , max-margin attention weight matrix W ∗ kq , and the trained attention weight matrix W ( t ) kq .

Step 1 (Lemma 5). By analyzing 〈 ∇ W kq L ( θ ( t ) ) , x ( n ) ℓ ( X ( n ) -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 :

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

6 Generalization Ability

In this section, we prove the generalization ability of the trained transformers. Recall that Theorem 1 shows that W ( 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 , and

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

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

7 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 ov and W kq by Algorithm 1, each with 900 iterations. The parameters are chosen as d = |V| , η 0 = 0 . 2 / √ d , and η = 0 . 05 / √ d . In Figure 2, the first three plots show the dynamics of the training stage 1, which indicates the convergence of the loss L 0 ( W ( t ) ov ) to its minimum value, the convergence of W ( t ) ov in direction to W ∗ ov , and the linear increase of the norm ∥ W ( t ) ov ∥ , 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 L ( θ ( t ) ) , the convergence of W ( t ) kq in direction to W ∗ kq , and the linear increase of the norm ∥ W ( t ) kq ∥ . These results verify Theorem 1 and Theorem 2. All experiments are conducted on a PC equipped with an i5-12400F processor and 16GB of memory.

2

1

3

W

M

￥

0.5

1.0

IGMI

20

40

30

10

20

W

￥

0.5

1.0

10

5

0

0

300600900

0.0

0

300600900

0

0

300600900

0

0

300600900

0.0

0

300600900

0

300600900

t

t

t

t

t

t

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

Acknowledgments and Disclosure of Funding

The work of R. 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-2134145.

References

Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. (2024). Transformers learn to implement preconditioned 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:2211.15661 .

Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. (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., et al. (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 arXiv: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.19442 .

Cheng, X., Chen, Y., and Sra, S. (2023). Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06528 .

Cui, Y., Ren, J., He, P., Tang, J., and Xing, Y . (2024). Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401.17426 .

Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. (2023). On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310.12680 .

Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 .

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 at scale. arXiv preprint arXiv:2010.11929 .

Dryer, M. S. (1991). Svo languages and the ov: Vo typology1. Journal of linguistics , 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 on 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, V. (2023). Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17086 .

Giannou, A., Rajput, S., Sohn, J.-y., Lee, K., Lee, J. D., and Papailiopoulos, D. (2023). Looped transformers as programmable computers. In International Conference on Machine Learning , pages 11398-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 meanfield 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 on 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 theoretical 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). Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics , pages 685-693. PMLR.

Lin, L., Bai, Y ., and Mei, S. (2023). Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint arXiv:2310.08566 .

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 arXiv:2307.03576 .

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

Phuong, M. and Lampert, C. H. (2020). The inductive bias of relu networks on orthogonally separable data. In International Conference on 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.16898 .

Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. (2023b). Max-margin token selection in attention mechanism. In 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-layer transformer. Advances in Neural Information Processing Systems , 36:71911-71947.

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

Vardi, G. (2023). On the implicit bias in deep-learning 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 preprint arXiv:2402.05738 .

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems , 30.

Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., 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 .

Contents

A Expression of Gradients

We first provide the general formula for the gradients of both layers.

B Proof of Proposition 1

Recall that we use the loss,

The updating rule of W ov is that

We know that L 0 is convex respect to W ov . Therefore, we have

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.

Proof. Using Equation (5), we have

̸

where ( a ) is due the constraints that W ∗ ov satisfies. On the other hand,

In this case, the result is trivial.

̸

Then, we consider the case when W ( 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

̸

where ( a ) follows from ∥ AB ∥ ≤ ∥ A ∥∥ B ∥ for any matrices A and B . Thus, we obtain

To lower bound the norm of W ov , we recall the updating rule (Equation (7)).

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.

Proof. First, we consider the case when W ( t ) ov = W ∗ ov ∥ W ( t ) ov ∥ ∥ W ∗ ov ∥ . Due to Equation (8) in Lemma 1, we have

̸

Therefore, the loss on W ( t ) ov can be lower bounded as follows.

where ( a ) is due to the fact that log(1 + x ) ≥ x/ 2 when 0 < x < 1 . On the other hand, let W ′ ov = ( ∥ W ( t ) ov ∥ ∥ W ∗ ov ∥ +2log(2 |V| ) ) W ∗ ov . Then, the loss on W ′ ov has the following upper bound.

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 have

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.

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 have

Therefore,

where ( a ) follows from Lemma 1 and ∥ W (2) ov ∥ ≤ 2 η 0 ≤ ∥ W ∗ ov ∥ , and ∥ W ( t ) ov ∥ ≥ tη 0 / (2 ∥ W ∗ ov ∥

C Proof of Theorem 1 and Theorem 2

C.1 Supporting Lemmas

Lemma 3 With zero initialization, under the updating rule Equation (7) , for any iteration t , W ( t ) ov satisfies that

̸

Proof. The proof follows directly from induction and the fact that

̸

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 ≥ 2 y ( x -y ) for any x, y ∈ R .

Summing over t starting from 2 , we have

Furthermore, due to Lemma 1,

Similarly,

)

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 )

̸

̸

Proof. The second equality follows directly from Lemma 3.

To show the first equation, we analyze

where ( a ) follows from Proposition 1, and ( b ) is due to Lemma 1. On the other hand, we also have

The 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 ) ˆ W 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

Lemma 4 (Projection of gradient of W kq ) If W ov satisfies that

̸

we have

̸

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( X ( n ) ) . In addition, we denote T θ ( X ( n ) ) by T ( n ) θ for simplicity. From Equation (6), we have

̸

where ( a ) is due to the fact that ∑ i ∈ [ |V| ] [T ( n ) θ ] i = 1 for any θ, n .

Main Steps

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 for W ( t ) kq is

C.2 Step 1

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.

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

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

where the last inequality is due to Assumption 3, the induction hypothesis and ∑ i ∈ [ |V| ] [T ( n ) θ ] i = 1 Therefore, for any n , we have

Proof.

We examine the gradient ∇ W kq L ( θ ) projected onto the optimal direction W ∗ kq / ∥ W ∗ kq ∥ .

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

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 W (0) kq and the updating rule Equation (10) , for each iteration t , the following inequality holds.

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,

Thus, we have

By the updating rule Equation (10), we have

In addition, by the triangle inequality,

The proof is completed.

C.4 Step 3

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 = ⌈ 8 L max ∥ W ∗ kq ∥ 2 η ⌉ . Then, for any t ≥ t 0 , we have

where

Proof. During the proof, we denote

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

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 problem

Therefore, if W ( 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 have

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

Let L n ( θ ) = -log e ⊤ In( X ( n ) ) T θ ( X ( n ) ) be the loss on sample X ( n ) .

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

We further let

and

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

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 hand

where ( a ) is due to Equation (11), and ( b ) is obtained by choosing T = log 2 L max ∥ 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(2 L max ∥ W ( t ) kq ∥ )) ∥ W ( t ) kq ∥ . Thus, we further have

Next, we analyze B ( n,t ) , and further divide B ( n,θ ) into B ( n,θ ) + and B ( n,θ ) -:

Due to Proposition 4, we have B ( n, ∗ ) -= 0 , and thus

We then analyze:

where ( a ) follows from that b ∗ ,ℓ ≤ -(1+ β 0 ) ∥ W ( t ) kq ∥ / ∥ W ∗ kq ∥ and Equation (11), and ( b ) is due to the fact that ∑ i ∈ [ |V| ] [T ( n ) θ ] i = 1 for any θ, n , and ( c ) follows from β 0 / 2 ≥ ∥ W ∗ kq ∥ ∥ W ( t ) kq ∥ log(2 L max ∥ W ( t ) kq ∥ )

For the term B ( n,t ) -, we have

Note that

and

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

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 φ ( n 0 ,t ) -≥ exp( -∥ W ( t ) kq ∥ / ∥ W ∗ kq ∥ ) Then, we aim to show that

Since

We further note that [T ( n 0 ) θ ( t ) ] In( x ( n 0 ) ℓ ∗ ) < [T ( n ) θ ( t ) ] In( x ( n ) ℓ ∗ ) due to φ ( n 0 ,t ) + < φ ( n,t ) + . Thus,

where ( a ) is due to that β 0 ≥ 2 ∥ W ∗ kq ∥ (1+2 ∥ W ∗ kq ∥ ) ∥ W ( t ) kq ∥ log(1 + ∥ W ( t ) kq ∥ 2 ∥ W ∗ kq ∥ ) , ∥ W ( t ) kq ∥ ≥ 2( e -1) ∥ W ∗ kq ∥ , and β 0 ≥ 2 ∥ W ∗ kq ∥ ∥ W ( t ) kq ∥ log 1 + β 0 +2 ∥ W ∗ kq ∥ β 0 .

In summary, we have that

where α t ≥ ( N +3) β 0 . The proof is finished.

Now, we are ready to prove Theorem 1.

C.5 Proof of Theorem 1

Proof of Theorem 1.

Let t 0 = ⌈ 8 L max ∥ W ∗ kq ∥ 2 η ⌉ be defined in Lemma 7. Summing over t from t 0 , we have

Furthermore,

where ( a ) follows from η ≤ ∥ W kq ∥ /e .

Therefore, we have

The proof is finished.

C.6 Proof of Theorem 2

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

Note that for all ℓ ∗ ∈ l ( n ) and ℓ / ∈ l ( n )

where ( a ) follows from t ≥ 1696 NL 4 max ∥ W ∗ kq ∥ 6 log 2 t/η . Therefore,

where the last inequality follows from t ≥ 4 L max ∥ W ∗ kq ∥ η log L max ϵ Hence, the loss on the sentence X ( n ) satisfies that

.

̸

̸

Thus, the average loss has upper bound, which is

for C 1 = η/ (4 L max ∥ W ∗ kq ∥ 2 ) , and ∆ = C 0 T for C 0 = η 0 / (4 ∥ W ∗ ov ∥ 2 ) .

D Proof of Proposition 2 and Theorem 3

D.1 Proof of Proposition 2

Proposition 4 (Restatement of Proposition 2) Under Assumption 2, if W ∗ kq satisfies Equation (3) , i.e.,

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 have

A direct result is that

Proof. Let U ∈ R d be the rotation matrix such that Ux = e I( x ) . Because U preserves Frobenius norm, the optimization problem in Equation (3) can be written as

Notably ˜ W ∗ = UW ∗ 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 ) ℓ ) ) ⊤ We 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

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 is

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

Thus,

Let

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

where x = w k +1 = . . . = w k + m . Hence, the optimal solution is w 1 = . . . = w k = m/ ( m + k ) , and w 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

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 have

D.2 Proof of Theorem 3

Proof. The proof follows similar logic to Theorem 2. By Equation (13), we have for any x, x ′ ∈ V

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(

Then, we have

for sufficiently large

̸

̸

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.

x

W

(

t

)

kq

x

)

t

= Ω(log(1

/ϵ

))

such that

∑

φ

1

-

ϵ

.

ℓ

L

ℓ

0

∈

l

0

ℓ

0

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

2. Limitations

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.

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: We provide the full proofs in 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 crossreferenced.

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

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

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

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

4. Experimental Result Reproducibility

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

Answer: [Yes]

Justification: We provide our experiment details in Section 7, including the model dimension, vocabulary size, the learning rate, and the iteration number. The experiment follows exactly of Algorithm 1.

Guidelines:

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

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

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

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

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide our code in the supplemental. We do not use open source data. Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

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

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

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

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

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

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: 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 to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

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

Answer: [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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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

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

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

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

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

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

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

8. Experiments Compute Resources

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

Answer: [Yes]

Justification: Our experiment is on synthetic data and is computationally lightweight, which does not require any GPU. The running 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 of compute required for each of the individual experimental runs as well as estimate the total compute.

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

9. Code Of Ethics

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

Answer: [Yes]

Justification: 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 a deviation from the Code of Ethics.

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

10. Broader Impacts

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

Answer: [NA]

Justification: This paper presents work whose goal is to advance the field of Machine 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).

11. Safeguards

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

Answer: [NA]

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

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

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

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

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

Answer: [NA]

Justification: 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 crowdsourcing nor research with human subjects.

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

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

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

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

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

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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…\n\nQUESTIONS:\nHow 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?\n\nLIMITATIONS:\nSee above, no negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

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

STRENGTHS:
1. The paper introduces novel theoretical frameworks for analyzing transformer training, focusing on non-asymptotic convergence.
2. The research is technically robust, with sound mathematical derivations and empirical validation.
3. 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.
2. While the theoretical results are strong, additional empirical studies, particularly on real-world datasets, could further strengthen the claims.

QUESTIONS:
What are the potential implications of the findings on transformer training efficiency and computational costs in practical settings?

LIMITATIONS:
NaN
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-16

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe approach described in the paper may not align with practical Transformer training.\nPlease see the questions below for more details.\n\nQUESTIONS:\nQ1: 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?\n\nLIMITATIONS:\nThe assumption of the existence of collocations seems far from real traninng dataset.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-18"}