word2vec

I reimplemented word2vec in PyTorch. I came to this straight from reimplementing Bengio et al. 2003, A Neural Probabilistic Language Model. Section 2 in that paper gives you the model architecture, the parameter set, the likelihood, the gradient structure and everything. You can sit down with it and write code. So I started reading Mikolov et al., Efficient Estimation of Word Representations in Vector Space (arXiv:1301.3781) to implement and learn about embeddings, expecting the same thing but the paper was very concise focusing upon only the error and other optimisations.

There is no objective function for CBOW. Skip-gram gets one line of average log probability before the paper hands off hierarchical softmax to Morin & Bengio 2005. There is no initialization scheme, no learning rate schedule, no description of how context windows are constructed, no mention of the sampling distribution which was one of the things that confused me.

You cannot implement word2vec from the word2vec paper (atleast It was the case for me as of now). What you can implement it from is a 700-line C file.

Where the model actually lives

Mikolov et al. 2013b (arXiv:1310.4546): This is where negative sampling appears, along with frequent-word subsampling and the noise distribution's 3/4 exponent.

Rong, word2vec Parameter Learning Explained (arXiv:1411.2738). Twenty pages of derivations, structured perfectly describing one-word context first, then CBOW, then skip-gram, every gradient written at index level. It exists because the original is underspecified, and it's the paper I actually implemented from and took reference.

Goldberg & Levy, word2vec Explained (arXiv:1402.3722). Five pages that make precise what the negative sampling objective is.

word2vec.c itself. The normative spec. Several things that materially change results appear only here, and they're the subject of a whole section below.

What negative sampling actually is

Skip-gram's stated objective is a softmax over the vocabulary:

That denominator touches all 71,290 output vectors for every training pair. It's the partition function, and it's the entire cost of the model.

Negative sampling doesn't approximate it. It replaces the question. Instead of "which of $V$ words is the context?" it asks "is this (word, context) pair real, or did I fabricate it?" For each observed pair you draw $k$ fabricated ones and train a binary classifier to separate them:

Cost per update: $k + 1$ dot products instead of $V$.

The negatives are load-bearing. Without them the objective has a trivial maximum which makes every vector identical and large, and $\sigma(\mathbf{v} \cdot \mathbf{v}) \to 1$ everywhere. The negatives are what prevent representational collapse. The model has to place vectors so real pairs score high and random pairs score low. This is worth naming, because the same trick shows up everywhere once you see it. Instead of computing a probability over every possible answer, you score the right answer against a few wrong ones. CLIP does this with images and captions. InfoNCE does it in representation learning generally. Energy-based models do it to avoid computing their own normalizing constant. One rule I can observe here is that when the denominator is too expensive to compute, stop computing it and compare against samples instead.

Here's the part Mikolov's paper obscures. Negative sampling is frequently described as an approximation to the full softmax. It isn't. NCE (noise contrastive estimation) is the consistent estimator: it converges to the true distribution as the number of noise samples $k \to \infty$. Negative sampling drops exactly the terms that make NCE consistent; Goldberg & Levy show the two coincide only in the special case where $k = |V|$ and the noise distribution is uniform. Outside that case it optimizes a genuinely different objective that happens to produce good vectors. Their paper exists largely to state this, and they're explicit that they do not explain why the resulting vectors have the properties they do.

The noise distribution

The exponent is pure empiricism. Mikolov reports it beat unigram (exponent 1) and uniform (exponent 0) on every task tried.

What it does is flatten. Take a word at 7% corpus frequency and one at 0.0018%: the raw frequency gap is $3944\times$, but after the 3/4 power it becomes $3944^{0.75} \approx 498\times$ in sampling probability. Rare words get drawn several times more often than their raw frequency warrants.

Both endpoints fail for opposite reasons. Under unigram sampling, your negatives are the, of, and on nearly every draw, so most of the vocabulary's output vectors barely train. Under uniform sampling, negatives are overwhelmingly obscure words that are trivially distinguishable from any real context, so $\sigma(-\mathbf{v}_n \cdot \mathbf{v}_w)$ is already near 1, gradient near 0, no learning signal. The 3/4 exponent interpolates: negatives common enough to be confusable, spread enough to cover the vocabulary.

The things that only exist in the C

The dynamic window. The window size is redrawn for every single center word: b = random % window, effective radius $\text{window} - b$. So the actual radius is uniform over $\{1, \dots, 5\}$, mean 3, not a fixed 5. This is implicit distance weighting (nearby words land in context more often), and it's word2vec's only mechanism for saying "adjacent matters more than five-away." No paper mentions it. Remove it and you get a uniform box filter; training still runs, loss still drops, vectors are quietly worse (my observation while implementing and testing for neighbours).

A subsampling formula that doesn't match the paper. The paper gives

The code computes

and compares it against a uniform draw. These are different functions. I used the paper's implementation.

The mechanism matters more than the formula. Subsampling deletes tokens from the stream before windowing, and the survivors close ranks. Deleting the doesn't just remove a useless pair; it pulls a content word into range that was previously out of reach. Over a corpus where function words are roughly 15% of tokens, the effective semantic radius widens considerably. That's why subsampling improves rare-word vectors rather than just speeding things up. It's also why the draw has to be redone every epoch: it's a stochastic augmentation, not a vocabulary filter, and materializing one subsampled corpus and reusing it throws that away.

Asymmetric initialization. Input vectors $\mathcal{U}(-0.5/D,\, +0.5/D)$, output vectors zeros. Note the divisor is $D$, not $\sqrt{D}$. This is not Xavier and isn't derived from any variance analysis.

Two things make it work. You only need to break symmetry on one side: with the output matrix at zero, every initial score is exactly 0, so $\sigma(0) = 0.5$ and the model starts perfectly uncommitted. That's the output matrix doing the work, not the input scale. The dot products are zero regardless of how the input matrix is initialized.

The $1/D$ scaling matters for what happens next. It puts input vector norms at

and the gradient on the output matrix is proportional to those vectors, so both matrices grow slowly and scores stay near 0, in the region where $\sigma'$ is largest, for a long stretch of early training. Compare against nn.Embedding's default of $\mathcal{N}(0, 1)$, which gives norms around $\sqrt{D} \approx 14$: initial dot products in the hundreds, deep in sigmoid saturation, gradient effectively zero. If you construct the embeddings and don't override the weights, training breaks while producing a plausible-looking loss curve.

(A naming note: syn0 is the input/word matrix, the one exported as your final vectors, and syn1neg is the output/context matrix, so named because it's the one used under the negative sampling objective.)

No collision filtering. Sampled negatives can include the true context word. The C doesn't check. Neither should you, and not just for fidelity: the expectation in the objective is over $P_n$ unconditionally, and filtering makes it $P_n(w \mid w \neq c)$, a different distribution.

What actually was kind of confusing while implementing

The math is easy. The data pipeline is 70% of the work and all of the bugs.

Learning rate. This is the one that cost me the most time and it has nothing to do with word2vec. The C is fully online, batch size 1, $\text{lr} = 0.025$. I batch at 1024 with the loss meaned over the batch, and mean-reduction divides each example's gradient by the batch size. The paper's 0.025 becomes an effective per-example step of $0.025 / 1024 \approx 2.4 \times 10^{-5}$. Nothing learns. Retuning to $\text{lr} = 1.0$ recovers a comparable per-example step. It looks absurd written down but it's correct.

Sparse gradients. nn.Embedding(..., sparse=True). Without it, every step materializes a dense $V \times D$ gradient. At $71{,}290 \times 200$ that's 57 MB of near-zeros per matrix per step, and it dominates runtime completely.

MPS lost. I benchmarked CPU against the MPS backend expecting the GPU to win and it was slower. This workload is memory-bound gather/scatter with $D = 200$, a handful of tiny dot products against data-dependent row indices. There's no arithmetic intensity to exploit, so there's nothing for a GPU to do. Final numbers: ~128k tokens/sec on CPU, ~66 s/epoch, ~11 minutes for 10 epochs, ~1.96 GB peak.

F.logsigmoid, never log(sigmoid(x)). In float32 the sigmoid saturates at both ends: large negative inputs give exactly 0.0, and $\log(0) = -\infty$; large positive inputs give exactly 1.0, and $\log(1) = 0$, which silently kills the gradient instead. logsigmoid computes the composition stably and avoids both.

The verification that caught real bugs:

• torch.autograd.gradcheck against my hand-derived Rong equations on a $D = 3$ float64 case
• deliberately overfitting a 500-token corpus then scaling
• Loss is useless as a quality signal at scale, because it's a discrimination loss against noise and it decreases whether or not the vectors are good. This was new for me.

Results

$D = 200$, $\text{window} = 5$, $k = 5$, \text{min_count} = 5, $t = 10^{-4}$, 10 epochs on text8, vocabulary 71,290.

france   -> austria .96, spain .96, germany .96, italy .957, hungary .951
physics  -> chemistry .928, mathematical .917, mathematics .908, quantum .901
god      -> spirit .956, divine .937, eternal .927, allah .926, heaven .917
computer -> digital .934, graphics .923, desktop .921, computing .921, unix .918
war      -> wars .879, invasion .875, vietnam .873, civil .873, battles .872

Analogies, 3CosAdd:

france:paris :: germany:? -> berlin .959, munich .911, moscow .905, vienna .904
man:king    :: woman:?    -> emperor .887, empress .881, augustus .871

berlin at rank 1. queen is not in the top 5 for the famous one. empress shows up at rank 2 instead. Looking at king's own neighbours explains it: viii, vii, elizabeth, constantine, crowned. The vector has absorbed regnal numbering and the general semantics of monarchy rather than the gendered ruler axis the analogy needs. text8 is 17M tokens; the original results used billions.

3CosMul gave identical rankings on both of these. That's worth noting rather than assuming. 3CosMul isn't a monotone rescaling of 3CosAdd and can genuinely reorder results, which is why Levy & Goldberg proposed it. It tends to help most on weaker, undertrained embeddings, and getting the same order here suggests these two queries aren't close calls.

SGNS is secretly factorizing a matrix

Here's the part that reframed the whole exercise. As this is like very interesting to see in practice after training the embeddings.

Levy & Goldberg (NIPS 2014) showed that if $D$ is large enough that $W C^\top$ can realize any $V \times V$ matrix, the terms $x_{wc} = \mathbf{v}_w \cdot \mathbf{v}_c$ in the SGNS objective become independent free parameters. So optimize each one on its own. With $a = \#(w,c)$ and $b = k \cdot \#(w) \cdot P_n(c)$, the per-cell objective is $a \log \sigma(x) + b \log \sigma(-x)$, and

Writing $|\mathcal{D}|$ for the number of observed pairs in the corpus (not to be confused with $D$, the embedding dimension) and substituting $P_n(c) = \#(c) / |\mathcal{D}|$:

SGNS's optimum is a factorization of the pointwise mutual information matrix, shifted down by $\log k$. Not an analogy but an identity, under the stated assumption.

Which means you can skip the neural network entirely. Count co-occurrences on the same corpus with the same window scheme, take

to keep the matrix sparse and finite, run a truncated SVD to rank 200, and you have word vectors.

They agree where the theory says they should:

And they diverge in a way that's more informative than the agreement:

SGNS returns words of the same kind. SVD returns rare words that happen to sit next to the target.

hauptbahnhof (German for "central station") appears almost exclusively adjacent to German place names, so its PMI with germany is enormous. And truncated SVD minimizes Frobenius error, which weights every cell of the matrix equally: a pair observed once and a pair observed ten thousand times get identical say in the reconstruction. So a handful of rare, high-PMI spikes dominate.

SGNS's objective weights each cell by $\#(w,c)$. It's a weighted matrix factorization. Those rare spikes get damped by the frequencies they actually occur at, and subsampling suppresses them further.

This is Levy & Goldberg's own stated caveat, and seeing it produce hauptbahnhof on my own corpus was more convincing than reading it. The theorem says the two methods target the same matrix. The gap between them is the weighting, and the weighting is what makes SGNS produce vectors that behave semantically rather than statistically.

What I'd tell someone starting this

Read Rong, not Mikolov. Implement from the C where the papers dont give much info, and write down every place you diverged. The deviations section in my README will give you some info.

And run the SVD baseline. It's the thing that changes how you think about the model. Word embeddings weren't a break from the count-based distributional semantics that preceded them; they're a stochastic, frequency-weighted way of computing the same object. The "neural" part is an optimization strategy, not a theory of meaning.

Code and full results:

• Codebase ocdbytes-ai/embeddings
• skip-gram model
• SVD baseline
• HF trained 200-d embeddings model