Simple explanation of NCE(Noise Contrastive Estimation) and InfoNCE

This post gives a concise explanation of InfoNCE, the loss introduced in Representation Learning with Contrastive Predictive Coding. Paper

InfoNCE is foundational to contrastive learning. It can be viewed as one way to find a shared representation space that can be learned across different tasks, helping open the new frontier of AI commonly called multimodality. Before examining InfoNCE itself, I will briefly introduce Noise Contrastive Estimation. If NCE does not interest you, feel free to skip ahead to the InfoNCE section. I referred to the following article: link.

Background

To explain the concept, we first need a concrete problem; an NLP task makes a useful example. Word2Vec is the best-known approach, so let us understand the problem through the following process.

Suppose we have the sentence, “The quick orange fox can jump.”

Use a sliding window to create a (context, target) pair for each word. The target is the middle word in the window, while the context consists of its surrounding words. To keep the example simple, assume one neighboring word on either side.

• ((quick, fox), orange),

• ((orange, can), fox),

• ((fox, jump), can)

Convert each context word into a vector, for example with a lookup table. For implementation details, see TensorFlow's tutorial. Embedding generally means converting an item into vector form; the context embedding that represents the full context can then be taken as the mean of the individual context-word embeddings.

Feed this mean embedding—the context vector—into an MLP (fully connected neural-network layers), and use softmax to obtain a probability mapping over target words. In other words, the output is a probability map over candidate target words.

Optimize cross-entropy loss against the one-hot encoding of the correct word.

This teaches the model the frequency and statistics of particular words within a sentence. Because the lookup table is learnable, once training finishes, similar words should occupy nearby positions in the embedding space and dissimilar words should be far apart.
Looking closely at the fourth step above, however, the network's dense layer has a weight matrix shaped (embedding dimension, vocabulary size). To predict each vocabulary word, we must first compute the layer output for every word ($\text{for } i : \text{ index of word, }z_i = Wx_i$), then map those outputs to probabilities with a softmax transformation.

$$p(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{\vert V \vert} e^{z_j}}$$

Naturally, this calculation requires a fixed vocabulary size. Cross-entropy loss is then computed from these probabilities.

$$L = -\sum\_j^{\vert V \vert} y_j \log (p_j) = -\log (p_{\text{target}})$$

The important point is that although the loss appears to sum over every prediction, the only nonzero term is where $y_i=1$ for the true label. In other words, it considers only the probability of the actual target word. Yet in $p(z_i)$, the probability for every word position shares the same denominator. That denominator is

$$\sum_{j=1}^{\vert V \vert} e^{z_j}$$

and because of it, every parameter receives a nonzero gradient term for every training example. You can think of this as noisy training.

Negative sampling

One way to solve this problem is to select only some incorrect vocabulary items (negative terms) instead of summing all of them. We call these selected non-target words negative samples. I will continue to use the terminology defined here throughout the post and in the discussion of InfoNCE.
The preceding process remains unchanged except for the addition of negative sampling. Written out again:

Suppose we have the sentence, “The quick orange fox can jump.”

Use a sliding window to create a (context, target) pair for each word. The target is the middle word in the window, while the context consists of its surrounding words. To keep the example simple, assume one neighboring word on either side.

• ((quick, fox), orange),

• ((orange, can), fox),

• ((fox, jump), can)

Convert each context word into a vector, for example with a lookup table. For implementation details, see TensorFlow's tutorial. The context embedding can be taken as the mean of the individual context-word embeddings.

Feed this mean-embedding context vector into an MLP and use softmax to obtain a target-word probability mapping. This time, rather than mapping probabilities over every word, treat the selected negative samples and the positive sample as the entire sample space.

Optimize cross-entropy loss against the one-hot encoding of the correct word.

Selecting only some samples means the denominator changes each time, so normalization may be imperfect. We assume it can be approximated over many training iterations. The key benefit of this training method is that it reduces the number of gradient updates:

$$\vert Embedding \vert \times \vert V \vert \rightarrow \vert Embedding \vert \times \vert N+1 \vert$$

For $N$ negative samples, optimization covers only this subset rather than the entire space of vocabulary samples. This is sensible for NLP tasks with very large vocabularies. Optimizing the embedding for “zebra” even in a context where it would never be used can harm training and consume additional memory.

Noise Contrastive Estimation (NCE) implements this negative-sampling idea with additional theoretical grounding.

Learning by comparison

In negative sampling, the true target is labeled 1 and the random samples are labeled 0. These labels naturally teach the network to distinguish real samples from noise. NCE answers this problem through logistic-regression modeling. Logistic regression models the log-odds that an input came from one class rather than another.

$$logit = \log (\frac{p_1}{p_2}) = \log (\frac{p_1}{1-p_1})$$

Here, we define the log-odds as the ratio between the probability of coming from the true word distribution $P$ and the probability of coming from the noise distribution $Q$.

$$logit = \log (\frac{P}{Q}) = \log (P)- \log (Q)$$

The name noise contrastive estimation reflects the fact that logit training estimates the real distribution relative to noise obtained through negative sampling.
The real distribution $P$ is intractable, but we can define the noise distribution $Q$ in various ways—for example, by sampling all vocabulary items uniformly or by accounting for word frequency in the training data. Regardless of the exact method, the important point is that this gives us a clear way to calculate $\log(Q)$.
Let us revisit the Word2Vec network above and see how $Q$ can be used.

We use the context vector as the network input. Instead of computing outputs for every vocabulary word, however, we compute them for words sampled randomly from our predefined distribution $Q$. The network output is therefore evaluated for the target word and $N$ words sampled from the noise distribution—$N+1$ items in total (random samples plus the target).
Because we define the noise distribution used for negative sampling and draw noise from it, the probability of each word can be calculated analytically under $Q$.
For example, if “cat” is sampled with probability 10% and “neko” with probability 90%, then $Q=0.1$ when “cat” is selected. We can thus think of $Q$ as the sampling frequency of each word. The probability $P$ in the preceding equation is the network's prediction. As in an ordinary logistic-regression task, training assigns 1 to the target word and 0 to the negative samples.

We have now taken a fairly long detour through NCE. I explained Noise Contrastive Estimation at length because the form and intuition of the loss proposed in this paper originate there.
The success of AI and deep learning has relied on gradient-based optimization, with algorithms such as SGD, RMSprop, and Adam providing excellent optimizers, and supervised learning on labeled datasets playing a central role. Yet deep learning remains limited in how readily it can be applied across different modalities. Here, a modality is the form in which something is represented—in other words, a type of dataset.
A dataset trained to distinguish human voices (speaker recognition) is difficult to apply to music-genre classification or translation. Methods such as transfer learning and domain adaptation assume similar representations have been learned, but that assumption fails for completely different modalities such as vision/audio or vision/text.
The same issue raises another question in unsupervised learning: can we learn a representation meaningful enough to capture high-level information without explicit supervision?
Because unsupervised learning must create a supervisory signal without labels, it often uses predictive coding. This method assumes causal data and trains by predicting missing parts or future values. The paper takes inspiration from an objective that formalizes predictive coding together with NCE.
The proposed method has three components:

• Map high-dimensional data into a more compact latent embedding space so the model can make conditional predictions more easily.
• Use a powerful autoregressive (AR) model to predict from this latent condition.
• Use NCE (Noise Contrastive Estimation) as the loss so the whole model can be trained end to end.

That concludes the long introduction. We can now examine Contrastive Predictive Coding itself.

Main Intuition and Motivation

What matters most in this paper is not the equations themselves, but the intuition behind the optimization design. If the shared information is a high-dimensional signal, the method learns how different parts of that signal are encoded. Its main objective is not to learn a representation of the modality itself—its low-level information or raw data—but to make the model learn relationships within each domain's latent space implicitly.
This called for a new kind of loss, and NCE supplied the approach. In prediction, especially with regressive structures, forecasting farther into the future may require global information; the paper calls this a slow feature. A slow feature might be intonation or mood in audio, the object itself in an image, or the plot or overarching theme in text—anything requiring inference over the data as a whole.

One clear difficulty with high-dimensional datasets is that MSE and cross-entropy loss are not especially effective. When an image-generation model produces a $1024 \times 1024$ sample, for example, a per-pixel loss accumulates a large penalty, causing predictions to blur rather than draw sharp edges. In a translation model, cross-entropy loss must propagate gradients through every class as the number of candidate words grows; training also degrades because a huge number of probability labels are sparsely distributed.

Another problem is the need for a powerful conditional generative model. During training it must reside on the GPU alongside the encoder ($E$) being learned—even if its own parameters are not updated—consuming hardware resources, while generation time introduces a training bottleneck. A further problem with conditional generative models is their tendency to ignore context $c$. When generating high-dimensional data such as images, a high-level latent variable such as a class label contains relatively little information. Directly modeling the conditional distribution $p(x \vert c)$ in a decoder may therefore train it to generate while simply ignoring $c$, rather than using information shared between $x$ and $c$.

To predict future information, the paper therefore learns the target $x$ (the future value) and context $c$ (the current value) in compact distributions, designing a nonlinear mapping that makes them share information. Using mutual information, this is written as

$$I(x; c) = \sum_{x,~c} p(x, c) \log \frac{p(x \vert c)}{p(x)}$$

Mutual information (MI) measures how similar the joint distribution $p(x,c)$ of its two arguments $(x,c)$ is to the product $p(x)p(c)$, and is derived as follows.

$$I(x, c) = \sum_c \sum_{x \in X} p(x,c) \log \frac{p(x, c)}{p(x)p(c)}$$

This expression is derived in the same way as KL divergence. Applying Bayes' rule gives

$$I(x, c) = \sum_c \sum_{x \in X} p(x,c) \log \frac{p(x, c)}{p(x)p(c)} = \sum_{x, c} p(x, c) \log \frac{p(x \vert c) p(c)}{p(x) p(c)} = \sum_{x, c} p(x, c) \log \frac{p(x \vert c)}{p(x)}$$

It measures the degree of dependence between $x$ and $c$: if $x$ is independent of $c$, the value is small; the more dependent they are, the larger it becomes. Unlike KL divergence, this expression is commutative and therefore symmetric.

Contrastive predictive coding

The figure above shows the structure of a CPC (Contrastive Predictive Coding) model. First, $g_{enc}$, a nonlinear encoder, converts the input sequence $x_t$ into a sequence of latent representations.

$$g_{enc}(x_0, x_1, \cdots, x_t) = (z_0, z_1, \cdots, z_t)$$

The autoregressive model $g_{ar}$ then uses causal inputs—the past relative to the current time—to produce a context latent representation.

$$g_{ar}(z_{\leq t}) = c_t$$

Up to this point, the flow is effectively identical to an ordinary AR architecture. But rather than predicting a future observation directly, the model estimates a density ratio proportional to the mutual information between the future prediction and the context vector.

$$f_k(x_{t+k,~c_t}) \propto \frac{p(x_{t+k} \vert c_t)}{p(x_{t+k})}$$

The quantity inside $\log$ in the preceding equation is the ratio that captures the difference between the joint distribution of $x,c$ and the product of their individual distributions. Building the model in proportion to this quantity automatically encourages the future prediction to attend meaningfully to context. By optimizing the density ratio, we can act like a tutor telling the model, “This is what you should look at and learn from.”

The density ratio $f$ is unnormalized. It may be a log-bilinear model as below, or it may be replaced by a nonlinear network or an RNN.

$$f_k(x_{t+k}, c_t) = exp(z_{t_k}^T, W_kc_t)$$

Using this density ratio to predict the next $z$ together with the encoder avoids requiring the model to predict the high-dimensional data $x$ itself. Predicting $z$ in the encoded latent space does not directly manipulate $p(x)$ or $p(x \vert c)$, but it lets us use Noise Contrastive Estimation and importance sampling, which draw samples from a distribution we can posit. At last, it is time to put all that work understanding NCE to use.

InfoNCE, Mutual information estimation

Looking at the architecture, we ultimately need to train two models: the encoder and the autoregressive model. Both are optimized by an NCE-based loss function, which the authors call InfoNCE.
As in the NCE described earlier, assume $N$ random samples. One is positive, and the other $N-1$ are negative samples drawn from a proposal distribution.

$$X = (x_1, x_2, \cdots, x_N)$$ $$\begin{cases} positive~x, & x \sim p(x_{t+k} \vert c_t) \newline negative, & x \sim p(x_{t+k}) \end{cases}$$

The loss to optimize is therefore

$$\mathcal{L} = -E_X (\log \frac{f_k(x_{t_k}, c_t)}{\sum_{x_j \in X} f_k(x_j,c_t)})$$

At first glance, this objective looks like ordinary categorical cross-entropy, so it may be unclear why training approaches density estimation. It certainly confused me, though some readers may see it immediately. Consider the optimal probability we ultimately want, $p(d=i \vert X,c_t)$.
As discussed above, the model predicts the relative likelihood that a sample came from the context rather than independently of it, and the logarithm of this value quantifies mutual information. We can therefore express the optimal probability for the given sample set $X$ and context vector $c$ as follows.

$$p(d = i \vert X, c_t) = \frac{p(x_i \vert c_t) \prod_{l \ne i} p(x_l)}{\sum_{j=1}^N p(x_j \vert c_t) \prod_{l \ne j} p(x_l)} = \frac{\frac{p(x_i \vert c_t)}{p(x_i)}}{\sum_{j=1}^N \frac{p(x_j \vert c_t)}{p(x_j)}}$$

Thus, the value predicted by network $f$ in the preceding expression is proportional to the probability ratio, independently of the number of negative samples. Training uses only the loss, but the mutual information can be evaluated through the following lower bound.

$$I(x_{t+k}, c_t) \geq \log(N) - L_N$$

This bound becomes tighter as $N$ increases. Here, tight means that the lower bound approaches the infimum of the actual mutual information—in other words, it better approximates the true value. The appendix contains more details, but expanding all of them here would be excessive, so I will stop the derivation at this point.

Experiments

The model's greatest advantages are that it can use an encoder of arbitrary architecture and that it establishes a loss objective for representation learning. Accordingly, the paper experiments across several modalities.

Audio dataset

LibriSpeech consists of audio files from 251 speakers, each accompanied by a text transcript. Because these transcripts are not aligned with the actual phone sequences, additional annotation is required. When I implemented this in a lab, I wrote a hopeless Python algorithm myself—only to discover there is a Kaldi Toolkit. I wish someone had told me earlier. OpenCV-Python is widely used in computer vision; Kaldi feels like its audio counterpart. At least I learned something useful.

To my surprise, the authors even shared their segmented and aligned dataset on Google Drive. Had I known it existed when training my speech transformer, I would have downloaded it instead. What a regret.
This dataset supports speaker classification, which identifies a speaker from audio, and phone classification, which classifies the phone (syllabic sound) at each time step. CPC performs well against supervised methods on both tasks. The table on the right presents several ablations of the CPC model for phone classification.

This may look like a simple loss curve, but it shows mean phone-prediction accuracy as a function of how many latent steps ahead of the present are predicted. There are 41 possible phone classes, so chance performance is roughly 0.025. One latent step spans 10 ms; the graph appears to show that useful prediction extends to about 20 steps.

Vision

The vision setup is somewhat unusual: it uses a ResNet-v2-101 backbone as the image encoder but removes batch normalization from the original architecture. One interpretation is that batch normalization stabilizes tasks such as classification, where outputs fall within a fixed probability range, but is better avoided in generation-related training that must learn representations. Take that as you will, though normalizing feature maps with batch normalization can indeed make distribution collapse more likely. After unsupervised learning, a linear layer is trained separately to predict ImageNet labels.

Training proceeds as follows.

• Extract a $7 \times 7$ grid of $64 \times 64$ crops from a $256 \times 256$ image. If you are wondering how that works when $64 \times 7=448$, each patch overlaps the next by 32 pixels: $256+(32 \times 6)/64=4+3=7$.

• Encode each crop, then mean-pool each channel to obtain a 1,024-dimensional vector. Since there are $7 \times 7$ crops, the final output is a $7 \times 7 \times 1024$ tensor.

• Perform unsupervised learning with a PixelCNN-style AR model that predicts subsequent pixels. In brief, PixelCNN is a generative model that looks at previous pixels to predict the next one; the figure above makes the setup clear.

• Train a linear classifier on the previously learned CPC feature map. The authors use Adam for CPC and SGD for the linear classifier.

The results appear very strong. Note that Top-1 and Top-5 accuracy here differ somewhat from those in a standard classification task: performance is evaluated from the feature maps learned without supervision in the procedure described above.

Natural Language

Now we return to NLP. Explaining experiments for every modality makes me feel as if I am becoming multimodal myself, but in today's wide-open AI landscape, sampling many fields may be the way to survive. In any case...

First, the model is trained on BookCorpus. Because NLP is inherently autoregressive, refer to the earlier Word2Vec/NCE discussion for the training process. To evaluate words not seen during training, the authors add a linear mapping between Word2Vec embeddings and the embeddings learned by the model.

The classification datasets include MR (movie-review sentiment), CR (customer product reviews), subjective/objective classification, opinion classification (MPQA), and question-type classification (TREC), among others.

The transfer-learning setup matches that of the skip-thought vectors paper, which also determines the comparison baselines. In short, the results show that the method works reasonably well.

Reinforcement learning

At last we have wandered all the way into reinforcement learning. By this point, one might conclude that the authors are unusually obsessive—but really, hats off to DeepMind.

Because RL differs somewhat from ordinary deep learning, it needs a different objective. The experiment uses an A2C agent as the base model and CPC as an auxiliary loss. I do not know the details well enough to elaborate, so I will move on; judging by the strong red curve, the authors appear to consider it effective.

Appendix

$$\begin{aligned} \mathcal{L}\_N^\text{opt} =& -\mathbb{E}\_X \log \left( \frac{\frac{p(x_{t+k} \vert c_t)}{p(x_{t_k})}}{\frac{p(x_{t+k} \vert c_t)}{p(x_{t+k})} + \sum\_{x_j \in X_\text{neg}}\frac{p(x_j \vert c_t)}{p(x_j)}} \right) \newline =& \mathbb{E}\_X \log \left( 1+\frac{p(x_{t+k})}{p(x_{t+k} \vert c_t)} \sum\_{x_j \in X_\text{neg}} \frac{p(x_j \vert c_t)}{p(x_j)} \right) \newline \approx& \mathbb{E}\_X \log \left( 1+\frac{p(x_{t+k})}{p(x_{t+k} \vert c_t)} (N-1) \mathbb{E}\_{x_j} \left( \frac{p(x_j \vert c_t)}{p(x_j)} \right) \right) \newline =& \mathbb{E}\_X \log \left( 1+\frac{p(x_{t+k})}{p(x_{t+k} \vert c_t)} (N-1) \right) \newline \geq& \mathbb{E}\_X \log \left( \frac{p(x_{t+k})}{p(x_{t+k} \vert c_t)} N \right) \newline =& -I(x_{t+k}, c_t)+\log(N) \end{aligned}$$

Derivation of the loss function's lower bound

$$\begin{aligned} \mathbb{E}\_X \left( \log \frac{f(x, c)}{\sum\_{x_j \in X} f(x_j,c)} \right) =& \mathbb{E}\_{(x, c)} \left( F(x, c) \right) - \mathbb{E}\_{(x, c)} \left( \log \sum\_{x_j \in X_\text{neg}} e^{F(x_j, c)} \right) \newline =& \mathbb{E}\_{(x, c)} \left( F(x, c) \right) - \mathbb{E}\_{(x, c)} \left( \log \left( e^{F(x, c)} + \sum\_{x_j \in X_\text{neg}} e^{F(x_j, c)} \right) \right) \newline \leq& \mathbb{E}\_{(x, c)} \left( F(x, c) \right) - \mathbb{E}\_c \left( \log \sum\_{x_j \in X_\text{neg}} e^{F(x_j,c)} \right) \newline =& \mathbb{E}\_{(x, c)} \left( F(x, c) \right) - \mathbb{E}\_c \left( \log \frac{1}{N-1} \sum\_{x_j \in X_\text{neg}} e^{F(x_j, c)} + \log (N-1) \right) \end{aligned}$$

This shows the relationship to neural estimation of mutual information. The equations appear to explain precisely how the authors drew their intuition from mutual information.