Contrastive Learning

Contrastive learning trains an encoder so that the embeddings of a positive pair are close and the embeddings of negatives are far apart. The positive pair can be two augmented crops of one photograph, two face images of the same person, or an image and its caption; the negatives are the rest of the batch. It is one family of self-supervised learning methods, beside non-contrastive and masked ones.

Purpose

Retrieval, feature matching, face verification and nearest-neighbor classification need embeddings in which distance means similarity, ideally without labels. The negatives are essential: Hadsell, Chopra and LeCun (2006) pointed out that a loss that only pulls similar pairs together is minimized by a constant mapping, which the term pushing dissimilar pairs apart prevents.

When to Use It

It fits when positives are cheap to define (augmentations of unlabeled images, a shared identity or label, image–text pairs) and embeddings will be compared by cosine similarity. It fits less well when compute is tight, since the methods below use thousands of negatives per step, or, in our reasoning, when random negatives would often show the same thing as the anchor (near-duplicates, few classes).

How It Works

From margins to a softmax over negatives

The first losses compared one pair or triplet at a time: the margin loss of Hadsell et al. (2006) and FaceNet’s triplet loss, both given in the loss function article. FaceNet chose semi-hard negatives, farther than the positive but inside the margin, because the hardest ones could collapse the model early in training (Schroff, Kalenichenko & Philbin, 2015).

Sohn (2016) argued that one negative per update makes convergence slow, and proposed the multi-class N-pair loss, which identifies the positive among $N-1$ negatives with a softmax and reuses the other pairs in a batch of $N$ pairs as negatives. Van den Oord, Li and Vinyals (2018) named the same softmax form InfoNCE and derived that, for $N$ samples containing one positive, it bounds the mutual information between the context and the sample: $I \geq \log N - \mathcal{L}_N$, a bound that tightens as $N$ grows. Wu et al. (2018) turned this into instance discrimination: every training image is its own class, its features are stored in a memory bank, and noise-contrastive estimation with 4,096 sampled negatives approximates the softmax over about 1.2 million ImageNet images.

The NT-Xent loss

In the two-view notation of the self-supervised learning article (views $v = t(x)$ with $t \sim \mathcal{T}$, $h = f_\theta(v)$, $z = g(h)$, cosine similarity $\mathrm{sim}$), a batch of $N$ images gives $2N$ embeddings, and SimCLR’s NT-Xent (“normalized temperature-scaled cross entropy”) loss for view $i$ with positive partner $j$ is

$$\ell_{i,j} = -\log \frac{\exp\big(\mathrm{sim}(z_i, z_j)/\tau\big)}{\sum_{k=1}^{2N} \mathbb{1}[k \neq i]\, \exp\big(\mathrm{sim}(z_i, z_k)/\tau\big)},$$

averaged over all $2N$ ordered positive pairs (Chen et al., 2020). It is a $(2N-1)$-way cross-entropy whose logits are similarities divided by the temperature $\tau$, with the partner as the correct class and the other $2(N-1)$ views as negatives. An encoder that cannot tell views apart pays $\log(2N - 1)$.

With $s_{ik} = \mathrm{sim}(z_i, z_k)$ and $P_{ik}$ the softmax probability of view $k$ in row $i$, as for any softmax cross-entropy

$$\frac{\partial \ell_{i,j}}{\partial s_{ik}} = \frac{P_{ik} - \mathbb{1}[k = j]}{\tau}.$$

The positive is pulled with strength $(1 - P_{ij})/\tau$ and each negative pushed with strength $P_{ik}/\tau$; the pushes sum to the pull.

Temperature and hard negatives

Because $P_{ik} \propto \exp(s_{ik}/\tau)$, the push concentrates on the negatives most similar to the anchor, the hard negatives. Chen et al. (2020) noted that margin and logistic losses lack this weighting and need semi-hard mining instead; in their comparison NT-Xent reached 63.9% ImageNet linear accuracy against 57.9% for the best of these, a logistic loss with semi-hard mining. Wang and Liu (2021) analyzed this as a hardness-aware loss: a smaller $\tau$ concentrates the pushes, and as $\tau \to 0^+$, $\tau\,\ell_{i,j}$ tends to $\max(0, s_{\text{neg}} - s_{ij})$ for the most similar negative, a zero-margin triplet loss (training on that limit learned nothing useful); as $\tau$ grows, all negatives are pushed about equally. They also described a uniformity–tolerance dilemma: a small $\tau$ spreads embeddings uniformly but also pushes apart semantically similar images, which a good representation should keep close.

The temperature must be tuned: in SimCLR’s ablation, $\tau = 0.1$ gave 64.4% linear accuracy against 59.7% at 0.05 and 58.0% at 1; without $\ell_2$ normalization the contrastive task was solved more often but the representation was worse (Chen et al., 2020, Table 5). MoCo v2 found $\tau = 0.2$ best with an MLP head (Chen, Fan, Girshick & He, 2020).

Alignment and uniformity

Wang and Isola (2020) named two properties the loss optimizes. Alignment measures how close positives are, $\mathcal{L}_{\text{align}} = \mathbb{E}\,\lVert f(x) - f(y) \rVert_2^2$ over positive pairs; uniformity measures how evenly embeddings spread, $\mathcal{L}_{\text{uniform}} = \log \mathbb{E}\, e^{-2\lVert f(x) - f(y) \rVert_2^2}$ over random pairs (their parameters $\alpha = 2$, $t = 2$). They proved that as the number of negatives grows, the contrastive loss (offset by $\log M$ for $M$ negatives) converges to a term minimized exactly by perfectly aligned encoders plus a term minimized by perfectly uniform ones, when those exist. Encoders with low values of both performed best downstream in their experiments, and optimizing the two terms directly outperformed the contrastive loss. Citing earlier work, they also noted that maximizing a tighter mutual-information bound often gives worse representations, so the bound does not explain the loss’s success.

How to Apply It

SimCLR

SimCLR takes its negatives from the batch, with no memory bank (Chen et al., 2020):

• Augmentations: random crop with resize (8–100% of the area) and flip, strong color distortion including random grayscale, and Gaussian blur on half the images.
• Encoder and head: a ResNet-50 gives $h$ (2,048-D) after average pooling; a two-layer MLP head gives a 128-D $z$.
• Loss and batch: NT-Xent with batches of 256 to 8,192 (16,382 negatives per positive pair at 8,192), the LARS optimizer for stable large-batch training, and batch normalization statistics aggregated over all devices (“global BN”).

A nonlinear head was 3% better than a linear one, and $h$ was more than 10% better than $z$. The authors conjectured that the loss makes $z$ invariant to the augmentations, discarding information such as color or orientation that $h$ keeps; in their Table 3, a rotation applied during pre-training could be predicted from $h$ with 67.6% accuracy, against 25.6% (chance 25%) from $g(h)$.

MoCo

MoCo decouples the number of negatives from the batch size (He et al., 2020): keys from recent batches stay in a queue of $K$ entries ($K = 65{,}536$ for the main results). The key encoder receives no gradient; its parameters follow the query encoder as a moving average, $\theta_k \leftarrow m\,\theta_k + (1 - m)\,\theta_q$, with $m = 0.999$, so keys from different batches stay consistent. With $K = 4{,}096$, accuracy was 59.0% at $m = 0.999$ and 55.2% at $m = 0.9$, and training failed at $m = 0$; a memory bank was 2.6 points worse. MoCo v2 added SimCLR’s MLP head, blur and a cosine learning-rate schedule; with 200 epochs at batch 256 it reached 67.5%, against SimCLR’s 61.9% at batch 256 and 66.6% at 8,192 (Chen, Fan, Girshick & He, 2020, Table 2).

Example: NT-Xent, its gradient, and the temperature

The script computes NT-Xent for two shifted, noisy views of four $16 \times 16$ shape images (a blur stands in for $g(f(\cdot))$), checks the hand-written gradient against finite differences, and shows how the push is shared among negatives at three temperatures; square B, a slightly smaller square, is the hard negative. It runs instantly with NumPy 2.3.

import numpy as np

rng = np.random.default_rng(0)
S = 16
yy, xx = np.mgrid[0:S, 0:S] - 7.5
shapes = {                                     # four 16x16 grayscale images
    "square A": (abs(xx) < 5) & (abs(yy) < 5),
    "square B": (abs(xx) < 4) & (abs(yy) < 4),   # a smaller square: a hard negative
    "bar":      (abs(xx) < 7) & (abs(yy) < 1.5),
    "ring":     (xx**2 + yy**2 < 49) & (xx**2 + yy**2 > 25),
}

def view(img):
    """Augmentation t ~ T: a random shift of up to one pixel plus pixel noise."""
    dy, dx = rng.integers(-1, 2, size=2)
    return np.roll(img.astype(float), (dy, dx), axis=(0, 1)) + 0.1 * rng.standard_normal((S, S))

g = np.exp(-0.5 * (np.arange(-4, 5) / 1.5) ** 2); g /= g.sum()
def encode(v):
    """Stand-in for g(f(v)): Gaussian blur, flatten, subtract a constant."""
    b = np.apply_along_axis(np.convolve, 0, v, g, "same")
    return np.apply_along_axis(np.convolve, 1, b, g, "same").ravel() - 0.2

names = [f"{k} view {i}" for k in shapes for i in (1, 2)]
Z = np.stack([encode(view(img)) for img in shapes.values() for _ in range(2)])  # 2N = 8 rows
partner = np.arange(len(Z)) ^ 1                 # views 2k and 2k+1 come from image k

def nt_xent(Z, tau):
    """Mean NT-Xent over all 2N views and its gradient w.r.t. the unnormalized embeddings."""
    n2 = len(Z)
    norms = np.linalg.norm(Z, axis=1, keepdims=True)
    U = Z / norms
    logits = U @ U.T / tau
    np.fill_diagonal(logits, -np.inf)          # the indicator k != i
    logits -= logits.max(1, keepdims=True)
    P = np.exp(logits)
    P /= P.sum(1, keepdims=True)               # row i: softmax over the other 2N - 1 views
    loss = -np.log(P[np.arange(n2), partner]).mean()
    G = P.copy()                               # d loss / d sim(i, k) = (P_ik - [k = j(i)]) / (tau 2N)
    G[np.arange(n2), partner] -= 1
    G /= tau * n2
    dU = (G + G.T) @ U                         # sim(i, k) depends on both u_i and u_k
    dZ = (dU - (dU * U).sum(1, keepdims=True) * U) / norms   # back through z / ||z||
    return loss, dZ, P

loss, dZ, _ = nt_xent(Z, 0.2)
num = np.zeros_like(Z)
for idx in np.ndindex(*Z.shape):
    Zp, Zm = Z.copy(), Z.copy()
    Zp[idx] += 1e-6
    Zm[idx] -= 1e-6
    num[idx] = (nt_xent(Zp, 0.2)[0] - nt_xent(Zm, 0.2)[0]) / 2e-6
print(f"NT-Xent loss at tau = 0.2: {loss:.4f}")
print(f"max |analytic - finite-difference gradient|: {abs(num - dZ).max():.1e}")

U = Z / np.linalg.norm(Z, axis=1, keepdims=True)
neg = np.arange(2, len(Z))                     # anchor: view 1 of square A (row 0), positive: row 1
print(f"\nanchor = {names[0]}; cosine similarity to the positive: {U[0] @ U[1]:.3f}")
print("share of the push on each negative:")
print("negative           cos    tau=0.05  tau=0.2  tau=1.0")
shares = []
for tau in (0.05, 0.2, 1.0):
    s = U[neg] @ U[0] / tau
    r = np.exp(s - s.max())
    shares.append(r / r.sum())                 # P_ik / sum of P over negatives
for n, k in enumerate(neg):
    print(f"{names[k]:17s} {U[0] @ U[k]:.3f}  " + "".join(f"{sh[n]:9.3f}" for sh in shares))
for tau in (0.05, 0.2, 1.0):
    print(f"tau = {tau}: P(positive) = {nt_xent(Z, tau)[2][0, 1]:.3f}, "
          f"anchor loss = {-np.log(nt_xent(Z, tau)[2][0, 1]):.3f}")

NT-Xent loss at tau = 0.2: 0.6092
max |analytic - finite-difference gradient|: 8.4e-10

anchor = square A view 1; cosine similarity to the positive: 0.927
share of the push on each negative:
negative           cos    tau=0.05  tau=0.2  tau=1.0
square B view 1   0.772      0.323    0.418    0.239
square B view 2   0.809      0.677    0.503    0.248
bar view 1        0.219      0.000    0.026    0.137
bar view 2        0.238      0.000    0.029    0.140
ring view 1       0.066      0.000    0.012    0.118
ring view 2       0.067      0.000    0.012    0.118
tau = 0.05: P(positive) = 0.879, anchor loss = 0.129
tau = 0.2: P(positive) = 0.476, anchor loss = 0.742
tau = 1.0: P(positive) = 0.218, anchor loss = 1.523

At $\tau = 0.05$ the two views of square B receive the whole push and the bar and ring none; at $\tau = 1$ every negative gets 12–25%, even the ring, which is nearly orthogonal to the anchor. This is the hardness-aware behavior Wang and Liu describe. Since the loss itself changes with $\tau$, loss values compare only at equal temperature.

Example: a tiny SimCLR in PyTorch

The second script trains SimCLR on 2,000 unlabeled $32 \times 32$ RGB images of a disk, square, triangle or cross with random size, position, brightness and tint, augmented with a random resized crop (30–100% of the area) and flip, color jitter and grayscale. A three-layer CNN gives a 64-D $h$ and an MLP head a 32-D $z$; 300 Adam steps use batches of 128 images and $\tau = 0.2$. A linear probe and a 5-nearest-neighbor classifier on $h$, fitted on only 40 labelled images, then classify 1,000 test images; the control is the same encoder untrained. It uses PyTorch 2.14 on the CPU and takes about 20 seconds.

import sys
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_num_threads(4)
H = 32
yy, xx = torch.meshgrid(torch.linspace(-1, 1, H), torch.linspace(-1, 1, H), indexing="ij")

def make_images(n, gen):
    """n RGB 32x32 images of a disk, square, triangle or cross with random size,
    position, brightness and tint. Returns images (n, 3, 32, 32) and shape labels."""
    labels = torch.randint(0, 4, (n,), generator=gen)
    imgs = torch.empty(n, 3, H, H)
    for i in range(n):
        cy, cx = (torch.rand(2, generator=gen) - 0.5) * 0.5
        r = 0.45 + 0.25 * torch.rand(1, generator=gen).item()
        y, x = yy - cy, xx - cx
        k = labels[i].item()
        if k == 0:
            m = x**2 + y**2 < r**2
        elif k == 1:
            m = (x.abs() < 0.85 * r) & (y.abs() < 0.85 * r)
        elif k == 2:
            m = (y < 0.8 * r) & (y > -r + 2 * x.abs())
        else:
            m = ((x.abs() < 0.3 * r) | (y.abs() < 0.3 * r)) & (x.abs() < r) & (y.abs() < r)
        bg = 0.4 * torch.rand(1, generator=gen) * (0.6 + 0.4 * torch.rand(3, generator=gen))
        fg = bg + (0.4 + 0.2 * torch.rand(1, generator=gen)) * (0.6 + 0.4 * torch.rand(3, generator=gen))
        imgs[i] = torch.where(m, fg[:, None, None], bg[:, None, None])
    return imgs + 0.05 * torch.randn(imgs.shape, generator=gen), labels

def augment(x, gen):
    """t ~ T: random resized crop (scale 0.3-1), horizontal flip, color jitter, grayscale."""
    n = len(x)
    s = 0.55 + 0.45 * torch.rand(n, generator=gen)          # crop side as a fraction of the image
    t = (1 - s)[:, None] * (2 * torch.rand(n, 2, generator=gen) - 1)
    flip = torch.where(torch.rand(n, generator=gen) < 0.5, -1.0, 1.0)
    theta = torch.zeros(n, 2, 3)
    theta[:, 0, 0], theta[:, 1, 1], theta[:, :, 2] = s * flip, s, t
    x = F.grid_sample(x, F.affine_grid(theta, x.shape, align_corners=False), align_corners=False)
    b = 1 + 0.8 * (torch.rand(n, 1, 1, 1, generator=gen) - 0.5)       # brightness
    c = 1 + 0.8 * (torch.rand(n, 1, 1, 1, generator=gen) - 0.5)       # contrast
    gain = 1 + 0.8 * (torch.rand(n, 3, 1, 1, generator=gen) - 0.5)    # per-channel color shift
    mean = x.mean((1, 2, 3), keepdim=True)
    x = ((x - mean) * c + mean) * b * gain
    gray = torch.rand(n, 1, 1, 1, generator=gen) < 0.2
    return torch.where(gray, x.mean(1, keepdim=True).expand_as(x), x).clamp(0, 1)

class Encoder(nn.Module):                                  # f: image -> h (64-D)
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1), nn.Flatten())
    def forward(self, x):
        return self.net(x)

def nt_xent(z, tau):
    """SimCLR loss for 2N embeddings; rows 2k and 2k+1 are the two views of image k."""
    z = F.normalize(z, dim=1)
    logits = z @ z.T / tau
    logits.fill_diagonal_(float("-inf"))
    partner = torch.arange(len(z)) ^ 1
    return F.cross_entropy(logits, partner)

@torch.no_grad()
def features(f, x):
    f.eval()
    h = F.normalize(f(x), dim=1)
    f.train()
    return h

def linear_probe(htr, ytr, hte, yte):
    W = torch.zeros(htr.shape[1], 4, requires_grad=True)
    b = torch.zeros(4, requires_grad=True)
    opt = torch.optim.Adam([W, b], lr=0.05)
    for _ in range(300):
        opt.zero_grad()
        (F.cross_entropy(htr @ W + b, ytr) + 1e-3 * (W**2).sum()).backward()
        opt.step()
    return ((hte @ W + b).argmax(1) == yte).float().mean().item()

def knn(htr, ytr, hte, yte, k=5):
    idx = (hte @ htr.T).topk(k, dim=1).indices
    pred = torch.mode(ytr[idx], dim=1).values
    return (pred == yte).float().mean().item()

def run(seed, steps=300, batch=128, tau=0.2):
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    unlabeled, _ = make_images(2000, gen)                 # labels thrown away
    f = Encoder()
    g = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 32))   # projection head
    f_random = Encoder()
    f_random.load_state_dict(f.state_dict())              # control: same initialization, untrained
    opt = torch.optim.Adam(list(f.parameters()) + list(g.parameters()), lr=1e-3)
    for step in range(steps):
        x = unlabeled[torch.randint(0, len(unlabeled), (batch,), generator=gen)]
        v = torch.stack([augment(x, gen), augment(x, gen)], 1).flatten(0, 1)  # interleave views
        loss = nt_xent(g(f(v)), tau)
        opt.zero_grad()
        loss.backward()
        opt.step()
    xtr, ytr = make_images(40, gen)                      # only 40 labelled images
    xte, yte = make_images(1000, gen)
    out = {"loss": loss.item()}
    for name, enc in (("random", f_random), ("simclr", f)):
        htr, hte = features(enc, xtr), features(enc, xte)
        out[name] = (linear_probe(htr, ytr, hte, yte), knn(htr, ytr, hte, yte))
    return out

if __name__ == "__main__":
    print("PyTorch", torch.__version__)
    seeds = [int(a) for a in sys.argv[1:]] or [3]
    for seed in seeds:
        r = run(seed)
        print(f"seed {seed}: NT-Xent after 300 steps {r['loss']:.3f} "
              f"(uniform guess: ln 255 = {math.log(255):.3f})")
        for name in ("random", "simclr"):
            print(f"  {name:6s} encoder: linear probe {r[name][0]:.3f}, 5-NN {r[name][1]:.3f}")

PyTorch 2.14.1
seed 3: NT-Xent after 300 steps 3.692 (uniform guess: ln 255 = 5.541)
  random encoder: linear probe 0.288, 5-NN 0.340
  simclr encoder: linear probe 0.956, 5-NN 0.826

Two runs of seed 3 printed identical output. Over seeds 0–9, the contrastively trained encoder reached 84–99% with the linear probe and 60–88% with 5-NN; the untrained encoder stayed at 24–30% and 24–34%, near chance (25%). Our reading: the augmentations change position, scale and color but never the shape, so encoding the shape is the cheapest way to match two views. Batch normalization here sees all 256 views at once, as with SimCLR’s global BN.

Variants

• Supervised contrastive learning (SupCon). Khosla et al. (2020) treat every image of the same class in the batch as a positive. Summing over the positives outside the logarithm worked far better than inside (78.7% against 67.4% ImageNet top-1, ResNet-50). With a linear classifier on the frozen encoder, a ResNet-200 reached 81.4% ImageNet top-1, 0.8 points above the best number previously reported for it, and training was more stable to hyperparameters.
• Image–text contrastive learning. CLIP trains an image encoder (a ResNet or a Vision Transformer) and a text encoder so that the $N$ true image–caption pairs in a batch score highest among the $N^2$ pairings, with a symmetric cross-entropy over images and over texts, a batch of 32,768 and a learned temperature initialized to the equivalent of 0.07 (Radford et al., 2021).
• Sigmoid loss. SigLIP replaces the softmax with an independent binary classification of every image–text pair, $-\log \sigma\big(z_{ij}(t\, \mathbf{a}_i^\top \mathbf{b}_j + b)\big)$ for normalized image and text embeddings $\mathbf{a}_i$, $\mathbf{b}_j$, with $z_{ij} = 1$ for matching pairs and $-1$ otherwise, a learned temperature $t$ and a learned bias $b$, so no normalization over the batch is needed (Zhai et al., 2023). The sign of $b$ follows the paper’s Algorithm 1 (its printed equation has $-b$); $b$ starts at $-10$, near the prior that almost every pair is negative. With a frozen image tower (SigLiT), it beat the softmax loss by a large margin below a batch of 16k, and both saturated around 32k.

Best Practices

• Normalize embeddings, tune $\tau$ (0.07–0.5 in the papers above) and report it.
• Train with a nonlinear projection head; use the layer before it downstream.
• Choose augmentations that remove shortcuts between instances (color histograms, position) but keep what the downstream task needs.
• Get many negatives from a large batch or a queue with a momentum encoder; with small batches, train longer.

Pitfalls

• False negatives. Instance discrimination pushes away other photographs of the same kind of object; random negatives rely on such false negatives being rare (Khosla et al., 2020), and a small $\tau$ tolerates them less (Wang & Liu, 2021).
• Batch-size dependence. SimCLR’s gains from large batches were largest at 100 epochs and shrank with longer training (Chen et al., 2020); MoCo v2 estimated that an end-to-end batch of 4,096 would need 93 GB per GPU (Chen, Fan, Girshick & He, 2020).
• Cheating through batch statistics. With per-device batch normalization, MoCo’s model found a low-loss solution that did not learn good features; shuffling samples across devices for the key encoder fixed it (He et al., 2020). SimCLR’s global BN is the alternative.

Evidence and Impact

Wu et al. (2018) reported 54.0% ImageNet top-1 with a linear SVM on conv5 features of a ResNet-50. Two years later SimCLR reached 69.3% with a linear classifier on a ResNet-50, against 76.3% for its supervised baseline (Chen et al., 2020): by our arithmetic the gap shrank from about 22 points to 7, though the protocols differ. In transfer learning, MoCo beat supervised pre-training on seven detection and segmentation tasks (He et al., 2020).