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

Thanks for your overall positive feedback. In the final version of the paper, we will make sure to address the clarity issues and typos pointed out. Thank you!

d>V: Unlike static models, contextualized models involve two levels of compression: mapping 1) V words, and 2) m contexts (sequences of words) to d-dim space. Despite the requirement d>V, we allow the number of contexts to be much larger, i.e. m>>d. So, from the perspective of contexts, it remains interesting and nontrivial to formalize how the model compresses them into a lower dimensional space. Besides, it is not a priori clear to us, even under the assumption of d>V, how the geometry depends on the language statistics as encoded in the sparsity pattern. That said, we provide preliminary results for d<V at https://ibb.co/Yb9S5Tr as motivation for future work.

We also remark that the rank-constrained logit relaxation (see Lem1) holds for any d. For example, this directly yields a necessary and sufficient condition for the loss to reach the entropy lower bound, which is the feasibility of the rank-constrained NTP-SVM problem. We also hypothesize that the sparsity pattern of language promotes low-rank solutions to the NTP-SVM problem which could allow extensions of the parameter convergence results to the setting d<V.

Visual inspections: Thank you for commenting on that. Following your suggestion, we now verify the correlation quantitatively as well by computing the normalized covariance of the heatmaps. The plot of this metric vs epoch is available at https://ibb.co/Yb9S5Tr . We hope this addresses your concern and thanks again for the suggestion.

Thanks for your review. Before addressing your concerns:

1. We are puzzled by your comment “only the one-hot case…”. In fact, we explicitly account for the fact that multiple tokens (labels) can follow a given sequence of tokens (context) in the training corpus which results in a soft-label setup.
2. Allow us to clarify our contributions: Our high-level goal is to identify how the information from the training corpus gets encoded in language models. This is influenced by several factors: 1) optimization objective, 2) model architecture, and 3) model/data scale. We decouple these factors by focusing on first one: interactions of context/word embeddings $H,W$ in the objective at optimality. We find them as matrix factorization of a logit that decomposes to a finite PMI-like matrix $L_{in}$ and a directional component $L_{mm}$, which is driven by the sparsity pattern of the conditional probability of next tokens. Overall, this a) gives a new perspective to classical static word embedding results such as [LG14], b) connects contextualized embeddings to the sparsity pattern of language, c) establishes conditions for NTP loss reaching lower-bound, d) confirms the log-bilinear NTP-UFM as a good analytical proxy for complex models, e.g. deep transformers.

d>V: Unlike static models, contextualized models involve two levels of compression: mapping 1) V words, and 2) m contexts (sequences of words) to d-dim space. Despite the requirement d>V, we allow for the number of contexts to be much larger, i.e. m>>d. So, from the perspective of contexts, it remains nontrivial to formalize how the model compresses them into a lower dimensional space. That said, we provide preliminary results for d<V at https://ibb.co/Yb9S5Tr as motivation for future work.

Practical insights: We believe that our perspective opens avenues to counteract the negative impacts of biases and extreme imbalances specific to language. Better understanding of the context/tokens distributions and how they correlate with the model embeddings may help combat unfairness in large language models. For example, we envision using the geometry of embeddings to reverse engineer the design of more robust loss or to devise posthoc robust sampling methods that account for imbalances in the token/context distributions of language data. Our work is the first step in making these connections more concrete but further investigations are needed to fully achieve these practical goals.

Thanks for your positive feedback. As a matter of fact, your summary very well captures the message of our work!

Repeating here for completeness: Our high-level goal is to identify how the information from the training corpus gets encoded in prediction-based language models. This is influenced by several factors: 1) optimization objective, 2) model architecture, and 3) scale of model and data. We decouple these factors by focusing on the first one: we study the interactions of context embeddings $H$ and word embeddings $W$ in the objective at optimality. We find that they occur as matrix factorization of a logit matrix that decomposes to a finite PMI-like matrix $L_{in}$ and a directional component $L_{mm}$. This decomposition is driven by the sparsity pattern of the conditional probability distributions of next tokens.

Overall, this a) gives a fresh perspective to classical matrix factorization results for static word embeddings such as [LG14], b) establishes a connection of the contextualized embeddings to the sparsity pattern of language, c) establishes conditions under which the NTP training loss reaches its lower-bound, d) confirms the log-bilinear model (NTP-UFM) is a good analytical proxy for more complex prediction-based models, e.g. deep transformers.

Looking ahead we believe that capturing such influences of how the geometry depends on language statistics such as the conditional probability distributions of next tokens, opens avenues to counteract the negative impacts of biases and extreme imbalances in the training set, as text data is known to follow a heavily long-tailed distribution by Zipf's law. Better understanding of the distribution of context/tokens and how they correlate with the model embeddings may help combat unfairness in large language models. For example, we envision using the geometry of embedding information to reverse engineer the design of more robust loss or to devise posthoc robust sampling methods that account for imbalances in the token/context distributions of language data. We acknowledge that further investigations are needed to fully achieve this goal. However, we consider our framework a first step in making these connections more concrete. Furthermore, we believe these foundations could help develop interpretability metrics for LLMs in order to better analyze their properties.

We thank you for your positive feedback. We would leverage the extra page of the camera-ready to address your concern about the density of information in the current form. We will also include a clearer statement on the motivation/implications, as briefly discussed below.

Our high-level objective is to understand what information from the training corpus gets encoded in the embedding space of contexts and words. In other words, which structures or statistics from the training set influence the similarity and arrangements of embeddings and how this correlates with the semantic similarity of words and contexts. This is influenced by several factors: 1) optimization objective, 2) model architecture, and 3) scale of model and data. We decouple these factors by focusing on the first one: we study the interactions of context embeddings $H$ and word embeddings $W$ in the objective at optimality.

We find that the statistical properties that are important are encoded in the sparsity patterns of the conditional probability of the next tokens. For example, this is reflected in our experiments by the similarities between the Gram matrices of word/context embeddings and $SS^T, S^TS$. Conceptually, these findings relate prediction-based models that operate on contextualized word embeddings to models based on co-occurrence statistics, much like the classical connection of word2vec, which operated on static word embeddings, to count-based models [LG14]. The latter classical study motivated and resulted in improved models such as GloVe. We anticipate that our work provides analogous opportunities for contextualized word embedding models.

We also believe that such connections can open avenues to counteract the negative impacts of biases and extreme imbalances in the training set, as text data is known to follow a heavily long-tailed distribution by Zipf's law. Better understanding of the distribution of context/tokens and how they correlate with the model embeddings may help combat unfairness in large language models. We acknowledge that further investigations are needed to fully achieve this goal. However, we consider our framework a first step in making these connections more concrete. Furthermore, we believe these foundations could help develop interpretability metrics for LLMs in order to better analyze their properties.

We appreciate your thorough review.

In-support tokens: For context j, a token z is in-support if $p_{j,z}>0$, i.e. it appears after the context within the dataset. This is important as it relates to a few of your comments. We apologize for the confusion.

d>V: Unlike static models, contextualized models involve two levels of compression: mapping 1) V words, and 2) m contexts (sequences of words) to d-dim space. Despite the requirement d>V, we allow the number of contexts to be much larger, i.e. m>>d. So, from the perspective of contexts, it remains interesting and nontrivial to formalize how the model compresses them into a lower dimensional space. That said, we provide preliminary results for d<V at https://ibb.co/Yb9S5Tr as motivation for future work.

Fresh perspective: The classic [LG14], interprets minimizers of word2vec as matrix factorization of PMI matrix. This relies on a subtle assumption: the PMI matrix has well-defined entries, i.e. co-occurrence probability of any word given context is non-zero, typically enforced heuristically by smoothing. Instead, we account explicitly for sparsity which drives our decomposition of L in two orthogonal components: 1) $L_{in}$, a finite component conceptually similar to PMI, 2) $L_{mm}$, a directional component that only depends on the sparsity pattern.

Correct prediction: This is subtle but important: Prop1 describes the dominant component $L_{mm}$. But, recall that the logits have a second finite component $L_{in}$ (living on an orthogonal subspace). This is the one that accounts for the conditional probability of in-support tokens ensuring a correct prediction. See Fig5 for numerical verification of convergence to the two components. We will re-emphasize this after Prop1. Thanks!

Beyond NTP Great question. The framework is indeed more general. MLM is an example: it applies directly to simplified setting of masking a single token, but for the practical setting of random masking of several tokens one might need to be more careful to capture all the intricacies. While interesting, we leave this for future work.

[JT19,T24]: The basis of our contribution is incorporating the UFM framework within the formalism of [T24]. Technically, training both $W$ and $H$ makes the problem non-convex, but convex relaxation (4) allows us to leverage regularization path ideas from prior works [JT19] back to [RZH03]. We also go beyond [T24] in our empirical evaluations on language data.