Lesson 48: Autoencoders¶
Every network up to this point learned from labeled examples: an image and its correct class, box, or mask. Autoencoders need none of that — they're a foundational example of unsupervised learning, where the only training signal comes from the data itself, never a human-provided label. Train a network to reconstruct its own input, through a deliberately narrow bottleneck, and the bottleneck is forced to discover a compressed, meaningful representation of the data — not because anyone labeled anything, but because reconstruction through a narrow-enough pipe is only solvable by throwing away what a shape looks like pixel-for-pixel and keeping what a shape is. This lesson builds one from scratch (Hinton & Salakhutdinov, 2006), compares it against the classical linear alternative (PCA), and shows what the bottleneck actually buys you: denoising, and a smoothly interpolable latent space.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
A shape dataset, and an encoder-decoder¶
Reuse the plus/circle shapes from earlier lessons. The encoder compresses a 16x16 image down to a small bottleneck vector; the decoder expands that vector back up to a full 16x16 reconstruction. Training minimizes reconstruction error — the input is also the target, so no labels are used anywhere in this lesson.
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)
rng = np.random.default_rng(6)
X_train, y_train = make_dataset(rng, 400)
X_test, y_test = make_dataset(rng, 150)
class Encoder(nn.Module):
def __init__(self, bottleneck=8):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 16 -> 8
nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 8 -> 4
)
self.fc = nn.Linear(16 * 4 * 4, bottleneck)
def forward(self, x):
return self.fc(self.conv(x).flatten(1))
class Decoder(nn.Module):
def __init__(self, bottleneck=8):
super().__init__()
self.fc = nn.Linear(bottleneck, 16 * 4 * 4)
self.up = nn.Upsample(scale_factor=2, mode='nearest')
self.conv1 = nn.Sequential(nn.Conv2d(16, 8, 3, padding=1), nn.ReLU()) # 4 -> 8
self.conv2 = nn.Conv2d(8, 1, 3, padding=1) # 8 -> 16
def forward(self, z):
h = self.fc(z).view(-1, 16, 4, 4)
h = self.conv1(self.up(h))
return torch.sigmoid(self.conv2(self.up(h)))
class Autoencoder(nn.Module):
def __init__(self, bottleneck=8):
super().__init__()
self.encoder = Encoder(bottleneck)
self.decoder = Decoder(bottleneck)
def forward(self, x):
return self.decoder(self.encoder(x))
def train_autoencoder(bottleneck, epochs=300, lr=0.01, seed=0):
torch.manual_seed(seed)
model = Autoencoder(bottleneck)
opt = torch.optim.Adam(model.parameters(), lr=lr)
Xt = torch.tensor(X_train).unsqueeze(1)
for _ in range(epochs):
opt.zero_grad()
loss = F.mse_loss(model(Xt), Xt)
loss.backward()
opt.step()
return model
model = train_autoencoder(bottleneck=8)
Xte = torch.tensor(X_test).unsqueeze(1)
with torch.no_grad():
recon_te = model(Xte)
test_mse = F.mse_loss(recon_te, Xte).item()
print(f'test reconstruction MSE (bottleneck=8): {test_mse:.4f}')
fig, axes = plt.subplots(2, 6, figsize=(11, 4))
for i in range(6):
axes[0, i].imshow(X_test[i], cmap='gray'); axes[0, i].axis('off')
axes[1, i].imshow(recon_te[i, 0].numpy(), cmap='gray'); axes[1, i].axis('off')
axes[0, 0].set_title('original', fontsize=9, loc='left')
axes[1, 0].set_title('reconstructed', fontsize=9, loc='left')
plt.show()
How much does the bottleneck matter?¶
An 8-number bottleneck for a 256-pixel image is a 32x compression. Sweep the bottleneck width and watch reconstruction quality trade off against it directly.
bottlenecks = [1, 2, 4, 8, 16, 32]
ae_mses = []
for b in bottlenecks:
m = train_autoencoder(bottleneck=b, seed=0)
with torch.no_grad():
mse = F.mse_loss(m(Xte), Xte).item()
ae_mses.append(mse)
print(f'bottleneck={b:>3}: test MSE = {mse:.4f}')
plt.figure(figsize=(5, 3.5))
plt.plot(bottlenecks, ae_mses, '-o')
plt.xscale('log', base=2)
plt.xlabel('bottleneck size'); plt.ylabel('test reconstruction MSE')
plt.title('Reconstruction quality vs. bottleneck width')
plt.show()
Nonlinear vs. linear compression: autoencoder vs. PCA¶
PCA is the classical version of this same idea: project onto the top-k directions of variance, then project back — but only ever through a linear map. Compare reconstruction quality at the same bottleneck size.
def pca_reconstruct(X_train_flat, X_test_flat, k):
mean = X_train_flat.mean(axis=0, keepdims=True)
_, _, Vt = np.linalg.svd(X_train_flat - mean, full_matrices=False)
components = Vt[:k] # (k, 256)
projected = (X_test_flat - mean) @ components.T
return projected @ components + mean
X_train_flat = X_train.reshape(len(X_train), -1)
X_test_flat = X_test.reshape(len(X_test), -1)
pca_mses = []
for k in bottlenecks:
recon_pca = pca_reconstruct(X_train_flat, X_test_flat, k)
pca_mses.append(np.mean((recon_pca - X_test_flat) ** 2))
print(f'{"bottleneck":>10} {"autoencoder MSE":>16} {"PCA MSE":>10}')
for b, ae_mse, pca_mse in zip(bottlenecks, ae_mses, pca_mses):
print(f'{b:>10} {ae_mse:>16.4f} {pca_mse:>10.4f}')
At every small bottleneck (1, 2, 4, 8) the nonlinear autoencoder reconstructs noticeably better than PCA at the same dimensionality — a linear projection genuinely cannot capture "shape identity plus a 2D position" as efficiently as a network that can bend its compression around the data's actual, curved structure. But look at what happens by bottleneck 32: the autoencoder's error has already plateaued (it stopped improving back at bottleneck 4), while PCA's keeps falling as it's handed more linear directions to spend — and by 32 dimensions PCA has actually overtaken it. This isn't a flaw in the nonlinear idea; it's a reminder that "nonlinear beats linear" only holds where nonlinearity is the binding constraint. Once the bottleneck is wide enough that capacity isn't the bottleneck anymore, the autoencoder's remaining error comes from training — optimizing a nonlinear network is a genuinely harder problem than PCA's exact, closed-form solution — not from any fundamental ceiling on what a network could represent.
Denoising: reconstructing the clean signal from a corrupted input¶
A denoising autoencoder is trained with a twist: feed it a corrupted input, but still ask it to reconstruct the clean original. Doing this well forces the bottleneck to encode signal, not noise, since noise itself has nothing predictable to preserve.
def add_noise(imgs, rng_local, level=0.3):
noisy = imgs + rng_local.normal(0, level, imgs.shape)
return np.clip(noisy, 0, 1).astype(np.float32)
noise_rng = np.random.default_rng(9)
X_train_noisy = add_noise(X_train, noise_rng)
X_test_noisy = add_noise(X_test, noise_rng)
def train_denoising_autoencoder(bottleneck=8, epochs=300, lr=0.01, seed=0):
torch.manual_seed(seed)
model = Autoencoder(bottleneck)
opt = torch.optim.Adam(model.parameters(), lr=lr)
Xt_noisy = torch.tensor(X_train_noisy).unsqueeze(1)
Xt_clean = torch.tensor(X_train).unsqueeze(1)
for _ in range(epochs):
opt.zero_grad()
loss = F.mse_loss(model(Xt_noisy), Xt_clean)
loss.backward()
opt.step()
return model
denoiser = train_denoising_autoencoder()
Xte_noisy = torch.tensor(X_test_noisy).unsqueeze(1)
with torch.no_grad():
denoised = denoiser(Xte_noisy)
noisy_vs_clean_mse = np.mean((X_test_noisy - X_test) ** 2)
denoised_vs_clean_mse = F.mse_loss(denoised, Xte).item()
print(f'noisy input vs. clean target MSE: {noisy_vs_clean_mse:.4f} (doing nothing)')
print(f'denoiser output vs. clean target MSE: {denoised_vs_clean_mse:.4f}')
fig, axes = plt.subplots(3, 6, figsize=(11, 5.5))
for i in range(6):
axes[0, i].imshow(X_test[i], cmap='gray'); axes[0, i].axis('off')
axes[1, i].imshow(X_test_noisy[i], cmap='gray'); axes[1, i].axis('off')
axes[2, i].imshow(denoised[i, 0].numpy(), cmap='gray'); axes[2, i].axis('off')
for r, name in enumerate(['clean', 'noisy input', 'denoised']):
axes[r, 0].set_title(name, fontsize=9, loc='left')
plt.show()
A structured latent space: interpolation¶
If the bottleneck genuinely encodes shape identity and position rather than memorizing individual pixels, decoding a point between two encoded images should produce something in between, not noise.
idx_a = 0
idx_b = next(i for i in range(len(X_test)) if y_test[i] != y_test[idx_a])
with torch.no_grad():
z_a = model.encoder(Xte[idx_a:idx_a+1])
z_b = model.encoder(Xte[idx_b:idx_b+1])
alphas = np.linspace(0, 1, 6)
interp_imgs = [model.decoder((1 - a) * z_a + a * z_b)[0, 0].numpy() for a in alphas]
shape_names = ['plus', 'circle']
fig, axes = plt.subplots(1, 6, figsize=(11, 2))
for ax, im, a in zip(axes, interp_imgs, alphas):
ax.imshow(im, cmap='gray'); ax.axis('off')
ax.set_title(f'{a:.1f}', fontsize=8)
plt.suptitle(f'Interpolating in latent space: {shape_names[y_test[idx_a]]} -> {shape_names[y_test[idx_b]]}')
plt.show()
Where this goes next¶
None of this lesson used a single label — every result came from the reconstruction objective alone. That pattern, a network forced to compress and then rebuild its own input, is the ancestor of a surprising amount of what the rest of this part does with unlabeled data: Lesson 49's contrastive learning and Lesson 50's self-distillation both replace "reconstruct the pixels" with other self-supervised objectives, but keep the same "no labels, invent the training signal from the data itself" structure. Lesson 51 builds a masked autoencoder — the same encoder-decoder shape as this lesson, but with randomly hidden patches as the information bottleneck instead of a narrow vector. And the denoising objective built here, applied one small noise-step at a time and chained together, is exactly the mechanism Lesson 60's diffusion models scale up to generate entirely new images from pure noise.
Exercises¶
- Try
bottleneck=64andbottleneck=128in the sweep above. Does reconstruction MSE keep improving all the way, or does it plateau — and what does a plateau suggest about how many degrees of freedom this dataset actually has (shape identity plus a 2D position)? - The denoising autoencoder above was trained and evaluated at
level=0.3. Retrain and re-evaluate atlevel=0.6. Does the denoiser still recover something recognizable, or does reconstruction collapse into an ambiguous blur? - Retrain an autoencoder with
bottleneck=2, encode all 150 test images, and scatter-plot the resulting 2D codes colored byy_test. Do plus and circle images separate into two visible clusters, despite the network never once seeing a label during training?