In §1.4 we got embeddings by counting contexts and compressing the counts. Now we take the other road. We predict contexts with a tiny neural model and let the embeddings fall out as its parameters. This is word2vec (Mikolov et al., 2013a,b), one of the most influential small models in the history of NLP. It's worth knowing down to the last gradient, because its ingredients reappear everywhere. They include a softmax over a huge vocabulary, contrastive "real versus fake" training, negative sampling, and the dot product as a learned compatibility score.

By the end of this section you'll have derived every gradient by hand, implemented the training step and the samplers that interviewers love to ask for, trained a model in your browser, and seen why this "neural" method secretly computes the same thing as §1.4.

From counting to predicting

The first neural word embeddings came from language modeling. Bengio et al. (2003) trained a network to predict the next word from the previous few. It looked up their embeddings, concatenated them, applied an MLP, and finished with a softmax over the vocabulary. The embedding table was trained jointly, and similar words ended up with similar vectors because that helped prediction. (Andrej Karpathy rebuilds exactly this model in makemore part 2, linked from §1.2.)

Mikolov's insight was that if you only want the embeddings, you can throw away almost all of the model. Make the prediction task as simple as possible, with no hidden layer and just a dot product. Then train on billions of words.

Skip-gram predicts the neighbors

The skip-gram task is simple. Given a center word $w_t$ at position $t$, predict each word within $c$ positions of it (here $c$ is the window half-width).

... the  quick  [brown]  fox  jumps ...        center = brown, c = 2
          ↑       ↑       ↑     ↑
  pairs: (brown, the) (brown, quick) (brown, fox) (brown, jumps)

Every word gets two vectors, $\mathbf{v}_w \in \mathbb{R}^d$ when it is the center (row $w$ of the input table $\mathbf{E}_{\text{in}}$) and $\mathbf{u}_c \in \mathbb{R}^d$ when it is a context (row $c$ of the output table $\mathbf{E}_{\text{out}}$). From here on, $c$ inside a subscript names a context word. The model says that the probability of context $o$ given center $w$ is a softmax over dot products,

In plain words, score every possible context word by how well its output vector aligns with the center's input vector, then turn the scores into probabilities.

Training maximizes the average log-probability of all observed (center, context) pairs in a corpus of $T$ tokens,

where $\theta$ is both tables. In plain words, reward the model for giving high probability to the words that actually appeared around each center word.

The full-softmax gradient is expected minus observed

Write the loss for one pair as $\mathcal{L} = -\log p(o\mid w) = -\mathbf{u}_o^\top\mathbf{v}_w + \log\sum_{c}\exp(\mathbf{u}_c^\top\mathbf{v}_w)$. Differentiating gives

The update pulls $\mathbf{v}_w$ toward the observed context's vector and away from the average of what the model currently predicts. When the prediction is right, the two terms cancel and learning stops.

import numpy as np

def softmax_sg(v, U, o):
    """returns loss (float), dv (d,), dU (V, d)"""
    raise NotImplementedError

The problem is that sum over the vocabulary

Look at the denominator again, $\sum_{c\in\mathcal V}\exp(\mathbf{u}_c^\top\mathbf{v}_w)$. Every single training pair touches every output vector. With $V = 10^6$ and $d = 300$, that is $3\times10^8$ multiply-adds for one pair, and a corpus of a billion tokens with a window of 5 produces ten billion pairs. Word2vec offers two escapes, hierarchical softmax and negative sampling. Negative sampling is the one that won, and its core idea (contrast the real thing against sampled fakes) is the ancestor of the contrastive losses behind modern embedding models and joint image and text spaces.

Negative sampling asks "real or fake?" instead of "which word?"

Change the question. Instead of "which of the $V$ words is the context?" (a $V$-way classification), ask "is this (center, context) pair real, or did I make it up?" That's a binary question, answered with a sigmoid,

For each real pair $(w, o)$ from the corpus, we draw $k$ negative words $n_1,\dots,n_k$ at random from a noise distribution $P_n$. Those pairs are (almost certainly) fake. The loss for that single example is

In plain words, make the real pair's score high, and make each fake pair's score low. (Note $\sigma(-x) = 1-\sigma(x)$, the probability of "fake".) The cost per pair is now $(k+1)d$ multiply-adds instead of $Vd$. With $k = 5$ and $d = 300$ that is 1,800 instead of 300 million, a factor of about 170,000.

The gradients

Using $\frac{d}{dx}\log\sigma(x) = 1-\sigma(x)$ and $\frac{d}{dx}\log\sigma(-x) = -\sigma(x)$, we get

Each coefficient is an error signal. $\sigma(s_o)-1 \in (-1, 0)$ is how far the real pair falls short of "definitely real". $\sigma(s_{n_i}) \in (0,1)$ is how much the model wrongly believes a fake pair. When the model is right, the errors vanish.

One update, by hand

Numbers make this stick. Take $d = 2$, one negative, and

The scores are $s_o = 0.15 - 0.08 = 0.07$ and $s_n = -0.05 - 0.12 = -0.17$, so $\sigma(s_o) = 0.5175$ and $\sigma(s_n) = 0.4576$. The loss is $-\log 0.5175 - \log(1 - 0.4576) = 0.6587 + 0.6118 = 1.2705$. The center gradient is

The output gradients are $(0.5175-1)\,\mathbf{v}_w = (-0.2413, 0.0965)$ for the positive and $0.4576\,\mathbf{v}_w = (0.2288, -0.0915)$ for the negative. Take one step with learning rate $\eta = 0.5$, updating all three vectors from these gradients. The positive score rises to 0.166, the negative score falls to $-0.283$, and the loss drops to 1.175.

Notice one more number. word2vec initializes the output table to zeros, so every score starts at 0 and $\sigma(0) = \tfrac12$. The initial loss is therefore exactly $(k+1)\log 2$, which is $6\log 2 \approx 4.159$ for $k = 5$. If your training curve doesn't start there, your initialization is not what you think it is.

import numpy as np

def sgns(v, u_pos, U_neg):
    """returns loss, dv (d,), du_pos (d,), dU_neg (k, d)"""
    raise NotImplementedError

Choosing the negatives with the $3/4$ power

Which words should be the "fakes"? Uniform sampling would mostly pick obscure words, which are too easy to reject. Sampling by raw frequency would mostly pick the, of, and. Mikolov et al. (2013b) found empirically that a compromise works best,

where $f(w)$ is the word's relative frequency in the corpus. In plain words, raising frequencies to the power $3/4$ flattens the distribution, so frequent words are sampled less than their frequency and rare words more. Two words with frequencies $0.9$ and $0.1$ get noise probabilities of about $0.84$ and $0.16$. Same order, gentler gap. It's the same flattening exponent we used for PPMI context smoothing in §1.4, and Levy, Goldberg and Dagan (2015) showed that the connection is not a coincidence (the derivation at the end of this section makes it exact).

How many negatives? The paper recommends $k = 5$ to 20 for small datasets and $k = 2$ to 5 for very large ones.

import numpy as np

def noise_dist(counts, power=0.75):
    raise NotImplementedError

def sample_negatives(P, shape, rng):
    raise NotImplementedError

Constant-time sampling with an alias table

Billions of negatives call for something faster than a binary search. Word2vec's C code precomputes a table of $10^8$ entries in which each word appears in proportion to $P_n$, so a sample is one random index into the table. That costs 400 MB of int32 and quantizes the probabilities.

The alias method (Walker, 1977; Vose, 1991) gets exact $O(1)$ sampling with only two arrays of length $V$. Picture $V$ columns of equal capacity 1, one per item, and scale every probability by $V$, so the heights $q_i = V p_i$ average exactly 1. Some columns overflow ($q_i > 1$) and some fall short ($q_i < 1$). Repeatedly take a short column $s$ and an overflowing column $\ell$, keep $q_s$ of column $s$ for itself, and fill its remaining $1 - q_s$ with mass taken from $\ell$, recording $\ell$ as $s$'s alias. Column $\ell$ shrinks by $1 - q_s$, and goes back on the short or overflow list. When the lists empty, every column is exactly full and holds at most two items. To sample, pick a column uniformly and flip a biased coin to choose between the column's own item and its alias.

A worked example with $P = (0.5, 0.3, 0.15, 0.05)$. Scaled by 4, the heights are $(2.0, 1.2, 0.6, 0.2)$. Item 3 (height 0.2) borrows 0.8 from item 0, leaving item 0 at 1.2. Item 2 (0.6) borrows 0.4 from item 1, leaving item 1 at 0.8, which is now short. Item 1 borrows 0.2 from item 0, leaving item 0 at exactly 1. The final table keeps probability $(1.0, 0.8, 0.6, 0.2)$ in place with aliases $(\cdot, 0, 1, 0)$. Check item 0. It gets the whole of column 0, plus $0.2$ of column 1 and $0.8$ of column 3, so $\tfrac14(1 + 0.2 + 0.8) = 0.5$. Alias tables are not unique. Pairing short and tall columns in a different order gives a different table with the same implied distribution, which is why the tests below check the distribution rather than the table.

import numpy as np

def build_alias(P):
    """-> (prob, alias): prob is float (V,), alias is int (V,)"""
    raise NotImplementedError

def alias_sample(prob, alias, shape, rng):
    raise NotImplementedError

Two more tricks that matter a lot

Subsampling frequent words. Words like the appear in millions of pairs, but "brown co-occurs with the" teaches almost nothing. Word2vec randomly discards each occurrence of word $w$ before training, with probability

where $t$ is the subsampling threshold. In plain words, a word is kept with probability $\sqrt{t/f(w)}$, so the more frequent it is, the more of its occurrences are thrown away. Some numbers with $t = 10^{-5}$. A word as common as the ($f \approx 0.05$, one token in twenty) is kept only 1.4% of the time. A word with $f = 10^{-4}$ is kept 32% of the time. Any word rarer than $t$ is always kept. This speeds training up a lot and improves the vectors of rarer words. It also effectively widens the window, because deleting the brings meaningful words closer together. (The released C code uses a slightly different formula with the same spirit, and its defaults differ from the paper too. It trains CBOW unless told otherwise, with a subsampling threshold of $10^{-3}$, a window of 5, and 5 negatives. Check the defaults of any implementation you use.)

Dynamic window. For each center word, the actual window size $b_t$ is drawn uniformly from $1$ to $c$. A context word at distance $j$ is therefore used with probability $(c-j+1)/c$, so nearer words count more. With $c = 5$, distances 1 through 5 are used with probabilities 1, 0.8, 0.6, 0.4, 0.2. GloVe gets a similar effect by weighting co-occurrence counts by $1/\text{distance}$.

import numpy as np

def skipgram_pairs(ids, b):
    """ids: (T,), b: (T,) window per center -> (centers, contexts)"""
    raise NotImplementedError

Implement the training step (the interview question)

Here is the version you'd write at a whiteboard or in a coding interview. It's a mini-batch SGNS step, fully vectorized, updating both tables in place. The crucial detail is the same one as in §1.2. The same word appears many times in a batch, so updates must be scatter-added.

import numpy as np

def sgns_step(Vin, Uout, w, c, neg, lr):
    raise NotImplementedError

Train word2vec in your browser

Now we put everything together, a synthetic corpus, subsampling, pair extraction (with a fixed window, for simplicity), the $f^{3/4}$ noise distribution, linear learning-rate decay (as in the original code), and the vectorized step. It runs in a few seconds.

The words never saw their category labels. They clustered because they share contexts. Try changing k, d, or the window, or set t so that nothing is subsampled, and compare. (Without subsampling, the, and, and "." dominate the batches, and because a word that appears 60 times in a batch receives the sum of 60 gradients at once, the learning rate that worked above can make training diverge.)

import numpy as np

def train_sgns(ids, V, d=16, k=5, window=2, epochs=3, lr0=0.05, batch=256, t=1e-2, seed=0):
    """-> (Vin of shape (V, d), list of per-batch losses)"""
    raise NotImplementedError

CBOW, hierarchical softmax, and the rest of the family

CBOW (continuous bag of words) flips skip-gram around. It averages the context vectors and predicts the center word. The math is identical once you replace $\mathbf{v}_w$ with the average $\bar{\mathbf{v}} = \frac{1}{2c}\sum_{0 < \lvert j\rvert \le c} \mathbf{v}_{w_{t+j}}$ (using the input table for contexts) and predict $w_t$ with the output table. It's faster, since there's one prediction per window instead of $2c$, and tends to do slightly better for frequent words. Skip-gram tends to do better for rare words and small datasets (per the word2vec authors' notes). Averaging also throws away word order inside the window, and blurs rare context words into the crowd.

GloVe regresses log counts

GloVe (Pennington, Socher & Manning, 2014) starts from the co-occurrence matrix $\mathbf{X}$ of §1.4 (GloVe's name for $\mathbf{C}$) and asks the vectors to explain the log counts directly,

with word vectors $\mathbf{w}_i$, context vectors $\tilde{\mathbf{w}}_j$, scalar biases $b_i, \tilde b_j$, $\alpha = 3/4$, and $x_{\max} = 100$. In plain words, it is a weighted least-squares fit that makes each word–context dot product (plus two biases) equal the log of how often they co-occurred. The weighting $f$ keeps very rare pairs (noisy) from mattering much and caps very frequent pairs (like of the), so a pair seen 10 times gets weight $0.1^{0.75} \approx 0.18$ and a pair seen 100 or more times gets weight 1.

The motivation is a nice piece of reasoning. Ratios of co-occurrence probabilities carry meaning. In the paper's 6-billion-token corpus, $P(\text{solid}\mid\text{ice}) / P(\text{solid}\mid\text{steam}) = 8.9$ and $P(\text{gas}\mid\text{ice})/P(\text{gas}\mid\text{steam}) = 0.085$, while for words related to both or neither the ratio is near 1 ($1.36$ for water, $0.96$ for fashion). If we want a function of vector differences to reproduce such ratios, $F\big((\mathbf{w}_i - \mathbf{w}_j)^\top\tilde{\mathbf{w}}_k\big) = P_{ik}/P_{jk}$, then $F = \exp$ is the natural choice, because it turns differences into ratios. That forces $\mathbf{w}_i^\top\tilde{\mathbf{w}}_k = \log P_{ik} = \log X_{ik} - \log X_i$, and absorbing $\log X_i$ into a bias (and adding a context bias for symmetry) gives the objective above. GloVe is a weighted matrix factorization that combines the global statistics of the counting approach with the trainable, scalable machinery of the predictive one. The final vectors are usually $\mathbf{W} + \tilde{\mathbf{W}}$.

fastText says words are made of pieces

Word2vec gives each word an atomic vector, so it knows nothing about a word it never saw, and nothing about the shared structure of run / runs / running. fastText (Bojanowski et al., 2017) represents a word as the sum of its character n-gram vectors (n from 3 to 6, with boundary markers) plus a vector for the whole word. For where with $n = 3$, the n-grams are <wh, whe, her, ere, re>, plus the special sequence <where>. Then skip-gram with negative sampling trains exactly as before, with $\mathbf{v}_w = \sum_{g\in G(w)}\mathbf{z}_g$, where $G(w)$ is the set of the word's n-grams and $\mathbf{z}_g$ the vector of n-gram $g$. To bound memory, the n-grams are hashed into 2 million buckets, so unrelated n-grams occasionally share a vector. Unseen words still get vectors (sum their n-grams), and morphology is shared. This idea of building words from subword pieces is exactly where Chapter 2 picks up.

The big reveal, SGNS is factorizing a PMI matrix

Why do count-based (§1.4) and predict-based methods end up so similar? Levy and Goldberg (2014) proved a beautiful result. If the dimension $d$ is large enough that every dot product can be set independently, then the SGNS objective is maximized exactly when

(with PMI computed using the noise distribution for the context marginal, so with $P_n\propto f^{3/4}$, this is the smoothed PMI from §1.4). In plain words, word2vec implicitly factorizes the PMI matrix, shifted down by $\log k$, into $\mathbf{E}_{\text{in}}\mathbf{E}_{\text{out}}^\top$. In practice $d$ is smaller, so SGNS finds a low-rank weighted approximation of it, much like the truncated SVD of §1.4 but with a different weighting of errors (frequent pairs count more, because they appear in more gradient steps).

You can check the one-cell claim numerically in a few lines. Pick counts, maximize the per-cell objective over a fine grid, and compare with the formula.

This connects the whole chapter. Counting and predicting are two routes to factorizing the same matrix. It also explains why the two approaches' tricks transfer (the $3/4$ smoothing, subsampling, and so on), and why $k$ acts like a shift that pushes rare, weakly associated pairs below zero. It also shows that $\mathbf{u}_c^\top\mathbf{v}_w$ is not a log-probability, since nothing forces $\sum_c e^{\mathbf{u}_c^\top\mathbf{v}_w}$ to equal 1, which is the NCE point from earlier made concrete.

import numpy as np

def shifted_pmi(C, k):
    raise NotImplementedError

In PyTorch

The production version uses two nn.Embedding tables with sparse gradients, logsigmoid for stability, and SparseAdam. On a laptop or a single GPU it trains on the text8 corpus (about 17 million tokens of Wikipedia) in minutes.

Read more

• Mikolov et al. (2013a). Efficient Estimation of Word Representations in Vector Space. CBOW and skip-gram.
• Mikolov et al. (2013b). Distributed Representations of Words and Phrases and their Compositionality. Negative sampling (§2.2) and subsampling (§2.3).
• Rong (2014). word2vec Parameter Learning Explained and Goldberg & Levy (2014), word2vec Explained. Every gradient, derived.
• Dyer (2014). Notes on Noise Contrastive Estimation and Negative Sampling. Exactly how negative sampling differs from NCE, in four pages.
• Pennington, Socher & Manning (2014). GloVe. Weighted least squares on log counts, and the co-occurrence ratio argument.
• Bojanowski et al. (2017). Enriching Word Vectors with Subword Information. fastText.
• Levy & Goldberg (2014). Neural Word Embedding as Implicit Matrix Factorization. SGNS equals shifted PMI.
• Zhang et al. Dive into Deep Learning, §15.1 to §15.7 and Lilian Weng, Learning Word Embedding. Two excellent complementary walkthroughs.

Watch next

▶

▶

▶