Autoencoders and VAEs

You phone a friend and describe a picture so they can draw it, but you are allowed to say only two numbers. You would agree on a system first: the first number says which way the line runs, the second says how far it sits from the middle. You squeeze the picture into two numbers, send them, and your friend rebuilds it.

An autoencoder is a neural network that invents that system by itself. It has two halves. The encoder squeezes a picture into a few numbers, called the code. The decoder rebuilds the picture from the code. Nobody tells it what the numbers should mean. The only instruction is "the rebuild must match the original", and to obey it through such a narrow gap the network has to discover what really varies in the data.

Then comes a second question. If your friend can draw from any two numbers, can you invent a new picture by making up two numbers? With a plain autoencoder, usually not: most made-up codes were never used for anything, and the drawing comes out as a smudge. A variational autoencoder (VAE) is trained so that a made-up code drawn from a known range draws something sensible. That step, from compressing to generating, is what this lesson is about, and it opens this part of the primer: GANs (primer.ml.generative.gans) and diffusion (primer.ml.generative.diffusion) are other answers to the same question.

Tiny example Every picture in this lesson is 8 × 8 = 64 pixels holding one soft pen stroke. The stroke runs in one of four directions (horizontal, vertical, diagonal down, diagonal up) and is shifted up to 2.5 pixels from the centre. The ink fades like a bell curve either side of the line: a pixel on the line has brightness 1.0, a pixel one pixel away about 0.25, a pixel two away almost nothing.

So each picture is 64 numbers, but only two facts change from picture to picture: the direction (one of four) and the offset (any amount). A good 2-number code has to rediscover both from the pixels alone. The 200 pictures are the top row of the first figure below.

In code: make_strokes draws the pictures and returns their hidden directions and offsets; STROKE_KINDS names the four directions.

Everyday picture A zip file squeezes a document and gets it back exactly. An autoencoder is a lossy, learned zip: it keeps what matters most for this kind of data and lets the rest go.

Tiny example: a one-number code you can check by hand. Take points that lie on the line y = 2x, such as (1, 2) and (2, 4). Each takes two numbers to write down, but one number is enough: how far along the line the point is.

• Encoder: code = u · x, where u = (1, 2)/√5 = (0.447, 0.894) is the line's direction scaled to length 1, and "·" is the dot product (multiply matching entries and add; see primer.notation).
• Decoder: rebuild = code × u.

For (1, 2): code = 0.447 + 1.789 = 2.236 (which is √5, its distance from the origin), and rebuild = 2.236 × (0.447, 0.894) = (1, 2). Nothing lost.

For (2, 3), which is off the line: code = (2 + 6)/√5 = 3.578, and rebuild = 3.578 × (0.447, 0.894) = (1.6, 3.2), the closest point on the line. The part that pointed away from the line, (0.4, −0.2), is gone. That loss is exactly what training measures and shrinks:

import math
# u: the line's direction, scaled to length 1
u = [1 / math.sqrt(5), 2 / math.sqrt(5)]
def f(x):
    return sum(u_m * x_m for u_m, x_m in zip(u, x))
def g(z):
    return [z * u_m for u_m in u]
def error(x):
    return sum((a - b) ** 2 for a, b in zip(x, g(f(x))))
# f(x): the encoder squeezes (2, 3) to one number
round(f([2, 3]), 3)  # → 3.578
# g(f(x)): the decoder rebuilds two numbers from it
[round(v, 2) for v in g(f([2, 3]))]  # → [1.6, 3.2]
# ‖x − x̂‖² for each point: (1, 2) is on the line, (2, 3) is not
[round(error(x), 2) for x in ([1, 2], [2, 3])]  # → [0.0, 0.2]
# L_rec: the average over n = 2 points
round((error([1, 2]) + error([2, 3])) / 2, 2)  # → 0.1

The real network has the same two halves, only bent. The encoder is a two-layer network like the one built in primer.ml.neural_net: 64 pixels into 32 hidden numbers (through tanh, which lets it bend), then out to a code of 2 numbers. The decoder mirrors it: 2 numbers into 32 hidden, then out to 64 pixels squashed between 0 and 1 by a sigmoid so they are valid brightnesses. Training is the loop from primer.ml.neural_net with the Adam optimizer from primer.ml.optimizers: rebuild every picture, measure , send the gradient back through both halves, adjust, 1,500 times.

flowchart LR
  X["picture x<br/>64 pixels"] --> E["encoder f<br/>64 → 32 → 2"]
  E --> Z["code z<br/>2 numbers"]
  Z --> D["decoder g<br/>2 → 32 → 64"]
  D --> XH["rebuild x̂<br/>64 pixels"]
  X --> L["loss<br/>squared error between x and x̂"]
  XH --> L
  L -. "gradients flow back<br/>through decoder, then encoder" .-> E

In code: worked_example_line is the hand example above; Autoencoder holds both halves and its hand-written backward pass in Autoencoder.loss_and_gradients; train runs Adam; trained_autoencoder is the lesson's trained network; reconstruction_error measures the rebuilds.

PCA is the straight-line special case

Everyday picture PCA (principal component analysis, built in primer.ml.embeddings.clustering) fits a flat sheet through the data and records where each point lands on the sheet. That is an autoencoder with no bends: the encoder is one matrix multiply, and so is the decoder.

Tiny example Scatter 100 points near the line y = 2x and train exactly that no-bend autoencoder, with a 1-number code, by plain gradient descent. The direction it learns is (0.447, 0.894), the line's own direction, and its rebuild error matches PCA's with one component. Baldi and Hornik proved in 1989 that this always happens: a linear autoencoder trained on squared error lands on the same flat sheet as PCA.

So why did PCA smudge the strokes? Because a stroke sliding across the image does not travel in a straight line through pixel space. Take a horizontal stroke at the top and another at the bottom. Their average, the point halfway along the straight line between them, is two faint lines, not one stroke in the middle. The real strokes lie on a curved surface in the 64-dimensional space of pictures, and a flat sheet can only cut through it. The autoencoder's tanh layers let its "sheet" bend to follow the curve.

Why it matters in practice. Before generation, autoencoders earned their keep as learned compressors. They compress, they denoise (train on noisy inputs, ask for clean outputs), and they detect anomalies: an input the network rebuilds badly is unlike anything it trained on, which is a standard way to flag fraud or a failing machine.

In code: train_linear_autoencoder is the no-bend autoencoder; pca_reconstruction rebuilds from the top principal components.

Everyday picture Imagine a town where houses were built only along four winding roads. Pick a random spot inside the town limits and you will most likely land in a field. Ask the decoder to draw the picture that "lives" in that field, and it answers anyway, with whatever its weights happen to produce, because nothing in training ever asked it about that spot.

Tiny example The obvious way to make up a random code is to draw each number from the standard normal distribution, the bell curve centred on 0 with a spread (standard deviation) of 1: most draws fall between −1 and 1, almost all between −3 and 3. Written for several numbers at once, it means "draw each number from that bell curve, independently".

Draw the code z = (1.2, −0.9) and decode it with the plain autoencoder. The result is a bright smear with about three times the ink of any real stroke; its squared distance to the nearest real stroke is about 7.5. Do that 500 times and the median distance is about 3, while three quarters of the draws land further than 1.0 from every real stroke. That 1.0 is this lesson's line for junk: about a ninth of a whole stroke's squared ink.

To be fair to the autoencoder, draw instead from the box that holds its own codes, from −24 to 1.6 across and −17 to 14 up. Still about 44% of those codes decode to junk. They fall in the holes between the strands.

Why it matters in practice. A compressor is not a generator. To sample new data you need a code space with a known shape that you can draw from, no holes inside that shape, and smoothness, so nearby codes decode to similar pictures. A plain autoencoder promises none of these, because its loss only ever looks at codes of real pictures.

In code: generate decodes codes drawn from the standard normal; codes_in_box draws from a model's own code range; distance_to_data measures how far each image is from the nearest real stroke, and JUNK_DISTANCE is the line between plausible and junk.

Everyday picture Instead of pinning each picture to one exact spot on the map, the encoder draws a small fuzzy circle: "somewhere around here". During training the decoder is handed a random point from inside that circle, so it must draw the right picture from anywhere nearby. A whole neighbourhood now decodes sensibly, not just one pin. A second rule stops the encoder from cheating: every circle pays rent, more the further it sits from the middle of the map and more if it shrinks towards a pin. Circles crowd towards the centre and overlap, and the fields between the roads fill in.

Tiny example One picture, one code number. The encoder says μ = 0.5 and σ = 0.5: "about 0.5, give or take 0.5". This step, a bell-curve roll gives ε = 1.2, so the decoder is handed z = 0.5 + 0.5 × 1.2 = 1.1. Next step the roll is ε = −0.4 and the decoder gets z = 0.5 − 0.2 = 0.3. Both must rebuild the same picture.

flowchart LR
  X["picture x"] --> E["encoder"]
  E --> MU["μ: centre of the region"]
  E --> LV["log σ²: size of the region"]
  EPS["ε drawn from N(0, I)"] --> Z["z = μ + σ·ε"]
  MU --> Z
  LV --> Z
  Z --> D["decoder"] --> XH["rebuild x̂"]
  XH --> R["rebuild error<br/>‖x − x̂‖²"]
  X --> R
  MU --> KL["KL rent<br/>pulls regions to N(0, I)"]
  LV --> KL
  R --> LOSS["loss = rebuild + β · KL"]
  KL --> LOSS

In code: VAE is the Autoencoder with a two-headed encoder; VAE.encode_distribution returns μ and log σ², and VAE.encode returns just the centre μ, the best single code for a picture.

The reparameterization trick: moving the dice outside

Everyday picture To pick a random seat in a row, you could close your eyes and point. If someone then moves the row one seat to the left, you have no idea how your pick would have changed. Or you could roll a die for "how many seats from the middle" and count from wherever the middle is. Now if the row moves one seat left, your seat moves exactly one seat left. Same kind of random seat, but you can say how it responds to the row moving.

Training needs exactly that. Backpropagation asks of every step: "if I nudge this number, how does the loss change?" (that rate of change is the gradient, or derivative; see primer.notation). A raw dice roll has no answer, so the gradient would stop at the sampling step and the encoder would never learn. The reparameterization trick rolls the dice first, as an ordinary input ε, and builds the code with arithmetic:

import math
mu, log_var, eps = 0.5, math.log(0.25), 1.2
round(log_var, 3)  # → -1.386
# σ = e^(½ log σ²)
sigma = math.exp(0.5 * log_var)
round(sigma, 3)  # → 0.5
# z = μ + σ ⊙ ε
z = mu + sigma * eps
round(z, 3)  # → 1.1
# nudge μ, then σ, by a tiny h and watch z: ∂z/∂μ = 1 and ∂z/∂σ = ε
h = 1e-6
round((mu + h + sigma * eps - z) / h, 3)  # → 1.0
round((mu + (sigma + h) * eps - z) / h, 3)  # → 1.2

flowchart LR
  subgraph A["Sampling directly: the gradient stops"]
    direction LR
    m1["μ, σ"] --> s1["draw z from N(μ, σ²)<br/>a dice roll"] --> d1["decoder"] --> l1["loss"]
    l1 -. "no route back<br/>through a dice roll" .-> s1
  end
  subgraph B["Reparameterized: the gradient flows"]
    direction LR
    e2["ε from N(0, I)<br/>just another input"] --> s2["z = μ + σ·ε<br/>plain arithmetic"]
    m2["μ, σ"] --> s2 --> d2["decoder"] --> l2["loss"]
    l2 -. "∂z/∂μ = 1, ∂z/∂σ = ε" .-> m2
  end

Why it matters in practice. The same trick lets gradients pass through random choices elsewhere too, for example in the noisy steps of diffusion models (primer.ml.generative.diffusion) and in reinforcement learning with continuous actions. It is a large part of why the VAE paper mattered.

In code: reparameterize is the formula; in VAE.loss_and_gradients the lines under "Through z = μ + σ·ε" are the two gradients above, and a test checks them against finite differences.

The KL penalty: rent for every region

Everyday picture The rent from the everyday picture has a formal name: the KL divergence (Kullback-Leibler divergence) from the region to the standard bell curve, a measure of how different two distributions are that is 0 only when they match (primer.ml.training_stages uses it for distillation). For a bell-curve region and the standard bell curve, it has a short closed form:

import math
def kl(mu, var):
    return 0.5 * (mu ** 2 + var - math.log(var) - 1)
# one code number with μ = 0.5, σ = 0.5 (σ² = 0.25)
round(kl(0.5, 0.25), 3)  # → 0.443
# the rent for the centre alone, then for the size alone
round(0.5 * 0.5 ** 2, 3)  # → 0.125
round(kl(0.0, 0.25), 3)  # → 0.318
# a code number that already is the standard normal pays nothing
kl(0.0, 1.0)  # → 0.0
# Σ over j: two code numbers, their rents add
round(kl(0.5, 0.25) + kl(0.0, 1.0), 3)  # → 0.443
# shrink σ to a pin (σ = 0.05) and the rent jumps
round(kl(0.5, 0.05 ** 2), 2)  # → 2.62

In code: kl_to_standard_normal is the formula, summed over the code numbers; a test checks it against a brute-force average over 400,000 samples.

The whole VAE loss

rebuild, kl_value, beta = 0.2, 0.443, 0.3
# L = ‖x − x̂‖² + β · KL
round(rebuild + beta * kl_value, 3)  # → 0.333

Where does this come from? The VAE paper derives the loss (with β = 1) as the negative of the evidence lower bound (ELBO): a quantity that is guaranteed never to exceed the log-probability the model gives to the real pictures, so pushing the bound up pushes that probability up. Choosing squared error amounts to assuming each pixel is the decoder's output plus bell-curve noise of spread s, and working that through gives β = 2s². This lesson uses β = 0.3, which assumes pixel noise of spread about 0.39.

In code: VAE.loss_and_gradients computes both terms and their gradients; trained_vae trains the lesson's VAE at a chosen β.

Everyday picture Once the map is packed, throw a dart at the middle of it and you land in a neighbourhood, not a field. Generating is exactly that: throw away the encoder, draw a code from the standard bell curve, and decode.

Tiny example Draw the same code as before, z = (1.2, −0.9), and decode it with the VAE. Out comes a horizontal stroke through the centre of the picture, with squared distance about 0.06 from a real stroke in the training set. The plain autoencoder turned that same code into a smear 7.5 away. Over 500 draws the VAE's median distance is about 0.56 against the plain autoencoder's 3.

flowchart LR
  N["draw z from N(0, I)<br/>2 random numbers"] --> D["decoder g"] --> X["a new picture<br/>never seen in training"]
  ENC["encoder"] -. "not used when generating" .-> N

Interpolation: walking between two codes

Everyday picture A morph between two faces in a film: every in-between frame should look like a face, not a double exposure.

Tiny example Encode two pictures to codes and , step along the straight line between them, and decode every stop:

z_a, z_b = [-1.0, 0.5], [1.0, 1.5]
# z_t = (1 − t) z_a + t z_b, a quarter of the way along
t = 0.25
[(1 - t) * a + t * b for a, b in zip(z_a, z_b)]  # → [-0.5, 0.75]
# the nine stops of a walk
[i / 8 for i in range(9)]  # → [0.0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0]

Why it matters in practice. Smooth latent spaces are what made "move this slider to add a smile" demos possible, and interpolation is still a standard sanity check for any learned code: if the walk is full of junk, the space has holes.

In code: generate draws and decodes; interpolate encodes two pictures and decodes the walk between them; largest_step measures its biggest jump.

Everyday picture Go back to the rent. Set it too low and every region shrinks to a pin far from the centre: rebuilds are sharp, but the holes come back. Set it too high and every region moves to the centre and swells to the full bell curve. Then every picture's region is the same region, the code carries no information at all, and the decoder can do nothing better than draw the average of all the strokes: a grey smudge. That failure is called posterior collapse.

Tiny example The lesson trains six VAEs, with β from 0.01 to 3:

(beta_sweep produces this table, including β = 0.03, and demo prints it.)

The samples are best in the middle. At β = 3 the KL is zero: the regions are the standard bell curve itself, and the rebuild error of 6.08 is what you get by drawing the average stroke every time.

Why VAE pictures are blurry

Even at a good β the regions overlap, so one code can stand for several different pictures. What should the decoder draw for it? Take a single pixel that is black (1) in one of those pictures and white (0) in the other, equally likely. Squared error has a clear answer:

outcomes = [0.0, 1.0]
def expected_error(c):
    return sum((x - c) ** 2 for x in outcomes) / len(outcomes)
# E[(x − c)²] for three guesses: white, grey, black
[expected_error(c) for c in (0.0, 0.5, 1.0)]  # → [0.5, 0.25, 0.5]
# E[x]: the average outcome, which is the winning guess
sum(outcomes) / len(outcomes)  # → 0.5

Across a whole picture, "average the pictures this code might mean" is a blur. The more the regions overlap, the more pictures each code must stand for, and the blurrier the output.

Why it matters in practice. β is a real dial, and the name β-VAE (Higgins and colleagues, 2017) comes from turning it up to get codes whose numbers line up with separate factors, like direction and offset here. The blur is the VAE's best-known weakness. Modern systems fix it by adding a second judge of realism to the loss (an adversarial loss, the trick behind primer.ml.generative.gans, as in VQGAN), or by letting a stronger generator do the generating while the autoencoder only compresses.

In code: beta_sweep trains and measures one VAE per β; expected_squared_error is the grey-pixel calculation.

Everyday picture The autoencoder became the zip format for images and sound that other generative models work inside.

Tiny example Stable Diffusion's autoencoder turns a 512 × 512 colour image (512 × 512 × 3 = 786,432 numbers) into a 64 × 64 × 4 grid of codes (16,384 numbers): 48 times fewer. The diffusion model does all its work on that small grid, and the decoder paints full-size pixels only at the very end. Its KL weight is tiny, so it is mostly a compressor, with just enough rent to keep the codes well-behaved.

A VQ-VAE (vector-quantized VAE) goes one step further: it snaps each code vector to the nearest entry of a learned codebook, so an image becomes a grid of whole numbers, the entries' positions in the codebook. Those numbers are tokens, and a transformer can read and write them exactly as it reads and writes words. Neural audio codecs do the same for sound.

flowchart LR
  subgraph LD["Latent diffusion"]
    direction LR
    I1["image<br/>512 × 512 × 3"] --> E1["VAE encoder"] --> L1["latent<br/>64 × 64 × 4"]
    L1 --> DF["diffusion model<br/>works here"] --> D1["VAE decoder"] --> O1["image"]
  end
  subgraph TK["Tokenizer for a transformer"]
    direction LR
    I2["image or audio"] --> E2["encoder"] --> Q["snap each vector to<br/>nearest codebook entry"]
    Q --> T["grid of token ids"] --> TR["transformer"]
    TR --> D2["decoder"] --> O2["image or audio"]
  end

Why it matters in practice. When a model "generates an image", there is very often an autoencoder at both ends of the pipeline. Its quality sets a ceiling: whatever detail the decoder cannot rebuild, no generator working in its latent space can produce.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.