Implicit Optimization Bias of Next-token Prediction in Linear Models

Thank you for your thoughtful review and the constructive questions raised. We appreciate the careful read. We hope that our responses below answer your questions.

Q1.  for the NTP-compatible and separable conditions to hold, one needs d > m.

Here are the key points to consider about the condition $d>m$ of Lemma 1:

1. First, this is only a sufficient condition, not a necessary one. In general, whether the conditions are satisfied depends on the geometry of embeddings and the sparsity pattern of the language training set.

2. Second, the result holds for linear models (fixed embeddings). It is an open question (Line 358-360) for future work to derive sufficient conditions for nonlinear models. The NTP-compatibility/separability conditions, which are themselves both necessary and sufficient for the NTP training loss to reach its lower bound, independent of linearity, provide the appropriate formulations for such an endeavor.

Specifically, in nonlinear models, the total number of parameters can be increased by augmenting the width or depth, rather than directly modifying the embedding dimension $d$ as in linear models. This is similar to known results for one-hot classification: For linear models, $d > m$ is a sufficient condition for separability of data (where $m$ is the number of samples and $d$ the ambient dimension), but this condition is relaxed when adding one hidden nonlinear MLP layer, where data can be separable provided $d > m/W$ (with $W$ being the width of the hidden layer). By increasing the width and making $dW$ the total number of parameters, the requirement on $d$ is relaxed. Deriving analogous results for the NTP setting is an interesting future direction. This remains non-trivial since an MLP layer shall be replaced by a sequence-to-sequence architecture (such as transformer), adding complexity to the analysis. Additionally, we expect that the result depends on the sparsity patterns of the support sets.

We appreciate the question and will elaborate on Lines 358-360 in the future work section towards relaxing the current limitation of $d>m$ due to linearity.

Q2. In lines 154-157, why does equation 4 constrain $W^p$ w.r.t. this subspace? Why is the solution $W^\star$ unique, assuming equation 4 has a solution?

Assuming Eq. (4) has a solution, say $W_1$, every other solution takes the form $W^p = W_1 + W_{\text{null}}$, where $W_{\text{null}}$ (the null-space component) is orthogonal to $(e_z - e_{z'})h_j^T : z \neq z' \in S_j, j \in [m].$ Thus, $W_{\text{null}} \in F^\perp$. We will reword Line 155 according to this.

Assuming it exists, the solution to Eq. (4) is unique when constrained to the data subspace $F$. To see this suppose $W_1$ and $W_2$ are both solutions in $F$. Since both are solutions, $W_1 - W_2$ is orthogonal to $F$. However, because $F$ is a subspace, $W_1 - W_2$ must belong to $F$. Therefore, $W_1 = W_2$, which is a contradiction. We will include this in the appendix for completeness.

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

Thank you for highlighting this. Upon review, we agree that our initial statement in line 148 requires more precise phrasing. It is straightforward to see that (3a) implies NTP-compatibility as described in Equation (4). However, the converse is not inherently true; NTP-compatibility alone does not ensure (3a) holds.

For (3a) to be valid, both NTP-separability conditions (Equations 6a and 6b) must also be satisfied in addition to Equation (4). Here’s a proof outline:

Assuming Equations (4), (6a), and (6b) hold for some $W^p$ and $W^d$, we define $W_\gamma=W^p+\gamma W^d$. Decompose the inverse of $S_z(W_\gamma h_j)$ as:

Using the properties from (4) and (6a) for in-support pairs:

leading to: $1+ \sum_{z'\neq z\in S_j }\exp\left( (e_z'-e_z)^TW_\gamma h_j \right) = 1+\sum_{z'\neq z\in S_j } p_{j,z'}/p_{j,z} = 1+(1-p_{j,z})/p_{j,z}=1/p_{j,z}.$

For out-of-support pairs, applying (6b) and considering the limit as $\gamma\rightarrow\infty$, we have:

Putting these together gives $1/S_z(W_\gamma h_j) \stackrel{\gamma\rightarrow \infty}{\longrightarrow} 1/p_{j,z}$ as desired.

We will include the above calculations within a formal proof of Proposition 1 in the appendix and we will remove the 'iff' statement from Line 148. Thanks again for catching this!

We are grateful for your encouraging feedback and for endorsing our paper.

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

Indeed, more work is needed to connect the structure of the weights to generalization. In fact, the question of generalization is at the center of what we envision as follow-up work. Although much of this is still in early stages, we would like to share some of our thinking around linking implicit-bias optimization results to generalization in the language setting below.

The first line of thinking relates to how to quantify generalization. One option is test loss (NTP loss with expectation over fresh contexts). In this case, we have reason to believe that it might be possible to extend techniques from standard classification that relate the generalization gap to training loss via algorithmic stability, e.g. [SK22]. The way the ‘structure of the weights’ would come into such a result is by ensuring that there exist weights, sufficiently close to initialization, such that training loss becomes $\epsilon$-close to its entropy lower bound. Since approximating the lower bound requires NTP-separability, we expect the max-margin weights to give the best candidate for selecting those weights. However, two challenges need to be resolved with this approach: the first is technical and involves extending the above-mentioned prior work to accommodate the imbalanced distribution of (distinct) contexts, which breaks the ‘identically distributed’ assumption. The second is more conceptual and relates to the second line of thinking discussed below: what is a ‘good’ model for the context embeddings and for the sparsity patterns? Beyond test loss, it is unclear to us what the right analogue of one-hot error in standard classification is. A possibility is to study specific tasks for which such a metric is clearly defined. A concrete example could be the bracket-tracing task Dyck, for which [Mur+23] recently demonstrated that NTP-trained transformers can generalize well, a phenomenon they call structural grokking, but only after extensive training. 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: (1) the model of NTP as sparse soft-label classification fits the synthetic language dataset Dyck, for which the experiments are reported in [Mur+23]. (2) It is experimentally suggested that structural grokking occurs only after extended training and only after the NTP loss saturates (i.e., when NTP-separability and compatibility conditions kick in). (3) The principle of margin-maximization has been very recently connected to the grokking phenomenon, although only in one-hot settings of modular arithmetic [Moh+23, Mor+24].

The second line of thinking involves identifying appropriate models for context embeddings and for sparsity patterns of the next-token distribution. Directly modeling context embeddings appears necessary if we are to hope that linear models can still provide (to some extent) insights on language generalization. Simply put the question becomes: what is the simplest (discrete) analogue of the Gaussian mixture model (features distributed normally around a mean vector for each class) which we often use to model image-classification data and produces insights on things like benign-overfitting and optimal loss functions even within linear models? Even if we decide we need to push beyond linear models, we still need to consider modeling the sparsity patterns of the next-token distributions, as this is critical in determining conditions for reaching the entropy lower bound and determining the structure of weights via the NTP max-margin program.

Finally, we hypothesize that considering the conceptual meaning that certain weights might carry could be a way forward. Concretely, decoder weights correspond to word embeddings and last-layer activations ($h_j$ in the paper) correspond to context embeddings. Identifying the structure of these weights translates to the structure of word and context embeddings (see Lines 361-7 for a possible way for future work to arrive at these). Although how to explicitly relate this structure to generalization is at the moment still an open question, at the very least, this direction could lead to insights on the functionality of large language models by understanding how NTP maps language sparsity patterns to word/context representations. Importantly, such characterizations could also investigate linguistic regularities such as word analogies (e.g., representations where king-man+woman=queen). This relates to generalization, since the ability of a model to produce representations for which such arithmetic operations hold has been linked experimentally to better downstream generalization [Mik+13].

[SK22] Stability vs Implicit Bias of Gradient Methods on Separable Data and Beyond

[Mur+23] Grokking of Hierarchical Structure in Vanilla Transformers

[Moh+23] Grokking modular arithmetic can be explained by margin maximization

[Mor+24] Feature emergence via margin maximization: case studies in algebraic tasks

[Mik+13] Efficient Estimation of Word Representations in Vector Space

Q: text too dense.

Well received. It makes sense to slightly shorten Sections 6 and 7 by moving some parts to the appendix.

Q: lines 32-34.

We appreciate the careful read. We will correct the typo by replacing "when" with "of" in Line 32.

Thank you for your time and for the positive feedback and score.

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

We agree that this paper is a starting point and indeed we are currently working on extending it further. Below we’d like to share some of our early-stage thinking to help convey our vision for this line of work and its relevance to language modeling.

A first idea is to use the problem formulation and convergence framework as a basis for characterizing the geometry of word and context representations of language models trained with NTP. In our paper, we take the initial step by fixing the context representations and characterizing word embeddings. An idea to also explicitly account for context representations is to assume that context embeddings are freely optimized together with word embeddings. We briefly discuss this in Lines 361-7. An advantage of this approach is that it circumvents the complexities of a specific architecture (e.g., transformers) and isolates NTP from it. However, this requires assuming a large enough model with sufficient expressive power to generate such unconstrained embeddings. Recall that our current framing of the NTP paradigm as a sparse soft-label classification over distinct contexts, together with identifying the necessary and sufficient conditions for reaching the entropy lower bound, continue to hold in this setting. Thus, the technical question that future work needs to address is extending the convergence analysis to the bilinear setting where both $W$ and $h_j, j\in[m]$ are optimized. Our conjecture, pending further analysis, is that the convergence result will be of similar flavor: when projected to an appropriate data subspace, $F$, the word and context embeddings converge in direction to the solution of an appropriately defined margin-maximization problem. But, unlike the linear case where $F$ depends both on the embeddings and on the sparsity patterns of the language data, in this new setting, it is natural to expect that $F$ would only depend on the latter.

It follows then that the (corresponding) margin maximization program would establish a direct mapping between language data statistics as encoded in the sparsity patterns of the training data and the geometry of representations. We deem this interesting in its own right. Moreover, down the road, we hope this might also help enhance model interpretability and explainability, as well as provide a way to algorithmically (e.g., by modifying the CE loss) mitigate unfavorable imbalances in language data (e.g., rare words/contexts). Indeed, characterizing the geometry of word representations has a long history in NLP literature. This dates back to at least the work [LV14], which studies word geometry of the Skip-gram with negative sampling objective in word2vec. This has been used to provide insights on ‘word analogies’ and inspire algorithms that modify the geometry of representations (e.g., making them more isotropic [Aro+16]) towards improving linguistic regularities [MBV17]. More recently, there are many works of similar flavor on modern transformer models trained with NTP, but to the best of our knowledge, most of these are heuristic/experimental, often resulting in contradictory claims. We see an opportunity for a theoretical framework to complement such work.

Additionally, we envision that the results can be used to gain insights into how language models generalize. To give a more concrete example, it could be possible to use the results to theoretically investigate the empirically observed phenomenon of ‘Grokking of Hierarchical structure’, i.e., the ability of models to infer hierarchical structures in language data when trained far beyond the point of saturating the training accuracy [Mur+23]. While [Mur+23] report this phenomenon in transformers, we conjecture that the source of this structural grokking phenomenon is actually the inductive bias of NTP training. Various reasons lead us to believe that the framing of NTP in our paper might provide tools to study this: (1) the model of NTP as sparse soft-label classification fits synthetic language datasets such as Dyck, for which the structural-grokking experiments are reported in [Mur+23]. (2) It is experimentally suggested that structural grokking occurs only after extended training and only after the NTP loss saturates (i.e., when NTP-separability and compatibility conditions kick in). (3) The principle of margin-maximization has been very recently connected to the grokking phenomenon, although only in one-hot settings of modular arithmetic [Moh+23,Mor+24].

Overall, while more work is needed to materialize the paper’s results into direct language modeling insights (e.g., the relation of representation geometry to language data statistics and how it impacts linguistic regularities, as well as generalization phenomena like structural grokking), we hope the above discussion (which we are happy to elaborate upon in the paper) convinces that it is a worthwhile endeavor.

[LV14] Neural Word Embedding as Implicit Matrix Factorization

[Aro+16] A latent variable model approach to psi-based embeddings

[MBV17] All-but-the-top: Simple and effective post processing for word representations

[Mur+23] Grokking of Hierarchical Structure in Vanilla Transformers

[Moh+23] Grokking modular arithmetic can be explained by margin maximization

[Mor+24] Feature emergence via margin maximization: case studies in algebraic tasks

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.

Thank you for your response and for shedding light on how future directions and practical implications. Although this paper does not lie squarely within my area of expertise, I think it is a good paper, and will be raising my score to an accept.

Thank you for your review.

Below, we clarify the distinctions from the references you mention and explain why our problem setting differs from studying self-attention/transformers, focusing instead on the NTP paradigm. While we have detailed these discussions in the submission, we repeat them here for your convenience, respecting the reviewing load.

We hope these clarifications lead to a re-evaluation of your score!

Q1: Comparison to works on implicit bias in transformers:

We have discussed these references in detail in the submission. For your convenience, we summarize them below.

[44] Mechanics of Next Token Prediction with Self-Attention - Li et al.

[79] Max-Margin Token Selection in Attention Mechanism - Tarzanagh et al.

[80] Transformers as Support Vector Machines - Tarzanagh et al.

In Lines 327-328, we explain that our convergence results relate to a conjecture in [79] about implicit optimization bias in transformers, though our findings are different in nature. This is further elaborated in Appendix B, Lines 648-667. The discussion there is self-contained and directly answers your questions. We reproduce it below:

“As already mentioned in Sec. 6, our work is closely related to [79], where the authors investigate the implicit bias of self-attention in transformers. The insight put forth in the prequel [80] is that softmax attention induces implicit-bias behaviors that bear similarities to vanilla implicit bias of one-hot prediction. Concretely, [79] studies GD optimization of one-layer self-attention with fixed decoder and one-hot binary classification. They show that, in the limit, GD finds attention weights that converge in direction to the solution of an SVM problem that separates optimal tokens from non-optimal ones. Their non-convex setting introduces locally optimal SVM directions to which GD may converge depending on initialization. Different to them, the NTP setting that we study involves predictions over multiple categories and is not one-hot. Also, while they fix the decoder, here, we fix the embeddings. In these respects their results are rather different. More similarities arise when [79] replace the linear decoder with a MLP, which they note can induce multiple optimal tokens per sequence. This leads them to formulate a more general token-separating SVM program, which similar to ours confines the separation on a certain data subspace. However, the operational nature of the programs remains different as theirs optimizes attention weights and separates tokens within a sequence, while ours optimizes decoder weights and separates context embeddings based on their respective support sets. More importantly, while [79] only conjectures the convergence of GD to their general SVM program, we leverage convexity in our setting to prove an analogous statement rigorously. Eventually, as we move lower in our top-down approach and consider architecture-specific embeddings generated by attention, we anticipate to see integration of our ideas with theirs.”

Additionally, Lines 677-684 compare our work to Li et al. [44]:

“Upon completing this paper, we became aware of independent contemporaneous research by Li et al. [44] that also examines the implicit bias of self-attention with a fixed linear decoder in next-token prediction scenarios. Unlike our study which utilizes the widely adopted CE loss, their approach is based on log-loss, which renders the training loss convex, a similarity shared with our model despite the inclusion of self-attention. Both our results and those of Li et al. substantiate the conjecture posited by Tarzanagh et al. [79], albeit in very distinct settings. Notably, contrary to both [79] and [44], we unveil the optimization intricacies of the NTP paradigm, even within the simplest linear settings.”

We believe these detailed comments clarify the distinctions to the above references. If you have further questions, please kindly let us know and we are happy to elaborate.

Q2: Could you explain why it is still valid to disregard attention in your analysis?

Our key message is that self-attention and next token prediction (NTP) are distinct. NTP, which involves predicting the next token given preceding tokens using cross-entropy loss, is used across transformer-based models, state-space models, and LSTMs (see footnote 3). Thus, studying NTP separately from transformers/self-attention is valid.

By isolating NTP, we highlight essential aspects of the problem often overlooked (e.g., [47,50]). For example, modeling NTP in language settings as soft-label classification over sparse probabilistic labels and deriving conditions for reaching the entropy lower bound. These foundational aspects remain valid regardless of the embedding architecture (see Lines 192-195).

Our convergence results assume fixed embeddings and training only classifier weights (aka word embeddings). In Section 7, we suggest two avenues to extend this analysis. First, studying architecture-specific embeddings, including those generated by self-attention (Lines 352-357). Second, exploring architecture-agnostic embeddings via the unconstrained features model (Lines 361-367). This approach leads to a model where both context and word embeddings are freely optimized, helping to understand the geometry of context and work embeddings once the NTP loss reaches the lower bound.

Q.3 distinct contexts might differ by only a single character and are expected to be interrelated. Problem with independence and identically distributed (i.i.d) data assumption in Eq. (2)?

There is no i.i.d. assumption in Eq. (2). Rather, Eq. (2) is a reformulation of the empirical NTP loss expressed in terms of distinct contexts, with the summation over these contexts. If we have two contexts that differ by a single character, they are still considered distinct. Thus, Eq. (2) applies to them as it corresponds to two different distinct embeddings, say $h_{j_1}$, $h_{j_2}$ with $j_1,j_2\in[n]$.

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.