Lesson 51: Masked Autoencoders¶
Lesson 48's autoencoder compressed a whole image through a narrow bottleneck vector. Masked autoencoders (MAE, He et al., 2021) create the information bottleneck a completely different way: chop the image into patches (Lesson 47's ViT patchify), hide most of them entirely — not noised, not blurred, just never shown to the encoder at all — and train the network to reconstruct the missing patches from whatever's left. This is the third self-supervised pretraining paradigm in this part, after contrastive learning (Lesson 49) and self-distillation (Lesson 50): no labels, no negative pairs, no teacher network — just a plausible-sounding fill-in-the-blank task invented directly from unlabeled images.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Patchify and randomly mask¶
Same 16x16 image, 4x4 patches as Lesson 47 — 16 patches total. For each image, pick a random subset to keep visible and discard the rest; the discarded patches are not part of the encoder's input in any form, not even as zeros or noise.
IMG = 16
PATCH = 4
N_PATCHES = (IMG // PATCH) ** 2 # 16
PATCH_DIM = PATCH * PATCH # 16 (grayscale)
def patchify(img, patch_size=PATCH):
B, C, H, W = img.shape
P = patch_size
patches = img.unfold(2, P, P).unfold(3, 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)
def unpatchify(patches, patch_size=PATCH, img_size=IMG):
B, N, D = patches.shape
P = patch_size
n_side = img_size // P
x = patches.view(B, n_side, n_side, P, P).permute(0, 1, 3, 2, 4).contiguous()
return x.view(B, 1, img_size, img_size)
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
def random_masking(x, mask_ratio, generator=None):
B, N, D = x.shape
n_visible = max(1, int(N * (1 - mask_ratio)))
noise = torch.rand(B, N, generator=generator)
ids_shuffle = torch.argsort(noise, dim=1) # random permutation of patch indices per image
ids_keep = ids_shuffle[:, :n_visible] # kept as VISIBLE, in original patch-index units
ids_mask = ids_shuffle[:, n_visible:] # HIDDEN entirely from the encoder
x_visible = torch.gather(x, 1, ids_keep.unsqueeze(-1).expand(-1, -1, D))
return x_visible, ids_keep, ids_mask
def make_image(shape_type, cx, cy, size=IMG):
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)
rng = np.random.default_rng(6)
X_train, y_train = make_dataset(rng, 400)
X_test, y_test = make_dataset(rng, 150)
Xt = torch.tensor(X_train).unsqueeze(1)
Xte = torch.tensor(X_test).unsqueeze(1)
print(f'{N_PATCHES} patches per image, {PATCH_DIM} pixels each')
An asymmetric encoder-decoder¶
The encoder is a small Transformer (Lesson 46) that only ever processes the visible patches — masked patches never enter it, so the encoder is fast and never wastes computation on tokens that get thrown away. The decoder is where the masked patches reappear: a shared, learned mask token vector stands in for every hidden patch, combined with that patch's original positional encoding (so the decoder knows where each blank belongs, even though it never saw what was there), and a lightweight Transformer predicts each masked patch's raw pixel values. Loss is computed only on the masked patches — the network gets no credit for copying pixels it was already shown.
class MAEEncoder(nn.Module):
def __init__(self, patch_dim=PATCH_DIM, embed_dim=32, n_heads=4, n_layers=2):
super().__init__()
self.embed = nn.Linear(patch_dim, embed_dim)
self.register_buffer('pe', positional_encoding(N_PATCHES, 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)
def forward(self, x_visible, ids_keep):
tok = self.embed(x_visible) + self.pe[ids_keep] # positional encoding at each patch's TRUE index
return self.encoder(tok)
class MAEDecoder(nn.Module):
def __init__(self, embed_dim=32, decoder_dim=16, patch_dim=PATCH_DIM, n_heads=4, n_layers=1):
super().__init__()
self.embed = nn.Linear(embed_dim, decoder_dim)
self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
self.register_buffer('pe', positional_encoding(N_PATCHES, decoder_dim))
layer = nn.TransformerEncoderLayer(decoder_dim, n_heads, dim_feedforward=decoder_dim * 2,
batch_first=True, dropout=0.0)
self.decoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.pred = nn.Linear(decoder_dim, patch_dim)
self.decoder_dim = decoder_dim
def forward(self, encoded_visible, ids_keep, ids_mask):
B, n_mask = ids_mask.shape
vis_tok = self.embed(encoded_visible) + self.pe[ids_keep]
mask_tok = self.mask_token.expand(B, n_mask, -1) + self.pe[ids_mask]
full = torch.zeros(B, N_PATCHES, self.decoder_dim)
full = full.scatter(1, ids_keep.unsqueeze(-1).expand(-1, -1, self.decoder_dim), vis_tok)
full = full.scatter(1, ids_mask.unsqueeze(-1).expand(-1, -1, self.decoder_dim), mask_tok)
return self.pred(self.decoder(full)) # (B, N_PATCHES, patch_dim), full sequence, original order
class MAE(nn.Module):
def __init__(self, mask_ratio=0.5, embed_dim=32, decoder_dim=16):
super().__init__()
self.mask_ratio = mask_ratio
self.encoder = MAEEncoder(embed_dim=embed_dim)
self.decoder = MAEDecoder(embed_dim=embed_dim, decoder_dim=decoder_dim)
def forward(self, img, generator=None):
patches = patchify(img)
x_visible, ids_keep, ids_mask = random_masking(patches, self.mask_ratio, generator=generator)
encoded = self.encoder(x_visible, ids_keep)
pred = self.decoder(encoded, ids_keep, ids_mask)
D = patches.shape[-1]
pred_masked = torch.gather(pred, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))
target_masked = torch.gather(patches, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))
loss = F.mse_loss(pred_masked, target_masked)
return loss, pred, ids_mask, ids_keep, patches
def train_mae(mask_ratio, epochs=600, lr=0.001, seed=0):
torch.manual_seed(seed)
model = MAE(mask_ratio=mask_ratio)
opt = torch.optim.Adam(model.parameters(), lr=lr)
for _ in range(epochs):
opt.zero_grad()
loss, *_ = model(Xt)
loss.backward()
opt.step()
return model
How much can be reconstructed from how little?¶
Sweep the fraction of patches hidden from 25% to 75%, and compare each model's masked-patch reconstruction error against a trivial baseline: predicting the training set's mean pixel value for every masked patch, regardless of image.
train_patches = patchify(Xt)
mean_pixel = train_patches.mean()
print(f'{"mask ratio":>10} {"MAE test MSE":>14} {"mean-pixel baseline":>20}')
mae_models = {}
for mr in [0.25, 0.5, 0.75]:
model = train_mae(mask_ratio=mr)
mae_models[mr] = model
with torch.no_grad():
test_loss, pred, ids_mask, ids_keep, patches = model(Xte, generator=torch.Generator().manual_seed(123))
D = patches.shape[-1]
target_masked = torch.gather(patches, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))
baseline_mse = F.mse_loss(torch.full_like(target_masked, mean_pixel), target_masked).item()
print(f'{mr:>10} {test_loss.item():>14.4f} {baseline_mse:>20.4f}')
Reconstruction gets harder as more of the image is hidden — unsurprising, there's simply less evidence to work with — but the model beats the naive baseline by a wide margin even at 75% masking, where only 4 of the 16 patches are ever shown to the encoder. That's only possible because these shapes have enormous internal redundancy: a circle's silhouette is entirely determined by its center and radius, so a handful of visible boundary patches, combined with attention across all of them, is enough to infer the rest. Real MAE finds the same thing at a much larger scale — natural photographs are also highly spatially redundant (a patch of sky, a patch of skin, a patch of brick predicts its neighbors) — which is why He et al. found the best pretraining mask ratio for images is around 75%, dramatically higher than the ~15% word-masking ratio BERT uses for language, where far less of any given sentence is predictable from the rest.
idx = 7
model75 = mae_models[0.75]
with torch.no_grad():
_, pred, ids_mask, ids_keep, patches = model75(Xte, generator=torch.Generator().manual_seed(42))
masked_view = patches.clone()
masked_view[torch.arange(len(patches)).unsqueeze(1), ids_mask] = 0.5 # gray out hidden patches
recon_view = patches.clone()
recon_view[torch.arange(len(patches)).unsqueeze(1), ids_mask] = pred[torch.arange(len(patches)).unsqueeze(1), ids_mask]
orig_img = unpatchify(patches)[idx, 0].numpy()
masked_img = unpatchify(masked_view)[idx, 0].numpy()
recon_img = unpatchify(recon_view)[idx, 0].numpy()
fig, axes = plt.subplots(1, 3, figsize=(7, 2.8))
for ax, im, title in zip(axes, [orig_img, masked_img, recon_img],
['original', 'what the encoder saw\n(75% hidden)', 'reconstructed']):
ax.imshow(im, cmap='gray', vmin=0, vmax=1); ax.set_title(title, fontsize=9); ax.axis('off')
plt.show()
From reconstruction to representation: a linear probe¶
The reconstruction task itself is thrown away after pretraining — what real MAE keeps is the encoder, exactly as Lessons 48 and 49 kept only the trained encoder/teacher and discarded the contrastive/distillation machinery around it. Freeze the mask-ratio-0.75 encoder, feed it every patch at once (no masking, mask_ratio=0 at evaluation time — the whole point of pretraining under a harsh bottleneck was to prepare the encoder for the easy case), average-pool its output into one feature vector per image, and train a linear probe on a handful of labeled examples.
def encode_full(encoder, img):
patches = patchify(img)
B, N, _ = patches.shape
ids_all = torch.arange(N).unsqueeze(0).expand(B, -1)
with torch.no_grad():
encoded = encoder(patches, ids_all)
return encoded.mean(dim=1) # average-pool over all 16 patches -> one vector per image
def make_noisy_dataset(rng_local, n, noise=0.2, 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)
probe_rng = np.random.default_rng(11)
X_probe_train, y_probe_train = make_noisy_dataset(probe_rng, 8)
X_probe_test, y_probe_test = make_noisy_dataset(probe_rng, 150)
Xpt = torch.tensor(X_probe_train).unsqueeze(1)
Xpte = torch.tensor(X_probe_test).unsqueeze(1)
def linear_probe_acc(encoder, seed):
torch.manual_seed(seed)
feat_train = encode_full(encoder, Xpt)
feat_test = encode_full(encoder, Xpte)
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()
mae_accs, random_accs = [], []
for seed in range(5):
mae_accs.append(linear_probe_acc(mae_models[0.75].encoder, seed=seed + 50))
torch.manual_seed(seed)
random_encoder = MAEEncoder()
random_accs.append(linear_probe_acc(random_encoder, seed=seed + 50))
print(f'linear probe on RANDOM (untrained) MAE encoder: {np.mean(random_accs):.1%} (+/- {np.std(random_accs):.1%})')
print(f'linear probe on MAE-pretrained encoder: {np.mean(mae_accs):.1%} (+/- {np.std(mae_accs):.1%})')
That's a genuinely disappointing number — the MAE-pretrained encoder's frozen, average-pooled features probe worse than a random, untrained encoder's. This isn't a bug; it's a real, documented property of masked-reconstruction pretraining, and it's worth understanding why. Contrastive learning (Lesson 49) and self-distillation (Lesson 50) both directly optimize a pooled, whole-image embedding to be useful on its own — that's the entire training signal. MAE's training signal never touches a pooled representation at all: the encoder only has to produce per-patch tokens rich enough that attention within the decoder can fill in missing patches. Nothing forces the average of those tokens to be a good summary of "what shape is this" — a perfectly good reconstruction-supporting representation can average out to something a linear probe finds useless. The original MAE paper reports exactly this pattern at real scale: substantially lower linear-probe accuracy than contrastive methods, with MAE's real payoff only showing up after full fine-tuning (updating every encoder weight on the downstream task, not just a linear head on top of frozen features) — a genuinely different deployment story from Lessons 48 and 49, not a worse one.
Three self-supervised paradigms, one shared pattern¶
Contrastive learning (Lesson 49), self-distillation (Lesson 50), and masked reconstruction (this lesson) invent three completely different training signals from the same raw material — unlabeled images — yet all three exist to answer the same question: what's the cheapest information-destroying task whose solution forces a network to learn something genuinely useful about images? Contrastive learning destroys identity across augmented views and asks the network to recover it; self-distillation destroys the teacher's certainty and asks the student to match it anyway; masked reconstruction destroys most of the image outright and asks the network to fill in the rest. But this lesson's probe result is a real warning against assuming all three produce interchangeable representations: MAE's training signal is architecturally cheap (the encoder never even processes masked tokens, unlike the other two, which always see the whole image) and excellent for downstream fine-tuning, but it earns that efficiency by never once forcing a pooled, linearly-useful summary vector to exist — which is exactly what contrastive learning and self-distillation are built around. Choosing among these three in practice is a genuine engineering tradeoff, not a strict ranking.
Exercises¶
- Increase
mask_ratioto0.9(only 1-2 patches visible). Does reconstruction MSE still beat the mean-pixel baseline, or does it collapse to roughly the baseline's error — and does that reveal a point past which this dataset's redundancy runs out? - The decoder here is deliberately tiny (
n_layers=1,decoder_dim=16, smaller than the encoder). Try making the decoder larger than the encoder (e.g.decoder_dim=64,n_layers=3) while keeping the encoder fixed. Does reconstruction quality improve much — and does that match real MAE's design choice of a small decoder, on the reasoning that the decoder is discarded after pretraining anyway? - Instead of freezing the encoder for the linear probe, fine-tune it end to end: let gradients from the 8-example classification loss update the encoder's own weights too, not just a linear head on top of frozen features. Does end-to-end fine-tuning recover the gap between MAE and the random baseline — matching the real MAE paper's finding that fine-tuned MAE features are excellent, even when frozen ones probe poorly?