Variational Autoencoders

A plain autoencoder gives you codes but no way to generate new ones: pick a random point in the code space and the decoder produces nonsense, because nothing ever asked the space to be filled in. A VAE fixes that by encoding to a distribution instead of a point — and the reparameterisation trick is what makes that trainable at all.

The problem a VAE solves

Train a standard autoencoder on faces and the codes land wherever the optimiser found convenient — clusters here, empty regions there, no particular scale. Sample a point between two clusters and the decoder has never seen anything near it. The space has holes, and a space with holes cannot be generated from.

The VAE’s answer is to stop mapping an input to a point. Map it to a small cloud, and add a penalty that keeps all the clouds near the origin and roughly unit-sized. Overlapping clouds force the decoder to produce something sensible everywhere in between, and a code space with a known shape can be sampled from directly.

The encoder outputs a distribution

encoder(x) -> mu(x)       : (k,)   the centre of the cloud
              log_var(x)  : (k,)   its log-variance, one per dimension

sigma = exp(0.5 * log_var)

z ~ Normal(mu, sigma^2)             sample a code
x_hat = decoder(z)

Two heads instead of one. The log-variance rather than the variance is predicted because it is unconstrained — a network can emit any real number and the exponential makes it positive, with no clamping and no risk of a negative variance.

And now the problem. There is a random draw in the middle of the forward pass, and the loss is computed after it. Gradient descent needs to know how the loss changes when mu changes. But z was sampled, and sampling is not a differentiable function of its parameters: z did not come out of an arithmetic expression you can take a derivative of, it came out of a random number generator. The chain rule has nowhere to go.

The reparameterisation trick

The trick is a change of variables, and once you see where the randomness ends up it stops being clever. Instead of drawing z from a distribution that depends on the parameters, draw a fixed standard normal and build z out of it arithmetically:

Before:  z ~ Normal(mu, sigma^2)         random node, depends on mu, sigma
After:   eps ~ Normal(0, 1)             random node, depends on NOTHING
         z = mu + sigma * eps           deterministic, differentiable

The two lines produce identically distributed samples. If eps is standard normal then mu + sigma*eps is normal with mean mu and standard deviation sigma — that is the definition of a location-scale family. Nothing about the model has changed. What has changed is where the randomness sits in the graph.

Now the derivatives exist and are embarrassingly simple. Treat eps as a constant, because for this backward pass it is one:

z = mu + sigma * eps

dz/dmu    = 1
dz/dsigma = eps

so  dL/dmu    = dL/dz
    dL/dsigma = dL/dz * eps

That is the whole trick. The gradient flows from the loss, through the decoder, through z, and straight into mu and sigma, because z is now an ordinary arithmetic expression in them. The sampler still runs, but it sits off the path the gradient takes, contributing a number that the backward pass treats as data rather than as a function.

The name for what it does to the estimator is worth one sentence: it turns a high-variance score-function gradient estimate into a low-variance pathwise one. In practice that is the difference between a model that trains and one that does not.

The KL term, written out

The loss has two parts: rebuild the input, and keep the cloud near a standard normal.

loss = reconstruction + beta * KL( q(z|x) || Normal(0, I) )

For diagonal Gaussians the KL has a closed form, per dimension:

KL = 0.5 * ( mu^2 + sigma^2 - log(sigma^2) - 1 )

Check it at the point where it should vanish. With mu = 0 and sigma = 1: 0.5 * (0 + 1 - 0 - 1) = 0. Good. Now push the cloud away: mu = 2, sigma = 1 gives 0.5 * (4 + 1 - 0 - 1) = 2. Shrink it instead: mu = 0, sigma = 0.1 gives 0.5 * (0 + 0.01 + 4.605 - 1) = 1.81. The term punishes both drifting away and collapsing to a point, which is precisely the pair of behaviours that would reintroduce the holes.

beta is the dial between the two objectives. Below 1 the model reconstructs better and the space gets less regular; above 1 the space gets smoother and more disentangled and the reconstructions get worse. There is no principled setting. It is a knob, and saying so is more useful than pretending otherwise.

Why the loss is called a lower bound

The two terms are not an arbitrary pairing. They are what falls out when you try to maximise the likelihood of the data and cannot.

What you want:   log p(x)
                 = log integral over z of p(x|z) p(z) dz
                 ...an integral over the whole latent space. Intractable.

What you get:    log p(x)  >=  E_q[ log p(x|z) ]  -  KL( q(z|x) || p(z) )
                              \_____________/       \________________/
                               reconstruction              the KL term

The gap between the two sides is exactly:
                 KL( q(z|x) || p(z|x) )

...the distance from the encoder's guess to the true posterior.

Three things follow, and each one explains a design decision that otherwise looks arbitrary.

• The bound tightens as the encoder improves. The slack is the encoder’s error as an approximation of the true posterior, so a more expressive encoder gives a bound closer to the real likelihood. That is the entire motivation for richer posterior families — normalising flows over q, for instance.
• The reconstruction term is a log-likelihood. It is not “an error measure someone picked”. Using squared error means asserting that p(x|z) is Gaussian with fixed variance, and the blurriness in the next section is that assertion being honoured. Choosing a different likelihood — a discretised logistic, a categorical over pixel values — is choosing a different failure mode.
• beta breaks the bound on purpose. At beta = 1 the loss is exactly the negative bound. At any other value it is no longer a bound on anything, which is worth knowing before treating the number as a likelihood. It is a regularisation weight from that point on, and should be described as one.

Two failure modes

Posterior collapse

If the decoder is powerful enough to model the data on its own — an autoregressive decoder over text is the classic case — the cheapest way to reduce the loss is to ignore z entirely. The encoder then outputs mu = 0, sigma = 1 for every input, the KL term goes to exactly zero, the reconstruction term is handled by the decoder alone, and the latent carries no information at all.

The diagnostic is unambiguous: KL near zero and reconstructions still reasonable means collapse. The usual mitigations are KL annealing (start beta at zero and raise it), free bits (do not penalise the first few nats per dimension), or simply weakening the decoder.

Blurriness

A squared-error reconstruction loss is the log-likelihood of a Gaussian with fixed variance, and the value that minimises expected squared error is the mean. When several outputs are plausible, the mean of several sharp images is a blurred image, and the model is being correctly optimised when it produces one.

This is not a bug that better training fixes. It is what the objective asks for, and it is the reason VAEs alone never produced photorealistic samples.

Why VAEs lost image generation and stayed anyway

As a standalone image generator the VAE lost to GANs and then to diffusion, for the reason above: the likelihood term averages, and averaging looks like fog.

It stayed because the compression is excellent even when the samples are not. In latent diffusion the VAE is not the generator; it is the compressor at both ends. It maps 512 by 512 by 3 down to 64 by 64 by 4 — 48 times fewer elements — diffusion does the generative work in that small space, and the decoder maps back. The blur problem is handled by training that decoder with perceptual and adversarial losses alongside the pixel one, which is a neat piece of history: the adversarial loss that lost the generation contest ended up inside the winner, doing the one job it was always good at.

The general lesson generalises past this architecture. A probabilistic latent space is worth having; a Gaussian pixel likelihood is not. Most of what came after keeps the first and replaces the second.