Implicit Optimization Bias of Next-token Prediction in Linear Models

Authors: christos  thrampoulidis
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 00:10:09

Abstract: 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 \emph{distinct} contexts, each tied with a \emph{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.

Main paper: 
 introduction: 
 Next-token prediction (NTP) has emerged as the go-to paradigm in training modern language models, revolutionizing various applications such as machine translation, text-summarization, and language generation [66] . In NTP, models are trained to predict the most probable token given a sequence of preceding tokens, commonly referred to as the context. Concretely, the objective is to learn a mapping from the input context to the probability distribution over the (finite) vocabulary of possible tokens, enabling the model to generate a token that is contextually appropriate [9, 8] . Recently, the NTP paradigm has witnessed remarkable empirical success through its utilization on large-scale deep-learning architectures trained on vast corpora of data [66, 67, 86] , leading to unprecedented advances in the field, and the swift integration of these advanced language models into society [62] . Concurrently, researchers have raised critical concerns about robustness, interpretability, and fairness-bias issues arising from our limited understanding of the fundamental operational principles of these models [10, 6] . Despite progress, a comprehensive theory that elucidates the fundamentals of modern language models-including key components like the NTP paradigm and transformer architecture, particularly in terms of optimization and generalization principles-is still lacking.We initiate an investigation when implicit optimization biases in training language models under the NTP paradigm, particularly in overparameterized regimes where the empirical-loss reaches its lower bound and there is many possible minimizers. To formalize the NTP paradigm, consider 38th Conference on Neural Information Processing Systems (NeurIPS 2024).autoregressive model q θ parameterized by θ trained to predict the next-token on sequences of length T using the cross-entropy (CE) loss:EQUATIONHere, sequences z = (z 1 , . . . , z T ) consist of tokens z t from a finite vocabulary V = {1, . . . , V } and Ê is expectation over training set T n of n such sequences sampled from some underlying true distribution over sequences. Typically, the model q θ outputs probability of the next token computed via softmax applied on output logits, which are computed by projecting d-dimensional embeddings h θ ′ to the V -dimensional space with a trainable linear decoder W ∈ R V ×d . Formally, 1q θ (z t | z 1 , . . . , z t-1 ) = S zt (W h θ ′ (z 1 , . . . , z t-1 )) = 1 1 + ∑ z ′ ∈V z ′ ≠zt exp ((e z ′ -e zt ) ⊺ W h θ ′ (z 1 , . . . , z t-1 )).The CE loss is then minimized over θ = (W , θ ′ ) using gradient-based methods, e.g. (S)GD, Adam.We pose the question: Given training set T n , what are the structural properties of the weights θ found by minimizing the NTP objective with gradient-based optimizers? As in prior research in one-hot supervised classification2 (e.g. [97, 7, 76, 34] ), we specifically target this question in an overparameterized setting, where the NTP objective (1) may have an infinite number of solutions, representing an infinite number of models θ that minimize the training loss. The central challenge is to discern the particular solution the optimizer is inherently biased towards. Since this 'bias' is not explicitly introduced through regularization but is instead ingrained in the training objective and algorithmic structure, it is termed 'implicit bias' [61] . The exploration of implicit bias has a long history in the traditional supervised one-hot classification (see Related Work in Sec. 6). In this traditional scenario, the training set comprises feature-label pairs (x, y), where x ∈ R p is a continuous feature, and y represents its unique label. The optimization process minimizes the following training objective (over W , θ ′ ): Ê(x,y) [-log (S y (W h θ ′ (x)))] .At first glance, excluding the sequential format of Eq. ( 1), the NTP training scenario might seem identical to traditional one-hot prediction: both aim to minimize the same CE loss across models that parameterize probabilities using the softmax of logits. Consider predicting the next token over fixed-length sequences, say sequences of length t -1, via optimizing: Êz [-log (S zt (W h θ (z 1 , . . . , z t-1 )))] .The context here acts as the feature, and the next token as the label. Recent works [49, 52] draw on such apparent similarities to the traditional one-hot classification paradigm to extrapolate known results from the latter to the NTP setting. However, this comparison overlooks a fundamental, yet critical difference in the nature of the training data that distinguishes these two paradigms (even when the sequential format of Eq. ( 1) is disregarded): In the traditional setting, each feature (e.g., image) is assigned a single label (e.g., image category). In contrast, in the NTP setting, contexts z 1 , . . . , z t-1 of finite length sampled from finite vocabularies are naturally repeated in a (vast) training set, potentially multiple times, each time followed by different tokens z t [73] . Consequently, the NTP paradigm involves training over m ≤ n distinct (non-repetitive) contexts, each followed by a multitude of possible next tokens, appearing at varying frequencies. For instance, the context "She is excellent at her role as a" may be followed by next tokens such as "doctor," "lawyer," "reviewer," or "mother," each with different frequencies. Importantly, certain vocabulary tokens may not appear after a given context; e.g., in the above example, tokens like "run," "and," etc., will not follow.Model. We study NTP training over a finite vocabulary employing the following model. Given a large training set of n total sequences, we identify m ≤ n distinct contexts. Each distinct context j ∈ [m] is linked to a V -dimensional empirical probability vector pj , which encodes the frequency with which each vocabulary token follows the context throughout its occurrences in the training set. Crucially, the probability vectors pj are sparse, i.e., the support set S j of pj satisfies |S j | ≪ |V| = V . In an extreme where |S j | = 1, ∀j ∈ [m], the probability vector pj becomes one-hot, leading to a scenario reminiscent of the traditional classification setting described earlier. However, such an extreme is essentially improbable in practical language modeling [73] . With this framing, the NTP paradigm is also related to supervised vision classification with soft labels, which advocates for training models on datasets where each example is associated with a vector of soft labels (rather than a one-hot vector), such as by averaging multiple annotators' hard labels [65] , knowledge distillation [32] or label smoothing [79] . With this connection, our analysis can also be interpreted (more broadly) as investigating the implicit bias of sparse soft-label classification.Formulation. Recognizing the differences between NTP and one-hot classification, we study the question of implicit optimization bias within the NTP setting. To facilitate this, we utilize the model outlined in the previous paragraph and detailed in Sec. 2. For concreteness, our analysis adopts a 'top-down' approach, training only the decoding (also referred to as word-embedding) matrix W ∈ R V ×d while keeping context-embeddings fixed. This approach mirrors foundational studies on implicit optimization bias in one-hot classification [76, 34] , which first focused on linear models. It allows exploring the complexities of the NTP training objective, distinct from the embedding architecture 3 , and while it renders the logits linear and the objective convex, it still poses a technical challenge in terms of determining parameter convergence [76, 34, 37, 60, 38] .Conditions for reaching entropy. In Sec. 3, we identify the necessary and sufficient conditions for the logits of the trained model to enable the CE loss to approach its lower bound, the empirical conditional entropy. We introduce two conditions: NTP H -compatibility and NTP-separability, which impose constraints on mutually orthogonal subspaces that are determined by the sparsity patterns of distinct contexts within the dataset. These conditions determine the necessary and sufficient overparameterization a model needs to achieve the empirical entropy lower bound during training.Margin in NTP setting. Motivated by the NTP-separability condition, we introduce a margin concept for NTP in Sec. 4, which extends the classical definition of margin used in one-hot supervised classification [88] . We further establish the relevance of this new margin notion for optimization by demonstrating that a decoder maximizing the NTP-margin, denoted as W mm , guides the directional convergence of the ridge-regularized CE minimizer, Ŵλ , as the regularization parameter λ → 0.Implicit bias of GD. We establish that W mm also determines the implicit bias of gradient descent (GD) iterates in Sec. 5. Specifically, in the limit of iterations k → ∞, the GD iterates grow undoubtedly in norm and converge to a finite W ⋆ within a data subspace F, while simultaneously aligning with W mm in the complementary subspace F ⊥ . The finite component W ⋆ ∈ F solves a system of linear equations associated with the NTP H -compatibility condition. Finally, we numerically verify these findings and discuss related and future work in Secs. 6 and 7.Additional experiments, further related work and detailed proofs are in the appendix. 
setup: 
 Let vocabulary V = [V ] ∶= {1, . . . , V } represent a set of V tokens (e.g. words) and z 1∶t = (z 1 , . . . , z t ) denote sequence of t tokens z t ∈ V. To simplify presentation, we focus on predicting the T -th token z T given contexts z <T ∶= z 1∶T -1 of fixed length, and we further let x = z <t denote the context and z denote the last token. See App. C for straightforward extension to the sequential format of Eq. ( 1).We assume access to a training set consisting of n sequencesT n ∶= {(x i , z i )} i∈[n] , with x i ∈ X ∶= V T -1 and z i ∈ V. Let h ∶ X → R dan embedding map that maps contexts (i.e., sequences of T -1 tokens) to d-dimensional embeddings. The map h can be parameterized (e.g. by a transformer [90] or an LSTM [5] ), but this paper assumes that it is fixed. The next-token is predicted via a linear modelf W ∶ X → R V parameterized by decoding matrix W ∈ R V ×d , such that f W (x) = W h(x).When the model output passes through a softmax, it defines the model's probability mass function for the next-token prediction, given as qW (⋅|x) = S(f W (x)), where S(⋅) ∶ R V → ∆ V -1 is the softmax and ∆ V -1 is the V -dimensional simplex. The decoder is trained by minimizing the empirical CE lossCE(W ) ∶= 1 n ∑ i∈[n] -log (q W (z i |x i )). Distinct sequences and next-token distributions. Given dataset T n we denote x1 , . . . , xm the m ≤ n distinct contexts among the (large number of) total n contexts x 1 , . . . , x n within T n . Let πj be the empirical probability of distinct context xj . That is, 1 ≤ n ⋅ πj ≤ n is the number of contexts x i that equal xj . Furthermore, for each distinct context xj , j ∈ [m] let pj ∈ ∆ V -1 denote the probability vector of conditional next-token distribution, i.e., pj,z ∶= p (z| xj ) , z ∈ V, j ∈ [m]. In other words, n ⋅ πj ⋅ pj,z is the number of occurences of token z as a follow-up to context xj . Finally, we denote the support set and size of the support set of these conditional distributions as S j ∶= {z ∈ V | pj,z > 0} and S j ∶= |S j |. Tokens z ∈ S j and v ∉ S j are referred to as 'in-support' and 'out-of-support' respectively. Onwards, we implicitly assume that "not all tokens are likely after every context," i.e. ∃j ∈ [m] such that S j < V . This mild assumption is naturally satisfied in language modeling under rich enough vocabulary. With this notation, 4 we can express the NTP training loss asEQUATIONwhere, in the last line we defined the shorthand hj = h( xj ). Similarly, we leth i = h(x i ), i ∈ [n].With some abuse of notation, we then obtain the following equivalent descriptions of the training set{(x i , z i )} i∈[n] =∶ T n ≡ T m ∶= {( hj , πj , pj,z∈V )} j∈[m]that emphasizes distinct contexts and their respsective sparse next-token probability distributions.Entropy. The empirical T -gram entropy (referred to hereafter as entropy for simplicity) of the dataset is [74, 73] : 3 When can the NTP-loss reach the entropy lower-bound?H T ∶= H ∶= Ê(x,z)∼Tn [-log (p(z|x))] = -∑ j∈[m] ∑ z∈Sj πj pj,z log (p j,zThe first question we ask is: Under what conditions on the training data can the CE loss reach its entropy lower-bound? By the entropy lower-bound, CE(W ) = H ⇔ KL ( p || qW ) = 0 iff for all j ∈ [m] and all z ∈ V: qW (z| xj ) = pj,z . Equivalently, for all j ∈ [m]:EQUATIONEQUATIONBeginning with (3a), this requires 5 the training data to satisfy the NTP H -compatibility condition defined below. Definition 1 (NTP H -compatible). Let e v denote the v-th standard basis vector in R V . We say that training data T m are NTP-entropy-compatible if there exists V × d matrix W p satisfying:EQUATIONWe comment on the independence of the constraints: Fix any j ∈ [m]. Then, the set of constraints (as expressed in Eq. ( 4)) for all z ≠ z ′ ∈ S j (yielding ( Sj 2 ) constraints in total) is equivalent to the set of the same constraints for any anchor z j ∈ S j and z ′ ≠ z j ∈ S j , i.e., an effective total of S j -1 linearly independent constraints for each j ∈ [m]. Additionally, note that the system of equations in Eq. ( 4) constrains W p with respect to a specific subspace of V × d matrices:EQUATIONthat is defined in terms of context embeddings and their respective support sets. Assuming Eqs. ( 4) have a solution, we denote the unique solution within the subspace F as W ⋆ ∈ F for later reference 6 .Next, we examine Eq. (3b), which requires softmax outputs be zero for tokens that never occur following a fixed context throughout the dataset. Due to the strict positivity of softmax, the constraint is never satisfied for finite W . Thus, for all finite W , there exists a gap between the cross-entropy loss and its lower bound, i.e., CE(W ) > H. Yet, it is possible to approach entropy as the norm of the weights W grows, provided that weights move in the appropriate direction formalized below.Definition 2 (NTP-separable). We say that training data T m are NTP-separable if there exists V × d matrix W d satisfying the following:∀j ∈ [m], z ≠ z ′ ∈ S j ∶ (e z -e z ′ ) ⊺ W d hj = 0 (6a) ∀j ∈ [m], z ∈ S j , v ∉ S j ∶ (e z -e v ) ⊺ W d hj ≥ 1 . (6b)As before, it is easy to see that the constraints in ( 6) can be equivalently expressed by enforcing (6a) and (6b) for an anchor z j ∈ S j and all z ′ ∈ S j / {z j } and v ∉ S j , respectively. Consequently, there exist effectively V -1 linearly independent constraints per context j ∈ [m].We now discuss the interpretation of these constraints. The subspace constraints in Eq. (6a) project W d onto the subspace F ⊥ , which is the orthogonal complement of the subspace F defined in (5) . This leaves the softmax probabilities of possible next tokens (in set S j ) intact, and fully determined by W p as per the NTP H -compatibility condition. Formally, W p + W d continues satisfying (4) .Moving on the halfspace constraints in (6b), we can interpret these using Kesler's construction as enforcing linear separability in the space R V ×d [30] : Each d-dimensional context embedding hj is mapped to S j (V -S j ) higher-dimensional points (e z -e v ) h⊺ j , z ∈ S j , v ∉ S j . These points collectively for all j ∈ [m] must lie within the interior of the same halfspace induced by the hyperplane ⟨W d , ⋅⟩ = 0. Refer to Fig. 1 (Left) and its caption for an alternative interpretation of the rows of W mm as word-embeddings in R d (illustration in d = 2).The impact of NTP-separability on the softmax probabilities can be understood algebraically by considering W γ ∶= γW d and v ∉ S j . We have:EQUATIONwhere the first inequality removes non-negative exponential terms and the second one follows from (6b). The upper bound above approaches 0 as γ → ∞, thus (3b) holds asymptotically in γ.Taking into account the observations made above, the satisfaction of both conditions guarantees convergence of the cross-entropy loss CE to H. This is formalized in the proposition below. Proposition 1. Assume training data T m is NTP H -compatible and NTP-separable, with the respective matrices W p and W d satisfying conditions (4) and (6) . While all finite W satisfy CE(W ) > H, it holds forW γ = W p + γ ⋅ W d that CE(W γ ) γ→+∞ → H.Hence, CE approaches its lower-bound in the limit of a direction W d ∶= W d /∥W d ∥ and offset W p satisfying the constraints of NTP-separability and NTP-compatibility, respectively. In other words, parameter weights W that minimize the CE loss consist of two components: a finite projection W F ∶= P F (W ) = W ⋆ onto the data subspace F and an infinite-norm component onto the orthogonal complement F ⊥ in the direction of W d .Finally, we note that while Defns. 1 and 2 are stated for linear models, they naturally extend to a more general formulation for nonlinear models. Specifically, consider NTP-separability (similar for NTP-compatibility): the general conditions require that both the decoder weights W and model weights θ, which parameterize the embeddings hj = h θ ( xj ), must satisfy Eq. ( 6) simultaneously. 
the role of overparameterization: 
 We show that overparameterization provides a sufficient condition for the solvability of Eqs. ( 4) and (6) . Start with the halfspace constraints in Eq. ( 4) for NTP H -compatibility. These can be compactly expressed as E j,zj W p hj = a j,z , where E j,zj ∈ R (Sj -1)×V has rows e zj -e z ′ and a j,zj ∈ R (Sj -1) has entries log (p j,zj /p j,z ′ ) for some anchor z j ∈ S j . Now, since the rows of E j,zj are linearly independent, the question becomes equivalently that of determining when W p [ h1 , . . . , hm ] = [E † 1,z1 a 1,z1 , . . . , E † m,zm a m,zm ] has a solution. This is always the case when d > m and the d × membedding matrix H = [ h1 , . . . , hm ] is full rank (m). Then, there exists W p such that condition (4) holds. In fact, H⊺ has a nullspace, implying the existence of an infinite number of solutions to (4). These solutions take the form W p = W ⋆ + W p ⊥ , where W ⋆ ∈ F is the unique solution onto the subspace, and W p ⊥ ∈ F ⊥ . In contrast to (4), the constraints in (6) involve linear inequalities. However, a sufficient proxy for feasibility in this case is that the corresponding system of equations (instead of inequalities) has a solution. By following the exact same argument as before, we arrive at the same sufficient conditions for the existence of a solution W d . We summarize these findings. Lemma 1 (Overparameterization implies NTP-separability). Assume overparameterization d > m and full-rank embedding matrix H ∈ R d×m . Then, there exists an infinite number of solutions W p and W d that satisfy conditions (4) and (6), respectively. Thus, d > m, 7 which also generically favors full-rankness of the embedding matrix [92] , implies both NTP H -compatibility and NTP-separability. Combined with Prop. 1, it also implies that there are infinitely many possible directions W d along which the NTP loss approaches H, motivating the implicit-bias question: For a specific iterative algorithm aimed at minimizing the NTP loss, which direction does it prefer? We will address this question in the remainder of the paper. Remark 1. In the trivial case where S j = 1, ∀j ∈ [m] (one-hot classification), the entropy lower bound is zero and is attained iff the data is linearly separable. Indeed, F reduces to the empty set, and NTP-separability simplifies to traditional multiclass separability. For binary classification, [20] showed that d/m > 1/2 is sufficient and necessary for data in general position to be linearly separable. More recently, several works have extended this analysis to structured (random) data, including [12, 71, 57, 54] . The exact threshold in corresponding mutliclass settings is more intricate, but [19, 81, 11] have made progress in this direction. An interesting question is determining exact thresholds for NTP-separability, which would improve upon the sufficient condition of Lemma 1. 
regularization path: 
 This section investigates the implicit bias of NTP by examining the minimization of CE loss through iterates defined as follows for an increasing sequence of positive regularization parameters B:EQUATIONThis involves minimizing a strictly convex function in a bounded domain; thus, ŴB is unique. This section's main result characterizes the limit of ŴB as B → ∞ under NTP-separability/compatibility. Before that, we first define the next-token prediction support-vector machines (SVM) problem. Definition 3 (NTP-SVM). Given NTP-separable training set T m , NTP-SVM solves the following:W mm ∶= arg min W ∥W ∥ subj. to W ∈ R V ×d satisfying (6a) and (6b). (NTP-SVM)This is a strongly convex quadratic program with mV -∑ j∈[m] S j linear inequality and ∑ j∈[m] S j -m linear equality constraints. Its solution can be also defined as the classifier that maximizes margin between in and out-of -support tokens while being constrained on the orthogonal compelemnt F ⊥ :W mm = arg max ∥W ∥=1,W ∈F ⊥ min j∈[m],z∈Sj ,v∉Sj (e z -e v ) ⊺ W hj .It turns out this direction determines the preferred limiting direction of the regularization path. Theorem 1 (Implicit bias of the regularization-path). Assume training data T m is NTP H -compatible and NTP-separable. Let ŴB be defined as in (8) . Then, it holds thatlim B→∞ ⟨ ŴB ∥ ŴB ∥ , W mm ∥W mm ∥ ⟩ = 1 .The proof sketch below illustrates how the NTP-separability/compatibility assumptions influence the outcome and why the regularization path induces an optimization bias toward the NTP-SVM direction.Complementing Thm. 1, we also show (see Lemma 4 in the appendix) that lim B→∞ P F (W B ) = W ⋆ . These together provide a complete characterization of the implicit optimization bias of (8) .Proof sketch (App. E.2 for details). We first show ŴB is on the boundary:∥ ŴB ∥ = B. If not, then ⟨∇ CE( ŴB ), W mm ⟩ = 0. But, few algebraic manipulations show ⟨-∇ CE( ŴB ), W mm ⟩ equals ∑ j∈[m] πj ∑ z∈Sj pj,z ( ∑ z ′ ∈Sj ,z ′ ≠z s j,z ′ (e z -e z ′ ) ⊺ W mm hj + ∑ v∉Sj s j,v (e z -e v ) ⊺ W mm hj ),where we denotes j,v ∶= S v ( ŴB hj ) > 0, v ∈ V, j ∈ [m].The first term in the parenthesis is zero by (6a), while the second term is strictly positive by (6b), leading to contradiction. Now, consider a 'genie' pointW ⋆ B = W ⋆ +R(B)⋅W mm , where W ⋆ ∈ F satisfies (4), and R = R(B) is chosen such that ∥W ⋆ B ∥ = B.We will show that W ⋆ B attains a small CE loss as B (hence, R) grows. To do this, denote for convenience the logitsℓ ⋆ j,v ∶= e ⊺ v W ⋆ hj and ℓ mm j,v ∶= e ⊺ v W mm hj for all for v ∈ V, j ∈ [m], and note that e ⊺ v W ⋆ B hj = ℓ ⋆ j,v + R ℓ mm j,v. By using ( 4) and (6a):EQUATIONMoreover, using (6b) and definingC ∶= V e ∥W ⋆ ∥M for M ∶= √ 2 ⋅ max j∈[m] ∥ hj ∥, gives: ∑ v∉Sj e -(ℓ ⋆ j,z +Rℓ mm j,z -ℓ ⋆ j,v -Rℓ mm j,v ) ≤ e -R ∑ v∉Sj e -(ℓ ⋆ j,z -ℓ ⋆ j,v ) ≤ C e -R .Combining the above within Eq. ( 2), using log(1 + x) ≤ x, x > 0 and the fact that πj , pj,z are probabilities, yields:EQUATIONNext, towards contradiction, we will show that if ŴB is not in the direction of W mm , then it incurs a loss that is larger than CE(W ⋆ B ). The trick here is to bound the KL divergence term:EQUATIONwhere we denote logits ℓ j,v ∶= e ⊺ v ŴB hj . Assume there exists ϵ > 0 and arbitrarily large B satisfying:EQUATION) Define Ŵ = ( ŴB -W ⋆ )/R ′ (B), where R ′ = R ′ (B) > 0 can be chosen so that ∥ Ŵ ∥ = ∥W mm ∥.Further choose B large enough so that Eq. ( 11) guarantees ∥ Ŵ -W mm ∥ ≥ ϵ ′ , for some ϵ ′ > 0.Since W mm is the unique minimizer of (NTP-SVM) and ∥ Ŵ ∥ = ∥W mm ∥, there exists δ ∈ (0, 1) and j ∈ [m] such that at least one of the following is true:(i) ∃z and z ′ ≠ z ∈ S j such that |(e z -e z ′ ) ⊺ Ŵ hj | ≥ δ (ii) ∃z ∈ S j , v ∉ S j such that (e z -e v ) ⊺ Ŵ hj ≤ 1 -δ.Case (i): Without loss of generality (e z -e z ′ ) ⊺ Ŵ hj ≤ -δ (otherwise, flip z, z ′ ). Thus, ignoring all but the (j, z, z ′ )-term in (10) and usingℓ j,z ′ -ℓ j,z ≥ R ′ δ + log ( pj,z ′ pj,z ) gives CE( ŴB ) -H ≥ πj pj,z log (p j,z e (ℓ j,z ′ -ℓj,z) ) ≥ 1 n log ( e R ′ δ n ).Comparing this to (9) for large enough B gives that CE( ŴB ) > CE(W ⋆ B ), a contradiction. Case (ii): We can assume Ŵ ∈ F ⊥ , since otherwise we are in Case (i). Now, again ignoring all but the (j, z) term in the CE loss for which the assumption holds for some v ∉ S j , we findCE( ŴB ) -H ≥ πj pj,z log (p j,z ( ∑ z ′ ∈Sj e (ℓ j,z ′ -ℓj,z)+ e (ℓj,v-ℓj,z) )).Using P F ( ŴB ) = W ⋆ and (4) yields ∑ z ′ ∈Sj e (ℓ j,z ′ -ℓj,z) = 1 pj,z . Moreover, by assumption of Case (ii):e ℓj,v-ℓj,z ≥ e -R ′ (1-δ) e ℓ ⋆ j,v -ℓ ⋆ j,z ≥ c ′ e -R ′ (1-δ) , for c ′ ∶= e -∥W ⋆ ∥M .Putting together yields:CE( ŴB ) -H ≥ πj pj,z log (1 + pj,z c ′ e -R ′ (1-δ) ) ≥ c ′ e -R ′ (1-δ) /2n 2 ,where the second inequality uses log(1 + x) ≥ x 1+x , x > 0. Compare this with (9): For large enough B, since R, R ′ grow at the same rate, it holds c ′ 2n 2 e -R ′ (1-δ) > Ce -R . Thus, CE( ŴB ) > CE(W ⋆ B ), a contradiction. In either case, we arrive at a contradiction, which completes the proof.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: See Sec. A. 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. , 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. 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). 
gradient descent: 
 This section studies the implicit bias of GD. Denote the GD iterates at time k byW k = W k-1 - η∇ CE (W k-1) for arbitrary initial point W 0 and constant step-size η > 0 small enough to guarantee descent. The first observation is that the norm of the GD iterates increases with iterations. Lemma 2 (Norm growth). If training data are NTP H -compatible and NTP-separable, thenlim k→∞ CE(W k ) = H and lim k→∞ ∥W k ∥ = ∞.This is intuitive because the CE loss is convex in W (thus, GD approaches the objective's infimum H), and, in view of Proposition 1, the CE loss at all finite W is bounded away from H. The relevant question then becomes that of determining the limit of the direction of the GD iterates. Theorem 2 (Implicit bias of GD). Assume NTP H -compatible and NTP-separable training data T m . Then, it holds thatlim k→∞ ⟨ W k ∥W k ∥ , W mm ∥W mm ∥ ⟩ = 1 . Moreover, lim k→∞ P F (W k ) = W ⋆ .The theorem establishes8 that in the limit of iterations:W k ≈ W ⋆ + ∥P ⊥ (W k )∥W mm, which is analogous to the result we obtained previously for the regularization path. Although its proof is more involved compared to the proof of Thm. 1, the proof of its main ingredient (Lem. 5 in the appendix) is conceptually similar: It involves comparing the loss CE(W k ) for large iterations k to the loss evaluated at a "genie" point that is chosen so that: (i) On the subspace F, it agrees with W k . This is because it is easy to show that P F (W k ) converges to W ⋆ by standard gradient descent analysis for convex functions; (ii) On the orthogonal subspace F ⊥ , it follows the optimal (with respect to accelerating loss decrease) max-margin direction W mm ∈ F ⊥ . To establish the loss comparison, the ideas is to compare the values of the adjusted loss CE ⊥ (W ) ∶= CE(W ) -CE (P F (W )).We validate our analysis with experiments on synthetic data in App. A. For illustration, Fig. 1 shows a 2D setting with m = 3 distinct contexts, each followed by S j = 3 tokens/words out of total V = 5 words in the vocabulary. The left subfigure illustrates: (i) In black markers, the context-embedding geometry along with the associated support sets for each context A, B, and C. (ii) In colored markers, the geometry of word-embeddings, that is the max-NTP-margin vectors (W mm ) ⊺ e v , v ∈ [5] , to which GD directionally converges. See caption for interpretation and Fig. 2  
related work: 
 We build on the literature on implicit optimization bias of CE loss in one-hot supervised classification. [76] show that for linear models and linearly-separable data, GD converges in direction to the maxmargin classifier. This result strengthens [68] that showed the regularization path of CE minimization converges to the same limit. Closer to us, [34, 37] extend the analysis to encompass general binary data as follows: the data are linearly separable only on a certain subspace, and they show that GD converges, in direction, towards the max-margin classifier confined within that subspace. On the orthogonal subspace, it converges to a finite point. While operationally similar, Thms. 1, 2 cannot ⊺ v ∶= e ⊺ v W mm , v ∈ [5]found by GD. Interpreting decoder vectors as word embeddings leads to intuitive findings on their geometry learned by NTP training. E.g., word embedding w 3 (almost) aligns with context-embedding A and the normal hyperplane it defines separates A from B and C, since word 3 only appears after context A. The rest of the words follow two contexts each and their word-representation naturally belongs to the cone defined by the embeddings of those respective contexts. The wider the cone, the larger the magnitude of the word embedding to compensate for the large angle between context-representations that share the same next-word. Note that geometry of depicted word embeddings only depends on support sets, but the conditional probabilities define another set of word representations on an orthogonal (matrix) subspace; see text for details and vis. Right: Upper/lower graphs confirm the predictions of Lemma 2 and of Theorem 2, respectively. be directly derived from theirs since our setting is neither binary nor one-hot. Nevertheless, our proofs extend the foundational work of [68, 34, 37] , akin to numerous other studies that explore extensions to nonlinear architectures [50, 35, 28, 29, 83, 89] , and to stochastic and adaptive algorithms [60, 64, 21, 47, 77, 3, 14, 2] . The implicit bias viewpoint has also created opportunities to study generalization in overparameterized settings. [31, 4, 57, 22] build a two-stage approach initially leveraging implicit bias to simplify the complexities of optimization before addressing generalization. This narrows the generalization question to the properties of the corresponding max-margin classifier [58, 13, 43, 78, 23, 100, 72, 94] . The same strategy has also been adopted to study model robustness to adversarial perturbations [33, 80, 16] , out-of-distribution data [87] , and imbalances [69, 15, 42] . Our results motivate such extensions in the richer NTP setting.Recent work [49] also studies forms of implicit bias for language models trained to reach the risk lower bound. However, they assume training with population loss and analyze implicit bias through Hessian-trace minimization without providing explicit parameter characterizations as in Thm. 2. Crucially, their results do not apply to CE loss 9 or to sparse support-sets. Another interesting work [52] studies learning abilities of autoregressive training and inference. However, their findings do not apply to NTP as they inherently assume each context is followed by a unique next token.Finally, although stemming from different perspectives, the form of our convergence results echoes a recent conjecture by [82] regarding implicit optimization bias in transformers. Unlike their conjecture, which focuses on binary classification, our results are rigorously proven and apply to the NTP setting. Further detailed discussion on related follow-up work on implicit optimization bias in self-attention architectures, as initiated by [83] , is deferred to Appendix B. In contrast to this line of work, we here focus on the optimization biases of the NTP training-paradigm itself, which is orthogonal to the intricacies of the specific architecture generating the context embeddings. 
conclusion, limitations and future work: 
 Towards characterizing implicit regularization effects, we highlight two key aspects of NTP training: (i) Formulating it as CE optimization over distinct contexts; this is long recognized in language modeling (e.g., [44, 63] ) since Shannon's initial work, yet seemingly overlooked in recent studies, such as [49, 52] . (ii) Accounting for sparsity in the matrix of next-token conditional probabilities. While traditional language modeling techniques often mitigate sparsity using smoothing heuristics that assign non-zero probabilities to unobserved next tokens [44, 63, 39] , we recognize sparsity as a critical factor in NTP optimization that influences parameter divergence 10 .As the first study of implicit biases in NTP training, our results are based on several assumptions essential for establishing an initial foundational understanding. The framework allows for various exciting promising research directions, some of which we outline below.Even within the assumed linear setting and GD, interesting directions involve:• NTP-separability thresholds: Identifying exact thresholds for NTP-separability under distributional assumptions, akin to previous work on one-hot separability (Remark 1). However, relaxing the overparameterization requirement that the embedding dimension d be proportional to the number of distinct contexts m would necessitate exploring non-convex architectures (see 'Memory capacity' below).• Generalization: Studying generalization in NTP settings by examining statistical properties of the NTP-SVM solution. Past research has successfully undertaken similar investigations for one-hot classification (see Sec. 6). While we acknowledge the importance of addressing specific challenges inherent to NTP -such as determining an appropriate measure of generalization, or establishing suitable statistical models for context-embeddings that respect the discrete nature of the underlying token subsequences-we believe this direction holds promise for further exploration.In addition to these, essential extensions include relaxing the linearity assumption.• Architecture-specific embeddings: A bottom-up approach considering architecture-specific embeddings could begin by modeling the embeddings produced by, for instance, a shallow transformer and analyzing the effects of optimization biases on the training of both the transformer and the decoder weights. This complements the works of [83, 82] , who investigate one-layer self-attention with a fixed decoder. A challenge in this approach is balancing the restriction to shallow transformers (for analytical tractability) with ensuring that the NTP loss reaches the entropy lower bound. This may require constraining the training data distribution, for example, to a Markov chain [51, 25] .• Memory capacity in NTP settings: Without imposing further restrictions on the data beyond the discrete nature of tokens from a finite vocabulary, there is a strong case for investigating the memory capacity of sequence-to-sequence architectures, such as transformers, in the context of NTP. Recent studies on transformer memory capacity [40, 41] do not apply here.• Unconstrained features: Extending the top-down approach, one could consider freely optimizing context embeddings together with decoder vectors (also known as word embeddings). The resulting log-bilinear model, reminiscent of wor2vec models [63, 55] , extends the unconstrained features model, which has recently been employed to investigate neural collapse geometry in one-hot classification settings [56] . This idea offers a promising avenue for uncovering structures in the geometries of context and word embeddings when learned jointly, potentially revealing new insights into the capabilities of sufficiently expressive language models (see Fig. 1 for cases involving only the latter).• Other optimizers: Exploring the NTP implicit bias of adaptive algorithms, such as Adam, potentially building on recent works in this area focused on one-hot classification [96, 95] .We hope this work inspires further research in the discussed directions, contributing to a deeper understanding of the intricacies involved and potentially yielding improvements in NTP training.  
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification: The paper involves foundational research on the optimization properties of next-token prediction, that has the potential to enable better understanding of operating regimes of language models with respect to optimization, generalization and robustness.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). 
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]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 effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.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] 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. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] 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. • 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. 
 
 Appendix: All experiments were conducted on a MacBook Pro equipped with a 2.3 GHz Quad-Core Intel Core i7 processor and 32 GB of memory. The experiments are of relatively small scale and were implemented in Matlab. The code is straightforward to reproduce, following the detailed specifications provided in the subsequent sections. For completeness, the code will be made publicly available on Github in the final version of the paper.A.1 Additional details on 2D example of Fig.1 Interpreting the V × d decoder matrix as the matrix of learned word embeddings, this provides a visualization of their geometry. As per our results, the two word-embedding matrices W ⋆ and W mm lie on orthogonal subspaces. The geometry of the first depends on the probabilities of in-support tokens, while that of the second depends only on the support set of these probabilities. See also caption of Fig. 2 .We examine the implicit bias of GD on NTP training with overparameterization on synthetic data generated as follows. We construct dataset with n = 5000 sequences involving m = 50 distinct contexts. Each distinct context gets mapped to a randomly generated embedding of dimension d = 60 > m. We set vocabulary size V = 10 and each context j ∈ [m] is followed by S j = 6, ∀j ∈ [m] possible next-tokens. The support sets S j ⊂ V and the probabilities pj,z , z ∈ S j are chosen randomly; see Fig. 3 for representative examples from the training dataset. For a fixed realization of the dataset (for which H ≈ 1.445nats), we run GD, normalized GD (NGD), and Adam from random LeCun initialization. For GD, we use learning rate η = 0.5 and for NGD and Adam η = 0.01. For Adam, we also set β 1 = 0.9, β 2 = 0.99. We run all algorithms for 1e4 iterations. For each case, we plot the following as a function of iterations: Fig. 4 shows an instance of these. As predicted by our analysis, in this overparameterized setting: CE loss converges to its lower-bound, parameter norm increases, iterates align in direction with W mm , and the subspace component converges to W ⋆ .Figure 3 : Eight randomly picked contexts with their associated next-token empirical conditional probabilities pj . The indices shown on the x-axis define the support set S j of each context.Figure 5 illustrates the same plots, but this time for training over the same dataset with NGD and Adam. We observe same implicit bias, but faster convergence. For NGD, this is consistent with analogous findings (rigorous in that case) for one-hot classification [60, 36] . Implicit bias in transformers. As already mentioned in Sec. 6, our work is closely related to [82] , where the authors investigate the implicit bias of self-attention in transformers. The insight put forth in the prequel [83] is that softmax attention induces implicit-bias behaviors that bear similarities to vanilla implicit bias of one-hot prediction. Concretely, [82] 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 [82] 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 [82] 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.Beyond [82] , there is growing recent research investigating optimization and generalization principles of transformers, e.g., [70, 24, 48, 93, 99, 1, 45, 83, 82, 84, 17] . These efforts predominantly employ a 'bottom-up' approach that involves isolating shallow transformers, often with simplifications such as removing MLPs, utilizing single heads instead of multiple, and fixing certain parts while training only a subset of trainable parameters. Most of these studies have focused on classical one-hot supervised settings, and only a handful (e.g., [84, 85] ) have seeked extending these 'bottom-up' analyses to NTP settings. Yet, their primary emphasis remains on uncovering the role of attention and how attention weights evolve during training. Instead, our approach uniquely emphasizes the NTP training paradigm itself, shifting the focus from the intricacies of specific transformer architectures.Upon completing this paper, we became aware of independent contemporaneous research by Li et al. [46] 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 and colleagues [82] , albeit in very distinct settings. Notably, contrary to both [83] and [46] , we unveil the optimization intricacies of the NTP paradigm, even within the simplest linear settings.Classification with soft labels. Unlike one-hot classification, soft-label classification associates each example with a probability vector, where each entry represents the likelihood of a corresponding label characterizing the example. Although arguably less prevalent than one-hot (or hard-label) classification, soft-label classification arises in various contexts, including modeling human confusion during crowd-sourcing [65, 75, 18] , knowledge distillation [32] , label smoothing [79] , and mixup [98] . Our model of last-token prediction also falls within this setting. Specifically, our approach is most closely related to soft-labels generated by averaging annotators' hard labels [65] , rather than following the winner-takes-all rule to assign labels. [65] and follow-up work have provided empirical evidence that using probabilistic soft labels generated from crowd annotations for training leads to improved performance in terms of model generalization, calibration, and robustness to out-of-distribution data. To the best of our knowledge, no prior work has investigated the implicit bias of gradient descent in this or other soft-label classification settings; thus, our results are of direct relevance to these contexts as well.For concreteness and simplified notation, in the paper's main body we focus on NTP over sequences of fixed length. We show here that this encompasses the autoregressive (i.e., sequential) setting with minimal changes. This also emphasizes the role played in our results by the sequence length.As pointed in (1), the full autoregressive NTP objective averages T individual losses (without loss of generality assume sequences of equal maximum length T ). In order to make our analysis applicable, we first need to express (1) in terms of unique contexts. Mirroring the notations in Sec. 2, define the following for t ∈ [T -1]:• m t , t ∈ [T -1] is the number of distinct contexts of size t. Note that m 1 ≥ m 2 ≥ ⋯ ≥ m T -1 . • m = ∑ T -1 t=1 m t is the total number of distinct contexts in the dataset • ht,j ∶= h θ ( xj,t ), t ∈ [T -1], j ∈ [m t ]is the embedding of the j-th (among all t-long contexts) distinct context xj,t . • πj,t is the empirical probability of xj,t .• pj,t,z is the empirical probability that context xj,t is followed by token z ∈ V. • S j,t is the support set of the next-token distribution of context xj,t .CE = -∑ t∈[T -1] ∑ j∈[mt]πt,j ∑ z∈Sj,t pt,j,z log (S z (W ht,j )) .To continue enumerate the multi-setI ∶= {i = (j, t) | t ∈ [T -1], j ∈ [m t ]}.We may then rewrite the above asCE = -∑ i∈I πi ∑ z∈Si pi,z log (S z (W hi )) .At this point note that this is of identical form to (2) . Consequently, the definitions (e.g., NTPseparability, NTP-margin) and results derived in the main body for sequences of fixed length are applicable to the AR setting, extending mutatis mutandis. Remark 2 (The role of sequence length.). Despite the above reduction of the AR setting to the fixed-length setting, it is crucial to recognize that sequence length remains a significant factor in the AR model. Specifically, it influences the formulation through support sets and their associated probabilities. As sequences extend in length, their corresponding support sets generally become sparser, indicative of less ambiguity in predicting the next token. This dynamic is captured by Shannon's inequality,H t ≥ H t+1 , where H t = -∑ j∈[mt] ∑ z∈S ℓ t,jπ t,j pt,j,z log(p t,j,z ), reflecting the incremental reduction in entropy as sequence length increases.Throughout, lowercase and uppercase bold letters (e.g., a and A) represent vectors and matrices, respectively. ⟨⋅, ⋅⟩ and ∥⋅∥ denote Euclidean inner product and norm, respectively. For matrix A, we denote its pseudoinverse as A † . All logarithms are natural logarithms (base e). We denote e v the v-th standard basis vector in R V . ∆ V -1 denotes the V -dimensional unit simplex and S() ∶ R V → ∆ V -1 the softmax map:S(a) = [S 1 (a), . . . , S V (a)] ⊺ , with S v (a) = e e ⊺ v a ∑ v ′ ∈[V ] e e ⊺ v ′ a .As explained in Section 2 we represent a training set asT m ∶= {( hj , πj , pj,z∈V )} j∈[m] .We assume that embeddings are bounded and denoteM ∶= √ 2 max j∈[m] ∥ hj ∥ .Given T m , letF = span ({(e z -e z ′ ) h⊺ j ∶ z ≠ z ′ ∈ S j , j ∈ [m]})a subspace of V × d matrices and F ⊥ its orthogonal complement. Denote P F , P ⊥ the orthogonal projections onto F and F ⊥ , respectively. For convenience, for W ∈ R V ×d , we denoteEQUATIONClearly, for all W ∈ R V ×d , it holds CE(W ) ≥ CE F (W ). Note also that for all W ∈ F and for all W d ∈ F ⊥ that satisfy Eq. (6a), it holds CE F (W ) = lim R→∞ CE(W + RW d ). Thus, under NTP compatibility and NTP separability,EQUATIONThroughout we assume GD is ran with step-size η ≤ 1/(2L) where L is the smoothness of CE loss. This condition is not explicitly mentioned thereafter.The following result follows from standard optimization analysis for smooth convex functions specialized to functions that do not attain their infimum. The version presented here is adopted from Lemma 2 in [37] .Lemma 3. It holds lim k→∞ CE(W k ) = inf W CE(W )and also lim k→∞ ∥W k ∥ = ∞.In the lemma below, we collect some useful and simple-to-show properties of the GD and regularization paths. These are adaptations of corresponding results for one-hot binary classification over general non-separable data established in [34] . Lemma 4. Suppose conditions (6) hold for some W d . Also, that there exists W p = W ⋆ ∈ F satisfying condition (4). The following hold:1. CE F (W ⋆ ) = inf W ∈F CE F (W ) = H,2. W ⋆ is the unique minimizer of CE F on the subspace F,3. lim k→∞ P F (W k ) = W ⋆ , where W k are GD iterates, 4. lim k→∞ ∥P ⊥ (W k )∥ = ∞, 5. lim B→∞ P F ( ŴB ) = W ⋆ ,where ŴB is the reguarlized solution (8),6. lim B→∞ ∥P ⊥ ( ŴB )∥ = ∞.Proof. It is easy to check by direct substitution of W ⋆ in ( 12) and use of ( 4) that CE F (W ⋆ ) = H. This and (13) show the first claim.The first claim shows W ⋆ is a minimizer. Suppose for the sake of contradiction there is a different minimizerW ⋆ ≠ W 1 ∈ F. Then, since CE F (W 1 ) = H, it also holds for W R ∶= W 1 + RW d that lim R→∞ CE(W R ) = H.In turn, this implies for all j ∈ [m]:EQUATIONlim R→∞ S v (W R hj ) = 0, ∀v ∉ S j .The first condition gives then that W 1 must satisfy (4). Since W ⋆ also satisfies these equations, denotingW ∆ = W ⋆ -W 1 ≠ 0, it holds: ⟨W ∆ , (e z -e z ′ ) ⊺ hj )⟩ = 0, ∀j ∈ [m], z ≠ z ′ ∈ S j .But W ∆ ∈ F, so this forms a contradiction. Hence, W ⋆ is unique solution in F of ( 4) and unique minimizer of CE F on the subspace F.The proof of the third claim follows the same way as the proof of part (1) of Thm. 15 of [37] .For completeness: It follows by the lemma's assumptions and Lemma 3 that lim k→∞ CE(W k ) = H. Combining with the first claim of the lemma yieldslim k→∞ CE(W k ) = CE F (W ⋆ ). Since CE F (W k ) ≤ CE(W k ), this finally gives lim k→∞ CE F (W k ) = lim k→∞ CE F (P F (W k )) = CE F (W ⋆ ).Since W ⋆ is unique by the second claim, the desired then follows.For the fourth claim, recall from Lemma 3 that lim k→∞ ∥W k ∥ = ∞. From the previous claim, we also have lim k→∞ ∥P F (W k )∥ < C for some constant C > ∥W ⋆ ∥. Thus, the desired follows by applying the fact that∥W k ∥ = ∥P F (W k )∥ + ∥P ⊥ (W k )∥.The proof of the last two claim is exactly same as that of the third and fourth claim. Only now use the facts that lim B→∞ CE(W B ) = H and lim B→∞ ∥W B ∥ = ∞ (see proof of Theorem 1).Lemma 5. Let W k denote the GD iterate at iteration k. Recall the decompositionW k = P F (W k ) + P ⊥ (W k ) = W k,F + W k,⊥ .Fix any α ∈ (0, 1). There exists large enough R = R(α) and k 0 = k 0 (R) such that for any k ≥ k 0 , it holds that ∥W k,⊥ ∥ ≥ R andEQUATIONProof. We drop the subscript k to lighten notation.First, note by Lemma 4.D that, for arbitrary R, we can pickk 1 = k 1 (R) such that for all k ≥ k 1 : ∥W ⊥ ∥ ≥ R.Thus next, we will prove the main claim, i.e. for large enough ∥W ⊥ ∥ inequality ( 14) holds. Denote R ′ = ∥W⊥∥ ∥W mm ∥ . Substituting in CE expression (2), and using the fact that W mm ∈ F ⊥ by (6a) yield:EQUATIONMoreover, decomposing W = W F + W ⊥ , and definingW⊥ ∶= ∥W mm ∥ ∥W ⊥ ∥ W ⊥ = 1 R W ⊥ ,we haveEQUATIONwhere we used that, by definition, W ⊥ ∈ F ⊥ . Thus, our goal becomes showing ( 15) ≤ ( 16), for large enough R. To do this, we consider two cases as follows below.For the remaining of the proof recallM ∶= max j∈[m] √ 2∥hj ∥ and use the logits shorthand:lj,v = e ⊺ v W⊥ hj and ℓ mm j,v = e ⊺ v W mm hj .Case 1: W ⊥ is well aligned with W mm . SupposeEQUATIONUsing this, linearity of logits, and Cauchy-Schwartz, yieldslj,z -lj,v ≤ ℓ mm j,z -ℓ mm j,v + ϵM, ∀j ∈ [m], z ∈ S j , v ∉ S j . Thus, ∑ v∉Sj e -R ′ (ez-ev) ⊺ W⊥ hj ≥ e -ϵM R ′ ∑ v∉Sj e -R ′ (ez-ev) ⊺ W mm hj = e -αR ′ ∑ v∉Sj e -R ′ (ez-ev) ⊺ W mm hjAlso recall by feasibility of W mm thatEQUATIONThus,∑ v∉Sj e -(1+α)R ′ (ez-ev) ⊺ W⊥ hj ≤ e -αR ′ ∑ v∉Sj e -R ′ (ez-ev) ⊺ W mm hjComparing the above two displays yields∑ v∉Sj e -(1+α)R ′ (ez-ev) ⊺ W⊥ hj ≤ ∑ v∉Sj e -R ′ (ez-ev) ⊺ W⊥ hj ,which implies the desired (15)≤( 16) for any value of R ′ (eqv. ∥W ⊥ ∥).Case 2: No alignment. Suppose now that (17) does not hold. Note that ∥ W⊥ ∥ = ∥W mm ∥ and since (NTP-SVM) has a unique solution it must be that W⊥ is not feasible. But W⊥ ∈ F ⊥ , thus it satisfies the equality constraints. This then means that there exist δ ∶= δ(ϵ) andEQUATION(Note the above holds for all z ∈ S j⋆ because lj⋆,z = lj⋆,z ′ since W⊥ ∈ F ⊥ .)To continue, we introduce the shorthand notationA j,z ∶= A j,z (W ) = ∑ v∈V e -(ez-ev) ⊺ W F hjas well asEQUATIONA max ∶= max j∈[m],z∈Sj A j,z .Using (19) we may lower bound (16) as follows:EQUATIONwhere in the last line we used πj ≥ 1/n, ∀j ∈ [m] as well as log(1 + x) ≥ x 1+x , x > 0. On the other hand, using property (18) for max-margin logits, we can upper bound (15) as follows:EQUATIONA min ,where in the last line we used log(1 + x) ≤ x, x > 0.In view of the two last displays, it suffices thatV e -R ′ (1+α) A min ≤ e -R ′ (1-δ) n(A max + 1) ⇐⇒ R ′ ≥ 1 δ + α log ( nV (A max + 1) A min ) .All it remains is obtaining bounds for A min , A max specifically showing that they do not depend on R. By Cauchy-Schwartz:EQUATIONThus, there existsk ⋆ = k ⋆ (∥W ⋆ ∥) such that for all k ≥ k ⋆ : V e -2M ∥W⋆∥ ≤ A min ≤ A max ≤ V e 2M ∥W⋆∥ .Hence, the desired ( 21)≤(20) holds providedEQUATIONSet R = R(α) = {RHS of (23)} and k 0 (R) ∶= max{k 1 (R), k ⋆ }. We have shown this guarantees for all k ≥ k 0 : ∥W ⊥ ∥ ≥ R and by choice of R also (21)≤ (20) . This in turn implies (15)≤( 16), as desired to complete the proof.For the subspace component, see Lemma 4.C. For the directional convergence, the key ingredient of the proof is Lemma 5. After that, the proof follows identically to Thm. 15(2) in [37] . We include the details for completeness, but there are no novel aspects in the rest of this section.Let any ϵ ∈ (0, 1) and choose α = ϵ/(1 -ϵ). By Lemma 5, there exists k 0 such that for any k ≥ k 0 , we have ∥W k,⊥ ∥ ≥ max{R(α), 1/2} and⟨∇ CE(W k ), W k,⊥ -(1 + α)∥W k,⊥ ∥W mm ⟩ = ⟨∇ CE(W k ), W k -(W k,F + (1 + α)∥W k,⊥ ∥W mm )⟩ ≥ CE(W k ) -CE(W k,F + (1 + α)∥W k,⊥ ∥W mm ) ≥ 0 ,where we also used convexity of the loss.Consequently,⟨W k+1 -W k , W mm ⟩ = ⟨-η∇ CE(W k ), W mm ⟩ ≥ (1 -ϵ)⟨-η∇ CE(W k ), W k,⊥ ⟩ ≥ (1 -ϵ)⟨W k+1,⊥ -W k,⊥ , W k,⊥ ⟩ ≥ (1 -ϵ)⟨W k+1,⊥ -W k,⊥ , W k,⊥ ⟩ = (1 -ϵ) 2∥W k,⊥ ∥ (∥W k+1,⊥ ∥ 2 -∥W k,⊥ ∥ 2 -∥W k+1,⊥ -W k,⊥ ∥ 2 ) ≥ (1 -ϵ) (∥W k+1,⊥ ∥ -∥W k,⊥ ∥ -2η(CE(W k,⊥ ) -CE(W k+1,⊥ )) ,where the last step used ∥W k,⊥ ∥ ≥ 1/2, the fact that x 2 -y 2 ≥ 2y(xy), ∀x, y and smoothness of the CE loss.Telescoping the above expression and rearranging yields⟨W k , W mm ⟩ ≥ (1 -ϵ) ∥W k,⊥ ∥ ∥W k ∥ - ⟨W k0 , W mm ⟩ -(1 -ϵ)∥w k0,⊥ ∥ -η CE(W k0 ) ∥W k ∥ ≥ (1 -ϵ) - ∥W k,F ∥ 2 + ⟨W k0 , W mm ⟩ -(1 -ϵ)∥w k0,⊥ ∥ -η CE(W k0 ) ∥W k ∥ Now recall from Lemma 4 that lim k→∞ ∥W k ∥ = ∞ and lim k→∞ ∥W k,F ∥ = ∥W ⋆ ∥. Thus, lim inf k→∞ ⟨W k , W mm ⟩ ≥ 1 -ϵ.Since ϵ is arbitrary, the desired follows.We provide a detailed proof of Theorem 1 filling in missing details from the proof sketch in the main paper.First, we show that ŴB is on the boundary, i.e. ∥ ŴB ∥ = B. Suppose not, then ⟨∇ CE( ŴB ), U ⟩ = 0 for all U ∈ R V ×d . Using the CE expression in (2) and a few algebraic manipulations, yields (24) where we denote the output probabilities at ŴB as (24) . Then, the first term in the parenthesis in (24) is zero by (6a), while the second term is strictly positive by (6b) and strict positivity of softmax entries, leading to contradiction. Now, consider point W ⋆ B = W ⋆ + R(B) ⋅ W mm , where, W ⋆ ∈ F satisfies (4), and⟨-∇ CE( ŴB ), U ⟩ = ∑ j∈[m] πj ∑ z∈Sj pj,z ( ∑ z ′ ∈S j z ′ ≠z s j,z ′ (e z -e z ′ ) ⊺ U hj + ∑ v∉Sj s j,v (e z -e v ) ⊺ U hj ),s j,v ∶= S v ( ŴB hj ), v ∈ V, j ∈ [m]. Choose U = W mm inR = R(B) is chosen such that ∥W ⋆ B ∥ = B. Concretely, for B > ∥W ⋆ ∥, set R = 1 ∥W mm ∥ √ B 2 -∥W ⋆ ∥ 2 .Note also that R/B → 1/∥W mm ∥ as B → ∞. We will show that W ⋆ B attains a small CE loss as B (hence, R) grows. To do this, denote for convenience the logits for all v ∈ V, j ∈[m] ∶ ℓ ⋆ j,v ∶= e ⊺ v W ⋆hj and ℓ mm j,v ∶= e ⊺ v W mm hj , and note that e ⊺ v W ⋆ B hj = ℓ ⋆ j,v + R ℓ mm j,v . By using ( 4) and (6a):∑ z ′ ∈Sj e -(ℓ ⋆ j,z +Rℓ mm j,z -ℓ ⋆ j,z ′ -Rℓ mm j,z ′ ) = 1 pj .Moreover, using (6b)∑ v∉Sj e -(ℓ ⋆ j,z +Rℓ mm j,z -ℓ ⋆ j,v -Rℓ mm j,v ) ≤ e -R ∑ v∉Sj e -(ℓ ⋆ j,z -ℓ ⋆ j,v ) ≤ C e -R ,where we define constant(independent of R) C ∶= V e ∥W ⋆ ∥M , for M ∶= √ 2 ⋅ max j/∈[m] ∥ hj ∥.Combining the above displays and using in Eq. ( 2), yieldsEQUATIONwhere, the second line uses log(1 + x) ≤ x, x > 0, and the third line uses πj , pj,z are probabilities.Next, towards arriving at a contradiction, we will show that if ŴB is not in the direction of W mm , then it incurs a loss that is larger than CE(W ⋆ B ). Concretely, assuming the statement of the theorem is not true, we we will upper boundEQUATIONBy our assumption, there exists ϵ > 0, such that there exists arbitrarily large B satisfying:EQUATION) Define Ŵ = 1 R ′ (B) ( ŴB -W ⋆ ),where, R ′ = R ′ (B) > 0 is chosen so that ∥ Ŵ ∥ = ∥W mm ∥. Concretely, for large enoughB ≥ 2∥W ⋆ ∥, set R ′ = 1 ∥W mm ∥ √ B 2 -2B⟨W B , W ⋆ ⟩ + ∥W ⋆ ∥ 2 .Note that it holds lim B→∞ R ′ /B = 1/∥W mm ∥. Thus, we can always choose B large enough so that Eq. ( 27) guarantees ∥ Ŵ -W mm ∥ ≥ ϵ ′ , for some ϵ ′ > 0. Since W mm is the unique minimizer of (NTP-SVM) and ∥ Ŵ ∥ = ∥W mm ∥, it follows that there exists δ ∈ (0, 1) and j ∈ [m] such that at least one of the following is true (i) ∃z and z ′ ≠ z ∈ S j such thatEQUATION(ii) ∃z ∈ S j , v ∉ S j such thatEQUATIONCase (i): Without loss of generality (e z -e z ′ ) ⊺ Ŵ hj ≤ -δ (otherwise, flip z, z ′ ). Thus, ignoring all but one term in (26) givesEQUATIONwhere we use ℓ j,v = e ⊺ v ŴB hj , v ∈ V to denote logits of ŴB . Using ( 4) and ( 28), yieldsℓ j,z ′ -ℓ j,z = (e z ′ -e z ) ⊺ (R ′ Ŵ + W ⋆ ) hj ≥ R ′ δ + log ( pj,z ′ pj,z ) .Put in (26) and using pj,z ≥ πj pj,z ≥ 1/n showsCE( ŴB ) ≥ H + 1 n log ( e R ′ δ n )Compare this with (25) . For large enough B, it is clear that πj pj,z log (p j,z c e R ′ δ ) > Ce -R . Thus, CE( ŴB ) > CE(W ⋆ B ), a contradiction. Case (ii): We can assume Ŵ ∈ F ⊥ , since otherwise we are in Case (i). Now, again ignoring all but the (j, z) term in the CE loss for which (29) holds for some v ∉ S j , we find CE( ŴB ) -H ≥ πj pj,z log (p j,z ( ∑ z ′ ∈Sj e (ℓ j,z ′ -ℓj,z) + e (ℓj,v-ℓj,z) )).EQUATIONMoreover, by ( 29):e ℓj,v-ℓj,z ≥ e -R ′ (1-δ) e ℓ ⋆ j,v -ℓ ⋆ j,z ≥ c ′ e -R ′ (1-δ) ,for constant (independent of B) c ′ ∶= e -∥W ⋆ ∥M . Putting the above together yield:CE( ŴB ) -H ≥ πj pj,z log (1 + pj,z c ′ e -R ′ (1-δ) ) ≥ c ′ e -R ′ (1-δ) 2n 2 .where the second inequality uses log(1 + x) ≥ x 1+x , x > 0. Compare this with (25) . For large enough B, (recall R, R ′ grow at the same rate) it holdsc ′ 2n 2 e -R ′(1-δ) > Ce -R . Thus, CE( ŴB ) > CE(W ⋆ B ), a contradiction. In either case, we arrive at a contradiction, which completes the proof.Justification: Detailed proofs of all results provided in Sec. E. 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.Throughout, ev ∈ R V is the v-th standard basis vector, and Sz(u) = e ⊺ z S(u) the z-th entry of softmax output.In NTP, the ground-truth next token is inherently embedded within the underlying text, thus strictly speaking, it falls under the self-supervised learning paradigm[66]. Yet, the utilization of the CE training objective resembles to supervised training. We leverage this resemblance and regard NTP training as an instance of supervised learning, while also emphasizing how it differs from one-hot encoding supervision.NTP is widely used across various modern language modeling architectures, including transformers[66,67], state-space models[26,27], and LSTMs[5].A complete list of notations is also given in Appendix D.It will be see below, and can be easily checked by the reader, this condition alone is insufficient; the NTPseparability condition in Defn. 2 is also needed.6 If Eqs. (4) have a solution, say W1, every other solution takes the form W p = W1 + Wnull, where Wnull is orthogonal to (ez -e z ′ ) hT j ∶ z ≠ z ′ ∈ Sj, j ∈ [m]. Thus, Wnull ∈ F ⊥ is in the orthogonal complement of F.The necessity for such large d can be mitigated through the utilization of non-linear architectures (such as an MLP decoder), in which 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.In line with observations in one-hot encoding[59], we anticipate the directional behavior remains unchanged under stochasticity, e.g. when using SGD to minimize(2). Yet, note a subtle but crucial difference in applying SGD to (1) vs (2), as the latter involves sampling distinct contexts in each iteration. In this latter case, we also point out that favorable interpolation conditions, such as strong-growth (e.g.,[91]), can be shown to hold.[49, Thm. 4.3] uses[47, Cor. 5.2], which applies to regression on scalar labels; thus is not applicable in NTP.Parameter divergence in transformer-based language models has been empirically observed in[53]. 


Decision: The authors investigate the implicit structural bias introduced by the gradient descent optimizer with the cross-entropy objective in the context of next token prediction (NTP), a common approach used in large language models (LLMs). This work extends the analysis of implicit bias from one-hot classification to the NTP setting, where the target output for a given context is a (sparse) distribution over the vocabulary rather than a single one-hot vector. This adjustment reflects the multiplicity of potentially correct tokens given the same context in NTP tasks. The paper's analysis is restricted to the linear layer in NTP while assuming a fixed context embedding. The results appear to follow a similar line to prior work on implicit bias in soft-max attention, though, the authors highlighted some notable conceptual as well as mechanistic distinctions. The impact of the work can also be further improved by extending their analysis to include the attention mechanism, which would offer a more complete perspective on implicit biases within NTP models.

Overall, the paper offers useful insights into implicit bias in NTP settings. The impact can be further enhanced by broadening the analysis to account for the attention mechanisms.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Just the remark in the weaknesses.
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: A minor suggestion: lines 32-34 appear to require rephrasing.
limitations: limitations: The authors thoroughly discuss the limitations of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See above.
limitations: limitations: See above.

Comments: comment: Thank you for your comments. I raised the score.
Did this answer your question?
Also, we would appreciate if you could kindly update your score, as you mentioned in your previous response.

Thank you for your time
Thank you for your kind words and for sharing your thoughts on the conjecture—this is encouraging! 
We appreciate your time and support.
Thank you again for your feedback and support. 
We appreciate your decision to raise the score.
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.
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.
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.
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!
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.
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?
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.
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.