Implicit Geometry of Next-token Prediction: From Language Sparsity Patterns to Model Representations

Authors: yize  zhao
tina  behnia
vala  vakilian
christos  thrampoulidis
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-23 02:13:04

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.

Main paper: 
 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 z <t ∈ V t-1 of t -1 preceding tokens, the cross-entropy (CE) loss between the model's predicted conditional probability distribution over potential next tokens from a vocabulary V and the one-hot encoded actual next token z t ∈ V. The model's conditional distribution is defined through a softmax map applied to logits ℓ <t (W, θ) = W h θ (z <t ), which are generated by mapping context embeddings h θ (z <t ) ∈ R d -the neural network's ddimensional representations of contexts z <t -using a matrix W ∈ R |V |×d 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 et 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 of NTP-so named because the NTP objective does not explicitly impose any such relationship-explores 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 geometry is key to understanding functional principles of large language models, since NTP is used for training across diverse architectures, from LSTMs to 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 state-of-the-art 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 on 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 NTP loss to its entropy lower-bound. This approach isolates the influence of NTP itself-rather than architectural nuances-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 as 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) .Our methodology integrates three foundational modeling concepts as follows:Firstly, we frame NTP as soft-label classification with CE loss applied to sparse probabilistic label vectors (Thrampoulidis, 2024) . This isolates m distinct contexts, which could be repeated multiple times throughout the training corpus, and assigns to each a sparse V = |V |-dimensional conditional-probability label vector pj , reflecting the frequency of each token following context j ∈ [m]. The sparsity of pj 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 a tractable analysis of context embeddings by assuming expressive (enough) neural networks can produce unconstrained embeddings h j ∈ R d , independent of the architecture's specific complexities (Yang et 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 ) = [ p1 , p2 , . . . , pm ] ∈ [0, 1] V×m , leading to the following log-bilinear model: min W ∈R V×d , H∈R d×m L NTP W H; P .(1)This way, our goal to study word-word, context-context, and word-context geometric relationships becomes that of characterizing the Gram matrices G W = WW ⊤ and G H = H ⊤ H, as well as the logit matrix L = W H 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 to zero. For the purpose of comparing the analysis outcomes to our numerical evaluations, we interpret the prediction obtained from the regularization path analysis as a proxy for the solution found by gradient-based optimization when the NTP loss approaches its empirical entropy lower bound.Building on the above methodology, our analytical framework leads to the following results:Formulation in logit space. Sec. 3 presents an equivalent formulation of the NTP objective L NTP W H; P + λ∥W ∥ 2 + λ∥H∥ 2 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 ≥ V1 , in Sec. 4, we demonstrate that the logit matrix L λ for regularization λ → 0 (which is proxy for training iteration k → ∞), behaves for some R(λ) → ∞ as shown by the following claim:(C1) Logits' sparse plus low-rank decomposition: As λ → 0,EQUATION)where the orthogonal components L in and L mm have distinct operational roles.L in encodes information regarding frequencies of in-support tokens, inheriting the sparsity of the data matrix P. In contrast, L mm depends only on the sparsity pattern S ∈ {0, 1} V×m 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 is to maximize the logit margin between in-and off-support tokens, 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 L mm .TF model. Panel (c) compares the TF output to our analytical predictions: Comparing TF's context/word embeddings and logits to H mm /W mm and L mm 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 a semi-definite program for L mm 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 S := (I V -1 V 1 V 1 ⊤ V )S be the (column-wise) centered support-set matrix. The directional component of word and context embedding matrices can be well-approximated by S as follows:EQUATIONPanel (d) of Fig. 1 suggests proxy (P) 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 a subset of the TinyStories, where computing L mm 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 exhibit low correlation. This finding is consistent with classical intuitions of the "distributional assumption" of words/contexts and is formalized by our framework via the implicit bias of NTP to promote low-rank logits subject to margin conditions for in/off-support tokens. 
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 a log-biliniear model, while the log-bilinear model in ( 1) is only used by us as 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 it leads to diverging embeddings, a scenario where setting the loss gradient to zero, as 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 L in + RL mm , where L in captures frequencies akin to 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 to 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 is inspired by the implicit bias/regularization research on the preferred solutions of optimizers like GD in overparameterized 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.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. 
formulation: 
 Notations. Throughout, lowercase and uppercase bold letters (e.g., a and A) represent vectors and matrices, respectively. We use a and A to denote vectors/matrices normalized by their Euclidean norm. We denote A[i, j] the (i, j)-th entry of matrix A and a j its jth column. We let R(A) and N (A) its range-space (aka column space) and null-space, respectively. ⟨•, •⟩ and ∥•∥ denote Euclidean inner product and norm, respectively. We use ∥ • ∥ * to denote the nuclear-norm (i.e. sum of singular values). I V represents the identity matrix of size V and 1 V , the all ones vector of size V × 1 (subscripts are removed when clear from context). ∆ V-1 denotes the V-dimensional unit simplex and S(•) :R V → ∆ V-1 the softmax map: S(a) = [S 1 (a), . . . , S V (a)] ⊤ with S v (a) = exp(e ⊤ v a) ∑ v ′ ∈[V] exp(e ⊤ v ′ a), where e v is the v-th standard basis vector in R V . We also denote e j the j-th standard basis vector in R m . All logarithms are natural logarithms (base e).We let V = [V] := {1, . . . , V} represent a finite vocabulary of tokens (we use the terms 'word' and 'token' interchangeably). We denote by z 1:t = (z 1 , . . . , z t ) a sequence of t tokens z t ∈ V and focus, for simplicity, on prediction of the last T-th token z := z T given context x := z 1:T-1 . For this, we assume access to a training set consisting of n sequences T n := {(x i , z i )} i∈ [n] , such that x i ∈ X := V T-1 and z i ∈ V. This is used to train model f θ ′ : X → V, f θ ′ (x) = W h θ (x) parameterized by θ ′ = {W, θ}, where W ∈ R V×d is a decoding matrix and θ parameterizes a map h θ : X → R d 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 a TF. We refer to row w v , v ∈ V of W as word embedding of token v and h θ (x) as context embedding of context x. The model is found by minimizing the empirical CE loss CE(θ′ ) = 1 n ∑ i∈[n] -log (S z i ( f θ ′ (x))). Following Thrampoulidis (2024) , we reframe the NTP training objective as classification over m ≤ n distinct contexts, each associated with a sparse probabilistic label vector pj ∈ ∆ V-1 . Concretely, we denote x1 , . . . , xm the m ≤ n distinct contexts among the (large number of) total n contexts T n . Also, we let πj= 1 n ∑ i∈[n] 1[x i = xj ]denote the empirical probability of distinct context xj . Accordingly, let pj ∈ ∆ V-1 denote the probability vector of conditional next-token distribution, i.e., for all z ∈ V: pj,z := 1 n ∑ i∈[n]:x i = xj 1[z i = z], j ∈ [m].In words, n • πj • pj,z is the number of occurrences of token z as a follow-up to context xj . Define the support sets of these probability vectors as S j := {z ∈ V | pj,z > 0} and let S j := |S j |.It is also convenient to define the probability matrix P = [ p1 , . . . , pm ] and its corresponding support matrix S ∈ {0, 1} V×m , such that S[z, j] = 1, iff z ∈ S j . In typical scenarios, S j < V; thus, P and S are sparse matrices. For given context index j ∈ [m], we say word z is in-support token if S[z, j] = 1 (eqv. z ∈ S j ); otherwise, we say z is off-support token.With the above notation, the training loss becomes (Thrampoulidis, 2024) ,CE(θ ′ ) = -∑ j∈[m] πj ∑ z∈V pj,z log S z (W h θ ( xj )) . (7)This is lower bounded by the empirical T-gram entropy (referred to hereafter as entropy) of the data (Shannon, 1948) , i.e., for all θ ′ : CE(θ ′ ) ≥ H := -∑ j∈[m] πj ∑ z∈V pj,z log pj,z .To gain insights into the geometry of solutions to CE minimization in Eq. ( 7), we assume sufficient model expressivity, allowing us to optimize embeddings freely, instead of abiding by their architecture-specific parameterization. We formalize this concept below.Definition 1 (NTP-UFM). The unconstrained features model (UFM) for NTP training over training set T m = { πj , pj } j∈[m] refers to the following log-bilinear optimization problem:min W,H CE(W H) + λ 2 ∥W ∥ 2 + λ 2 ∥H∥ 2 , (NTP-UFM)where minimization is over word and context embedding matrices W ∈ R V×d andH := [h 1 , . . . , h m ] ∈ R d×m , and CE(W H) := -∑ j∈[m] πj ∑ z∈S j pj,z log S z (W h j ) .The key difference of NTP-UFM compared to Eq. ( 7) is that embeddings h θ ( xj ) are now optimized unconstrainedly through free variables h j . 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 a proxy to examine the behavior of gradient descent.2 Specifically, we use the vanishingregularization solution of NTP-UFM as a proxy for the parameters (W k , H k ) learned in large iteration k → ∞ through GD on the unregularized objective (7).We now introduce a reformulation of NTP-UFM utilizing the logit matrix L = WH. This is grounded on the well-known fact concerning the nuclear norm of a matrix (Srebro et al., 2004; Fazel, 2002) :∥L∥ * = min L=W H 1 2 ∥W ∥ 2 + 1 2 ∥H∥ 2 . Lemma 1. Denote logit matrix L = [ℓ 1 , . . . , ℓ m ] ∈ R V×m and let L λ be a minimizer of min L : rank(L)≤d -∑ j∈[m] πj ∑ z∈S j pj,z log S z (ℓ j ) + λ∥L∥ * , (8)with SVD L λ = UΣV ⊤ , where U ∈ R V×r , Σ ∈ R r×r , V ∈ R m×r and r = rank (L λ ) ≤ d. Then:(i) The optimal cost of NTP-UFM is the same as that of Eq. (8).(ii) (W λ , H λ ) minimizes NTP-UFM if and only if there exists a minimizer L λ of Eq. (8) such that:W λ H ⊤ λ = L λ , W λ W ⊤ λ = UΣU ⊤ , and H ⊤ λ H λ = V ΣV ⊤ .The constraint rank (L) ≤ d in Eq. ( 8) ensures the logit matrix can be factorized as W H with inner factor dimension being d. Thus, in general, Eq. ( 8) is non-convex. 
analysis of unconstrained features model for ntp training: 
 This section analyzes the regularization path of Eq. ( 8) and, subsequently, of NTP-UFM. Throughout, we assume d ≥ V, in which case the constraint rank (L) ≤ d becomes redundant. All proofs are deferred to the appendix and numerical evaluations to App. D.1.The following nuclear norm-minimization problem plays a key role in our results.Definition 2. Define the NTP-SVM logit matrix L mm as solution to the following optimization:EQUATIONEQUATION)The optimization in NTP-SVM ⋆ is a semidefinite program (SDP) and always feasible, as shown by the straightforward feasibility of the centered support matrix S = (IV -1 V 1 V 1 ⊤ V )S.Its solution L mm depends solely on the support matrix S of the data. Interestingly, we prove in App. C that S solves NTP-SVM ⋆ under symmetry assumptions on S. More generally, our empirical results show that S closely approximates the true solution L mm , supporting the use of Proxy (P); see App. C.3.Finally, when S j = 1, ∀j ∈ [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.We now characterize the solutions of (8) when regularization vanishes and d ≥ V.Theorem 1. Suppose d ≥ V. In the limit of vanishing regularization, the solution L λ of optimization (8) diverges in norm and converges in direction to L mm . Formally, lim λ→0 ∥L λ ∥ = +∞, and, there exists minimizer L mm of NTP-SVM ⋆ such that lim λ→0L λ ∥L λ ∥ * -L mm ∥L 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 pj,z influence the learned logit matrix and what implications this has for the loss? The theorem below complements Thm. 1 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, j) and (z ′ , j), which equal 1 and -1 respectively for all context indices j ∈ [m] and in-support word indices z, z ′ ∈ S j :F = F(S) = span (e z -e z ′ ) e ⊤ j : z ̸ = z ′ ∈ S j , j ∈ [m] ⊂ R V×m .Note for any L, the projection P F (L) onto F is sparse with support matching S.Theorem 2. Under the setting of Theorem 1, the loss approaches its lower bound lim λ→0 CE(L λ ) = H. Additionally, lim λ→0 P F (L λ ) = L in , where L in ∈ F is the unique L ∈ F that guarantees differences of logits equal log-odds, i.e., L[z, j] -L[z ′ , j] = log( pj,z / pj,z ′ ), ∀z ̸ = z ′ ∈ S j , j ∈ [m] .The loss reaches the lower-bound H because on the subspace F, 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. 1 because L mm belongs to the orthogonal complement F ⊥ (which follows from Eq. ( 9); see App. B.3.1).Combining the two theorems, L λ converges to L in in F and diverges in F ⊥ , where it converges directionally to L mm . This justifies Claim (C1).In view of Lemma 1 we 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 λ , H λ ) any minimizer of NTP-UFM. Consider the SVD L mm = UΣV ⊤ . Then, using • notation to denote normalized quantities, in the limit of vanishing regularizationW λ W ⊤ λ → UΣU ⊤ =: G mm W , H ⊤ λ H λ → V ΣV ⊤ =: G mm H and W λ H λ → UΣV ⊤ .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. 
experiments: 
 We empirically validate our analysis on text data. We start with two small-scale synthetic datasets to examine Claims (C1)-(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 -1 = 5. The dataset is constructed so that all support-sets satisfy S j = 3.Simplified TinyStories. For a more realistic but still controlled setup, we curate the dataset from the TinyStories corpus: We derive contexts xj and support sets S j by choosing the most frequent word-level contexts with length T -1 = 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. Metrics. We define the following metrics for verifying our theoretical results:Visualization of embeddings' geometry. To visualize the geometry of the word and context embeddings, we plot the heatmap of the normalized Gram matrix CORR(W ⊤ ) and CORR(H), where for matrix X, CORR(X ) is a matrix whose (i, j)-th entry is the cosine similarity between the respective columns of X, i.e., CORR(X ) i,j = x ⊤ i x j /(∥x i ∥∥x j ∥). To check Proxy (P), we compare the context and word embeddings with CORR( S) and CORR( S ⊤ ).Quantifying geometric similarity. To measure the structural similarity between W, H and S ⊤ , S, we use SIM X, Y = (σ XY + ϵ)/(σ X σ Y + ϵ), as the correlation matrix between two matrices. 3 Here, σ XY is the covariance, and σ X and σ Y are the standard deviations of X and Y, respectively, and ϵ is a 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 H, S := SIM CORR(H), CORR(S) , SIM W ⊤ , S ⊤ := SIM CORR(W ⊤ ), CORR(S ⊤ ) .Recovery of soft-labels. To verify Thm. 2, we measure the distance ∥P F (L k ) -L in ∥ of L in from the projection of L k onto the subspace F . Next, we discuss Fig. 3 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 (H mm , W mm ). Note that all similarity 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.EQUATIONFinally, 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/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. 
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-path analysis to the non-convex realm of problem ( 8), potentially revealing a 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 constraints 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 experiments, 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 over-parameterized enough to reach the entropy lower bound and ensure expressiveness in the embedding space. For a brief discussion, see App. D.3. Third, there is a 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 mm . 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 influences the learning of linguistic patterns and regularities.This section expands upon the section on related work in the main body.Neural probabilistic language models are designed to predict the probability distribution of a target word v ∈ V within a given vocabulary V, based on a sequence of preceding or surrounding words of (variable or fixed) length t known as the context z ∈ V t . 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 Word2Vec 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 alignment between vectors of target words and the context words surrounding the target, while disassociating those of randomly drawn pairs. This is achieved by optimizing a 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 a formal examination of embeddings generated by the SGNS objective, framing it as a 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 a measure of association between words and contexts by having entries PMI[z, j] = log Pr(z|j)Pr(z) = log pj,z Pr(z) for target 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[z, j] = -∞). To address the possibility of negative infinity values, an approximation involving a positive, sparse, and non-negative PMI (PPMI) 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 NTP as a 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-bilinear model, is in spirit to the log-bilinear models trained by Word2Vec, we only use such log-bilinear models as an analytically-tractable proxy to expressive deep neural probabilistic language models such as transformers. (iii) Thirdly, we confront the sparsity in the probabilistic labels of each context head-on, unlike other matrix-factorization 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 to diverging embeddings, a 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-we unveil that word and context embeddings emerge from the factorization of a matrix L in + RL mm . Here, L in captures the frequency of word-context occurrences akin to the PMI, setting nonoccurring 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 context-word probability matrix.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 anisotropy of the learned word/context 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 on geometric analysis and our findings could similarly motivate further research in this domain, extending the understanding and application of language model embeddings.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 hardmargin 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 et al., 2018a; b) ) including recent studies on transformers (Tarzanagh et al., 2023b; a) ; for a 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-dimensional contexts where traditional margin-based bounds are inadequate, revealing scenarios where no regularization yields optimal outcomes, a phenomenon known as 'harmless interpolation' or 'benign overfitting.' Most of these findings, whether on optimization or generalization, are confined to the one-hot classification framework.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 a 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.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 equiangular 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 overparameterized 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 a 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 classification-specifically, when all relevant tokens are 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 to 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.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 as 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 a 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.We apply our geometry analysis to Word2Vec. Consider the following Skip-gram objective:1 T ∑ t∈[T] ∑ -c≤j≤c,j̸ =0 log p z t+j |z t = 1 T ∑ t∈[T] ∑ -c≤j≤c,j̸ =0 log p S z t+j (Hw z t ) ,where H ∈ R V×d denotes the embedding matrix of context words z t+j , j ̸ = 0 and W ∈ R V×d denotes the embedding matrix of target words z t . As presented in Mikolov et al. (2013c) , skipgram actually implements a heuristic of this objective that aims to simplify the computation of the normalizing factor in the softmax. Second, consider the CBOW objective:1T ∑ t∈[T] log p z t |{z t+j } {|j|∈[c],j̸ =0} = 1 T ∑ t∈[T] log p S z t (W • 1 2c -1 ∑ -c≤j≤c,j̸ =0 h z t+j ) , where 1 2c-1 ∑ -c≤j≤c,j̸ =0 h z t+j is the continuous-bags-of-word embedding of the context {z t+j } {|j|∈[c],j̸ =0} .Suppose the CBOW model with context consisting of only the previous word. To keep notation consistent with the rest of the paper, let j ∈ [m] = [V] denote index of distinct contexts (aka previous words) and z ∈ [V] denote index of (next) words. Note the embedding matrix H ∈ R d×V is now of same dimension as W ⊤ . Let πj denote the marginal of context words and pj,z denote the conditional probability of word z following context word j. Then, the CBOW objective becomes equivalent to∑ j∈[V] πj ∑ z∈[V] pj,z log S z W h j .This now fits exactly our framework and the results for NTP apply mutatis mutandis.In Mikolov et al. (2013c) , instead of minimizing the CE objective, they actually use the Negative-Sampling (SGNS) objective. Levy & Goldberg (2014) rewrites the SGNS objective as follows:∑ j∈[m] ∑ z∈[V] -Pr(j, z) log σ(w ⊤ z h j ) -k Pr(j) Pr(w) log 1 -σ(w ⊤ z h j ) ,where Pr(j, z) := πj pj,z , Pr(j) := πj , Pr(z) = ∑ j∈[m] πj pj,z and k ∈ N is a hyperparameter. Defining q j,z := Pr(j,z) Pr(j) Pr(z) k + Pr(j,z) Pr(j) Pr(z)∈ [0, 1] ,we can rewrite the SGNS objective as∑ j∈[m] ∑ z∈[V] Pr(j) Pr(w) k + Pr(j, z) Pr(j) Pr(z) -q j,z log σ(w ⊤ z h j ) -(1 -q j,z ) log 1 -σ(w ⊤ z h j ) .For each (j, z)-pair, the rightmost term in the expression above is lower bounded by H(q j,z ) := -q j,z log(q j,z ) -(1q j,z ) log(1q j,z ).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 matrixL 1 ∈ R V×m (of rank ≤ d) such that L 1 [z, j] = PMI(z, j) -log(k) = log Pr(j, z) Pr(j) Pr(z) -log(k), ∀j ∈ [m], z ∈ S j . . B. There exists logit matrix L 2 ∈ R V×m (of rank ≤ d) such that L 2 [z, j] = 0, ∀j ∈ [m], z ∈ S j , L 2 [v, j] < 0, ∀j ∈ [m], v / ∈ S j .The logit matrix L 1 is the shifted PMI matrix as also characterized in Levy & Goldberg (2014, Eq. ( 7)). This component plays a similar role to the finite component L in in the NTP setup of our paper. The second component, which is only non-zero if there exists a target-context pair (z, j) with zero co-occurrence in the training set (Pr(j, z) = 0), specifies the directional component, akin to the role of L mm in NTP. This component is absent in the analysis of Levy & Goldberg (2014) , as they implicitly assume no zero entries in the co-occurence 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, L 2 identifies the infinite direction that drives the loss to its lower-bound. Levy & Goldberg (2014) address sparsity indirectly by discussing alternative matrix factorizations such as shifted positive-PMI (PPMI).Lemma 2. For any λ > 0, any minimizer W of NTP-UFM satisfies 1 ⊤ V W = 0. That is, word embeddings are centered.Proof. The proof is the same as Thrampoulidis et al. (2022, Lem. E.1) and is thus omitted.There is two consequences of Lemma 2:A. The word embedding matrix W is centered motivating the use of the centered support matrix S (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 d ≥ V -1. Finally, note this also implies the constraint 1 ⊤ L = 0 can be added without any change in Eq. ( 8).We prove a slightly stronger version of Lemma 1 which assumes the additional constraint 1 ⊤ 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 ∥L∥ * = minW H=L 1 2 ∥W ∥ 2 + 1 2 ∥H∥ 2 .Firstly, fix any L such that r := rank (L) ≤ d and1 ⊤ V L = 0. Let L = UΣV ⊤ be its SVD with Σ ∈ R r×r . Choose Ŵ = UΣ 1/2 R and Ĥ = R ⊤ Σ 1/2 V ⊤ for some partial orthonormal matrix R ∈ R r×d (i.e., R ⊤ R = I d ). It then holds clearly that Ŵ Ĥ = L and ∥ Ŵ ∥ 2 + ∥ Ĥ∥ 2 = tr UΣU ⊤ + tr V ΣV ⊤ = 2 tr (Σ) = 2∥L∥ * , Thus, CE(L) + λ∥L∥ * = CE( Ŵ Ĥ) + λ 2 ∥ Ŵ ∥ 2 + ∥ Ĥ∥ 2 .Because we can apply this for any L, it shows that4 minEQUATION)where we define R d := {L ∈ R V×m : rank (L) ≤ d} to be the manifold of rank-constrained matrices.Secondly, consider any W ∈ R V×d , H ∈ R d×m such that 1 ⊤ V W = 0. Let L = W H and denote L = UΣV ⊤ its SVD. We then have 1 ⊤ L = 0 andEQUATION)where in the first line we used ∥L∥ * = tr (Σ) = tr U ⊤ LV , in the second line we used Cauchy-Schwartz, and, in the third line we used that ∥U∥ 2 ≤ 1, ∥V ∥ 2 ≤ 1. Thus, this proves minEQUATION)where the first equality follows by Lemma 2. Moreover, note that equalities in (13) hold if and only if the following three conditions holdW H = L = UΣV ⊤ =⇒ U ⊤ W HV = Σ U ⊤ W = V ⊤ H ⊤ ∥U ⊤ W ∥ 2 = ∥W ∥ 2 , ∥V ⊤ H ⊤ ∥ 2 = ∥H∥ 2 .The first two conditions combined giveU ⊤ WW ⊤ U = V ⊤ H ⊤ HV = Σ .Thus, WW ⊤ = UΣU ⊤ + U ⊥ Σ ⊥ U ⊤ ⊥ for some non-negative diagonal matrix Σ ⊥ and U ⊥ the orthogonal complement of U. But, the third condition requirestr U ⊤ WW ⊤ U = tr WW ⊤ =⇒ tr (Σ) = tr (Σ) + tr (Σ ⊥ ) =⇒ tr (Σ ⊥ ) = 0 =⇒ Σ ⊥ = 0.EQUATION)and similarlyEQUATION)Combining Eqns. ( 12) and ( 14) proves that minW,H CE(W H) + λ 2 ∥W ∥ 2 + ∥H∥ 2 = min L∈R d 1 ⊤ L=0 CE(L) + λ∥L∥ * .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 λ H ⊤ λ as stated in the lemma.The proof involves several components, and so, we proceed in steps.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 ̸ = z ′ ∈ S j is equivalent to the set of the same constraints for any anchor z j ∈ S j and z ′ ̸ = z j ∈ S j . That is, there is an effective total of S j -1 linearly independent constraints for each j ∈ [m].Definition 3 (E in,j ). For each j ∈ [m] fix an anchor index z j ∈ S j and define matrix E in,j ∈ R (S j -1)×V with S j -1 independent rows (e z je z ) ⊤ , z ̸ = z j , z ∈ S j .With this definition, we can rewrite Eq. ( 9) sinceEQUATION)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 z j ∈ S j and reformulate Eq. ( 10) asE out,j ℓ j ≥ 1, j ∈ [m].Here, E out,j ∈ R (V-S j )×V has V -S j independent rows (e z je v ) ⊤ , v / ∈ S j . In total, there are V -1 constraints per context in NTP-SVM ⋆ .With the above definition of the matrices E in,j ∈ R (S j -1)×V , j ∈ [m], we can define the subspace F of V × m matrices that have their j-th column living on the range space R(E ⊤ in,j ) for all j ∈ [m]:F := L ∈ R V×m : ℓ j ∈ R(E ⊤ in,j ) , j ∈ [m] = span (e z -e z ′ ) ⊤ e j : z ̸ = z ′ ∈ S j , j ∈ [m] .Note that this set depends on the support matrix S, but we suppress the dependence to simplify notation. Let P F : R V×m → R V×m be the projection operator on this subspace. For concreteness, note that this can be explicitly defined as the matrix operator projecting each column ℓ j of L onto R(E ⊤ in,j ), i.e., P F (L) = [P F (L)1 P F (L) 2 • • • P F (L) m ] with P F (L) j := E ⊤ in,j (E in,j E ⊤ in,j ) -1 E in,j ℓ j , j ∈ [m].Note that P F is a projection operator since it is clearly linear and P F (P F (L)) = P F (L). Also, it is clear that P F (L) ∈ F for all L ∈ R V×m . Operationally, note that P F (L) has zero-entries on off-support elements. We will use this property indiscriminately throughout this section.Finally, define the orthogonal complement to F asEQUATION)Clearly, R V×m = F ⊕ F ⊥ . Operationally, note that L ∈ F ⊥ implies that on-support elements of every column are identical, that is, L[z, j] = L[z ′ , j] for all z ∈ S j and j ∈ [m]. In other words, it holds L ∈ F ⊥ ⇐⇒ L satisfies (9), (18) which we will use without further explicit reference.We also let P ⊥ denote the projection operator to F ⊥ .For word/context embeddings W, H and logits L = W H, recall the NTP loss is written asEQUATIONand is lower bounded by the empirical T-gram entropy of the data (Shannon, 1948), i.e., for all θ ′ : CE(W H) ≥ H := Ê(x,z)∼T n [log ( p(z|x))] .We now state explicit conditions on the logit matrix L for which the lower bound is attained by a softmax model L = S(W H) 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 a slight extension of Thrampoulidis (2024, Prop. 1) .Proposition 1. The NTP objective of a softmax model S(W H) with W ∈ R V×d , H ∈ R d×m can reach the empirical entropy lower-bound if the following three conditions hold simultaneouslyEQUATION)There exists logit matrix L 2 ∈ R V×m that is feasible in NTP-SVM ⋆ .C. [Rank constraint] For all arbitrarily large positive scalars ρ,rank (L 1 + ρ • L 2 ) ≤ d.Concretely, if the above conditions hold, then for all arbitrarily large ρ, there exist W ρ ∈ R V×d andH ρ ∈ R d×m such that lim ρ→∞ CE(W ρ H ρ ) = H .Proof. It suffices to prove that forEQUATION)To see that this is sufficient, note from condition C of the theorem that for all arbitrarily large ρ, L ρ admits a factorization L ρ = W ρ H ρ where W ρ ∈ R V×d and H ρ ∈ R d×m . Since this holds for all large ρ, Eq. ( 21) implies the desired.Thus, we now prove Eq. ( 21). For any j ∈ [m], it is easy to check, using conditions A and B of the theorem, that the following two statements are true for all ρ > 0:A. For all z ̸ = z ′ ∈ S j :L ρ [z, j] -L ρ [z ′ , j] = (L 1 [z, j] -L 1 [z, j]) + ρ • L 2 [z, j] -L 2 [z ′ , j] = log( pj,z / pj,z ′ ) B. For all z ∈ S j , v ̸ ∈ S j : L ρ [z, j] -L ρ [v, j] = (L 1 [z, j] -L 1 [v, j]) + ρ • (L 2 [z, j] -L 2 [v, j]) ≥ (L 1 [z, j] -L 1 [v, j]) + ρ ≥ -2∥L 1 ∥ 2 + ρ.In the last inequality above, we used the loose bound∥L 1 ∥ 2 ≥ max j∈[m],z∈[V] |L 1 [z, j]| for the spectral norm of L 1 .Using these in Eq. ( 19) gives:EQUATIONThe last inequality uses log(1 + x) ≤ x, x > 0.Clearly the bound above converges to H = ∑ j∈[m] πj ∑ z∈S j pj,z log 1 pj,z as ρ → ∞. From this and CE(L ρ ) ≥ H, the desired Eq. ( 21) follows. This completes the proof.Note that the first two conditions of the theorem are easy to satisfy.On the one hand, Eq. ( 20) always has a solution. For later use, we also note that the solution is unique on the subspace F and we denote this solution as L in . This is easy to check: for any j ∈ [m], ℓ j is a solution of Eq. ( 20) if and only if it takes the formℓ j = E ⊤ in,j E in,j E ⊤ in,j -1 a j + v, v ∈ N (E in, )where a j ∈ R S j -1 has entries log( pj,z / pj,z ′ ), z ′ ̸ = z ∈ S j . Thus, L in is explicitly defined as follows.Lemma 3 (L in ). Let L in be the unique solution of Eq. (20) on the subspace F. Then, L in has columnsℓ in j = E ⊤ in,j E in,j E ⊤ in,j -1 a j , ∀j ∈ [m],where a j ∈ R S j -1 has entries log( pj,z j / pj,z ′ ) for z ′ ̸ = z j ∈ S j .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 ≥ V, then the rank condition C of the theorem is trivial. Thus, we have the following corollary. Corollary 2. Suppose d ≥ V, then the entropy lower bound can be asymptotically attained by the NTP loss.Instead of the regularized problem (8), it is more convenient (but, equivalent due to convexity) to study the following ball-constrained version:EQUATIONparameterized by B > 0. We will study properties of the minimizers L B of (23) as B → ∞.The theorem below subsumes Theorems 1 and 2. Theorem 3. Suppose d ≥ V. Denote L B any solution of (23) for regularization B. Then, in the limit of diverging regularization the following statements are true.(i) lim B→∞ CE( L B ) = H.(ii) lim B→∞ P F ( L B ) = L in , where L in is defined in Lemma 3.(iii) The solution diverges in norm: lim B→∞ ∥ L B ∥ = +∞.(iv) The solution converges in direction to the optimal set of NTP-SVM ⋆ . That is, there exists minimizer L mm of NTP-SVM ⋆ such thatlim B→∞ L B ∥ L B ∥ * - L mm ∥L mm ∥ * = 0.The in-support NTP loss and its minimizer L in Define the in-support NTP loss as follows:CE in (L) = ∑ j∈[m] πj ∑ z∈S j pj,z log   1 + ∑ z ′ ̸ =z∈S j e -(L[z,j]-L[z ′ ,j])  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) ≥ CE in (L). In fact, because of the sparsity ansatz (∃j ∈ [m] such that S j < V) and strict nonegativity of exponential, the inequality is strict for all finite L, i.e.EQUATION)Moreover, observe that for any L, it holds thatEQUATIONTo see this, decompose L in orthogonal components onto F and F ⊥ , respectively, i.e., L = P F (L) + P ⊥ (L). By Eq. ( 18), L ⊥ := P ⊥ (L) satisfies Eq. ( 9), i.e.L ⊥ [z, j] -L ⊥ [z ′ , j] = 0 for all z, z ′ ∈ S j and j ∈ [m].The lemma below uses this to show that CE in has a unique minimizer in F.Lemma 4. The in-support loss CE in (L) has a unique finite minimizer in F, which we denote L in :EQUATION)The minimizer L in satisfies the NTP H -compatibility condition in Eq. (20) andCE in,⋆ := CE in (L in ) = H.Proof. Note that L ∈ F is a column-wise separable constraint ℓ j ∈ R(E ⊤ in,j ) that restricts the non-zero entries L[z, j] = ℓ j,z to indices z ∈ S j . The objective is also separable with respect to columns. Without loss of generality fix j = 1 and assume S j = [S j ]. It then suffices proving that the function f : R S 1 → R:EQUATION) has a unique minimizer in R( E ⊤ in,1 ), where E in,1 := E in,1 [:, S 1 ] ∈ R (S 1 -1)×S 1 is the non-zero block of E in,1 .By non-negativity of the KL divergence,f (ℓ) ≥ -∑ z∈S 1 p1,z log ( p1,z )with equality if and only if,EQUATIONIn turn, this is equivalent toEQUATION)To see the equivalence note the following: On the one hand, starting with Eq. ( 28) for z, z ′ ∈ S 1 , 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∑ z ′ ∈S 1 exp(ℓ z ′ ) exp(ℓ z ) = ∑ z ′ ∈S 1 exp(ℓ z ′ -ℓ z ′ ) = ∑ z ′ ∈S 1 p1,z ′ p1,z = 1 p1,z .Now, observe that ( 29) is the same as the NTP H -compatibility equation for j = 1. Moreover, note that a logit vector ℓ ∈ R S 1 solves (29) if and only if forℓ = E ⊤ in,1 E in,1 E ⊤ in,1 -1 a 1 + v, v ∈ N E in,1 ,where a ∈ R S 1 -1 with entries log( p1,z 1 / p1,z ′ ) for z ′ ̸ = z 1 ∈ S 1 . This proves that there is a unique minimizer ℓ in1 := E ⊤ in,1 E in,1 E ⊤ in,1 -1 a 1 on R E ⊤ in,1 .Since the choice of j = 1 above was arbitrary, this proves the lemma.This section proves a sequence of auxiliary lemmas that are used to prove Thm. 3 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 ∈ R V×m and all L ′ ∈ F ⊥ that additionally satisfy E out,j ℓ ′ j ≥ 1, ∀j ∈ [m] (i.e., L ′ is NTP-SVM ⋆ feasible), it holdsCE in (L) = lim R→∞ CE(L + RL ′ ).Using this, it holds for all L ∈ R V×m thatEQUATION)Proof. To see the first claim note thatL ′ ∈ F ⊥ implies ∀j ∈ [m], z ̸ = z ′ ∈ S j : L ′ [j, z] = L ′ [j, z ′ ] =⇒ ∀j ∈ [m] , z ∈ S j : ∑ z ′ ∈S j e -(L[z,j]-L[z ′ ,j]+R(L ′ [j,z]-L ′ [j,z ′ ])) = ∑ z ′ ∈S j e -(L[z,j]-L[z ′ ,j])andE out,j ℓ ′ j ≥ 1, ∀j ∈ [m] implies ∀j ∈ [m], z ∈ S j , v / ∈ S j : L ′ [j, z] -L ′ [j, v] > 0 =⇒ ∀j ∈ [m] , z ∈ S j : ∑ v/ ∈S j e -(L[z,j]-L[v,j]+R(L ′ [j,z]-L ′ [j,v])) R→∞ -→ 0 .The second claim applies the first claim for L = L in 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 islim B→∞ ∥ L B ∥ * = ∞Proof. Suppose on the contrary that lim B→∞ ∥ L B ∥ * < ∞. From Eq. ( 30), we haveCE( L B ) > lim R→∞ CE(L in + RL mm ).Moreover, by triangle inequality for nuclear norm ∥L R ∥ * ≤ ∥L in ∥ * + R∥L mm ∥ * . Thus, for any arbitrarily large R > 0, there exists large enough B such that ∥L R ∥ * ≤ B, making L R := L in + RL mm feasible in ( 23). These together, contradict optimality of L B , proving that lim B→∞ ∥ L B ∥ * = ∞. Finally, recalling the norm inequality∥ L B ∥ ≥ ∥ L B ∥ * / √ V , it also holds that lim B→∞ ∥ L B ∥ = ∞.The last lemma constructs a sequence of logits approaching the entropy lower bound while remaining feasible in problem ( 23) as B → ∞. Lemma 7. The objective function approaches the loss lower bound, i.e.,lim B→∞ CE( L B ) = CE in,⋆ = H .Concretely, there exists logit matrix L R parameterized by R = R(B) that is feasible in (23) and sastisfies:CE(L R ) ≤ CE in,⋆ +Ve 2∥L in ∥ 2 e -R .(31)Moreover, lim B→∞ R = ∞ and lim B→∞ R ∥ L B ∥ * = 1 ∥L mm ∥ *where ∥L mm ∥ * is the optimal cost of NTP-SVM ⋆ .Proof. We evaluate the loss achieved by the following candidate good pointL R := L in + R(B) • L mm .Here, (for large enough B > ∥L in ∥ * ) we setR := R(B) = ∥ L B ∥ * -∥L in ∥ * ∥L mm ∥ * . Note this is chosen so that ∥L R ∥ * ≤ ∥ L B ∥ * ≤ B towards making L R feasible in (23). Also, recall from Lemma 6 that ∥ L B ∥ * → ∞ as B → ∞; thus, also R → ∞.It remains to prove Eq. ( 31). This follows directly from the sequence of equations in (22).We prove each statement separately.Since for all L, it holds CE(L) ≥ H, it suffices to prove that H can be asymptotically attained. Since d ≥ V, this follows from Corollary 2. (For a more explicit proof, see Lemma 7).Since for all L it holds CE(L)≥ CE in (L) ≥ CE in,⋆ = H, it follows from Statement (i) that lim B→∞ CE in ( L B ) = CE in,⋆ .But CE in (L) = CE in (P F (L)) (see Eq. ( 25)). Thus, limB→∞ CE in (P F ( L B )) = CE in,⋆ .The desired now follows ince L in is unique minimizer of CE in on T by Lemma 4.Follows directly from Lemma 6.From Lemma 7, there exists L R for which the loss is close to H (see Eq. ( 31)).Towards arriving at a contradiction, we will show that if L B is not in the direction of L mm , then it incurs a loss that is larger than the upper bound of CE(L R ) computed in Eq. ( 31).To do this, assuming the statement of the theorem is not true, we will lower bound j,v]) .CE( L B ) -H = ∑ j∈[m] πj ∑ z∈S j pj,z log pj,z ∑ z ′ ∈S j e -( L B [j,z]-L B [j,z ′ ]) + ∑ v/ ∈S j e -( L B [j,z]-L B [(32)By our assumption, there exists ϵ > 0, such that there exists arbitrarily large B satisfying:EQUATION)for all minimizers L mm of NTP-SVM ⋆ . DefineL = 1 R ′ (B) L B -L in ,where, R ′ := R ′ (B) > 0 is chosen so that ∥ L∥ * < ∥L mm ∥ * . Concretely, this can be ensured by setting (to see this use the triangle inequality for nuclear norm):R ′ = (∥ L B ∥ * + ∥L in ∥ * + ζ) ∥L mm ∥ * , for some ζ > 0. Recall that ∥ L B ∥ * → ∞ as B → ∞. On the other hand ∥L in ∥ * is not dependent on B. Thus, in the large limit B → ∞, it holds ∥L mm ∥ * ∥ L B ∥ * -1 R ′ → 0 and ∥L in ∥ * ∥ L B ∥ * → 0.These combined with (33) show that we can always choose B large enough so that Eq. ( 33) guarantees, for some ϵ ′ > 0, that∥ L -L mm ∥ ≥ ϵ ′ .But L achieves the optimal cost of NTP-SVM ⋆ , since ∥ L∥ * < ∥L mm ∥ * . Thus, there exists δ ∈ (0, 1) and j ∈ [m] such that at least one of the following is true (i) ∃z and z ′ ̸ = z ∈ S j such thatEQUATIONEQUATION)Case (i): Without loss of generality (e ze z ′ ) ⊤ ℓ j ≤ -δ (otherwise, flip z, z ′ ). Thus, ignoring all but one term in (32) givesEQUATION) But, L B [j, z] -L B [j, z ′ ] = R ′ ( L[j, z] -L[j, z ′ ]) + (L in [j, z] -L in [j, z ′ ]) ≤ -R ′ δ + 2∥L in ∥ 2 (37)Put in (32) and using pj,z≥ 1/V , πj ≥ 1/m shows CE( L B ) ≥ H + 1 mV log 1 V e R ′ δ exp(-2∥L in ∥ 2 ) ≥ H + 1 mV log e R ′ δ V exp(2∥L in ∥ 2 ) Compare this with (31), it is clear that as B → ∞, so that R ′ → ∞, it holds 1 mV log e R ′ δ V exp(2∥L in ∥ 2 ) > 1 > V exp(2∥L in ∥ 2 ) e -R . Thus, CE( L B ) > CE(L in R ), a contra- diction.We can assume E in,j ℓ j = 0 for all j ∈ [m], since otherwise we are in Case (i). Now, again ignoring all but the (j, z) term in the CE loss for which (35) holds for some v / ∈ S j , we find j,z]) .CE( L B ) -H ≥ πj pj,z log p j,z ∑ z ′ ∈S j e -( L B [j,z]-L B [j,z ′ ]) + e ( L B [j,v]-L B [On the one hand, using E in,j ℓ j = 0, we have∑ z ′ ∈S j e -( L B [j,z]-L B [j,z ′ ]) = ∑ z ′ ∈S j e -(L in [j,z]-L in [j,z]) = ∑ z ′ ∈S j pj,z ′ pj,z = 1 pj,z ′ .On the other hand, by ( 35):e L B [j,v]-L B [j,z] ≥ e -R ′ (1-δ) e L in [j,v]-L in [j,z] ≥ e -R ′ (1-δ) exp(-2∥L in ∥ 2 ),Putting the above together yield:CE( L B ) -H ≥ πj pj,z log 1 + e -R ′ (1-δ) exp(-2∥L in ∥ 2 ) V ≥ e -R ′ (1-δ) mV 2 exp(2∥L in ∥ 2 ) + mV .where the second inequality uses log(1 + x) ≥ x 1+x , x > 0. Compare this with (31). As B → ∞, and noting that R, R ′ grow at the same rate, it holdse -R ′ (1-δ) mV 2 exp(2∥L in ∥ 2 )+mV > V exp(2∥L in ∥ 2 ) e -R . Thus, CE( L B ) > CE(L R ), a contradiction.In either case, we arrive at a contradiction, which completes the proof.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.Proposition 2 (Subspace collapse). Assume NTP-SVM ⋆ is feasible. There exists minimizer H = [h 1 , . . . , h m ] of NTP-UFM such that for all contexts j, j ′ ∈ [m] with same support set S j = S j ′ , their embeddings are same in direction in the limit of vanishing regularization, i.e., lim λ→0 h jh j ′ = 0. This is consistent with our empirical observations in Figs. 1 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) heatmaps.The proposition can be proved by inspecting NTP-SVM ⋆ : Since the affine constraints (9) for each sample j ∈ [m] depend solely on the support set S j , it can be shown that there exists a minimizerL mm = [ℓ mm 1 , • • • , ℓ mm m ] such that ℓ mm j = ℓ mm j ′whenever S j = S j ′ . This implies that the corresponding columns j, j ′ of V ⊤ are identical. As a result, by Cor. 1,lim λ→0 h j -h j ′ = 0.It is important to reiterate that Prop. 2 does not prevent the minimizer L λ = W λ H λ from accurately recovering the distinct soft-label values pj , pj ′ and achieving the empirical entropy lower-bound H as λ → 0: The collapse occurs on subspace F ⊥ , orthogonal to the data subspace F on which the logit matrix satisfies the log-odds constraints indicated in Thm. 2; See Secs. D.1 and 5 for numerical verifications. To gain insights, we start with an idealized symmetric setting.Proposition 3. Fix k ∈ [V -1] and suppose S contains all m = ( V k ) support sets of size k. Then, (i) The logit matrix takes the form L mm = (I V -1 V 11 ⊤ )S ≜ S. (ii) Word embeddings form equiangular tight frame (ETF) being equinorm and maximally separated.(iii) Context embeddings are equinorm and the embedding h j of context j is co-linear to ∑ z∈S j w z .In the symmetric setting of the proposition, we can analytically solve NTP-SVM ⋆ giving L mm 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:EQUATION)For k = 1 the embeddings recover the ETF geometry that has been previously shown for the setting of one-hot classification (Papyan et 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 j forms the same acute angle with all in-support word vectors w v , v ∈ S j and the same obtuse angle with all out-of-support word vectors w v , v / ∈ S j . Additionally, for k > 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.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 EQUATION)Computing the matrix square roots in Eq. ( 39) is as expensive as an SVD decomposition which can be prohibitive for large matrices. Instead, we suggest the following more computationally efficient proxies for estimating the directional component of the implicit geometries:EQUATIONThis leads to Proxy (P) introduced in Sec. 1.2.Ignoring for simplicity the projection to the 40) gives a simple explanation for similarity between embeddings: For given context j ∈ [m], the corresponding embedding has higher correlation with contexts j ′ whose support set intersection |S j ∩ S j ′ | is larger. Similarly, word embeddings are closer if the respective words appear together in more support sets: word embedding w z is more similar to w z ′ when |{j : w z , w z ′ ∈ S j , j ∈ [m]}| is larger. The projection then fixes the bias of the parameters: since ridge-regularization is strongly convex and softmax is invariant to shift, any minimizer of NTP-UFM satisfies 1 ⊤I V -1 V 1 V 1 ⊤ V subspace, Eq. (V W = 0 ( see Lemma 2).We will prove that for j, j ′ ∈ [m] with same support set S j = S j ′ the j-th and j ′ -th columns of L mm = [ℓ 1 , . . . , ℓ m ] are same. For the sake of contradiction, suppose that this is not the cases, i.e. ℓ j ̸ = ℓ j ′ . Consider the following three candidate solutions of NTP-SVM ⋆ :L mm = [ℓ 1 , . . . , ℓ j , . . . , ℓ j ′ , . . . ℓ m ] L mm flip = [ℓ 1 , . . . , ℓ j ′ , . . . , ℓ j , . . . ℓ m ] L mm avg = ℓ 1 , . . . , ℓ j + ℓ j ′ 2 , . . . , ℓ j + ℓ j ′ 2 , . . . ℓ m .By assumption L mm is optimal in NTP-SVM ⋆ . We now show that the other two matrices are also optimal.Firstly, note that L mm flip 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 L mm flip = rank (L mm ) ≤ d since permuting columns is rank-preserving operation. Moreover, permuting columns also preserves nuclear-norm of a matrix. To see this note that L mm flip = L mm P = UΣ(P ⊤ V ) ⊤ for some permutation matrix P ∈ R m×m . Since V is partial orthonormal, the same is true for P ⊤ V . Thus, ∥L mm flip ∥ * = tr(Σ) = ∥L mm ∥ * , which proves optimality of rank L mm flip .Similarly, we can argue that L mm avg 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 L mm avg = rank (L mm ) ≤ d since adding columns and multiplying them by scalars are rank-preserving operations.But now, from convexity of the objective function, and noting that L mm avg = 1 2 L mm + L mm flip we have∥L mm avg ∥ * ≤ 1 2 ∥L mm ∥ * + 1 2 ∥L mm flip ∥ * = ∥L mm ∥ * .Thus, L mm avg is also optimal. Now, from Lemma 1, we know that for some partial orthonormal matrix R, in the limit of vanishing regularization:H λ ∝ R ⊤ Σ 1 2 V ⊤ where V , Σ are SVD factors of L mmavg . Since the j, j ′ columns of L mm avg are the same by construction, the same is true for the j, j ′ columns of V ⊤ . This proves the desired.To prove the proposition, we first derive the dual of NTP-SVM ⋆ .We begin with a more convenient, but equivalent, formulation of NTP-SVM ⋆ . One way to do this is by recalling that NTP-SVM ⋆ is a convex relaxation to minW,H 1 2 ∥W ∥ 2 F + 1 2 ∥H∥ 2 F , s.t. ∀j ∈ [m], z j ̸ = z ∈ S j : (e z j -e z ′ ) ⊤ W h j = 0 ∀j ∈ [m], v / ∈ S j : (e z j -e v ) ⊤ W h j ≥ 1 .where, for each sample j, we choose an anchor label z j ∈ S j . We can now relax this in terms of the matrix X = W H ⊤ W ⊤ H . (Note that L essentially corresponds to the off-diagonal blocks of X). This gives us the following reformulation of NTP-SVM ⋆ . (Formally, it is easy to show that (41) is 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 X≽0 1 2 tr(X) s.t. ∀j ∈ [m], z j ̸ = z ∈ S j : X[z j , V + j] -X[z ′ , V + j] = 0 (41) ∀j ∈ [m], v / ∈ S j : X[z j , V + j] -X[v, V + j] ≥ 1 . Define A V×m such that A[z, j] = -λ j,z , z ̸ = z j and A[z j , j] = ∑ z̸ =z j λ j,z . (1 ⊤ V A = 0). The dual problem is max A min X≽0 ∑ j∈[m] ∑ v̸ ∈S j λ j,v + 1 2 tr I -A -A ⊤ I X s.t. 1 ⊤ V A = 0 A[v, j] ≤ 0, v ̸ ∈ S j = max A ∑ j∈[m] ∑ v̸ ∈S j λ j,v s.t. 1 ⊤ V A = 0 A[v, j] ≤ 0, v ̸ ∈ S j , ∥A∥ 2 ≤ 1 = max A tr(A ⊤ S) s.t. 1 ⊤ V A = 0 A[v, j] ≤ 0, v ̸ ∈ S j , ∥A∥ 2 ≤ 1where here S is defined in Prop. 3. The last line above holds since,∑ j∈[m] ∑ v̸ ∈S j λ j,v = ∑ j∈[m] ∑ v̸ ∈S j λ j,v + ∑ j∈[m] ∑ v∈S j v̸ =z j λ j,v -∑ j∈[m] ∑ v∈S j v̸ =z j λ j,v = ∑ j∈[m] ∑ v̸ =z j λ j,v -∑ j∈[m] ∑ v∈S j v̸ =z j λ j,v = ∑ j∈[m] ( 1 -|S j | V + |S j | V ) ∑ v̸ =z j λ j,v -∑ j∈[m] ∑ v∈S j v̸ =z j λ j,v = ∑ j∈[m] 1 -|S j | V A[z j , j] -∑ j∈[m] |S j | V ∑ v̸ =z j A[v, j] + ∑ j∈[m] ∑ v∈S j v̸ =z j A[v, j] = ∑ j∈[m] S[z j , j]A[z j , j] + ∑ j∈[m] ∑ v̸ ∈S j S[v, j]A[v, j] + ∑ j∈[m] ∑ v∈S j v̸ =z j S[v, j]A[v, j] = ∑ j∈[m] ∑ v∈[V] S[v, j]A[v, j] .We remark the following about the dual problem derived above. First, because S is feasible in NTP-SVM ⋆ and constraints are linear, Slater's conditions hold, thus, we have strong duality. Second, by complementary slackness, if the nonegativity conditions A[v, j] ≤ 0, v / ∈ S j 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[v, j] ≤ 0, v / ∈ S j , then the solution to the dual is simply A = UV ⊤ where UΣV ⊤ is the SVD of S. In fact, it is not hard to show that this is the unique solution of the relaxed dual; see for example (Thrampoulidis et al., 2022, Lemma C.1) for an analogous result. Therefore, if UV ⊤ 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 of NTP-SVM ⋆ must satisfy the inequality constraints with equalities. Formally, the optimal solution L satisfies simultaneously:L⊤ 1 V = 0 m and ∀j ∈ [m] : E in,j lj = 1 S j , E out,j lj = 0 V-1-S J .Note that this is a system of m + ∑ j∈[m] (S j + V -1 -S j ) = Vm linear equations. It can also be easily verified that these equations are linearly independent. Thus, their unique solution is S. This completes the proof of the proposition.At the heart of proving both Propositions 3 and 4 is the following result, which determines the SVD factors U, Σ, and V of S = (I V -1 V 11 ⊤ )S. 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 ∈ [V] and suppose S contains all ( V k ) support sets of size k, i.e. m = ( V k ). Let P ∈ R V×(V-1) be orthonormal basis of the subspace orthogonal to 1 V . Then, the U, Σ factors of the SVD of S are: U = P and Σ = ( V-2 k-1 ) I V-1 . Thus,UΣU ⊤ ∝ I V - 1 V 11 ⊤ .Moreover, the matrix V ∈ R m×(V-1) of right-singular vectors is such that its rows v j , j ∈ [m] can be expressed as follows with respect to the rows u z , z ∈ [V] of U:v j = 1 ( V-2 k-1 ) ∑ z∈S j u z j ∈ [m].Proof. We will compute explicitly S S ⊤ . Start with computing SS ⊤ . For diagonal elements, the dot product is between a row of S and itself, counts the number of 1s in that row. Since each element is included in ( V-1 k-1 ) support sets (choosing the remaining k -1 elements from the other V -1 elements), each diagonal entry of SS ⊤ is ( V-1 k-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 -2 elements from V -2 elements, giving ( V-2 k-2 ) for each off-diagonal entry of SS ′ . Overall, after algebraic simplification, we find thatSS ⊤ = V -2 k -1 I V + k -1 V -k 11 ⊤ .We may now use the closed form of S and few more algebra simplifications to find thatS S ⊤ = V -2 k -1 I V - 1 V 11 ⊤ .Let P ∈ R V×(V-1) denote an orthonormal basis of the subspace orthogonal to 1 V . SinceS S ⊤ = ( V-2 k-1 )PP ⊤ and S S ⊤ = UΣ 2 U ⊤ , we have shown that U = P and Σ = ( V-2 k-1 ) I V .Using 1 ⊤ S = k1 ⊤ m , we can computeS ⊤ S = S ⊤ S - 1 V (S ⊤ 1 V )(S ⊤ 1 V ) ⊤ = S ⊤ S - k 2 V 1 m 1 ⊤ m .From this and the fact that diag(S ⊤ S) = kI m , and recalling that S ⊤ S = V Σ 2 V ⊤ , we find that all context embeddings are equinorm with squared norm (i.e. diagonal entries of V ΣV ⊤ ) equal to∥h j ∥ 2 = k -k 2 V ( V-2 k-1 ) , ∀j ∈ [m].Recall also that the diagonal entries of UΣU ⊤ are∥w v ∥ 2 = 1 - 1 V V -2 k -1 , ∀v ∈ [V].From the off-diagonal entries of S ⊤ S, we can infer the angles between different context embeddings.Let j ̸ = j ′ ∈ [m], then the (j, j ′ ) of S ⊤ S is |S j ∩ S j ′ | -k 2 V . Thus, the (j, j ′ ) off-diagonal entries of V ΣV ⊤ are h ⊤ j h j ′ = |S j ∩ S j ′ | -k 2 V ( V-2 k-1 ) =⇒ cos h j , h j ′ = |S j ∩ S j ′ | -k 2 V k -k 2 V j ̸ = j ′ ∈ [m] .Recalling that W H ⊤ = L mm , we may now compute the angle between word and context embeddings as follows:EQUATION)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 S = (I V -1 V 11 ⊤ )S is the unique minimizer of NTP-SVM ⋆ .To prove this, we appeal to Proposition 5. First, it is straightforward checking that L mm satisfies the SVM constraints (in fact, with equality). Thus, from Proposition 5, it suffices that the matrix A := UV ⊤ satisfies A[v, j] < 0 for all v / ∈ S j and all j ∈ [m]. But, in our caseA = UV ⊤ = ( V-2 k-1 ) -12 S because Σ = I V-1 . Thus, the desired condition holds for A since S[v, j] = -k/V < 0 for all j / ∈ S j and j ∈ [m].Here, our 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 S is shown in Fig. 7-(c ). We set πj = 1/m and for this fixed support matrix S, we generate random soft labels pj ∈ ∆ V-1 such that pj,v = 0, v ̸ ∈ S j .We display the training evolution in Fig. 8 . We train NTP-UFM with SGD and small weight decay λ = 10 -5 until reaching the empirical entropy lower-bound H ≈ 1.04 within an order of 10 -4 as shown in Fig. 8-(a ). This decrease in the loss is accompanied by a consistent increase in the parameter norms, as displayed in Panel (b). Despite the increase in norms, the projection of L k onto the subspace F remains close to the finite component L in specified by Thm. 2. This is illustrated in Panel (c) and ensures the recovery of the soft-labels on the in-support tokens.In Fig. 7 , we visualize the logits L k = W k H k , and cosine similarities CORR(W ⊤ k ) and CORR(H k ) 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 visualizes the prediction made by Thm. 1 and Cor. 1. In Panel (c), we also visualize the cosine similarity of S and S ⊤ , which we introduced as a less expensive proxy of W mm and H mm (see App. C.4). While the theory and proxy are 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, S and SIM W ⊤ , S ⊤ , where SIM X, Y measures the structural correlation between the two matrices (see Sec. 5). Fig. 8-(d ) confirms the high correlation between the proxy (P) and the embeddings' implicit geometry. We also note that in the special case where the 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 mm = S (we verified that the dual-certificate condition in Prop. 5 holds).We use a total of three datasets. We first employ two smaller-scale datasets, which are well-suited for verifying our theoretical solutions. Then, we use one larger-scale dataset designed to investigate the geometric properties of text in conditions that approximate real-world text scenarios. Structural correlation between the learned embeddings and the proxy, i.e., SIM H, S , SIM W ⊤ , S ⊤ . The correlation with H mm and W mm⊤ (instead of S, S ⊤ ) is displayed with dashed lines for reference.Synthetic. We generate each context xj as follows: We select T -1 = 5 words (tokens), and manually come up with the T = 6-th tokens that are consistent with the given context. To model the probabilities πj and pj , we sample each support independently to a number of repeats to emulate the behavior of repeated context in natural language. For instance, the context xj = "Lily wants to try the" is followed by S = {soup, game, movie} with respective probabilities [0.5, 0.25, 0.25] . The dataset consists of n = 116 samples, containing m = 16 distinct contexts, and a vocabulary size of V = 30. Each context has a fixed support set length of |S j | = 3. In Fig. 9 , we show the m = 16 unique contexts and the soft labels on their next token. In 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 xj and support sets S j from the TinyStories dataset and we create the training data from the most frequent word-level contexts with length 5, such as xj = ["once", "upon", "a", "time", ","]. For the support sets S j , we record all next tokens of xj in the original dataset. Then, we replace the words in the contexts with their synonyms to generate new contexts with identical support sets. This allows us to have multiple contexts sharing common support sets, while controlling the vocabulary size, which is set to V = 104. The final dataset consists of n ≈ 3050 samples with m ≈ 400 unique contexts. The frequencies pj are determined by independently sampling each next token from S j several times. Here, H = 1.0166. TinyStories. To experiment with more a 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 = 6. 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 m ∼ 10 5 and H = 0.3112.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 -4 , 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 5 epochs to increase the learning rate to 5 × 10 -4 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 .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 embeddings similarity in Fig. 2 , we choose 10 unique support sets in the dataset with |S j | > 2. For each of the chosen support sets S, we choose 10 distinct contexts j such that S j = S (a total of 100 samples), and illustrate the Gram matrix of the centered support matrix S and context embeddings H for this subset of the contexts. Given that the support sets are sparse, we choose to use contexts with the largest support set size which would result in a more pronounced pattern on the correlation matrices. We skip visualizing the word embeddings W in this experiment, as the heatmaps did not display any noticeable visual structures.For the metrics displayed in Fig. 11 , due to the larger number of context embeddings, we only compute the metrics such as ∥P F (L k ) -L in ∥, 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 a consequence of layer normalization in the final embedding layer. rank among all minimizers depends on 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 d TF = 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 d = 64, for a 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.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.The equivalence between GD and regularization paths has been rigorously established in certain convex settings(Ji et al., 2020), but it may not hold generally. Our setting is non-convex, so we use this equivalence as a heuristic proxy; see future work discussion in Sec. 6.This is the Structural Similarity Index Measure (SSIM) used to measure the similarity between two images(Wang et al., 2004).Here, use the fact that the minimizer of the LHS of (12) is attained because of coercivity of the nuclear norm. 
 
 Appendix:  


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

Reviews: review: paper_summary: paper_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. 

The 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.
strength_weakness: strength_weakness: **Model is unrealistic**

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

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

[1] Carl Allen and Timothy Hospedales. Analogies explained: Towards understanding word embeddings. In International Conference on Machine Learning, pp. 223–231. PMLR, 2019.

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

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

**Writing and formulation are hard to follow**

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

**Minimal insight can be derived from theoretical results**

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

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

2. 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: questions: 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.

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

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

I 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!
paper_summary: paper_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.
strength_weakness: strength_weakness: How this work could be further strengthened:

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

* 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.
paper_summary: paper_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.

This 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
small-scale real language datasets

I 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.
strength_weakness: strength_weakness: 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: questions: Please see above.
paper_summary: paper_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.
I 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.
strength_weakness: strength_weakness: - 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)
- Notation used is not just explained but even just introduced in the appendix
- Relevance to practical language modeling and potential for useful insights derived from the described characterization are not given
questions: questions: - please use parentheses around citations where appropriate (\citep{})
- typo: empiriical -> empirical
paper_summary: paper_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.

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

Clarity:  The paper is generally very clear.

Originality:  As far as I know, this is original work.

Significance:  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.
strength_weakness: strength_weakness: - The findings are somewhat anecdotal, based on toy settings and qualitative results
- 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: questions: - 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).

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

- Spell out "ETF" (Section 3.4.1).

- I believe the reference to Harris 1954 should be Firth, J. R., Studies in Linguistic Analysis, 1957. (Section 4.1)

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

- There are a couple of missing citations:    
  * "the linear relationships underlying word analogies" (Section 5)
  * "recent promising studies ... that leverage the neural-collapse geometry" (Section 6)

- Typos, grammar, etc.:  
  * "their synonymous" --> "their synonyms".  How did you derive the synonyms?  (Section 4.2)
  * "Different to them" --> "Unlike them" (Section 5)

Comments: title: Thank you for the additional insights
comment: Thank you again for your detailed review and feedback. We appreciate it. We have made every effort to respond to your questions and clarify misunderstandings in terminology (such as the meaning of in-support words) that may have led to misconceptions. We hope you are satisfied with our responses and look forward to your reply. We are also happy to engage in further discussion and answer any questions before the discussion period ends tomorrow.
Thank you very much for your positive feedback! We appreciate your time.
Thank you for your careful read and suggestions, and for endorsing our paper. We appreciate your time.
Thank you for acknowledging our response.

We are glad that the misunderstanding about our main contribution—explicitly accounting for the fact that multiple tokens (labels) can follow a given sequence of tokens (context) in the training corpus, resulting in a **sparse** soft-label setup—has been addressed. 

We reiterate that combining this viewpoint with unconstrained features offers a new perspective in studying word/context embeddings. Given the originality and timeliness (word embedding geometry studies date back to the word2vec era and have recently surged in the context of LLMs), we believe our contribution is valuable to the community.

Additionally, beyond the motivational figures on embedding geometry, some of our theoretical analysis already applies to the regime d<V. For instance, in Appendix C.3.1 (see definition 4 and in the paragraph below), we show that the feasibility of NTP-SVM with an additional constraint rank(L)≤d is necessary and sufficient for the model to reach the entropy lower bound. We consider this a fundamental result. Besides, the results remain non-trivial even when d>V. Therefore, we respectfully disagree with the assessment that our work lacks in becoming a “more worthwhile” contribution.
Re 1: Oh, you're right, I must have misunderstood that!

Re d>V: Additional results on this seen to a proper conclusion like in the submitted paper may raise this to be a more worthwhile contribution, but I don't think that would be ready in time for a camera-ready...

With the clarified practical insights, I am raising my score slightly.
Thank you for the responses!  The additional explanation and results do help.  If the paper is accepted, I hope that you can add the additional results.  That being said, the experimental settings are still toy-ish, and that more realistic experiments could have been done (if that is not the case, I would be happy to be corrected!).  For this reason I continue to feel that the paper is only marginally above the acceptance threshold.
These follow up comments help give me a better understanding of your contributions.
Dear Reviewers,

Thank you for reviewing our submission. We have tried to address all your comments while staying within the character limit.

We hope our responses are clear. If so, we would appreciate if this is reflected to your evaluation.
If you have further questions, we are more than happy to elaborate on our initial responses during the discussion period.

Thank you again.
Thank you very much for acknowledging our response. If any more questions arise, we are happy to discuss!
This explanation looks good to me. Thank you for your response!