On this page
Contrastive learning trains an encoder, the network that turns an image into a vector, its embedding. Every training step, each view in the batch answers one multiple-choice question: which candidate is its positive, the other view (randomly edited copy) of the same image? (Views are introduced in Self-Supervised Learning.) The loss rewards the right pick, which pulls each positive pair together and pushes other images apart. SimCLR does this with one big batch, and MoCo with a queue of past batches.
In each batch, the two views of one image are a positive pair and every other image is a negative. The view asking the question is the anchor.
SimCLR (Chen et al. 2020) turns a batch of N images into one loss. It edits each image twice, embeds every view, and trains the encoder so each view picks its positive out of all the others. One encoder f serves every view, followed by a small projection head g.
Figure 3. A toy batch of N = 4; each step redraws the part of SimCLR it names.
InfoNCE is the loss behind the multiple-choice question. An anchor faces M candidates, its positive plus M−1 negatives. InfoNCE (named after noise-contrastive estimation, from Contrastive Predictive Coding, van den Oord et al. 2018) scores every candidate, turns the scores into probabilities and applies cross entropy, the standard loss for picking the right answer from a set, with the positive as the right answer.
| Method | Anchor | Positive | Candidates M |
|---|---|---|---|
| CPC2018 | past speech, summarized | the true next step | the true next step plus random other steps |
| MoCo2020 | a view | its other view | the positive plus 65,536 queued past views |
| SimCLR2020 | a view | its other view | the other 2N−1 views |
| CLIP2021 | a photo, then a caption | its caption, then its photo | all N in the batch |
The snippet computes L for one toy batch, where z1 and z2 hold the unit-length embeddings of the same N images’ two views, row for row.
import torch
import torch.nn.functional as F
N, D, TAU = 4, 128, 0.1
z1 = F.normalize(torch.randn(N, D), dim=1)
z2 = F.normalize(torch.randn(N, D), dim=1)
z = torch.cat([z1, z2]) # [2N, D]
s = z @ z.T / TAU # [2N, 2N]
is_self = torch.eye(2 * N, dtype=torch.bool)
s = s.masked_fill(is_self, -torch.inf) # never pick itself
t = torch.arange(2 * N).roll(N) # positive of view i
L = F.cross_entropy(s, t)Figure 5. Arrow length is each negative’s share of the push.
More candidates make the question harder, so a low loss says more about how much the two views share. Guessing at random among M candidates gives the positive probability 1/M, a loss of log M; a perfect model drives it toward 0.
A loss stuck at log M means every candidate scores the same, the signature of collapse (every image gets the same embedding).
SimCLR edits every view with a random crop, color distortion and blur. Crops of one photo share the same mix of colors, so with crops alone the network matches color histograms (how often each color appears) and ignores shapes. Adding color distortion blocks that shortcut.
The snippet builds SimCLR’s edit pipeline in torchvision, and each call of augment draws fresh random edits.
import torch
import torchvision.transforms as T
SIZE, STRENGTH = 224, 1.0
BLUR_K = 23 # odd, about 10% of SIZE
jit = 0.8 * STRENGTH
color = T.ColorJitter(jit, jit, jit, 0.2 * STRENGTH)
blur = T.GaussianBlur(BLUR_K)
augment = T.Compose([
T.RandomResizedCrop(SIZE),
T.RandomHorizontalFlip(),
T.RandomApply([color], p=0.8),
T.RandomGrayscale(p=0.2),
T.RandomApply([blur], p=0.5),
T.ToTensor(),
])
img = T.ToPILImage()(torch.rand(3, 256, 256)) # toy photo
view1, view2 = augment(img), augment(img)STRENGTH up to its default 1.0 lifts SimCLR from 59.6% to 63.2% while supervised falls from 77.0% to 75.7% (SimCLR Table 1).Weak edits teach little, because matching pixels is enough. Strong edits force the model to match content, until the crops share no scene and the two views no longer show the same thing. Edits that keep the label fight overfitting the same way they do in Data Augmentation for Wearables.
Figure 8. Pixels agree falls first; scene shared falls only for strong crops.
The encoder f outputs h, the embedding you keep. A small projection head g, an MLP (two linear layers with a ReLU, which zeroes negative values, between them), maps h to z, and the loss reads only z. SimCLR throws g away after training. The head learns to drop details the loss ignores, such as color and orientation, so h still holds them.
MoCo (Momentum Contrast, He et al. 2020) gets many negatives from a small batch by keeping recent batches' embeddings in a queue. SimCLR's negatives come from the batch, so it trains at a batch of 4096 (Figure 6 shows the lower ceiling at batch 256). In MoCo the anchor, the view that gets gradients, is called the query; its other view is the positive key, and 65,536 queued keys are the negatives. The loss is the same InfoNCE with M = 65,537 candidates, at τ = 0.07.
This snippet runs one SimCLR training step, where both views pass through the encoder f and the head g of Section 5 and the loss reads only z.
import torch
import torch.nn as nn
import torch.nn.functional as F
N, D, TAU = 256, 128, 0.1
f = nn.Sequential(nn.Flatten(), nn.Linear(3 * 32 * 32, 512), nn.ReLU())
g = nn.Sequential(nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, D))
params = list(f.parameters()) + list(g.parameters())
opt = torch.optim.SGD(params, lr=0.3, momentum=0.9, weight_decay=1e-6)
x1, x2 = torch.rand(N, 3, 32, 32), torch.rand(N, 3, 32, 32)
z = F.normalize(g(f(torch.cat([x1, x2]))), dim=1) # [2N, D]
s = z @ z.T / TAU
s = s.masked_fill(torch.eye(2 * N, dtype=torch.bool), -torch.inf)
t = torch.arange(2 * N).roll(N)
loss = F.cross_entropy(s, t)
opt.zero_grad()
loss.backward()
opt.step()This snippet runs one MoCo training step, then moves the key encoder toward the query encoder by EMA (exponential moving average, Section 6 eq line 1) and puts the batch’s keys in the queue.
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
D, K, M_EMA, TAU = 128, 65536, 0.999, 0.07
enc_q = nn.Sequential(nn.Flatten(), nn.Linear(3 * 32 * 32, D))
enc_k = copy.deepcopy(enc_q).requires_grad_(False)
queue = F.normalize(torch.randn(K, D), dim=1)
opt = torch.optim.SGD(enc_q.parameters(), lr=0.03, momentum=0.9, weight_decay=1e-4)
x1, x2 = torch.rand(256, 3, 32, 32), torch.rand(256, 3, 32, 32)
q = F.normalize(enc_q(x1), dim=1) # queries, [256, D]
with torch.no_grad():
k = F.normalize(enc_k(x2), dim=1) # keys, no gradient
pos = (q * k).sum(dim=1, keepdim=True) # [256, 1]
neg = q @ queue.T # [256, K]
logits = torch.cat([pos, neg], dim=1) / TAU
labels = torch.zeros(len(q), dtype=torch.long) # positive is column 0
loss = F.cross_entropy(logits, labels)
opt.zero_grad()
loss.backward()
opt.step()
with torch.no_grad():
for p_k, p_q in zip(enc_k.parameters(), enc_q.parameters()):
p_k.mul_(M_EMA).add_(p_q, alpha=1 - M_EMA)
queue = torch.cat([k, queue[:-len(k)]]) # newest in, oldest outx1 and x2, as augment in Section 4 makes them, give query q and key k, the two encoders in Figure 10.| SimCLR | MoCo | |
|---|---|---|
| Acts | every step, on the batch | every step; 256 keys in, the oldest 256 out |
| Keeps | only the weights | the queue and the key encoder |
| Settings | LARS (a large-batch optimizer), 100 to 1000 epochs | SGD, lr 0.03, 200 epochs |
| Cost | about 1.5 hours per 100 epochs on 128 TPU v3 cores | about 53 hours on 8 GPUs; the queue holds 65,536 × 128 numbers, about 34 MB |
| Code | google-research/simclr | facebookresearch/moco |
| Gotchas | small batches mean few negatives; BatchNorm (which normalizes each feature with its batch’s mean and variance) must use statistics from all devices (Global BN) | m = 0 does not converge; a query and its key on one GPU share BatchNorm statistics and can match by them, so MoCo shuffles samples across GPUs for the keys |
| Knob, default | Raise it | Lower it |
|---|---|---|
| τ divides every score si,k; 0.1 SimCLR, 0.07 MoCo | shares flatten; very high, every negative is pushed about equally | the push gathers on the closest negatives; near 0, only the single closest |
| N, the batch, gives each view 2N − 2 negatives in C; 4096 SimCLR | more negatives, a higher bound log M, better accuracy, more memory | fewer negatives and lower accuracy; at N = 1 there are none |
| K, the queue, adds K negatives per query; 65,536 MoCo | accuracy rises (MoCo Fig. 3); costs memory, 128 numbers per key | back toward a batch of negatives; at 0 there are none |
| m blends the key encoder each step; 0.999 MoCo | steadier keys; at 1 the key encoder never changes | queued keys disagree; m = 0.9 gives 55.2%, against 59.0% at 0.999, and 0 fails (MoCo, K = 4096) |
| D, the size of z; 128 SimCLR | results stay similar at any size (SimCLR Fig. 8); probe h either way | |
From SimCLR sec. 2.2 and Table 5, and MoCo Algorithm 1 and App. A.9.
Image-Text Contrastive Learning (CLIP) comes next, where the same loss runs in both directions between photos and their captions.