Variational Autoencoders

A plain autoencoder's latent space has holes, so sampling it gives nothing. Make the encoder emit a distribution instead of a point and the gaps fill in.

The encoder emits a mean and a spread per latent dimension, noise drawn outside the model joins them through addition and multiplication, and that join keeps a differentiable path through sampling

Why Does This Exist?

A ceramics studio holds 4,000 photographs of vase silhouettes and wants twenty new shapes by Monday. Not copies, not blends of two chosen pots. New shapes that belong with the rest.

An autoencoder gets close. It squeezes each silhouette into eight numbers and rebuilds it faithfully, so those eight numbers evidently capture what a vase is. Feed the decoder eight numbers of your own and you should get a vase back.

You get noise. Every time.

That is the wall, and it is not a bug. Nothing in training asked the encoder to use the space between the silhouettes it saw. It scattered 4,000 codes wherever was convenient, and the decoder learned what to do only at those exact spots. A random point lands in a void the decoder never visited.

So the fix targets the space, not the decoder. Two terms first. A latent space is where those codes live. A divergence measures how far one distribution sits from another, and here that is KL divergence.

Think of It Like This

Pinning butterflies against spraying them

A collector pins 4,000 butterflies to a board, each at one exact spot. Precise, and the board is mostly bare cork. Point anywhere between two pins and you point at nothing.

Now do it differently. For each butterfly, spray a soft patch of paint centred where the pin would have gone, wide enough that neighbouring patches overlap. The board fills in. Every point sits inside somebody's patch, so every point means something.

Two pressures set the patch size, pulling opposite ways. Keeping butterflies distinguishable wants patches small and far apart. Leaving no bare cork wants them broad and overlapping. Where those settle is a board you can point at blind and still hit a butterfly.

Where it stops: paint patches never explain how training reaches through sampling. The board shows why distributions fill space, not how gradients survive randomness. That second half needs the trick below, which has no counterpart in paint.

How It Actually Works

Two terms, permanently in tension

The encoder stops emitting a code. It emits a mean and a spread, one pair per latent dimension, and the code is drawn from that distribution. The loss gains a second term:

\mathcal{L} = \underbrace{\lVert x - \hat{x} \rVert^2}_{\text{reconstruction}} + \underbrace{D_{\mathrm{KL}}\!\left(q(z \mid x) \,\|\, p(z)\right)}_{\text{keep the patches spread out}}

$x$ is the silhouette, $\hat{x}$ the rebuild, $q(z \mid x)$ the distribution the encoder produced for it, and $p(z)$ a standard normal, the shape you want the whole board to take.

Those two terms are the collector's two pressures. Reconstruction wants each vase's patch narrow and far from the others, so codes stay distinguishable. The divergence term wants every patch centred at the origin with unit width, so they overlap and fill the space. Neither wins. What you get is organised and continuous, which makes sampling work: draw eight numbers from a standard normal, decode, and land inside a region the decoder was trained on.

One dimension, with numbers

Take a single latent dimension where the encoder produced a mean $\mu = 0.5$ and a spread $\sigma = 0.8$ . For two Gaussians the divergence has a closed form, and against a standard normal it is:

D_{\mathrm{KL}} = -\log \sigma + \frac{\sigma^2 + \mu^2 - 1}{2} = 0.2231 + \frac{0.64 + 0.25 - 1}{2} \approx 0.168

About 0.168 nats. Small, because this patch already sits near the origin at near unit width, which is what the divergence term rewards. Now sample: draw $\epsilon = 0.3$ from outside and compute $z = 0.5 + 0.8 \times 0.3 = 0.74$ . That 0.74 is the code the decoder receives. During training the loss adds thousands of these small divergence bills, one per dimension per example, and the sum is what drags the whole board toward the standard normal.

The trick that makes it trainable

One hole remains. Training needs gradients to reach the encoder, and a sampling step sits in the way. You cannot differentiate through a random draw.

The reparameterisation trick moves the randomness out of the path. Instead of sampling from the encoder's distribution, sample $\epsilon$ from a standard normal outside the model, then compute:

z = \mu + \sigma \epsilon

Same distribution, different graph. Now $\epsilon$ is an input, like the image, and $z$ is an ordinary differentiable function of $\mu$ and $\sigma$ , so the derivative of anything downstream flows straight back through. That is the accented join in the diagram, and the single move the method rests on.

What it cost

Reconstructions come out blurry. Squared error on pixels rewards the average of everything that could plausibly have been there, and the average of several sharp edges is soft. That is the loss behaving correctly, not the model failing.

And one term can win outright. If the divergence pressure overwhelms reconstruction, the encoder gives up and returns the standard normal for every input. The patches merge and the board carries no information. That failure has a name, and it is below.

A VAE optimises a bound on the data's likelihood, so it carries an explicit if approximate density. GANs never write a probability down.

Code

The trick, measured. Both routes estimate the same derivative, and only one is cheap enough to train with.

1import numpy as np2
3rng = np.random.default_rng(0)4MU, SIGMA, N = 1.5, 0.7, 100_0005
6
7def decoder_loss(z: np.ndarray) -> np.ndarray:8    return z ** 2  # a stand-in for whatever the decoder charges for a code9
10
11eps = rng.normal(0.0, 1.0, N)12z = MU + SIGMA * eps  # the noise is an input now, so the derivative of z by mu is just 113reparam = 2.0 * z  # chain rule straight through the sample14score = decoder_loss(z) * (z - MU) / SIGMA ** 2  # no path, so estimate it from scratch15print(f"exact           {2 * MU:.4f}")16print(f"reparameterised {reparam.mean():.4f}   spread {reparam.std() / np.sqrt(N):.4f}")17print(f"score function  {score.mean():.4f}   spread {score.std() / np.sqrt(N):.4f}")18# -> exact           3.000019# -> reparameterised 2.9987   spread 0.004420# -> score function  2.9925   spread 0.0222

Both land on 3.0, so both are valid. The spread is the story: five times noisier without the trick, from the same 100,000 samples. Multiply that noise across millions of parameters and training stops working.

Watch Out For

Posterior collapse, where the divergence term wins

The symptom is oddly cheerful. Total loss falls nicely, the divergence term drops to near zero, and reconstructions look decent. Then you sample and every draw decodes to roughly the same vase.

The encoder has stopped reading its input. It returns the standard normal for every silhouette, the divergence term is perfectly satisfied, and a strong decoder carries reconstruction alone by memorising an average pot. Watch the divergence term as its own number, not inside the total: a healthy run holds it clearly above zero. Both standard fixes weaken the pressure temporarily. Ramp the divergence weight up from zero over the first few thousand steps, or hold it below a floor so the encoder is never rewarded for going silent.

Sampling from the encoder instead of from the prior

Generation code reads an image, takes the encoder's mean and spread, samples, decodes. Lovely output, and it generated nothing. Those are variations on a silhouette you already had, which is what an autoencoder with noise does.

Sample from the prior, the standard normal, with no input image anywhere in the call. That is the only version that tests whether the space actually filled in. The wrong way round hides the exact failure you needed to catch, because a model with a half-empty latent space still reconstructs beautifully.

The Quick Version

A plain autoencoder's latent space has voids, so decoding a point you never encoded returns nothing. This page is the fix.
The encoder emits a mean and a spread per dimension instead of a code, and the code is sampled from that distribution.
The loss gains a divergence term pulling every distribution toward a standard normal, fighting reconstruction on purpose. That tension fills the space.
Sampling blocks gradients, so reparameterise: draw the noise outside the model and compute mean plus spread times noise. The derivative flows straight through.
Squared-error reconstruction averages over plausible outputs, which is why samples come out soft.
If the divergence term wins outright the encoder ignores its input. Watch that term separately, and sample from the prior when you evaluate.