Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yize Zhao, Tina Behnia, Vala Vakilian, Christos Thrampoulidis

TL;DR

The paper investigates how next-token prediction (NTP) training shapes the geometry of context and word representations. It introduces a sparse soft-label classification view (NTP-UFM) with unconstrained embeddings and analyzes the logit space via a regularization-path approach, revealing a sparse plus low-rank decomposition of the logit matrix L = WH near the entropy lower-bound. The authors prove that, as ridge regularization vanishes, Lλ decomposes into a sparse in-support component Lin and a dominant low-rank component Lmm whose structure is determined solely by the support pattern of next-token distributions, predicting subspace-collapse of contexts sharing the same next-token sets. They further connect this theory to neural-collapse-like geometry, provide practical proxies (centered centered-support tilde S) to approximate Lmm, and validate the framework with synthetic data and TinyStories experiments, showing that embeddings converge directionally to the predicted manifolds and align with the proposed proxies. This work offers a principled, data-driven lens to relate linguistic statistics to model representations under NTP, with implications for interpretability, bias analysis, and training dynamics.

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

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

TL;DR

The paper investigates how next-token prediction (NTP) training shapes the geometry of context and word representations. It introduces a sparse soft-label classification view (NTP-UFM) with unconstrained embeddings and analyzes the logit space via a regularization-path approach, revealing a sparse plus low-rank decomposition of the logit matrix L = WH near the entropy lower-bound. The authors prove that, as ridge regularization vanishes, Lλ decomposes into a sparse in-support component Lin and a dominant low-rank component Lmm whose structure is determined solely by the support pattern of next-token distributions, predicting subspace-collapse of contexts sharing the same next-token sets. They further connect this theory to neural-collapse-like geometry, provide practical proxies (centered centered-support tilde S) to approximate Lmm, and validate the framework with synthetic data and TinyStories experiments, showing that embeddings converge directionally to the predicted manifolds and align with the proposed proxies. This work offers a principled, data-driven lens to relate linguistic statistics to model representations under NTP, with implications for interpretability, bias analysis, and training dynamics.

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 findings 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.
Paper Structure (52 sections, 18 theorems, 103 equations, 18 figures)

This paper contains 52 sections, 18 theorems, 103 equations, 18 figures.

Table of Contents

  1. Introduction
  2. Methodology
  3. Summary of findings
  4. Related work
  5. Formulation
  6. NTP objective as sparse soft-label classification
  7. Unconstrained features model for NTP training
  8. Reformulation in terms of logits
  9. Analysis of unconstrained features model for NTP training
  10. NTP-SVM logits
  11. Regularization-path analysis of logits
  12. Geometry of embeddings
  13. On the solution of NTP-SVM*
  14. Special case: support sets of equal sizes
  15. A simple candidate solution
  16. ...and 37 more sections

Key Result

Lemma 1

Denote logit matrix $\bm{L}=[\boldsymbol{\ell}_1,\ldots,\boldsymbol{\ell}_m]\in\mathbb{R}^{V\times m}$ and let $\bm{L}_{\lambda}$ be a minimizer of with SVD $\bm{L}_\lambda=\bm{U}\boldsymbol\Sigma{\bm{V}}^\top$, where $\bm{U}\in\mathbb{R}^{V\times r}, \boldsymbol\Sigma\in\mathbb{R}^{r\times r}, {\bm{V}}\in\mathbb{R}^{m\times r}$ and $r=\operatorname{rank}\left(\bm{L}_\lambda\right)\leq d$. Then:

Figures (18)

  • Figure 1: Can we predict the geometry of a model's context and word representations based on the structure of the training data? This figure demonstrates that it is possible. Panel (a) shows the cosine similarity of context and word embeddings, $\textsc{corr}({\bm{H}})$ and $\textsc{corr}({\bm{W}}^\top)$, at the end of training a four-layer transformer on the Simplified TinyStories dataset, long enough for the NTP loss to converge to its empirical entropy lower-bound. Panel (b) illustrates $\textsc{corr}({\bm{H}})$ and $\textsc{corr}({\bm{W}}^\top)$ found by training our analysis model, NTP-UFM (see Eq. \ref{['eq:NTP intro']}), on the same data. This model abstracts the neural network and allows to predict the geometry analytically. Panel (c) shows this analytical prediction of the geometry that is obtained purely from the data. Specifically, ${\bm{H}}^{\rm{mm}}$ and ${\bm{W}}^{\rm{mm}}$ are determined by the right/left singular factors (Claim \ref{['P directional']}) of the low-rank (max-margin) component $\bm{L}^{\rm{mm}}$ of logits (Claim \ref{['P logits']}). Finally, Panel (d) provides an easy to compute heuristic proxy for the embeddings' geometry based on computing overlaps between the support sets of context’s next-words. Details in Secs. \ref{['sec:summary']} and \ref{['sec:exp_main']}.
  • Figure 2: In the limit of vanishing regularization, the ridge-regularized NTP objective solution $\bm{L}_{{\lambda}}$ (aka the logit matrix) decomposes into two orthogonal components. The component $\bm{L}^{\rm{in}}$ inherits the sparsity of the matrix ${\bm{P}}$ of next-token conditional probabilities. In the orthogonal complement, $\bm{L}_{\lambda}$ diverges in norm but converges in direction to $\bm{L}^{\rm{mm}}$, which results from a low-rank-promoting nuclear-norm minimization that depends only on the sparsity pattern of ${\bm{P}}$.
  • Figure 3: Similar to Fig. \ref{['fig:simplified_intro']}, this time on a 12-layer TF trained on a subset of $100$ stories from the TinyStories dataset. Here, computing the theoretical prediction ${\bm{H}}^{\rm{mm}}$ is computationally expensive. Thus, we compare the embeddings geometry with the Proxy \ref{['P proxy']}. Details in Sec. \ref{['sec:exp_main']}.
  • Figure 4: Evolution of \ref{['eq:ufm']} parameters $\bm{L}_k$, $\textsc{corr}({\bm{W}}_k^\top)$ and $\textsc{corr}({\bm{H}}_k)$ when training close to convergence to the empirical entropy $\mathcal{H}$ (See Fig. \ref{['fig:UFM_curves']}). At the end of the training, the parameters align with the prediction of Thm. \ref{['thm:reg path']} ((a) vs (b)). Additionally, the correlation patterns between the embeddings closely follow the similarities between the support sets ((a) vs (c)). See Sec. \ref{['sec:ufmexp']} for details.
  • Figure 6: Experiments on Synthetic and Simplified TinyStories datasets. (a): CE approaches $\mathcal{H}$, (b): Parameters' norms grow during training. (c):$\bm{L}_k$'s projection on the data subspace $\mathcal{F}$ converges to the sparse component $\bm{L}^{\rm{in}}$ specified by the soft-labels $\hat{\bm{p}}_j$. (d):$\textsc{sim}({\bm{H}}_k,\widetilde{\bm{S}})$ (solid) and $\textsc{sim}({\bm{W}}_k^\top,\widetilde{\bm{S}}^\top)$ (dashed). At the final stage the learned embeddings exhibit high similarity with proxy \ref{['P proxy']}. (e): Same as (d), this time comparing with theory, i.e., $\textsc{sim}({\bm{H}}_k,{\bm{H}}^{\rm{mm}})$ and $\textsc{sim}({\bm{W}}_k^\top,{{\bm{W}}^{\rm{mm}}}^\top)$. Details in Sec. \ref{['sec:exp_main']}.
  • ...and 13 more figures

Theorems & Definitions (29)

  • Definition 1: NTP-UFM
  • Definition 2: Regularization path
  • Lemma 1
  • Definition 3
  • Theorem 1
  • Theorem 2
  • Corollary 1
  • Proposition 1: Subspace collapse
  • Proposition 2
  • Proposition 3
  • ...and 19 more