Autoencoders: Learning by Squeezing

An autoencoder is a neural network given a strange job: reproduce its input. The catch is a narrow middle layer, the bottleneck, that forces it to squeeze each input into a few numbers and rebuild it from those alone. To succeed it has to learn what matters about the data and what can be thrown away. Whatever survives the squeeze is a compact summary the network discovered by itself. Here you’ll train one in your browser on handwritten digits, make new digits from its summaries, and use it to spot the odd ones out.

Learning by squeezing#

An autoencoder has two halves. The encoder takes an image, here 64 pixels of a handwritten digit, and turns it into a short list of numbers called its code. The decoder takes the code and tries to draw the original image back. Training compares the drawing with the original, pixel by pixel, and adjusts every weight in both halves to shrink the difference, using the same gradient descent as any other network.

No labels are needed: the input is its own answer. That makes autoencoders a kind of unsupervised (or self-supervised) learning. Nobody tells the network that it’s looking at a 3 or a 7. It just has to find codes that work.

With 2 numbers the held-back digits come back blurry but mostly readable. Give it 4, 8 or 16 numbers and they sharpen. The comparison row uses PCA, which finds the best straight-line summary with the same number of values. At 2, 4 and 8 numbers the autoencoder’s error is a fifth to a quarter lower than PCA’s, because its curved, nonlinear layers can fold more information into each number. By 16 numbers there’s little left to gain, and the two are about level.

That’s no coincidence. An autoencoder with no curves in it at all, just straight-line layers, ends up rebuilding images exactly as PCA does. The nonlinear layers are what let it beat PCA when the squeeze is tight. In 2006 Geoffrey Hinton and Ruslan Salakhutdinov showed this at scale: a deep autoencoder compressed images of faces and handwritten digits far better than PCA could.

Inside the bottleneck#

With a 2-number bottleneck, every digit’s code is a point you can plot. Pick “2 numbers” in the figure and look at the map: the network has sorted the digits into regions, 0s with 0s and 6s with 6s, though it never saw a label. Digits that are easy to confuse, such as 5s, 8s and 9s, overlap, and the round 0s sit next to the round 6s. With only two numbers, about four digits in five have a nearest neighbour of the same kind.

Now tap the map. The decoder will draw an image from any point, including points where no real digit sits. Tap inside a group and you get a convincing digit nobody wrote. Tap between two groups and you get a cross between them. Tap far outside, and you get a smudge. The decoder is a generator, but only where the codes are dense.

Spotting the odd one out#

An autoencoder learns to rebuild the kind of data it was trained on. Show it something different and it rebuilds it badly, so a large reconstruction error is a sign that an input is unusual. That makes autoencoders a popular tool for anomaly detection: train on normal machine readings, normal transactions or normal product photos, then flag whatever the network can’t rebuild.

At the default threshold of 0.05, about one ordinary digit in a hundred is a false alarm. Every patch of random noise is flagged, and so is every digit moved two pixels to the right. Only about half the upside-down digits are caught, and the upside-down 6, a 9 to the network, never is. Ink blots are hit and miss: a blot on a stroke that was already dark changes little.

Notice what that means. The off-centre digits are perfectly readable, just not where the training digits sat, and they’re all flagged. Some upside-down digits really are wrong, and they pass. An autoencoder flags what’s unlike its training data, which is not the same as what’s wrong. Make sure your training data covers every kind of normal, and test the detector on the specific faults you care about.

Where autoencoders are used#

• Anomaly detection in factories, networks and payments, where faults are rare and varied but normal data is plentiful.
• Denoising: train the network to rebuild clean images from deliberately corrupted ones, and it learns to remove noise.
• Compression and features: the codes are compact summaries that other models can use, a learned cousin of PCA.
• Generation: VAEs, and the autoencoders inside latent diffusion models like Stable Diffusion.
• Looking inside language models: researchers, including teams at Anthropic, train sparse autoencoders on a model’s internal activations to pull out human-readable features.

In code#

The figure’s network in PyTorch (pip install torch scikit-learn), on all 1,797 of scikit-learn’s digits, with PCA for comparison:

import torch
from sklearn.datasets import load_digits
from sklearn.decomposition import PCA
from torch import nn

torch.manual_seed(0)
X = torch.tensor(load_digits().data / 16, dtype=torch.float32)   # 1,797 digits, 64 pixels each
train, test = X[:1500], X[1500:]

model = nn.Sequential(
    nn.Linear(64, 32), nn.Tanh(), nn.Linear(32, 2),                  # encoder: 64 pixels -> 2 numbers
    nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 64), nn.Sigmoid(),    # decoder: 2 numbers -> 64 pixels
)
opt = torch.optim.Adam(model.parameters(), lr=0.005)
for step in range(3000):
    batch = train[torch.randint(len(train), (32,))]
    loss = ((model(batch) - batch) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()

with torch.no_grad():
    ae_error = ((model(test) - test) ** 2).mean().item()
pca = PCA(n_components=2).fit(train.numpy())
pca_error = ((pca.inverse_transform(pca.transform(test.numpy())) - test.numpy()) ** 2).mean()
print(f"autoencoder: {ae_error:.4f} per pixel   PCA: {pca_error:.4f} per pixel")
# autoencoder: 0.0396 per pixel   PCA: 0.0526 per pixel

Your exact numbers may differ a little between machines, but the autoencoder should beat PCA by about a quarter, as it does in the figure. model[:3] is the encoder on its own: run it on any digit to get its 2-number code. For anomaly detection, compute each new image’s error and flag the ones above a threshold chosen on normal data.