Lesson 47: Self-Supervised and Contrastive Learning¶
Lesson 36's transfer learning still needed a labeled source task to pretrain a useful backbone. Self-supervised learning removes even that requirement: pretrain on unlabeled images by inventing a task from the data itself, using no human annotations at all. This lesson builds the dominant recipe for that, contrastive learning (SimCLR-style, Chen et al., 2020): two randomly augmented views of the same image should produce similar embeddings; views of different images should not.
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Two views, one identity¶
Take 500 unlabeled shape images (the labels exist in this synthetic dataset only so accuracy can be measured later — the pretraining step below never looks at them). For each image sampled during pretraining, create two independently augmented views (Lesson 8/34's flips and rotations, plus pixel noise).
def make_image(shape_type, cx, cy, size=16):
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
img[cy-1:cy+2, cx-3:cx+4] = 1.0
img[cy-3:cy+4, cx-1:cx+2] = 1.0
else:
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
return img
def make_dataset(rng_local, n, position_range=(4, 12)):
imgs, labels = [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
imgs.append(make_image(shape_type, cx, cy))
labels.append(0 if shape_type == 'plus' else 1)
return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)
def augment(img, rng_local):
if rng_local.random() < 0.5:
img = np.fliplr(img).copy()
angle = rng_local.uniform(-20, 20)
M = cv2.getRotationMatrix2D((8, 8), angle, 1.0)
img = cv2.warpAffine(img, M, (16, 16))
return np.clip(img + rng_local.normal(0, 0.1, img.shape), 0, 1).astype(np.float32)
rng = np.random.default_rng(5)
X_unlabeled, y_unlabeled = make_dataset(rng, 500)
aug_rng = np.random.default_rng(1)
fig, axes = plt.subplots(2, 4, figsize=(9, 4.5))
for i in range(4):
axes[0, i].imshow(augment(X_unlabeled[i], aug_rng), cmap='gray'); axes[0, i].axis('off')
axes[1, i].imshow(augment(X_unlabeled[i], aug_rng), cmap='gray'); axes[1, i].axis('off')
axes[0, 0].set_title('view 1', fontsize=9, loc='left')
axes[1, 0].set_title('view 2', fontsize=9, loc='left')
plt.suptitle('Two augmented views of the same 4 images')
plt.show()
The contrastive loss (NT-XEnt)¶
For a batch of N images, produce 2N embeddings (two views each). For each embedding, its one true positive is the other view of the same image; every other one of the 2N-2 embeddings is a negative. Treat this as a classification problem — "which of the other 2N-1 embeddings is my positive pair?" — and minimize cross-entropy over cosine similarities. This is the normalized temperature-scaled cross-entropy (NT-Xent) loss.
class Encoder(nn.Module):
def __init__(self, out_dim=16):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1),
)
self.proj = nn.Linear(32, out_dim)
def forward(self, x):
return self.proj(self.conv(x).flatten(1))
def nt_xent_loss(z1, z2, temperature=0.5):
z1 = F.normalize(z1, dim=1)
z2 = F.normalize(z2, dim=1)
z = torch.cat([z1, z2], dim=0) # (2N, D)
sim = z @ z.T / temperature # (2N, 2N) cosine similarities
N = z1.shape[0]
mask = torch.eye(2 * N, dtype=torch.bool)
sim.masked_fill_(mask, float('-inf')) # exclude comparing an embedding to itself
targets = torch.cat([torch.arange(N, 2 * N), torch.arange(0, N)]) # each row's true positive index
return F.cross_entropy(sim, targets)
def train_contrastive(seed, epochs=200, lr=0.01, batch_size=64):
torch.manual_seed(seed)
encoder = Encoder()
opt = torch.optim.Adam(encoder.parameters(), lr=lr)
local_aug_rng = np.random.default_rng(seed + 100)
n = len(X_unlabeled)
for epoch in range(epochs):
idx = np.random.default_rng(epoch).permutation(n)[:batch_size]
batch = X_unlabeled[idx]
view1 = np.stack([augment(im, local_aug_rng) for im in batch])
view2 = np.stack([augment(im, local_aug_rng) for im in batch])
z1 = encoder(torch.tensor(view1).unsqueeze(1))
z2 = encoder(torch.tensor(view2).unsqueeze(1))
loss = nt_xent_loss(z1, z2)
opt.zero_grad()
loss.backward()
opt.step()
return encoder, loss.item()
encoder, final_loss = train_contrastive(seed=0)
print(f'final NT-Xent loss after 200 steps: {final_loss:.3f} (random-chance loss would be ln(2*64-1) = {np.log(2*64-1):.3f})')
Evaluating with a linear probe¶
The standard way to check whether self-supervised pretraining actually learned something useful: freeze the pretrained encoder (Lesson 36's frozen-backbone pattern) and train only a small linear classifier on top of its features, using a tiny number of labeled examples — exactly the low-label regime self-supervised pretraining is meant to help with. Compare against the same linear probe on a randomly initialized (never trained) encoder.
def make_noisy_dataset(rng_local, n, noise=0.35, position_range=(4, 12)):
imgs, labels = [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
img = make_image(shape_type, cx, cy)
img = np.clip(img + rng_local.normal(0, noise, img.shape), 0, 1).astype(np.float32)
imgs.append(img)
labels.append(0 if shape_type == 'plus' else 1)
return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)
X_probe_train, y_probe_train = make_noisy_dataset(rng, 8) # only 8 labeled examples
X_probe_test, y_probe_test = make_noisy_dataset(rng, 150)
def linear_probe_acc(enc, seed):
torch.manual_seed(seed)
with torch.no_grad():
feat_train = enc(torch.tensor(X_probe_train).unsqueeze(1))
feat_test = enc(torch.tensor(X_probe_test).unsqueeze(1))
probe = nn.Linear(feat_train.shape[1], 2)
opt = torch.optim.Adam(probe.parameters(), lr=0.05)
ytr = torch.tensor(y_probe_train)
for _ in range(300):
opt.zero_grad()
loss = F.cross_entropy(probe(feat_train), ytr)
loss.backward()
opt.step()
with torch.no_grad():
preds = probe(feat_test).argmax(1).numpy()
return (preds == y_probe_test).mean()
random_accs, contrastive_accs = [], []
for seed in range(5):
enc_contrastive, _ = train_contrastive(seed=seed)
contrastive_accs.append(linear_probe_acc(enc_contrastive, seed=seed + 50))
torch.manual_seed(seed)
enc_random = Encoder()
random_accs.append(linear_probe_acc(enc_random, seed=seed + 50))
print(f'linear probe on RANDOM (untrained) features: {np.mean(random_accs):.1%} (+/- {np.std(random_accs):.1%})')
print(f'linear probe on CONTRASTIVE-pretrained features: {np.mean(contrastive_accs):.1%} (+/- {np.std(contrastive_accs):.1%})')
With only 8 labeled examples to train the probe, the contrastively-pretrained encoder gives a real (if noisy, given how few labels are involved) improvement over random features — the contrastive objective, despite never seeing a label, has pushed same-shape images toward nearby points in embedding space and different-shape images apart, purely by requiring that augmented views of the same image agree.
This is the mechanism behind why self-supervised pretraining matters at scale: models like DINOv2 (Lesson 48) are pretrained this way on hundreds of millions of unlabeled images — impossible to hand-label at that scale — and the resulting features transfer to downstream tasks with only a handful of labeled examples, exactly as demonstrated here in miniature.
Exercise¶
- Increase
out_diminEncoderfrom 16 to 64. Does a higher-dimensional embedding space improve the linear probe's accuracy, hurt it, or make little difference at this data scale? - Try
temperature=0.1andtemperature=2.0innt_xent_lossinstead of0.5. Temperature controls how sharply the loss penalizes near-miss negatives — does either extreme change the final probe accuracy noticeably? - Weaken the augmentations in
augment(e.g. rotation range(-5, 5)instead of(-20, 20), no noise). Does contrastive pretraining still beat the random baseline — and what does the answer suggest about why augmentation strength is one of the most-tuned hyperparameters in real contrastive learning pipelines?