Autoencoders

Every model so far learned to map inputs to labels. Autoencoders drop the labels entirely: the network's target is its own input. That sounds useless — copy x to x — until you add one constraint that changes everything: the copy must pass through a narrow bottleneck. In this lesson you'll see why that constraint forces the network to discover structure, build a full MNIST autoencoder, and use the same trick for denoising and anomaly detection.

Learning to compress

An autoencoder has three parts:

• Encoder — squeezes the input down, e.g. a 784-pixel MNIST image through layers of 256 → 64 → 32 numbers.
• Bottleneck (latent space) — the 32-number summary, also called the latent code z.
• Decoder — a mirror image that inflates 32 numbers back to 784 pixels.

Training minimizes reconstruction loss — how far the decoder's output is from the original input, typically mean squared error over pixels. No labels anywhere: the data supervises itself, which is why this is called self-supervised (or classically, unsupervised) learning.

The bottleneck is the whole point. If the latent space were as wide as the input, the network could learn the identity function and reconstruct perfectly while learning nothing. Forced through 32 numbers, it cannot memorize pixels — it has to find the 32 most useful facts about a digit image (roughly: which digit, how slanted, how thick the stroke, ...) and learn to redraw from them. Compression pressure is what turns copying into understanding.

The linear ancestor: PCA

You've met this idea before. PCA is exactly a linear autoencoder: project onto the top k principal components (encode), then project back (decode) — and among all linear maps, PCA's reconstruction error is optimal. An autoencoder with no activation functions and MSE loss learns the same subspace as PCA; the deep, nonlinear version is what earns its keep on complex data.

Watch reconstruction quality degrade as we shrink PCA's "bottleneck" on the digits dataset — this runs in your browser:

With 32 components the digits are nearly perfect; at 8 they're blurry but recognizable; at 2 almost everything is lost. A deep autoencoder plays the same game, but its nonlinear encoder/decoder can pack far more structure into the same number of latent dimensions.

A full MNIST autoencoder in PyTorch

Now the real thing. PyTorch doesn't run in the browser, so these cells belong in the downloadable notebook — Colab with a GPU trains this in a couple of minutes. Encoder and decoder are plain nn.Sequential stacks, mirror images of each other; the final Sigmoid keeps outputs in the 0–1 pixel range:

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
 
device = "cuda" if torch.cuda.is_available() else "cpu"
 
train_set = datasets.MNIST("data", train=True, download=True,
                           transform=transforms.ToTensor())
test_set = datasets.MNIST("data", train=False, download=True,
                          transform=transforms.ToTensor())
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=128)
 
class Autoencoder(nn.Module):
    def __init__(self, latent_dim=32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),                      # (batch, 1, 28, 28) -> (batch, 784)
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent_dim),         # the bottleneck: (batch, 32)
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid(), # back to 0-1 pixels
        )
 
    def forward(self, x):
        z = self.encoder(x)
        recon = self.decoder(z).view(-1, 1, 28, 28)
        return recon, z

The training loop is the standard one with a twist: the loss compares the reconstruction to the input, and the labels are thrown away:

model = Autoencoder().to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
for epoch in range(10):
    model.train()
    total = 0.0
    for images, _ in train_loader:        # labels ignored!
        images = images.to(device)
        recon, z = model(images)
        loss = criterion(recon, images)   # reconstruction vs original
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total += loss.item() * len(images)
    print(f"epoch {epoch + 1}: reconstruction MSE {total / len(train_loader.dataset):.4f}")

And the payoff — originals on top, reconstructions from 32 numbers below:

import matplotlib.pyplot as plt
 
model.eval()
images, _ = next(iter(test_loader))
with torch.no_grad():
    recon, z = model(images.to(device))
 
fig, ax = plt.subplots(2, 8, figsize=(14, 4))
for i in range(8):
    ax[0, i].imshow(images[i].squeeze(), cmap="gray");        ax[0, i].axis("off")
    ax[1, i].imshow(recon[i].squeeze().cpu(), cmap="gray");   ax[1, i].axis("off")
ax[0, 0].set_title("original", loc="left")
ax[1, 0].set_title("reconstructed from 32 numbers", loc="left")
plt.show()

A 784-pixel image squeezed through 32 numbers and redrawn — a 24× compression learned purely from the data.

Application 1: denoising

Here's a beautiful variation: corrupt the input, but keep the clean image as the target. The network can no longer succeed by copying — it must learn what digits should look like in order to repair them. This is the denoising autoencoder, and it takes a two-line change:

def add_noise(images, noise_factor=0.5):
    noisy = images + noise_factor * torch.rand_like(images)
    return noisy.clamp(0.0, 1.0)
 
# inside the training loop, replace the forward pass with:
#     recon, z = model(add_noise(images))
#     loss = criterion(recon, images)     # target is still the CLEAN image

Retrain, then feed it noisy test images: the reconstructions come out clean. The same recipe — corrupted input, clean target — powers document cleanup (removing coffee stains and shadows from scanned pages), audio denoising, and image inpainting where whole patches are masked out and repainted.

Application 2: anomaly detection

An autoencoder only learns to reconstruct the kind of data it was trained on. Show it something from a different distribution and the reconstruction fails — the error jumps. That gives a simple anomaly detector: score every sample by reconstruction error and flag the outliers.

Since a linear autoencoder is PCA, we can demonstrate the whole idea in the browser: fit PCA on real digits, inject one fake "image" of pure random noise, and watch its reconstruction error stand out:

The impostor's error is several times larger than any real digit's — it ranks first out of 1,798. In production this pattern detects fraudulent transactions, failing machines, and network intrusions: train on normal data only, then alert on whatever the model can't redraw.

The latent space, and a teaser

The bottleneck isn't just small — it's organized. Encode all of MNIST and look at the 32-dimensional codes: images of the same digit cluster together, similar handwriting styles sit near each other, and walking in a straight line between the code for a 3 and the code for an 8 decodes into images that morph smoothly from one into the other. The encoder has arranged concepts geometrically, without ever seeing a label.

Next up: the final lesson — two networks locked in a forgery contest, better known as generative adversarial networks.