Lesson 50: Self-Supervised and Contrastive Learning¶
Lesson 38'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()
A small Vision Transformer as the encoder¶
The encoder that turns each augmented view into an embedding is a ViT (Lesson 47's exact pattern: patchify, linearly embed, a [CLS] token gathers information via self-attention, project its final state to the embedding dimension) rather than a CNN. Real SimCLR's own paper (Chen et al., 2020) used a CNN (ResNet) backbone instead — also a perfectly legitimate, arguably more sample-efficient choice at this scale — but this part of the course is specifically about Transformer-based architectures, so a ViT is used here instead.
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.
def patchify(img, patch_size):
B, C, H, W = img.shape
P = patch_size
patches = img.unfold(2, P, P).unfold(3, P, P) # (B, C, H/P, W/P, P, P)
patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)
return patches.reshape(B, -1, C * P * P) # (B, num_patches, C*P*P)
def positional_encoding(T, D):
pos = torch.arange(T).unsqueeze(1).float()
i = torch.arange(D).unsqueeze(0).float()
angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))
angles = pos * angle_rates
pe = torch.zeros(T, D)
pe[:, 0::2] = torch.sin(angles[:, 0::2])
pe[:, 1::2] = torch.cos(angles[:, 1::2])
return pe
class Encoder(nn.Module):
def __init__(self, out_dim=16, img_size=16, patch_size=4, embed_dim=32, n_heads=4, n_layers=2):
super().__init__()
self.patch_size = patch_size
n_patches = (img_size // patch_size) ** 2
patch_dim = patch_size * patch_size
self.embed = nn.Linear(patch_dim, embed_dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.proj = nn.Linear(embed_dim, out_dim)
def forward(self, x):
p = patchify(x, self.patch_size)
tok = self.embed(p)
cls = self.cls_token.expand(x.shape[0], -1, -1)
tok = torch.cat([cls, tok], dim=1) + self.pos_enc
out = self.encoder(tok)
return self.proj(out[:, 0]) # project the [CLS] token's final state to out_dim
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 38's frozen-backbone pattern) and train only a small linear classifier on top of its features — a technique called linear probing — 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 shows a modest but consistently positive improvement over random features across repeated runs — not the dramatic gap a CNN encoder shows on this same setup, but a real one. That's an honest consequence of the architecture choice, not a bug: Lesson 47 already showed that a ViT has no built-in translation invariance and needs far more data than a CNN to learn useful structure from scratch, and the same limitation shows up here — the contrastive objective still pushes same-shape images together and different-shape images apart, but a few hundred unlabeled toy images and a couple hundred training steps is a lot less signal for a Transformer to work with than it would be for a CNN's stronger inductive bias. (This is also why SimCLR's own original paper used a CNN/ResNet backbone, not a ViT — a legitimate, arguably more sample-efficient choice at this scale; the ViT is used here instead because this part of the course is specifically about Transformer-based self-supervised pretraining.)
This is the mechanism behind why self-supervised pretraining matters at scale: models like DINOv2 (Lesson 51) 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. At real scale, the data-hunger gap this lesson's toy experiment exposes stops being a limitation at all.
Exercises¶
- 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?