What exactly does word2vec learn?

What exactly does word2vec learn, and how? Answering this question amounts to understanding representation learning in a minimal yet interesting language modeling task. Despite the fact that word2vec is a well-known precursor to modern language models, for many years, researchers lacked a quantitative and predictive theory describing its learning process. In our new paper, we finally provide such a theory.

We prove that there are realistic, practical regimes in which the learning problem reduces to unweighted least-squares matrix factorization. We solve the gradient flow dynamics in closed form; the final learned representations are simply given by PCA. Learning dynamics of word2vec. When trained from small initialization, word2vec learns in discrete, sequential steps. Left: rank-incrementing learning steps in the weight matrix, each decreasing the loss.

Right: three time slices of the latent embedding space showing how embedding vectors expand into subspaces of increasing dimension at each learning step, continuing until model capacity is saturated. Before elaborating on this result, let’s motivate the problem. word2vec is a well-known algorithm for learning dense vector representations of words.

These embedding vectors are trained using a contrastive algorithm; at the end of training, the semantic relation between any two words is captured by the angle between the corresponding embeddings. In fact, the learned embeddings empirically exhibit striking linear structure in their geometry: linear subspaces in the latent space often encode interpretable concepts such as gender, verb tense, or dialect.

This so-called linear representation hypothesis has recently garnered a lot of attention since LLMs exhibit this behavior as well, enabling semantic inspection of internal representations and providing for novel model steering techniques. In word2vec, it is precisely these linear directions that enable the learned embeddings to complete analogies (e. g. , “man: woman:: king: queen”) via embedding vector addition.

Maybe this shouldn’t be too surprising: after all, the word2vec algorithm simply iterates through a text corpus and trains a two-layer linear network to model statistical regularities in natural language using self-supervised gradient descent. In this framing, it’s clear that word2vec is a minimal neural language model. Understanding word2vec is thus a prerequisite to understanding feature learning in more sophisticated language modeling tasks.

The Result With this motivation in mind, let’s describe the main result. Concretely, suppose we initialize all the embedding vectors randomly and very close to the origin, so that they’re effectively zero-dimensional. Then (under some mild approximations) the embeddings collectively learn one “concept” (i. e. , orthogonal linear subspace) at a time in a sequence of discrete learning steps. It’s like when diving head-first into learning a new branch of math.

At first, all the jargon is muddled — what’s the difference between a function and a functional? What about a linear operator vs. a matrix? Slowly, through exposure to new settings of interest, the words separate from each other in the mind and their true meanings become clearer. As a consequence, each new realized linear concept effectively increments the rank of the embedding matrix, giving each word embedding more space to better express itself and its meaning.

Since these linear subspaces do not rotate once they’re learned, these are effectively the model’s learned features. Our theory allows us to compute each of these features a priori in closed form – they are simply the eigenvectors of a particular target matrix which is defined solely in terms of measurable corpus statistics and algorithmic hyperparameters. What are the features?

The answer is remarkably straightforward: the latent features are simply the top eigenvectors of the following matrix: where $i$ and $j$ index the words in the vocabulary, $P(i,j)$ is the…