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

We thank the reviewer for the helpful comments, and would like to provide the following responses.

Q1: How is the loss function between lines 227 and 228 obtained? Why is its form different from Eq. 1?

A1: The loss in Eq.\ (1) is for the general-length case. The loss between lines 227 and 228 is for the training stage 1 (i.e., training of feed-forward layer), where the inputs are one-token sentences. We note that there was a typo in Eq.\ (1), i.e., missing $\exp$ in the first bracket. The correct form should be
\log \left(\sum_v \exp( e_v^\top \bar{T}(X) ) \right) - e_i^\top \bar{T}(X) ,\

where $i$ is the index of the next token of $X$. Such a form can specialize to the loss for stage 1 as follows. Let $X = [x_1]$. Then we have $\bar{T}(X) = W_{ov}x_1$. Thus, the loss becomes
\log\left( \sum_v \exp(e_v^\top W_{ov}x_1) \right) - \log\exp(e_i^\top W_{ov} x_1) = -\log\frac{ \exp(e_i^\top W_{ov} x_1) }{ \sum_v \exp(e_v^\top W_{ov}x_1) }.\

Thank you for your question and we will fix the typo in the revision.

Q2: In Section 6, the authors theoretically prove the generalization ability of the trained Transformer. Can the generalization ability be proved in the experimental part?

A2: Yes. Please see the attached PDF file in the global response for the test loss result. For the test dataset, we randomly construct test data as described in Theorem 3, and plot the average test loss for 20 trials. As shown in the figure, the test loss also converges almost to 0. In addition, since the generalization ability is determined by the parameter $W_{kq}$ and $W_{ov}$, alignment of these parameters with their corresponding max-margin solutions as shown in Figure 2 of the paper implies the generalization ability as well.

Q3: Why do the authors design a two-stage training algorithm? Is it valuable in practical applications compared to single-stage training?

A3: Thanks for the question. Our characterization of the realizable dataset is via two steps: (i) identifying collocation and (ii) using collocation to define partial ordering and optimal token. This naturally motivates the two training stages to learn the corresponding information: (i) stage 1 learns collocation, i.e., next-token prediction for length-one sentences; and (ii) stage 2 learns partial ordering to identify optimal token. Thus, it is also meaningful to apply such a two-stage algorithm in practice. We also expect that stage 1 of pre-processing value matrix using one-token sentences could help to stablize the entire training process of transformers. Single-stage training may still yield good solution simply from optimization perspective, and we are currently working on analytically characterizing the dynamics of such a training process.

Q4: This manuscript simplifies the Transformer into a single layer for analysis. While I understand the need for this simplification, it would be nice to include a brief paragraph outlining the views on generalizing to multi-layer Transformers.

A4: Thanks for the suggestion. A possible approach to generalize it into multi-layer transformers is to view our single layer as the last layer of a multi-layer transformer. Therefore, the input $X$ in the paper can be viewed as the output of a multi-layer network. Then, we can combine our technique for the last layer and the technique of neural tangent kernel (NTK) for other layers (Emami et al., 2021), and analyze the training dynamics of such a multi-layer transformer. We expect that the last layer evolves similarly to our findings given that other layers output $X$ satisfies certain conditions, which may be potentially shown by exploiting the realizability of the dataset and the properties of NTK.

Emami et al. "Implicit bias of linear RNNs." ICML 2021.

---

We thank the reviewer again for the highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

We thank the reviewer for the helpful comments, and would like to provide the following responses.

Q1: How would you change the manuscript to make it more easy to understand your definitions and motivation behind the mathematics? For example, the example dataset provides some additional insights, but does not clarify why these properties are important, how this leads to zero training loss in the limit, etc.

A1: Thanks for the suggestion. We will provide the following high-level description of our dataset and their properties.

The goal of introducing the collocation and query-dependent partial order is to make training transformer provable and interpretable. We breakdown the insights as follows.

(i) Since the model should apply to all subsequences of sentences in the training set including length-two sentences, then collocation between two words naturally becomes a property of the dataset.

(ii) It can be observed that the input of our loss function is a convex combination of $L$ points each being a projection of one token embedding in the sentence, where the weights in the convex combination are the attention scores of each token with the query token. In this way, the loss values of these tokens determine an order among them, and the token that achieves the minimum loss value serves as the optimal token. In other words, some tokens has larger attention scores than others under a specific query. These facts motivate us to introduce query-dependent partial order.

Combining (ii) with (i), the collocated token of the optimal token becomes the predicted next token. Finally, due to the injection property of collocation and the structure of cross-entropy, the loss function can be trained to approach zero.

Q2: How would you change the manuscript to better assess/discuss the limitations of your work? For example, provide discussion of how general the obtained insights are; i.e., what the paper contributes beyond the setting of the simple network architecture, custom dataset and custom training method used to derive the results.

A2: Thanks for the suggestion. We provide the following discussion on the generality of our findings and will add it to our revision.

We remark that our work establishes the basic framework to analyze the convergence performance of transformers for NTP task. Several insights can be extended to more complex settings.

Regarding multi-layer transformers, for example, one possible idea is to view the single layer in our work as the last layer of a multi-layer transformer. This is promising if we consider the last layer training regime (Kirichenko, et al., 2022). In this regime, if the output of other layers satisfies certain conditions such as near-orthogonality, which is often satisfied in high dimensional setting, our convergence and generalization results can be directly applied to multi-layer transformers.

To generalize dataset, for example, our technique can be extended to study a noisy version of collocation. We expect that most of our proof steps can still be applied by incorporating appropriate noise bound.

To generalize training method, an alternative natural method is to update the linear layer and the attention layer simultaneously instead of two-stage training. A popular way to implement such an algorithm is via two timescale stepsizes so that the attention update can be in a faster scale and the final convergence can be determined by the feedforward layer by incorporating the suboptimality error of the attention layer. Such suboptimality error can be bounded by our current techniques for training stage 2.

Kirichenko et al. "Last layer re-training is sufficient for robustness to spurious correlations." arXiv:2204.02937, 2022.

---

We thank the reviewer again for the highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

We thank the reviewer for the helpful comments, and would like to provide the following responses.

Q1: Some of the mathematical concepts, particularly the query-dependent partial orders, are complex and could be better explained.

A1: We appreciate the feedback on the complexity of the mathematical concepts. To address this, we will provide more intuitive explanations and examples in the revised manuscript.

Specifically, query-dependent partial order formalizes the concept of importance or relevance of tokens in a given sentence. If one token is more important than others or more relevant to the next token, it is considered as an optimal token and thus is ''greater'' than other tokens. For example, assume that the sentence Machine learning is a popular has next word area, and popular is the most important word. Then, under the query popular, popular itself is the optimal token, and hence it is ''greater'' than other tokens.

In addition, such a partial order is query-dependent so that the order can be different given different queries. For example, consider a new sentence A popular area is machine with next word learning. Then, under the query machine, machine itself is more important or ''greater'' than popular.

Finally, the motivation for defining these concepts is detailed in Response A1 to Reviewer ZtZ3.

Q2: While the theoretical results are strong, additional empirical studies, particularly on real-world datasets, could further strengthen the claims.

A2: Thanks for the suggestion. We are actively seeking suitable real-world datasets and experiments to further illustrate our results.

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

A3: Thanks for the question. Our results have the following implications. (1) The collocation training can be separated from the attention training to reduce the complexity of joint training, and at the same time achieve stable and fast convergence rate for each separate training. (2) Our result shows that the loss function converges fast (with potentially linear convergence rate), and hence the training does not require many iterations. The main computational costs lie in the gradient computations in each iteration, which scales with the number of training parameters. Thus, employing techniques such as LoRA to reduce the number of training parameters is crucial to achieve scalable training.

---

We thank the reviewer again for the highly inspiring comments. We hope that our responses have resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.

We thank the reviewer for the helpful comments.

Q1: Deterministic ground truth model seems uncommon. Besides assuming n, where is the existence of $p_L^*$ theoretically necessary?

A1: In recent line of theoretical research on transformers, deterministic models are often adopted such as in Li et al., (2024) on next-token prediction and Tarzanagh et al., (2023) on binary classification, mainly because deterministic models can offer valuable theoretical insights due to their simplicity and clean structure. In particular, it provides a clear cause-and-effect relationship between inputs and outputs, which makes it appealing to facilitate the interpretation of the model. Note that deterministic model of ``separable data'' is also commonly studied in theoretical works on neural networks such as in Soudry et al., (2018).

The ground truth model $p^*_L$ serves as a theoretical baseline for analyzing an algorithm so that the performance of a training algorithm is well defined. As a notation, $p^*_L$ can be replaced by $\mathrm{n}(X)$, where $X$ contains $L$ tokens. As a realizable model, it is captured by collocation and partial order. With the existence of such a model $p^*_L$, the training algorithm can have provable performance.

Tarzanagh et al. Transformers as support vector machines. arXiv:2308.16898, 2023.

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.

Q2: Existence of collocations seems far from reality; realistic data where such assumptions hold? How to prepare n() in real datasets?

A2: Thanks for the question. The existence of collocations can be observed in many natural languages, where certain words frequently appear together (Lehecka 2015). Taking such an assumption enables us to derive interesting theoretical results, while remaining arguably reasonable in practice. Note that the assumption of separable data, as a special case of this, has formed an important line of theoretical research in deep learning to explore the implicit bias and training dynamics of neural networks (Soudry et al. 2018, Chizat and Francis 2020).

Given a dataset, $\mathrm{n}$ can be prepared by the counting method described in Lehecka, T. (2015). More generally, it can be extracted by simply including all two-token sequences from the dataset, and each second token $y$ is the next token of the first token $x$, i.e., $n(x)=y$. Furthermore, leveraging domain-specific knowledge, such as medical dictionaries, can also enrich collocations. These approaches allow for the practical extraction of $\mathrm{n}()$ in a scalable and efficient manner.

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

Chizat and Francis Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. COLT 2020.

Q3: How to understand definition of partial order.

A3: Thanks for the question. We clarify that the partial order is query-dependent. Therefore, if $x>\_{x_q} x'$ and $n(x)=x_{L+1}\neq n(x')$ in some sentences ending with query $x_q$, it is not possible that $n(x)\neq x_{L+1} = n(x')$ in another sentence ending with the same query $x_q$.

This is because of the structure of the transformer where the attention weights are proportional to $\exp(x^\top W_{kq} x_q)$. If $n(x)=x_{L+1}\neq n(x')$ under query $x_q$, then $x^\top W_{kq} x_q > (x')^\top W_{kq}x_q$ as long as the sentence ends with $x_q$. Thus, $x$ is always more important than $x'$ in $x_q$-ended sentences. However, $n(x)\neq x_{L+1} = n(x')$ is still possible in other sentences ending with different queries than $x_q$. These facts motivates us to introduce query-dependent partial orders.

Q4: How commonly is the normalized gradient used? AdamW is typically used.

A4: Thanks for the question. Although normalized gradient descent (NGD) is often favored in theoretical studies, there are also empirical works showing that the achieved accuracy of NGD is slightly better than Adam in training transformers (Cutkosky and Mehta 2020), where momentum is employed to improve the performance of NGD. Cutkosky and Mehta (2020) also theoretically proved that NGD achieves a fast convergence rate for generic optimization, which aligns with our theoretical findings that NGD achieves a linear convergence rate for NTP with one-layer transformers.

We acknowledge that AdamW is indeed more commonly used in practice. However, it is much more challenging to analyze the training dynamics under AdamW theoretically. So far, all theoretical analysis on transformers have been focused on GD and its variants. It is somewhat necessary to first develop techniques for GD and its variants, before these tools can be further advanced to study AdamW. Our work takes a first step to understanding the NGD dynamics of Transformers for NTP, and we are actively exploring the analysis for Adam-based Transformer training.

Cutkosky and Mehta, Momentum Improves Normalized SGD, ICML 2020.

Q5: Assumption 3 may not hold in situations where attention is sparse.

A5: Thanks for the question. We apologize that the statement of Assumption 3 may cause confusion. It actually requires that the number of optimal tokens is greater than or equal to the number of each individual non-optimal tokens, not the summation of all non-optimal tokens. Hence, such an assumption is not in conflict with sparse attention. One simple example satisfying Assumption 3 is a sentence where every token is distinct. The transformer's attention needs to focus only on the single optimal token, which is clearly sparse.

---

We thank the reviewer again for the highly inspiring comments. We hope that our responses resolved your concerns. If so, we wonder if the reviewer could kindly consider to increase your score. Certainly, we are more than happy to answer your further questions.