Contrastive learning

Deep dive

How it works

Parts:

• sampler — builds (a, b): two augmentations (Data augmentation), image and caption, adjacent sentences, query and click, a sensor window and a later one (CPC).
• encoder f — shared (SimCLR) or one per modality (CLIP); MoCo adds a momentum copy.
• projection head g — small MLP; the loss sees g(f(x)), downstream uses f(x).
• negatives — the rest of the batch, or a queue of past keys (MoCo).

1. Sample B pairs, make views a_i, b_i.
2. Encode z_i = g(f(a_i))/‖·‖, z'_i = g(f(b_i))/‖·‖: two forward passes per example.
3. Score S_ij = z_i·z'_j / τ, a B×B matrix; the diagonal holds the positives.
4. Loss ℓ_i = -S_ii + log Σ_j exp(S_ij); add the column term for two modalities (CLIP).
5. Backprop: ∂ℓ/∂S = softmax - 1[i=j], so hard negatives push hardest.
6. Scale. All-gather embeddings so each device sees all B. MoCo: keys from a momentum encoder (θ_k ← m·θ_k + (1-m)·θ_q, m = 0.999) fill a FIFO queue of K = 65,536; only queries get gradients.
7. Discard g, keep f.

 a_i ─ f ─ g ─ z_i  ─┐                        ┌─ row softmax (a→b)
                     ├─ S = Z Z'ᵀ / τ  (B×B) ─┤
 b_i ─ f ─ g ─ z'_i ─┘   diagonal = positives └─ column softmax (b→a), symmetric

Temperature. Small τ sharpens the softmax so a few close negatives dominate; large τ weighs all alike. CLIP learns it (clipped, scale ≤ 100).

Collapse. Rewarding only closeness makes a constant embedding optimal; negatives block it. Without negatives: BYOL regresses a slow moving-average target from an online predictor (74.3 % linear top-1, ResNet-50); SimSiam uses predictor and stop-gradient only; VICReg penalises low variance and correlated dimensions.

Alignment and uniformity (Wang & Isola). Asymptotically InfoNCE makes positives close (alignment) and features spread over the hypersphere (uniformity); optimising both directly matched it downstream.

Evaluation. Linear probe: freeze f, train a linear classifier (SimCLR 76.5 % ImageNet top-1). Retrieval: recall@k of the partner in a held-out gallery (HNSW at scale). Zero-shot (CLIP): embed class names as text prompts, pick the nearest per image; it matched a supervised ResNet-50 on ImageNet without its labels.

Code

Linear encoder, symmetric InfoNCE with analytic gradients. Views: one of 64 latent prototypes plus Gaussian noise (σ = 0.7); 32 distinct prototypes per batch. Asserts: loss falls; held-out top-1 retrieval among 64 exceeds 0.9; smaller τ sharpens the softmax. The last loop times B×B similarity + loss, d = 128. exp(s - 1/τ) needs no max pass: cosines are at most 1 (τ ≳ 0.025 in f32).

use std::time::Instant;
struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        self.0.wrapping_mul(0x2545F4914F6CDD1D)
    }
    fn gauss(&mut self) -> f32 {
        let (u, v) = (((self.next() >> 40) as f32 + 0.5) / 16777216.0, (self.next() >> 40) as f32 / 16777216.0);
        (-2.0 * u.ln()).sqrt() * (6.2831853 * v).cos()
    }
}
fn dot(a: &[f32], b: &[f32]) -> f32 { 
    let mut acc = [0f32; 8];
    for (x, y) in a.chunks_exact(8).zip(b.chunks_exact(8)) { for l in 0..8 { acc[l] += x[l] * y[l]; } }
    acc.iter().sum()
}
/// Symmetric InfoNCE on L2-normalised rows (b x d): S = z1 z2^T / tau, positives on the diagonal,
/// loss = 1/2b sum_i [CE(row i) + CE(column i)]. Returns (loss, dL/dz1, dL/dz2).
fn info_nce(z1: &[f32], z2: &[f32], b: usize, d: usize, tau: f32, grad: bool) -> (f32, Vec<f32>, Vec<f32>) {
    let m = 1.0 / tau; // |cos| <= 1, so exp(s - m) <= 1: no max pass
    let mut e = vec![0f32; b * b];
    let (mut rs, mut cs) = (vec![0f32; b], vec![0f32; b]);
    for i in 0..b { for j in 0..b {
        let x = (dot(&z1[i * d..][..d], &z2[j * d..][..d]) / tau - m).exp();
        e[i * b + j] = x; rs[i] += x; cs[j] += x;
    } }
    let loss = (0..b).map(|i| rs[i].ln() + cs[i].ln() - 2.0 * e[i * b + i].ln()).sum::<f32>() / (2 * b) as f32;
    let (mut g1, mut g2) = (vec![], vec![]);
    if grad {
        (g1, g2) = (vec![0f32; b * d], vec![0f32; b * d]);
        for i in 0..b { for j in 0..b {
            let g = (e[i * b + j] / rs[i] + e[i * b + j] / cs[j] - if i == j { 2.0 } else { 0.0 }) / (2 * b) as f32 / tau;
            for k in 0..d { g1[i * d + k] += g * z2[j * d + k]; g2[j * d + k] += g * z1[i * d + k]; }
        } }
    }
    (loss, g1, g2)
}
const DI: usize = 32; // input dim
const DO: usize = 32; // embedding dim
fn encode(w: &[f32], x: &[f32], b: usize) -> (Vec<f32>, Vec<f32>) {
    let (mut z, mut n) = (vec![0f32; b * DO], vec![0f32; b]);
    for i in 0..b {
        for o in 0..DO { z[i * DO + o] = dot(&w[o * DI..][..DI], &x[i * DI..][..DI]); }
        n[i] = dot(&z[i * DO..][..DO], &z[i * DO..][..DO]).sqrt();
        for o in 0..DO { z[i * DO + o] /= n[i]; }
    }
    (z, n)
}
fn step(w: &[f32], x1: &[f32], x2: &[f32], b: usize, tau: f32) -> (f32, Vec<f32>) {
    let ((z1, n1), (z2, n2)) = (encode(w, x1, b), encode(w, x2, b));
    let (loss, g1, g2) = info_nce(&z1, &z2, b, DO, tau, true);
    let mut dw = vec![0f32; DO * DI];
    for (z, n, g, x) in [(&z1, &n1, &g1, x1), (&z2, &n2, &g2, x2)] {
        for i in 0..b {
            let zg = dot(&z[i * DO..][..DO], &g[i * DO..][..DO]);
            for o in 0..DO { // through the normalisation
                let du = (g[i * DO + o] - z[i * DO + o] * zg) / n[i];
                for k in 0..DI { dw[o * DI + k] += du * x[i * DI + k]; }
            }
        }
    }
    (loss, dw)
}
fn views(rng: &mut Rng, protos: &[f32], ids: &[usize], sigma: f32) -> (Vec<f32>, Vec<f32>) {
    let mut v = |_| ids.iter().flat_map(|&p| (0..DI).map(move |k| (p, k))).map(|(p, k)| protos[p * DI + k] + sigma * rng.gauss()).collect();
    (v(0), v(1))
}
fn main() {
    let mut rng = Rng(0x9E3779B97F4A7C15);
    let (k, b, sigma, tau) = (64usize, 32usize, 0.7f32, 0.1f32);
    let protos: Vec<f32> = (0..k * DI).map(|_| rng.gauss()).collect();
    let mut w: Vec<f32> = (0..DO * DI).map(|_| rng.gauss() / (DI as f32).sqrt()).collect();
    // held-out top-1: query view of prototype p must pick gallery view of p
    let all: Vec<usize> = (0..k).collect();
    let top1 = |w: &[f32], rng: &mut Rng| {
        let (q, g) = views(rng, &protos, &all, sigma);
        let (zq, zg) = (encode(w, &q, k).0, encode(w, &g, k).0);
        let sim = |i: usize, j: usize| dot(&zq[i * DO..][..DO], &zg[j * DO..][..DO]);
        (0..k).filter(|&i| (0..k).all(|j| j == i || sim(i, i) > sim(i, j))).count() as f32 / k as f32
    };
    let before = top1(&w, &mut rng);
    // b distinct prototypes per batch: no false negatives
    let (mut pool, mut first, mut last) = ((0..k).collect::<Vec<usize>>(), 0.0, 0.0);
    for it in 0..600 {
        for i in 0..b { let j = i + (rng.next() % (k - i) as u64) as usize; pool.swap(i, j); }
        let (x1, x2) = views(&mut rng, &protos, &pool[..b], sigma);
        let (loss, dw) = step(&w, &x1, &x2, b, tau);
        for (p, g) in w.iter_mut().zip(&dw) { *p -= g; } // SGD, lr 1
        if it == 0 { first = loss }
        last = loss;
    }
    let after = (0..5).map(|_| top1(&w, &mut rng)).sum::<f32>() / 5.0;
    println!("loss {first:.3} -> {last:.3}; top-1 retrieval {before:.3} -> {after:.3}");
    assert!(last < first * 0.8 && after > 0.9);
    let maxp = |t: f32| { let e = [1.0f32, 0.8, 0.2, -0.3].map(|s| (s / t).exp()); e[0] / e.iter().sum::<f32>() };
    assert!(maxp(0.05) > maxp(0.5) && maxp(0.5) > maxp(1.0));
    for b in [256usize, 1024, 4096] {
        let d = 128; // unit rows e_(i mod d); the cost does not depend on the values
        let z: Vec<f32> = (0..b * d).map(|i| if i % d == i / d % d { 1.0 } else { 0.0 }).collect();
        let t = Instant::now();
        info_nce(&z, &z, b, d, 0.1, false);
        println!("B={b}: similarity + loss {:.1} ms", t.elapsed().as_secs_f64() * 1e3);
    }
}

Measured (Apple M4 Pro, one thread, release build): loss 0.545 → 0.209; top-1 retrieval 0.688 → 0.966 (mean of 5 held-out draws); similarity + loss 0.8 ms (B = 256), 12 ms (1024), 202-209 ms (4096) over five runs. Gradient checked by central differences in a scratch run.

Complexity and constants

Loss: 2·B²·d multiply-adds plus B² exps; 4·B² bytes in f32 (262 KB, 4.2 MB, 67 MB at B = 256, 1024, 4096). 4.3 GFLOP at B = 4096 in 0.2 s: order of 20 GFLOP/s with the exps.

Efficiency

• Cost: M4 Pro, one thread, B = 4096, d = 128: about 0.2 s and 67 MB for similarity + loss (measured above), about 16× less per 4× smaller B. On accelerators the encoder dominates (SimCLR: ~1.5 h for 100 epochs, ResNet-50, batch 4096, 128 TPU v3 cores). Bound: encoder FLOPs, then S memory.
• Levers:
  • Local loss on gathered embeddings — each device builds B_local × B_global, not B² (open_clip local_loss) — multi-device.
  • Key queue + momentum encoder — 65,536 negatives from a small batch, no backward pass on keys (MoCo) — few GPUs.
  • Sigmoid pairwise loss — no global softmax normalisation, batch scales further (SigLIP) — very large batches.
  • Mixed precision and gradient checkpointing — fit large batches (CLIP) — large encoders.
• Naive → standard → state of the art:
  • One random negative per positive · O(B), weak signal.
  • In-batch InfoNCE, full S · O(B²d), B of a few thousand.
  • Sharded or sigmoid loss, or key queue · B of 32k (CLIP), 65k keys (MoCo), no B² matrix per device.

Implementation notes

• Batch norm leaks: positives on one device share statistics; SimCLR aggregates across devices.
• Mask duplicate ids: a repeated item is a false negative (Negative sampling).
• The listing uses cross-view negatives only; SimCLR also uses the other 2B-2 views.

Pitfalls

• Small batch, no queue: few negatives.
• Views too easy: the encoder matches a shortcut (colour, a timestamp); too hard: no shared content. SimCLR finds the composition of augmentations critical.
• Evaluating the projection output instead of f; one τ for every dataset and batch size.

Variants and alternatives

In production

• Vision pretraining, zero-shot classification, image-text search: SimCLR, MoCo, CLIP.
• Speech: wav2vec 2.0 (contrastive task over quantised latents). Medical imaging with few labels: ConVIRT (X-rays with reports).
• Recommendation and text search: in-batch InfoNCE (Two-tower retrieval, sentence-transformers).
• Sensors and wearables: time-shifted windows as CPC-style positives (no paper cited).

Implementations

Sources

• Chen, Kornblith, Norouzi, Hinton, A Simple Framework for Contrastive Learning of Visual Representations, 2020. paper · code (archived)
• He, Fan, Wu, Xie, Girshick, Momentum Contrast for Unsupervised Visual Representation Learning, 2019. paper · code
• Radford et al., Learning Transferable Visual Models From Natural Language Supervision, 2021. paper
• van den Oord, Li, Vinyals, Representation Learning with Contrastive Predictive Coding, 2018. paper
• Grill et al., Bootstrap your own latent: A new approach to self-supervised Learning, 2020. paper
• Chen, He, Exploring Simple Siamese Representation Learning, 2020. paper
• Bardes, Ponce, LeCun, VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning, 2021. paper
• Wang, Isola, Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere, 2020. paper
• Gao, Yao, Chen, SimCSE: Simple Contrastive Learning of Sentence Embeddings, 2021. paper
• Zhai, Mustafa, Kolesnikov, Beyer, Sigmoid Loss for Language Image Pre-Training, 2023. paper
• Baevski, Zhou, Mohamed, Auli, wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations, 2020. paper
• Zhang et al., ConVIRT: Contrastive Learning of Medical Visual Representations from Paired Images and Text, 2020. paper