Implicit Optimization Bias of Next-Token Prediction in Linear Models

We initiate an investigation into the optimization properties of next-token prediction (NTP), the dominant training paradigm for modern language models. Specifically, we study the structural properties of the solutions selected by gradient-based optimizers among the many possible minimizers of the NTP objective. By framing NTP as cross-entropy minimization across distinct contexts, each tied with a sparse conditional probability distribution across a finite vocabulary of tokens, we introduce "NTP-separability conditions" that enable reaching the data-entropy lower bound. With this setup, and focusing on linear models with fixed context embeddings, we characterize the optimization bias of gradient descent (GD): Within the data subspace defined by the sparsity patterns of distinct contexts, GD selects parameters that equate the logits' differences of in-support tokens to their log-odds. In the orthogonal subspace, the GD parameters diverge in norm and select the direction that maximizes a margin specific to NTP. These findings extend previous research on implicit bias in one-hot classification to the NTP setting, highlighting key differences and prompting further research into the optimization and generalization properties of NTP, irrespective of the specific architecture used to generate the context embeddings.

Paper

Similar papers

Peer review

Reviewer k4Ko5/10 · confidence 3/52024-06-23

Summary

This paper studies the implicit bias of the gradient descent on the Next-Token Prediction (LTP) problem in linear models. They first formulate this NTP problem as minimizing the cross-entropy (CE) loss over distinct contexts, each tied with a sparse conditional probability over the token space. They then provide the necessary conditions for the CE loss to reach the entropy lower bound, i.e., the NTP-compatible condition and the NTP-separable condition. Then, they prove one sufficient condition for those two conditions is oevrparameterization, i.e., the dimension of the embedding space d is larger than the number of distinct contexts in the dataset. Assuming both compatible and separable conditions, they then prove the directional convergence of the minimizer of the CE loss within a certain range and the directional convergence of the GD iterate towards the direction of the solution of an NTP-SVM. In general, I think this paper delves into a good and important problem: the optimization path and implicit bias of NTP mechanism. The authors provided a good formulation, and the proof is solid.

Strengths

1. They investigate an interesting and important problem: the optimization path and the implicit bias of NTP. 2. Their formulation of NTP into the CE minimization over distinct contexts is novel. 3. They provide rigorous theoretical results and the proofs are solid, to my knowledge.

Weaknesses

1. The main issue of this paper is that, for the NTP-compatible and separable conditions to hold, one needs d > m. Does this overparametrization condition usually hold in practice or not? To my knowledge, in practice, the embedding dimension d is much smaller than the number of training data. Since m is not the number of training data and can be much smaller than that, it is not clear to me whether this assumption is possible in practice. 2. There are some paragraphs that are not very clearly written. For example, in lines 154-157, why does equation 4 constrain W^p w.r.t. this subspace? Why is the solution W* unique, assuming equation 4 has a solution? I think those can be expressed as lemmas to make them clearer. In line 148, the authors claim that (3a) holds if and only if the data satisfies the NTP-compatible condition. The 'if' direction is trivial, but the other direction needs a more rigorous proof.

Questions

See above.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

See above.

Reviewer kgdu7/10 · confidence 3/52024-07-13

Summary

This work studies the implicit bias of optimization in next token prediction tasks by analyzing the structure of the decoding matrix at infinite time. The paper introduces two novel conditions under which the loss reaches its minimum theoretical value and demonstrates that if these conditions hold (which can be, for example, the case when the model is overparameterized), then after GD training, the decoding matrix will converge (in direction) to a matrix reminiscent of the maximum-margin matrix in "standard" classification.

Strengths

This work studies a timely topic (next token prediction) and approaches it from a learning theoretic perspective (implicit bias of optimization), which has proven to be very fruitful in "standard" classification. The assumption of sparse contexts is clever and should be of wider applicability. The results are novel and analogous to similar results that were proven for "standard" classification. Furthermore, the presentation is comprehensive, with many pointers to related work, which help contextualize this paper's contributions.

Weaknesses

A weakness, which the authors do acknowledge in their work, that prevented me from giving a higher score is that there is no clear connection between the structure of the weights and generalization, as there exists in "standard"/one-hot classification. As a result, it is unclear how much insight can be derived from the current result. I would appreciate the authors' thoughts on this. Minor: The text is too dense in places, with the authors trying to include more details than what the space permits. I would suggest moving some of the discussion in Sections 6 and 7 to the Appendix to facilitate a smoother flow.

Questions

A minor suggestion: lines 32-34 appear to require rephrasing.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors thoroughly discuss the limitations of their work.

Reviewer 6D5d7/10 · confidence 2/52024-07-14

Summary

This paper studies the structural properties of the solutions selected by gradient-based optimizers among the many possible minimizers of the NTP objective, the central challenge being to discern the "implicit bias" of the optimizer towards particular solutions.

Strengths

- The paper is generally well written, and the notation is very clear. - The paper provides a a very interesting starting point for studying the solutions found by gradient descent in NTP settings

Weaknesses

While the paper provides a a very interesting starting point for studying the solutions found by gradient descent in NTP settings, it's not very clear whether margin maximization practically corresponds to any meaningful takeaway in language modeling.

Questions

Just the remark in the weaknesses.

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Authorsrebuttal2024-08-13

Dear Reviewer, We appreciate your support of our submission. We hope that our response regarding the potential takeaways of our study for language modeling sparked some interest. If you have any questions about these ideas, we are happy to elaborate. In any case, we look forward to hearing your feedback before the discussion period ends.

Reviewer YBdf5/10 · confidence 4/52024-07-17

Summary

This study investigates the structural properties of solutions chosen by gradient-based optimizers for next-token prediction (NTP), framing NTP as cross-entropy minimization across various contexts with sparse conditional probability distributions over a finite vocabulary. It focuses on the optimization bias of gradient descent (GD), characterizing how GD selects parameters that equate the logits’ differences of supported tokens to their log-odds.

Strengths

This study enables deriving the data-entropy lower bound in NTP for understanding the optimization and generalization properties of NTP models.

Weaknesses

The study's focus on linear models analyzing CE loss for NTP may limit its novelty and applicability, making its contributions to the field appear unclear compared to existing research.

Questions

Q.1 Please clarify the differences and advantages of your study compared to the following existing research. What new insights does this study provide, and why are they important? Specifically, while these existing studies highlight the critical role of attention in NTP, your study omits this aspect. Could you explain why it is still valid to disregard attention in your analysis? Mechanics of Next Token Prediction with Self-Attention Yingcong Li, Yixiao Huang, M. Emrullah Ildiz, Ankit Singh Rawat, Samet Oymak Max-Margin Token Selection in Attention Mechanism Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, Samet Oymak Transformers as Support Vector Machines Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, Samet Oymak Q.2 When considering next-token prediction (NTP) using sequence data, distinct contexts might differ by only a single character and are expected to be interrelated. Does the assumption of independence and identically distributed (i.i.d) data in Eq.(2) not pose a problem in this scenario?

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

The limitations of this study include its reliance on the simplicity of the analyzed model, unclear distinctions and advantages over existing research, and its omission of key aspects such as the properties of attention mechanisms.

Authorsrebuttal2024-08-13

Dear Reviewer, Did our response address your questions? Specifically, did it resolve your concern about how our work positions itself in relation to the papers you mentioned? We look forward to your feedback before the discussion period ends. Thank you for your time.

Reviewer YBdf2024-08-13

Thank you for your response, I will raise my score. On Q2. Let me confirm the motivation. Does your response mean that next token prediction can be separated from a specific model to understand its properties? If so, can such a separated theory provide a unified explanation of the phenomena of a specific model? In that case, it it better to experimentally demonstrate the findings obtained in this paper by using multiple different models. On Q3. When inputting a context in a sliding window, the sequence that corresponds to the previous answer is partially included in the context. Isn't this not i.i.d even if it is distinct?

Authorsrebuttal2024-08-13

Thank you for your response! And we appreciate your willingness to raise the score. **On Q2:** Yes, that’s right. The thesis we aim to communicate is that NTP as a training framework has intrinsic properties worth investigating, independent of the underlying architecture. This is not to say that architecture doesn’t matter, but significant insights can be gained by focusing on NTP itself. Our paper highlights this idea by examining a fixed embeddings setting, which allows us to isolate NTP properties. We provide experiments in the paper for this setting. More broadly, we envision using the proposed problem formulation and convergence framework to characterize the geometry of word and context representations in language models trained with NTP. Specifically, we aim to answer: how do the statistics of language data map to the geometry of representations during training? In our paper, we take an initial step by fixing the context representations and characterizing word embeddings. To also account for context representations, one idea is to assume that context embeddings are freely optimized along with word embeddings. We briefly discuss this in Lines 361-367. An advantage of this approach is that it circumvents the complexities of a specific architecture (e.g., transformers) and isolates NTP from it, provided the model is large enough with sufficient expressive power to generate such unconstrained embeddings. Our current framing of the NTP paradigm as a sparse soft-label classification over distinct contexts, combined with identifying the necessary and sufficient conditions for reaching the entropy lower bound, still holds in this setting. Thus, the technical question for future work is extending the convergence analysis to the bilinear setting where both $W$ and $h_j, j \in [m]$, are optimized. We conjecture that this will lead to an appropriate margin maximization program that establishes a direct mapping between language data statistics, as encoded in the sparsity patterns of the training data, and the geometry of representations. We find this inherently interesting, and such characterizations could also investigate linguistic regularities, such as word analogies (e.g., representations where king-man+woman=queen). Additionally, they might help enhance model interpretability and explainability, and provide a way to algorithmically (e.g., by modifying the CE loss) mitigate unfavorable imbalances in language data (e.g., rare words/contexts). **On Q3:** You are correct. The contexts are not i.i.d., but this does not create any problems in the current formulation. E.g., as mentioned: There is no i.i.d. assumption in Eq. (2) and such contexts are still viewed by the optimization as distinct contexts. What the implicit bias viewpoint of NTP suggests is that the relationship between learned context representations depends on the sparsity patterns of their next-tokens. That is, if two contexts (regardless of their degree of partial inclusion) are followed by a similar set of words (even if the probabilities of each word differ for the two contexts), their representations will tend to be more aligned, and vice versa.

Authorsrebuttal2024-08-12

Dear Reviewers, With all due respect to your reviewing load, we would appreciate your feedback on our responses. We hope that we have adequately addressed your questions and concerns, and we would be delighted to see this reflected in your evaluations. However, if you have any remaining questions, we are more than happy to respond while the discussion period is still ongoing. Thank you again for your time.

Authorsrebuttal2024-08-13

Dear Reviewer, Thank you again for your thoughtful comments and careful read. Did our responses answer your questions? We look forward to your feedback before the discussion period ends!

Reviewer kgdu2024-08-13

I apologise for my delayed response, which was due to force majeure. > We conjecture that the source of inductive bias for this structural grokking phenomenon is actually the NTP training, and that the framing of NTP in our paper might provide the right tools to study this since: [...] In case my input is useful, I agree with this conjecture. I would like to sincerely thank you for your comment. It taught me many things. In general, I believe that the paper is a solid contribution and should be accepted.

Authorsrebuttal2024-08-13

Thank you for your kind words and for sharing your thoughts on the conjecture—this is encouraging! We appreciate your time and support.

Authorsrebuttal2024-08-13

Thank you again for your feedback and support. We appreciate your decision to raise the score.

Reviewer YBdf2024-08-14

Thank you for your comments. I raised the score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC