Noise-Contrastive Estimation and InfoNCE: How Partition Function Elimination and Mutual Information Lower Bounds Power Modern AI

In statistical machine learning and generative modeling, evaluating the exact probability of an observed event frequently requires calculating a normalizing constant known as the partition function. For continuous spaces or discrete spaces with high cardinality, such as a natural language vocabulary spanning over one hundred thousand tokens or high-dimensional pixel distributions, computing this partition function directly is computationally intractable.

Noise-Contrastive Estimation (NCE), introduced by Michael Gutmann and Aapo Hyvarinen in 2010, transformed unnormalized statistical modeling by converting density estimation into a supervised binary classification task between true data and synthetic noise. In 2018, Aaron van den Oord, Yazhe Li, and Oriol Vinyals generalized this framework into the InfoNCE loss within Contrastive Predictive Coding (CPC), establishing a mathematical bridge between multi-class classification and variational lower bounds on mutual information.

From word embeddings and audio pre-training to contrastive vision-language architectures like CLIP and dense text retrievers, NCE and InfoNCE remain fundamental mathematical primitives of modern self-supervised representation learning.

The Partition Function Bottleneck

Consider an energy-based model or a probabilistic language model parameterized by $\theta$. The probability of a data point $x \in \mathcal{X}$ is expressed via an unnormalized score function $s_\theta(x)$ or energy $E_\theta(x) = -s_\theta(x)$:

$$p_\theta(x) = \frac{\exp(s_\theta(x))}{Z(\theta)}$$

where the partition function $Z(\theta)$ is defined as:

$$Z(\theta) = \sum_{x' \in \mathcal{X}} \exp(s_\theta(x')) \quad \text{or} \quad Z(\theta) = \int_{\mathcal{X}} \exp(s_\theta(x')) \, dx'$$

When training under standard Maximum Likelihood Estimation (MLE), the objective minimizes the negative log-likelihood:

$$\mathcal{L}_{\text{MLE}}(\theta) = -\mathbb{E}_{x \sim p_{\text{data}}} [\log p_\theta(x)] = -\mathbb{E}_{x \sim p_{\text{data}}} [s_\theta(x)] + \log Z(\theta)$$

Taking the gradient with respect to $\theta$ yields:

$$\nabla_\theta \mathcal{L}_{\text{MLE}}(\theta) = -\mathbb{E}_{x \sim p_{\text{data}}} [\nabla_\theta s_\theta(x)] + \mathbb{E}_{x \sim p_\theta} [\nabla_\theta s_\theta(x)]$$

The second expectation requires computing expectations over the model distribution $p_\theta$ at every single optimization step. In continuous spaces, this necessitates Markov Chain Monte Carlo (MCMC) sampling or Contrastive Divergence, which is notoriously slow, unstable, and prone to mode collapse. In discrete spaces with vocabulary size $|V|$, computing $\log Z(\theta)$ requires evaluating the full softmax denominator over all $|V|$ classes, which becomes a severe computational bottleneck during training.

Noise-Contrastive Estimation (NCE)

Gutmann and Hyvarinen (2010) resolved this bottleneck by redefining the learning objective. Instead of directly modeling the normalized density $p_\theta(x)$, NCE trains a non-linear binary classifier to discriminate between genuine samples drawn from the true data distribution $p_{\text{data}}$ and synthetic noise samples drawn from a known, tractable noise distribution $q_n$.

Suppose for every genuine data sample $x \sim p_{\text{data}}$, we draw $\nu$ noise samples $y_1, \dots, y_\nu \sim q_n$. The joint dataset contains data points labeled with $D=1$ (true data) with prior probability $P(D=1) = \frac{1}{1 + \nu}$ and noise points labeled with $D=0$ with prior probability $P(D=0) = \frac{\nu}{1 + \nu}$.

Applying Bayes' theorem, the posterior probability that a given sample $u$ came from the true data distribution is:

$$P(D=1 | u; \theta) = \frac{p_\theta(u) P(D=1)}{p_\theta(u) P(D=1) + q_n(u) P(D=0)} = \frac{p_\theta(u)}{p_\theta(u) + \nu q_n(u)}$$

Letting $p_\theta(u) = \exp(s_\theta(u) - c)$, where $c$ is a learned scalar representing the log partition function $\log Z$, the posterior becomes:

$$P(D=1 | u; \theta) = \sigma(s_\theta(u) - c - \log(\nu q_n(u)))$$

where $\sigma(z) = \frac{1}{1 + \exp(-z)}$ denotes the standard sigmoid activation function.

The NCE objective function maximizes the log-posterior likelihood of the binary classification task:

$$\mathcal{J}_{\text{NCE}}(\theta, c) = \mathbb{E}_{x \sim p_{\text{data}}} [\log P(D=1 | x; \theta, c)] + \nu \mathbb{E}_{y \sim q_n} [\log P(D=0 | y; \theta, c)]$$

Expanding this into sigmoid form yields:

$$\mathcal{J}_{\text{NCE}}(\theta, c) = \mathbb{E}_{x \sim p_{\text{data}}} \left[ \log \sigma(s_\theta(x) - c - \log(\nu q_n(x))) \right] + \nu \mathbb{E}_{y \sim q_n} \left[ \log(1 - \sigma(s_\theta(y) - c - \log(\nu q_n(y)))) \right]$$

Key Mathematical Properties of NCE

1. Consistent Parameter Estimation: Gutmann and Hyvarinen proved that under mild regularity conditions, if the model family contains the true data distribution $p_{\text{data}}$, the NCE estimator $\hat{\theta}_{\text{NCE}}$ converges almost surely to the true parameter $\theta^*$ as sample size $N \to \infty$.
2. Asymptotic Convergence to MLE: As the noise ratio $\nu \to \infty$, the asymptotic variance of the NCE estimator approaches the Cramer-Rao lower bound, matching the statistical efficiency of exact Maximum Likelihood Estimation without ever summing over the state space.
3. Self-Normalization in Neural Language Modeling: In 2012, Andriy Mnih and Yee Whye Teh applied NCE to train neural probabilistic language models. They discovered that by fixing $c = 0$ (setting the partition function $Z = 1$), neural networks naturally learn self-normalizing output distributions with negligible loss in modeling perplexity, completely bypassing the softmax normalization during inference.

       True Data Distribution p_data(x) --------> Positive Sample (D=1)
                                                              |
                                                              v
       Noise Distribution q_n(y) ---------------> Negative Samples (D=0)
                                                              |
                                                              v
                              [ Binary Logistic Classifier ]
                                             |
                                             v
                           L_NCE = Log-Loss(Data vs Noise)

Negative Sampling vs. Formal NCE

In 2013, Tomas Mikolov and colleagues introduced Skip-Gram with Negative Sampling (SGNS) in Word2Vec. While SGNS is frequently described as NCE, it is a simplified heuristic approximation.

SGNS simplifies the objective by omitting the explicit noise density ratio $\log(\nu q_n(u))$:

$$\mathcal{L}_{\text{SGNS}} = \log \sigma(v_{w_O}'^\top v_{w_I}) + \sum_{i=1}^k \mathbb{E}_{w_i \sim P_n(w)} \left[ \log \sigma(-v_{w_i}'^\top v_{w_I}) \right]$$

Because SGNS ignores the noise distribution terms inside the sigmoid, it does not estimate a normalized probability distribution. As proved by Omer Levy and Yoav Goldberg in 2014, SGNS instead computes an implicit low-rank matrix factorization of the Pointwise Mutual Information (PMI) matrix shifted by $\log k$.

Multi-Class Generalization: InfoNCE and Contrastive Predictive Coding

While classical NCE operates as a binary classification problem between one data point and $\nu$ noise points, modern self-supervised learning frequently structures the problem as categorical selection across multiple candidate samples.

In 2018, Aaron van den Oord, Yazhe Li, and Oriol Vinyals introduced Contrastive Predictive Coding (CPC) and formulated the InfoNCE objective. Given a context representation $c$ and a set of $K$ candidate samples $X = \{x_1, x_2, \dots, x_K\}$ containing exactly one positive target $x_+$ drawn from the conditional distribution $p(x|c)$ and $K-1$ negative samples drawn from the proposal distribution $p(x)$, the InfoNCE loss is defined as:

$$\mathcal{L}_{\text{InfoNCE}} = -\mathbb{E}_{X} \left[ \log \frac{f_\theta(x_+, c)}{\sum_{i=1}^K f_\theta(x_i, c)} \right]$$

Here, $f_\theta(x, c)$ is a parameterized density ratio critic. Commonly, $f_\theta(x, c) = \exp\left( \frac{\text{sim}(g(x), h(c))}{\tau} \right)$, where $\text{sim}(\cdot, \cdot)$ is cosine similarity, $g$ and $h$ are neural encoders, and $\tau > 0$ is a learnable or fixed temperature parameter.

Proof: InfoNCE as a Lower Bound on Mutual Information

A fundamental theoretical contribution of CPC is proving that minimizing $\mathcal{L}_{\text{InfoNCE}}$ maximizes a variational lower bound on the mutual information $I(X; C)$ between the context $C$ and the target variable $X$.

The mutual information between two random variables $X$ and $C$ is defined as:

$$I(X; C) = \sum_{x, c} p(x, c) \log \frac{p(x, c)}{p(x) p(c)} = \mathbb{E}_{p(x, c)} \left[ \log \frac{p(x|c)}{p(x)} \right]$$

Let $X = \{x_1, \dots, x_K\}$ be a set of $K$ samples where index $j$ is the true positive sample generated from $p(x|c)$, and all remaining $K-1$ samples are drawn independently from $p(x)$. The probability that index $j$ is the positive sample given set $X$ and context $c$ is:

$$P(j = \text{pos} | X, c) = \frac{p(x_j | c) \prod_{l \neq j} p(x_l)}{\sum_{i=1}^K p(x_i | c) \prod_{l \neq i} p(x_l)} = \frac{\frac{p(x_j | c)}{p(x_j)}}{\sum_{i=1}^K \frac{p(x_i | c)}{p(x_i)}}$$

Setting the optimal density ratio critic $f^*(x, c) \propto \frac{p(x|c)}{p(x)}$, the expected categorical cross-entropy loss becomes:

$$\mathcal{L}_{\text{InfoNCE}}^* = -\mathbb{E}_{X, c} \left[ \log \frac{\frac{p(x_+ | c)}{p(x_+)}}{\frac{p(x_+ | c)}{p(x_+)} + \sum_{i=1}^{K-1} \frac{p(x_{i,-} | c)}{p(x_{i,-})}} \right]$$

Rewriting the argument of the logarithm:

$$\log \left( \frac{\frac{p(x_+ | c)}{p(x_+)}}{\frac{p(x_+ | c)}{p(x_+)} + \sum_{i=1}^{K-1} \frac{p(x_{i,-} | c)}{p(x_{i,-})}} \right) = -\log \left( 1 + \frac{p(x_+)}{p(x_+ | c)} \sum_{i=1}^{K-1} \frac{p(x_{i,-} | c)}{p(x_{i,-})} \right)$$

Taking the expectation over the negative samples $x_{i,-} \sim p(x)$, since $\mathbb{E}_{x \sim p(x)}\left[\frac{p(x|c)}{p(x)}\right] = \int p(x) \frac{p(x|c)}{p(x)} dx = \int p(x|c) dx = 1$:

$$\mathbb{E}_{x_{1 \dots K-1}} \left[ \sum_{i=1}^{K-1} \frac{p(x_{i,-} | c)}{p(x_{i,-})} \right] = K - 1$$

Applying Jensen's inequality to the concave logarithm function yields:

$$\mathbb{E}_X \left[ \log \frac{f(x_+, c)}{\sum_{i=1}^K f(x_i, c)} \right] \le -\mathbb{E}_{p(x, c)} \left[ \log \left( 1 + \frac{p(x)}{p(x|c)} (K - 1) \right) \right] \approx \mathbb{E}_{p(x, c)} \left[ \log \frac{p(x|c)}{p(x)} \right] - \log K$$

Therefore:

$$\mathcal{L}_{\text{InfoNCE}} \ge -\left( I(X; C) - \log K \right) \implies I(X; C) \ge \log K - \mathcal{L}_{\text{InfoNCE}}$$

  Context c ---------> Encoder h(c) ------------\
                                                  Cosine Similarity Matrix
  Target x_+ --------> Encoder g(x_+) ----------/        (scaled by 1/tau)
                                                         |
  Distractors x_- ---> Encoder g(x_-) -----------------> |
                                                         v
                                              Categorical Softmax
                                                         |
                                                         v
                                              Loss = -log P(Positive)
                                                         |
                                                         v
                                         Mutual Information Bound:
                                            I(X; C) >= log(K) - L

The $\log K$ Mutual Information Bottleneck

The inequality $I(X; C) \ge \log K - \mathcal{L}_{\text{InfoNCE}}$ exposes a crucial structural limitation analyzed by Ben Poole et al. (2019) and David McAllester and Karl Stratos (2020): the InfoNCE lower bound is mathematically capped at $\log K$.

Even if the true mutual information between representations is exceptionally high (e.g., in paired audio or image crops where $I(X;C) \gg 20$ nats), an InfoNCE estimator trained with $K=256$ negative samples can never report or leverage a mutual information bound greater than $\log(256) \approx 5.54$ nats.

This mathematical property explains the evolutionary trajectory of modern contrastive architectures:

• SimCLR (Chen et al., 2020): Relied on massive batch sizes ($N = 4096 \implies K = 8190$) on specialized TPU clusters to push the $\log K$ boundary.
• MoCo (He et al., 2020): Decoupled dictionary size from mini-batch size using a dynamic memory queue ($K = 65,536 \implies \log K \approx 11.09$ nats).
• CLIP (Radford et al., 2021): Utilized synchronized distributed matrix multiplications across hundreds of GPUs to compute contrastive losses across batches of 32,768 text-image pairs.

Contrastive Learning in Modern Architectures

The principles established by NCE and InfoNCE govern several major branches of contemporary AI systems:

1. Vision-Language Alignment (CLIP and SigLIP)

OpenAI's CLIP trains dual encoders $g_\theta$ (image) and $h_\phi$ (text) across a batch of $N$ image-text pairs using symmetric InfoNCE:

$$\mathcal{L}_{\text{CLIP}} = \frac{1}{2N} \sum_{i=1}^N \left( -\log \frac{\exp(\langle u_i, v_i \rangle / \tau)}{\sum_{j=1}^N \exp(\langle u_i, v_j \rangle / \tau)} - \log \frac{\exp(\langle u_i, v_i \rangle / \tau)}{\sum_{j=1}^N \exp(\langle u_j, v_i \rangle / \tau)} \right)$$

In 2023, Xiaohua Zhai and colleagues introduced SigLIP (Sigmoid Loss for Language-Image Pre-training), which replaces the categorical InfoNCE multi-class softmax with pairwise binary classification. By evaluating each pair independently via sigmoid loss, SigLIP eliminates the global normalization denominator, removing cross-device communication overhead during distributed training and improving scaling efficiency.

2. Dense Information Retrieval and Embedding Models

Modern embedding models (e.g., BGE, E5, Sentence-Transformers, ColBERT) optimize dense text retrieval using InfoNCE with hard negative mining:

$$\mathcal{L}_{\text{Retrieval}} = -\log \frac{\exp(\text{sim}(q, d^+) / \tau)}{\exp(\text{sim}(q, d^+) / \tau) + \sum_{j=1}^M \exp(\text{sim}(q, d_j^-) / \tau)}$$

By pairing standard in-batch negatives with hard negatives mined via lexical BM25 search or cross-encoder re-rankers, InfoNCE forces the embedding space to separate semantically distinct documents with identical keyword profiles.

3. Self-Supervised Speech and Audio (wav2vec 2.0)

Meta's wav2vec 2.0 uses InfoNCE over masked continuous latent speech representations. The model quantizes latent speech features into discrete codebook entries and optimizes an InfoNCE critic to select the true quantized representation corresponding to masked audio frames among distractors sampled from other time steps in the same utterance.

Practical Implementation: InfoNCE in PyTorch

Below is a complete, numerical-precision-stabilized PyTorch implementation of the symmetric InfoNCE loss with in-batch negatives:

import torch
import torch.nn as nn
import torch.nn.functional as F

class InfoNCELoss(nn.Module):
    """
    Symmetric InfoNCE Loss for dual-encoder representation learning.
    Computes contrastive loss using in-batch negatives.
    """
    def __init__(self, temperature: float = 0.07, learnable_temp: bool = False):
        super().__init__()
        if learnable_temp:
            self.log_inv_tau = nn.Parameter(torch.tensor(torch.log(torch.tensor(1.0 / temperature))))
        else:
            self.register_buffer("log_inv_tau", torch.tensor(torch.log(torch.tensor(1.0 / temperature))))

    def forward(self, query_embeddings: torch.Tensor, key_embeddings: torch.Tensor) -> torch.Tensor:
        """
        Args:
            query_embeddings: Tensor of shape (batch_size, embedding_dim)
            key_embeddings: Tensor of shape (batch_size, embedding_dim)
        Returns:
            Scalar tensor representing symmetric contrastive cross-entropy loss.
        """
        # Step 1: Normalize embeddings onto unit hypersphere (L2 normalization)
        queries = F.normalize(query_embeddings, p=2, dim=-1)
        keys = F.normalize(key_embeddings, p=2, dim=-1)

        # Step 2: Compute cosine similarity matrix: shape (batch_size, batch_size)
        inv_tau = torch.exp(self.log_inv_tau)
        logits = torch.matmul(queries, keys.transpose(0, 1)) * inv_tau

        # Step 3: Ground-truth target indices lie along the main diagonal
        batch_size = queries.size(0)
        labels = torch.arange(batch_size, device=queries.device, dtype=torch.long)

        # Step 4: Compute symmetric cross-entropy loss across rows and columns
        loss_queries = F.cross_entropy(logits, labels)
        loss_keys = F.cross_entropy(logits.transpose(0, 1), labels)

        return (loss_queries + loss_keys) / 2.0

Comparison of Optimization Approaches

• Standard Maximum Likelihood Estimation (MLE): Solves exact generative density estimation. Requires computing the full partition function sum or evaluating an intractable integral over the state space. Minimizes forward KL divergence $D_{\text{KL}}(p_{\text{data}} \parallel p_\theta)$. Used primarily in autoregressive language generation and small-vocabulary classification.
• Noise-Contrastive Estimation (NCE): Converts density estimation into binary classification discriminating true data from synthetic noise. Replaces the partition function with a learned parameter $c = \log Z$ or fixes $c=0$ in neural language models. Eliminates the partition function sum while maintaining asymptotic consistency with MLE as the noise sample count grows.
• InfoNCE and Contrastive Predictive Coding (CPC): Generalizes NCE to multi-class categorical selection across positive targets and negative distractors. Cancels the partition function entirely in the probability ratio. Maximizes a variational lower bound on mutual information $I(X; C) \ge \log K - \mathcal{L}_{\text{InfoNCE}}$. Powers multimodal alignment (CLIP), self-supervised vision (SimCLR, MoCo), audio pre-training (wav2vec 2.0), and dense text retrieval.

By transforming intractable partition function integrals into classification objectives and establishing rigorous variational bounds on mutual information, Noise-Contrastive Estimation and InfoNCE provided the mathematical foundation that unlocked large-scale self-supervised learning across modalities.

Sources

• Noise-contrastive estimation: A new estimation principle for unnormalized statistical models (Gutmann & Hyvarinen, AISTATS 2010)
• A Fast and Simple Algorithm for Training Neural Probabilistic Language Models (Mnih & Teh, ICML 2012)
• Representation Learning with Contrastive Predictive Coding (van den Oord, Li, & Vinyals, 2018)
• On Variational Bounds of Mutual Information (Poole et al., ICML 2019)
• Learning Transferable Visual Models From Natural Language Supervision (Radford et al., 2021)
• Sigmoid Loss for Language-Image Pre-Training (Zhai et al., ICCV 2023)
• Neural Word Embedding as Implicit Matrix Factorization (Levy & Goldberg, NeurIPS 2014)