Contrastive Learning: InfoNCE, SimCLR, MoCo

The contrastive idea

Take an image, apply two random augmentations (a crop, a color jitter, a blur), and you get two views. They look different pixel-by-pixel but depict the same scene, so their representations should be close — that is a positive pair. Any other image in the batch is a negative: its representation should be far away. Contrastive learning simply optimizes both desiderata at once. The deep insight is that negatives are what prevent collapse: if you only pulled positives together, the model would happily map everything to one point.

InfoNCE: the objective and its meaning

The workhorse loss is InfoNCE (also called NT-Xent). For an anchor view, you compute its similarity to the one positive and to K negatives, divide by a temperature τ, and apply a softmax: the loss is the negative log-probability assigned to the positive. Minimizing it is equivalent to classifying which of the K+1 candidates is the true partner. Crucially, InfoNCE is a lower bound on the mutual information between the two views — so contrastive learning is, formally, maximizing how much one view tells you about the other.

\mathcal{L}_{\text{InfoNCE}} = -\log \frac{\exp\!\big(\mathrm{sim}(z_i, z_i^{+})/\tau\big)}{\displaystyle\sum_{k=0}^{K}\exp\!\big(\mathrm{sim}(z_i, z_k)/\tau\big)}

The InfoNCE (NT-Xent) loss: a temperature-scaled softmax that pulls the one positive together and pushes the K negatives apart.

z1, z2 = proj(f(view1)), proj(f(view2))   # L2-normalized
logits = (z1 @ z2.T) / tau                  # cosine sims, scaled
labels = arange(batch)                      # positive is the diagonal
loss = cross_entropy(logits, labels)        # InfoNCE / NT-Xent

InfoNCE in four lines: the positive is the matching row/column; every off-diagonal entry is a negative.

The temperature τ matters more than newcomers expect. A small τ sharpens the softmax and concentrates the gradient on the hardest negatives — the ones closest to the anchor — which speeds learning but can over-penalize semantically similar images. This same InfoNCE machinery first appeared for sequences as contrastive predictive coding (CPC), which predicts future latent states rather than augmented views.

SimCLR: simple, hungry for batch size

SimCLR stripped contrastive learning to its essentials and showed it could rival supervised pretraining. Its three lessons were: (1) the composition of augmentations matters enormously, especially random crop plus color distortion; (2) a nonlinear projection head between the backbone and the loss substantially improves the features you keep (you probe the backbone, not the projection); and (3) negatives come for free from other images in the batch, so you need a large batch — thousands — to have enough hard negatives.

MoCo: decoupling negatives from batch size

Needing a 4096-sized batch is a hardware tax most labs cannot pay. Momentum Contrast (MoCo) solves this with two ideas. First, a queue: keep a rolling buffer of thousands of recently-seen embeddings to serve as negatives, decoupling the negative count from the current batch. Second, a momentum encoder: the keys are encoded by a slowly-moving exponential average of the query encoder, so the queued embeddings stay consistent even as the model updates. Without the momentum encoder, old queued keys would be stale and training destabilizes.

k = momentum_encoder(view2)             # no gradient through keys
for p_q, p_k in zip(q_enc.params, k_enc.params):
    p_k.data = m * p_k.data + (1 - m) * p_q.data   # m ~ 0.999
logits = cat([q @ k.T, q @ queue.T]) / tau
queue = enqueue_dequeue(queue, k)       # FIFO buffer of negatives

MoCo: a FIFO queue supplies negatives; a momentum-averaged encoder keeps them consistent.

What a good contrastive embedding looks like

A clarifying theory decomposes the contrastive objective into two measurable properties on the unit hypersphere: alignment and uniformity. Alignment says positive pairs should land close together; uniformity says the whole set of embeddings should spread out evenly, preserving maximal information. Collapse is exactly the failure of uniformity. This pair of metrics gives you a diagnostic independent of any downstream label: you can watch alignment and uniformity during training to see whether your model is learning or quietly collapsing.

\begin{aligned}\mathcal{L}_{\text{align}} &= \mathbb{E}_{(x,x^{+})}\big\lVert f(x)-f(x^{+})\big\rVert_2^2,\\[2pt] \mathcal{L}_{\text{uniform}} &= \log\,\mathbb{E}_{x,y}\;e^{-t\lVert f(x)-f(y)\rVert_2^2}.\end{aligned}

A good embedding balances alignment (positive pairs close) and uniformity (features spread evenly on the unit hypersphere).