A Discussion on ``Noise Contrastive Estimation'': The Beauty of a Winding Path

Exponential Family Distributions

Exponential family distributions appear in many problems. For the probability p(\boldsymbol{x}) of a variable \boldsymbol{x}, we can 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 type of distribution is also known as the “Boltzmann distribution.”

In machine learning, there are two main sources of exponential family distributions. The first source is softmax: when making classification predictions, we usually use softmax to activate the output of the final fully connected layer; this is a discrete Boltzmann distribution over a finite number of points. The second source is the Principle of Maximum Entropy: when we introduce a feature and can estimate its expected value, the maximum entropy model tells us that its distribution should take the exponential form of that feature. (Refer to “Too Much Entropy: From Entropy and the Maximum Entropy Principle to the Maximum Entropy Model (II)”).

The Hard-to-Calculate Partition Function

Overall, exponential family distributions are a very practical class of distributions encountered in machine learning, mathematics, and physics. However, they have a major problem: they are difficult to calculate—specifically, the partition function is difficult to calculate.

There are two main reasons for this difficulty. One is the sheer volume of computation. For example, in language modeling (including Word2Vec), predicting the distribution of the current word based on context requires summing over hundreds of thousands or even millions of terms (depending on the vocabulary size) to calculate the normalization factor. In this case, it’s not that it cannot be calculated, but that the computational cost is prohibitively high. The other situation is when it simply cannot be calculated analytically. For instance, if p(x)=\frac{e^{-ax^2-bx^4}}{Z}, then: Z = \int e^{-ax^2-bx^4} dx \tag{2} This integral cannot be solved easily, let alone for more complex functions. From this perspective, we can perhaps appreciate why the Gaussian distribution is so commonly used—because if we switch to another distribution, the math simply stops working...

In machine learning, if we are only doing classification or prediction, it doesn’t matter whether the normalization factor is calculated, because we only need to find the maximum relative value. However, before prediction, we face the problem of training, which is parameter estimation. Specifically, G(\boldsymbol{x}) contains some unknown parameters \boldsymbol{\theta}, so it should be written as G(\boldsymbol{x};\boldsymbol{\theta}). The probability distribution is: p(\boldsymbol{x})=\frac{e^{G(\boldsymbol{x};\boldsymbol{\theta})}}{Z(\boldsymbol{\theta})} \tag{3} We need to infer \boldsymbol{\theta} from samples of \boldsymbol{x}. Usually, we use Maximum Likelihood Estimation (MLE), but without calculating Z(\boldsymbol{\theta}), we cannot compute the likelihood function, and thus we cannot proceed.

Turning it into a Binary Classification Problem

The idea of NCE is simple: it hopes that we compare real samples with a set of “noise samples” to discover the patterns of the real samples.

Specifically, the energy function remains G(\boldsymbol{x};\boldsymbol{\theta}), but we no longer calculate the probability p(\boldsymbol{x}) directly because the normalization factor is hard to compute. Instead, we calculate: 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.

Then, the loss function for NCE 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}) represents the real samples and U(\boldsymbol{x}) is some “uniform” distribution or another fixed distribution that is easy to sample from.

In short, the approach of NCE is to transform the problem into a binary classification problem, where real samples are classified as 1 and samples from another distribution are classified as 0.

Equivalence to the Original Distribution

The question now is whether the \boldsymbol{\theta} estimated from equation (5) is the same as the result of a direct Maximum Likelihood Estimation from equation (3) (assuming it were theoretically feasible).

The answer is that they are essentially the same. We can 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 \tilde{p}(\boldsymbol{x}) and U(\boldsymbol{x}) are independent of the parameters \boldsymbol{\theta} and \gamma, changing the loss to the following form will not affect 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 the KL divergence. Since the KL divergence is non-negative, when the “assumed distribution form is satisfied and fully optimized,” equation (7) should be 0, which gives us \tilde{p}(y|\boldsymbol{x})= p(y|\boldsymbol{x}), meaning: \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} Solving for \tilde{p}(\boldsymbol{x}), we get: \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 a uniform distribution, then U(\boldsymbol{x}) is just a constant. Thus, the final result shows that \gamma - \log U(\boldsymbol{x}) plays the role of \log Z, while the distribution remains the original distribution (3) and \boldsymbol{\theta} remains the original \boldsymbol{\theta}.

This demonstrates that NCE is a clever scheme for indirectly optimizing equation (3): it seems roundabout, but the results are equivalent, and the computational cost of equation (5) is greatly reduced because it only depends on the number of samples.

Brief Overview of NCE and Negative Sampling

The systematic proposal of NCE was in the 2010 paper “Noise-contrastive estimation: A new estimation principle for unnormalized statistical models.” Since then, training large-scale neural language models has basically adopted NCE or similar losses. The title of the paper actually highlights the key point of NCE: it is a “parameter estimation principle” for “unnormalized models,” specifically designed for scenarios where the normalization factor is hard to compute.

But in fact, the idea of “negative sampling” had been used much earlier. For example, in the 2008 ICML paper “A Unified Architecture for Natural Language Processing: Deep Neural Networks with Multitask Learning,” Ronan Collobert and Jason Weston already used negative sampling to train word vectors. Keep in mind that this was four or five years before Word2Vec was released! For more on the history of word vectors and language models, refer to licstar’s “Word Vectors and Language Models.”

Based on the same need to reduce computational complexity, Google’s Word2Vec later adopted the negative sampling technique. In many tasks, it even performed better than Huffman Softmax, especially in the “word analogy” experiments. We will analyze the secret behind this shortly.

Word2Vec

Now let’s look at Word2Vec to analyze a few things. Taking the Skip Gram model as an example, the goal of Word2Vec 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 parameters to be optimized, representing word vectors in two different spaces (center word and context). Clearly, the problem here is the large computational cost of the partition factor. Solutions include Huffman Softmax and negative sampling. We won’t focus on Huffman Softmax here; just know it’s an approximation of the standard Softmax. Let’s look at negative sampling. Word2Vec changes the optimization objective to: \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 formula might look a bit dizzying, but it essentially expresses the idea: “Skip Grams appearing in the corpus are treated as positive samples, and randomly sampled words are treated as negative samples.”

First, it is obvious that compared to equations (4) and (5), equation (12) omits the training parameter \gamma, or effectively sets \gamma=0. Is this allowed? It is said that some have conducted comparative experiments showing that the trained \gamma indeed fluctuates around 0, so this default operation is basically reasonable.

Secondly, for negative samples, Word2Vec does not “sample every word uniformly”; instead, it samples according to the overall frequency of each word. Consequently, 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 effect 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 side is the mutual information of the two words! Originally, our fitting goal was for the inner product of two words to equal the (log) conditional probability \tilde{p}(w_j|w_i); now, through negative sampling in Word2Vec, the inner product of two words is their mutual information.

This explains why Word2Vec’s negative sampling might perform better than Huffman Softmax. Huffman Softmax is just an approximation of Softmax, essentially fitting \tilde{p}(w_j|w_i), while the negative sampling technique fits the mutual information \log\frac{\tilde{p}(w_j|w_i)}{\tilde{p}(w_j)}. We know that Word2Vec relies on word co-occurrence to reflect meaning, and mutual information reflects the “true” co-occurrence relationship better than the conditional probability \tilde{p}(w_j|w_i). 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 know you,” which better captures semantic relationships.

In a previous word vector model I constructed, “A More Unique Word Vector Model (III): Models Describing Correlation,” I also showed that models based on mutual information can theoretically explain many experimental results like “word analogy.” This indirectly confirms that the combination of “Skip Gram + Negative Sampling” based on mutual information is an excellent pairing for Word2Vec. Therefore, the fundamental reason is not whether Huffman Softmax or negative sampling is inherently superior, but that their optimization objectives are different.