Implicit Optimization Bias of Next-token Prediction in Linear Models
Authors: Christos Thrampoulidis
Venue: NeurIPS 2024 poster (2024)
Keywords: entropy
gradient descent
SVM
next-token prediction
linear models
separability
language models
word embeddings
Primary area: natural_language_processing
TL;DR: By framing next-token prediction training as sparse soft-label classification, we characterize the implicit optimization biases of gradient descent in language models trained to achieve their entropy lower bound.

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.

Implicit Optimization Bias of Next-token Prediction in Linear Models
Christos Thrampoulidis Department of Electrical and Computer Engineering University of British Columbia
Vancouver, Canada cthrampo@ece. .bc.
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 distinct contexts, each tied with a sparse conditional probability distribution across a finite vocabulary of tokens, we introduce "NTP-separabilitt 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. the orthogonal subspace, the GD parameters diverge in norm and select the direction that maximizes 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.
Introduction
Next-token prediction (NTP) has emerged as the go-to paradigm in training modern language models, revolutionizing various applications such 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 mapping from the input context to the probability distribution over the (finite) vocabulary of possible tokens, enabling the model generate a token that is contextually appropriate [9, 8]. Recently, the NTP paradigm has witnessed remarkable empirical success through its utilization 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 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-los reaches its lower bound and there is many possible minimizers. To formalize the NTP paradigm, consider
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

autoregressive model qe parameterized by 0 trained to predict the next-token on sequences of length T using the eross-entropy (CE) loss:
min ExerT. E log (go (1.....-11)) 0
(1)
Here, sequences x (Z1, 2T) consist of tokens 2; from a finite vocabulary V (1,..., .V
and É is expectation over training set Tn of n such sequences sampled from some underlying true distribution over sequences. Typically, the model qo outputs probability of the next token computed via softmax applied on oulput logits, which are computed by projecting d-dimensional embeddings ho to the V-dimensional space with a trainable linear decoder W R Rxd, Formally,
1
go 2 , 10-1) Sz, ( (Z1,...,,...1))
1+ Σ zov exp ((ez -e..) Woo (Z1, (......, 2--1) 244
The CE loss is then minimized over 0 = (W. 0') using gradient- graseent-based methods, e.g. (S)GD, Adam.
We pose the question: Given training set Tru what are the structural properties of the weights O found by minimizing the NTP objective with gradient-based optimizers? As in prior research in one-hot supervised classification 2 (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 l 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). this traditional scenario, the training set comprises feature-label pairs (I.y), where æ É RP is a continuous feature, and y represents its unique label. The optimization process minimizes the following training objective (over W,0'): )'): [<...) log (WW (z=)
At first glance, excluding the sequential format of Eq. (1), the NTP training scenario might seem identi- cal to traditional one-hot prediction: both aim minimize the same CE loss across models that param- eterize probabilities using the softmax of logits. Consider predicting the next token over fixed-length
sequences, say sequences of length f 1, via optimizing: Eg [- log (Sz, Whhe(Z1, 22-1))) 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 21,..., =t- of finite length sampled from finite vocabularies are naturally repeated a (vast) training set, potentially multiple times, each time followed by different tokens [73]. Consequently, the NTP paradigm involves training over m S n distinct (non-repetitive) contexts, each followed by a multitude of possi- ble 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," "rad,""and" etc., will not follow.
We study NTP training over finite vocabulary employing the following model. Given a large training set of n total sequences, we identify m S 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 P, are sparse, i.c., the support set Sj of P, satisfies [S,] << | V. In an extreme where [S = 1, Vj [m]. the probability vector P. becomes one-hot, leading to a scenario reminiscent of the traditional classification setting described earlier. However, such extreme is essentially improbable in practical language modeling [73]. With this framing, the NTP paradigm is
Model.
'Throughout, e., e R is the v-th standard basis vector, and S_(u) = (e)( the 2-th entry of softmax output. II NTP. the ground-truth next token is inherently embedded within the underlying text, thus strictly speaking, 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.

---

also related to supervised vision classification with soft labels, which advocates for training models on datasets where each example 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.
1.1 Contributions and Organization
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 e RVxd while keeping confext-embeddings fixed. This approach mirrors foundational studies on implicit optimization bias in one-hot classification [76, 34], which first focused on linear models. allows exploring the complexities of the NTP training objective, distinct from the embedding architecture", and while it renders the logits linear and the objective convex, it still poses 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 approach its lower bound, the empirical conditional entropy. We introduce two conditions: NTPpa-compatibility and NTT--spparbbility, 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 decoder maximizing the MTP-margin, denoted as W guides the directional O.
convergence of the ridge-regularized CE minimizer, Wx. as the regularization parameter
Implicit bias of GD. We establish that mm also determines the implicit bias of gradient descent (GD) iterates in Sec. 5. Specifically, in the limit of iterations k undoubtedly in norm and converge to a finite W' within a data subspace I, while simultaneously aligning with W in the complementary subspace Ft. The finite component W e I solves system of linear equations associated with the NTP compatiiiiyy condition.
the GD iterates grow
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.
2 Sctup
......... V} represent a set of V tokens (e.g. words) and Z1st = (21..... -,) denote sequence of tokens 2 e V. To simplify presentation, 2r given contexts Zer := æ1:T-1 of fixed length, and we further let x = zet denote the context and a denote the last token. See App. C for straightforward extension to the sequential format of Eq. (I).
Let vocabulary V [V]] := (1,....
focus predicting the T-th token
We assume access to a training set consisting of n sequences T, := {(x,, =;) Hetm, with æ; 6 := TT and 2; E V. Let /: X' - Rd an embedding map that maps contexts (i.c., sequences of T - 1 tokens) to d-dimensional embeddings. The map 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 model fw X the model output passes through a softmax, it defines the model's probability mass function for the next-token prediction, given as qw (-|x=) (ffw)))). where S(-)= RR V 1 is the softmax and A V-1 is the V V-ddimensional simplex. The decoder is trained by minimizing the empirical CE loss CE(W):= 1 Σid(n] log (ww22(x)))
+ R parameterized by decoding matrix W R R/xdd such that fw (x) Wh(x). When
Distinct sequences and next-token distributions. Given dataset T, we denote T T...... Z.m the m S n distinct contexts among the (large number of) total n contexts 21,..., En within T, Let ft.
NTPP is widely used across various modern language modeling architectures. including transformers [66, 67]. state-space models [26, 27J, and LSTMs [5].
3

---

be the empirical probability of distinct context ... That is, 1 s n ū, S n is the number of contexts
that equal ē.. Furthermore, for each distinct context J e [m] let P. e A V-1 denote the probability vector of conditional next-token distribution, i.e., Pj.: := () e j.j= e [m]] ]. In other words, n- - n- PP.: is the number of occurences of token 2 as a follow-up context J. Finally, we denote the support set and size of the support set of these conditional distributions as Sj := (z { | x > 0} and S and v S are referred to 'in-support' and "out-of-support" respectively. Onwards, we implicitly assume that "not all loknss are likely after every context," i.e. 3j e [m]] such that S3- V. This mild assumption is naturally satisfied in language modeling under rich enough vocabulary. With this notation, we can express the NTP training loss
Sj== ==== [- Tokens
CE(W) = Σ ; je[m]] 2EV
Pj. Po.- log (S<(W((z))) -- Σ ū E Ds.e log (S,(W)))). SES,
(2)
Je me[m]]
where, in the last line we defined the shorthand h; = (.))) Similarly, we let h; h(a:;); i [m]. With some abuse of notation, we then obtain the following equivalent descriptions of the training set
(<x;; zi)
24)) [e[m] T, Tm (hhj # (.,,ev)) je[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]: HT := H := ((e,))-T, [- log (f(z/=)) = Eje[m] Exes, RzPo,= log (fj.z) It lower bounds the CE loss since CE(W)= = H. KL (pl gw) and the KL divergence is nonnegative.
3 When can the NTP-loss reach the entropy lower-bound?
The 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 <> (L qw) 0 iff for all j [m] and all 2 V: @w (z(z;)= Pj.z. Equivalently, for all j {[m]:
S² Whi)) Ú3.20 Vz c
(3a)
S (Wh,) =0, Vo V
(3b)
Beginning with (3a), this requires the training data to satisfy the NTPy-compatiblity condition defined below.
Definition 1 (NTP compai-om....ll.. Let denote the w-th standard basis vector in RV. We say that training data Tm are NYP-entropy-compatibe if there exists V x matrix WP satisfying:
Vj E [m]] = eS (e WPh; log (D5.=/53..)
(4)
We comment on the independence of the constraints: Fix any j [m]] Then, the set of constraints (as expressed in Eq. (4)) for all = # z' e S, (yielding (S,) constraints total) is equivalent to the set of the same constraints for any anchor 2, E S, and z' zj E Sy. i.e., an effective total of S₁ 1 linearly independent constraints for each j [m]. Additionally, note that the system of equations in Eq. (4) constrains WP with respect to a specific subspace of d matrices:
F = span ee 2 z' Sz E [mm]]),
(5)
that 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 I for later reference
Next, 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.
'A complete list of notations is also given in Appendix D.
'It will he see helow, and can be casily checked by the reader, this condition alone is insufficient; the NTP. separability condition in Defn. 2 is also needed.
'If Eqs. (4) have a solution, say Wa, every other solution takes the form W = W. Walli where Winll is orthogonal to (e. e. )h3 z' Sj.j [m.]. Thus, Will e 3 is in the orthogonal complement of

---

Definition 2 (NTP-separable). We say that training data Tm are -TP-separable if there exists V x d matrix W satisfying the following:
Vj E [m] (e:-e.) W h= =0
(6a)
E [m] u ee) " h, 21.
(6b)
As before, is easy to see that the constraints in (6) can be equivalently expressed by enforcing (6a)
and (6b) for an anchor 2j e Sj and all z' c Sj \ {2²} and é Sj. respectively. Consequently, there exist effectively V linearly independent constraints per context j E [m]]
We now discuss the interpretation of these constraints. The subspace constraints in Eq. (6a) project W onto the subspace F which is the orthogonal complement of the subspace I defined in (5),
This leaves the softmax probabilities of possible next tokens (in set Sy) intact, and fully determined by WP as per the compattiiiy condition. Formally, WP + Wd continues satisfying (4). Moving the halfspace constraints in (6b), can interpret these using Kesler's construction as enforcing linear separability in the space Rvxd [30]: Each d-dimensional context embedding h is mapped Sj Sÿ) higher-dimensional points (e2 e, en)h,,= = Sj, 0 Sj- These points collectively for all e [ml] must lie within the interior of the same halfspace induced by the hyperplane (W 0. Refer to Fig. I(Left) and its caption for an alternative interpretation of the rows of Wmm as word-embeddings in Rd (illustration in d 2).
The impact of -Tp-saparabyllty on the softmax probabilities can be understood algebraically by considering W y and v Sy. We have:
S. (w (E ( 'w",,
Σ
en (exe hw
zc5
$ Σ e (c,) "W"
2c5,
< <e
where the first inequality removes non-negative exponential terms and the second one follows from (6b). The upper bound above approaches as 7 - co, thus (3b) holds asymptotically in y.
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 Tm is NTP cTompccoetee and NTP- separable, with the respective matrices WP and W' satisfying conditions (4) and (6). While all finite W satisfy CE(W) H, 70++00
holds for W WP + +.. Wd that CE(W")
H.
Hence, CE approaches its lower-bound in the limit of a direction W = Wd / and offset WP satisfying the constraints of NTP- -TP-separy and NTP- STP-compatibiliity respectively. In other words, parameter weights W that minimize the CE loss consist of two components: a finite projection Wy := P3 (W) W onto the data subspace I and infinite-norm component onto the orthogonal complement 3 in the direction of W
Finally, we note that while Defns. and are stated for linear models, they naturally extend to a more general formulation for nonlinear models. Specifically, consider NTP-separability (similar for NoP--oopatiiilltyy: the general conditions require that both the decoder weights W and model weights 0, which parameterize the embeddings h, = ho (T.;), must satisfy Eq. (6) simultaneously.
3.1 The role of overparameterizzation
We show that overparameterizztion provides sufficient condition for the solvability of Eqs. (4) and (6). Start with the halfspace constraints in Eq. (4) for NTP, compatibiiyy These can be compactly expressed as Ej,e, WPh, = aj,z. where Ej.,, E R(S,-1)xV has rows es,- E2' and aj e (R())-1) has entries log (6.,,, /P,.2'))))))))))(() for some anchor 2j Sy. Now, since the rows of Ej,z, independent, the question becomes equivalently that of determining when W wyih....., hm]]
linearly
[E] (11.111.... E Em am. has a solution. This is always the case when d > m and the d x m
5

---

embedding matrix H = [hi..... hmn] is full rank (m). Then, there exists WP such that condition (4) holds. In fact, HT has nullspace. implying the existence of infinite number of solutions to (4). These solutions take the form WP W WP, where W e &F is the unique solution onto the subspace, and WP e gt.
In contrast (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 solution. By following the exact same argument as before, we arrive at the same sufficient conditions for the existence of solution W We summarize these findings.
Lemma (Overparameteriiztion implies NTP-separability). Assume overparumeterization > m and full-rank embedding matrix H Ruxm, Then, there exists an infinite number of solutions WP and W that satisfy conditions (4) and (6), respectively.
Thus, > m. 7 which also generically favors full-rankness of the embedding matrix [92]. implies both compatioit and -VP-aapara..llity Combined with Prop. 1, it also implies that there are infinitely many possible directions W along which the NTP loss approaches H. motivating the implicit-bie -bia question: For a specific iterative algorithm aimed at minimizing the NTP loss, which direction does prefer? We will address this question in the remainder of the paper.
Remark 1. In the trivial case where S₁ 1, V E [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- separabbryyy 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 structured (random) data, including [12, 71, 57, 54). The exact threshold in corresponding muiliclass settings is more intricate, but [19, 81, 11] have made progress in this direction. An interesting question is determining exact thresholds for NTP STP-separabillty, which would improve upon the sufficient condition of Lemma 1.
4 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:
W F == arg min [ CE(W).
(8)
This involves minimizing a strictly convex function in a bounded domain; thus, We is unique. This section's main result characterizes the limit of WB as B + oo under ppp-separailliyycoomaiibnt Before that, we first define the next-token prediction support-vector machines (SVM) problem. Definition 3 (NTP-SVM), Given NTP-separable training set Tm. NTP-SVM solves the following:
W : arg minw [W]
subj. to W VRyd satisfying (6a) and (6b).
(NTP-SVM)
This is a strongly convex quadratic program with mV- -Ej{(m] Sj linear inequality and Eje[m) Sy-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 Fi:
Wmm arg max|
||| W-.- minje(m].
(e. - (e)- WWh,
It turns out this direction determines the preferred limiting direction of the regularization path.
Theorem 1 (Implicit bias of the regularization-patie Assume training data Tm is NTP XTPx-compatibe and sep-saaaee Let WB be defined in (8). Then, holds that lim
B Wmmm
Wal
The proof sketch below illustrates how the NTP pillyppirabilyyyyyopabilllllttyy 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 - se Po ( (W)) W*, These together provide complete characterization of the implicit optimization bias (8).
"The necessity for such large d can be mitigated through the utilization of non-linee 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.

---

Proof sketch (App. E.2 for details). We first show WB is on the boundary: WB = B. If not, then (V CE(WB) Wmm) )- O. But, few algebraic manipulations show (-V CE( WB). W equals
Σ J=[m
n
Σ 255,
P3.=
Σ 2'8j,a't2
$3.2' (e= -ex)) Wimmy Σ
8j. (e= e.) " w vES;
where we denote $j := Se Se(WBl Bh,) 0, V.je E [m]] The first term in the parenthesis is zero by (6a), while the second term strictly positive by (6b), leading to contradiction.
Now, consider a 'genie' point W = W R(B) Wmmm where W E F satisfies (4), and R R(B) chosen such that Wi| B. We will show that W attains small CE loss B (hence, Rt) grows. To do this, denote for convenience the logits
= h and fum j.v E e,WWimm W
for all for V,je e [m]. and note that e. Wihh =,,, R 2. By using (4) and (6a):
Σ
(%.a* Rem
Σ
Σ
Pi 1
E'E
Moreover, using (6b) and defining C' Ve Vel IM for M v2. maxje[m]] klyl gives:
Σ
+ -f3,v- RR
S se Σ c (5 -G,, SCER.
D$S;
Combining the above within Eq. (2), using log(1 + x) s W. and the faet that Ry. are probabilitties, yields:
py. log CE(WE) & Σ n je[m] St,;
+ CC CH
(9)
p3.z
Next, towards contradiction, will show that if WB B is not in the direction of Wmm. then it incurs a loss that is larger than CE(
(W)). The trick here to bound the KL divergence term:
CE(W n) H- Σ xj Σ Pj.e log (p..)( Σ 25S,
eiaebii + Σ cl-bb)))
(10)
3c[m]
158
where we denote logits Ej.0 == e, W Whj.. Assume there exists e 0 and arbitrarily large B satisfying:
( (1I I/B) WB Wmm
(11)
Define W = (WB (B), where R' (BBB) > 0 can be chosen so that W = W
||
Further choose B large enough so that Eq. (11) guarantees W Wmmm 2 é', for some e' 0.
Since Wmm is the unique minimizer of (NTP-SVM) and [ | [ |, there exists d E (0, 1) and j [m] such that least one of the following is true: (i) 3z and z'
S, such that
(ee - Whyl 2 (ii) 3ze E S,, v S, such that (e, - e,) Wh, h, s 21-8.
Case (i): Without loss of generality (e, - Ex,) ) TWh; h $ - (otherwise, flip z, z'), Thus, ignoring all
but the (j,2,2))-ten term in (10) and using ly,- E,,2 2 R'8 + log Ps.s
(Paee) gives
R'b
CE(WB) -H t,. log p.. (( 2 n
H
1
log
n
Comparing this (9) for large enough gives that CE(WB) CE (W)), contradiction.
Case (ii): We can assume W e FL, since otherwise we in Case (i). Now, again ignoring all but
the (j, 2) term in the CE loss for which the assumption holds for some v & SJ. we find
CB(WB) R,Pi, log (pra( Σ cllaarlaa)) + eel +1,))) 2'E3,
7

---

Using Ps (WB) W and (4) yields Eves, (. 1 Moreover, by assumption of Case Pz.
(ii): ebw i it' (1-6) the de 12 (1-1), for d == - IWIM
|M Putting together yields:
CE Wn) - H 7ンpリ,= log (1 P3.cC' - 2 R
where the second inequality uses log( + x) 142 I O. Compare this with (9): For large enough B, since R, Ht' grow at the same rate, it holds contradiction. In either case, we arrive at contradiction, which completes the proof.
2n 2 1 > Ce-R, Thus, CE(WB) > CE(W)).
5 Gradient Descent
This section studies the implicit bias of GD. Denote the GD iterates at time k by WN W 7V CE 1-1 for arbitrary initial point Wo and constant step-size > 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 NTPx -oPpaccmpal and NTP NTP-separable, then limg--oo c H and lim W& = c0.
This is intuitive because the CE loss is convex in W (thus, GD approaches the objective's 's infimum HO, 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 Tm Wmmm
Then, it holds that lim
W
= Moreover, limg--s Pa (W+) = W
[
[Wmm]]
The theorem establishes 8 that in the limit of iterations: WE W + |P_ (P,(WW))) W um 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: involves comparing the loss (E(W)) for large iterations & to the loss evaluated at "genie" point that is chosen that: (i) On the subspace F, it agrees with W.. This is because is easy to show that P: (W converges to W by standard gradient descent analysis for convex functions; (ii) On the orthogonal subspace FLL it follows the optimal (with respect to accelerating loss decrease) max- max-margin direction Wmmm E FF To establish the loss comparison, the ideas is to compare the values of the adjusted loss CE2 (W) := CE(W)- CE (P))
We validate our analysis with experiments synthetic data in App. A. For illustration. Fig. 1 shows 2D setting with m distinct contexts, each followed by S, =3 3 tokens/words out of total V 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) colored markers, the geometry of word-embeddings, that is the max- max-Nn--margin vectors (Wmm) "E, v e [5]. to which GD directionally converges. See caption for interpretation and Fig. 2 in the App. for vis. of the finite component of word-embeddinggs on the subspace F. The right subfigure shows results of GD training with respect to training loss, norm growth, alignment with Wum, and convergence to W on F. See App. A for further implementation details and additional experiments.
(Wmm)]ee,
6 Related work
We build on the literature implicit optimization bias of CE loss in one-hot supervised classification. [76] show that for linear models and lincarly-separabll data, GD converges in direction to the max- margin 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 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 finite point. While operationally similar, Thms. 1, 2 cannot
$In line with observations in one-hot 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 (I) vs (2). as the latter involves sampling distinet contexts in each iteration. In this latter case, we also point out that favorable interpolation conditions, such strong-growth (e.g.. [91]), be shown to hold.
8

---

A
234
CE loss
[aal
CE(W))
28
2
1.2
Entropy #
10
26-
W
04
10"
10
10a
10"
10
10
1d
(W,
W W
W. 4.1
W
|2/ W
10"
125
10-2
C
W
B
145
10-2
WF
10
10"
10a
10
10"
10
10
-05
-5.6
Iterations
Iterations
Figure 1: Vis. of NTP implicit optimization bias in setting with m distinct contexts, embedding dimension d = 2, vocabulary of [V] 5 words and support sets of length |S 3,j [3]. Left: Vis. of context embeddings h in circle black markers (marked as B, and of their associated support sets Sj (colored text below each marker). Colored vectors (star markers) represent max-NTP-mane vectors w e, ".mmm, , é [5] found by GD. Interpreting decoder vectors as word embeddings leads intuitive findings on their geometry learned by NTP training. E.g., word embedding w3 (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-representation 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 and of Theorem 2, respectively.
that share the same next-word. Note that geometry
be directly derived from theirs since our setting is neither binary nor one-hot. Nevertheless, proofs extend the foundational work of [68, 34, 37], akin to numerous other studies that explore extensions nonlinear architectures(50, 35, 28, 29, 83, 89], and 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- 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],
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 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 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 unique next token.
Finally, although stemming from different perspectives, the form of our convergence results echoes recent conjecture by [82] regarding implicit optimization bias transformers. Unlike their conjecture, which focuses 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, initiated by [83], is deferred to Appendix B. In contrast this line of work, here focus on the optimization biases of the NTP training- paraiingg itself, which is orthogonal to the intricacies of the specific architecture generating the context embeddings.
1[49, Thm. 4.3] uses [47, Cor. 5.2), which applies to regression on scalar labels; thus is not applicable in NTP,

---

7 Conclusion, limitations and future work
Towards characterizing implicit regularization effects, we highlight two key aspects of NTP training: (i) Formulating as CE optimization over distincl contexts; this long recognized in language modeling (e.g.. [44, 63]) since Shannon's '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
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:
-eP-separyyilitty thresholds: Identifying exact thresholds for NTP-separabiiis under distribu- tional 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-embedding that respect the discrete nature of the underlying token subsequences we believe this direction holds promise for further exploration.
In addition these, essential extensions include relaxing the linearity assumption.
Architecture-speeifii embeddings: A bottom-up approach considering architecture-specae 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 fixed decoder. A challenge in this approach is balancing the restriction to shallow transformers (for ytalytical 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-sequencc architectures, such 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 (5), extends the unconstrained features model, which has recently been employed to investigate neural collapse geometry in one-hot classification settings [56]. This idea offers 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, poten- tially 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.
Acknowledgements
Thank you Tina Behnia, Yize Zhao, Vala Vakilian, and Puneesh Deora for inspiring discussions that contributed to this work and for their valuable suggestions on the manuscript. am also grateful Gautam Goel for his careful reading and for pointing out several typos. Thanks to the anonymous reviewers for their feedback. This work is supported by the NSERC Discovery Grant No. 2021-03677, the Alliance Grant ALLRP 581098-22, 20FFTE-0033-0033 and a CIFAR AI Catalyst Grant.
Parameterr divergence in transformer- based language models has been empirically observed in [53].
10

---

References
[1] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023, URL not//foruu211--0g0XXXHM/My111
[2] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-ave aware minimization. In International Conference on Machine Learning, pages 639-668. PMLR, 2022.
[3] Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameterized nonlinear models. IEEE Transactions on Neural Networks and Learning Systems, 33(12): 7717-7727, 2021.
[4] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. arXiv preprint arXiv:1906.1 11300, 2019.
[5] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prud- nikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04517, 2024.
[6] Mikhail Belkin. The necessity of machine learning theory in mitigating ai risk. ACM/JMS Journal of Data Science, 2024.
Mikhail Belkin, Alexander Rakhlin, and Alexandre B Tsybakov. Does data interpolation contradict statistical optimality? arXiv preprint arXiv: 18X6.09806.0941,,, 2018.
Samy Bengio and Yoshua Bengio. Taking on the curse of dimensionality in joint distributions using neural networks. IEEE Transactions on Neural Networks, 11(3):500-557, 2000.
[9] Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. Advances in neural information processing systems, 13, 2000.
[10] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx. Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.072555 2021.
[11] Burak Çakmak, Yue M Lu, and Manfred Opper. A convergence analysis of approximate message passing with non-separable functions and applications to multi-class classification. arXiv preprint arXiv:2402.08670 2024.
[12] Emmanuel Candès and Pragya Sur. The phase transition for the existence of the maximum likelihood estimate in high-dimensional logistic regression, arXiv preprint arXiv:1804.09755 2018.
[13] Yuan Cao, Quanquan Gu, and Mikhail Belkin. Risk bounds for over- pwrrpparamterrized maximum margin classification sub-gat -o--gaussian mixtures. Advances in Neural Information Processing Systems, 34:8407-8418, 2021.
[14] Matias D Cattaneo, Jason M Klusowski, and Boris Shigida. On the implicit bias of adam. arXiv preprint arXiv:2309.00073 2023.
[15] Niladri S Chatterji, Philip M Long, and Peter L Bartlett. When does gradient descent with logistic loss find interpolating two-layer networks? The Journal of Machine Learning Research, 22(1):1135-7182, 2021.
[16] Jinghui Chen, Yuan Cao, and Quanquan Gu. Benign overfitting in adversarially robust linear classification. In Uncertainty in Artificial Intelligence, pages 313-323, PMLR, 2023.
[17] Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. arXiv preprint arXiv:2402. 0400024000044084, 2024.

---

[18] Katherine M Collins, Umang Bhatt, and Adrian Weller. Eliciting and learning with soft labels from every annotator. In Proceedings of the AAAI Conference on Human Computation and Crowdsourcing, volume 10, pages 40-52, 2022.
[19] Elisabetta Cornacchia, Francesca Mignacco, Rodrigo Veiga, Cédric Gerbelot, Bruno Loureiro, and Lenka Zdeborová. Learning curves for the multi-class teacher-student perceptron. Machine Learning: Science and Technology. 4(1)::15015) 2023.
[20] Thomas M Cover. Geometrical and statistical properties of systems of linear inequalities with applications in pattern recognition. IEEE transactions electronic computers, pages 326-334, 1965.
[21] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449-27461, 2021.
[22] Zeyu Deng, Abla Kammoun, and Christos Thrampoulidis. A model of double descent for dighnnimensionll binary linear classification. Information and Inference: A Journal of the IMA, 11(2):435 495, 2022.
[23] Konstantin Donhauser, Nicolo Ruggeri, Stefan Stojanovic, and Fanny Yang. Fast rates for noisy interpolation require rethinking the effect of inductive bias. In International Conference Machine Learning, pages 5397-5428. PMLR, 2022.
[24] Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. arXiv preprint 1rXiv:2110.10090 2021.
[25] Benjamin Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv e-prints, pages arX2v-2402 2024.
[26] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.1400 14052, 2022.
[27] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. 2023.
[28] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages 1832-1841. PMLR, 2018.
[29] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems, 11:9461-9477 2018.
[30] Peter E Hart, David G Stork, and Richard O Duda. Pattern classification. Wiley Hoboken, 2000.
[31] Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J Tibshirani. Surprises in high-dimensional ridgeless least squares interpolation. arXiv preprint 183.00990..08560, 2019.
[32] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv: 15XXiv11003.000 02531, 2015.
[33] Adel Javanmard and Mahdi Soltanolkotabi. Precise statistical analysis of classification accura- cies for adversarial training. The Annals of Statistics, 50(4):2127-2156, 2022.
[34] Ziwei and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXiv: /00::8003000 803.07300, 2018.
[35] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment deep learning. Advances in Neural Information Processing Systems, 33:17176-17186, 2020.
12

---

[36] Ziwei Ji and Matus Telgarsky. Characterizing the implicit bias via a primal analysis. In Algorithmic Learning Theory, pages 772-804, PMLR, 2021.
[37] Ziwei Ji, Miroslav Dudík, Robert Schapire, and Matus Telgarsky. Gradient descent follows the regularization path for general losses. In Conference on Learning Theory, pages 2109-210000000 2136. PMLR, 2020.
[38] Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference on Machine Learning, pages
PMLR, 2021.
[39] Daniel Jurafsky and James H. Martin. Speech and Language Processing. Draft, 3 edition, 2023. URI. https: /wwb. stanford. staanforrddddu/"/uaaastt/////// httpe/afskkk/slpp////ooooo/o pdf.
[40] Tokio Kajitsuka and Issei Sato. Are transformers with one layer self aelf-attention using low-rank weight matrices universal approximators? 2024.
[41] Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. 2023.
[42] Ganesh Ramachandra Kini, Orestis Paraskevas, Samet Oymak, and Christos Thrampoulidis. Label-imbalanced and group-sensitive classification under overparameterization. Advances in Neural Information Processing Systems, 34:18970-18983, 2021.
[43] Frederic Koehler, Lijia Zhou, Danica Sutherland, and Nathan Srebro. Uniform convergence of interpolators: Gaussian width, norm bounds and benign overfitting, Advances in Neural Information Processing Systems, 34:20657-20668, 2021.
[44] Omer Levy and Yoav Goldberg. Neural word embedding as implicit matrix factorization. Advances in neural information processing systems, 27, 2014.
[45] Yingcong Li, M. Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning, 2023.
[46] Yingcong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pages 685-693. PMLR, 2024.
[47] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss? mathematical framework. arXiv preprini 1rXiv:21100069) 2021.
[48] Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller, On the expressive power of self-attention matrices, 2021.
[49] Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on PMLR, 2023.
Machine Learning, pages
[50] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020.
[51] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: framework for principled analysis of transformers via markov chains. arXiv preprini arXiv:2402.04410 04161, 2024.
[52] Eran Malach. Auto-regressive next-token predictors are universal learners. arXiv preprint ND0:23309066799 2023.
[53] William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah Smith. Effects of parameter norm growth during transformer training: Inductive bias from gradient descent. In Proceedings of the 2021 Conference on Empirical Methods Natural Language Processing, pages 1766-1781, 2021.
[54] Francesca Mignacco, Florent Krzakala, Yue M Lu, and Lenka Zdeborová. The role of regularization classification of high-dimensional noisy gaussian mixture. arXiv preprint 1rXiv:2002.11544, 2020.
13

---

[55] Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations vector space. arXiv preprint arXiv:1301.3781, 2013.
[56] Dustin Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint 1rXiv:201
2020.
[57] Andrea Montanari, Feng Ruan, Youngtak Sohn, and Jun Yan. The generalization error of max-margin linear classifiers: High- nigh-dimensional asymptotics in the overparametrized regime. arXiv preprint arXiv:1911.0154: 2019.
[58] Vidya Muthukumar, Adhyyan Narang, Vignesh Subramanian, Mikhail Belkin, Daniel Hsu, and Anant Sahai. Classification vs regression in overparameterized regimes: Does the loss function matter? arXiv preprint 1xXiv:2005.08054 2020.
[59] Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with fixed learning rate. arXiv preprint 1rX.v:1806.0179, 2018.
[60] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019.
[61] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.660 2014,
[62] OpenAL. Openai: Introducing chatgpt, 2022. URL 2022.
wtt.llcccccclcc/eo chatppt,
[63] Jeffrey Pennington,
Global vectors for word representation. In Proceedings of the 2014 conference empirical methods in natural language processing (EMNLP), pages 1532-1543, 2014.
Richard Socher, and Christopher Manning. Glove:
[64] Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity, Advances in Neural Information Processing Systems, 34:29218-29233 2021.
[65] Joshua C Peterson, Ruairidh M Battleday, Thomas L Griffiths, and Olga Russakovsky. Human uncertainty makes classification more robust. Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9617-9626, 2019.
[66] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. OpenAl blog, 2018
[67] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[68] Saharon Rosset, Ji Zhu, and Trevor Hastie. Margin maximizing loss functions. In NIPS, 2003.
[69] Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. International Conference on Machine Learning, pages 8346-8356. PMLR, 2020.
[70] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. International Conference on Machine Learning, 2022.
[71] Fariborz Salchi, Ehsan Abbasi, and Babak Hassibi. precise analysis of phasemax in phase retrieval. In 2018 IEEE International Symposium on Information Theory (ISIT), pages 976-980, IEEE, 2018.
[72] Ohad Shamir. The implicit bias of benign overfitting. In Conference on Learning Theory, pages 448-478 PMLR, 2022.
14

---

[73] Claude E Shannon. Prediction and entropy of printed english. Bell system technical journal, 30(1):50-64, 64, 1951.
[74] Claude Elwood Shannon, A mathematical theory of communication. The Bell system technical journal, 27(3):379-423, 1948.
[75] Viktoriia Sharmanska, Daniel Hernández-Lobato Jose Miguel Hermandez-Lobato and Novi Quadrianto. Ambiguity helps: Classification with disagreements in crowdsourced annotations. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 2194-2202, 2016.
[76] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):2222-2877) 2018.
[77] Haoyuan Sun, Kwangjun Ahn, Christos Thrampoulidis, and Navid Azizan. Mirror descent maximizes generalized margin and can be implemented efficiently. Advances in Neural Information Processing Systems, 35:3 19391089-11011. 2022.
[78] Pragya Sur and Emmanuel J Candès. A modern maximun-likelihood theory for high- dimensional logistic regression. Proceedings of the National Academy of Sciences, page 2018 201810420, 2019.
[79] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 2818-2826, 2016.
[80] Hossein Taheri, Ramtin Pedarsani, and Christos Thrampoulidis. Asymptotic behavior of adversarial training in binary linear classification. IEEE Transactions on Neural Networks and Learning Systems, 2023.
[81] Kai Tan and Pierre C Bellec. Multinomial logistic regression: Asymptotic normality on null covariates in high-dimensions. arXiv preprint urXiv:2305.17822 2023.
[82] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Trans- formers as support vector machines, 2023.
[83] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023.
[84] Yuandong Tian, Yiping Wang. Beidi Chen, and Simon Du Scan and snap: Understanding training dynamics and token composition in l-layer transformer, 2023.
[85] Yuandong Tian, Yiping Wang. Zhenyu Zhang, Beidi Chen, and Simon Du Joma: Demys- tifying multilayer transformers via joint dynamics of mip and attention. arXiv preprint arXiv:2310.00535, 2023.
[86] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
[87] Nilesh Tripuraneni, Ben Adlam, and Jeffrey Pennington. Overparameterizzaton improves robustness to covariate shift in high dimensions. Advances in Neural Information Processing Systems, 34:1
2021.
[88] Vladimir N Vapnik and Alexey Ya Chervonenkis. A note on one class perceptrons. Automa- tion and Remote Control, 25:774-780, 1964.
[89] Bhavya Vasudeva, Puncesh Deora, and Christos Thrampoulidis. Implicit bias and fast conver- gence rates for self-attention. arXiv preprint 4rXiv:2402.05738 2024.
[90] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
15

---

[91] Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of sgd for over- parameterized models and accelerated perceptron. The 22nd international conference on artificial intelligence and statistics, pages 199--1204. PMLR, 2019.
[92] R. Vershynin. Lectures in geometric functionnl analysis. Unpublished manuscript. Available at http://www-personal. umich. cdtb/ronnnv/aprrsGG---ooooCC
pdf. 2011.
[93] Johannes von Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ArXiv, abs/2212.07677, 2022.
[94] David Wu and Anant Sahai. Precise asymptotic generalization for multiclass classification with overparameterized linear models. Advances in Neural Information Processing Systems, 36, 2024.
[95] Shuo Xie and Zhiyuan Li. Implicit bias of adamw: Vell_iinfty norm constrained optimization. arXiv preprint arXiv:2404. 04454, 2024.
[96] Chenyang Zhang, Difan Zou, and Yuan Cao. The implicit bias of adam on separable data. arXiv preprint 24Xi:24406.166000 2024.
[97] Chiyuan Zhang. Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization, 2017.
[98] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint rrXiv:771000941., 2017.
[99] Ruiqi Zhang, Spencer Frei, and Peter L.. Bartlett. Trained transformers learn linear models in-context, 2023.
[100] Lijia Zhou, Danica Sutherland, and Nati Srebro. On uniform convergence and low-norm interpolation learning. Advances in Neural Information Processing Systems, 33:6867-6877, 2020.
16

---

234
A 234
B
44
B 45
Maide
Figure 2: Same setup as Fig. 1. Left: Matrix P of conditional probabilities words (cols.) per context (rows). Each row corresponds to the conditional probability vectors Pj.J e [m]] Black entries correspond to off-support words. Middle: Shown as w. z [5]. the rows of the NTP-SVM solution Wmom which GD directionally converges. Right: Shown w2, [5], the the finite parameter W to which GD iterates projected on F converge to. The geometry of Wmm depends only on the support-set of P. On the other hand, the geometry of W depends on the entries of P for in-support tokens/words, As seen from visualization of P. the words 1 and 5 have the same support pattern (i.e.. both follow the same contexts A and B). Thus, wi ww in the Middle plot. However, on the subspace F corresponding to the Right plot, wi + Ws, which allows matching the different conditional probabilities with which each follows contexts and B.
of
A Experiments
All experiments were conducted on 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 made publicly available on Github the final version of the paper.
A.1 Additional details on 2D example of Fig.
Figure illustrates toy 2d example where the embeddings and the hyperplanes defined by each row of Wmum can be visualized. We used d 2, m 3, V and S1 Sg S3 3. The support sets of each embedding are shown in the figure color-coded match the respective decoder hyperplane. Probabilities are assigned randomly. The empirical conditional entropy evaluates to H. 0.8811 and the matrix of conditional probabilities visualized Figure 2. the same figure, we also visualize the rows of the directional component W (Middle) and of the finite component W (Right). Interpreting the V d decoder matrix the matrix of learned word embeddings, this provides visualization of their geometry. As per our results, the two word-embedding matrices W and Wm 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.
A.2 Overparameterized setting
We examine the implicit bias of GD on NTP training with overparameterizzaton synthetic data generated as follows. We construct dataset with n 5000 sequences involving m 50 distinct contexts. Each distinct context gets mapped to randomly generated embedding of dimension 60 > m. We set vocabulary size
and each context j [m] is followed by S, = 6, V.j [m] possible next-tokens. The support sets S₁ c V and the probabilities Pj.z, = Sj are chosen randomly: see Fig. 3 for representative examples from training dataset. For fixed realization of the dataset (for which H 1.445nats), initialization. For GD, we use learning rate 77 = 0.5 and for NGD and Adam 71 = 0.01. For Adam, also set 31 0.9, 82 0.99. We run all algorithms for le4 iterations. For each case, plot the following as a function of iterations:
GD, normalized GD (NGD), and Adam from random LeCun
1. Upper Left: CE loss versus entropy lower bound
2. Upper Right: parameter norm growth
17

---

3. Lower Left: correlation of Wm with iterates Wt and of "corrected" iterates We W after substracting the component
4. Lower Right: convergence the subspace component W (W....)
Fig. 4 shows an instance of these. As predicted by analysis, in this overparameterized setting: CE loss converges to its lower-bound, parameter norm increases, iterates align in direction with Wmm and the subspace component converges to W
P
0.4
0.4
0.3
0.2
0.2
0.2
0.1
0
0
0
8910
436
900
5
910
34
6.78
0.6
P
0.6
Pg
0.8
0.6
0.3
0.4
0.2
0.2
0.2
0.1
2
56
8810
234
300
12
5
78
34
67
900
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 Sj 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 consistent with analogous findings (rigorous in that case) for one-hot classification [60, 36).
CE loss
[W1
CR(W,)
Eutroger
10
5
1.5
0
10
10a
10
10"
10P
10
1
(W, W": W
W.r W 1
10"
W W
WP
10
10
10P
10
10a
10"
10F
10
Iterations
Iterations
Figure 4: Experimental illustration of the implicit bias of GD in NTP over synthetic data with overparameterizztion.. See App. for detailed description of the experimental setting. The upper two graphs confinm the predictions of Lemma 2, while the lower two graphs adhere to the predictions of Theorem 2.
18

---

CM
CR((s))
2.6
2.5
Eonnropyy X
15
10
5
10
1.5
109
10
10'
10
10P
10c
101
10'
10P
10"
107
corr (W,- W W
[W.- W/ W/ WT
102
3
-oort (Wa W., Wamme
1Wa
10
W
W
10-1
10a
10 Iterations
10
10'
10
10P
10
10a
10P
10c
10*
NGD
Iterations
Iteeations
Adam
Toerations
Figure 5: Implicit bias of normalized GD (Left) and of Adam (Right) in NTP over synthetic data with overparameterization. Both exhibit the same implicit bias, but converge faster than GD, with Adam being slightly faster than NGD.
B Additional related work
Implicit bias in transformers. As already mentioned 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 them, the NTP setting that we study involves predictions over multiple categories and is not one-hot. Also, while they fix the decoder, here, fix the embeddings. In these respects their results are rather different. More similarities arise when [82] replace the linear decoder with 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 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 setting to prove an analogous statement rigorously. Eventually, as we move lower in our top-down approach and consider architecture-ssecific 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 MIPss utilizing single heads instead of multiple, and fixing certain parts while training only subset of trainable parameters. Most of these studies have focused on classical one-hot supervised settings, and only a handful (c.g., [84, 85]) have secked extending these 'bottom-up' analyses to NTP settings. Yet, their primary emphasis remains 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 prediction scenarios. Unlike our study which utilizes the widely adopted CE loss, their approach based on log-loss, which renders the training loss convex, similarity shared with our model despite the inclusion of self-attention. Both our results and those of Li al. substantiate the conjecture posited by Tarzanagh and colleagues [82], albeit in very distinct settings. Notably, contrary to both
with fixed linear decoder in next-token
19

---

[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-soureing [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 improved performance
terms of model generalization, calibration, and robustness to out-of-ue of-distributtin data. To the best of our knowledge, no prior work has investigated the implicit bias of gradient descent this or other soft-label classification settings; thus, our results are of direct relevance to these contexts as well.
C Autoregressive setting
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) setring with minimal changes. This also emphasizes the role played in our results by the sequence length.
As pointed in (I), the full autoregressive NTP objective averages I individual losses (without loss of generality assume sequences of equal maximum length 7). In order to make our analysis applicable, first need to express (1) in terms of unique contexts. Mirroring the notations in Sec. 2, define the following for t [T -T--1 1]:
mz, f. [T - 1] is the number of distinct contexts of size t. Note that m, m2 2 ... mT-1-
- m, is the total number of distinct contexts in the dataset
hr,z h e []_ e [mæ] is the embedding of the j-th (among all t-long contexts)
distinct context F2,--
åz,t is the empirical probability of æz.--
Pya, is the empirical probability that context æy, is followed by token 2 e V.
Sj,, is the support set of the next-token distribution of context Fz,t-
With this notation, the NTP objective becomes
CE Σ Σ tnd Σ Pr.j log (WWh))) 205;.
EE+ te[T [<{T-]]j((ma]
To continue enumerate the multi-set I= =i= (i,t)lte [T 1].j [m+]]. We may then rewrite the above as
CE n Σ Pr. log log (s.
VEI 26S:
At this point note that this is of identical form to (2). Consequently, the definitions (c.g., NTP- separability, 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 captured by Shannon's 's inequality,
H. Heal, where H. =
Σ Σ je[ma] 2ES!
TesPr log(pr..z))
reflecting the incremental reduction in entropy sequence length increases.
20

---

D Notations
Throughout, lowercase and uppercase bold letters (e.g., a and A) represent vectors and matrices, respectively. (:-)) and I-ll denote Euclidean inner product and norm, respectively. For matrix A, denote its pseudoinverse .A". All logarithms natural logarithms (base e). We denote e. the e-th standard basis vector in R" V-1 denotes the V-dimensional unit simplex and S(): RVV + A V-1 the softmax map:
e
S(a) = [S))
with (a(a)= = Every E
As explained in Section we represent a training set as Tm {h2, n Pg. We assume that embeddings are bounded and denote
M 22 max hyll
5c(m]
Given Tm, let
F =pan (((e=-ez)) ez 2 E j [m]}]
a subspace of matrices and F its orthogonal complement. Denote Py, Pı the orthogonal projections onto F and Ft, respectively. For convenience, for W Rxdd we denote
W5 = PW
and W1 WI-P((W)
Define
CEr(W)+
Σ ū [=[m]
pJ. log 1+ Σ (eseee)) Wh, Σ 265,
(12)
242
Clearly, for all W RVxd it holds CF(W) CF, (W). Note also that for all W F and for all w e 3 that satisfy Eq. (6a), it holds CEz(W) limg--oo CE(W RWd). Thus, under NTP compatibility and NTP separability,
inf WeF
CEy(W) = inf CE(W) H.
(13)
W
E
Proofs
E.1
Gradient Descent
Throughout we assume GD is ran with step-size 17 s 1/(2L) where L is the smoothness of CE loss. This condition is not explicitly mentioned thereafter.
E.1.1 Auxiliary Lemmata
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 CE(WE)) = inf CE W) R=
W
and also lim Wxll
In the lemma below, we collect some useful and simple-to-show properties of the GD and regular- ization paths. These adaptations of corresponding results for one-hot binary classification general non-separable data established in [34].
Lemma Suppose conditions (6) hold for some Wd Also, that there exists WP W F satisfying condition (4), The following hold:
21

---

1. (Ey(W") wey CEy W) = H.
2. W" the unique minimizer of CEp on the subspace F,
3. limg-ro Py(W = W where Wt are GD iterates,
4. oi [P)(W 00,
5. lim Py (WB) W' where WB is the reguarlized solution (8).
6. p PP W B)] 00.
Proof. It is easy to check by direct substitution of W in (12) and use (4) that (Eg(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 different minimizer W * W1 F. Then, since (E>(W1) H. it also holds for WR == W1 + RW4 that lim CE( (RR) H. In turn, this implies for all j [m]]:
lim (WWw), Pries E R
and
lim R
(
The first condition gives then that W, must satisfy (4). Since W also satisfies these equations, denoting WA W W1 + 0, it holds:
(Wa (e, - ex) Vje [m], e
But WA F. this forms contradiction. Hence, W is unique solution in F of (4) and unique minimizer of CE. on the subspace T.
The proof of the third claim follows the same way the proof of part (1) of Thm. 15 of [37]. For completeness: follows by the lemma's assumptions and Lemma that limgg-oo CE(W: H. Combining with the first claim of the lemma yields lim CF(W;) CEg(W*) Since CWx) $ CE(Ws), this finally gives
lim CEs CEs (W:) 20000
lim CE. (P=(WL)) CEy(W").
A DO
Since W is unique by the second claim, the desired then follows.
For the fourth claim, recall from Lemma 3 that limg--oo [&ll oo. From the previous claim, also
have lim |P (W+) < C for some constant C [' Thus, the desired follows by applying
the fact that W = |P (&)l [[P.
The proof of the last two claim is exactly same as that of the third and fourth claim. Only now use the facts that p- CF(WB) H and lim [Wn oo (see proof of Theorem 1).
E.1.2 Key Lemma
Lemma Let Wa denote the GD iterate at iteration Recall the decomposition W.= PP(WW))
Pi(We)= Wk.y WK.1- Fix any (0. 1). There exists large enough R = R(a) and ku ku ku(R)
such that for any le 2 Ka, it holds that W 2 R and
CE WR.3 + (1 a) Wt., Wimm CE(WW).
(14)
Proof. We drop the subscript & lighten notation.
First, note by Lemma 4.D that, for arbitrary R, we can pick kn = ki (R) such that for all k: 2 k; [Wl > R.
22

---

Thus next, we will prove the main claim, i.c. for large enough WWll inequality (14) holds. Denote W.
R²'
Substituting in CE expression (2), and using the fact that Wmm g by (6a) yield:
W
CE
W + RI+m)
Σ ]d[m]
j
Σ 2
Pj.e log
Σ
W(ooo + Σ
( + Σ (ltn) Whihh
(e,-o,))
z'e5,
05S,
085,
Σ ]4[m]
J
Σ z
Pj.z2 log
Σ (e:-t)) Waah
Σ
(La)) Wh
(15)
DEV
Moreover, decomposing W- W3 + W1, and defining
W
W
Wi
W
W
R
have
CE CE(W)= = Σ
ū Σ
p3. e log
Σ
(e==ey)"Wssh
+ Σ
(ere--o) + Σ
-l((es-e)) Wiih
s4[m]
xc5,
z'eS,
015,
015,
Σ ]c[m]
R Σ S
pa. log
(c.-oo))) "We,, Σ e4S
Wyh
e
(t. W.h:
(16)
NcV
where we used that, by definition, W1 F Thus, our goal becomes showing (15) s (16), for large enough R. To do this, we consider two cases as follows below.
For the remaining of the proof recall M == maxje[m] v h, and use the logits shorthand:
T,,00 e hi
and
grum 1.0 e, W hj.
Case 1: W1 is well aligned with W Suppose
W Wilse== a M yields + eM, Vj E [m]] S j,0
(17)
Using this, linearity of logits, and
E,,2 Z, S j,z
Thus,
Σ (e. )..)
W,,, WW, 2 er
Σ
ee-R'e a.) wmmp (mmm- an' Σ of
VFS,
U4S
Also recall by feasibility of Wmm that
pmm
amm
2 1, [m], 2 U 1.0
(18)
Thus,
Σ (e. Wih,, aR Σ e " Wimmh
etS;
D$S,
Comparing the above two displays yields
Σ
e (ee.)) s Σ e Wih,,
e
o4S,
which implies the desired (15)s(16) for any value of R' (eqv. W.I..
Case 2: No alignment. Suppose now that (17) does not hold. Note that W [ and since (NTPPSVM) has a unique solution must be that W is not feasible, But W, F thus satisfies the equality constraints. This then means that there exist ö == ö(c) and j. [m], v. S,. such that
T, Tn $ 1 31-6. 8, Vo Sy..
(19)
23

---

(Note the above holds for all 2 é Sj. because T, - since W FI.) To continue, we introduce the shorthand notation
A,,2 == (W) Σ c (es-ees)'Wsh
vc
as well as
Amin
min Aj.s Je[m]].asS,
and
Amax max Aj,z je[m]],eS,
Using (19) we may lower bound (16) follows:
e (el W. Baja
CE(W)-
Σ
ūj Σ
Pia log Σ
(e,-ew) WW,h,
fj..
Σ ze5]
Pioe log
1+
Ajaiz
]=[m]
=65;
key
1+ c (-R((--6)
jT. Σ S
PJ.e log
Amax
Rl' (1-6)
(20)
TI( Amax + 1)
where in the last line we used ft. 2 1/n, Vj [[m] as well as log(1 (1 + x) 2 On the other hand, using property (18) for max-margin logits, we can upper bound (15) as follows:
__ x > 0.
V (R
CE (W + +1
Σ cc[m]
Pi. log Σ e (er-eee) log Σ zc5,
Amin
ney
V Ve =
(21)
Amin
where in the last line we used log(1 + x) x 0. In view of the last displays, it suffices that
Rt' (1+0))
e (-R((-0))
mV (Amax + 1)
R log
Amin
"m(mmxx* + 1)
8
Amin
All it remains is obtaining bounds for Amin, Amax specifically showing that they do not depend on R.
By Cauchy-Schwartz:
Ve-M Ve-M $ Amin S Amax S MIess
Further recall by Lemma 4.C that l is large enough then
[W W I [|
|Wa] s 2/ W
(22)
Thus, there exists k, k. (I) W.Wi []) such that for all k k.:
Ve 2.. < Amin S Amax V W...11
Hence, the desired (21)s(20) holds provided
Wmmm| W
log (2
(23)
Set R= R(a)) (RHS of (23)) and ko (Rt) max ( (R), k. We have shown this guarantees for
all k 2 ki:: ||W 2 Rt and by choice of Rt also (21)s(20). This in turn implies (15)s(16), as desired complete the proof.
E.1.3 Proof of Theorem 2
For the subspace component, see Lemma 4.C. For the directional convergence, the key ingredient of the proof is Lemma After that, the proof follows identically Thm. 15(2) in [37]. We include the details for completeness, but there are novel aspects in the rest of this section.
24

---

Let any e 10, and choose = c/(1 -). c). By Lemma 5, there exists ko such that for any lo 2 ko-
we have
[Wa,l 2 max (a(()(())))))) 1/2)
and
CEWW), Wk., a1 Wmm) = (V We (Wa,g (1+ -
W Wmm)
2 -E (1+ W W 0,
where we also used convexity of the loss.
Consequently,
(W k+1 W. W
= CB(W), Wmm
2 (1 WA,1)
2 (11-)(WWk-L) K+1,1 WR. Wx,1)
2 (1-c)(Wkk))) 1+1,1 WK.1, WK.1)
(1-c) (I
(.Wlll Wk+1.1 Wa.1l²)
2 Wa.1
2 (1 k W 27 2n( CE( CCE(W+...))
where the last step used W 1/2, the fact that x²² - y2² 2y(x v), Vx, y and smoothness of the CE loss. Telescoping the above expression and rearranging yields
Wr.1l
(Weg>
W (mm) (1 6) Who (K.)).))
(WA k Wmmm mmm
2 (1-
W
[W&ll
W.,llll + (Was kot Wmm) (1 e) Wkk. CE(
2 (1-c)
WWall
Now recall from Lemma 4 that lim | oo and oi Wi
|W Thus,
lim inf, onf (W ks Wum 21-e. Since e arbitrary, the desired follows.
E.2 Regularization Path
We provide detailed proof of Theorem filling in missing details from the proof sketch in the main paper.
E.2.1 Proof Theorem
First, we show that WB on the boundary, i.e. [BB B. Suppose not, then (V (BB)B).UU (U-=1 for all U RVxd Using the CE expression in (2) and few algebraic manipulations, yields
(-V (C(WWB))BUU
Σ n jc[m]
Σ S
P..=
Σ 2448
83,2 (e. - Uh; + Σ
$j. (e: - c. Uh,). (24)
uRS,
where denote the output probabilities WB B S, S. W Whh,), v V.j [m]. Choose U = Wmmm in (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 W + R(B) Wmm where, W e F satisfies (4), and R(B) is chosen such that W B. Concretely, for [W set
B I
W
Note also that R/B - Wm B We will show that W attains small CE loss B (hence, R) grows. To do this, denote for convenience the logits for all v E V,j [m]
G ee w. and fyv == e mm h
25

---

and note that w. = 2,0 + R 3,0'
By using (4) and (6a):
Σ
+G - Re
Dj
Moreover, using (6b)
Σ
e- (G Rey, 11.:--10 -- Re
sc- Σ e (G
$ CE
btS
vtS,
where define constant (independent of R) C == Vel W IM for M == v2. max3//[m] [hgl-
Combining the above displays and using in Eq. (2), yields
log _
CE(W S Σ J
Σ z6S;
Ps.e log
+ +c
s Σ ūj Σ
Di
je[m]
jm(ml
265
PJ.2
H Ce
(25)
where, the second line uses log(1 + x) $ x > 0, and the third line uses f..P3. are probabilities.
Next, towards arriving at a contradiction, we will show that if Wp is not in the direction of Wmm, then it incurs loss that larger than CE( (W)). Concretely, assuming the statement of the theorem is not true, we we will upper bound
PJ.= Sz (wWn,,,) h,))))
CE W B)
H
Σ ]e[m]
Σ Ds.e log 265;
(26)
By our assumption, there exists O, such that there exists arbitrarily large B satisfying:
mm
W >6
(27)
B
Define
1
W
W W RBB
where, It' =R= RB) > 0 is chosen so that W set
mm
Concretely, for large enough 2 2| W
R'
B 2B W B, W + 11.
W mm
Note that it holds lim B- R' R = 1/ Wmm Thus, we can always choose large enough so that
Eq. (27) guarantees W Wmm 2 é', for some > O. Since Wmmm is the unique minimizer of
(NTP-SVM) and [W| [ it follows that there exists 8 E (0, 1) and j e [m]] such that at least one of the following true
(i) 32 and 2' # e S, such that
le. (er) (er Why 20,
(28)
(ii) 3z Uj u Sy such that
(e, - (e_ $ 1.
(29)
Case (i): Without loss of generality (e. -z)) Wh; s (otherwise, flip 2, z²). Thus, ignoring all but one term in (26) gives
pj.= (<
CE Wn) H R. log
2 カ log
(30)
where we use ly.0 = e, WBh,, v E to denote logits of W Using (4) and (28), yields
8 - l,,5 = (ez -e))-e)) (R' W h R'5 + log P
Dia
26

---

Put in (26) and using PJ.= 2 k,Ps,s 2 1/n shows
1
R26
CF(WB) (B) 2 2 log
n
Compare this with (25). For large enough B, it is clear that RyPy, log (iy coR > Ce-R, Thus, CE (WBB)) > CE Wi), a contradiction.
Case (ii): We can assume W E F1, since otherwise we are Case (i). Now, again ignoring all but the (j, 2) term in the CE loss for which (29) holds for some v $ Sj. we find
CE (B) - H R,Pj,a log [ Σ
e(),)16)))))) + c(()))))))))
2'ES,
Using PT(WB) W yields
Σ ell 2e
Σ
PJ.
2'eS,
P3.z
Pj.²
Moreover, by (29):
e. 2 R- (1-5) -,3 2 2e Rt (2-6)
for constant (independent of B3 c': == I Putting the above together yield:
d'E-R(
CE - , log D (((0)))
2n1
x where the second inequality uses log(1 1 + x) 2
14. x > 0.
Compare this with (25). For large enough B, (recall R, R' grow at the same rate) it holds 2n² c > Ce-R, R Thus, (E(WB)> > CE(W;;) a contradiction.
In either case, we arrive at a contradiction, which completes the proof.
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Sec. 2-5.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Sec.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approxinations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or a speech-to-tee system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
28

---

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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper 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 the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless 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 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 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 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 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 NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, may be that access the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
29

---

Answer: [No]
Justification: As mentioned in Sec. A: the code is straightforward to reproduce, following the detailed specifications provided in the same section. For completeness, the code will be made publicly available online the final version of the paper.
Guidelines: As mentioned in Sec.
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (kttps://nips.ce public/guides: weeeddeeeoooeubbisseeeeeeeeeeeee wwsssonoolcyyyl for more details.
While we encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips..... wwtppppwlccggceeeeeeeeeooee httgggudddd//oduuummttttc Set////ooeeuisssseeeeeee ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Sec. A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We use deterministic full-batch optimization in the experiments starting from zero initialization.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
30

---

The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have a 96% CI, the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Sec. A
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / h/weurirss...//// cc
Answer: [Yes]
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 deviation from the Code of Ethics.
The authors should make sure preserve anonymity (e.g., if there is special considera- tion due to laws or regulations in their jurisdiction).
10. 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 enable better understanding of operating regimes of language models with respect to optimization, generalization and robustness.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does 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.
31

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there 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 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 addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse 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 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 encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [NA]
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URI..
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 pack- age should be provided. For popular datasets, paperswithcode. thcode. com/ 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- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
32

---

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 obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
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 the institution conducting the review.
33


Reviews:
[{"review_id": "4q4R2OWFrz", "reviewer": "Reviewer_YBdf", "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.", "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?\n\nMechanics of Next Token Prediction with Self-Attention\nYingcong Li, Yixiao Huang, M. Emrullah Ildiz, Ankit Singh Rawat, Samet Oymak\n\nMax-Margin Token Selection in Attention Mechanism\nDavoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, Samet Oymak\n\nTransformers as Support Vector Machines\nDavoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, Samet Oymak\n\nQ.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": "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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This study enables deriving the data-entropy lower bound in NTP for understanding the optimization and generalization properties of NTP models.", "weaknesses": "The study's focus on linear models  analyzing CE loss for NTP may limit its novelty and applicability, making its contributions to the field appear unclear compared to existing research.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. \n\nBelow, 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. \n\nWe hope these clarifications lead to a re-evaluation of your score!\n\n**Q1: Comparison to works on implicit bias in transformers:**\n\nWe have discussed these references in detail in the submission. For your convenience, we summarize them below.\n\n[44] Mechanics of Next Token Prediction with Self-Attention - Li et al. \n\n[79] Max-Margin Token Selection in Attention Mechanism - Tarzanagh et al. \n\n[80] Transformers as Support Vector Machines - Tarzanagh et al.\n\nIn 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:\n\n“*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 \nfrom 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.*”\n\nAdditionally, Lines 677-684 compare our work to Li et al. [44]: \n\n“*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.*” \n\nWe 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. \n\n**Q2: Could you explain why it is still valid to disregard attention in your analysis?**\n\nOur 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.\n\nBy 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).\n\nOur 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.\n\n**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)?**\n\nThere 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.\nIf 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]$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ldp7YiBmwf", "reviewer": "Reviewer_6D5d", "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.", "questions": "Just the remark in the weaknesses.", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is generally well written, and the notation is very clear.\n\n- The paper provides a a very interesting starting point for studying the solutions found by gradient descent in NTP settings", "weaknesses": "While the paper provides a a very interesting starting point for studying the solutions found by gradient descent in NTP settings, it's not very clear whether margin maximization practically corresponds to any meaningful takeaway in language modeling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and for the positive feedback and score. \n\n**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.**\n\nWe 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.\n\nA 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. \n\nIt 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.\n\nAdditionally, 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].\n\nOverall, 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. \n\n[LV14] Neural Word Embedding as Implicit Matrix Factorization\n\n[Aro+16] A latent variable model approach to psi-based embeddings\n\n[MBV17] All-but-the-top: Simple and effective post processing for word representations\n\n[Mur+23] Grokking of Hierarchical Structure in Vanilla Transformers\n\n[Moh+23] Grokking modular arithmetic can be explained by margin maximization\n\n[Mor+24] Feature emergence via margin maximization: case studies in algebraic tasks", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ciVXvf4wb0", "reviewer": "Reviewer_kgdu", "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.", "questions": "A minor suggestion: lines 32-34 appear to require rephrasing.", "limitations": "The authors thoroughly discuss the limitations of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work studies a timely topic (next token prediction) and approaches it from a learning theoretic perspective (implicit bias of optimization), which has proven to be very fruitful in \"standard\" classification. The assumption of sparse contexts is clever and should be of wider applicability. The results are novel and analogous to similar results that were proven for \"standard\" classification. Furthermore, the presentation is comprehensive, with many pointers to related work, which help contextualize this paper's contributions.", "weaknesses": "A weakness, which the authors do acknowledge in their work, that prevented me from giving a higher score is that there is no clear connection between the structure of the weights and generalization, as there exists in \"standard\"/one-hot classification. As a result, it is unclear how much insight can be derived from the current result. I would appreciate the authors' thoughts on this.\n\nMinor: 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your encouraging feedback and for endorsing our paper. \n\n**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.**\n\nIndeed, 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.\n\nThe 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?\nBeyond 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].\n\nThe 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.\n\nFinally, 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].\n\n[SK22] Stability vs Implicit Bias of Gradient Methods on Separable Data and Beyond\n\n[Mur+23] Grokking of Hierarchical Structure in Vanilla Transformers\n\n[Moh+23] Grokking modular arithmetic can be explained by margin maximization\n\n[Mor+24] Feature emergence via margin maximization: case studies in algebraic tasks\n\n[Mik+13] Efficient Estimation of Word Representations in Vector Space\n\n**Q: text too dense.**\n\nWell received. It makes sense to slightly shorten Sections 6 and 7 by moving some parts to the appendix. \n\n**Q: lines 32-34.**\n\nWe appreciate the careful read. We will correct the typo by replacing \"when\" with \"of\" in Line 32.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4g0jhptnRd", "reviewer": "Reviewer_k4Ko", "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.\n\nIn 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.", "questions": "See above.", "limitations": "See above.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. They investigate an interesting and important problem: the optimization path and the implicit bias of NTP.\n\n2. Their formulation of NTP into the CE minimization over distinct contexts is novel.\n\n3. They provide rigorous theoretical results and the proofs are solid, to my knowledge.", "weaknesses": "1. The main issue of this paper is that, for the NTP-compatible and separable conditions to hold, one needs d > m. Does this overparametrization condition usually hold in practice or not? To my knowledge, in practice, the embedding dimension d is much smaller than the number of training data. Since m is not the number of training data and can be much smaller than that, it is not clear to me whether this assumption is possible in practice.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and the constructive questions raised.  We appreciate the careful read.\nWe hope that our responses below answer your questions.\n\n**Q1.  for the NTP-compatible and separable conditions to hold, one needs d > m.**\n\nHere are the key points to consider about the condition $d>m$ of Lemma 1:\n\n1. 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. \n\n2. 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.\n\nSpecifically, 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. \n\nWe 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.\n\n**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?**\n\nAssuming Eq. (4) has a solution, say $W_1$, every other solution takes the form\n$W^p = W_1 + W_{\\text{null}}$, where\n$W_{\\text{null}}$ (the null-space component) is orthogonal to\n$(e_z - e_{z'})h_j^T : z \\neq z' \\in S_j, j \\in [m].$\nThus,\n$W_{\\text{null}} \\in F^\\perp$. We will reword Line 155 according to this.\n\nAssuming it exists, the solution to Eq. (4) is unique *when constrained to the data subspace $F$*. To see this suppose $W_1$ and\n$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$.\nTherefore, $W_1 = W_2$, which is a contradiction. We will include this in the appendix for completeness.\n\n\n**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.**\n\nThank 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.\n\nFor (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:\n\nAssuming Equations (4), (6a), and (6b) hold for some  $W^p$ and $W^d$, we define $W_\\gamma=W^p+\\gamma W^d$. \nDecompose the inverse of $S_z(W_\\gamma h_j)$ as:\n$$\n1/S_z(W_\\gamma h_j) = 1+ \\sum_{z'\\neq z\\in S_j }\\exp\\left( (e_z'-e_z)^TW_\\gamma h_j \\right) + \\sum_{v\\not\\in S_j }\\exp\\left( (e_v-e_z)^TW_\\gamma h_j \\right)\n$$\n\nUsing the properties from (4) and (6a) for in-support pairs:\n$$\n(e_z'-e_z)^TW_\\gamma h_j = \\log\\left(p_{j,z'}/p_{j,z}\\right)\n$$\nleading to: \n$$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}.$$\n\nFor out-of-support pairs, applying (6b) and considering the limit as $\\gamma\\rightarrow\\infty$, we have:\n$$\n\\exp\\left( (e_v-e_z)^TW_\\gamma h_j \\right) = \\exp\\left( \\gamma (e_v-e_z)^TW^d h_j \\right)\\cdot \\exp\\left( (e_v-e_z)^TW^p h_j \\right) \\leq \\exp\\left(-\\gamma\\right)\\cdot \\exp\\left( (e_v-e_z)^TW^p h_j \\right)  \\stackrel{\\gamma\\rightarrow \\infty}{\\longrightarrow} 0\n$$\n\nPutting these together gives $1/S_z(W_\\gamma h_j) \\stackrel{\\gamma\\rightarrow \\infty}{\\longrightarrow} 1/p_{j,z}$ as desired. \n\nWe 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!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)