A Miscellany on "Noise Contrastive Estimation": The Beauty of an Indirect Route
Speaking of noise contrastive estimation, or "negative sampling," people probably immediately think of Word2Vec. In fact, its significance goes far beyond that. Noise Contrastive Estimation (NCE) is a roundabout yet remarkably elegant trick: it lets us estimate the parameters of a probability distribution even when we can't directly compute the normalization factor (also called the partition function). This post is an appreciation of the winding-path beauty of NCE.
Note: because of a different starting point, what this post calls "noise contrastive estimation" actually leans more toward what's often called the "negative sampling" trick, but the two are essentially the same thing, so we won't distinguish between them here.
Origins of the Problem
The root of the problem is that thorny, hard-to-shake-off exponential probability distribution.
Exponential family distributions
Exponential family distributions show up in many problems. That is, for the probability $p(\boldsymbol{x})$ of some variable $\boldsymbol{x}$, we write it as
$$p(\boldsymbol{x}) = \frac{e^{G(\boldsymbol{x})}}{Z}\tag{1}$$
where $G(\boldsymbol{x})$ is some "energy" function of $\boldsymbol{x}$, and $Z=\sum_{\boldsymbol{x}} e^{G(\boldsymbol{x})}$ is the normalization constant, also called the partition function. This kind of distribution is also known as a "Boltzmann distribution." more
In machine learning, exponential family distributions mainly arise from two sources. The first is softmax: when we do classification, we typically activate the output of the final fully-connected layer with softmax, which is exactly a discrete Boltzmann distribution over a finite set of points. The second source is the maximum entropy principle: once we introduce some feature and can estimate its expectation, the maximum entropy model tells us the resulting distribution should take the exponential form of that feature. (See "Entropy" Uncut: From Entropy, the Maximum Entropy Principle, to Maximum Entropy Models (Part 2)].)
The intractable partition function
Broadly speaking, exponential family distributions form an extremely useful class of distributions, one you'll run into whether in machine learning, mathematics, or physics. However, they come with a fairly serious problem: they're not easy to compute — more precisely, the partition function is not easy to compute.
Specifically, there are two possible reasons for this difficulty. One is that the computation is simply too expensive — for instance, in language model scenarios (including Word2Vec), where we predict the distribution of the current word from its context, computing the normalization factor requires summing over hundreds of thousands or even millions of terms (depending on vocabulary size). It's not that it's impossible, but that the computational cost is unbearable. The other case is that it genuinely can't be computed at all — for example, suppose $p(x)=\frac{e^{-ax^2-bx^4}}{Z}$, then we'd have
$$Z = \int e^{-ax^2-bx^4} dx\tag{2}$$
This integral simply cannot be worked out in closed form, let alone for more complicated functions. Perhaps from this angle we can now appreciate why the Gaussian distribution is used so often: because, because, because — switch to a different distribution and you get stuck.
In machine learning, if all we're doing is classification or prediction, it doesn't really matter whether we can compute the normalization factor, since we only need to compare relative values to find the maximum. But before prediction, we still face the problem of training, i.e., parameter estimation. Specifically, $G(\boldsymbol{x})$ actually contains some unknown parameters $\boldsymbol{\theta}$; to be precise, we should write it as $G(\boldsymbol{x};\boldsymbol{\theta})$, so the probability distribution becomes
$$p(\boldsymbol{x})=\frac{e^{G(\boldsymbol{x};\boldsymbol{\theta})}}{Z(\boldsymbol{\theta})}\tag{3}$$
We want to infer $\boldsymbol{\theta}$ from samples of $\boldsymbol{x}$. Normally we'd use maximum likelihood, but without being able to compute $Z(\boldsymbol{\theta})$ we can't even write down the likelihood function, and so we're stuck.
Enter NCE
Fortunately, NCE was born, and it successfully sidesteps this difficulty. For cases where the partition function can't be computed at all, it offers a way to keep going; for cases where the partition function is just too expensive to compute, it offers a way to cut down the computational cost.
Turning it into a binary classification problem
The idea behind NCE is simple: it has us compare real samples against a batch of "noise samples," and from that contrast, discover the pattern underlying the real samples.
Specifically, the energy is still the original energy $G(\boldsymbol{x};\boldsymbol{\theta})$, but now instead of directly computing the probability $p(\boldsymbol{x})$ — since the normalization factor is hard to compute — we instead compute
$$p(1|\boldsymbol{x})=\sigma\Big(G(\boldsymbol{x};\boldsymbol{\theta})-\gamma\Big)=\frac{1}{1+e^{-G(\boldsymbol{x};\boldsymbol{\theta})+\gamma}}\tag{4}$$
Here $\boldsymbol{\theta}$ is still the original parameter to be optimized, while $\gamma$ is a newly introduced parameter to be optimized as well.
Then, the NCE loss function becomes
$$\mathop{\text{argmin}}_{\boldsymbol{\theta},\gamma} - \mathbb{E}_{\boldsymbol{x}\sim \tilde{p}(\boldsymbol{x})}\log p(1|\boldsymbol{x})- \mathbb{E}_{\boldsymbol{x}\sim U(\boldsymbol{x})}\log p(0|\boldsymbol{x})\tag{5}$$
where $\tilde{p}(\boldsymbol{x})$ is a real sample, and $U(\boldsymbol{x})$ is some "uniform" distribution, or some other fixed distribution that's convenient to sample from.
Put simply, what NCE does is turn the problem into binary classification: label real samples as 1, and samples drawn from another distribution as 0.
Equivalence to the original distribution
The question now is whether the $\boldsymbol{\theta}$ estimated from equation $(5)$ is the same as what we'd get directly from maximum likelihood estimation on equation $(3)$ (which is theoretically feasible).
The answer is: essentially yes. Let's rewrite the loss in equation $(5)$ as
$$-\int \tilde{p}(\boldsymbol{x})\log p(1|\boldsymbol{x}) d\boldsymbol{x}- \int U(\boldsymbol{x})\log p(0|\boldsymbol{x})d\boldsymbol{x}\tag{6}$$
Since neither $\tilde{p}(\boldsymbol{x})$ nor $U(\boldsymbol{x})$ depends on the parameter $\boldsymbol{\theta},\gamma$, we can rewrite the loss in the following form without affecting the optimization result:
$$\begin{aligned}&\int \big(\tilde{p}(\boldsymbol{x})+U(\boldsymbol{x})\big) \left(\tilde{p}(1|\boldsymbol{x}) \log \frac{\tilde{p}(1|\boldsymbol{x})}{p(1|\boldsymbol{x})} + \tilde{p}(0|\boldsymbol{x})\log \frac{\tilde{p}(0|\boldsymbol{x})}{p(0|\boldsymbol{x})}\right)d\boldsymbol{x}\\ =&\int \big(\tilde{p}(\boldsymbol{x})+U(\boldsymbol{x})\big) KL\Big(\tilde{p}(y|\boldsymbol{x})\Big\Vert p(y|\boldsymbol{x})\Big) d\boldsymbol{x}\end{aligned}\tag{7}$$
where
$$\tilde{p}(1|\boldsymbol{x})=\frac{\tilde{p}(\boldsymbol{x})}{\tilde{p}(\boldsymbol{x})+U(\boldsymbol{x})}\tag{8}$$
Equation $(7)$ is the integral of a KL divergence, and since KL divergence is non-negative, when "the assumed form of the distribution is satisfied and optimization is carried out sufficiently," equation $(7)$ should equal 0, giving us $\tilde{p}(y|\boldsymbol{x})= p(y|\boldsymbol{x})$, i.e.,
$$\frac{\tilde{p}(\boldsymbol{x})}{\tilde{p}(\boldsymbol{x})+U(\boldsymbol{x})}=\tilde{p}(1|\boldsymbol{x})=p(1|\boldsymbol{x})=\sigma\Big(G(\boldsymbol{x};\boldsymbol{\theta})-\gamma\Big)\tag{9}$$
From which we can solve for
$$\begin{aligned}\tilde{p}(\boldsymbol{x})=&\frac{p(1|\boldsymbol{x})}{p(0|\boldsymbol{x})}U(\boldsymbol{x})\\ =&\exp\Big\{G(\boldsymbol{x};\boldsymbol{\theta})-\gamma\Big\}U(\boldsymbol{x})\\ =&\exp\Big\{G(\boldsymbol{x};\boldsymbol{\theta})-\big(\gamma-\log U(\boldsymbol{x})\big)\Big\}\end{aligned}\tag{10}$$
If $U(\boldsymbol{x})$ is taken to be a uniform distribution, then $U(\boldsymbol{x})$ is just a constant, so the net effect is that $\gamma - \log U(\boldsymbol{x})$ plays the role of $\log Z$, while the distribution itself remains the original $(3)$, and $\boldsymbol{\theta}$ remains the original $\boldsymbol{\theta}$.
This shows that NCE is an ingenious scheme for indirectly optimizing equation $(3)$: seemingly roundabout, yet the result is equivalent, and moreover the computational cost of equation $(5)$ is greatly reduced, since it now only depends on the number of sampled points.
A Few Digressions
A few topics related to NCE, gathered here.
A brief note on NCE and negative sampling
NCE was systematically proposed in the 2010 paper Noise-contrastive estimation: A new estimation principle for unnormalized statistical models], and afterward, essentially all training of large-scale neural language models adopted NCE or similar losses. The title of the paper actually spells out the key point of NCE: it's an "estimation principle" for "unnormalized statistical models," designed specifically to handle scenarios where the normalization factor is hard to compute.
But in fact, the idea of "negative sampling" had already been in use well before that — for instance, in the 2008 ICML paper A Unified Architecture for Natural Language Processing: Deep Neural Networks with Multitask Learning] by Ronan Collobert and Jason Weston, negative sampling was already used to train word embeddings. Keep in mind, that was four or five years before Word2Vec was even released! For the story of word embeddings and language models, see licstar's Word Embeddings and Language Models].
Driven by the same need to cut down computational cost, Google's Word2Vec later adopted the negative sampling trick as well, and on many tasks it even outperformed Huffman-tree-based softmax — especially in the famous "word analogy" experiments. Let's now dig into the mystery behind this.
Word2Vec
Let's now bring this down to earth with Word2Vec. Taking the Skip-Gram model as an example, Word2Vec's objective is
$$p(w_j|w_i)=\frac{e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle}}{Z_i}\tag{11}$$
where $\boldsymbol{u}_i, \boldsymbol{v}_j$ are the parameters to be optimized, representing two separate sets of word embeddings — one for center words and one for context words. Obviously, the issue here is that the normalization factor is expensive to compute, and the standard remedies are Huffman softmax and negative sampling. We won't worry about Huffman softmax here — it suffices to know it's just an approximation of the standard softmax. Let's instead look at negative sampling. Word2Vec turns the optimization objective into:
$$\mathop{\text{argmin}}_{\boldsymbol{u},\boldsymbol{v}} - \mathbb{E}_{w_j\sim \tilde{p}(w_j|w_i)}\log \sigma\Big(\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\Big) - \mathbb{E}_{w_j\sim \tilde{p}(w_j)}\log \Big[1-\sigma\Big(\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\Big)\Big]\tag{12}$$
This expression may look a bit dizzying, but in short it expresses the idea that "Skip-Gram pairs that actually occur in the corpus are treated as positive samples, and randomly sampled words are treated as negative samples."
The first obvious thing is that, compared with equation $(4),(5)$, equation $(12)$ drops the training parameter $\gamma$ — in other words, it implicitly assumes $\gamma=0$. Is that allowed? Apparently people have run comparison experiments, and the results show that the trained $\gamma$ does indeed hover around 0, so this default choice is basically justified.
Second, for negative samples, Word2Vec doesn't "sample every word uniformly" — instead, it samples according to each word's overall frequency in the corpus. As a result, equation $(10)$ becomes
$$\tilde{p}(w_j|w_i)=\frac{p(1|w_i, w_j)}{p(0|w_i, w_j)}p(w_j)=e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle}\tilde{p}(w_j)\tag{13}$$
That is to say, the final fitting target is
$$\log \frac{\tilde{p}(w_j|w_i)}{\tilde{p}(w_j)} = \langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\tag{14}$$
As you can see, the left-hand side is exactly the mutual information of the two words! Our original fitting target was to make the inner product of two word vectors equal to the (log) conditional probability $\tilde{p}(w_j|w_i)$; but after negative sampling, the inner product of two word vectors in Word2Vec instead equals the mutual information between the two words.
This roughly explains why Word2Vec's negative sampling tends to outperform Huffman softmax. Huffman softmax is merely an approximation of softmax — it's still essentially fitting $\tilde{p}(w_j|w_i)$ — whereas the negative sampling trick is instead fitting mutual information $\log\frac{\tilde{p}(w_j|w_i)}{\tilde{p}(w_j)}$. As we know, Word2Vec relies on word co-occurrence to reflect word meaning, and mutual information better reflects the "genuine" co-occurrence relationship between words than conditional probability $\tilde{p}(w_j|w_i)$ does. In other words, $\tilde{p}(w_j|w_i)$ might reflect a relationship like "I know Jay Chou, but Jay Chou doesn't know me," whereas mutual information reflects "you know me, and I also know you" — and the latter better captures semantic relatedness.
Another word-embedding model I built earlier, A More Elegant Word Embedding Model (III): A Model Describing Correlation], also showed that models built from mutual information can theoretically explain many experimental results, including the "word analogy" phenomenon. This indirectly confirms that the combination of "Skip-Gram + negative sampling," grounded in mutual information, is an excellent pairing for Word2Vec. So the real reason isn't a matter of which of Huffman softmax or negative sampling is inherently better — it's that their optimization objectives are simply different to begin with.
The Train Has Reached Its Final Stop
The goal of this post was to introduce NCE as an elegant parameter-estimation technique, showing that it can be used to estimate the parameters of a probability distribution even when normalization is intractable. In principle this is a general-purpose method, and quite possibly, in some scenarios, it's the only viable option.
Finally, using Word2Vec as a concrete example, we did a bit of analysis, touched on some of the finer details of using NCE, and along the way explained why negative sampling works so well.
Related link: Word Embedding Blog Series Part 2: Comparing Several Approximate Softmax Methods for Language Modeling]
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.