Implicit Geometry of Next-token Prediction: From Language Sparsity Patterns to Model Representations
Authors: Yize Zhao
Tina Behnia
Vala Vakilian
Christos Thrampoulidis
Venue: COLM (2024)
Keywords: language models, neural embeddings, optimization, implicit regularization, low-rank matrix factorization, support-vector machines
TL;DR: We reframe next-token-prediction as soft-label classification problem and provide a framework to characterize word and context embeddings by making connections to a nuclear-norm minimization problem.

Abstract: Next-token prediction (NTP) over large text corpora has become the go-to paradigm to train large language models. Yet, it remains unclear how NTP influences the mapping of linguistic patterns to geometric properties of the resulting model representations. We frame training of large language models as soft-label classification over sparse probabilistic label vectors, coupled with an analytical approximation that allows unrestricted generation of context embeddings. This approach links NTP training to rank-constrained, nuclear-norm regularized optimization in the logit domain, offering a framework for analyzing the geometry of word and context embeddings. In large embedding spaces, we find that NTP implicitly favors learning logits with a sparse plus low-rank structure.  While the sparse component captures the co-occurrence frequency of context-word pairs, the orthogonal low-rank component, which becomes dominant as training progresses, depends solely on the sparsity pattern of the co-occurrence matrix. Consequently, when projected onto an appropriate subspace, representations of contexts that are followed by the same set of next-tokens collapse—a phenomenon we term subspace-collapse. We validate our theory on synthetic and small-scale real language datasets. Finally, we outline potential research directions aimed at deepening the understanding of NTP's influence on the learning of linguistic patterns and regularities.

Published as a conference paper at COLM 2024
Implicit Geometry of Next-token Prediction: From Language Sparsity Patterns to Model Representations
Yize Zhao*, Tina Behnia', Vala Vakilian & Christos Thrampoulidis Department of Electrical & Computer Engineering
University of British Columbia (zhaoyize, tina. behnia, vaalaa, cthrampo jêece hww.:/eeeeeebb..o..
Abstract
Next-token prediction (NTP) over large text corpora has become the go- paradigm train large language models. Yet, remains unclear how NTP influences the mapping of linguistic patterns to geometric proper- ties of the resulting model representations. We frame training of large language models as soft-label classification over sparse probabilistic label vectors, coupled with an analytical approximation that allows unrestricted generation of context embeddings. This approach links NTP training to cank-constrained, nuclear-norm regularized optimization in the logit do- main, offering a framework for analyzing the geometry of word and context embeddings. In large embedding spaces, we find that NTP implicitly fa- vors learning logits with sparse plus low-rank structure. While the sparse component captures the co-occurrence frequency of context-word pairs, the orthogonal low-rank component, which becomes dominant as training progresses, depends solely on the sparsity pattern of the co-occurrence matrix. Consequently, when projected onto an appropriate subspace, rep- resentations of contexts that are followed by the same set of next-tokens collapse-a phenomenon we term subspace-collapse. We validate our the- ory on synthetic and small-scale real language datasets. Finally, we outline potential research directions aimed at deepening the understanding of NTP's influence the learning of linguistic patterns and regularities.
1 Introduction
Next-token prediction (NTP) is the preferred training paradigm for state-of-the-art language models. The process, elegantly simple, uses a large training corpus to minimize, for each context Zet E yl-1 of preceding tokens, the cross-entropy (CE) loss between the model's predicted conditional probability distribution over potential next tokens from vocabulary and the one-hot encoded actual next token 2; E V. The model's conditional distribution is defined through softmax map applied to logits (Wr 0) Whe(zer), which are generated by mapping context embeddings hę(zet) E Ra_the neural network's d- dimensional representations of contexts z<r using matrix W E E [ of word embeddings.
Rooted in the foundational works of Shannon (1948) and inspired by the "co-occurrence statistics" and "distributional hypothesis" (Harris, 1954), the NTP paradigm suggests that a word's meaning is defined by its context. This principle underlies both classical vector-space models (Schütze et al., 2008) and neural language models (Turian et al., 2010; Bengio et al., 2000; Baroni et al., 2014; Bengio & Bengio, 2000; Mikolov al., 2013b), propelling the development of today's sophisticated large language models (Radford et al., 2018; 2019; Brown et al., 2020). However, a fundamental question remains:
How does the NTP learning objective shape the relationship between language statistics and the geometry of model representations?
This inquiry, which we term the implicit geometry NTP-so named because the NTP objective does not explicitly impose any such certoonssip-xxpllorre how distances and angles within the neural network's d-dimensional representational space correlate with linguistic patterns at the end of training. We postulate that understanding this implicit
1

---

Published as a conference paper at COLM 2024
(a) Deep Network CORR(H)
NTP-UFM CORR(H) H)
(c) Theory (
Proxy CORR(S)
200
200
0.0
300
300
300
300
-0.5
400
400
-1.0
200
400
250
250
400
200
4:00
CORR [W]]
CORR W
CORR SORRWWIm WWmm
CORR(S
1.0
20
20-
20-
20
0.5
40
0.0
60
60-
00
80-
80
00
100
1000
-1.0
50
100
100
50
100
50
100
1.00
L
L
mm
S
0.75
0.50
0.25
100 200 400 200 400 200 400
100
200
400
0.00
-025
-050
Figure 1: A 4-layer transformer (TF) trained on the Simplified TinyStories dataset. (a) Cosine similarity of TF's context and word embeddings CORR(H) and CORR(W at the end of training (when NTP loss converges to its empirical entropy lower-bound). (b) CORR(H) and CORR W found by training the log-bilinear model (NTP-UFM) of Eq. (1) on the soft-labels P and support sets S of the original training set. (c) Geometry of context and word embeddings specified by our analysis. H and Wimmt are determined by the right/ left singular factors (Claim (C2)) of the low-rank component I. mm of logits (Claim (C1)). (d) An easy to compute heuristic proxy for the embeddings' geometry based on the training support set S (Proxy (P)). Details in Secs. 1.2 and 5.
geometry is key to understanding functional principles of large language models, since NTP used for training across diverse architectures, from LSTMs transformers and state-space models. Specifically, exploring how optimization under NTP shapes representations of words and contexts, which empirically mirror complex human-like patterns, not only fosters scientific interest but could also enhance model interpretability and explainability. Further, revealing how implicit geometry correlates with language statistics could lead to refined training and inference methods, addressing challenges like statistical imbalances in language data. Conversely, understanding how caate-of-the-as models internalize language to form representations might also enhance our grasp of language itself.
This paper develops an analytical framework to characterize the implicit geometry induced by NTP training on language datasets. Drawing inspiration from seminal studies the geometry of deep model representations in image recognition (Papyan et al., 2020), our framework distinguishes itself by not concentrating on specific architectures such as transformers, which have been the focus of previous studies on language representations. Instead, we assume that the model has adequate representation capacity and undergoes effective optimization, making it possible to minimize the NIP loss to its entropy lower-bound. This approach isolates the influence of NTP itself-rather than architectural nuances-in in shaping the implicit geometry of the language model. Our analysis also identifies the key role of the sparsity pattern in language statistics (Thrampoulidis, 2024). Concretely, we show that the recurrence of only a few tokens next-tokens in particular contexts leads to an implicit bias towards a matrix of logits that develops a sparse plus low-rank structure during training. The sparse component of this matrix captures the probabilities of co-occurring words and contexts, while the dominant low-rank component encodes the sparsity pattern of the co-occurrence matrix. Overall, this framework introduces a novel perspective, markedly distinct from traditional analyses of word representations like those in the Word2Vec model (Mikolov et al., 2013b; Levy & Goldberg, 2014).
1.1 Methodology
Our methodology integrates three foundational modeling concepts as follows:
2

---

Published as a conference paper at COLM 2024
Firstly, we frame NTP soft-label classification with CE loss applied to sparse probabilistic label vectors (Thrampoulidiss, 2024). This isolates m distinct contexts, which could be repeated multiple times throughout the training corpus, and assigns to each a sparse == dimensnnill conditional-probablii label vector p,, reflecting the frequency of each token following context E [mm. The sparsity of p; indicates that certain tokens, which we refer to as off-support tokens for the specific j-th context, never follow this context.
The second concept facilitates tractable analysis of context embeddings by assuming expressive (enough) neural networks can produce unconstrained embeddings h RR, independent of the architecture's specific complexities (Yang al., 2017; Mixon et al., 2020). This redefines NTP as a minimization, over word and context embedding matrices W and H, of the NTP loss across a training set T determined by the matrix of conditional probabilities P P(T) = [p , E [0, 1 leading to the following log-bilinear model:
weR HER NTT (W;
(1)
This way, our goal to study word-word, context-context, and word-context geometric rela- tionships becomes that of characterizing the Gram matrices Gw WW and Gμ = H H, as well as the logit matrix L = WH at the minimizers of Eq. (1). This task is complicated by the non-convex nature of the minimization and the sparsity of the probabilistic label vectors. Specifically, as we show, the sparsity of P may lead to multiple minimizers, potentially making geometric characterization ambiguous.
To address this, we leverage a third concept: focusing on specific minimizers identified through the regularization path (Rosset et al., 2003). This entails following the solution trajectory of the empirical risk minimization in Eq. (1) as an additive ridge regularization for W and H diminishes zero. For the purpose of comparing the analysis outcomes our numerical evaluations, we interpret the prediction obtained from the regularization path analysis as proxy for the solution found by gradient-basee optimization when the NTP loss approaches its empirical entropy lower bound.
1.2 Summary of findings
Building on the above methodology, our analytical framework leads to the following results:
Formulation in logit space. Sec. 3 presents equivalent formulation of the NTP objective CNTP (WH; + A w + ||| into the logit space, given in terms of a rank-constrained and nuclear-norm regularized minimization. From this, word and context matrices can be obtained through matrix factorization of logits L.
Logit Convergence. Focusing on d > Vl, in Sec. 4, we demonstrate that the logit matrix LA for regularization A + 0 (which is proxy for training iteration k + oo), behaves for some R(A) + co as shown by the following claim:
(C1) Logits' sparse plus low-rank decomposition: As A + O,
LA
in
+ R(A)
mm
(2)
sparse companent: sets in-support token probabiliries
low-rank component:
separates from off support tokens
where the orthogonal components Ln and L mm have distinct operational roles.
Li encodes information regarding frequencies of in-support tokens, inheriting the sparsity of the data matrix P. In contrast, mm depends only on the sparsity pattern S E {0, 1 of P. Thus, its value is guided by both the observed 'company' of each context (in-support tokens), but not their frequency, and, by the 'company' it lacks (off-support tokens). The role of this second component measured with respect to the low-rank promoting nuclear norm. This decomposition in (2) also shows the logits become unbounded in norm when projected on the subspace of mm
maximize the logit margin between in- and off-support tokens,
This does not constrain d relative to the number of distinct contexts m, which could be significantly larger than the vocabulary size V. Also, despite large embedding dimensions, our results still reveal how the geometry of learned word embeddings encodes fine-grained similarities between words as materialized in the patterns of language data.
3

---

Published as a conference paper at COLM 2024
(a) Training checkpoints
(b) Proxy
Epoch
Epoch 100
Epoch 600
$
Epoch 0
Epoch 100
Epoch 600
CORR(S)
1OS Figure 2: Similar to Fig. 1, this time on a 12-layer TF trained on subset of 100 stories from the
CORR (HL)
TinyStories dataset. Here, computing the theoretical prediction Hmm is computationally expensive. Thus, we compare the embeddings geometry with the Proxy (P). Details in Sec. 5.
Context and word embeddings' geometry. Similar to logits, word/ context embeddings W and H also grow in norm as À 0. This occurs way that simultaneously guarantees the resulting logit matrix abides by Claim (C1), which leads to the following additional claims:
(C2) Directional convergence: Word and context embedding matrices converge direction- ally to matrices Wmmm := uz1/2R, Hmm := R 1//2VT, where R is a rotation matrix and uEv is the singular value decomposition (SVD) of the low-rank max-margin component m in Eq. (2). Concretely, letting A A/|| for any matrix A:
Gw uEu T and Gи VEV
(3)
(C3) Subspace collapse:
Since L mm depends only on the sparsity patterns of next-token distributions (not their frequencies), consequence of (C2) is property termed subspace collapse: embeddings H,, hy of contexts j * [m] that are followed by the same set next-tokens (although their frequencies may differ), converge to the same limiting direction, i.e., for all j. E [m]
support( Pj) = support(fy) = cos (hij kty) h1 h by 1) - 1. (4)
(C4) Soft-label interpolation:
When projected on the subspace of in-support tokens, the logits WH interpolate the corresponding soft-labels ensuring that the NTP loss reaches the entropy lower bound. Concretely, for all) [m],
(w= (0=) logg (Piz///1))2) Vz,z's e support(p;).
Since the subspace of in-support tokens orthogonal to the subspace mmm the subspace collapse (C3) does not prevent logits w h and w. hp of in-support tokens z E support (;) = support( to interpolate (in the sense of Eq. (5)) potentially different conditional probabilities (soft-labels) ÚJ,2 and pr. respectively.
Experiments. Sec. 5 presents experiments on controlled settings that validate our findings. An example is shown in Fig. 1: The experiment involves training a 4-layer transformer (TF) on a training set extracted and curated from the TinyStories dataset (Eldan Li, 2023), with the following characteristics: a vocabulary size V = 104 and total number of contexts n x 3050, out of which m 400 are distinct. We choose TF embedding dimension d = 128 and train until the empirical entropy lower bound is reached within an order of 10-4 The leftmost Panel (a) depicts the geometry of the context embeddings H (Top), word embeddings W (Middle), and logits L (Bottom) learned by the TF at the end of training. Panel (b) compares these with the same quantities learned by the log-bilinear model in Eq. (1) evaluated on the same training dataset. The apparent resemblance of the patterns validates that the proposed log-bilinear analysis model is good proxy for the
4

---

Published as a conference paper at COLM 2024
TF model. Panel (c) compares the TF output to our analytical predictions: Comparing TF's context/w word embeddings and logits to Hmm /Wmm and L mi confirms Claims (C2) and (C1). The bright yellow regions in the Gram matrices of embeddings help visualize the subspace collapse Claim (C3). See Sec. 5 for additional details and verification of Claim (C4).
Finding the analytical prediction (C2), illustrated in Panel (c), requires solving semi- program for L and computing its SVD. In large scales, this can be computationally prohibitive, motivating the following heuristic proxy for the embedding geometry.
(P) Empirical proxy for directional convergence: Let ;
== (Ty - 4 Bl 1 be the (column-wise) centered support-set matrix. The directional component of word and context embedding matrices can be well-approximated by as follows:
W Wmm T
SS
and
H mm
S S.
(6)
Panel (d) of Fig. 1 suggests proxy (PP) as a good enough approximator of the embeddings structure that simply only depends on the sparsity pattern S of the next-tokens' conditional probability matrix.
A second example supporting proxy (P) is shown in Fig. 2: Here, we train a larger TF on subset of the TinyStories, where computing Lmm and its SVD is expensive. Instead, we observe that the structure of the support set closely captures the structure of the learned embeddings: Along the diagonal blocks, where contexts with the exact same support sets are situated, the context embeddings align closely (subspace collapse in Claim (C3)). Conversely, when support sets have zero intersection (dark entries in Panel (b)), context embeddings ex- hibit low correlation. This finding is consistent with classical intuitions of the " "distributional assumption" of of NTP to promote low-rank logits subject to margin conditions for in/ mfoff-support tokens.
and is formalized by our framework via the implicit bias
2 Related work
We pinpoint three main related research areas; see App- A for more in-depth discussion.
First, our research conceptually mirrors the seminal work by Levy & Goldberg (2014) who framed the Skip-Gram with Negative Sampling (SGNS) training objective of Word2Vec (Mikolov et al., 2013b;a) as weighted matrix factorization. Specifically for large d, they demonstrated that SGNS implicitly factorizes the pointwise mutual information (PMi) matrix. Their analysis relies on the fact that Word2Vec architecture already is log-biliniear model, while the log-bilinear model in (1) is only used by us an analytical proxy for more complex architectures. Also, different to them, we focus on NTP, which employs softmax instead of sigmoids as in SGNS. (However, our analysis applies also to SGNS; see App. A.4). More importantly, we contribute a fresh perspective on this line of inquiry by: (i) confronting the sparsity of probabilistic labels head-on, identifying that leads to diverging embeddings, scenario where setting the loss gradient to zero, in Levy & Goldberg (2014), is infeasible; (ii) examining embeddings through the lens of the regularization path-a surrogate for gradient descent optimization- which unveils that embeddings emerge from the factorization of in + RL mm, where in captures frequencies akin the PMI, while the directional component L mm becoming dominant as weights diverge, reflects the explicit sparsity patterns of the context-word co-occurrences. We envision this fresh perspective could similarly motivate further research leveraging the geometric insights of embeddings uncover linguistic phenomena such as the linear relationships underlying word analogies (Mikolov et al., 2013a; Pennington et al., 2014).
Second, our exploration of NTP's regularization path inspired by the implicit bins/regularization research on the preferred solutions of optimizers like GD in overpa- rameterized systems. Ji & Telgarsky (2020) showed that GD's trajectory in linear one-hot encoding models aligns with the regularization path, providing a lens for examining GD dynamics. Thrampoulidis (2024) recently extended this analysis to NTP, framing it as sparse soft-label classification. By lifting their assumption of fixed context embeddings, we delve into the more complex non-convex domain, establishing explicit connections between the implicit geometry and linguistic patterns, which we also verify empirically.
5

---

Published as a conference paper at COLM 2024
Finally, our research intersects with the study of neural-collapse (NC), which delves into the geometry of last-layer features and weights in deep networks trained in the interpolating regime (Papyan et al., 2020), utilizing the unconstrained festures model (UFM) for analysis (Mixon et al., 2020; Fang et al., 2021; Zhu et al., 2021). Our work is particularly aligned with Thrampoulidis et al. (2022): We also examine the UFM's regularization path but within the NTP framework, extending their one-hot classification findings, which can be seen as special cases of ours. Additionally, our findings have some parallels with the multilabel NC geometry explored in Li et al. (2023); Fisher et al. (2024). However, our setting is more general raising stringent assumptions on the label distribution. For detailed comparison, see App. A.
3 Formulation
Notations. Throughout, lowercase and uppercase bold letters (e.g., and A) represent vectors and matrices, respectively. We use a and A to denote vectors/ pectors/matricess matrices normalized by their Euclidean norm. We denote A the (i, entry of matrix A and a its j- th column. We let (A) and N (A) its range-space (aka column space) and null-space, respectively. (.-)) and ||-|| denote Euclidean inner product and norm, respectively, We use ||+ I to denote the nuclear-norm (i.e. sum of singular values). BV represents the identity matrix of size V and Ly, the all ones vector of size V x 1 (subscripts are removed when clear from context). V denotes the V-dimensional unit simplex and S(.) R + V-1 the softmax map: 5(a)
with (a) exp(e, a)/ ///E[[]] exp(e_a), where eu is the v-th standard basis vector in IRY. We also denote ē the j-th standard basis vector in R All logarithms are natural logarithms (base e).
3.1 NTP objective as soft-label classification
We let V = [V] := {I,..., } represent a finite vocabulary of tokens (we use the terms 'word' and 'token' interchangeably). We denote by Z
..i,...,2,)) sequence of t tokens Z₁ E V and focus, for simplicity, on prediction of the last T-th token z == ZT given context == 21:T-1- For this, we assume access to a training set consisting of π sequences Tn = (1),)))))())) such that x; E X == VT and z; V. This is used to train model
for X' - V. (x) Whg(x) parameterized by 0' (W, where W RVxd is decoding matrix and 0 parameterizes a map ho X - Rd from contexts to d-dimensional embeddings. We impose no restrictions on the specific form of the embedding map, which may, for example, be produced by an MLP, an LSTM, or TF. We refer to row Wp, V E of W as word embedding of token V and h as context embedding of context x. The model is found by minimizing the empirical CE loss CF((') 1 Lisnn] log (S,, (fo (x))).
Following Thrampoulidis (2024), we reframe the NTP training objective classification over m S distinct contexts, each associated with a sparse probabilistic label vector P) V Concretely, we denote 8,,..., the m S distinct contexts among the (large number of) total contexts Tn- Also, we let A. = π Ei(p] Il [x: = xj] denote the empirical probability of distinct context x.. Accordingly, let Pj E A V-1 denote the probability vector of conditional next-token distribution, i.e., for all z E V: Pj= == Eij 1[2; = 2), j [mr] In words, At Pj.z is the number of occurrences of token as follow-up to context xy. Define the support sets of these probability vectors S := f- V P1, > 0} and let S == |Sy].
It also convenient define the probability matrix p [p,.... Pm]] and its corresponding support matrix S E (0, 1) Vxm such that $iz. = 1, iff z E S,. In typical scenarios, S < V; thus, P and S are sparse matrices. For given context index j E [m], we say word z is in-support token if S(z, 1 (eqv. Z E S;); otherwise, we say 2 is off-support token.
With the above notation, the training loss becomes (Thrampoulidis, 20244)
CE(0') Eje(m At Piz log (Shhha(())
(7)
This is lower bounded by the empirical T-gram entropy (referred to hereafter entropy) of the data (Shannon, 1948), i.e., for all 0': CE(0') > H == Ejc(m) A; Exev P1z log (f)z)-

---

Published as a conference paper at COLM 2024
3.2 Unconstrained features model for NTP training
To gain insights into the geometry of solutions to CE minimization in Eq. (7), we assume sufficient model expressivity, allowing optimize embeddings freely, instead of abiding by their architecture-specffic parameterization. We formalize this concept below.
Definition 1 (NTP-UFM), The unconstrained features model (UFM) for NTP training over training set Tm (Āg, Pj) jc(ml refers to the following log- oo-billinear optimization problem:
A 111WI] + 1
minw,1 CE(WH) + +
(NTP-UFM)
where minimization is over word and context embedding matrices W E Rvdd and = Nmm] E Rdxm and CE(WH) == -jj(mm] &g LeEs, Р1, log (S2
The key difference of NTP-UFM compared to Eq. (7) is that embeddings hg)(8) are now optimized unconstrainedly through free variables 11. Furthermore, we have introduced ridge regularization. Although our focus is on understanding the geometry arising from minimizing the unregularized CE objective in Eq. (7), we utilize ridge regularization as proxy to examine the behavior of gradient descent.? 2 Specifically, regularization solution of NTP-UFM as proxy for the parameters (W,, H₂)) learned in large iteration k through GD the unregularized objective (7).
use the vanishing-
3.3 Reformulation in terms of logits
We now introduce reformulation of NTP-UFM utilizing the logit matrix WH. This is grounded on the well-known fact concerning the nuclear norm of matrix (Srebro et al., 2004; Fazel, 2002):||L]] min
[|]l] + ||||
Lemma 1. Denote logit matrix L = [lt,...,ee lm] E Vxm and let LA be a minimizer of
Lje(m] tt Eres, Piz log (S.(e)) + AlI
min L:rank(L)sd
(8)
with SVD L, LEV where U E R RRV
EEV E R V E R and r rank (La) S d. Then:
(i) The optimal cost of NTP-UFM is the same as that of Eg. (8).
(ii) (Wa, H₂)) minimizes NTP-UFM if and only if there exists minimizer LA of Eq. (8) such that:
W LA, W.W, =Euu and H, VEVT
The constraint rank (L) S Eq. (8) ensures the logit matrix can be factorized as WH with inner factor dimension being d. Thus, in general, Eq. (8) is non-convex.
4 Analysis of unconstrained features model for NTP training
This section analyzes the regularization path of Eq. (8) and, subsequently, NTP-UFM. Throughout, we assume d 2 V, in which case the constraint rank (L) < d becomes redun- dant. All proofs are deferred to the appendix and numerical evaluations to App. D.1.
4.1 NTP-SVM logits
The following nuclear norm-minimization problem plays key role in our results. Definition 2. Define the NTP-SVM logit matrix mmm as solution to the following optimization:
L mm E arg min ||
(NTP-SVM.)
subj. to L[z, f] L[z', = 0, [m), # z' ES. L L - L[ 2 1, Vj [m], E S,, 4 S,.
(9)
(10)
Thhe equivalence between GD and regularization paths has been rigorously established in certain convex settings (li et al., 2020), but it may not hold generally. Our setting is non-convex, use this equivalence heuristic proxy; future work discussion in Sec.
7

---

Published as a conference paper at COLM 2024
The optimization in NTP-SVM, is semidefinite program (SDP) and always feasible, as shown by the straightforward feasibility of the centered support matrix Š (Ml b1y Its solution m depends solely on the support matrix S of the data. Interestingly, we prove in App. C that $ solves NTP-SVM, under symmetry assumptions on S. More generally, our empirical results show that 5 closely approximates the true solution L supporting the use of Proxy (P); see App. C.3.
Finally, when S = vj E [m],, then Eq. (9) vanishes and Eq. (10) lower bounds the margin between the correct token and the rest. This constraint appears in the classical hard-margin SVM, which explains our naming. Instead, the NTP setting results in additional equality constraints and minimizes nuclear-norm rather than Euclidean norm.
4.2 Regularization-path analysis of logits
We now characterize the solutions of (8) when regularization vanishes and d 2 V.
Theorem 1. Suppose d 2 V. In the limit of vanishing regularization, the solution L, of optimization (8) diverges in norm and converges in direction to Lmm, Formally, lim
|| L +oa, and, there
exists minimizer mm of NTP-SVM, such that lim >0
L
m.mm
0.
Thus, with vanishing regularization, NTP-SVM, captures the structure of the logits. By Lemma 1, it also captures the structure of W and H. We explore this further in Sec. 4.3.
Now, recall that NTP-SVM, only depends on the support matrix of the training data, raising the questions: How do the probabilities Piz influence the learned logit matrix and what implications this has for the loss? The theorem below complements Thm. and addresses these questions. To state the result, define the following matrix subspace, as the span of rank-one matrices each being zero except entries (z, and (z',j), which equal 1 and -1 respectively for all context indices j E [m] and in-support word indices z, E Sj:
F F= span (f (ez ez) =
Z' S E [m] C R Vxm
Note for any L, the projection (L( onto F is sparse with support matching S. Theorem 2. Under the setting of Theorem 1, the loss approaches its lower bound lim CE(L,)
H. Additionally, lim, -0 P L where i I is the unique L E I that guarantees differences of logits equal log-odds, iee, L[z,]] [[=', = log fizz/ //zz)))))))) Vz f E S₁, E [m]]
The loss reaches the lower-bound H. because on the subspace T, L, satisfies the log-odds linear equations of the theorem. This gives rise to Claim (C4): the model outputs the correct probabilities for in-support tokens. There is no contradiction of this claim to Thm. because mm belongs to the orthogonal complement F1 (which follows from Eq, (9); see App. B.3.1).
Combining the two theorems, LA converges to in in F and diverges in F1, where converges directionally to L mm, This justifies Claim (C1).
4.3 Geometry of embeddings
In view of Lemma can obtain the embeddings by factorizing the logit matrix. Since logits diverge, the same is true for the embeddings, which also converge directionally.
Corollary 1. Denote (W,, H2) any minimizer of NTP-UFM. Consider the SVD mm = U.V Then, using notation to denote normalized quantities, in the limit of vanishing regularization
W.W. -EuT =: G: H HA + VΣV G: m and W,H, UΣVT
This corollary supports Claim (C2). Another direct consequence of it is Claim (C3): embeddings of contexts whose support sets are identical asymptotically collapse to the same embedding as À - 0. We formalize this in Proposition 2 in App. C.1.
8

---

Published as a conference paper at COLM 2024
5 Experiments
We empirically validate our analysis on text data. We start with two small-scale synthetic datasets to examine Claims (CI)-(C4). We then experiment with a larger-scale dataset, where SVD calculation are computationaly expensive, we examine the Proxy (P).
Datasets. For a complete description of the datasets, see App. D.2.
Synthetic. We manually create simple (context, next-token) pairs, with context length of size T - 11- = 5. The dataset is constructed so that all support-sets satisfy Sj = 3.
Simplified TinyStories. For more realistic but still controlled setup, we curate the dataset from the TinyStories corpus: We derive contexts &; and support sets Sj by choosing the most frequent word-level contexts with length
= 5.
TinyStories. For more standard data, we use 100 stories sampled from TinyStories. Unlike the other two datasets, we do not sample over the contexts and support sets.
Models. We train decoder-only TF with layers and d 128 for Synthetic and Simplifion fied TinyStories, and 12 layers and d = 256 for TinyStories. For training details, see App. D.2.
We verify two claims made in Sec. 4: 1) We assess whether numerically optimizing both the TF and NTP-UFM in each setup yields the same embedding geometries, 2) We evaluate the global solution of NTP-UFM, as specified by Thm. 1 and Cor. 1, and check whether it respects the same implicit geometry as the TF and the NTP-UFM. For the first one, we train NTP-UFM directly on the label distribution of each dataset; this is computationally feasible only for the first two datasets. For the second claim, we find the global solution L mum of NTP-SVM, on each dataset using CVXPY (Diamond & Boyd, 2016) and Hmm, Wmm using Cor. 1. If not computationally possible, we only use the Proxy (P) as reference.
Metrics. We define the following metrics for verifying our theoretical results:
Visualization of embeddings' geometry. To visualize the geometry of the word and context em- beddings, we plot the heatmap of the normalized Gram matrix CORR(W (W and CORR (H), where for matrix X, CORR(X) is matrix whose (i, entry is the cosine similarity between the respective columns of X, i.e., |CORR (X) i, = x x xI |1y))) To check Proxy
(P), we compare the context and word embeddings with COR and (
Quantifying geometric similarity. To measure the structural similarity between W, H and ST, 5, we use SIM(X Xx Y) = (oxy + (Txxy e), as the correlation matrix between two matrices. 3 Here, yxy is the covariance, and Ø'x and Qy are the standard deviations of X and Y, respectively, and e is small constant for stable division. A value of SIM (X, Y) =1 indicates perfect structural correlation. To simplify the presentation, and with some abuse of notation, we denote the similarity metric for context and word embeddings as follows,
SIM SIM(H, 5) == SIM (IM(COR(HH)
SIM(WT, ,5)) == SIM(CORR(W) "), CORR(S )).
Recovery of soft-labels. To verify Thm. 2, we measure the distance ||P3 - [in| of Lin from the projection of Lt onto the subspace F.
Results. In Sec. 1.1, we discussed Fig. 1, where we compared the logits, word and context em- beddings trained on Simplifies TinyStories by (a) TF, (b) NTP-UFM, and those predicted by (c) Claims (C3), and (d) Proxy (P). Similar observations for the TinyStoria ies and Synthetic datasets are shown in Figs. 2 and 10. For visualization details, see App. D.2.
Next, we discuss Fig. that tracks our metrics during training for the Synthetic and Simplified TinyStories datasets. Fig. 3-(a) confirms that the loss is close to the empirical entropy H. Fig. 3-(b) confirms the norm growth of the parameters. Note the slower growth of the context embeddings H in the TF model, which we suspect is due to the layer normalization at the final layer. We also confirm Claim (C4), the recovery of the soft labels, in Fig. 3-(c). At last, Fig. 3-(d) and (e) track the correlation of the learned geometries with the proxy (P) and with the theoretical predictions (Hmm) Wmm). ). Note that all similarity
"hiis is the Structural Similarity Index Measure (SSIM) used to measure the similarity between two images (Wang et al., 2004).

---

Published as a conference paper at COLM 2024
(a) Loss convergence
(b) Norm growth
(c) Pa (L.)
Ppoch:k)
Ppoch())
(d) Correlation with proxies
(e) Correlation with thcory
(poch(2)
Spach(s)
Figure 3: Experiments on Synthetic and Simplified TinyStories datasets. (a): CE approaches H, (b): Parameters' norms grow during training. (c): L projection on the data subspace F converges to the sparse component L specified by the soft-labels fj. (d): (IM (), (solid) and SIM (W,,W... )") (dashed). At the final stage the learned embeddings exhibit high similarity with proxy (P). (e): Same as (d), this time comparing with theory, ie,, SIM H2,, Hmm) and SIM (W.T, Wmm J. Details in Sec. 5.
metrics increase towards 1 as training progresses. In Figs. 12 and 13 in the appendix, we display the evolution of the embeddings by capturing their heatmaps at various training checkpoints. Interestingly, the final patterns emerge early in the training process. Also, see App. D.2.3 for analogous plots for TinyStories.
Finally, in Figs. 14-17, we visualize the text values of the contexts and words along with their geometry heatmaps. We see that the embeddings' similarities are consistent with the linguistic properties of words/ words/contexts. contexts. Embeddings of contexts such as "boy named Timmy, Timmy" and "kid called Lilly. She" align closely as they are likely followed by similar words. Also, word embeddings tend to cluster by their grammatical categories, with verbs and adjectives often showing high similarity. For instance, the verbs go, eat, sing, wear, play, and the adjectives clever, fat, furry, lively, little have high similarity.
6 Concluding remarks: Limitations and future work
Looking forward, we see several promising avenues for expanding upon our framework and addressing its present limitations. First, we believe it is feasible to broaden the regularization-pa analysis to the non-convex realm of problem (8), potentially revealing rank-constrained variant of NTP-SVM, We hypothesize that the support matrices typical of natural language data might accommodate low-rank solutions that adhere to the con- straints specified in Eqns. (9)-(10). This would allow us to relaxing the d > V constraint in our analysis. See App. D.5 for more discussions. Second, while our initial focus was on establishing the framework, conducting preliminary analyses and proof-of-concept ex- periments, we encourage more extensive experimental exploration, especially with larger datasets and more complex model architectures with larger context window. Experiments in this paper focused on TF models due to their recent success in language. However, our theoretical results are independent of network design. Instead, they require the model to be wver-parameterized enough to reach the entropy lower bound and ensure expressivenes eness in the embedding space. For a brief discussion, see App. D.3. Third, there is compelling opportunity to delve into how the distributional characteristics of word/ context pairs in natural language, represented by the probability matrix P, influence key properties like low-rankness and symmetry in the SVD factors of L Fourth, although the framework encompasses full autoregressive training, our experimental study has focused on prediction of the last token prediction, with the exception of preliminary experiments in App. D.4. Deeper studies are required with sequences of variable length to empirically evaluate the brittleness of our model's assumption, which enumerates distinct contexts regardless of their length or constituent token composition.
Ultimately, our goal is to spur further exploration into the geometries of context and word embeddings. More broadly, we aim to foster interest in deeper investigations into how the inductive bias of NTP in influences the learning of linguistic patterns and regularities.
10

---

Published as a conference paper at COLM 2024
Acknowledgments
This work was funded by NSERC Discovery Grant No. 2021-03677, Alliance Grant ALLRP 581098-22 and an Alliance Mission Grant. TB also gratefully acknowledges support by UBC's Four Year Doctoral Fellowship and scholarship by the Advanced Machine Leaning Training Network at UBC.
References
Carl Allen and Timothy Hospedales. Analogies explained: Towards understanding word embeddings. In International Conference Machine Learning, pp. 223 233-331. PMLR, 2019.
Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. A latent variable model approach to pmi-based word embeddings. Transactions of the Association for Computational Linguistics, 4:385-399, 2016.
Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameter- ized nonlinear models. IEEE Transactions on Neural Networks and Learning Systems, 33(12):
7717-7727, 2021.
Marco Baroni, Georgiana Dinu, and Germán Kruszewski. Don't count, predict! a systematic comparison of context-counting context-predicting p semantic vectors. In Proceedings of the 52nd Annual Meeting of the Association for Computational Linguistics (Volune Long Papers), pp. 238-247, 2014.
Tina Behnia, Ganesh Ramachandra Kini, Vala Vakilian, and Christos Thrampoulidis. On the implicit geometry of cross-entrope parameterizations for label-imbalanced data. In International Conference on Artificial Intelligence and Statistics, pp. 10815-10838. PMLRR, 2023.
Samy Bengio and Yoshua Bengio. Taking on the curse dimensionality in joint distributions using neural networks. IEÉE Transactions on Neural Networks, 11(3):550-557, 2000.
Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. Advances in neural information processing systems, 13, 2000.
Daniel Bis, Maksim Podkorytov, and Xiuwen Liu. Too much in common: Shifting of embeddings in transformer language models and its implications. In Proceedings of the 2021 conference of the North American chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5117-5130, 2021.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Matias D Cattaneo, Jason M Klusowski, and Boris Shigida. On the implicit bias of adam. arXiv preprint arXiv:2309.0007 2023.
Steven Diamond and Stephen Boyd. Cvxpy: A python-embeddeed modell language for convex optimization. Journal of Machine Learning Research, 17(83):1-5, 2016.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? Xiv preprint xrXie:2305.07759 2023.
Kawin Ethayarajh. How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings. arXiv preprint arXiv:1909.00512, 2019.
Kawin Ethayarajh, David Duvenaud, and Graeme Hirst. Towards understanding linear word analogies. Xiv preprint ar
2018.
11

---

Published as a conference paper at COLM 2024
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118(43), 2021.
Maryam Fazel. Matrix rank minimizationn with applications. PhD thesis, PhD thesis, Stanford University, 2002.
Quinn Fisher, Haoming Meng, and Vardan Paapyan. Pushing boundaries: Mixup's ss influence on neural collapse. arXiv preprint arXio:2402.06177, 2024.
Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie-Yan Liu. Representation degeneration problem in training natural language generation models, 2019.
Nathan Godey, Éric de la Clergerie, and Benoît Sagot. Anisotropy is inherent to self-attention in transformers. Xrviv preprint arXiv:2401.12244 121433 2024.
Michael Grant and Stephen Boyd. CVX: Matlab software for disciplined convex program- ming, version 2.1. http://cvxr com/cvx, cvx, March 2014.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pp. 1832-1841. PMLR, 2018a.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient déscent on linear convolutional networks." Advances in Neural Information Processing Systems, 31:9461-9471, 2018b.
Zellig $ Harris. Distributional structure. Word, 10(2-3):146-16,, 1954.
Tatsunori B Hashimoto, David Alvarez-Melis, and Tommi S Jaakkola. Word embeddings as metric recovery in semantic spaces. Transactions of the Association for Computational Linguistics, 4:273-222- 286, 2016.
Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. Xrv preprint ar trXie:1803.07300, 2018.
Ziwei li and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176-1716,, 2020.
Ziwei Ji, Miroslav Dudfk, Robert E Schapire, and Matus Telgarsky. Gradient descent follows the regularization path for general losses. In Conference on Learning Theory, pp. 2109-2136. PMLR, 2020.
Daniel Jurafsky and James H. Martin. Speech and Language Processing. Prentice Hall, 3 edition, 2023. Draft.
Omer Levy and Yoav Goldberg. Neural word embedding as implicit matrix factorization. Advances in neural information processing systems, 27, 2014.
Bohan Li, Hao Zhou, Junxian He, Mingxuan Wang, Yiming Yang, and Lei Li. On the sentence embeddings from pre-trained language models. arv preprint arXiv:2011.05864, 2020.
Pengyu Li, Yutong Wang, Xiao Li, and Qing Qu. Neural collapse in multi-label learning with pick-all-label loss. arxiv preprint ar axXo::3300115003 2023.
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 Machine Learning, pp. 22188-22214. PMLR, 2023.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. ar Xrx preprint ar arXta:1301.3388, 2013a.
12

---

Published as a conference paper at COLM 2024
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. Advances in neural information processing systems, 26, 2013b.
Tomáš Mikolov, Wen-tau Yih, and Geoffrey Zweig. Linguistic regularities in continuous space word representations. In Proceedings of the 2013 conference of the north american chapter of the association for computational linguistics: Human language technologies, pp. 746-751, 2013c.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained
features. Xio preprint ar arXit:2011.11619, 2020.
Jiaqi Mu, Suma Bhat, and Pramod Viswanath. All-but-thetop:: Simple and effective post- processing for word representations. arXiv preprint arXio:1702.01417, 2017.
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 3420-3428. PMLR, 2019.
Vardan Papyan, XY Han, and David I. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):24652-244 24663, 2020.
Jeffrey Pennington, Richard Socher, and Christopher D Manning, Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014.
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-29230, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, al. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Benjamin Recht, Maryam Fazel, and Pablo A Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM review, 52(3))471-501, 2010.
Saharon Rosset, Ji Zhu, and Trevor Hastie. Margin maximizing loss functions. In NIPS, PP. 1237-1244, 2003.
Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. A mathematical exploration of why language models help solve downstream tasks. arviv preprint ar arXit:2010.0366 2020.
Hinrich Schütze, Christopher D Manning, and Prabhakar Raghavan. Introduction to informa- tion retrieval, volume 39. Cambridge University Press Cambridge, 2008
Claude Elwood Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3):379-42)) 1948.
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):2822-22) 2878, 2018
Nathan Srebro, Jason Rennie, and Tommi Jaakkola. Maximum-margin matrix factorization. Advances in neural information processing systems, 17, 2004.
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:31089-31101, 2022.
13

---

Published as a conference paper at COLM 2024
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Trans- formers as support vector machines, 2023a.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023b.
Christos Thrampoulidis. Implicit bias of next-token prediction. arXiv:2402.18551, 2024.
arXiv preprint
Christos Thrampoulidis, Ganesh R Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. arXiv preprint ar arXi::2208.05512, 2022.
William Timkey and Marten Van Schijndel. All bark and no bite: Rogue dimensions in trans- former language models obscure representational quality. arXiv preprint ArX:::2009. X4t:2209.04404, 2021.
Joseph Turian, Lev Ratinov, and Yoshua Bengio. Word representations: simple and general method for semi-supervised learning. In Proceedings of the 48th annual meeting of the association for computational linguistics, pp. 384-394, 2010.
Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference on Learning Theory, Pp. 4224-42555 PMLR, 2021.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assess- ment: from error visibility to structural similarity. IEEE transactions image processing, 13 (4):600-612, 2004.
Robert Wu and Vardan Papyan. Linguistic collapse: Neural collapse in (large) language models. Xrviv preprint 2rXio:2406.1766,2000 2024.
Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, and William W Cohen. Breaking the softmax bottleneck: A high-rank rnn language model. arXiv preprint arXio:1711.033: 103955,, 2017.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu A geometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34, 2021.
14

---

Contents
1 Introduction
1.1
Methodology
1.2
Summary of findings
2 Related work
Formulation
3.1
NTP objective as soft-label classification
3.2
Unconstrained features model for NTP training
3.3
Reformulation in terms of logits
4
Analysis of unconstrained features model for NTP training
4.1
NTP-SVM logits
4.2
Regularization-pate analysis of logits
4.3
Geometry of embeddings
Experiments
6
Concluding remarks:
Limitations and future work
10
A Related work: Detailed version
16
A.1
Word embeddings as matrix factorization
16
A.2
Overparameterization and Implicit bias / regularization of GD
17
A.3
Neural collapse geometry
18
A.4
Connections to Word2Vec
18
Proofs
20
B.1
Centering of word embeddings
20
B.2
Proof of Lemma
20
B.3
Proof of Theorems and 2
21
B.3.1
Notations recap
22
B.3.2
When does the NTP loss attain its entropy lower-bound.
22
B.3.3
Ball-constrained CE minimization and main result
24
B.3.4
The in-support NTP loss and its minimizer I
25
B.3.5
Auxiliary lemmas
26
B.3.6
Proof of Theorem
27
C On the solution of NTP-SVM.
29
C.1 Subspace-collapse.
29
C.2 Special case: support sets of equal sizes
................................
30
C33 A simple candidate solution
30
C.4 A proxy for NTP-SVM, solutions
31
C.5 Proofs
32
C.5.1 Proof of Proposition
32
C.5.2 Proof of Proposition 5
33
C.5.3 Proof of Propositions and
34
D Numerical experiments
36
D.1 Numerical simulation: NTP-UFM...
36
D.2 Additional details and results: Sec.
D.2.1 Datasets
36
D.2.2 Training details
38
D.2.3 Additional results
38
D.3 Other architectures
44
D.4 Auto-regressive training
44
D.5 Discussion on large vocabulary setting
44
15

---

Published as a conference paper at COLM 2024
A
Related work: Detailed version
This section expands upon the section on related work in the main body.
A.1 Word embeddings as matrix factorization
Neural probabilistic language models are designed to predict the probability distribution of a target word v E within given vocabulary V, based on a sequence of preceding or surrounding words of (variable or fixed) length known as the context z E V*. In these models, word and context representations (aka embeddings) are parameterized by their internal architecture (Bengio et al., 2000). The emergence of energy-based models like the Word2 family, particularly through the introduction of training methodologies such as the Skip-Gram with Negative Sampling (SGNS) objective by Mikolov et al. (2013b;a), marked a significant advancement in learning high quality embeddings; e-g. see Jurafsky & Martin (2023). SGNS aims to optimize the alignument between vectors of target words and the context words surrounding the target, while disassociating those of randomly drawn pairs. This is achieved by optimizing log-bilinear model that learns unconstrained embeddings for words in two roles: as targets and as contexts. While intuitive, the objective remained somewhat mysterious until the seminal analysis by Levy & Goldberg (2014) offered formal examination of embeddings generated by the SGNS objective, framing it as form of weighted matrix factorization. Specifically for embeddings of large dimensionality, Levy and Goldberg demonstrated that SGNS implicitly factorizes a shifted version of the Pointwise Mutual Information (PMI) matrix, denoted as PMI, which is measure of association
between words and contexts by having entries PMI j] = log Pr(zi)
Pi:= for target
Pr(z) log Pr(z]
word z and context word j. This result establishes a fundamental connection between count-based and predictive models (Pennington et al., 2014). However, this approach presumes non-zero occurrences of all context-word pairs in the training set, a condition not always met, leading to undefined factorization for pairs with zero co-occurrence (resulting in PMI 2 jl = - oo). To address the possibility of negative infinity values, an approximation involving positive, sparse, and non-negative PMT (PPMT) matrix was proposed. While this adaptation offered an intriguing alternative to the original training objectives, the geometric characteristics of the resultant embeddings in comparison to those derived from the original loss remain unclear.
Our approach mirrors that of Levy & Goldberg (2014) to some extent (see also Pennington et al. (2014)), especially in treating NIP as soft-label classification across distinct contexts and presuming unconstrained context embeddings (although these concepts are not clearly mentioned in those original works). At a very basic level our work distinguishes itself by focusing on the NTP objective, which, although akin to SGNS and NC objectives, employs cross-entropy loss instead of sigmoid functions. Aside from this distinction, we identify the following conceptual divergences, contributing to our novel perspective on this line of inquiry: (i) Firstly, we aim to characterize the geometry of embeddings as learned directly from the loss function, rather than through approximations. (ii) Secondly, although NTP- UFM, as a log-biling billeeiin model, is in spirit to the log-bilinear models trained by Word2 proxy to expressive deep
we only use such log-bilinear models as an analytically-tractan neural probabilistic language models such as transformers. (iii) Thirdly, we confront the sparsity in the probabilistic labels of each context head-on, unlike other natrrixffattorzzttion methods that circumvent sparsity in a heuristic manner through smoothing or weighting techniques (Levy & Goldberg, 2014; Pennington et al., 2014). This helps us identify that sparsity leads diverging embeddings, scenario where setting the loss gradient to zero is infeasible with finite weights. (iv) Fourthly, by examining embeddings through the lens of the regularization path-a surrogate for gradient descent optimization unveil that word and context embeddings emerge from the factorization of a matrix Lin + mm Here, in captures the frequency of word-context occurrences akin to the PMI, setting non- occurring pairs to zero while retaining positive values for occurring pairs, much like the PPMI matrix. The directional component L mm becoming dominant as weights diverge, reflects the explicit sparsity patterns of the condext-word probabiitty matrix.
16

---

Published as a conference paper at COLM 2024
The above-mentioned conceptual novelties can also be directly applied to the SGNS objective generalizing the result of Levy & Goldberg (2014) to sparse language patterns. We show this in App. A.4.
The seminal contribution of Levy & Goldberg (2014) has inspired numerous subsequent studies, leveraging the geometric insights of embeddings to uncover linguistic phenomena such as the linear relationships underlying word analogies (Allen & Hospedales, 2019; Hashimoto et al., 2016; Arora et al., 2016; Ethayarajh et al., 2018; Mikolov et al., 2013c). Other works, investigate the amisotropy of the learned word/contex. embeddings for both static models such as Word2Vec (Mu et al., 2017) and contextualized models such as pre-trained BERT and GPT-2 (Ethayarajh,, 2019; Li et al., 2020). This property of the embeddings has been attributed to the imbalanced long-tail distribution of language (Gao et al., 2019; Li et al., 2020), shifts in the embedding space (Biś et al., 2021), and characteristic of the network architecture (Godey et al., 2024). Various techniques have been proposed to improve the isotropy of the embeddings, which yields embeddings with better language regularities (e.g, (Mu et al., 2017; Gao et al., 2019; Biś et al., 2021; Timkey & Van Schijndel, 2021)), We envision that our fresh perspective geometric analysis and our findings could similarly motivate further research in this domain, extending the understanding and application of language model embeddings.
A.2 Overparameterization and Implicit bias/regularization of GD
Recent studies on gradient-based optimization methods, such as SGD and its variants, have focused on overparameterized settings in empirical risk minimization without explicit regularization. This interest stems from observations that deep neural networks, especially in image classification, perform exceptionally well even with minimal or no weight decay. This research area, known as " "implicit bias" or "implicit regularization" of GD emerged to explore which among the numerous potential minimizers in overparameterized settings are preferred by algorithms like GD. Notably, Soudry et al. (2018); Ji & Telgarsky (2018) showed that for linear logistic regression with one-hot labels, GD aligns with the hard- margin SVM solution, essentially becoming the max-margin classifier. Additionally, Ji & Telgarsky (2020) indicated that GD's optimization path in linear models is analogous to the regularization path, which looks at risk minimizers as regularization diminishes. This path is analytically more accessible and offers a useful framework to study GD's behavior. In our work, we apply this idea to investigate the regularization path of NTP training. These foundational results on implicit bias have spurred further research into stochastic and adaptive gradient methods (Nacson et al., 2019; Pesme et al., 2021; Sun et al., 2022; Azizan et al., 2021; Cattaneo et al., 2023), and more intricate architectures (Lyu & Li, 2020; Ji & Telgarsky, 2020; Gunasekar transformers (Tarzanagh et al., 2023b;a); for review, refer to Vardi & Shamir (2021). These insights have also paved the way for generalization analysis, linking GD's generalization to that of SVM-like solutions. This connection has been explored in high-ddimnssiona contexts where traditional margin-based bounds are inadequate, revealing scenarios where regularization yields optimal outcomes, a phenomenon known as hharmless interpolation' or benign overfitting.' Most of these findings, whether on optimization or generalization, are confined to the one-hot classification framework.
al., 2018a;b)) including recent studies on
In the language modeling literature, several theoretical works suggest that achieving the empirical entropy lower-bound at the pre-training stage results in better down-stream performance (Liu et al., 2023; Saunshi et al., 2020). Liu et al. (2023) attributes this better down-stream performance to the implicit bias of the pre-training algorithms, by showing that the optimizers prefer the more transferable models among the ones that achieve the same minimal loss value. A formal deviation from one-hot to soft-label classification for NTP is the work by (Thrampoulidis, 2024), which partially inspired our research. While their analysis, assuming fixed context embeddings, applies to linear models, our work extends this by optimizing both word and context embeddings, leading to nonlinear model. We discover that this nonlinearity yields a preference towards minimizing the nuclear norm of the logits over the Euclidean norm of the word embedding matrix, as seen in (Thrampoulidis, 2024). Our theoretical contributions thus expand on their findings into the nonconvex domain.
17

---

Published as a conference paper at COLM 2024
A.3 Neural collapse geometry
A third area of research closely related to our work is the investigation of neural-collapse (NC) geometry, which seeks to understand the geometry of last-layer features and classifier weights in deep networks optimized in the interpolating regime. Initiated by Papyan et al. (2020), this line of research uncovered a universal phenomenon across various architectures and datasets, where classifiers and last-layer features consistently converge to an equiangu- lar tight frame (ETF). This convergence subsumes the aggregation of same-class features around their class mean, a phenomenon termed NC. The theoretical underpinnings for these empirical findings were later provided by analyzing the unconstrained features model (UFM) (Mixon et al., 2020; Fang et al., 2021; Zhu et al., 2021), suggesting that highly overpa- rameterized networks can produce arbitrary last-layer features. This concept, reminiscent of the early framework used in Word2Vec models (Levy & Goldberg, 2014) as discussed above (see also Yang et al. (2017)), has since spurred extensive research in diverse contexts.
Among the vast body of work on NC geometry-an exhaustive reference of which is beyond our scope-, two studies stand out for their relevance to our research. Thrampoulidis et al. (2022) first adapted the UFM to describe geometries diverging from symmetric ETF by moving away from the balanced data assumption, demonstrating that logits optimizing the UFM with vanishing regularization align with a nuclear-norm SVM problem; see also Behnia et al. (2023). This approach also bridges NC geometry with the implicit bias literature previously discussed. Our work extends their principles to the NTP objective, differing from the one-hot encoding classification focus of their study. We take a similar approach, but this time applied to the NTP objective. This is unlike their result which only applies to one-hot encoding classification. In fact, our formulation in NTP-SVM, extends their findings, encompassing them as a specific instance when each context is followed by only one token-a scenario less applicable to the linguistic contexts we focus on. Moreover, we demonstrate (see Theorem 2) that the NTP framework introduces finite projection component to the regularization path, a feature absent in one-hot settings.
Our findings also resonate to certain extent with the recent multilabel NC geometry study in Li et al. (2023), where each example, typically an image, is tagged with multiple correct labels, represented as a k-hot vector. In contrast, our NTP framework, modeled as soft-label classification, associates each context with a probability vector. This distinction makes our framework more versatile, applicable beyond the restrictive case where soft-label classification aligns with multi-label paasification-sseciffeea equally likely. There are also crucial differences in the outcomes. Their analysis assumes an equal number of training samples of multiplicity k = 1 for each category, a condition that, in our context, would necessitate every vocabulary word be followed by an identical number of contexts that are only followed by that word-a significantly limiting constraint. Our findings, however, are not bound by this assumption.
when all relevant tokens are
In an attempt to extend this line of work to language setups, contemporaneous work Wu & Papyan (2024), through experiments, suggests a correlation between the degree of convergence to the NC geometry and the validation loss for different language models trained on TinyStories for a few epochs. Different to us, they implicitly treat the NTP objective one-hot classification and suggest that better per-class (i.e., next-token) collapse of the context embeddings leads to improved validation loss. In contrast, our analysis reveals more intricate finding: within a sparse soft-label formulation of NTP training with language, it is not per-class collapse that minimizes the training objective. Instead, the objective is minimized when context embeddings with the same support set collapse after being projected onto subspace F which is defined by the sparsity pattern of the context-word co-occurrence matrix (see Sec. 4.3). Our approach also comes with an analysis framework rather than relying solely on experiments.
A.4 Connections to Word2Vec
We apply our geometry analysis to Word2Vec. Consider the following Skip-gram objective:
1 Σ Σ logr(z+)(l1)
Σ Σ
log pa P(2,,, (Ha,,(HW,,)))
T
T [C[T] -cS/scj#0
[C[T] -0
18

---

Published as a conference paper at COLM 2024
where H E RV denotes the embedding matrix of context words Z++j/ j * 0 and W E R denotes the embedding matrix of target words As presented in Mikolov et al. (2013c), skip- gram actually implements heuristic this objective that aims to simplify the computation of the normalizing factor in the softmax. Second, consider the CBOW objective:
1 Σ logp Σ log ((z) (1 clcl)/))) (()))))))) = T te
1
H(m=))) -ESjEc,j0
Sz, (W
Σ
T
2c
tE[T]
where I- I-I---- E- May is the of embedding of the context {Zr1) ([((e()))=0)-
Suppose the CBOW model with context consisting of only the previous word. To keep nota- tion consistent with the rest of the paper, let j [m]] [V] denote index of distinct contexts (aka previous words) and E [V] denote index of (next) words. Note the embedding matrix H E Rd is now of same dimension as W Let ftj denote the marginal of context words and مi denote the conditional probability of word following context word j. Then, the CBOW objective becomes equivalent to
Σ jE(V)
Σ Pjz log (Sz (Whj)) ZE(V]
This now fits exactly our framework and the results for NTP apply mutatis mutandis. In Mikolov et al. (2013c), instead minimizing the CE objective, they actually use the Regative-Sampling (SGNS) objective. Levy & Goldberg (2014) rewrites the SGNS objective as follows:
Σ jE ZF(V]
Σ - Pr(j,z) log (o(w]))) - k Pr( Pr(w) log (1- (1-o 10w hi))
where Pr(j.z) = iz Pr(j)= = Aj,, Pr(z) Lje(m) at k E IN is a hyperparameter. Defining
Pr(is)
Pr(j) Pr(z)
9jz² ==
E [0,1],
Pr(iz)
k+ Pr(z)
we can rewrite the SGNS objective as
[++ Pr(j.z) ( Pr(j) Pr(z)) (-9iz log (o(w) - (1 9iz))oog (1 o(w] h1)))
Σ Σ Pr(j) [c[m]
For each (jar), the rightmost term in the expression above is lower bounded by H(1;2) == - 41 log(q;z) - (1 - qj.2) log(1 -qiz).
By following the same arguments developed for NTP, it can be checked that the lower bound is attained if and only if the following conditions hold:
A. There exists logit matrix L1 IR Vxm (of rank < d) such that
Pr(j,z) ) log(k), Vje [m], E Sj. Pr(j) Pr(z))
L PMI(z,) log(k) log
(t)))
B. There exists logit matrix L2 E Vxm (of rank < d) such that
L2 il 0, [2 < 0,
Vje E 2 S,, Vje E [m], & Sj.
The logit matrix L is the shifted PMI matrix as also characterized Levy & Goldberg (2014, Eq. (7)). This component plays similar role to the finite component in in the NTP setup of our paper. The second component, which is only non-zero if there exists target-context
19

---

Published as a conference paper at COLM 2024
pair (z, with zero co-occurrence in the training set (Pr(j.z) = 0), specifies the directional component, akin to the role L in NTP. This component is absent in the analysis of Levy & Goldberg (2014), they implicitly assume matrix when setting the derivative of the loss to zero. However, when the co-occurrence matrix is sparse, zero derivation of the loss is not achieved at any finite solution. Instead, Le identifies the infinite direction that drives the loss to its lower-bound. Levy & Goldberg (2014) address sparsity indirectly by discussing alternative matrix factorizations such shifted positive-PMI (PPMI).
entries in the co-occurence
B Proofs
B.1 Centering of word embeddings
Lemma For any X O, any minimizer W NTP-UFM satisfies 1 W O. That is, word embeddings are centered.
Proof. The proof is the same as Thrampoulidis al. (2022, Lem. E.1) and thus omitted.
There two consequences of Lemma
A. The word embedding matrix W is centered motivating the use of the centered support matrix Š (rather than simply S) in Proxy (P).
B. The constraint 1 W = 0 can be added to NTP-UFM without changing its optimal set. Thus, the non-convex rank constraint can be ignored, making the problem convex, provided V Finally, note this also implies the constraint 1 L can be added without any change in Eq. (8).
B.2 Proof of Lemma
We prove slightly stronger version of Lemma 1 which assumes the additional constraint
L L = 0 in Eq. NTP-UFM; see Lemma 2. The proof follows identically if we ignore this centering property.
As mentioned in the main body, the proof of the lemma essentially relies on the well-known fact that
||| min 1 W|| 2 |||2. WH-L2
1
Firstly, fix any such that := rank (L) S d and 1 0. Let UIV be its SVD with Rxx, Choose
W = = 1/2 and f = R I IT
for some partial orthonormal matrix R E R (i.e., R R - Ia). It then holds clearly that
WA =
and
[|W]] + A tr [
(VEVT) 2 (E) 2|| II
Thus,
CE(L) + AI CE(WA) + A (1W1]+11((13)
Because can apply this for any L, it shows that"
À [[[W]|2 + [1H]1)).
min CE(L) + +|| > min CE(WH)
+ 2
(12)
LERa
W.H
I T
Heer, use the fact that the minimizer of the H.HS of (12) is attained because of coercivity of the nuclear
20

---

Published as a conference paper at COLM 2024
where we define Ra {L RVxm rank (L) d) to be the manifold of rank-constrained matrices.
Secondly, consider any W E Rxdd H E R such that 1
0. Let L WH and
denote L uEV its SVD. We then have LL 0 and
CE(L) + +||l CE(WH) + À U WHV)
CE(WH) + À (IIT w/l2 [[HV12 2
CE(WH) + À (IW]) + |1112),
(13)
where in the first line used [| tr (Z) tr (u' LV), in the second line used Cauchy-Schwartz, and, in the third line used that | - V S 1. Thus, this proves
min CE(WH) + 2 ([W] +1
À
min CE(WH) + 2 [|W]]] + +|H|??
W.H
W.F
W==0
min IER
CE(L) + A |
(14)
L
where the first equality follows by Lemma 2. Moreover, note that equalities (13) hold and only if the following three conditions hold
WH L UEV T
u WHV I
l W = v H
u w/| = [|] H 1 2-= H|
The first two conditions combined give
u WwTu=vV HV=E
Thus, WW u + + Σ U, for some nen-negative diagonal matrix Σ and U, the orthogonal complement of U. But, the third condition requires
tr u wwwT) = tr (WWT) = tr (E) = tr (Z) + tr = tr = = E E,= = =0
This gives
WW uzu
(15)
and similarly
H VIV
(16)
Combining Eqns. (12) and (14) proves that
min CE(HH) + || min CE(L) + ||I
À
W.H
LER
Moreover, the requirements for equalities in Eq. (13) established in Eqs. (15) and (16) show the desired about the form of gram matrices W, W. and H, as stated in the lemma.
B.3 Proof of Theorems 1 and 2
The proof involves several components, and so, we proceed in steps.
21

---

Published as a conference paper at COLM 2024
B.3.1 Notations recap
We begin by recaping the essential notation used throughout this section.
Recall NTP-SVM,: The constraints in Eq. (9) require in-support logits to be equal to each other for every context. For fixed j [m²), the set of equality constraints for all * z' E Sj is equivalent to the set of the same constraints for any anchor 2j E Sj and z' + 2j E Sy. That
there is an effective total of ; 1 linearly independent constraints for each j E [m]. Definition 3 (Ein.)- For each E [m]] fix an anchor index 2j E Sj and define matrix Ein E () with ; independent rows (e)) ,2, + 2j, 2 E Sj.
With this definition, we can rewrite Eq. (9) since
Ein e 0, Vj E
[m] + L [e,,..., em] satisfies Eq. (9)
The inequality constraints in Eq. (10), impose the requirement that in-support logits must be greater than out-of-support logits. Considering Eq. (9), we can again select any anchor 2j E Sj and reformulate Eq. (10) as Eout;4j 2 1, j E [m]. Here, Eoutj E (V-) has Sj
independent rows (ez, -e) & Sj. In total, there are V constraints per context in NTP-SVM..
With the above definition of the matrices Enj (1) ,j [m], we can define the subspace I of V x m matrices that have their j-th column living on the range space R (Ein))
for allj j E [m]:
= LE IREVV VXm : ( E [m]] span (f(ez - -z) (z)==
e
# E E E [m]] })
Note that this set depends on the support matrix S, but we suppress the dependence to simplify notation. Let Ps: Rvmm + V be the projection operator on this subspace. For concreteness, note that this can be explicitly defined as the matrix operator projecting each column 2; of L onto R(E R(E,,)),e,,
(L) = [Pr(L)1 PL(L)2
(L) with P= == Ein) En E [m].
Note that Pyr is projection operator since it is clearly linear and (P:((P:())) = (L) Also, it is clear that PL( F for all L IR Vxm Operationally, note that (L) has zero-entries on off-support elements. We will use this property indiscriminately throughout this section.
Finally, define the orthogonal complement to F as
5 == L E Vxm - E (N((E)))) E [m] = LERV =mnje; =, [m]
(17)
Clearly, Vvm I F1 Operationally, note that L E Fi implies that on-support elements of every column are identical, that is, [zz = L[z',f] for all z e Sj and j E [m]] In other words, it holds
LEF > satisfies (9),
(18)
which we will use without further explicit reference.
We also let P1 denote the projection operator to T1.
B.3.2 When does the NTP loss attain its entropy lower-bound?
For word/context context embeddings W, and logits
WH, recall the NTP loss is written as
CE(WH) CE(L)
Σ A Σ Pj
log
1+
Σ
Σ
jE
ml
=ES;
DS
(19)
22

---

Published as a conference paper at COLM 2024
and is lower bounded by the empirical T-gram entropy of the data (Shannon, 1948), ie,, for all 0:: CE(WH) 2 H = (x.2)-TT [- log (f(z)x))
We now state explicit conditions on the logit matrix L for which the lower bound is attained by a softmax model = S(WH) with unconstrained context embeddings H. Due to sparsity of the conditional probability matrix P, the lower bound is not attained for any finite logit matrix, but it can be reached asymptotically provided the conditions stated in the proposition below hold.
The statement below is slight extension of Thrampoulidis (2024, Prop. 1).
Proposition 1. The NTP objective of softmax model S(WH) with W E R Vxd H E Rd can reach the empirical entropy lowerr-bund if the following three conditions hold simultaneously
A. INTP -NTPp-compatbbilitty There exists logit matrix L1 E RV xm such that
Piz Pj
Vj E 2 # [.l.2]] Lu = log
(20)
B. INTP-separability) There exists logit matrix L2 E RR Vxm that is feasible in NTP-SVM..
C. [Rank constraint] For all arbitrarily large positive scalars pr
rank (L1 + L2) s. d.
Concretely, if the above conditions hold, then for all arbitrarily large p. there exist Wp E RVxd and Hp E Rd such that
lim (E(WHH,)) -H. p-000
Proof. It suffices to prove that for Lp == L1 + pp L2:
lim CE(Lp) = H. 100
(21)
To see that this sufficient, note from condition C of the theorem that for all arbitrarily large Lp admits a factorization Lp = Hp where Wp E Vxd and Hp E Rdxm Since this holds for all large p, Eq. (21) implies the desired.
Thus, we now prove Eq. (21). For any j E [m], is easy to check, using conditions A and of the theorem, that the following two statements true for all p > 0:
A. For all + E Sj:
Lp|..11 Lp|z',11 = (L1 L( + (L2[z.fl 12(2',1)
=o0g(p/z/ (00)/////1))
B. For all Sy.0 & Sj:
Lp Lp(D,
(L| [b,, + (L2[z, L2(0,
> (L [2, L [v, +
> 2 L1l2 +
In the last inequality above, we used the loose bound || maxje(m)z=(() 1L [z, for the spectral norm of L1-
23

---

Published as a conference paper at COLM 2024
Using these in Eq. (19) gives:
PJæt piz
CB(Lp) <
Σ
Pjz log
1
E ='#IES
+ Ve Pe- -e-1lal))))))
jE m
ZES
Σ jE(m]
Σ z
P1z
log
1
+ Ve - 2112)
Piz
Σ jc[ml
Σ
Pjz log
E
n1 Σ
Piz log ( PizVe Ve Pe 21Ibilla).
Pjz
zCS;
jc|ml
ZCS
$ H + Ve
Ve- 1
(22)
The last inequality uses log(1 + x) S x, 0.
Clearly the bound above converges to H = Eje(m) Fj Excs, Pjz log this and CE(Lp) H. the desired Eq. (21) follows. This completes the proof.
1 - co. From
Piz
Note that the first two conditions of the theorem are easy to satisfy.
On the one hand, Eq. (20) always has solution. For later use, also note that the solution unique on the subspace I and we denote this solution i This easy to check: for any
j [m], Lj is solution of Eq. (20) if and only if it takes the form ej Einj (Ein aj v, v E N (Ein.) where aj E R has entries log(P):= / ), & z E Sj. Thus, iin is
explicitly defined as follows.
Lemma 3 ). Let I'm be the unique solution of Eq. (20) on the subspace F. Then, Lin has columms
- =ETE] (Einj ETn
aj Vj [m].
where a; E R has entries log( / PIzz) for z' # 2j E S..
On the other hand, as mentioned in Sec. 4.1, NTP-SVM, is always feasible.
The challenge is in satisfying conditions A-C simultaneously.
However, when d 2 V, then the rank condition C of the theorem is trivial. Thus, we have the following corollary.
Corollary Suppose d > V, then the entropy lower bound can be asymptotically atlained by the NTP loss.
B.3.3 Ball-constrained CE minimization and main result
Instead of the regularized problem (8), it is more convenient (but, equivalent due to convex- ity) to study the following ball-constrained version:
Lg arg min CE(L),
(23)
parameterized by B > We will study properties of the minimizers LB of (23) as B co. The theorem below subsumes Theorems 1 and 2.
Theorem 3. Suppose 2 V. Denote L any solution of (23) for regularization B. Then, in the limit of diverging regularization the following statements true.
(i) limp CE(LB) = H.
(ii) limg Pg(LB)
Lim where Lin is defined in Lemma
(ii) The solution diverges in norm: limg yo0 |اااا||
+
24

---

Published as a conference paper at COLM 2024
(it) The solution converges in direction to the optimal set of NTP-SVM,. That is, there exists minimizer mm of NTP-SVM, such that
LB
L L
lim -0000
O.
[B
B.3.4 The in-support NTP loss and its minimizer L
Define the in-support NTP loss as follows:
CE (L)
Σ Rj
Σ ZES;
Pjz
log
1+ Σ
JE|nz
2/EE
Contrary to (19), the summation inside the logarithm is only over word indices z that belong to the support set of context j.
By non-negativity of log, for all L it holds CE(L) 2 CEin (L). In fact, because of the sparsity
ansatz (3j [m] such that S < V) and strict nonegativity of exponential, the inequality is strict for all finite L,i.e
CE(L) > CEn (L,(LL,, 7L E R
Vxd
(24)
Moreover, observe that for any holds that
CEnn (L) = CEn (Ps(()))))
(25)
To see this, decompose L in orthogonal components onto I and F1, respectively, i.e.,
L= P (L) P1pP1111 (L). By Eq. (18), L1 = Pi(L) satisfies Eq. (9), L [, L [z', for all Σ E S and E [m]]
The lemma below uses this to show that CEin has a unique minimizer in in
Lemma 4. The in-support loss CEin (L) has a unique finite minimizer in 3, which we denote in;
L arg min CE (L). LEF
(26)
The minimizer L in satisfies the NTP compatibility condition in Eq. (20) and
C = (Z.) H.
Proof. Note that L E T is column-wise separable constraint l; E R(ET))) that restricts the non-zero entries [[ f to indices Sj. The objective is also separable with respect columns. Without loss of generality fix = and assume S = [Sy]. It then suffices proving that the function f R R:
exp(fz) EvES, exp((x)
f(l)
Σ log zCSI
(27)
has a unique minimizer in R(E where Ein.1 == Ein (1 E (S]-)) is the non-zero
block of Em,1-
By non-negativity of the KI. KL. divergence,
f(e) 2
Đ1 log (f1,2) zeSi
Σ
with equality if and only if,
exp(lz)
Az E S Ezes, exp( = P1.2-
(28)
25

---

Published as a conference paper at COLM 2024
In turn, this is equivalent to
Vz, E S1 1- lz log(P1.//1.z))
(29)
To see the equivalence note the following: On the one hand, starting with Eq. (28) for z, E S₁, dividing both sides of the two equalities, and taking the logarithm, we arrive at Eq. (29). On the other hand, if Eq. (29) holds, then Eq. (28) holds since
Iees, exp(L2) exp((2)
Ee'es, PI.2'
exp(ey (x)) 2'ES
E
Pıæ
P1,z
Now, observe that (29) is the same as the y equation for = 1. Moreover, note that a logit vector l E R$ solves (29) if and only if for
E =
(Ein)Ema)
a1 + a DE N (Em.).
where E R$ with entries log(P1,7 /P1,') for z' + + 2 E S1. This proves that there is
unique minimizer e := Ei, (Eim, a on R (E).).
Since the choice of j = above was arbitrary, this proves the lemma.
B.3.5 Auxiliary lemmas
This section proves a sequence auxiliary lemmas that are used to prove Thm. in the next section.
The first lemma shows that entropy can only be approached asymptotically, i.e,, provided the logit matrix diverges.
Lemma 5. For all L RV Rxm and all L' J1 that additionally satisfy Eout, 1, Vj E [m]] (i.e, L' is NTP-SVM, feasible,, it holds
CBn (L)
lim CE(L+ + RL'). R-+00
Using this, it holds for all E RV that
lim CE(L in + RL')= CE, (L CEi,++ H. (L) CE(L). R-00
(30)
Proof. To see the first claim note that L' E Fi implies
VjE [m], + z' S; L' L'j
Vje mm 2 Sj
Σ e (Lie LE L
Σ (L|a LFEAA)
2'ES
2'E
and Eout e 2 1, Vj E [m] implies
Vj [m],z E S $ S L'(j,z] - ['
Vje z E S Σ e -e v4S
R- 0.
The second claim applies the first claim for L = L and uses Eq. (24).
We now use the previous lemma to show that the ball-consrtained minimizer diverges as the ball constraint approaches infinity. Lemma 6. The norm of the ball-constrained minimizer diverges, that is
lim 8-
[[
26

---

Published as a conference paper at COLM 2024
Proof. Suppose the contrary that lim
|B From Eq. (30), have
CE(LB) lim CE(L + RL m R-+00
Moreover, by triangle inequality for nuclear norm | S R Thus, for any arbitrarily large R > 0, there exists large enough B such that ||< S B, making LR L + RL feasible in (23). These together, contradict optimality of LB, proving that limg- |[Bl|++ co. Finally, recalling the norm inequality |[ z |[B||+/ it also holds that limg- [ co.
The last lemma constructs sequence of logits approaching the entropy lower bound while remaining feasible in problem (23) as B - co.
Lemma 7. The objective function approaches the loss lower bound, i.e,,
lim CE(LB) Clin,+ H.
B-+00
Concretely, there exists logit matrix LR parameterized by R = R(B) that is feasible in (23) and sastisfies:
CE(LR) S CEin, + 2
(31)
Moreover, limg-co R oo and limg [[a NTP-SVM,
where
L mm [|. is the optimal cost of
Proof. We evaluate the loss achieved by the following candidate good point
LR In + R(B) mm
Here, (for large enough
|| |L (I.) we set
| mm
R R(B)
Note this is chosen so that || [xI|+ S LBll= < towards making LR feasible in (23). Also, recall from Lemma 6 that [L + + co as B - co; thus, also R + co.
remains to prove Eq. (31). This follows directly from the sequence of equations in (22).
B.3.6 Proof of Theorem 3
We prove each statement separately.
Statement (i). Since for all L, it holds CE(L) > H. it suffices to prove that H can be asymptotically attained. Since d 2 V, this follows from Corollary 2. (For more explicit proof, see Lemma 7).
Statement (ii). Since for all L it holds CE(L) 2 CEin (L) 2 CEin, = H, it follows from Statement (i) that
lim CEnn (LB) CEi,, 300
But CE, (L) = (L)() (see Eq. (25)). Thus,
lim CE (P (Ln) CEin+ B3000
The desired now follows ince iin is unique minimizer of CBnn on T by Lemma 4.
Statement (iii). Follows directly from Lemma 6.
Statement (iv). From Lemma 7, there exists Lg for which the loss is close to H (see Eq. (31)).
27

---

Published as a conference paper at COLM 2024
Towards arriving at contradiction, we will show that if La is not in the direction of Lmm, then it incurs a loss that is larger than the upper bound of CR() computed in Eq. (31).
To do this, assuming the statement of the theorem is not true, will lower bound
CB(LB) H
Σ
Σ ZCS)
Pjz log (P1=
Σ
[-(Lalia)-Lalier
Σ -((alii))-aali>)))
Z'ES;
vÉS
(32)
By our assumption, there exists > 0, such that there exists arbitrarily large B satisfying:
|m
LB L. mm
c,
(33)
|B
for all minimizers mm of NTP NTP-SVM. Define
(LB ")
('(B)
where, ;= (B) > 0 is chosen so that || + L mm ||.. Concretely, this can be ensured by setting (to see this use the triangle inequality for nuclear norm):
(|( x + L 1n +
R'
mmm
for some ? > 0. Recall that [[s||- + co as - co.
On the other hand
|| is not
dependent on B. Thus, the large limit B - co, it holds 1
きे +
and
-
[Lg
Lgll.
These combined with (33) show that we can always choose B large enough so that Eq. (33)
guarantees, for some e' > 0, that
L L mm
d'
But L achieves the optimal cost of NTP-SVM,, since || |[||+ < | ||x. Thus, there exists
5 E 10, and] j [m] such that at least one of the following is true (i) 3z and z' + # S such that
[| - =|j - 2 8,
(34)
(ii) 3z Sj & S such that
- Lij 2) < 8.
(35)
Case (i): Without loss of generality (ez-ez) e 8 s -5 (otherwise, flip z,2'). Thus, ignoring all but one term in (32) gives
CE(LR) - 2 2 log Pi (Taalii -Teliz)))
(36)
But,
L Lo E0]j.21 R' ([ljjz - LIj,=')) + - 1,2)) - -R'5 2| (37) Put in (32) and using Pi, 2 1/V At N 1/ m shows
CE(LB) 2 H + mV
log (-
Re''s
1 R's exp( 1xp(-211 112) 2 76 + mV log V eexp(2)
1
eR's
Compare this with (31), it is clear that B co, that R' - co, it holds wt's MV log (eespielLmm,)) > V e eR Thus, CE(Ls) CE Lim), contra-
diction.
28

---

Published as a conference paper at COLM 2024
Case (ii): We can assume Einj& = 0 for all E [m]) since otherwise are in Case (i). Now, again ignoring all but the (j,z) term in the CE loss for which (35) holds for some Sj, we find
CE(LB) - 2 R,Pi2 log (iiz ( Σ
e ([nlia))[a(iz>))) + e(Laliaa-Laliin
Z'ES,
On the one hand, using Einje = O, we have
Σ
-e-(Loli=
Σ
Σ
Pjæt
Piz
Z'ES;
2'E
2'co)
On the other hand, by (35):
Jubisl Palial-Talike Isliel e (1 e - [ia 2 -R' R((1-))
112),
Putting the above together yield:
11+ e (1((1-0) exp(-2211 exp(-2)(((m12))
e' e-R((1-0) m V2 exp i (2) mV
CE(LB) H. A= log
V
where the second inequality uses log + x) 2 1 > O.
Compare this with (31). As B + co, and noting that R, R' grow at the same rate, it holds (1 mVV 22 > V exp (1 12) Thus, CE(LB) > CE(Lg), contradiction.
In either case, we arrive at a contradiction, which completes the proof.
C On the solution of NTP-SVM,
To better understand how the SVD factorization informs the geometry, we investigate further how the structure of L mm (solution to NTP-SVM.) depends on the support matrix S.
C.1 Subspace-collapse
Proposition 2 (Subspace collapse). Assume NTP-SVM, is feasible. There exists minimizer hm] of NTP-UFM such that for all contexts j.j' E [m] with same support set
H hh,.... .m]. S Sp, their embeddings are same direction in the limit of vanishing regularization, i.e.,
lim hp = 0.
This is consistent with our empirical observations in Figs. and 2: As the loss approaches its lower bound, context embeddings with the same support set converge in the same direction and exhibit maximum correlation. This is indicated by the bright diagonal blocks in CORR(H) H) heatmaps.
The proposition can be proved by inspecting NTP-SVM, Since the affine constraints (9) for each sample j E [m] depend solely on the support set S,, it can be shown that there exists a minimizer L mm = [ {mm) such that p = ém whenever Sj = SSy. This
implies that the corresponding columns j. of V are identical. As result, by Cor. 1, lim -0
= 0.
important to reiterate that Prop. does not prevent the minimizer L, = W,H, from accurately recovering the distinct soft-label values p; Py and achieving the empirical entropy lower-bound H. as À + 0: The collapse occurs subspace F1, orthogonal to the data subspace I on which the logit matrix satisfies the log-odds constraints indicated in Thm. 2; See Secs. D.1 and 5 for numerical verifications.
29

---

Published as a conference paper at COLM 2024
V-10
V-dn
0..6
-0.5
0.2
-1
0.5
0.5
Figure 4: Geometry properties illustration for Proposition 3. Shown two values of V for varying values of k C [V 1]. (Left) angles between context and word embeddings. (Right) normalized norm ratio of word to context embeddings.
C.2 Special case: support sets of equal sizes
To gain insights, start with idealized symmetric setting.
Proposition 3. Fix k E [V - 1] and suppose S &contins all m = (K) support sets of size k. Then,
(i) The logit matrix takes the form L (Iv 111
(ii) Word embeddings form equiangular tight frame (ETF) being equinorm and maximally separated.
(iil Context embeddings are equinorm and the embedding h of context j is co-linear Ezes w2-
In the symmetric setting of the proposition, we can analytically solve NTP-SVM, giving m as expressed in the statement (i). In fact, it is possible to obtain an explicit characterization of the SVD of L mm which enables the precise definition of embeddings' geometries in statements (ii) and (iii). This explicit characterization leads to closed-form formulas for the angles and norms of embeddings, detailed in Prop. 4.
Proposition 4. Assume the setting of Proposition 3. The geometry of context and word embeddings is described by the following relations:
-1
Vo # d' E [V] cos (os(2w,, w. V 1
and
||w,|| = ||
(38a)
18 Sp V
vj j E [m] cos cos( kty)
and
|| = 1
(38b)
V E S
Vj E [m] v E [V] cos(wp.) 11))
R(V
(38c)
E S
VK(V-R((V))
|wo 112
(V 1) 3
EjE [m]] E [V]
(38d)
h k(V - k)
For k = 1 the embeddings recover the ETF geometry that has been previously shown for the setting of one-hot classification (Papyan al., 2020). For general k 1, word embeddings continue forming an ETF, but the geometry of context embeddings changes although they remain equinorm. Embedding h; forms the same acule angle with all in-support word vectors Wo, v E S and the same obtuse angle with all out-of-support word vectors Wo, v & S,. Additionally, for 1, norms of word embedddings become larger than norms of context embeddings. See also Fig. 4 for a visualization of these properties. We defer the proof of these two propositions to App. C.5.3.
C.3 A simple candidate solution
Achieving an analytic expression for L mm as was possible in Prop. 3, might not always be possible. This challenge primarily arises from the combinatorial complexity of the
30

---

Published as a conference paper at COLM 2024
TevxPY
5.75
8.75
1,00
4.00
-2
-275
Figure Comparison between the solution found by CVXPY for NTP-SVM, and defined App. for two realizations of the support set matrix See text for details.
constraints that adhere to the sparsity pattern of S. To mitigate this issue and potentially circumvent the need to solve the semidefinite program in NTP-SVM,-whic becomes computationally intensive for large V and n-we introduce a strategy that allows for numerical verification of whether Š of Prop. solves NTP-SVM..
Š readily satisfies the feasibility conditions for NTP-SVM,, meeting all inequality constraints with equality. Below, we further introduce dual-certificate condition that, when fulfilled by 5, ensures optimality.
Proposition 5. Let Š = UEV denote the SVD of = (Ly - f11 Define, == uV Vj [m] and all Sj, it holds that Alu, i] 0, then solves NTP-SVM..
Verifying the certificate's 's conditions only requires an SVD of Š. In situations where these SVD factors can be analytically determined, the certificate also enables formal proof of optimality (e.g. Prop. 3).
In more complex settings than that of Prop. 3, whether Š satisfies the dual certificate in Prop. 5 or not, and thus its optimality, depends on the sparsity pattern S. Our empirical evaluations show that Š successfully solves NTP-SVM. in numerous realizations of S we have examined. However, this not the case in all realizations of the support set structure.
In Fig. 5, we present an example of two realizations of S, where the condition in Prop. met in one instance but not in the other. The second and third column of the figure display the heatmap of Š and the optimal solution of NTP-SVM, as determined by CVXPY (Grant & Boyd, 2014) for the support sets shown in the first column. Notably, although Š not optimal in the second row, provides a close and simple approximation the optimal mm which almost captures the overall structure of the optimal solution found by CVXPY. Additionally, in Fig. 6, we show similar heatmaps the optimal logits for the setup of the Simplified TinyStories dataset in Sec. 5. Once again, while is not necessarily optimal, serves as close proxy for the optimal solution. Note that in the case of one- hot classification under STEP imbalanced data, is previously shown to be the unique minimizer of NTP-SVM, (Thrampoulidis et al., 2022).
Although not necessarily optimal for all configurations of the support set, we conjecture that š, gives an approximation with almost similar geometric properties of the optimal solution m This conjecture motivates more efficient proxy for predicting the implicit geometry in the next section.
C.4 A proxy for NTP-SVM, solutions
Considering (Iy 111 S, ie., the column-wise centered support set, as a close proxy of the optimal solution of NTP-SVM, we provide connection between the support sets
31

---

Published as a conference paper at COLM 2024
0.75
CVVYY
0.8
0.5
0.00
0.00
200 402 130 203
420
110
300
400
-0.50
-550
-100
Figure 6: Same Fig. for the support set configuration in the Simplified TinyStories of Sec. 5.
and embeddings geometry. Suppose Lmm = S. Then, by Thm. 1,
wmm.pmm T mm L mm ) = ($S )3,
Hmm Hmm (L.mm L ($ $)
(39)
Computing the matrix square roots in Eq. (39) is as expensive which can be prohibitive for large matrices. Instead, we suggest the following more computa- tionally efficient proxies for estimating the directional component of the implicit geometries:
SVD decomposition
Wmmm Wim T
~ SS
1 (Iv 1 1v (Iv ,
v (v1l),
1
1
H
5 S S (Iv ~ 1v1b)s
(40)
This leads to Proxy (P) introduced in Sec. 1.2.
Ignoring for simplicity the projection the Iy - + 1 subspace, Eq. (40) gives a simple ex- planation for similarity between embeddings: For given context E [m], the corresponding embedding has higher correlation with contexts j' whose support set intersection S is larger. Similarly, word embeddings are closer if the respective words appear together in more support sets: word embedding w; more similar to w.r when I W2, w. E Sj, E [m]} larger. The projection then fixes the bias of the parameters: since fidgereguulrriztion is strongly convex and softmax is invariant to shift, any minimizer of NTP-UFM satisfies W see Lemma 2).
C.5
Proofs
C.5.1 Proof of Proposition 2
We will prove that for j. E [m] with same support set S - the j-th and j'tth columns of L mm Ll.... ... Lm] are same. For the sake of contradiction, suppose that this is not the cases, i.e. E + ly. Consider the following three candidate solutions of NTP-SVM.
L mm
[e,,... ey................. ep............. Lml
I. mm flip
[e...., lgp....
ey....em
L t.
l + +
I Lav
[l....,,
P.
By assumption L mm is optimal in NTP-SVM.. We now show that the other two matrices are also optimal.
mm is feasible in NIP-SVM, because: (i) the affine constraints only Firstly, note that fip depend on the support sets, which are same for j.j's, (ii) rank mm 'lip
rank (L S
since permuting columns is rank-preserving operation. Moreover, permuting columns also
32

---

Published as a conference paper at COLM 2024
preserves nuclear-norm of matrix. To see this note that Lip = L UE(P v) for some permutation matrix P E R= Since V is partial orthonormal, the same is true for P V. Thus, || mp l| = tr(z) = L | which proves optimality of rank (Lmp)
Similarly, we can argue that Lmm is feasible in NTP-SVM, because: (i) the affine constraints
only depend on the support sets, which are same for j.J's, (ii) rank (Laun d since adding columns and multiplying them by scalars are rank-preserving operations.
= rank (Lmm) <
But now, from convexity of the objective function, and noting that Lavg [m + flip we have
mm
1 mm
L
mm avg S
+
mm 2 nip
I| = =| mm
2
Thus, Lavg is also optimal.
Now, from Lemma 1, we know that for some partial orthonormal matrix R, in the limit of vanishing regularization: HA, R Σ VT where Σ are SVD factors of L mm Since the j, columns of Lavg are the same by construction, the same is true for the j. columns of V This proves the desired.
C.5.2 Proof of Proposition? 5
To prove the proposition, we first derive the dual of NTP-SVM,.
We begin with more convenient, but equivalent, formulation of NTP-SVM,- One way to do this is by recalling that NTP-SVM, is a convex relaxation to
2 | W
2 || s.t. Vj 2, = (e,,-e) Wh;
min
VjE mm] E S (ez,- -b) Whj 2 1.
where, for each sample we choose an anchor label 2, E S,. We can now relax this in terms
W
of the matrix X
W
H]. (Note that L essentially corresponds to the off-diagonal
H
blocks of X). This gives us the following reformulation of NTP-SVM.. (Formally, it is easy show that (41) the dual of the dual of NTP-SVM.. This result about nuclear minimization is classic (e.g., Recht et al. (2010)) and we omit the details.):
min
tr) s.t. VjE SZE X=0 2
X1z,, V + - X V + +1=0 =0
(41)
VjE [m]. Sj X(z, V + X(0, V j| 21.
Define Ayxm such that A(z, = - - z + zj and A A(zz. = Latzy A/z- (1,AA = 0).
The dual problem is
1
I ([A
-A -Al
I T
max min X=0
Σ
ja
tr
-A-
x)
s.t. 1 -0 Al0,100,00
2
jE
D
max A
Σ
Σ
Ajo
s.t. 1/ 0
A 0, VJ
||A||2 $
jE[m) S
-ma tr(AT 5)
s.t. 1 0
Alo < v SS |
A
33

---

Published as a conference paper at COLM 2024
where here $ is defined in Prop. 3. The last line above holds since,
Σ E je[m]
Σ DgS
1j.0
+
E
E
A
Σ
Σ
Ajo
j m² PES;
je/m]
DES;
je/m
PES; oF2
L
E jæ
Σ
Σ
1j.
jF(m] oF2
jelm
DES,
Jtzzz
|
S
Σ
Σ
Σ
Σ
Ajo
jc(m]
DCS; je(m]
DF22
- |S
S
Σ
A |
E
E
A
E
Σ
A[o,j]
V
jc(m]
je
je m DES
Dz
= Σ
Σ
Σ
Alv.1 +
Σ Σ
jE[m]
JE[m PES
[MMl UC
D222
E jE[m] [E
S(u,;|A(u,ji
We remark the following about the dual problem derived above. First, because S is feasible in NTP-SVM, and constraints are linear, Slater's 's conditions hold, thus, we have strong duality. Second, by complementary slackness, if the nonegativity conditions A [v, jl S 0, v & Sj are strict, then, at optimality, L satisfies the inequality constraints (10) with equality.
Now, regarding solving the dual, observe that if we remove the constraint A Alv, S 0, & Sj,
then the solution to the dual is simply A = uv where UΣV is the SVD of S. In fact, is not hard to show that this is the unique solution of the relaxed dual; see for example (Thrampoulidis al., 2022, Lemma C.1) for an analogous result. Therefore, UVV satisfies the condition of the proposition, then it is clearly the unique optimal solution to the dual problem.
From this and complementary slackness discussed above, it must be that any minimizer solution i satisfies simultaneously:
of NTP-SVM, must satisfy the inequality constraints with equalities. Formally, the optimal
L Iy Om and Vj E [m] En = 1s, Eoutie OV-1-5
Note that this is system of + +Ejejm) be easily verified that these equations are linearly independent. Thus, their unique solution S. This completes the proof of the proposition.
Emj (S; - - 1 S,) = Vm linear equations. It can also
C.5.3 Proof of Propositions and
At the heart of proving both Propositions 3 and 4 is the following result, which determines the SVD factors u, Σ, and V of Š (Iv 1 This shows immediately that word embeddings form an ETF. It also forms the basis for computing the geometry of context embeddings. We show this after proving the lemma.
Lemma 8. Fix any k E [V] and suppose S contains all (%) support sets of size k, i.e. m = (2) Let P VV--V) be orthonormal basis of the subspace orthogona to lv. Then, the U, E factors of the SVD of Š are: andd Σ = v (V== -1v Thus,
u (Iv V (Iv-v 11)).
34

---

Published as a conference paper at COLM 2024
Moreover, the matrix V E (V-((V-1) of right-singular vectors is such that its rows Dj. [m]] can be expressed follows with respect to the rows uz, E [V] of of
Dj
uz jE m|. ZESj
Σ
Proof. We will compute explicitly SST .Strrt with computing SS For diagonal elements, the dot product is between a row of S and itself, counts the number of is in that row. Since each element is included in (K=1) support sets (choosing the remaining k elements from the other V - elements), each diagonal entry of SS is 1 For off-diagonal elements, the dot product counts the number of support sets in which both elements of the corresponding rows appear. This requires choosing the remaining k elements from
V - 2 elements, giving ( for each off-diagonal entry of SS'. Overall, after algebraic simplification, we find that
V
ss'
Tv
V V-k
We may now use the closed form of Š and few more algebra simplifications find that
SS = V - 2 k (TV-
1 T 11
Let P E RVX(V denote an orthonormal basis of the subspace orthogonal to fl.. Since
SST -K P and 5S -u=u we have shown that U P and I / 3 lv. O
Using 1 S = k1 mi we can compute
5 S (s T ST
k2
V 1mmn
From this and the fact that diag (S T S) = klm, and recalling that ST VV2VT, we find that all context embeddings are equinorm with squared norm (i.e. diagonal entries of VEVT) equal to
k Vj [m].
|1,112
Recall also that the diagonal entries of uEu are
|1wl12 = -1-1 (V- -V-2) 22 k-1) Vo vo [V].
From the off-diagonal entries of ŠTS, we can infer the angles between different context embeddings. Let * [m], then the (i.j") of ŠT is |s, n &. Thus, the (j. f') off-diagonal entries of VEVT are
h hy = S
cos(h,, = 18 Sp V
i E [m]]
(
k
Recalling that WHT = mm may now compute the angle between word and context embeddings as follows:
V V
VES, S 060 OES
cos W hi)
(42)
VRIV R)(V 1)
35

---

Published as a conference paper at COLM 2024
This completes the proof of Proposition 4 and the proof of statements (ii) and (iii) of Proposition 3.
It remains to prove statement (i) of Proposition 3, i.e. proving that Š = (Iv } s is the unique minimizer of NTP-SVM..
To prove this, we appeal to Proposition 5. First, it is straightforward checking that m satisfies the SVM constraints (in fact, with equality). Thus, from Proposition it suffices that the matrix := uV satisfies A [,,
/] < 0 for all v g S and all E [m]. But, our case
A = uV = (==) 2 because Σ = Ty-1. Thus, the desired condition holds for A since $|0, k/V < 0 for all j & S and [m].
D Numerical experiments
D.1 Numerical simulation: NTP-UFM
Here, aim is to confirm, through numerical experiments, that the solution derived from (stochastic) gradient descent when optimizing NTP-UFM is consistent with the analysis presented in Sec. 4. We train NTP-UFM with vocabulary size V = 10, embedding dimension d 10, and m 95 training samples whose (centered) support matrix Š is shown in Fig. 7-(c). We set ;, = 1/m and for this fixed support matrix S, we generate random soft labels pj E A V- such that Pi = 0, g Sj.
We display the training evolution in Fig. 8. We train NTP-UFM with SGD and small weight decay A 10-5 until reaching the empirical entropy lower-bound H. 1.04 within an order of 10 as shown in Fig. 8-(a). This decrease in the loss accompanied by a consistent increase in the parameter norms, displayed in Panel (b). Despite the increase in norms, the projection of L onto the subspace F remains close to the finite component L specified by Thm. 2. This is illustrated Panel (c) and ensures the recovery of the soft-labels on the in-support tokens.
In Fig. 7, visualize the logits L, WyH,, and cosine similarities CORR( W;)) and CORR(H; (Hş) of word and context embeddings, respectively. As shown in Panel (a) of Fig. 7, as training continues the parameters recover the implicit geometry in Panel (b), which vi- sualizes the prediction made by Thm. and Cor. 1. In Panel (c), we also visualize the cosine similarity of Š and ST, which we introduced as a less expensive proxy of Wmm and Hmm (see App. C.4). While the theory and proxy not exactly the same, their heatmaps display similar structure. To verify the close relationship between the embeddings and the proxies
quantitatively, we measure SIM H Š) and SIM (MM.WT S where SIM Y) measures the structural correlation between the two matrices (see Sec. 5). Fig. 8-(d) confirms the high corre- lation between the proxy (P) and the embeddings' implicit geometry. We also note that in the special case where support sets of two contexts are identical, the heatmaps also verify the subspace collapse Claim (C3): Along the diagonal blocks of CORR (H), where the contexts with similar support sets lie, the samples have maximum correlation and are aligned.
We finally note that in this experiment L = S (we verified that the dual-certificate condition in Prop. holds).
D.2
Additional details and results: Sec. 5
D.2.1 Datasets
We use a total of three datasets. We first employ two smaller-scale datasets, which are well-suited for verifying theoretical solutions. Then, use one larger-scale dataset designed to investigate the geometric properties of text in conditions that approximate real-world text scenarios.
36

---

Published as a conference paper at COLM 2024
(a) Training checkpoints
(b) Theory
(c) Proxy
Ppach
Epoch
Epoch 1000
Epach
O
Epoch
Epoch 1000
Epuch
CORR(Wam
CORR(
CORR
Epoch
Epoch 100
Epoch 1000
Ppach
(ORR(
COaR(S)
CORRI
64
Figure 7: Evolution of NTP-UFM parameters Lg, CORR(W, and CORR(H.) when training close to convergence to the empirical entropy H (See Fig. 8). At the end of the training, the parameters align with the prediction of Thm. ((a) vs (b)). Additionally, the correlation patterns between the embeddings closely follow the similarities between the support sets ((a) vs (c)). See App. D.1 for details.
convergence
Norm growth
(c) || [1)) 1
Correlation with proxies
CE(L)) - H
13a
13
0.00
Epoch (8)
Epoch
Epoch (R)
Epoch (R)
Figure 8: Numerical expermients NTP-UFM. (a): CE converges closely the empirical entropy H.
(b): Norms the parameters grow, (c): While the parameters converge directionally to mm (Fig. 7),
the projection of L on F converges the finite component Z specified by the soft-labels Pj. (d):
Structural correlation between the learned embeddings and the proxy, i.e., SIM SIM (WT, ,.T).
The correlation with Hmm and Wmmn T (instead of Š, 5T) is displayed with dashed lines for reference.
Synthetic. We generate each context žÿ as follows: We select T 5 words (tokens), and manually come up with the T 6-th tokens that are consistent with the given context. To model the probabilities ft and p,, we sample each support independently to number of repeats to emulate the behavior of repeated context in natural language. For instance, the context ; = "Lily wants try the" followed by S {soup, game, movie) with respective probabilities [0.5, 0.25,,2,000 0.25). The dataset consists of 116 samples, containing 16 distinct contexts, and vocabulary size of V 30. Each context has a fixed support set length of S, = In Fig. 9, we show the m unique contexts and the soft labels on their next token. this dataset, the tokenization is done at the word level and the empirical entropy lower-bound is H. 1.6597.
Simplified TinyStories. Advancing towards a more realistic but still controlled dataset, we use contexts and support sets from the TinyStories corpus. We derive contexts &, and sup- port sets S from the TinyStories dataset and we create the training data from the most fre- quent word-level contexts with length 5, such Xi ["once", "upon", "a", "time", " ". For the support sets S;, record all next tokens of x; in the original dataset. Then, replace the words in the contexts with their synonyms to generate new contexts with identi- cal support sets. This allows us have multiple contexts sharing common support sets, while controlling the vocabulary size, which is set to V 104. The final dataset consists
37

---

Published as a conference paper at COLM 2024
.995 20 J9...n
Jy. 2h0
J0H Wr 2 mil
Jos
Jol. Jacase 2l0
Jork Jes Jesie 2a6
ce Jeraen Jks
Figure 9: The contexts, support sets, and soft labels in the Synthetic dataset.
of n =: 3050 samples with m ~ 400 unique contexts. The frequencies P; are determined by independently sampling each next token from Sj several times. Here, 7 = 1.0166.
TinyStories. To experiment with more standard dataset, we use 100 stories sampled from TinyStories. Here, we do not make any sample selections for contexts and support sets, unlike the previously mentioned datasets. We use a tokenizer with vocabulary size V = 128, and fixed context length T 1 We choose a small context window and vocabulary size to make tracking the distinct contexts and their support sets computationally manageable. In this setup, the number of distinct contexts 10 and H 0.3112.
D.2.2 Training details
Synthetic and TinyStories experiemtns. We train the models long enough to ensure that the loss approached the empirical entropy lower-bound H within an order of 10-4, In all experiments, we use AdamW optimizer with a weight decay À = 10-6. The learning rate is initially set to 10 with a step-decay schedule.
TinyStories experiemtns. We train TF with AdamW optimizer until it reaches the empirical entropy lower-bound within an order 10 3. We use warm-up over the first epochs to increase the learning rate to 1 and use a cosine learning rate scheduler to decay its values over the course of the training. We set the weight decay to 10-5.
D.2.3 Additional results
Figs. 10-17 complement the discussions and experimental results in Sec. 5. Below we only discuss the missing details of the TinyStories experiments.
Visualization details of TinyStories experiments. For visualizing the heatmap of embed- dings similarity in Fig. 2, we choose 10 unique support sets the dataset with |S|| > For each of the chosen support sets S, we choose 10 distinct contexts j such that S S (a total of 100 samples), and illustrate the Gram matrix of the centered support matrix Š and context embeddings H for this subset of the contexts. Given that the support sets are sparse, choose to contexts with the largest support set size which would result in more pronounced pattern on the correlation matrices. We skip visualizing the word embeddings W in this experiment, the heatmaps did not display any noticeable visual structures.
38

---

Published as a conference paper at COLM 2024
For the metrics displayed in Fig. 11, due to the larger number of context embeddings, only compute the metrics such as (PPr(Lz) - norm growth and correlation with proxies using a sample of 1000 token /context pairs chosen randomly from the training set. We observe qualitatively similar behavior compared to other experiments, with the exception of the norm growth of H, which we conjecture to be consequence of layer normalization in the final embedding layer.
(a) Deep Network
(b) NTP-UFM
(c) Theory
(d) Proxy
CORR(H)
CORR (H)
CORR (DIR(HmI))
CORR(S)
03
5-
8.0
10-
10-
-1.
10
10
CORR(W
CORR W
CORR WWmm
CORR(S
3.0
10
0.0
0.5
25
25
10
20
10
20
10
20
10
Z
10
35
20
05
Figure 10: Same as Fig. 1, this time the Synthetic dataset.
39

---

Published as a conference paper at COLM 2024
(a) Loss convergence
(b) Norm growth
(c) [|px(Le) 1
(d) Correlation with proxies
CE(LA) N
30
12'1
Epach (8)
Epoch (K)
Epach (R)
Epoch (4)
Figure 11: Similar to Fig. 7, this time on a 12-layer TF trained on a subset of 100 stories from the TinyStories dataset.
ineetaatiation similrtty
Potto
260
V.Belodkee dedrits
Sep cn dimteity
W7 issiidads
meices sedety
35-
630
630
-55
-150
Figure 12: Evolution of H₂, W₂ and L to NTP-SVM, solution Hmm Wmm and L throughout train- ing. The embeddings trained by the transformer the Simplified TinyStories dataset. See Sec. 5.
intetaataatinn similrtty.
Fale 20 sirilttt
I1stez 230 simiedtty
Trateo similerty
Aas
Figure 13: Same as Fig. 12, this time for the Synthetic dataset. See Sec. 5.
40

---

Published as a conference paper at COLM 2024
CORR(H)
_oye to tre
wants_ tostt_____ttot start the
Jlly was saa_ewss__p_ttt playynn__utth with per
Jil_and _and _tr__an_oo____aaaeee _ao__ooeeeeeee playing
Wy popyya In
Jily _ikessss to _
Jly liee to hap hor
_yy_ Hiss_io jo_ree__he
fy sr___ssn____eee fimet___e___eee jime_iith with
my is happy betause she
myy_aad fir____________ fir_ htr_heramne_tte she
Ail _nceddss____ ly__eedss_ help because__sste. thes
_ayy_aat__ boc__bccowe_____eeeeeee_________ thes
Ji was playkeg Lawt Juu
Jby and tom, wore glaylii
Joy prie to maaa the
Dr Bas to cork the
Iy thas to till har
Hord in pot her
Jyy yoon to most Jus
191 Jork sin_ pecuuse the
gal ar was so
JBy nests hel, becouse phes Jly was sod because prd
Jby. lezz to slay
Myy spart frow with
JMr in pappy, because
Ji92 wond to
Figure 14: Geometry of context embeddings for Synthetic dataset. A lighter color indicates higher similarity in the embedding space.
CORR (WT)
Jobb
.one
Jughss
sopedencee
uud
Figure 15: Geometry of word embeddings for Synthetic dataset. A lighter color indicates higher similarity in the embedding space.
41

---

Published as a conference paper at COLM 2024
(ORR(H)
Figure 16: Geometry of context embeddings for Simplified TinyStories dataset. lighter color indicates higher similarity in the embedding space.
42

---

Published as a conference paper at COLM 2024
CORR(W) WT)
Figure 17: Geometry of word embeddings for Simplified TinyStories dataset. A lighter color indicates higher similarity in the embedding space.
43

---

Published as a conference paper at COLM 2024
D.3 Other architectures
The results in Thms. and 2 hold provided the model is expressive enough to generate (approximately) unconstrained embeddings and the loss can be minimized close the lower-bound, Depending on the expressiveness of specific network design, the size of the network required for achieving the entropy lower-bound can vary significantly.
To explore this, we hereby repeat our experiments on Simplifiee TinyStories by replacing the TF model with multi-layer perceptron (MLP). We use an MLP consisting of 20 layers organized into four blocks, each containing five layers, with hidden dimension sizes 1024, 512, 256, and 128. This leads to network with around 8 times more parameters than the TF model. The final embeddings geometry is depicted in Fig. 18. We observe that the geometrical patterns appearing in CORR H) and CORR(W at a coarse level are similar to those of the TF model. However, the MLP, even with ~ 10 times larger size, struggles to recover the fine-grained patterns. In terms of loss convergence, we find the MLP converges to the empirical entropy in the order of 10-2 but TF converges in the order of 10-4 on the same dataset. We conjecture that an even larger MLP might be able to achieve better convergence to our theoretical prediction, but we also caution of potential optimization bottlenecks. We encourage additional experiments with other architectures such as LSTMs or state-space models as future work.
CORR S)
CORR H)
CORR ( )
CORR W
Figure 18: Embedding geometries trained by ML.P Left: Geometry comparison between CORR(S) and CORR(H), Right: Geometry comparison between CORR(S") and CORR(W See App. D.3 for details.
D.4 Auto-regressive training
Here, we train the model rute-regressively. In the previous settings, we focused on training samples that were all of fixed length T- 5. In this section, we let the model learn the embeddings for different sequence lengths 1 < T - 1 < 16. Note that the length of the context does not have an impact on our theoretical analysis other than affecting the sparsity pattern S of the training set which in turn influences the implicit geometry as have seen.
We train an 8-layer TF on 200 stories from TinyStories using character-level tokenizer, which limits the vocabulary to approximately 40 characters and promotes higher entropy in next- token distribution. We denote the loss value across different positions T E {2, ... 17] at iteration k by CE(LTx). In Fig. 20, we display the distance of each loss component CE(LT,) from its empirical entropy lower-bound H7, iee, the T-gram entropy of the training set. We observe generally better convergence for shorter sequence lengths T. In Fig. 19, we illustrate the context embeddings similarities. For visualization, we select 5 contexts per each sequence length that end in token "y_ or "_t". The context embeddings and the proxy (P) show similar patterns at coarse level. However, at finer scale, the learned context embeddings that end with the same tokens exhibit strong correlation on average, even when their support sets do not align. We defer further investigation into whether the embeddings' dependence is due to insufficient network capacity or an optimization bottleneck to future work.
D.5 Discussion on large vocabulary setting
Throughout our analysis, we require the embedding dimension to be larger than the vo- cabulary size, i.e., d > V. This condition was necessary to make the non-convex problem in (8) convex. However, if there exists a low-rank optimal solution in NTP-SVM, i.e, rank (m V, the condition d > V can be relaxed in the analysis. In general, the smallest
44

---

Published as a conference paper at COLM 2024
CORR S)
CORR (H)
0.5
0.6
0.2
0.0
Figure 19: Geometry of context embeddings CORR(H) and the heuristic proxy (P), CORR(S), in the autoregressive experiment App. D.4.
Loss convergence across positions
3it Lost
Poution 1m2.177
CE(LTA) HT
Iwd 640
Protion boc
104
Fostban 20, hvo 317
beo
Faalban
hw0
Podttion 24, 0-0-22
Podition h-o
320
330
202
Figure 20: TF trained autoregressively on subset of TinyStories dataset: Loss convergence to its
empirical entropy lower-bound for each sequence length 2, ,77 See App. D.4 for details.
rank among all minimizers depends the sparsity pattern of the language. It is intriguing to investigate this dependence further.
This section examines the network's performance when the decoder dimension d is smaller than the vocabulary size V. To ensure a fair comparison between different setups, we fix the transformer blocks' inner embedding dimension dTF 64 to maintain comparable expressivity for all networks. To vary the final layer embedding dimension d, we add an additional linear layer on top of the network to adjust the context embedding dimension.
In Fig. 21, we show the learned context embeddings for d = 128 and 64, for 10-layer TF trained on the Simplified TinyStories dataset with V 104. We observe that with smaller d, the speed of convergence to the entropy lower-bound H decreases. However, the learned context embeddings exhibit relatively similar geometry with moderate values of d V. We leave more theoretical and empirical exploration of this setting to future works.
45

---

Published as a conference paper at COLM 2024
CORR S)
CORR H), 128
CORR H)), d 64
302
10
2002
3120
(a) cosine similarity heatmaps
(b) Loss convergence
Figure 21: Varying embedding dimension d in the Simplified TinyStories experiment. (a) Left: Proxy (P) for context embeddings. Middle and Right: Geometry of context embeddings trained with TF with d 128 > V and 64 < V, respectively, (b) Loss convergence to the empirical entropy lower bound H. See App. D.5 for details.
46


Reviews:
[{"review_id": "NDR68J8Dml", "reviewer": "Reviewer_6jDj", "summary": "This paper studies the implicit bias of next-token prediction (i.e., language modeling) with the goal of characterizing how the contextual and word embeddings will behave. The theoretical characterization relies on tracing the evolution of the logit matrix L of the model, which is |V| x (number of contexts), over the course of gradient descent. In particular, they formulate language modeling as a semidefinite optimization problem, and then they identify that the logit matrix will converge to something that maximizes the margin between words co-occurring in a given context and words excluded from a given context. \n\nThe logit matrix L is the product of W, the word embeddings, and H, the contextual embeddings, so this result on L can be translated to describe how W and H behave. The conclusion then is that the embedding for a given context forms an acute angle (i.e. has a positive dot product with) the word embeddings in the context and an obtuse angle (i.e. has a negative dot product with) the word embeddings outside of the context. These theoretical results are confirmed through carefully controlled experiments on TinyStories and synthetic data.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thorough review.  \n\n**In-support tokens**: For context j, a token z is in-support if $ p_{j,z}>0 $, i.e. it appears after the context within the dataset. This is important as it relates to a few of your comments. We apologize for the confusion.\n\n**d>V**:  Unlike static models, contextualized models involve two levels of compression: mapping 1) V words, and 2) m contexts (sequences of words) to d-dim space. Despite the requirement d>V, we allow the number of contexts to be much larger, i.e. m>>d. So, from the perspective of contexts, it remains interesting and nontrivial to formalize how the model compresses them into a lower dimensional space. That said, we provide preliminary results for d<V at https://ibb.co/Yb9S5Tr as motivation for future work.\n\n**Fresh perspective**: The classic [LG14], interprets minimizers of word2vec as matrix factorization of PMI matrix. This relies on a subtle assumption: the PMI matrix has well-defined entries, i.e. co-occurrence probability of any word given context is non-zero, typically enforced heuristically by smoothing. Instead, we account explicitly for sparsity which drives our decomposition of L in two orthogonal components: 1) $L_{in}$, a finite component conceptually similar to PMI, 2) $L_{mm}$, a directional component that only depends on the sparsity pattern. \n\n**Correct prediction**: This is subtle but important: Prop1 describes the dominant component $L_{mm}$. But, recall that the logits have a second finite component $L_{in}$ (living on an orthogonal subspace). This is the one that accounts for the conditional probability of in-support tokens ensuring a correct prediction. See Fig5 for numerical verification of convergence to the two components. We will re-emphasize this after Prop1. Thanks!\n\n**Beyond NTP**  Great question. The framework is indeed more general. MLM is an example: it applies directly to simplified setting of masking a single token, but for the practical setting of random masking of several tokens one might need to be more careful to capture all the intricacies. While interesting, we leave this for future work. \n\n**[JT19,T24]**: The basis of our contribution is incorporating the UFM framework within the formalism of [T24]. Technically, training both $W$ and $H$ makes the problem non-convex, but convex relaxation (4) allows us to leverage regularization path ideas from prior works [JT19] back to [RZH03]. We also go beyond [T24] in our empirical evaluations on language data.", "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": "1. Understanding the geometry of the contextual and word embeddings is a well-studied area and is generally useful for interpreting and predicting language model's behavior. \n2. The paper clearly considers many related works and thoroughly discusses how their results fit in with prior ones. I would, however, suggest that the authors include more works on interpretability in language models.\n3. The theory can accommodate arbitrarily complex data and the training of transformers, which are the popular choice for language models these days (there are some caveats, see below).\n4. This approach to analyzing next-token prediction is thought-provoking. I found the future work section to be very interesting, and I think that extending these ideas to accommodate those more realistic settings could be valuable.", "reasons_to_reject": "**Model is unrealistic**\n\nThe theoretical results are in an unrealistic setting where the model's embedding dimension $d$ is larger than the vocabulary size $V$. As I understand it, this allows for the logit matrix $L$ to be uniquely factorized into $WH$, which is essential to the conclusion about the embedding geometry. This does not reflect the standard transformer construction -- in fact, the most recently released language model Gemma exhibits a massive gap between vocabulary size and hidden dimension. \n\nI also want to point out how this assumption weakens the conceptual takeaways from the paper. In most cases, we are interested in studying the model's low-dimensional representation to see how it compresses high-dimensional information -- the assumption that $d\\geq V$ removes the opportunity to do this. For example, the takeaway from Proposition 1 is that the contextual embedding will have the same (positive) dot product with all words contained in the context (Fig 4). In that case, how can the model even predict the next token correctly? It would assign equal probabilities to all tokens it has seen in this context (agnostic to their ordering), and that makes it substantially weaker than the standard next-token prediction setting. This setting thus does not provide any more insight than past works that have modeled contextual embeddings as averages of word embeddings [1, 2, 3] (note: the authors do cite these works in passing in their related works). One can also see this by observing that $L^{mm}$ places uniform weights on all terms in the vocabulary.\n\n[1] Carl Allen and Timothy Hospedales. Analogies explained: Towards understanding word embeddings. In International Conference on Machine Learning, pp. 223–231. PMLR, 2019.\n\n[2] Tatsunori B Hashimoto, David Alvarez-Melis, and Tommi S Jaakkola. Word embeddings as metric recovery in semantic spaces. Transactions of the Association for Computational Linguistics, 4:273–286, 2016.\n\n[3] Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. A latent variable model approach to pmi-based word embeddings. Transactions of the Association for Computational Linguistics, 4:385–399, 2016.\n\n**Writing and formulation are hard to follow**\n\nThe value of building such a complex formalism is predicated on it being comprehensible and easy to extend. The writing in this paper is unfortunately quite difficult to follow. For example, \"in-support tokens\" are never defined -- I have just assumed them to be tokens that occur within a given context. But they could also feasibly be \"all tokens that could follow a given context\". This is just one example of the haphazard writing in the paper that made it difficult for me to follow, even though I have experience reading and working with these theoretical topics.\n\n**Minimal insight can be derived from theoretical results**\n\nI have already mentioned the issue with the embedding dimension and the vocabulary size, but I want to raise other issues with deriving insights from this theory. \n\n1. The geometry of the embeddings can only be described when the support matrix (i.e., the set of contexts used for training) contains all possible sequences of length k that can be constructed from the vocabulary. Realistically, if the vocabulary was a standard (maybe even sub-word) one, most of these sequences would be gibberish, containing no semantic or syntactic coherence.\n\n2. The generative model used for language essentially amounts to treating every language model as a n-gram model. While this has some merit, I think it again goes against the ideas we have about how representation learning works. Relatedly, the max-margin result is pretty standard from the regularization path analysis, but max-margin does not shed much insight when considering _generative_ models instead of the usual classifiers.", "questions_to_authors": "1. Is any part of the analysis unique to next-token prediction? I believe it's possible to formulate masked language modeling in the same way, though one would have to be a little bit careful about what to do with the other masked tokens in a sequence when forming a prediction.\n\n2. Can you elaborate on what the technical challenges are of using a sparse high-dimensional label vector? I am not sure how the original regularization path analysis in Ji & Telgarsky would struggle to accommodate the NTP problem. Similarly, can you clarify how this paper is different from (Thrampoulidis, 2024)? It seems that paper did most of the work in formulating the next-token prediction problem and even discussed the same ideas around the regularization path. From my understanding, the key difference is that that paper had fixed $H$ and only trained $W$, whereas this paper seeks to train both together.\n\n3. Is there any insight into word embedding geometry that has not been shown by prior papers? As I see it, this idea that word embeddings are related to the SVD of the co-occurrence matrix is not too surprising (and your formulation essentially describes the logit matrix as a co-occurrence matrix). I would ideally have liked to have seen something more related to the generative capabilities of really large models (hallucinations, in-context learning, etc), so that one could use this framework of analysis to quantify why next-word prediction can cultivate surprising behaviors.\n\nI am an active reviewer and will take your answers into account to update my score! Also, it is entirely possible that the weaknesses I raised are incorrect -- please correct me if that's the case!", "extra_scores": {}}, {"review_id": "ZXaaHuGTXQ", "reviewer": "Reviewer_4aLd", "summary": "This paper aims to understand the geometric properties of word/context embeddings learned in the process of next-token-prediction in training large language models. To do this, the paper designs a formulation of the next-token-prediction task, that mimics the original model but provides a framework for analyzing the geometry of the embeddings. Through this framework the authors of this work are able to draw specific conclusions about the embedding's geometry. Experiments conducted on two (one synthetic, one real, small-scale) datasets to validate the theories from this paper.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your positive feedback. We would leverage the extra page of the camera-ready to address your concern about the density of information in the current form. We will also include a clearer statement on the motivation/implications, as briefly discussed below.\n\nOur high-level objective is to understand what information from the training corpus gets encoded in the embedding space of contexts and words. In other words, which structures or statistics from the training set influence the similarity and arrangements of embeddings and how this correlates with the semantic similarity of words and contexts. This is influenced by several factors: 1) optimization objective, 2) model architecture, and 3) scale of model and data. We decouple these factors by focusing on the first one: we study the interactions of context embeddings $H$ and word embeddings $W$ in the objective at optimality.\n\nWe find that the statistical properties that are important are encoded in the sparsity patterns of the conditional probability of the next tokens. For example, this is reflected in our experiments by the similarities between the Gram matrices of word/context embeddings and $SS^T, S^TS$. Conceptually, these findings relate prediction-based models that operate on contextualized word embeddings to models based on co-occurrence statistics, much like the classical connection of word2vec, which operated on static word embeddings, to count-based models [LG14]. The latter classical study motivated and resulted in improved models such as GloVe. We anticipate that our work provides analogous opportunities for contextualized word embedding models.\n\nWe also believe that such connections can open avenues to counteract the negative impacts of biases and extreme imbalances in the training set, as text data is known to follow a heavily long-tailed distribution by Zipf's law. Better understanding of the distribution of context/tokens and how they correlate with the model embeddings may help combat unfairness in large language models. We acknowledge that further investigations are needed to fully achieve this goal. However, we consider our framework a first step in making these connections more concrete. Furthermore, we believe these foundations could help develop interpretability metrics for LLMs in order to better analyze their properties.", "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": "How this work benefits the community:\n\n* This paper clearly has some interesting ideas on a methodology for analyzing word embeddings, which would be valuable to the community.\n\n* The paper provides strong theoretical support for claims and hypotheses, that would be very interesting to researchers working on analyzing large language models.\n\n* Based on the observations and learnings from this work, potentially impactful research directions are outlined for the audience.", "reasons_to_reject": "How this work could be further strengthened:\n\n* This paper would greatly benefit from more clarity in its exposition of the work. There is a *lot* packed into this work. Much of it is better understood by referring to details in the appendix. As such, this paper appears to be more like a long journal article -- which has been modified to fit this format. As a result, much of the interesting and important content is split out into the appendix.\n\n* While the focus of the paper has been clearly stated as being interested in uncovering the geometric properties of word embeddings, the motivation/benefit of this is not explicitly spelled out. It is unclear, at the outset, why and how the geometric properties of embeddings could help further research in this field.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aOTHkfhmPT", "reviewer": "Reviewer_HSEc", "summary": "Similar to Levy & Goldberg (2014), who framed the Skip-Gram with Negative Sampling (SGNS) training objective of Word2Vec (Mikolov et al., 2013b;a) as weighted matrix factorization, this work frame next-work prediction (NTP) training as soft-label classification over sparse probabilistic label vectors coupled with an analytical approximation that allows for the unrestricted generation of context embeddings by the language model. This combination establishes a link between NTP training and rank-constrained, nuclear-norm regularized optimization in the logit domain, offering a framework to analyze the embeddings’ geometry.\n\nThis framework allows the examination of how the embeddings’ geometry varies with the distributional properties of the training data, connecting geometric features to textual structures. The authors validate their theoretical insights through experiments on synthetic and\nsmall-scale real language datasets\n\nI think this paper leans a bit towards theory, which is somewhat outside my area of expertise. Therefore, please take my review with a grain of salt.", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive feedback. As a matter of fact, your summary very well captures the message of our work! \n\nRepeating here for completeness: Our high-level goal is to identify how the information from the training corpus gets encoded in prediction-based language models. This is influenced by several factors: 1) optimization objective, 2) model architecture, and 3) scale of model and data. We decouple these factors by focusing on the first one: we study the interactions of context embeddings $H$ and word embeddings $W$ in the objective at optimality. We find that they occur as matrix factorization of a logit matrix that decomposes to a finite PMI-like matrix $L_{in}$ and a directional component $L_{mm}$. This decomposition is driven by the sparsity pattern of the conditional probability distributions of next tokens.  \n\nOverall, this a) gives a fresh perspective to classical matrix factorization results for static word embeddings such as [LG14], b) establishes a connection of the contextualized embeddings to the sparsity pattern of language, c) establishes conditions under which the NTP training loss reaches its lower-bound, d) confirms the log-bilinear model (NTP-UFM) is a good analytical proxy for more complex prediction-based models, e.g. deep transformers. \n\n\nLooking ahead we believe that capturing such influences of how the geometry depends on language statistics such as the conditional probability distributions of next tokens, opens avenues to counteract the negative impacts of biases and extreme imbalances in the training set, as text data is known to follow a heavily long-tailed distribution by Zipf's law. Better understanding of the distribution of context/tokens and how they correlate with the model embeddings may help combat unfairness in large language models. For example, we envision using the geometry of embedding information to reverse engineer the design of more robust loss or to devise posthoc robust sampling methods that account for imbalances in the token/context distributions of language data. We acknowledge that further investigations are needed to fully achieve this goal. However, we consider our framework a first step in making these connections more concrete. Furthermore, we believe these foundations could help develop interpretability metrics for LLMs in order to better analyze their properties.", "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": "I think this work can deepen our understanding of why the next-word prediction objective works so well.", "reasons_to_reject": "This is not really a reason to reject the paper, but could you briefly summarize your work for a researcher who mainly works in empirical NLP rather than in the ML theory domain? Additionally, I am curious about the potential practical implications and applications of this work.", "questions_to_authors": "Please see above.", "extra_scores": {}}, {"review_id": "AV4KfQMkQk", "reviewer": "Reviewer_wYxA", "summary": "The paper sets up the autoregressive nature of language models as a prediction task from complete contexts (a very sparse huge set) to words in the vocabulary (~only the one-hot case of gold data is discussed~ including probabilistic targets) and characterizes that prediction through a series of assumption-driven decompositions into matrices and objectives.\nI found the paper very hard to follow, though that may be due to my lack of background in the area, but as the reasons to reject below show, even so I think there are definitive weaknesses standing in the way of publication at this conference.", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review. Before addressing your concerns: \n1) We are puzzled by your comment “only the one-hot case…”. In fact, we explicitly account for the fact that multiple tokens (labels) can follow a given sequence of tokens (context) in the training corpus which results in a soft-label setup.\n2) Allow us to clarify our contributions: Our high-level goal is to identify how the information from the training corpus gets encoded in language models. This is influenced by several factors: 1) optimization objective, 2) model architecture, and 3) model/data scale. We decouple these factors by focusing on first one: interactions of context/word embeddings $H,W$ in the objective at optimality. We find them as matrix factorization of a logit that decomposes to a finite PMI-like matrix $L_{in}$ and a directional component $L_{mm}$, which is driven by the sparsity pattern of the conditional probability of next tokens. Overall, this a) gives a new perspective to classical static word embedding results such as [LG14], b) connects contextualized embeddings to the sparsity pattern of language, c) establishes conditions for NTP loss reaching lower-bound, d) confirms the log-bilinear NTP-UFM as a good analytical proxy for complex models, e.g. deep transformers. \n\n**d>V**:  Unlike static models, contextualized models involve two levels of compression: mapping 1) V words, and 2) m contexts (sequences of words) to d-dim space. Despite the requirement d>V, we allow for the number of contexts to be much larger, i.e. m>>d. So, from the perspective of contexts, it remains nontrivial to formalize how the model compresses them into a lower dimensional space. That said, we provide preliminary results for d<V at https://ibb.co/Yb9S5Tr as motivation for future work.\n\n**Practical insights**: We believe that our perspective opens avenues to counteract the negative impacts of biases and extreme imbalances specific to language. Better understanding of the context/tokens distributions and how they correlate with the model embeddings may help combat unfairness in large language models. For example, we envision using the geometry of embeddings to reverse engineer the design of more robust loss or to devise posthoc robust sampling methods that account for imbalances in the token/context distributions of language data. Our work is the first step in making these connections more concrete but further investigations are needed to fully achieve these practical goals.", "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": "The clear stating of assumptions and rigorous treatment are appreciated, as is the extensive appendix.", "reasons_to_reject": "- The focus on d > V - 1 makes the entire analysis not fit the standard language modeling setup where d <<< V - 1 (in fact the whole motivation for neural modeling)\n- Notation used is not just explained but even just introduced in the appendix\n- Relevance to practical language modeling and potential for useful insights derived from the described characterization are not given", "questions_to_authors": "- please use parentheses around citations where appropriate (\\citep{})\n- typo: empiriical -> empirical", "extra_scores": {}}, {"review_id": "bfJ5bhqD96", "reviewer": "Reviewer_AKL7", "summary": "The authors provide an analysis of the geometry of word and context embeddings in language models learned with next-token prediction.  Specifically they model the learning problem as a soft-label classification of the next token, and assume that context embeddings can be learned in an unrestricted way (i.e. not constrained to the parametric form imposed by the model).  This leads to an interpretation of the problem as a nuclear-norm regularized optimization, similar to support vector machines, in which the geometry (similarity matrices) of the learned word and context embeddings align with that of the support sets of the corresponding contexts.  This is a nice and intuitive result, and it is also demonstrated empirically (qualitatively) on a synthetic toy dataset and on the Tiny Stories dataset.\n\nQuality:  This is generally high-quality, well-motivated work.  The empirical work is lower-quality, in the sense that the results are qualitative and limited to the restrictive setting where the model dimensionality is higher than the vocabulary size.\n\nClarity:  The paper is generally very clear.\n\nOriginality:  As far as I know, this is original work.\n\nSignificance:  The significance is limited by the restrictive assumptions and limited empirical work in the paper.  Like many papers on the science of LMs, it is not clear what the implications of the work on how practitioners train or use LMs might be.  Although an analysis with more obvious implications would have higher significance, I think this purely scientific analysis is significant enough.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your overall positive feedback. In the final version of the paper, we will make sure to address the clarity issues and typos pointed out. Thank you!\n\n**d>V**:  Unlike static models, contextualized models involve two levels of compression: mapping 1) V words, and 2) m contexts (sequences of words) to d-dim space. Despite the requirement d>V, we allow the number of contexts to be much larger, i.e. m>>d. So, from the perspective of contexts, it remains interesting and nontrivial to formalize how the model compresses them into a lower dimensional space. Besides, it is not a priori clear to us, even under the assumption of d>V, how the geometry depends on the language statistics as encoded in the sparsity pattern.  That said, we provide preliminary results for d<V at https://ibb.co/Yb9S5Tr as motivation for future work. \n\nWe also remark that the rank-constrained logit relaxation (see Lem1) holds for any d. For example, this directly yields a necessary and sufficient condition for the loss to reach the entropy lower bound, which is the feasibility of the rank-constrained NTP-SVM problem. We also hypothesize that the sparsity pattern of language promotes low-rank solutions to the NTP-SVM problem which could allow extensions of the parameter convergence results to the setting d<V.\n\n**Visual inspections**: Thank you for commenting on that. Following your suggestion, we now verify the correlation quantitatively as well by computing the normalized covariance of the heatmaps. The plot of this metric vs epoch is available at https://ibb.co/Yb9S5Tr . We hope this addresses your concern and thanks again for the suggestion.", "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": "+ The idea is well-motivated and the paper is generally well-written.\n+ The results, while preliminary, suggest that the analysis approach is promising.", "reasons_to_reject": "- The findings are somewhat anecdotal, based on toy settings and qualitative results\n- The assumption d > V (the hidden dimensionality is larger than the vocabulary) is very unrealistic, and it's not clear how it might be relaxed in future work (the authors don't address this)", "questions_to_authors": "- There seems to be a typo in the definition of $f_{\\theta'}$, as its output is not in $\\mathcal{V}$ as stated (Section 2, first paragraph).\n\n- In the definition of CE(\\theta'), what is $n$?  I believe there is a missing subscript $i$ on $x$, and $S_{z_i}$ seems to be undefined (Section 2, first paragraph).\n\n- Spell out \"ETF\" (Section 3.4.1).\n\n- I believe the reference to Harris 1954 should be Firth, J. R., Studies in Linguistic Analysis, 1957. (Section 4.1)\n\n- I believe the finding that \"Contexts with a larger intersection in the suppost set ... exhibit larger correlation\" (Section 4.2) is a qualitative statement based on visual inspection of Fig. 3. Is that correct?  If so, then that should be stated explicitly.  The mention of \"larger correlation\" suggests that a correlation was actually measured.\n\n- There are a couple of missing citations:    \n  * \"the linear relationships underlying word analogies\" (Section 5)\n  * \"recent promising studies ... that leverage the neural-collapse geometry\" (Section 6)\n\n- Typos, grammar, etc.:  \n  * \"their synonymous\" --> \"their synonyms\".  How did you derive the synonyms?  (Section 4.2)\n  * \"Different to them\" --> \"Unlike them\" (Section 5)", "extra_scores": {}}]

Author rebuttal:

Meta-review: I find the paper to be very interesting. While I mostly agree with Reviewer 6jDj on the weaknesses of the paper, I think we can expect the paper's result being useful for future work. For instance, while now the assumption d > V seems unrealistic, the size of language models (and thus d) is increasing whereas V is unlikely to increase by a lot (in monolingual setting). So when the LM becomes large enough, this assumption will hold true. Given that the theoretical part of deep learning and LMs is in a very early stage, I personally think we should approach the field in a bottom-up manner just like other science areas in their early days; that is, we first enumerate findings even if they may not be immediately useful, hoping that some future work will be able to connect them. Also, given the nature of CoLM, which is a more experimental conference than existing venues in NLP and machine learning, I think there will be sufficient audience who will be interested in this paper.
Decision: Accept