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

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.

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

+ The idea is well-motivated and the paper is generally well-written. + The results, while preliminary, suggest that the analysis approach is promising.

Reasons to reject

- The findings are somewhat anecdotal, based on toy settings and qualitative results - The assumption d > V (the hidden dimensionality is larger than the vocabulary) is very unrealistic, and it's not clear how it might be relaxed in future work (the authors don't address this)

Questions to authors

- There seems to be a typo in the definition of $f_{\theta'}$, as its output is not in $\mathcal{V}$ as stated (Section 2, first paragraph). - 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)

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.

Rating

4

Confidence

2

Ethics flag

1

Reasons to accept

The clear stating of assumptions and rigorous treatment are appreciated, as is the extensive appendix.

Reasons to reject

- The focus on d > V - 1 makes the entire analysis not fit the standard language modeling setup where d <<< V - 1 (in fact the whole motivation for neural modeling) - 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 to authors

- please use parentheses around citations where appropriate (\citep{}) - typo: empiriical -> empirical

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.

Rating

6

Confidence

1

Ethics flag

1

Reasons to accept

I think this work can deepen our understanding of why the next-word prediction objective works so well.

Reasons to reject

This is not really a reason to reject the paper, but could you briefly summarize your work for a researcher who mainly works in empirical NLP rather than in the ML theory domain? Additionally, I am curious about the potential practical implications and applications of this work.

Questions to authors

Please see above.

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.

Rating

7

Confidence

2

Ethics flag

1

Reasons to accept

How this work benefits the community: * This paper clearly has some interesting ideas on a methodology for analyzing word embeddings, which would be valuable to the community. * The paper provides strong theoretical support for claims and hypotheses, that would be very interesting to researchers working on analyzing large language models. * Based on the observations and learnings from this work, potentially impactful research directions are outlined for the audience.

Reasons to reject

How this work could be further strengthened: * 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.

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.

Rating

4

Confidence

4

Ethics flag

1

Reasons to accept

1. Understanding the geometry of the contextual and word embeddings is a well-studied area and is generally useful for interpreting and predicting language model's behavior. 2. The paper clearly considers many related works and thoroughly discusses how their results fit in with prior ones. I would, however, suggest that the authors include more works on interpretability in language models. 3. The theory can accommodate arbitrarily complex data and the training of transformers, which are the popular choice for language models these days (there are some caveats, see below). 4. This approach to analyzing next-token prediction is thought-provoking. I found the future work section to be very interesting, and I think that extending these ideas to accommodate those more realistic settings could be valuable.

Reasons to reject

**Model is unrealistic** 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 to authors

1. Is any part of the analysis unique to next-token prediction? I believe it's possible to formulate masked language modeling in the same way, though one would have to be a little bit careful about what to do with the other masked tokens in a sequence when forming a prediction. 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!

Thank you for the additional insights

These follow up comments help give me a better understanding of your contributions.