Contrastive Learning Intuition: InfoNCE, Temperature, Alignment and Uniformity

Contrastive learning has a reputation for being simple to state and hard to reason about: pull two views of the same thing together, push everything else apart, and somehow useful features emerge. This post makes the "somehow" concrete. It writes down the InfoNCE loss and its gradient, explains the role of the temperature, introduces the alignment–uniformity view that makes the loss interpretable, and then trains a two-dimensional embedding in NumPy where the whole representation can be drawn on a circle.

The loss

Take a batch of $N$ samples and produce two augmented views of each, giving $2N$ embeddings $z_1,\dots ,z_{2N}$ that are ${\ell }_2$-normalized. For view $i$ with partner $j$, the normalized-temperature cross-entropy (NT-Xent, the SimCLR form of InfoNCE) is

Read as a classifier, this is a $(2N-1)$-way softmax over cosine similarities with inverse temperature $1/\tau$, where the correct class is the partner view. Two consequences follow:

• A lower bound on mutual information. With $K$ candidates in the denominator, $I(z_i;z_j)\ge \log K-\mathcal{L}$. The bound is tight only when the loss reaches zero, and it is capped at $\log K$, which is why large batches help: $\log 255\approx 5.5$ nats is the most a batch of 128 pairs can certify.
• A gradient with a clean form. Writing $p_{ik}$ for the softmax weight of candidate $k$ for query $i$, the gradient with respect to the similarity matrix is $(p_{ik}-\mathbf{1}[k=j(i)])/2N$. Positives are pulled, negatives are pushed in proportion to how similar they already are, and the temperature decides how steeply that proportion rises with similarity.

The temperature

Because the embeddings are normalized, $z_i\cdot z_k\in [-1,1]$ and the logits live in $[-1/\tau ,1/\tau ]$. With $\tau =1$ the softmax is nearly uniform, every negative is pushed with roughly the same force, and the loss can barely distinguish a hard negative from an easy one. With $\tau =0.05$ the logits span $[-20,20]$, the softmax is dominated by the few most similar negatives, and the gradient concentrates on them: a small $\tau$ is an automatic hard-negative miner. That is also its risk, since the hardest negatives are often semantically positive samples from the same class, and pushing those apart hurts.

Alignment and uniformity

Wang and Isola showed that the contrastive loss asymptotically optimizes two separable quantities on the unit sphere:

Alignment asks that positive pairs map close together; uniformity asks that the marginal distribution of embeddings spread over the sphere as much as possible (the Gaussian-potential form is minimized by the uniform distribution). The decomposition explains both why the loss works and when it fails: a collapsed encoder that maps everything to one point has perfect alignment and the worst possible uniformity, and the negatives in the denominator are exactly what prevents it.

A NumPy experiment

The data are eight-dimensional vectors drawn from four latent classes with intrinsic within-class spread; each "augmentation" adds independent Gaussian noise, so the two views of a sample share the latent class and the sample's own offset but not the noise. The encoder is a tiny MLP (8 → 32 → 2, tanh) followed by ${\ell }_2$ normalization, so every embedding sits on the unit circle and the whole representation can be drawn. The class labels are never shown to the loss; they are only used afterwards to colour the plots and to measure a 10-nearest-neighbour accuracy.

def nt_xent(z, tau):
    """z: (2N, d) unit-norm rows, views i and i+N are partners. Returns loss and dL/dz."""
    N2 = z.shape[0]; N = N2 // 2
    pos = np.r_[np.arange(N, N2), np.arange(0, N)]          # index of each row's partner
    S = z @ z.T / tau
    np.fill_diagonal(S, -np.inf)                            # never compare a view with itself
    S = S - S.max(1, keepdims=True)
    P = np.exp(S); P /= P.sum(1, keepdims=True)             # row-wise softmax
    loss = -np.mean(np.log(P[np.arange(N2), pos]))
    Y = np.zeros_like(P); Y[np.arange(N2), pos] = 1.0
    dS = (P - Y) / N2                                       # dL / dS
    dz = (dS + dS.T) @ z / tau                              # S is symmetric in z
    return loss, dz

The gradient has to pass back through the normalization and the MLP. The normalization step is the one people get wrong, so it is worth writing out:

Training used Adam (learning rate $5\times {10}^{-3}$) for 600 steps with batches of 128 pairs, for four temperatures. Metrics are computed on 800 held-out pairs.

Three things in this table are worth reading carefully.

Alignment improves by a factor of 20 to 40 at every temperature, and the nearest-neighbour accuracy goes from 0.65 to essentially 1.0 without a single label. That is the headline result of contrastive learning reproduced in miniature.

The loss does not go to zero, and it cannot. With four latent classes and 128 pairs, each view has about 63 other views of its own class among its 254 "negatives". Once the encoder has collapsed each class to a point, those 63 are indistinguishable from the true partner, so the softmax can give the partner at most about $1/64$ of the mass and the loss floors near $\log 64\approx 4.16$. The measured floor, 4.17 at $\tau =0.1$, is within a hundredth of it. On real data the same effect is called the false-negative problem, and it is the reason supervised contrastive and debiased contrastive losses exist.

Uniformity barely moves, and the lowest temperature is the least uniform. On a circle with four classes, the uniformity metric is dominated by how the four arcs are laid out rather than by how tightly they are packed; at $\tau =0.05$ the arcs are tightest and the within-class pairs contribute $e^0=1$ each to the average, which raises the metric. The trend reported on real datasets, where lower temperature gives better uniformity, needs many more latent factors than four. This is a reminder that toy experiments show mechanisms, not benchmarks.

WARNING

The number of negatives, the temperature and the augmentation strength are not independent knobs. A small $\tau$ with weak augmentations lets the encoder solve the task by memorizing low-level details shared between views (exact pixel values, JPEG artefacts, colour histograms), which are perfectly aligned and perfectly useless. The augmentations define what the representation should be invariant to; the loss only enforces it.

Further reading

• Chen et al. (2020), A simple framework for contrastive learning of visual representations (SimCLR): https://arxiv.org/abs/2002.05709
• van den Oord, Li and Vinyals (2018), Representation learning with contrastive predictive coding, where InfoNCE was introduced: https://arxiv.org/abs/1807.03748
• Wang and Isola (2020), Understanding contrastive representation learning through alignment and uniformity on the hypersphere: https://arxiv.org/abs/2005.10242
• He et al. (2019), Momentum contrast for unsupervised visual representation learning (MoCo), on getting many negatives without a huge batch: https://arxiv.org/abs/1911.05722
• Poole et al. (2019), On variational bounds of mutual information, for what the $\log K-\mathcal{L}$ bound does and does not say: https://arxiv.org/abs/1905.06922
• Radford et al. (2021), Learning transferable visual models from natural language supervision (CLIP), the same loss across two modalities: https://arxiv.org/abs/2103.00020