What word2vec learns: A closed-form theory of feature learning

BAIR researchers have provided a quantitative and predictive theory explaining the learning process of word2vec, proving that under specific practical regimes, the learning problem reduces to unweighted least-squares matrix factorization. The study reveals that the final learned representations are given by Principal Component Analysis (PCA) of a target matrix derived from corpus statistics.

Word2vec as a Minimal Language Model

word2vec serves as a foundational algorithm for learning dense vector representations of words through a contrastive algorithm. The resulting embeddings capture semantic relations via the angle between vectors, often exhibiting a linear structure where subspaces encode interpretable concepts like gender, verb tense, or dialect. This "linear representation hypothesis" is a key driver for the model's ability to complete analogies (e.g., "man : woman :: king : queen") and is a behavior also observed in modern Large Language Models (LLMs).

Because word2vec is essentially a two-layer linear network trained via self-supervised gradient descent to model statistical regularities, it acts as a minimal neural language model. Understanding its feature learning is considered a prerequisite for understanding more complex language modeling tasks.

The Theory of Discrete Learning Steps

The research proves that when embedding vectors are initialized randomly and very close to the origin, the model learns "concepts" (orthogonal linear subspaces) one at a time in a sequence of discrete steps.

Rank-Incrementing Dynamics

Each new realized linear concept increments the rank of the embedding matrix. This process allows word embeddings to expand into subspaces of increasing dimension, providing more space to better express meaning. Because these linear subspaces do not rotate once learned, they function as the model's learned features.

Closed-Form Feature Computation

These features can be computed a priori in closed form as the eigenvectors of a target matrix $M^*$, defined by:

$$M^*{i,j} = \frac{P(i,j) - P(i)P(j)}{\sqrt{2(P(i,j) + P(i)P(j))}}$$

Where:

  • $i$ and $j$ index the words in the vocabulary.
  • $P(i,j)$ is the co-occurrence probability for words $i$ and $j$.
  • $P(i)$ is the unigram probability for word $i$.

Applying this to Wikipedia statistics shows that the top eigenvectors correspond to interpretable topic-level concepts, such as celebrity biographies, government administration, and geographical descriptors.

Theoretical Approximations and Empirical Validation

To achieve these closed-form solutions, the researchers applied four mild approximations:

  1. A quartic approximation of the objective function around the origin.
  2. Specific constraints on algorithmic hyperparameters.
  3. Sufficiently small initial embedding weights.
  4. Vanishingly small gradient descent steps.

Crucially, these approximations do not involve the data distribution, making the theory distribution-agnostic.

Performance Comparison

The theory's validity is supported by empirical results on the standard analogy completion benchmark:

  • Original word2vec: 68% accuracy
  • Approximate model (studied in theory): 66% accuracy
  • Standard classical alternative (PPMI): 51% accuracy

Implications for Abstract Representations

The theory was used to analyze the emergence of abstract linear representations (e.g., binary concepts like masculine/feminine). The researchers found that word2vec builds these representations in a sequence of noisy learning steps, with geometry described by a spiked random matrix model. While semantic signal dominates early in training, noise can dominate later in training, potentially degrading the model's ability to resolve these linear representations.

Sources