Negative Sampling and Training Word2Vec

The Word2Vec Skip-gram model defines, for a center word $w_c$ and a context word $w_o$, a probability

where $\mathbf{v}_{w_c}$ is the "input" embedding of the center word and $\mathbf{v}_w$ ranges over "output" embeddings of every word in the vocabulary $V$. Training by maximum likelihood requires computing the denominator — the partition function — over the entire vocabulary at every step, since its gradient touches every output vector. This is the full softmax bottleneck: each single training example costs $O(|V|)$ time and touches $O(|V|)$ parameters.

Negative sampling (introduced alongside skip-gram as "Skip-gram with Negative Sampling", SGNS) replaces this multi-class objective with a set of independent binary classification problems. For each observed (center, context) pair, treat it as a positive example with label $1$. Then sample $k$ random words from the vocabulary (that are not the true context word) as negative examples with label $0$. Train a logistic classifier — using the dot product of the two embeddings as the score — to distinguish real pairs from fake ones. The key terminology: the true pair is called a positive sample, the randomly drawn words are negative samples, and $k$ (typically 5–20 for small corpora, 2–5 for very large ones) is the negative sampling rate. A second, complementary idea is subsampling of frequent words. Words like "the", "a", "is" occur so often that they contribute little information about any specific center word's meaning, yet they dominate training time and drown out the gradient signal from rarer, more meaningful words. Subsampling randomly drops instances of a word from training with a probability that grows with its corpus frequency, controlled by a threshold hyperparameter $t$ (commonly $10^{-5}$). This both speeds up training (fewer tokens processed) and improves embedding quality (better-balanced gradient signal across the vocabulary).

The negative sampling objective

For a positive pair $(w_c, w_o)$, define the sigmoid score

The per-pair SGNS objective, sampling $k$ negatives $w_{n_1}, \dots, w_{n_k}$ from a noise distribution $P_n(w)$, is

We maximize this quantity (or equivalently minimize its negation) via stochastic gradient ascent/descent. Note the sign flip $-\mathbf{v}_{w_{n_i}}^\top \mathbf{v}_{w_c}$ inside $\sigma$ for negatives: we want the dot product for a fake pair to be small (score close to $0$), which is equivalent to wanting $\sigma(-\text{score})$ to be close to $1$.

The noise distribution

Negatives are not drawn uniformly nor from the raw unigram frequency $f(w)$, but from a smoothed unigram distribution:

Raising frequencies to the power $0.75$ flattens the distribution: very frequent words become relatively less likely to be sampled compared to their raw frequency, and rare words become relatively more likely. This empirically produces better embeddings — negatives that are neither trivially easy (extremely common, obviously unrelated words) nor unhelpfully obscure.

Subsampling of frequent words

Before pairs are even generated, each token $w_i$ in the training stream is kept with probability

a common formulation used in the reference implementation (an equivalent simplified variant is $P(\text{keep}) = \min\!\big(1, \sqrt{t / f(w_i)}\big)$ used in some expositions). Both forms share the essential shape: keep-probability decreases smoothly as frequency $f(w_i)$ grows, with negligible discarding below the threshold $t$ and aggressive discarding well above it.

Step-by-step training loop

The full pipeline, combining skip-gram context generation, subsampling, and negative sampling: ``` INPUT: corpus (stream of tokens), vocabulary V with counts, embedding dim d, window size m, negatives k, subsample threshold t, learning rate eta, epochs E 1. Build vocabulary V from corpus; compute frequency f(w) for each w. 2. Precompute noise distribution: P_n(w) = f(w)^0.75 / sum_{w'} f(w')^0.75 for all w in V Build an alias table (or cumulative table) for O(1)-ish sampling from P_n. 3. Initialize input embeddings U[w] and output embeddings W[w] for all w in V, each a d-dimensional vector (small random init). 4. for epoch in 1..E: for each token position i in corpus: w_c = token at position i // --- Subsampling step --- keep_prob = compute_keep_probability(f(w_c), t) if random() > keep_prob: continue // discard this token instance // --- Dynamic window --- window = random_integer(1, m) // shrink window randomly per example context_positions = [i-window, ..., i-1, i+1, ..., i+window] (clipped to valid corpus bounds) for j in context_positions: w_o = token at position j // positive context word // --- Positive update --- score = dot(U[w_c], W[w_o]) g_pos = (1 - sigmoid(score)) // gradient coefficient grad_U = g_pos W[w_o] grad_W_pos = g_pos U[w_c] // --- Negative samples --- grad_U_total = grad_U for s in 1..k: w_n = sample_from(P_n) such that w_n != w_o score_n = dot(U[w_c], W[w_n]) g_neg = -sigmoid(score_n) grad_U_total += g_neg W[w_n] W[w_n] += eta g_neg U[w_c] // update negative's vector // --- Apply updates --- W[w_o] += eta grad_W_pos U[w_c] += eta * grad_U_total 5. return U // (or average of U and W, or U alone, per convention) ``` The gradient coefficients above follow directly from differentiating $J$. For the positive term $\log \sigma(s)$ with $s = \mathbf{v}_{w_o}^\top \mathbf{v}_{w_c}$,

and for a negative term $\log \sigma(-s_n)$,

Both use the identity $\sigma'(x) = \sigma(x)(1-\sigma(x))$ and the chain rule; multiplying by $\partial s/\partial \mathbf{v}_{w_c} = \mathbf{v}_{w_o}$ (or the symmetric term for the output vector) gives exactly the vector updates in the pseudocode.

Time complexity per training pair. Full softmax requires computing $\mathbf{v}_{w_c}^\top \mathbf{v}_w$ for every $w \in V$ to normalize, costing $O(|V| \cdot d)$ per (center, context) pair, where $d$ is the embedding dimension. Negative sampling instead computes only $k+1$ dot products (one positive, $k$ negatives), costing $O(k \cdot d)$, independent of $|V|$. Since $k$ is a small constant (commonly 5–20) while $|V|$ can be $10^5$ to $10^6$, this is a reduction from linear-in-vocabulary to constant-in-vocabulary per pair — the entire reason SGNS scales to web-sized corpora. For an epoch over a corpus of $N$ tokens with average window size $m$, the total cost is $O(N \cdot m \cdot k \cdot d)$ for negative sampling versus $O(N \cdot m \cdot |V| \cdot d)$ for full softmax. With $N = 10^9$, $m=5$, $k=10$, $d=300$, $|V| = 10^6$: negative sampling costs on the order of $10^9 \times 5 \times 10 \times 300 \approx 1.5\times10^{13}$ floating point multiply-adds per epoch, while full softmax costs roughly $10^9 \times 5 \times 10^6 \times 300 \approx 1.5\times10^{18}$ — a factor of $10^5$ difference, matching the ratio $|V|/k$.

Space. Both approaches store the same parameters: two $|V| \times d$ matrices (input and output embeddings), so space is $O(|V| \cdot d)$ regardless of the training method — negative sampling only changes the per-step compute, not the parameter count.

Why the approximation is sound. SGNS is not merely "softmax but cheaper" — it is a different but related objective with its own justification. Consider a single negative sample per positive pair for simplicity, drawn from the true unigram distribution. It has been shown (Levy & Goldberg, 2014) that maximizing the SGNS objective in expectation, for the case $k$ negative samples drawn from the unigram distribution, drives the dot product between vectors to satisfy

where $\mathrm{PMI}(w,c) = \log \frac{P(w,c)}{P(w)P(c)}$ is the pointwise mutual information between word $w$ and context $c$ over the corpus. This means SGNS is, in expectation, implicitly performing a low-rank factorization of a shifted PMI matrix — the same statistical quantity that count-based distributional semantics methods have long used, but learned via cheap stochastic binary classification rather than an explicit matrix decomposition. This connection is exactly what motivates comparing SGNS to GloVe, which factorizes a log-count matrix explicitly.

Correctness sketch via gradient balance. At a stationary point of the SGD dynamics, for each (center, context) pair the expected gradient must vanish. Setting the derivative of $J$ with respect to the score $s = \mathbf{v}_w^\top \mathbf{v}_c$ to zero and using $P_n(w) \propto$ smoothed frequency, one finds the balance condition

which after algebra recovers $s \approx \log\frac{P(w,c)}{P(w) P_n(c)} - \log k$, i.e., the shifted-PMI identity above (with the caveat that in the real algorithm $P_n$, not the raw unigram $P$, is used for negatives, and the $0.75$ exponent slightly perturbs this ideal relationship in practice, which empirically improves embedding quality on rare words).

Honest tradeoffs. Negative sampling is a biased, approximate objective, not a strict lower bound on the true softmax log-likelihood in the way that e.g. hierarchical softmax approximations sometimes are; its quality depends on the number of negatives $k$, the smoothing exponent, and the noise distribution shape. Too few negatives leads to underfitting the "what is not a valid context" signal; too many negatives increases cost roughly linearly with little added benefit past a point (diminishing returns are well documented empirically, with $k \approx 5$–$20$ typically sufficient). Subsampling is also a lossy heuristic: it discards information, and its keep-probability formula, while empirically well-tuned, is not derived from a clean statistical principle — it is a practical corpus-frequency correction, not a probabilistically unbiased importance-sampling scheme.

Consider a tiny toy corpus (ignore real-world scale, this is for tracing the mechanics): ``` "the cat sat on the mat the dog barked at the cat" ``` Vocabulary (with counts): the:4, cat:2, sat:1, on:1, mat:1, dog:1, barked:1, at:1. Total tokens $N=13$.

Step 1 — noise distribution. Compute $f(w)^{0.75}$ for each word (raw frequency $f(w) = \text{count}(w)/N$):

Note "the" drops from raw share $0.3077$ to smoothed share $0.257$ — still likely to be sampled as a negative, but proportionally less dominant than its raw frequency would suggest, exactly the intended flattening effect.

Step 2 — subsampling decision for "the". Using the simplified rule $P(\text{keep}) = \min(1, \sqrt{t/f(w)})$ with $t = 10^{-5}$:

So roughly 99.4% of occurrences of "the" are dropped from training — it will almost never be used as a center word or (in most implementations) removed from context windows either, freeing training time for informative pairs like (cat, sat) or (dog, barked). Contrast with "cat" ($f=0.1538$): $P(\text{keep}) = \sqrt{10^{-5}/0.1538}\approx 0.0081$ — still heavily subsampled since even "cat" is common in this tiny corpus, illustrating that subsampling is aggressive whenever $t$ is small relative to corpus size; in real multi-billion-token corpora, $t=10^{-5}$ produces much milder, more realistic discarding for mid-frequency words while still hammering stopwords.

Step 3 — one SGNS update. Suppose after subsampling we still process the pair center = "cat", context = "sat" (positive), with $k=2$ negatives sampled from $P_n$: suppose we draw "barked" and "on". Let embeddings be 2-dimensional for hand-tracing, with current values

Positive score: $s = \mathbf{v}_{sat}\cdot\mathbf{v}_{cat} = (0.05)(0.10) + (0.30)(-0.20) = 0.005 - 0.06 = -0.055$. $\sigma(-0.055) \approx 0.4863$. Gradient coefficient $g_{pos} = 1 - 0.4863 = 0.5137$ (we want this dot product to increase). Negative 1 ("barked"): $s_{n1} = (-0.10)(0.10) + (0.10)(-0.20) = -0.01 - 0.02 = -0.03$. $\sigma(-0.03)\approx 0.4925$, so $g_{neg1} = -0.4925$ (push this dot product down further — it should be even more negative). Negative 2 ("on"): $s_{n2} = (0.20)(0.10) + (0.05)(-0.20) = 0.02 - 0.01 = 0.01$. $\sigma(0.01)\approx 0.5025$, $g_{neg2} = -0.5025$ (this one is positive, meaning the model currently thinks "on" is a plausible context for "cat" — the gradient will push it down harder). With learning rate $\eta = 0.1$:

New $\mathbf{v}_{cat} \approx [0.0975, -0.1920]$. The update nudged "cat" only slightly since the sigmoid values were all near $0.5$ (an uninitialized, nearly-random embedding stage) — this is expected early in training and is the correct behavior: large, confident updates only occur once dot products diverge from zero, which is exactly what repeated SGD steps are designed to produce over many epochs.

Negative sampling's core idea — replace an expensive full-normalization objective with cheap binary discrimination against sampled negatives — extends far beyond Word2Vec. It underlies contrastive learning broadly, including modern sentence and document embedding methods, retrieval systems (dense passage retrieval for search engines, where a query must be scored against millions of documents using in-batch or hard negatives instead of a full corpus softmax), and recommendation systems (candidate generation models that must rank against catalogs of millions of items). The Skip-gram-style objective and its negative sampling training recipe are reused, with modifications, in graph embedding methods (node2vec, DeepWalk) that treat graph random walks the same way skip-gram treats sentences. Pretrained static word embeddings trained this way (word2vec, fastText) are still used as fast, low-resource baselines and as initialization or auxiliary features in some production NLP pipelines where full contextual transformer models are too costly to serve. Studying this mechanism directly prepares the ground for GloVe's explicit matrix-factorization alternative and for understanding Embedding Geometry and Analogies, since the shifted-PMI interpretation of SGNS explains why linear analogy structure (king − man + woman ≈ queen) emerges from these training objectives at all.

Ask the tutor: How does changing the negative sampling exponent from 0.75 to 1.0 (raw frequency) change the embeddings learned for rare words?