Deep Learning · Representation Learning

InfoNCE and the Temperature Parameter

How the contrastive objective turns similarity scores into a softmax over the batch, why the temperature controls effective negative count and gradient shape, and what breaks when it is set wrong.

A contrastive objective takes two views of the same input, computes similarity between their representations, and asks the model to make that similarity large relative to the similarity between unrelated pairs. The family of losses that does this is usually written as InfoNCE, and almost every design decision in a contrastive training run, including batch size, temperature, and how negatives are drawn, is a consequence of how that loss behaves.

The loss is a classification problem in disguise

Given an anchor representation q, one positive representation k+, and N-1 negative representations k_i, the loss is the negative log of a softmax over similarity scores:

import torch
import torch.nn.functional as F

def info_nce(q, k, temperature=0.07):
    # q: (n, d) anchors; k: (n, d) matching positives, rows aligned
    q = F.normalize(q, dim=1)
    k = F.normalize(k, dim=1)
    logits = q @ k.T / temperature          # (n, n)
    labels = torch.arange(len(q))
    return F.cross_entropy(logits, labels)

The (n, n) score matrix has the positive pair on its diagonal and n² - n off-diagonal entries that are all treated as negatives. Row i is a classification problem: given anchor i, pick which of the n candidates is its positive. The model’s parameters are trained to make that classification correct. Nothing in the loss refers to a label in the original dataset; the label is the index itself, which is why the method is self-supervised.

The normalization step matters. Without it, the loss can be reduced by shrinking all representations toward zero, since scaling both q and k by a small constant reduces every logit toward zero and makes the softmax uniform. Normalizing to the unit sphere removes that direction of degeneracy and turns the objective into a problem about angles.

Written out for one pair, the loss is

L_i = -log( exp(sim(q_i, k_i)/τ) / Σ_j exp(sim(q_i, k_j)/τ) )

and the gradient with respect to the positive score has a simple form: it is -(1 - p_i)/τ, where p_i is the softmax probability the model currently assigns to the correct candidate. When the positive is already well separated, p_i is near one, the gradient nearly vanishes, and the pair stops contributing. When the positive is confused with a negative, the gradient is large. The loss concentrates effort on the pairs that are currently wrong, which is the behaviour of any softmax cross-entropy and the reason hard negatives are valuable.

Temperature is a scale on the logits, with three effects

The temperature τ divides the similarity scores before the softmax. Because cosine similarities lie in [-1, 1], dividing by τ = 0.07 produces logits spanning roughly [-14, 14], a range comparable to the logits of a classifier over a modest number of classes. Dividing by τ = 1 produces logits in [-1, 1], so the softmax is nearly uniform and every negative receives almost the same weight.

The first effect is on the effective number of negatives. The softmax’s entropy is set by the spread of the logits; a smaller τ sharpens the distribution, so the loss is dominated by the few negatives with the highest similarity. Raising τ flattens the distribution and makes the loss average over many negatives. Because the information-theoretic reading of InfoNCE bounds a mutual-information estimate by log(N) minus the loss, a sharper softmax extracts more from the same batch, which is why small temperatures are associated with being able to use smaller batches.

The second effect is on the gradient distribution. With a sharp softmax, the gradient flows almost entirely to the hardest negatives and their positives. This is efficient when the hard negatives are genuine, and destructive when they are false negatives, a case treated in False Negatives and Batch Composition. With a flat softmax, gradients spread over all negatives, which is less sensitive to mislabelled pairs but wastes capacity on pairs the model already separates.

The third effect is numerical. In half precision, exp of a large positive logit overflows, and the softmax implementations that subtract the maximum logit before exponentiating handle it. Dividing similarities by a very small τ pushes logits to the edge of the representable range, and the resulting gradients can exceed the range of the format. Frameworks that compute the loss in reduced precision usually upcast the logits to 32-bit floats for the softmax for this reason.

τ is a hyperparameter, not a learned quantity in the standard formulation, and it is not interchangeable between datasets. A value tuned for one batch size and augmentation policy does not transfer directly to another, because the temperature’s effect is relative to the distribution of similarities that the encoder produces, and that distribution shifts as the encoder improves.

Symmetrization and the two-encoder variants

The loss above is asymmetric: row i uses anchor i against all of k. Applying the same computation transposed, with k as anchors against q, and averaging the two gives a symmetric loss. The symmetric version doubles the number of terms and makes the objective invariant to which view is called the anchor, which matters when the two views are produced by different transformations.

Two-encoder designs, where the two views pass through separate networks, appear in several forms. One shares parameters between the branches, so the same encoder processes both views and the loss is computed on a single forward pass with the batch doubled. Another uses a momentum-updated copy of the encoder for one branch, with gradients blocked through it, which stabilizes the targets; the mechanism is described in Momentum Encoders and Stop-Gradients. A third keeps separate encoders for genuinely different modalities, where parameter sharing is impossible.

Batch size and temperature trade against each other. Doubling the batch adds N new negatives, which raises the log(N) ceiling on what the loss can extract, and this is one reason large batches are associated with better representations. The cost is quadratic in the score matrix: a batch of 4,096 produces a 4,096 × 4,096 matrix of 16.8 million entries per forward pass, which at 32 bits is 67 MB before any gradient storage. Memory-efficient implementations compute the loss in chunks over the batch dimension so the full matrix never exists at once.

Global batch versus local queue

Negatives need not come from the current batch. A memory bank or queue stores representations from recent batches and samples negatives from it, which decouples the number of negatives from the batch size. The trade is freshness: queued representations were produced by an older version of the encoder, so they are not exactly the outputs the current model would assign to those inputs. If the encoder changes quickly, queued negatives become stale and the loss compares current positives against outdated negatives, which weakens the signal. Momentum encoders address this by producing the queued representations with a slowly updated copy, so the queue stays consistent with the target network even as the online network moves.

The queue also changes what counts as a negative. An item that was the positive for a different anchor in an earlier batch is a negative now, and if the two items are semantically related, the loss pushes apart representations that should be close. Deduplication, either by removing exact duplicates from the batch or by masking out known-similar pairs, is the standard mitigation.

Failure modes that look like optimization problems

A loss that decreases while retrieval quality stagnates is the most common symptom. It usually means the model is separating the specific pairs in the batch without building a space that generalizes, which happens when the negatives are too easy: the softmax saturates on pairs that are trivially different, the gradient vanishes, and the remaining hard cases never receive pressure. Monitoring nearest-neighbour consistency on held-out data, rather than the loss value, is the way to detect this.

A loss that spikes and recovers repeatedly can come from a batch in which a single hard negative dominates the softmax for many anchors at once. Because the gradient for that negative is summed over all rows, its representation moves a large distance in one step, and the next batch sees a different geometry. Gradient clipping and a larger effective temperature both reduce the amplitude.

Collapse to a constant representation is the other extreme. If the encoder maps every input to nearly the same vector, all similarities are near one, the softmax is uniform, and the loss settles at log(N) with zero gradient. This is a local optimum of the loss as written, and preventing it requires either asymmetry between the branches, a stop-gradient, or an explicit term that penalizes the degenerate solution. The mechanisms that produce and prevent it are the subject of Collapse and Redundancy Reduction in Self-Supervised Training, and the complementary question of whether the resulting space is well spread is covered by Alignment and Uniformity as Separate Objectives.

References

  1. Representation Learning with Contrastive Predictive Coding — van den Oord, Li & Vinyals, 2018
  2. A Simple Framework for Contrastive Learning of Visual Representations — Chen et al., ICML 2020
  3. What Makes for Good Views for Contrastive Learning? — Tian, Sun, Poole & Krishnan, NeurIPS 2020