Autoencoders and latent spaces

This site is called Latent. Latent means hidden: present, but not directly observed. Inside every network in this track, between the input and the answer, there are hidden layers full of numbers that nobody chose and nobody labeled. This lesson makes one of those hidden spaces small enough to see, and lets you walk around in it.

Compress, then rebuild

Take the handwritten digits from Reading handwriting: 28 × 28 pixels, 784 numbers each. Now train a network to do something that sounds pointless. Its output should be its own input, the same 784 pixel values.

The catch is a bottleneck. The network narrows from 784 numbers down to just 2, then widens back to 784. Whatever the network needs to rebuild a digit has to pass through those two numbers. This is an autoencoderAutoencoderA network trained to rebuild its own input after squeezing it through a narrow middle layer. The encoder half compresses; the decoder half reconstructs; the middle layer's values are a learned code.Open in glossary: an encoder that compresses, a decoder that rebuilds, and in between a code.

Training is ordinary backpropagation. The loss is a reconstruction lossReconstruction lossHow far an autoencoder's output is from its input, such as squared error or binary cross-entropy summed over the pixels. Training an autoencoder means making it small.Open in glossary: how different the rebuilt image is from the original, pixel by pixel. No labels appear anywhere, so this is unsupervised learningUnsupervised learningLearning from examples that have no answers attached, by finding structure such as groups or directions of variation in the data itself.Taught in Finding groups: k-meansOpen in glossary. The network is never told that a 3 is a 3.

The two numbers form a coordinate system the network invented for itself. Every digit becomes a point in that plane, and the plane is a latent spaceLatent spaceThe space of learned, hidden coordinates a model uses to represent things internally, such as word vectors or a network's hidden layers. Nearby points stand for similar things.Open in glossary. Most real autoencoders use many more than two dimensions; two is the most that fits on a flat screen. In 2006, Geoffrey Hinton and Ruslan Salakhutdinov showed that a deep autoencoder could squeeze MNIST digits into two numbers that kept the digit classes far better separated than principal component analysisPrincipal component analysisA method that finds the directions along which data varies the most, often used to project many-dimensional data down to two or three dimensions for viewing.Taught in Seeing high dimensionsOpen in glossary, the classic straight-line method, could.

A map you can walk

The demo below runs two autoencoders, trained on the 60,000 MNIST training digits, in your browser. Both have the shape in the diagram. The one you see first is a variational autoencoder, explained further down. Every point on the map is a pair of numbers, so every point decodes to an image, and the map shows a grid of them.

What the map shows

The encoder grouped the digits without being told what they are. On average, about 6 of each test digit’s 10 nearest neighbors on the map are the same digit. The reason is the decoder. It is a smooth function of the two numbers, so nearby codes decode to similar images. To rebuild a thousand slightly different 0s well, the encoder has to give them nearby codes, and to keep 0s apart from 1s it has to put them in different places.

Two numbers are not much room, and the map shows the strain. Some digits get large, clean regions: 0s, 1s, and 7s each hold more than 15% of the map. Others crowd the middle, where 4s, 8s, and 9s overlap and each hold less than 2%. Between regions, the decoder draws in-between shapes, part one digit and part another. Nothing in the network represents “digit” as a category; there is only a smooth function from two numbers to 784 pixels.

Why “variational”

The demo has a second model, the Plain autoencoder. It has the same layers and was trained on the same data, but its encoder outputs exactly one point per image and its loss is reconstruction alone. Switch to it and look at its map.

Nothing told the plain autoencoder where to put its codes, and it shows. Its codes run over tens of units with no natural center. It stretches 1s across more than 40% of its map, while 3s, 4s, 5s, and 9s each get well under 1%. If you wanted a new digit, you would not know where to look. The plain autoencoder’s decoder still draws something in its empty stretches, often a plausible digit, because it is a smooth function too, but nothing in training checked those places.

A variational autoencoderVariational autoencoderAn autoencoder whose encoder outputs a probability distribution instead of a single code, trained with a KL term that pulls codes toward a standard normal prior, so any point near the center decodes to something plausible.Open in glossary (VAE), introduced by Diederik Kingma and Max Welling in 2013, and independently by Danilo Rezende, Shakir Mohamed, and Daan Wierstra, changes two things.

1. A cloud instead of a point. For each image, the encoder outputs a mean and a spread for each of the two numbers. During training the decoder rebuilds the image from a random point in that cloud, not from its center. So the decoder must draw a sensible digit for a whole neighborhood of points, not just one.
2. A pull toward the center. The loss adds a second term, the KL divergenceKL divergenceA measure of how much one probability distribution differs from another, in bits or nats. It is zero only when the two are identical. In RLHF a KL penalty keeps the trained model close to where it started.Taught in Quantization: smaller numbersOpen in glossary from each image’s cloud to a standard normal distribution, the bell curve centered on zero with spread 1. That bell curve is the model’s priorPriorWhat you believe about a hypothesis before seeing the evidence, as a probability. Bayes' rule combines it with the likelihood to give the posterior.Taught in Probability and Bayes' ruleOpen in glossary: the distribution codes should follow before you look at any particular image. The KL term grows when a cloud drifts away from zero or its spread moves away from 1.

Together they make the codes fill one compact disk around zero, with neighboring clouds overlapping. Now generating is easy: draw a point from the bell curve and decode it. That is what Sample a point from the prior does.

This order has a price. On the 10,000 test digits, the plain autoencoder’s reconstruction loss averages 134.5 nats per digit and the VAE’s 137.7 (plus a KL term of 6.5). The VAE rebuilds digits slightly worse in exchange for a map you can sample from.

import numpy as np

rng = np.random.default_rng(0)

# What an encoder outputs for one image: a mean and a log-variance per latent axis.
mu = np.array([0.8, -1.2])
logvar = np.array([-2.0, -1.5])
sigma = np.exp(0.5 * logvar)

# Reparameterization trick: sample noise separately, then shift and scale it.
# The randomness sits outside the network, so gradients reach mu and sigma.
eps = rng.standard_normal(2)
z = mu + sigma * eps
print("z =", np.round(z, 3))

# KL term: how far N(mu, sigma^2) is from the prior N(0, 1), in nats.
kl = 0.5 * np.sum(mu**2 + sigma**2 - 1 - logvar)
print("KL =", round(kl, 3))

# Reconstruction term: binary cross-entropy between the image and the decoder's output.
x = np.array([0.0, 1.0, 0.9, 0.0])       # four pixel intensities
x_hat = np.array([0.1, 0.8, 0.7, 0.2])   # what the decoder drew
bce = -np.sum(x * np.log(x_hat) + (1 - x) * np.log(1 - x_hat))
print("reconstruction =", round(bce, 3))
print("loss =", round(bce + kl, 3))

z = [ 0.846 -1.262]
KL = 1.969
reconstruction = 0.993
loss = 2.962

Morphing and arithmetic

Pick two digits, encode both, and walk the straight line between their codes, decoding as you go. The result is a morph: one digit gradually becomes the other. You can also draw your own digit, see where the encoder puts it, and use it as one end.

Straight paths often cross a third digit’s region. Of the 90 start and end pairs in the demo, 71 pass through another digit’s territory on the way. In a flat map with ten digits, most pairs cannot be neighbors.

Arithmetic is the morph’s cousin: instead of walking toward another digit, you add a direction. The demo’s directions come from averages. “More slant” for 0s is the average code of the most slanted quarter of test 0s minus the average code of the least slanted quarter. Each direction was then checked on test digits it was not computed from, using the digit recognizer from Reading handwriting to judge whether the result is still the same digit.

The results are mixed, and the demo shows them. Adding the slant direction made 0s and 1s more slanted while keeping them 0s and 1s about 95% of the time. For 5s and 9s it failed: only 3% of 5s and 12% of 9s came out both more slanted and still the same digit, and most turned into other digits. One shared “thicker” direction, averaged over all digits, thickened the strokes 9 times in 10 but kept the same digit only about 1 time in 6. With just two numbers, a digit’s identity and its style share the same coordinates, so moving in a “style” direction usually also moves you toward a different digit. Word embeddings have hundreds of dimensions, which leaves room for a direction like “plural” or “capital city” to mean roughly the same thing everywhere, though analogies there are approximate too.

Latent spaces everywhere

The same idea, a learned space where nearby points mean similar things, runs through the rest of this course.

• Word embeddings are latent spaces learned without a decoder: word2vec trains a small network to predict the words around each word, and the vectors it learns place related words near each other. You will meet them in Word embeddings.
• Image generators use autoencoders directly. Latent diffusion (Rombach and colleagues, 2022), the method behind Stable Diffusion, first trains an autoencoder with a slight KL penalty toward a standard normal, much like a VAE. Stable Diffusion’s version turns a 512 × 512 color image, 786,432 numbers, into a 64 × 64 × 4 latent, 16,384 numbers, 48 times fewer. Diffusion then runs in that latent space, and the decoder turns the result back into pixels.
• Interpretability research uses a different kind, sparse autoencoders, to pull apart the features a language model packs into its activations.