On the Generalization Ability of Next-Token-Prediction Pretraining

Thank you for your insightful comments and questions. We have carefully considered them and have added supplementary explanations in the relevant sections of the paper. Additionally, we have incorporated the suggested references into our paper. Below are our responses:

A1: This is indeed a significant question, and we fully understand your concerns. We must acknowledge that the sequence length $m$ is constrained by available training resources. However, from a theoretical standpoint, we assume that the sequence length $m$ can approach infinity. Firstly, from the intrinsic nature of human language, language sequences can indeed extend indefinitely, much like time series. Furthermore, based on the two papers [1, 2] we are aware of that also consider the generalization bounds of language models at the token level, their bounds include $\mathcal{O}(\sqrt{1/m})$, although they do not specifically target the NTP pre-training task. For instance, in Theorem 4.3 of paper [1], there is a term $\mathcal{O}(\sqrt{\frac{1}{T_p}})$, where $T_p$ represents sequence length, similar to $m$. Paper [2] also considers dependencies between tokens, and thus in their Theorem 1, the generalization bound is $\mathcal{O}(\sqrt{1/m})$. We hope this explanation alleviates your concerns. Additionally, we will enhance the discussion on sequence length in our paper and look forward to reaching a consensus with you.

A2: It is true that larger models often perform better in practice. However, according to the scaling laws for large models [3], larger models also require more tokens for pre-training. In the third experiment depicted in Figure 2 of our paper, we observed that when the total number of tokens is fixed, larger models tend to generalize worse. This is primarily because the limited number of tokens leads to overfitting in larger models. This observation aligns with our theoretical results.

A3: For our detailed explanation regarding disc(U), please refer to ​Response A3 in our reply to Reviewer axCZ. Furthermore, $C_{\varphi,r}$ represents the upper bound of $\Delta_{m}$ ( Remark 4.10) and also varies with $\phi$, measuring the diversity of the token sentences. We aim for a smaller $C_{\varphi,r}$ to enable the model to learn more diverse knowledge. We will further explore disc(U) and $\Delta_{m}$ in the discussion section of the paper.

A4: In Section 3.1 of the paper, we clarify that all tokens $\mathbf{t}^i_j \in \mathbb{R}^{n_v}$ are vectors, meaning that the first parameter of the loss function $\ell$ is the probability vector computed by Equation (7).

A5: We apologize for not being able to experimentally demonstrate in a short time frame that human language sequence is a mixing process. However, we provide two theoretical supports: (1) In numerous existing NLP studies, language sequences are often modeled as Markov processes [5,6]. Under certain conditions, Markov chains are indeed mixing processes [7]. (2) Under certain conditions, autoregressive processes are equivalent to mixing processes [8].

A6: Regarding the calculation of the error bound, we followed the approach outlined in [9], our computed generalization bound is $\sqrt{\frac{\Theta \tau_{1}}{Nm}}$. (1) We set the $G_{\pi}$ to 1 based on Lemma A.9 in [10]; (2) $B_{l}$ and $C_{l}$ were calculated by extracting the model parameter matrix and the attention matrix during training. We set the epoch to 2000 to ensure the model converges as much as possible [11], as smaller datasets require more epochs for generalization [12]. The primary purpose of this experiment was to demonstrate that even with a limited number of training tokens, our generalization bound remains valid and does not collapse as discussed in [9]. Additionally, we are actively working to supplement our study with more experiments.

Reference：

[1]Gong.(2025) Towards Auto-Regressive Next-Token Prediction: In-context Learning Emerges from Generalization.

[2]Lotfi.(2024) Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models.

[3]Open AI.(2020) Scaling Laws for Neural Language Models.

[4]Penedo, M.(2023) The Refined Web Dataset for Falcon LLM: Outperforming Curated Corpora with Web Data, and Web Data Only.

[5]Zhang, Y.(2023) What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization.

[6]Li, Y.(2023)Transformers as algorithms: Generalization and stability in in-context learning.

[7]Meyn.(2012) Markov chains and stochastic stability.

[8]Krishna.(1986) Mixing properties of harris chains and autoregressive processes.

[9]Galanti.(2023) Norm-based Generalization Bounds for Sparse Neural Networks.

[10]Edelman.(2022) Inductive Biases and Variable Creation in Self-Attention Mechanisms.

[11]Madden.(2024) Upper and lower memory capacity bounds of transformers for next-token prediction.

[12] Power.(2022) Grokking: Generalization beyond overfitting on small algorithmic datasets.

First, thank you very much for your recognition and support of our work. We have incorporated your valuable suggestions by adding a comparative discussion of related past and recent work in the appendix and have made every effort to supplement our experiments. Additionally, we are more than willing to address any questions you have about this work:

A1: The expression $\ln(1+x/e)^{0.5}$ can be seen as $g=f^{0.5}$ where $f=\ln(1+x/e)$. The function $g$ is "continuously concave (downward)."

A2: This uses a property of the softmax function: $\sum\_{j\ne i} e\_j=1-e\_i$, so the expression inside the parentheses becomes $2(1-e\_i)$.

A3: Please note the definitions of $\mathcal{C}\_{Q}$ and $\mathcal{C}\_{K}$ in line 1093, where the input data is any $\mathbf{Z}$. Here, $||\mathbf{Z}^*||\_F$ corresponds to $||\mathbf{X}||\_F$ in Lemma C.4, as we take the maximum value of the norm for the upper bound. Meanwhile, $||\mathbf{Z}\_{[N]}||\_F$ results from substituting the value of $\epsilon\_{Q}$.

A4: This is an excellent question! Since we are discussing the pre-training phase, as described in our paper and the paper[6], this phase is not autoregressive. Autoregression is the mechanism used during model reasoning. In the pre-training phase, $m$ queries (tokens) are input simultaneously, and $m$ output tokens are produced simultaneously. Due to the masking mechanism, each query can only use preceding information, mimicking the autoregressive mechanism without seeing future information.

A5: In line 1618, the right side of the inequality is a lower bound. Theoretically, $\alpha$ should be chosen to minimize this expression, i.e., where the derivative is zero. However, since deriving the expression with respect to $\alpha$ is complex, we set $\alpha=\frac{1}{\sqrt{Nmd}}$. Although this does not reach the lower bound, it simplifies derivation and analysis.

A6: The training data used is too limited, so dropout is added to prevent overfitting, similar to the experimental setup in [1].

A7: This does not involve any parameter optimization issues; it is solely to analyze the impact of parameter changes on model performance based on experimental results. We have clarified this in the paper.

A8: Thank you for your question. To explore the optimal selection of parameters $N$ and $m$ under a fixed total number of tokens, as you suggested, we should vary the sequence length while keeping the total number of tokens constant. Our experiments aimed to verify theoretical results, but indeed, the total number of tokens increased. We are working on this exploratory experiment and aim to present the results in the final version.

A9: Many empirical studies indicate that increasing model parameters excessively, while keeping the number of pre-training tokens constant, can harm model performance and generalization ability. Both Chinchilla's Law [2] and Scaling Laws [3] suggest that token numbers should increase alongside model parameters. This may be because larger models with fixed training tokens are more prone to over-parameterization and overfitting. As research in Nature paper [5] shows, for simple tasks with fewer tokens, large models may "memorize" noise or spurious patterns in the training data rather than learning underlying general rules, leading to poorer performance.

A10: Your understanding is correct. For example, in the sentence "I went to the restaurant today, the food was delicious, and after eating, I went to the park, which was crowded," let event A = "I went to the restaurant today," B = "the food was delicious," and C = "the park was crowded." According to the definition of $\varphi$-mixing, the time step $k\_1$ between A and B is much smaller than the time step $k\_2$ between A and C, so we have $\varphi(k\_1)>\varphi(k\_2)$, indicating that the dependency between A and B is stronger than between A and C. The dependency decreases as the time step $k$ increases.

A11: For the response regarding disc(U), please refer to response A3 in our reply to Reviewer axCZ.

Reference:

[1]Edelman, B.(2022) Inductive biases and variable creation in self-attention mechanisms.

[2]DeepMind.(2022) Training Compute-Optimal Large Language Models.

[3]Open AI.(2020) Scaling Laws for Neural Language Models.

[4]Penedo, M.(2023) The Refined Web Dataset for Falcon LLM: Outperforming Curated Corpora with Web Data, and Web Data Only.

[5]Zhou, L.(2024) Larger and more instructable language models become less reliable.

[6]Bachmann, G.(2024) The pitfalls of next-token prediction.

We are delighted to receive your recognition and support for our work, and we appreciate your careful attention to details we might have overlooked. We have made corrections to the relevant sections of the paper based on your suggestions. Additionally, we plan to release the code upon acceptance of the paper. Below are our responses to your questions, which we hope will address any concerns:

A1: Our work primarily focuses on the proof details, and we initially lacked a thorough technical comparison with previous works. We have now enhanced these sections. Compared to earlier works such as [1,2], our main technical challenges and innovations include:

(1) Considering both the independence between sequences and the dependency between tokens within sequences. Unlike previous works that only considered independent and identically distributed token sequences, we introduced the $\varphi$-mixing tool for analysis. This increased the complexity of error decomposition but also enhanced the interpretability in the NTP pre-training scenario, specifically the bound's decrease with sequence length.

(2) We introduced the Rademacher complexity of composite function spaces and proposed its decomposition theorem (Proposition 1), enabling separate analysis of different function spaces. This greatly simplifies the generalization analysis when changing DOMs task heads or transferring pre-trained DOMs to different downstream tasks.

(3) We performed a detailed analysis of the covering number for masked-self-attention-based decoder-only models. Our approach maintains a high structural consistency with the original Transformer [3] and analyzes under the F-norm, which was not discussed in previous works. The main challenge in considering the mask matrix is defining the upper bound of the attention matrix's norm, as seen in Lemma C.8. This results in an additional $\sqrt{\ln(m)}$ factor compared to analyses that do not consider the mask matrix's norm upper bound.

A2: Thank you for your question. We compare our work in two aspects, denoting the input data matrix as $\mathbf{X}_{[N]} \in \mathbb{R}^{Nmd}$: (1) Regarding the dependence on $\epsilon$ and data norm, their bound is $\sqrt{\frac{||\mathbf{X}\_{[N]}||^2}{N}}$, while ours is $\sqrt{\frac{\ln(||\mathbf{X}\_{[N]}||\_{F}/\sqrt{Nmd})}{Nm}}$. Given that the order of $||\mathbf{X}\_{[N]}||$ is $\sqrt{Nmd}$, their bound has a greater dependence on the data norm, whereas our bound almost eliminates this dependence. (2) Regarding the dependence on model dimension $d$, since the model parameter count $\Theta=\mathcal{O}(d^2)$, their bound has a logarithmic dependence on $\Theta$ due to its logarithmic dependence on $d$. Our bound has a polynomial dependence on $\Theta$. According to the Scaling Laws for language models [4], the dataset size should grow sub-linearly with the model parameter count, making our bound more consistent with the Scaling Laws.

A3: In our work, disc(U) primarily measures the overall quality of pre-training data. We assume all sequences are generated by a $\varphi$-mixing process, defined as an infinitely long sequence, implying that most sequences may originate from the same mixing process. On one hand, two different sequences might be extracted from the same corpus or even the same article. On the other hand, due to the nature of human language, even seemingly unrelated sentences like "Large language models benefit humanity" and "The weather is nice today" can appear in the same context when read together. This indicates that any high-quality sentence conforming to human language rules can be interconnected through language. Conversely, sequences like gibberish, grammatical errors, unclear expressions, or incorrect knowledge struggle to establish connections with high-quality human language. This underscores the importance of data cleaning quality [5]. Fewer low-quality sequences in the pre-training data result in a smaller disc(U), thereby enhancing the model's generalization performance.

A4： Regarding Lemma 4.16, we acknowledge that the notation $||W_Q,W_K,W_V||\_F \le B$ was indeed a space-saving simplification in the main text. The complete formulation, specifying $||W_Q||\_F ≤ B$, $||W_K||\_F ≤ B$, and $||W_V||\_F ≤ B$ individually, is provided in Lemma C.10. We have now added explicit clarification of this point in our paper. Additionally, we confirm that the issues raised regarding line 186 and line 427 have been addressed. We thank the reviewer again for bringing these important details to our attention!

Reference:

[1]Edelman, B.(2022) Inductive biases and variable creation in self-attention mechanisms.

[2]Deng, Y.(2024) On the generalization ability of unsupervised pertaining.

[3]Vaswani,A.(2017) Attention is all you need.

[4]Open AI.(2020) Scaling Laws for Neural Language Models.

[5]Penedo, M.(2023) The Refined Web Dataset for Falcon LLM: Outperforming Curated Corpora with Web Data, and Web Data Only.

Thank you for your recognition and support of our work. Below, we provide a detailed comparative analysis of our work with the two most relevant previous studies, [1] and [2]. We sincerely hope this will address your concerns. All of the following content has been added to the discussion section of our paper:

Our approach shares the overall methodology and thought process with [1] and [2], where we first decompose the excess risk error to derive an upper bound on the Rademacher complexity, and then obtain a more refined generalization error bound by analyzing the model capacity, i.e., the covering number. However, there are several key differences:

1. Task Context: We focus on NTP pre-training, using $\varphi$-mixing to model token sequences to capture dependencies between tokens within a sequence, whereas [1] and [2] only consider independence between sequences.

2. Rademacher Complexity: Since DOM can be viewed as composed of two parts, its function space is considered as a composite of two function spaces. Therefore, we introduced the Rademacher complexity of composite function spaces and derived the corresponding decomposition theorem, Proposition 1. Compared to the direct analysis of a single function space in [1] and [2], our method facilitates the analysis of pre-training generalization when changing DOM task heads and makes it easier to transfer to different downstream tasks.

3. Covering Number Analysis: Like [1] and [2], we analyzed the covering number of Transformers, which is why we share Assumption 4.14. Similar assumptions are common in covering number analyses, such as in [3] and [4]. However, [1] and [2] analyze the covering number of encoder-only models under the spectral norm, while we analyze decoder-only models under the F-norm. Our challenge includes considering the masked self-attention mechanism, with the main impact of the mask matrix detailed in Lemma C.8.

4. Final Results: [1] and [2] show polynomial dependence on data norms and parameter bounds, and logarithmic dependence on the number of model parameters, which is the opposite of our findings. Our results nearly eliminate dependence on data norms and show polynomial dependence on model parameters, aligning more closely with the Scaling Laws for large models [5], which suggest linear growth of data size with model parameter count. This is mainly because we cleverly applied the convexity and concavity of functions, as stated in Lemma C.6, to Lemma C.5, with detailed applications in Appendix D.

Additionally, Assumption 4.2 is crucial for considering dependencies between tokens within a sequence. Based on Assumption 4.2, we can use Lemma B.4 to analyze the model's generalization ability on a single token sequence.

Reference:

[1]Edelman, B.(2022) Inductive biases and variable creation in self-attention mechanisms.

[2]Deng, Y.(2024) On the generalization ability of unsupervised pertaining.

[3]Bartlett, P.(2017) Spectrallynormalized margin bounds for neural networks.

[4]Lin, S. and Zhang, J. Generalization bounds for convolutional neural networks.

[5]Open AI.(2020) Scaling Laws for Neural Language Models.