Word2Vec, in depth: skip-gram, negative sampling and its gradients, subsampling, a tested NumPy trainer and the PMI view

Predicting context: skip-gram and CBOW

The distributional hypothesis says a word is characterised by the company it keeps. Word2Vec makes that operational with a prediction task. Slide a window over the text. In skip-gram, the word in the middle (the center) predicts each word within the window (the contexts). In CBOW (continuous bag of words), the average of the context vectors predicts the center. Both use two parameter matrices of shape V x d: W_in holds one vector per word when it plays the center, and W_out holds one vector per word when it plays a context.

The model's score for a pair is a dot product, u_c . v_w. With a full softmax, the probability of context c given center w is exp(u_c . v_w) / sum over all V words of exp(u_j . v_w). The denominator touches every word in the vocabulary for every training pair, so with V = 1,000,000 and d = 300 a single pair costs 300 million multiply-adds. That cost is what the two approximations below remove.

CBOW makes one prediction per window position, so it is several times faster; skip-gram makes one per (center, context) pair, gives rare words more updates, and did better on rare words in the original papers. This article uses skip-gram.

Avoiding the softmax: negative sampling

Hierarchical softmax arranges the vocabulary as the leaves of a binary tree and replaces one V-way choice with about log2 V binary choices along the path to the word. The original tool builds a Huffman tree, so frequent words get short paths; the Huffman coding article shows the construction. It is exact in the sense that probabilities still sum to one, but its cost is per tree node, and it tends to help infrequent words.

Negative sampling (SGNS when combined with skip-gram) drops the normalised probability entirely. For each observed pair (w, c) it draws k noise words n_1..n_k and trains a logistic classifier to tell real pairs from noise. The objective for one pair, to be maximised, is:

log sigmoid(u_c . v_w) + sum_{i=1..k} log sigmoid(-u_{n_i} . v_w)

Differentiate with a label y = 1 for the real context and y = 0 for each noise word, and the gradient with respect to any score s = u . v is simply (y - sigmoid(s)). That gives the whole update: with learning rate lr and g = (y - sigmoid(u . v)) * lr for each of the k + 1 rows, add g * v to each output row and add the sum of g * u to the center row. Cost per pair is O((k + 1) d), independent of V. The papers suggest k of 5 to 20 for small corpora and 2 to 5 for large ones.

Noise words are drawn from the unigram distribution raised to the power 0.75 and renormalised. The exponent flattens the distribution: in the toy corpus below, the word 'the' is 10.0% of tokens but gets 8.7% of the noise draws, and rare words gain. The C tool samples from a large precomputed table; the alias method gives O(1) draws from the same distribution with far less memory. The C code also skips a negative that happens to equal the positive context.

The data pipeline: counts, subsampling, windows

Three data steps matter as much as the objective. First, a minimum count (gensim's min_count, default 5) drops rare words: they get too few updates to learn anything and they bloat both matrices. Second, subsampling discards frequent tokens before windows are formed. The paper gives a discard probability of 1 - sqrt(t / f) for a word with corpus frequency f; the C tool and gensim instead keep a word with probability sqrt(t / f) + t / f, with default sample=1e-3. In the toy corpus 'the' has f = 0.10, so it is kept 11% of the time. Subsampling speeds training and also widens the effective window, because removed function words no longer occupy window slots.

Third, the dynamic window: for each center the trainer draws b uniformly from 1 to the window size and uses b words on each side. Near contexts are therefore sampled more often than far ones, an implicit distance weighting. Small windows (2 to 5) favour syntactic and functional similarity; large ones (10 or more) favour topical relatedness.

A tested trainer in NumPy

The trainer below implements skip-gram with negative sampling, subsampling, the 0.75 noise distribution, the dynamic window and linear learning-rate decay. It is written for clarity and runs a few thousand pairs per second; production trainers are compiled and run many threads that update shared matrices without locks.

import numpy as np
from collections import Counter

def build_vocab(tokens, min_count=1):
    counts = Counter(tokens)
    vocab = [w for w, n in counts.most_common() if n >= min_count]
    idx = {w: i for i, w in enumerate(vocab)}
    freq = np.array([counts[w] for w in vocab], dtype=np.float64)
    return vocab, idx, freq

def keep_prob(freq, t=1e-3):               # the C-tool / gensim form
    f = freq / freq.sum()
    return np.minimum(1.0, np.sqrt(t / f) + t / f)

def noise_dist(freq, power=0.75):
    p = freq ** power
    return p / p.sum()

def train_sgns(tokens, dim=50, window=5, neg=5, epochs=2, lr0=0.025, t=1e-3, seed=0):
    rng = np.random.default_rng(seed)
    vocab, idx, freq = build_vocab(tokens)
    V = len(vocab)
    W_in = (rng.random((V, dim)) - 0.5) / dim
    W_out = np.zeros((V, dim))
    keep, noise = keep_prob(freq, t), noise_dist(freq)
    ids = np.array([idx[w] for w in tokens])
    total, step = epochs * len(ids), 0
    for _ in range(epochs):
        sub = ids[rng.random(len(ids)) < keep[ids]]
        for pos, center in enumerate(sub):
            lr = max(lr0 * (1 - step / total), lr0 * 1e-4)
            step += 1
            b = rng.integers(1, window + 1)
            ctxs = sub[max(0, pos - b):pos].tolist() + sub[pos + 1:pos + 1 + b].tolist()
            for ctx in ctxs:
                targets = np.concatenate(([ctx], rng.choice(V, neg, p=noise)))
                labels = np.zeros(neg + 1); labels[0] = 1.0
                v = W_in[center].copy()            # old value, used by both updates
                score = 1 / (1 + np.exp(-(W_out[targets] @ v)))
                g = (labels - score) * lr
                grad_in = g @ W_out[targets]
                np.add.at(W_out, targets, np.outer(g, v))   # repeated ids accumulate
                W_in[center] += grad_in
    return vocab, idx, W_in

Two lines are easy to get wrong. W_in[center] is a NumPy view, so without .copy() the output update would see the already-updated center vector. And W_out[targets] += ... silently drops updates when an id repeats among the targets, which np.add.at handles. Both bugs still produce plausible neighbours, which is why they survive.

Worked example: two sealed topics

To see the mechanics without a large corpus, generate 3,000 ten-token sentences, each drawn from one of two topics with no shared content words: royalty (king, queen, prince, crown, throne, palace...) and baking (bread, butter, flour, oven, sugar...), plus the words 'the' and 'a' in every sentence. That is 30,000 tokens and 18 types. Training with d = 50, window 5, k = 5 and two epochs gives these nearest neighbours by cosine similarity:

The topics separate, which is the point of the exercise. But look at the cross-topic pair: cosine(king, bread) is 0.777, not near zero. All the vectors share a large common direction, an artefact of training with shared function words and a tiny vocabulary. Subtract the mean vector from every row first and the same pairs read 0.988 for king and queen and -0.937 for king and bread. Real embeddings show the same anisotropy in milder form, and centering (plus removing a few dominant principal components) is a standard post-processing step. Do not read these toy numbers as evidence about analogies; a corpus with two sealed topics is a unit test, not a benchmark.

Why it works: shifted PMI

Why does predicting co-occurrence produce useful geometry? Levy and Goldberg showed in 2014 that SGNS, at its optimum, implicitly factorises a matrix whose entries are PMI(w, c) - log k, where PMI is the pointwise mutual information log(P(w, c) / (P(w) P(c))). The dot product of a word and a context vector approximates how much more often they co-occur than chance, shifted by the number of negatives. This explains several observations: words with similar PMI rows get similar vectors, k acts as a regulariser that shifts the matrix, and count-based methods (a truncated SVD of a positive PMI matrix, or GloVe) with matching hyperparameters reach comparable quality.

It also explains analogy arithmetic: if 'king' differs from 'queen' in co-occurrence roughly as 'man' differs from 'woman', the offsets line up and king - man + woman lands near 'queen', approximately.

Evaluating embeddings honestly

Evaluate on the task you care about first; intrinsic benchmarks are a sanity check. The standard ones are word-similarity sets (WordSim-353, SimLex-999), scored by Spearman correlation between human ratings and cosine similarities, and the analogy set released with the original tool (about 19,500 questions such as 'Athens is to Greece as Oslo is to ?').

Know the trap in the analogy protocol. The usual 3CosAdd method returns the word nearest to b - a + c excluding a, b and c themselves. Without the exclusion the answer to king - man + woman is very often 'king', because the offset is small relative to the vector. Several papers have since shown that much of the analogy accuracy comes from that exclusion rather than clean linear structure. Treat analogy scores as a regression test for your own pipeline.

For retrieval use, measure recall at k of nearest-neighbour search on labelled pairs from your domain, and if the vocabulary is large serve the vectors from an approximate index such as HNSW, normalising rows to unit length so inner product equals cosine.

Operational guidance

• Tokenise identically at training and serving. Lower-casing, Unicode normalisation and phrase joining (the paper scores bigrams with (count(ab) - delta) / (count(a) count(b)) and merges high scorers into tokens like New_York) must be the same code path in both places, or lookups silently miss.
• Budget memory. Training holds two V x d float32 matrices: for V = 1,000,000 and d = 300 that is 1.2 GB each. Typical settings are d = 100 to 300, window 5, k = 5, and 5 or more epochs; gensim's defaults are vector_size=100, window=5, min_count=5, sg=0, negative=5, sample=1e-3, ns_exponent=0.75, alpha=0.025, epochs=5. Note sg=0 means CBOW; pass sg=1 for skip-gram.
• Version vectors with their vocabulary and preprocessing. Two training runs, even on the same data, produce different coordinate systems. Vectors from different runs are not comparable, and a downstream model trained on one set breaks on the next. If you must compare, align with orthogonal Procrustes on shared words first.
• Plan for out-of-vocabulary words. Word2Vec has no vector for an unseen word. Use an explicit unknown vector, back off to subword models such as fastText (which sums character n-gram vectors), or retrain on a schedule.
• Expect multithreaded nondeterminism. Lock-free parallel updates make runs differ even with a fixed seed; test properties such as neighbour sets, not exact numbers.

Failure modes

• One vector per word type. 'bank' gets a single vector that averages the river and finance senses, weighted by frequency. If senses matter, you need a contextual encoder.
• Inherited bias. Vectors encode the associations in the corpus, including social stereotypes; Bolukbasi and colleagues documented gender analogies in 2016. Audit before using vectors in anything that ranks or scores people.
• Frequency leaks into geometry. Vector norms grow with frequency, and very frequent words cluster together. Normalise rows and consider centering before similarity search.
• Silent training bugs. The view and repeated-index bugs above, a learning rate that never decays, or forgetting subsampling all yield vectors that look reasonable on three hand-picked queries. Keep a fixed fixture with known topic structure and test separation on every change.

Trade-offs

What to do next

1. Copy the trainer, rebuild the two-topic fixture and confirm the topics separate; then delete the .copy() and watch what changes.
2. Train gensim with sg=1 on a corpus from your domain, and compare its neighbour lists with a PMI plus SVD baseline on 20 words you know well.
3. Build an evaluation set from your task, such as pairs of items users treat as substitutes, and score recall at 10 before tuning anything.
4. Write down the tokeniser version, vocabulary and vector file hash together, and fail the deploy if they disagree.
5. Read the token embedding lookup article to see how the same lookup table becomes the first layer of a contextual model.