Self-Supervised Learning

Week 14 of 18 · Generative and self-supervised · 7 days

Full curriculum
Week 14 · Generative and self-supervised

Self-Supervised Learning

Week 14 · Day 1 of 7

Learning From Data Nobody Labelled

The encoder, the linear probe, and the two baselines that decide everything

By 2366 words

Weeks 12 and 13 generated data without labels. The images were the product. This week keeps the label-free training and throws the generated output away, because what is actually wanted is the encoder: a network that turns an input into a vector, trained on data nobody annotated, and then reused for a task where labels are expensive.

This is where the large models get their general ability. Nobody labelled the text a language model is pretrained on. The label was manufactured out of the data itself, by hiding part of it and asking the model to fill it in. The same trick applies to images, and this week builds three versions of it and measures all of them.

Self-supervised learning: Training on a task whose labels are generated automatically from the unlabelled data, chosen so that solving it requires understanding something you actually care about. The task itself is discarded afterwards. Only the encoder is kept.

The setup

Twelve thousand MNIST images whose labels we will pretend not to have, and a held-out test set. Every experiment this week trains on the same pool and is scored the same way.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
print('unlabelled pool %s' % (tuple(unlab_x.shape),))
print('test set %s' % (tuple(test_x.shape),))
print('batches per epoch %d' % len(pool))
unlabelled pool (12000, 1, 28, 28)
test set (2000, 1, 28, 28)
batches per epoch 46

The encoder

One small convolutional network, used unchanged by every method below, so that differences in the results are differences between the methods rather than between architectures.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
torch.manual_seed(0)
enc = Encoder()
print(enc)
print()
print('%d parameters' % sum(p.numel() for p in enc.parameters()))
print('output shape %s' % (tuple(enc(unlab_x[:4]).shape),))
Encoder(
(net): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): GroupNorm(8, 32, eps=1e-05, affine=True)
(2): ReLU()
(3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(5): GroupNorm(8, 64, eps=1e-05, affine=True)
(6): ReLU()
(7): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(8): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(9): GroupNorm(8, 64, eps=1e-05, affine=True)
(10): ReLU()
(11): Flatten(start_dim=1, end_dim=-1)
(12): Linear(in_features=3136, out_features=64, bias=True)
)
)

256832 parameters
output shape (4, 64)

One choice in there is deliberate and worth explaining: group normalisation rather than batch normalisation. This encoder gets trained under one regime and evaluated under another, often on batches of a completely different size, and week 4 spent a day on what BatchNorm does in that situation. Since the choice is easy to argue about and cheap to settle, here it is settled:

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
def build(norm):
"""The same encoder twice, differing only in the normalisation."""
def n(c):
return nn.BatchNorm2d(c) if norm == 'batch' else nn.GroupNorm(8, c)
return nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), n(32), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), n(64), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), n(64), nn.ReLU(),
nn.Flatten(), nn.Linear(64 * 7 * 7, 64), nn.ReLU(),
nn.Linear(64, 10))

def compare(norm, steps=1500, seed=0):
torch.manual_seed(seed)
model = build(norm)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x, unlab_y), batch_size=128,
shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
with torch.no_grad():
model.eval()
proper = (model(test_x).argmax(1) == test_y).float().mean()
# what the same weights score using each batch's own statistics
model.train()
batched = (model(test_x).argmax(1) == test_y).float().mean()
return float(proper), float(batched)
print('%-14s %18s %18s' % ('normalisation', 'running stats',
'batch stats'))
for norm in ['batch', 'group']:
proper, batched = compare(norm)
print('%-14s %18.4f %18.4f' % (norm, proper, batched))
normalisation running stats batch stats
batch 0.9870 0.9845
group 0.9860 0.9860

On this architecture they are the same to within a rounding error, and batch normalisation's two columns agree, so there is no drama to report. Group normalisation is kept anyway, on the grounds that it cannot develop a dependence on batch composition later, and it costs nothing to have that guarantee.

How everything will be scored

Linear probe: Freeze the encoder completely, compute its features once, and fit a single linear layer on top using however many labels you are willing to spend. The accuracy that layer reaches is the measure of the representation, because a linear layer can only exploit structure the encoder already made explicit.

Two baselines matter, and skipping either one is how people convince themselves a pretext task worked when it did not. The first is training the whole encoder supervised on the same labels, which is what you would do if you never heard of any of this.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
print('%10s %14s' % ('labels', 'test accuracy'))
for n in [100, 500, 2000, 12000]:
print('%10d %14.4f' % (n, supervised(n)))
labels test accuracy
100 0.7320
500 0.9275
2000 0.9610
12000 0.9860

Budget in steps, not epochs

The first version of this comparison ran a fixed thirty epochs at every label count. With 100 labels that is 30 optimiser steps and with 12000 it is 2820, so the small settings were not underperforming for want of data, they were underperforming for want of training. Every run here gets 1500 steps regardless of how much data it has, which is the only version of the comparison that means anything.

The second baseline is the one that is almost always missing: the same encoder with its weights left exactly as initialised, never trained on anything at all.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
torch.manual_seed(0)
enc = Encoder()
print('%10s %14s' % ('labels', 'probe accuracy'))
for n in [100, 500, 2000, 12000]:
print('%10d %14.4f' % (n, linear_probe(enc, n)))
labels probe accuracy
100 0.6340
500 0.7780
2000 0.8205
12000 0.8390

A random convolutional stack, with no training whatsoever, reaches 0.839 with a linear layer on top. That is not a quirk of MNIST. Random convolutions are a genuinely strong feature extractor, because a random filter bank still measures edges and textures at every position, and pooling still summarises them.

The number every result this week has to beat

Not zero, and not chance. Any pretext task that leaves the encoder below 0.839 at 12000 labels has made the representation worse than never training it. Two of the three tasks in this week do exactly that.
Week 14 · Day 2 of 7

A Pretext Task That Made Things Worse

Rotation prediction, solved in three epochs, and what that cost

By 1654 words

The first pretext task to try is the one that sounds cleverest. Rotate each image by a quarter, a half or three quarters of a turn, and train the encoder to say which. The label is free, since it is whatever rotation the code just applied, and the argument for it is appealing: to know that a digit has been turned upside down, surely you have to know something about what the digit is.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def rotate_batch(x):
"""Four copies of every image, one per quarter turn, with the turn
itself as the label. No human wrote any of these labels down."""
outs, labels = [], []
for k in range(4):
outs.append(torch.rot90(x, k, dims=(2, 3)))
labels.append(torch.full((len(x),), k, dtype=torch.long))
return torch.cat(outs), torch.cat(labels)

def train_rotation(epochs=12, seed=0, trace_every=0):
torch.manual_seed(seed)
enc = Encoder()
head = nn.Linear(64, 4)
opt = torch.optim.AdamW(list(enc.parameters()) + list(head.parameters()),
lr=1e-3)
lf = nn.CrossEntropyLoss()
for ep in range(1, epochs + 1):
enc.train()
right = seen = 0
for (xb,) in pool:
rx, ry = rotate_batch(xb)
logits = head(enc(rx))
opt.zero_grad()
lf(logits, ry).backward()
opt.step()
right += (logits.argmax(1) == ry).sum().item()
seen += len(ry)
if trace_every and ep % trace_every == 0:
print('%6d %20.4f' % (ep, right / seen))
return enc
x = unlab_x[:2]
rx, ry = rotate_batch(x)
print('2 images became %d, with labels %s'
% (len(rx), ry.reshape(4, 2)[:, 0].tolist()))
2 images became 8, with labels [0, 1, 2, 3]

Training it

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def rotate_batch(x):
"""Four copies of every image, one per quarter turn, with the turn
itself as the label. No human wrote any of these labels down."""
outs, labels = [], []
for k in range(4):
outs.append(torch.rot90(x, k, dims=(2, 3)))
labels.append(torch.full((len(x),), k, dtype=torch.long))
return torch.cat(outs), torch.cat(labels)

def train_rotation(epochs=12, seed=0, trace_every=0):
torch.manual_seed(seed)
enc = Encoder()
head = nn.Linear(64, 4)
opt = torch.optim.AdamW(list(enc.parameters()) + list(head.parameters()),
lr=1e-3)
lf = nn.CrossEntropyLoss()
for ep in range(1, epochs + 1):
enc.train()
right = seen = 0
for (xb,) in pool:
rx, ry = rotate_batch(xb)
logits = head(enc(rx))
opt.zero_grad()
lf(logits, ry).backward()
opt.step()
right += (logits.argmax(1) == ry).sum().item()
seen += len(ry)
if trace_every and ep % trace_every == 0:
print('%6d %20.4f' % (ep, right / seen))
return enc
print('%6s %20s' % ('epoch', 'rotation accuracy'))
enc = train_rotation(epochs=12, trace_every=3)
print()
print('%10s %14s' % ('labels', 'probe accuracy'))
for n in [100, 500, 2000, 12000]:
print('%10d %14.4f' % (n, linear_probe(enc, n)))
epoch rotation accuracy
3 0.9931
6 0.9961
9 0.9978
12 0.9984

labels probe accuracy
100 0.6020
500 0.7315
2000 0.7740
12000 0.7970

Read the two halves of that output against each other, because together they are the whole lesson of the day. The pretext task is solved: 99.3 percent by epoch three, 99.8 by the end. And the representation it produced is worse than the untrained encoder at every single label count, by four to five points.

Twelve epochs of training made the encoder worse than not training it

0.797 against 0.839 at 12000 labels, 0.602 against 0.634 at 100. This is not a small effect and it is not noise. The encoder spent its capacity becoming good at something, and that something was not useful, so the random features it started with were degraded rather than improved.

Why it failed, and how you could have known

The diagnostic is sitting in the first column. A pretext task that reaches 99.3 percent after three epochs is not asking the encoder for much. Deciding which way up an MNIST digit is turned can be done from where the ink sits and which way the strokes lean, and neither of those requires any notion of which digit it is. The encoder found the cheapest solution available, exactly as it was asked to.

So the rule is not that rotation prediction is a bad idea. It is that a pretext task is only as good as the shortcut it fails to leave open, and the way you find the shortcut is to watch how fast the task is solved. If your pretext accuracy saturates almost immediately, the task is too easy and the representation will be disappointing.

This is the honest history of the field

Rotation prediction, jigsaw puzzles and colourisation were all real published methods, and they were all overtaken. The reason was not that anyone was careless. It was that the relationship between a pretext task and the representation it produces is not something you can reason your way to, and the field found this out by measuring.

Week 14 · Day 3 of 7

Filling In What You Hid

Masked patches, and why the same idea is stronger for text

By 1620 words

The second task takes the language model idea directly. Hide part of the input, ask the model to reconstruct it, and the hidden part is its own label. Here that means blanking six of the sixteen 7 by 7 patches of each image.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def mask_batch(x, patch=7, blank=6):
"""Blank six of the sixteen 7 by 7 patches, independently per image."""
n, grid = len(x), 28 // patch
# argsort of uniform noise is a batched randperm
cells = torch.rand(n, grid * grid).argsort(1)[:, :blank]
keep = torch.ones(n, grid * grid)
keep.scatter_(1, cells, 0.0)
keep = keep.reshape(n, 1, grid, grid)
keep = keep.repeat_interleave(patch, 2).repeat_interleave(patch, 3)
return x * keep

class MaskDecoder(nn.Module):
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, 256), nn.ReLU(),
nn.Linear(256, 784), nn.Sigmoid())

def forward(self, h):
return self.net(h).reshape(-1, 1, 28, 28)

def train_masked(epochs=25, seed=0, trace_every=0):
torch.manual_seed(seed)
enc, dec = Encoder(), MaskDecoder()
opt = torch.optim.AdamW(list(enc.parameters()) + list(dec.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
loss = ((dec(enc(mask_batch(xb))) - xb) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc
torch.manual_seed(0)
x = unlab_x[:1]
for img in [x, mask_batch(x)]:
t = img[0, 0]
for row in t[::2]:
print(''.join(' .:-=+*#%@'[min(9, int(v * 9.999))]
for v in row.tolist()))
print()



..-**@@@@@%*@@#:
%@@@@@##@@
*@-
#@:
-@@@=
-@@#
.+#@@%
=%@@@@#-
*%@@@@#-
+@@@%++





..-**@
%@@@@@##@@
*@-
#@:
-@@@=
-@@#
.+#@@%
=%@@@@#-


This is a harder question than which way up it is. To fill a blank square in the middle of a digit you have to know what kind of stroke was passing through it, and that depends on the rest of the image.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def mask_batch(x, patch=7, blank=6):
"""Blank six of the sixteen 7 by 7 patches, independently per image."""
n, grid = len(x), 28 // patch
# argsort of uniform noise is a batched randperm
cells = torch.rand(n, grid * grid).argsort(1)[:, :blank]
keep = torch.ones(n, grid * grid)
keep.scatter_(1, cells, 0.0)
keep = keep.reshape(n, 1, grid, grid)
keep = keep.repeat_interleave(patch, 2).repeat_interleave(patch, 3)
return x * keep

class MaskDecoder(nn.Module):
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, 256), nn.ReLU(),
nn.Linear(256, 784), nn.Sigmoid())

def forward(self, h):
return self.net(h).reshape(-1, 1, 28, 28)

def train_masked(epochs=25, seed=0, trace_every=0):
torch.manual_seed(seed)
enc, dec = Encoder(), MaskDecoder()
opt = torch.optim.AdamW(list(enc.parameters()) + list(dec.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
loss = ((dec(enc(mask_batch(xb))) - xb) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc
print('%6s %14s' % ('epoch', 'pixel MSE'))
enc = train_masked(epochs=25, trace_every=5)
print()
print('%10s %14s' % ('labels', 'probe accuracy'))
for n in [100, 500, 2000, 12000]:
print('%10d %14.4f' % (n, linear_probe(enc, n)))
epoch pixel MSE
5 0.0254
10 0.0182
15 0.0157
20 0.0145
25 0.0137

labels probe accuracy
100 0.6025
500 0.7890
2000 0.8585
12000 0.8910

Better, and still not good

0.891 at 12000 labels, against 0.839 for the untrained encoder. So it did learn something, which rotation prediction did not. It is also still far behind the 0.986 that plain supervised training reaches on the same labels, and at 100 labels it is level with the random encoder and behind supervised training by thirteen points.

The reason is the loss. Reconstructing pixels rewards getting the background right, and most of an MNIST image is background. A great deal of the capacity goes into a skill that has nothing to do with telling digits apart, which is the same complaint week 12 ended with and it has the same cause.

Why it works so much better for text

Masked prediction is the method behind an enormous amount of modern language modelling, so its middling showing here needs explaining rather than generalising.

Filling in a missing word requires the meaning of the sentence, because there is no low-level cue that gets you there. There is no equivalent of a blank background in text. Filling in a missing patch of image very often only requires knowing that the neighbourhood was dark. The method did not change between the two, the redundancy of the data did.

Week 14 · Day 4 of 7

Same Thing, Different View

Contrastive learning, the augmentations, and the NT-Xent loss

By 2059 words

The third approach stops asking the encoder to reconstruct anything. Take one image, make two different corrupted views of it, and ask only that those two views land near each other and away from every other image in the batch. Nothing is predicted and nothing is reconstructed. The only requirement is on the geometry of the representation.

Contrastive learning: A loss that pulls two views of the same input together and pushes different inputs apart. The views come from augmentation, so what the encoder is being told is exactly this: whatever these transformations changed is not part of the identity of the input.

The augmentations are the whole design

This is the part that decides everything. The encoder learns to ignore whatever the augmentation varies, so the choice of augmentation is a direct statement about what you consider irrelevant. Shift the digit and the encoder learns position does not matter. Scrub a square out of it and the encoder learns not to depend on any one region.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def augment(x):
"""Two different views of the same digit: shift it, scrub an 8 by 8
square out of it, and add a little noise.

Written with advanced indexing rather than a loop over the batch. The
loop version is easier to read and about twenty times slower, and
this runs twice per batch for every epoch.
"""
n = len(x)
idx = torch.arange(28)
pad = torch.nn.functional.pad(x, (3, 3, 3, 3))
rows = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
cols = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
out = pad[torch.arange(n).reshape(n, 1, 1), 0,
rows.reshape(n, 28, 1), cols.reshape(n, 1, 28)]
out = out.unsqueeze(1)
cy, cx = torch.randint(0, 21, (2, n, 1, 1))
ys, xs = idx.reshape(1, 28, 1), idx.reshape(1, 1, 28)
hole = (ys >= cy) & (ys < cy + 8) & (xs >= cx) & (xs < cx + 8)
out = out * (~hole).unsqueeze(1)
return (out + 0.05 * torch.randn_like(out)).clamp(0, 1)

class Projector(nn.Module):
"""The head the loss is applied to, thrown away afterwards."""
def __init__(self, dim=64, out=32):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(),
nn.Linear(dim, out))

def forward(self, h):
return nn.functional.normalize(self.net(h), dim=1)

def nt_xent(z1, z2, temperature=0.5):
"""Each view should be closer to its partner than to anything else
in the batch."""
z = torch.cat([z1, z2])
n = len(z1)
sim = z @ z.t() / temperature
sim.fill_diagonal_(-1e9)
target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
return nn.functional.cross_entropy(sim, target)

def train_contrastive(epochs=30, seed=0, trace_every=0, temperature=0.5,
negatives=True):
torch.manual_seed(seed)
enc, proj = Encoder(), Projector()
opt = torch.optim.AdamW(list(enc.parameters()) + list(proj.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
z1, z2 = proj(enc(augment(xb))), proj(enc(augment(xb)))
if negatives:
loss = nt_xent(z1, z2, temperature)
else:
# pull the pair together and ask for nothing else
loss = ((z1 - z2) ** 2).sum(1).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc, proj
torch.manual_seed(0)
views = augment(unlab_x[:1].repeat(2, 1, 1, 1))
for img in views:
for row in img[0][::2]:
print(''.join(' .:-=+*#%@'[min(9, int(v * 9.999))]
for v in row.tolist()))
print()
.
. .
#:#@@=
:@@ :=-.
-+
+@*
@%*= .
:#@@*.
@@@:
.+@@@@@%
:%@@@@#-..
:*%@@@@@+.

. .




::-=#@@@@@%%@@%.
%@@@@@##@@


. =.
@@#
.+#@@%
. =%@@@@#:
#%@@@@#:
.+@@@%*+. .
.

An augmentation that changes the answer is a bug

Week 4 made this point about supervised training and it is sharper here. A vertical flip would teach this encoder that a 6 and a 9 are the same thing, and nothing in the loss would object, because the loss never sees a digit label. The failure would only appear later, as a probe that cannot separate two classes and no explanation of why.

The loss

For a batch of n images there are 2n views. Each one has exactly one partner and 2n minus 2 non-partners, so scoring which of the others is the partner is an ordinary classification problem with 2n minus 1 options, and cross entropy handles it. The similarities are dot products of normalised vectors, divided by a temperature.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def augment(x):
"""Two different views of the same digit: shift it, scrub an 8 by 8
square out of it, and add a little noise.

Written with advanced indexing rather than a loop over the batch. The
loop version is easier to read and about twenty times slower, and
this runs twice per batch for every epoch.
"""
n = len(x)
idx = torch.arange(28)
pad = torch.nn.functional.pad(x, (3, 3, 3, 3))
rows = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
cols = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
out = pad[torch.arange(n).reshape(n, 1, 1), 0,
rows.reshape(n, 28, 1), cols.reshape(n, 1, 28)]
out = out.unsqueeze(1)
cy, cx = torch.randint(0, 21, (2, n, 1, 1))
ys, xs = idx.reshape(1, 28, 1), idx.reshape(1, 1, 28)
hole = (ys >= cy) & (ys < cy + 8) & (xs >= cx) & (xs < cx + 8)
out = out * (~hole).unsqueeze(1)
return (out + 0.05 * torch.randn_like(out)).clamp(0, 1)

class Projector(nn.Module):
"""The head the loss is applied to, thrown away afterwards."""
def __init__(self, dim=64, out=32):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(),
nn.Linear(dim, out))

def forward(self, h):
return nn.functional.normalize(self.net(h), dim=1)

def nt_xent(z1, z2, temperature=0.5):
"""Each view should be closer to its partner than to anything else
in the batch."""
z = torch.cat([z1, z2])
n = len(z1)
sim = z @ z.t() / temperature
sim.fill_diagonal_(-1e9)
target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
return nn.functional.cross_entropy(sim, target)

def train_contrastive(epochs=30, seed=0, trace_every=0, temperature=0.5,
negatives=True):
torch.manual_seed(seed)
enc, proj = Encoder(), Projector()
opt = torch.optim.AdamW(list(enc.parameters()) + list(proj.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
z1, z2 = proj(enc(augment(xb))), proj(enc(augment(xb)))
if negatives:
loss = nt_xent(z1, z2, temperature)
else:
# pull the pair together and ask for nothing else
loss = ((z1 - z2) ** 2).sum(1).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc, proj
torch.manual_seed(0)
z1 = nn.functional.normalize(torch.randn(4, 32), dim=1)
print('a random pairing scores %.4f' % nt_xent(z1, z1[torch.randperm(4)]))
print('a perfect pairing scores %.4f' % nt_xent(z1, z1))
print('chance for a batch of 4 is ln(7) = %.4f'
% float(torch.tensor(7.0).log()))
a random pairing scores 2.1051
a perfect pairing scores 0.5204
chance for a batch of 4 is ln(7) = 1.9459

Two details in that loss are easy to get wrong. The diagonal has to be removed, or every view trivially matches itself and there is nothing to learn. And the vectors have to be normalised before the dot product, or the model can raise every similarity by making its outputs longer instead of by arranging them better.

Week 14 · Day 5 of 7

Where It Finally Pays

The crossover against supervised training, and reading a loss that stays hard

By 1103 words

Train it on the same pool, for the same kind of budget, and probe it the same way.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def augment(x):
"""Two different views of the same digit: shift it, scrub an 8 by 8
square out of it, and add a little noise.

Written with advanced indexing rather than a loop over the batch. The
loop version is easier to read and about twenty times slower, and
this runs twice per batch for every epoch.
"""
n = len(x)
idx = torch.arange(28)
pad = torch.nn.functional.pad(x, (3, 3, 3, 3))
rows = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
cols = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
out = pad[torch.arange(n).reshape(n, 1, 1), 0,
rows.reshape(n, 28, 1), cols.reshape(n, 1, 28)]
out = out.unsqueeze(1)
cy, cx = torch.randint(0, 21, (2, n, 1, 1))
ys, xs = idx.reshape(1, 28, 1), idx.reshape(1, 1, 28)
hole = (ys >= cy) & (ys < cy + 8) & (xs >= cx) & (xs < cx + 8)
out = out * (~hole).unsqueeze(1)
return (out + 0.05 * torch.randn_like(out)).clamp(0, 1)

class Projector(nn.Module):
"""The head the loss is applied to, thrown away afterwards."""
def __init__(self, dim=64, out=32):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(),
nn.Linear(dim, out))

def forward(self, h):
return nn.functional.normalize(self.net(h), dim=1)

def nt_xent(z1, z2, temperature=0.5):
"""Each view should be closer to its partner than to anything else
in the batch."""
z = torch.cat([z1, z2])
n = len(z1)
sim = z @ z.t() / temperature
sim.fill_diagonal_(-1e9)
target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
return nn.functional.cross_entropy(sim, target)

def train_contrastive(epochs=30, seed=0, trace_every=0, temperature=0.5,
negatives=True):
torch.manual_seed(seed)
enc, proj = Encoder(), Projector()
opt = torch.optim.AdamW(list(enc.parameters()) + list(proj.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
z1, z2 = proj(enc(augment(xb))), proj(enc(augment(xb)))
if negatives:
loss = nt_xent(z1, z2, temperature)
else:
# pull the pair together and ask for nothing else
loss = ((z1 - z2) ** 2).sum(1).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc, proj
print('%6s %14s' % ('epoch', 'NT-Xent loss'))
enc, proj = train_contrastive(epochs=30, trace_every=5)
torch.save(enc.state_dict(), 'con_enc.pt')
torch.save(proj.state_dict(), 'con_proj.pt')
print()
print('%10s %14s' % ('labels', 'probe accuracy'))
for n in [100, 500, 2000, 12000]:
print('%10d %14.4f' % (n, linear_probe(enc, n)))
epoch NT-Xent loss
5 4.6325
10 4.5591
15 4.5258
20 4.5084
25 4.4926
30 4.4806

labels probe accuracy
100 0.7860
500 0.9330
2000 0.9530
12000 0.9635

At 100 labels this beats training the encoder supervised on those same 100 labels, 0.786 against 0.732. At 500 it is ahead again, 0.933 against 0.928. At 2000 and above supervised training pulls back in front, ending 0.986 against 0.964.

That crossover is the entire practical case for the method

Self-supervised pretraining is not a way to beat supervised learning. It is a way to spend unlabelled data, which is usually abundant, to reduce how much labelled data you need, which usually is not. When labels are plentiful it loses, and it is supposed to.

The crossover here sits somewhere between 500 and 2000 labels. Where it sits for your problem depends on how hard the task is and how much unlabelled data you have, and the only way to find it is the table above.

Reading the loss

The NT-Xent loss went from 4.63 to 4.48 over thirty epochs. Chance for a batch of 256 is ln(511), about 6.24, so the model is well clear of guessing, and it is also nowhere near solving the task. Compare that with rotation prediction, which was at 99.3 percent after three epochs. The task that stayed hard is the task that produced the useful representation, and the loss was still falling when the budget ran out.

Week 14 · Day 6 of 7

Collapse, Negatives and the Discarded Head

The degenerate solution, what prevents it, and why the loss sees a different layer

By 2261 words

There is an obvious way for a contrastive encoder to cheat. If it maps every input to the same vector, then the two views of an image are certainly close together, and the loss that asks for closeness is perfectly satisfied. The representation is worthless.

Representation collapse: When an encoder outputs nearly the same vector regardless of input. It satisfies any loss built only from pulling similar things together, and it is invisible unless you measure the spread of the outputs or probe them.

What actually prevents it

The negatives. Every other image in the batch is a thing this view has to be far from, so a constant output is now the worst possible answer rather than the best. Take that term away and keep everything else, and here is what happens:

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def augment(x):
"""Two different views of the same digit: shift it, scrub an 8 by 8
square out of it, and add a little noise.

Written with advanced indexing rather than a loop over the batch. The
loop version is easier to read and about twenty times slower, and
this runs twice per batch for every epoch.
"""
n = len(x)
idx = torch.arange(28)
pad = torch.nn.functional.pad(x, (3, 3, 3, 3))
rows = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
cols = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
out = pad[torch.arange(n).reshape(n, 1, 1), 0,
rows.reshape(n, 28, 1), cols.reshape(n, 1, 28)]
out = out.unsqueeze(1)
cy, cx = torch.randint(0, 21, (2, n, 1, 1))
ys, xs = idx.reshape(1, 28, 1), idx.reshape(1, 1, 28)
hole = (ys >= cy) & (ys < cy + 8) & (xs >= cx) & (xs < cx + 8)
out = out * (~hole).unsqueeze(1)
return (out + 0.05 * torch.randn_like(out)).clamp(0, 1)

class Projector(nn.Module):
"""The head the loss is applied to, thrown away afterwards."""
def __init__(self, dim=64, out=32):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(),
nn.Linear(dim, out))

def forward(self, h):
return nn.functional.normalize(self.net(h), dim=1)

def nt_xent(z1, z2, temperature=0.5):
"""Each view should be closer to its partner than to anything else
in the batch."""
z = torch.cat([z1, z2])
n = len(z1)
sim = z @ z.t() / temperature
sim.fill_diagonal_(-1e9)
target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
return nn.functional.cross_entropy(sim, target)

def train_contrastive(epochs=30, seed=0, trace_every=0, temperature=0.5,
negatives=True):
torch.manual_seed(seed)
enc, proj = Encoder(), Projector()
opt = torch.optim.AdamW(list(enc.parameters()) + list(proj.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
z1, z2 = proj(enc(augment(xb))), proj(enc(augment(xb)))
if negatives:
loss = nt_xent(z1, z2, temperature)
else:
# pull the pair together and ask for nothing else
loss = ((z1 - z2) ** 2).sum(1).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc, proj
def spread(enc):
"""How different are the embeddings from each other at all."""
f = features(enc, unlab_x[:2000])
return f.std(0).mean().item()

torch.manual_seed(0)
fresh = Encoder()
print('%-34s %12s %14s' % ('', 'spread', 'probe at 500'))
print('%-34s %12.4f %14.4f'
% ('untrained encoder', spread(fresh), linear_probe(fresh, 500)))
for negatives in [True, False]:
enc, _ = train_contrastive(epochs=15, negatives=negatives)
name = 'with negatives' if negatives else 'positives only'
print('%-34s %12.4f %14.4f'
% (name, spread(enc), linear_probe(enc, 500)))
spread probe at 500
untrained encoder 0.1877 0.7780
with negatives 1.9468 0.9110
positives only 0.0812 0.2300

The positives-only run does exactly what the argument predicts. Its embeddings have a spread of 0.081, less than half the untrained encoder's 0.188 and a twenty-fourth of what the proper loss produces, which is to say the outputs have become nearly constant. Its probe lands at 0.230, against 0.778 for an encoder that was never trained at all. Fifteen epochs of training destroyed the representation, and the objective it was given was satisfied throughout.

A falling loss is not evidence of anything here

This is the same warning as week 13, arriving from a different direction. In self-supervised training the loss is over a task you invented, and a model can drive it down by finding a degenerate solution to your invented task. Print the spread of the embeddings, and probe against the untrained encoder. Neither costs anything and together they catch this immediately.

Methods that avoid negatives

Negatives are expensive: the loss wants large batches, because a batch of 256 offers 510 things to be unlike and a batch of 32 offers 62. A family of later methods removed them, using an asymmetry between the two branches instead, typically a slowly-updated copy of the encoder on one side and a stop-gradient so that only the other side learns. The point to carry away is not the mechanism but that every one of them exists to solve the problem measured above, and any of them can be checked the same way.

The head you throw away

The loss was not applied to the encoder's output. It was applied to a small projection on top, which is then discarded, and the encoder underneath is what gets used. That looks like a pointless extra layer until you probe both:

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

def stack(dataset, n, start=0):
"""One pass, so images and labels stay together."""
xs = torch.stack([dataset[i][0] for i in range(start, start + n)])
ys = torch.tensor([dataset[i][1] for i in range(start, start + n)])
return xs, ys

# 12000 images whose labels we pretend not to have, and a held out set
unlab_x, unlab_y = stack(train_all, 12000)
test_x, test_y = stack(test_all, 2000)
pool = DataLoader(TensorDataset(unlab_x), batch_size=256, shuffle=True,
drop_last=True)
class Encoder(nn.Module):
"""The part we are trying to train without labels.

GroupNorm rather than BatchNorm, deliberately. This encoder gets
trained one way and evaluated another, often on batches of a very
different size, and BatchNorm's running statistics do not survive
that. Day 1 measures what it costs.
"""
def __init__(self, dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.GroupNorm(8, 32), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 64, 3, padding=1), nn.GroupNorm(8, 64), nn.ReLU(),
nn.Flatten(),
nn.Linear(64 * 7 * 7, dim))

def forward(self, x):
return self.net(x)
@torch.no_grad()
def features(enc, x, batch=512):
enc.eval()
return torch.cat([enc(x[i:i + batch]) for i in range(0, len(x), batch)])

def linear_probe(enc, n_labels, epochs=300, seed=0):
"""Freeze the encoder, fit one linear layer on n_labels examples."""
torch.manual_seed(seed)
tr_x, tr_y = unlab_x[:n_labels], unlab_y[:n_labels]
ftr, fte = features(enc, tr_x), features(enc, test_x)
head = nn.Linear(ftr.shape[1], 10)
opt = torch.optim.AdamW(head.parameters(), lr=1e-2, weight_decay=1e-4)
lf = nn.CrossEntropyLoss()
for _ in range(epochs):
opt.zero_grad()
lf(head(ftr), tr_y).backward()
opt.step()
with torch.no_grad():
return (head(fte).argmax(1) == test_y).float().mean().item()

def supervised(n_labels, steps=1500, seed=0):
"""The same encoder trained from scratch on the same labels.

Budgeted in optimiser steps rather than epochs, so that 100 labels
and 12000 labels each get the same amount of training. A fixed epoch
count gives the small settings almost no steps and makes them look
far worse than they are.
"""
torch.manual_seed(seed)
model = nn.Sequential(Encoder(), nn.Linear(64, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(unlab_x[:n_labels],
unlab_y[:n_labels]),
batch_size=128, shuffle=True)
done = 0
while done < steps:
model.train()
for xb, yb in loader:
opt.zero_grad()
lf(model(xb), yb).backward()
opt.step()
done += 1
model.eval()
with torch.no_grad():
return (model(test_x).argmax(1) == test_y).float().mean().item()
def augment(x):
"""Two different views of the same digit: shift it, scrub an 8 by 8
square out of it, and add a little noise.

Written with advanced indexing rather than a loop over the batch. The
loop version is easier to read and about twenty times slower, and
this runs twice per batch for every epoch.
"""
n = len(x)
idx = torch.arange(28)
pad = torch.nn.functional.pad(x, (3, 3, 3, 3))
rows = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
cols = torch.randint(0, 7, (n, 1)) + idx.reshape(1, 28)
out = pad[torch.arange(n).reshape(n, 1, 1), 0,
rows.reshape(n, 28, 1), cols.reshape(n, 1, 28)]
out = out.unsqueeze(1)
cy, cx = torch.randint(0, 21, (2, n, 1, 1))
ys, xs = idx.reshape(1, 28, 1), idx.reshape(1, 1, 28)
hole = (ys >= cy) & (ys < cy + 8) & (xs >= cx) & (xs < cx + 8)
out = out * (~hole).unsqueeze(1)
return (out + 0.05 * torch.randn_like(out)).clamp(0, 1)

class Projector(nn.Module):
"""The head the loss is applied to, thrown away afterwards."""
def __init__(self, dim=64, out=32):
super().__init__()
self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(),
nn.Linear(dim, out))

def forward(self, h):
return nn.functional.normalize(self.net(h), dim=1)

def nt_xent(z1, z2, temperature=0.5):
"""Each view should be closer to its partner than to anything else
in the batch."""
z = torch.cat([z1, z2])
n = len(z1)
sim = z @ z.t() / temperature
sim.fill_diagonal_(-1e9)
target = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
return nn.functional.cross_entropy(sim, target)

def train_contrastive(epochs=30, seed=0, trace_every=0, temperature=0.5,
negatives=True):
torch.manual_seed(seed)
enc, proj = Encoder(), Projector()
opt = torch.optim.AdamW(list(enc.parameters()) + list(proj.parameters()),
lr=1e-3)
for ep in range(1, epochs + 1):
enc.train()
total = n = 0.0
for (xb,) in pool:
z1, z2 = proj(enc(augment(xb))), proj(enc(augment(xb)))
if negatives:
loss = nt_xent(z1, z2, temperature)
else:
# pull the pair together and ask for nothing else
loss = ((z1 - z2) ** 2).sum(1).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
return enc, proj
enc, proj = Encoder(), Projector()
enc.load_state_dict(torch.load('con_enc.pt'))
proj.load_state_dict(torch.load('con_proj.pt'))

class Both(nn.Module):
def __init__(self, enc, proj):
super().__init__()
self.enc, self.proj = enc, proj

def forward(self, x):
return self.proj(self.enc(x))

print('%-40s %14s' % ('features taken from', 'probe at 500'))
print('%-40s %14.4f' % ('the encoder, 64 wide',
linear_probe(enc, 500)))
print('%-40s %14.4f' % ('the projection the loss saw, 32 wide',
linear_probe(Both(enc, proj), 500)))
features taken from probe at 500
the encoder, 64 wide 0.9330
the projection the loss saw, 32 wide 0.8530

The layer the loss was applied to is the worse representation. It has been shaped to discard everything the augmentations varied, which is precisely its job, and some of what the augmentations varied turns out to matter for classification. Keeping a layer between the encoder and the loss gives the encoder somewhere to retain that information while still satisfying the objective.

Week 14 · Day 7 of 7

The Whole Table

Four encoders compared, and the protocol worth reusing

By 497 words

Four ways of getting an encoder, one architecture, one pool of data, one probe.

labelsuntrainedrotationmasked patchescontrastivesupervised
1000.6340.6020.6030.7860.732
5000.7780.7320.7890.9330.928
20000.8210.7740.8590.9530.961
120000.8390.7970.8910.9640.986

Three findings, and the second is the one worth carrying furthest.

  1. Contrastive learning is the only one of the three that clearly helped, and it helped exactly where the theory says it should: below about a thousand labels, where it beat training the same encoder supervised on the same labels.
  2. Rotation prediction was worse than not training at all, at every label count. A pretext task can actively damage a representation, and nothing in its own loss will tell you.
  3. The untrained encoder is a serious competitor. It beat one of the three methods outright and matched a second at low label counts.

What this table is not

MNIST is close to the worst possible dataset for showing self-supervised learning at its best. It is small, clean, low resolution and nearly linearly separable in raw pixels, so the supervised baseline is very strong and there is little structure left for a pretext task to discover. On natural images at scale the gaps run the other way and the crossover point sits far higher.

The findings that do transfer are the method, not the numbers: always include an untrained encoder, watch how quickly the pretext task saturates, and check the spread of the embeddings.

Choosing a pretext task

  • Ask what shortcut the task leaves open. If there is a low-level cue that solves it, the encoder will find that cue and learn nothing else.
  • Watch the pretext metric. Saturating in the first few epochs is the single clearest sign that the task is too easy to be worth training on.
  • For contrastive methods, the augmentations are the design. They are a statement of what you consider irrelevant, and an augmentation that changes the true label is a silent bug.
  • Keep a projection head between the encoder and the loss, and probe the encoder rather than the projection.
  • Measure the spread of the embeddings, and treat a falling loss with suspicion until you have.

The protocol, written out

  1. Fix one architecture and one data pool for the whole comparison.
  2. Probe the untrained encoder first, and write the number down.
  3. Probe a supervised model trained on each label count you care about, budgeted in optimiser steps rather than epochs.
  4. Train each pretext task, then probe with the encoder frozen.
  5. Report the whole curve across label counts. A single label count can support whichever conclusion you were hoping for.

What the week actually establishes

Self-supervised learning is a way to convert unlabelled data into fewer required labels, and whether it succeeds depends entirely on the pretext task denying the encoder every shortcut but the one you want. Two of the three tasks here failed that test, and only a control nobody usually runs made that visible.