Closed-Form Training Dynamics Reveal Learned Features and Linear Structure in Word2Vec-like Models

Thank you for your assessment and review, we appreciate it. We are glad that you found our results clear and our main idea effective. Some comments regarding your concerns and questions:

What’s the point of this analysis? Isn’t word2vec outdated? Actually, the core mechanistic idea behind word2vec is back in vogue! Many recent papers empirically report linear structure in the latent representations in LLMs – a direct analog of the word2vec embedding geometry. (For example, see https://arxiv.org/abs/2305.16130, https://arxiv.org/pdf/2303.04245, and the references we mention in Section 4).

Trying to understand this behavior directly in large language models is an admirable goal to strive towards. However, we believe it is too difficult to tackle directly (for now). Previous efforts in this direction have not gotten very far. A useful way for science to proceed in this situation is to focus on first understanding the mysterious behavior in the simplest possible system that exhibits it. Doing this is one of our primary contributions.

In general, one of the main outstanding questions of deep learning theory is: what is the mechanism and character of representation learning? Theorists have very few solvable cases, and these analyses tend to rely on unrealistic assumptions on either the data or the learning task (or both). We believe our result is powerful because it answers this question in a relevant natural language task under very few, relatively weak assumptions.

(Also, it’s a bit surprising that we lack a theory of learning dynamics for one of the foundational algorithms of the field! Even if it weren’t directly connected to understanding modern models, for posterity’s sake we should try to solve and deeply understand the small handful of truly breakthrough algorithms in the history of neural learning.)

On our assumptions. We respectfully disagree that our solution relies on a lot of unverified assumptions. On the contrary, our result relies on rather few assumptions, especially compared to what is par for the course in theory. We tried to emphasize this point throughout the manuscript. For example, our solution does not rely on any distributional assumptions on the data, which is uncommon in theoretical literature. We clearly list the set of approximations and assumptions in Figure 2 – there are only four in total. The first two are approximations to word2vec which are supported by our empirical checks (see Figure 1, Figure 3, the discussion in Appendices A.3, and Figure 5 in Appendix B). The other two are technical conveniences that simplify the analysis; we expect that they may be relaxed (or possibly eliminated completely) with some additional effort, as is suggested by Figure 2.

In particular, we do not believe that the similarity between QWEMs and word2vec is sensitive to the choice of dataset or the model size. (Please do point us to where in the paper this was unclear, as we’d like to clarify this in the writing.) We have done some empirical checks confirming this, as we mention in Appendix A; we will include more details in the final manuscript.

As a side note, since our results are data distribution agnostic, we expect them to carry over to a broad class of linearized contrastive algorithm. For example, if your contrastive loss has the form $\\mathcal{L}(W) = \\mathbb{E}\_\\text{positive pairs}[f^{+}(w^T w’)] + \\mathbb{E}\_\\text{negative pairs}[f^{-}(w^T w’)]$, then you can simply take the quadratic Taylor approximation of your $f^{+}(w^Tw’)$ and $f^{-}(w^T w’)$, complete the square, and obtain a target matrix in terms of the positive and negative distributions for your learning problem of interest. In this sense, the idea behind our Theorem 1 is actually very general. See Appendix D.2 for an example using the SimCLR loss.

Why does gradient descent in word2vec lead to embeddings similar to QWEM rather than directly approximating PMI? This is a good question, and the answer is a key insight of our result. The crux of the issue is in adequately handling the rank constraint. (In general, rank-constrained optimization is NP-hard, so handling this properly is a subtle business.) Let us compare the two approaches.

• The idea behind SVD-approximating PMI is to say “Let us first solve the unconstrained minimizer of the loss (i.e., the PMI matrix) and then, at the end, apply the rank constraint by choosing the closest (in Frobenius norm) rank-$d$ matrix.” This fails. The reason is because the loss basin is extremely wide and shallow, so the global unconstrained minimum is actually very far from where the model actually ends up via gradient descent in finite time.
• Our insight is to change the approach: we say “Let us first approximate the loss landscape itself; then, we can account for the rank constraint throughout the entire optimization process.” With the quartic approximation of the loss, this turns out to be solvable. As a result, our prediction for what word2vec learns is significantly more accurate. We discuss this at the end of Section 2. We will include more details to clarify this point, as it is central to our result.

Why does decreasing SNR sometimes still yield high performance? Although the relationship between SNR and task performance is qualitative rather than quantitative, it’s still quite interesting that the relative ordering of analogy categories is preserved. In other words, SNR (which is computed from the corpus statistics and hyperparameters alone) is a useful measure of how difficult a particular class of analogies is to learn. To our knowledge, this is the first time that a measurable a priori proxy for analogy difficulty has been proposed.

This still leaves the question of why, within an analogy category, larger dimension yields better performance despite the degradation in SNR. Our leading hypothesis (corroborated by several empirical checks we did) is that this is an ancillary effect of the benchmark using top-1 accuracy as the error metric. As a concrete example, consider the analogy “France : Paris :: Japan : Tokyo”, satisfying japan+(paris-france) ≈ tokyo. As the effective embedding dimension increases, the SNR may decrease; at the same time, all the embeddings become more spaced out (due to the increased available volume in latent space), meaning there are fewer nearby “competitors” for tokyo. In the large $d$ regime, we empirically found that the embeddings are sufficiently spaced, and tokyo is still the nearest embedding to japan+(paris-france) despite an increase in absolute error. Thus, the measured top-1 performance does not degrade with the increased (relative) noise. Note that unintuitive effects associated with using top-k accuracy have been observed in LLMs as well (see https://arxiv.org/abs/2304.15004). We plan to include this discussion in the final version of the manuscript.

Thanks for the thoughtful comments on our manuscript. Let us know if you have further questions.

Thanks very much for your review, we appreciate it. We’re glad you found the result clear and convincing. Some comments regarding your concerns and questions:

What’s the point of this analysis? Isn’t word2vec outdated? Though the algorithm itself has now been superseded by larger context-aware embeddings, the core mechanistic idea behind word2vec is back in vogue! Many recent papers empirically report linear structure in the latent representations in LLMs – a direct analog of the word2vec embedding geometry. (For example, see https://arxiv.org/abs/2305.16130, https://arxiv.org/pdf/2303.04245, and the references we mention in Section 4).

Trying to understand this behavior directly in large language models is an admirable goal to strive towards. However, we believe it is too difficult to tackle directly (for now). Previous efforts in this direction have not gotten very far. A useful way for science to proceed in this situation is to focus on first understanding the mysterious behavior in the simplest possible system that exhibits it. Doing this is one of our primary contributions.

In general, one of the main outstanding questions of deep learning theory is: what is the mechanism and character of representation learning? Theorists have very few solvable cases, and these analyses tend to rely on unrealistic assumptions on either the data or the learning task (or both). We believe our result is powerful because it answers this question in a relevant natural language task under very few, relatively weak assumptions.

(Also, it’s a bit surprising to lack a theory of learning dynamics for one of the foundational algorithms of the field! Even if it weren’t directly connected to understanding modern models, for posterity’s sake we should try to solve and deeply understand the small handful of truly breakthrough algorithms in the history of neural learning.)

What about GloVe and FastText? We expect the main idea of our proofs to hold for both these models. FastText is essentially the same architecture and loss as word2vec, and its extra engineering tricks will not change the core story much. GloVe is very similar to word2vec; while word2vec uses a contrastive loss, GloVe instead performs explicit least-squares factorization. Incorporating something like our Setting 3.1 in the analysis of the GloVe algorithm would yield results exactly analogous to the ones presented here. We plan to include this discussion and some supporting experiments in the final version of the manuscript.

How can I apply these results to other self-supervised learning tasks? Since our results are distribution agnostic, they carry over to pretty much any linearized contrastive algorithm. To be specific, if your contrastive loss has the form $\\mathcal{L}(W) = \\mathbb{E}\_\\text{positive pairs}[f^{+}(w^T w’)] + \\mathbb{E}\_\\text{negative pairs}[f^{-}(w^T w’)]$, then you can simply take the quadratic Taylor approximation of your $f^{+}(w^Tw’)$ and $f^{-}(w^T w’)$, complete the square, and obtain a target matrix in terms of the positive and negative distributions for your learning problem of interest. In this sense, the idea behind our Theorem 1 is actually very general. See Appendix D.2 for an example using the SimCLR loss. To get closed-form learning dynamics, you’d need to find an equivalent of our Setting 3.1. This may be difficult depending on the form of your input data; for instance, it may require a whitening data preprocessing step. The idea is to choose these hyperparameters in such a way that the weighted matrix factorization becomes unweighted. For more details, see Proposition 2 and the subsequent discussion.

Let us know if you have further questions.

Thank you for your response! For completeness, here we include some follow-up.

We are curious as to which aspect you found unsatisfying. (This feedback would be useful to us when we discuss our result with other colleagues.) It seems to us that there are several "kinds" of analysis one could hope for regarding LLMs:

1. Direct analysis of full optimization dynamics under full transformer architecture and under no assumptions on the data distribution. To our knowledge, this is very difficult and the community does not yet know how to do this.
2. Analysis of convergence (without characterizing the optimization dynamics) under simplified transformer architecture.
   • https://arxiv.org/pdf/2308.16898
   • https://arxiv.org/pdf/2402.05738
3. Analysis of some aspects of optimization dynamics under simplified 1-layer transformer architecture and restrictive synthetic data distribution.
   • https://arxiv.org/pdf/2411.10830
   • https://arxiv.org/pdf/2305.16380
   • https://arxiv.org/abs/2410.09605
   • https://arxiv.org/pdf/2504.12916
4. Analysis of some aspects of final learned representations under simplified 1-layer transformer architecture and restrictive synthetic data distribution.
   • See many works studying algorithms implemented by transformers for ICL
   • https://arxiv.org/pdf/2403.03867
5. Analysis of full optimization dynamics under very simplified attention-like architecture and arbitrary natural data distribution.
   • Our current work. (See end of this comment for an argument that word2vec dynamics are derivable from linear transformer dynamics.)

Thus, there are many avenues to making progress in understanding representation learning in LLMs. We simply argue that ours is a promising but under-explored one. We note that, to our knowledge, theoretically understanding the intermediate representations of a large language model is still a very open problem.

We emphasize that understanding LLMs is not a primary goal of this paper. We frame our results purely in terms of understanding word2vec. We believe this alone carries sufficient merit. However, we do believe that our results may provide insight in this broader research project of understanding LLMs.

Heuristic argument that word2vec dynamics are closely related to transformer learning dynamics. Consider a next-token prediction task being learned by a single-layer linear attention transformer with no positional encoding and weight tying. The feedforward output is

$$\mathrm{logits}(X) = W ((X W Q K^\top W^\top X^\top) X W V)^\top$$

where $X\in\mathbb{R}^{L\times V}$ is a prompt matrix with one-hot tokens, $W \in\mathbb{R}^{V\times d}$ is the embedding layer (and unembedding, by weight tying), and $Q,K,V$ are the query/key/value weights. The relevant dimensions are $d$ for the embedding dim, $V$ for the vocabulary size, and $L$ for the prompt length.

Let us assume that the Q, K weights are shared (commonly done in practice, see https://arxiv.org/pdf/2001.04451). We will further assume that Q and K are wide and initialized i.i.d. Gaussian with large variance $\sigma^2$. Then we have that $QK^\top \approx \sigma^2 \mathbb{I}\_{d}$ from concentration of Marchenko-Pastur spectrum in the wide limit, and $\frac{d}{dt} (QK^\top) \approx 0$ (which follows from directly computing gradients). Finally, let us assume for simplicity that we fix $V = \mathbb{I}$.

Now things simplify a lot. At initialization we have

$$\mathrm{logits}(X) \approx W ((X W W^\top X^\top) X W)^\top = W (W\_{[L]} W^\top\_{[L]} W\_{[L]})^\top$$

where $W\_{[L]} \in \mathbb{R}^{L \times d}$ are the embeddings of the prompt words. Focusing on the final-token logits (which predict the next word)

$$\mathrm{finallogits}(X) \approx W W^\top\_{[L]} W\_{[L]} w\_{final}$$

Computing the gradient of the logit w.r.t. the $i$th word in the context, we find that

$$\frac{d \mathrm{finallogits}}{d w\_i} \approx W (w\_i^\top w\_{final}\mathbb{I}\_d + w\_i w\_{final} ^\top)$$

Finally, we recall that the standard cross-entropy loss partitions the vocabulary into positive (true next word) and negative (all other words) distributions. Computing the gradient wrt the loss, we find that the embeddings roughly evolve according to the following rule:

• words within the prompt $w\_i$ evolve to align towards the true next word.
• Simultaneously, all the $w\_i$ are pushed away from (the softmax-weighted mean of) all other words in the vocabulary.

This qualitatively recovers the word2vec update rule.

Thank you very much for your review! We’re glad you found the paper readable. Thank you for your suggestions regarding clarity – they are very useful to us and we’ll make the requisite changes. Some comments follow regarding your other questions.

Does early mixing and delayed convergence imply that QWEM promotes sparsity? Both QWEMs and word2vec exhibit a low-rank bias in their learning dynamics, although the effect is more pronounced in QWEMs, as you observe. We don’t believe that this low-rank bias is necessarily good or bad; that depends on the downstream task of interest. Our theory is simply descriptive of the learning dynamics.

The reason that clear stepwise learning occurs in QWEMs is that our Setting 3.1 turns the weighted factorization problem into an unweighted factorization – as a consequence, the singular value dynamics become decoupled (i.e., “untangled”). Without our Setting 3.1, there is “mixing” between the singular directions (as you observe), even when the corresponding singular values have grown to be $O(1)$ in magnitude. Exactly characterizing this mixing is NP hard (see Proposition 2 and the subsequent discussion). The somewhat surprising fact is that word2vec is “close enough” to the unweighted factorization – though the singular value dynamics do exhibit some early mixing, the overall learning dynamics are well-described by the simplified setting. This agreement is partly because the hyperparameters chosen in word2vec coarsely approximate our Setting 3.1 (see Appendix A.3).

As for the slightly delayed convergence, we don’t think this is an important effect. The two algorithms optimize different objectives, so the magnitudes of the gradients are not exactly the same; one would not expect their learning timescales to match exactly. Our theory predicts the convergence time with relatively small error – this is already much better than previous results, which only give the scaling behavior (or loose bounds) for the convergence time, rather than a direct and accurate prediction.

Shouldn't the task vectors be exactly equal rather than rank-1? Strictly speaking, for exact analogy recovery, you’re right that rank-1 task vectors are necessary but not sufficient. In practice though, when performance is measured with top-1 accuracy, the length of the task vectors is rather unimportant in the high-dimensional regime. The angular information is far more relevant. (One can intuitively see this by considering the toy setting of Gaussian random vectors in the high-dimensional limit – their lengths concentrate and only angles matter.) This is especially true since the embeddings are normalized to be unit norm before evaluating the analogy task. We discuss this subtlety in Appendices A.4 and A.6. We will clarify this in the main text.

Let us know if you have further questions!

Thank you very much for your review! We’re glad you found the paper clear and interesting. Some comments follow.

Why does decreasing SNR sometimes still yield high performance? This is a very interesting question and we spent some time exploring this. Our leading hypothesis (corroborated by several empirical checks we did) is that this is an ancillary effect of the benchmark using top-1 accuracy as the error metric. As a concrete example, consider the analogy “France : Paris :: Japan : Tokyo”, satisfying japan+(paris-france) ≈ tokyo. As the effective embedding dimension increases, the SNR may decrease; at the same time, all the embeddings become more spaced out (due to the increased available volume in latent space), meaning there are fewer nearby “competitors” for tokyo. In the large $d$ regime, we empirically found that the embeddings are sufficiently spaced, and tokyo is still the nearest embedding to japan+(paris-france) despite an increase in absolute error. Thus, the measured top-1 performance does not degrade with the increased (relative) noise. Note that unintuitive effects associated with using top-k accuracy have been observed in LLMs as well (see https://arxiv.org/abs/2304.15004). We plan to include this discussion in the final version of the manuscript.

Where is the clear, discrete, rank-incrementing behavior in bona fide word2vec? You correctly observe that the rank-incrementing steps in word2vec do not appear clearly separated (especially early in training). The reason that clear stepwise learning occurs in QWEMs is that our Setting 3.1 turns the weighted factorization problem into an unweighted factorization – as a consequence, the singular value dynamics become decoupled (i.e., “untangled”). Without our Setting 3.1, there is “mixing” between the singular directions, even when the corresponding singular values have grown to be $O(1)$ in magnitude. Exactly characterizing this mixing is NP hard (see Proposition 2 and the subsequent discussion). The somewhat surprising fact is that word2vec is “close enough” to the unweighted factorization – though the singular value dynamics do exhibit some mixing (as you observe), the overall learning dynamics are well-described by the simplified setting. This agreement is partly because the hyperparameters chosen in word2vec coarsely approximate our Setting 3.1 (see Appendix A.3).

Regarding the link between Section 4 and the theoretical results. Though the empirical results in Section 4 do not rely on the correctness of the preceding theory, the two are related. Our first result (left and top right of Fig. 4) suggests that it may be possible to obtain a sparse effective theory of semantic linear representations by studying their decomposition in terms of the semantic basis identified by our theory. Our second result (bottom right) identifies a potential benefit of early stopping in large models – it may prevent the realization of noisy semantic directions in the latents. This insight follows from the duality between model size and training time identified by our theory (i.e., the low-rank bias). In general, all the empirics in Section 4 were relatively easy to iterate on because they only involved factorizing a large matrix once, rather than re-training many word embedding models. This computational saving follows directly from the theory (we discuss it in Section 3.3).

Let us know if you have further questions!