Kingma and Welling made random sampling differentiable

A generative model trained by plain gradient descent

A generative model learns to produce new examples that resemble a training set. Diederik Kingma and Max Welling, at the University of Amsterdam, built theirs around two networks and a hidden space of factors, the latent space, and showed it could be trained end to end by ordinary gradient descent.¹ An encoder reads a data point and places it in that space; a decoder reads a point from the space and rebuilds the data. An encoder feeding a decoder, trained to reproduce its own input, is an autoencoder, the same kind of component that compresses images inside latent diffusion.

Two choices made this autoencoder generative rather than a mere compressor. The encoder outputs a whole distribution over the latent space rather than a single code, so the space fills in and new points can be drawn from it. And the model trains against a single objective a gradient can climb. The one obstacle was that the objective runs through a random sampling step, and a gradient does not normally pass through a coin toss. A move they call the reparameterization trick removed it.¹

The encoder describes a point as a distribution

In a plain autoencoder the encoder maps each input to one code. The VAE's encoder, which the paper calls a recognition model or a probabilistic encoder, instead reports a distribution: for every dimension of the latent space it emits a mean and a spread.¹ A digit lands not on one point but on a small blur of nearby points, any of which the decoder should rebuild into that digit.

The spread is what ties the space together. The model asks every example's blur to resemble one fixed reference distribution, the prior. That prior is a centered, standard bell curve in every dimension, written p(\mathbf{z}) = \mathcal{N}(0, I).¹ Because the blurs overlap and share that shape, the space has no dead gaps: draw a point near the center and the decoder, having seen its neighborhood, turns it into a plausible new example.

The latent space is not something to visualize. The experiments swept spaces from two dimensions up to two hundred, and no one pictures a two-hundred-dimensional space, professional mathematicians included.¹ The mean and spread are just lists of numbers the encoder prints, one pair per dimension.

Two terms pulling in different directions

Training maximizes one quantity, the evidence lower bound, or ELBO. It is a lower bound on how likely the data is under the model, and it splits into two terms that want different things.¹

The second term is the reconstruction reward. The decoder receives a code drawn from the encoder's blur and is scored on how well it rebuilds the original. The first term has a name worth defining: the KL divergence. Before the symbol, the plain question it answers is how different two distributions are. If you expected draws from one distribution but they actually come from another, the KL divergence measures how badly the first predicts the second. Here it measures the gap between one example's blur and the prior, and the objective subtracts it, which pushes every blur toward that shared bell curve.

When the prior and the encoder's blur are both bell curves, the gap has a short closed form, one that needs no sampling. For each latent dimension it adds up \tfrac{1}{2}\bigl(\mu_j^2 + \sigma_j^2 - 1 - \log\sigma_j^2\bigr) from the encoder's mean and spread.¹ Set the mean to zero and the spread to one and every piece cancels: the blur already matches the prior, and the penalty is nothing. Push the mean out to three, or shrink the spread to a tenth, and the penalty climbs.

The two terms are a tug of war. Reconstruction alone would let the encoder scatter each example to its own far corner and memorize the set, leaving the space between corners empty. The KL term alone would crush every blur onto the prior and throw away what makes one example different from another. The useful models live in the balance, packed tight enough to sample from and loose enough to tell the data apart.

How much weight the balance gives each side is a dial. Higgins and colleagues at DeepMind multiplied the KL term by a single knob they called β. At β of one it is the plain VAE. Turning β above one presses the blurs harder toward the prior, trading reconstruction sharpness for latent factors that vary more independently, the effect they named disentanglement.²

The balance is also fragile. Bowman and colleagues trained a VAE with a strong recurrent decoder on text. Most runs quietly set the encoder's blur equal to the prior, driving the KL term to zero, and the decoder learned to model the data on its own and ignore the latent code entirely.³ The failure now goes by posterior collapse. Their fixes were to ramp the KL weight up slowly during training and to weaken the decoder so it had to use the code. The tidy two-term picture holds as the objective, but in practice the terms can collapse into each other.

Sending a gradient through a random draw

One step blocks the whole plan. To train by gradient descent you take the gradient of the objective with respect to the encoder's parameters. But the objective is an average over codes drawn from the encoder's blur, and a gradient cannot pass through a random draw. The obvious workaround, estimating the gradient from samples, gives an estimate the paper reports as so noisy it "exhibits very high variance and is impractical."¹

The trick is to move the randomness out of the path the gradient travels. Rather than draw a code directly from the encoder's blur, draw a fixed unit of noise from a standard bell curve and build the code from it and the encoder's outputs.

Now the code is a plain function of the encoder's mean and spread plus outside noise, and the random part no longer sits between the objective and the parameters.¹ Fix the noise at a half, and the code becomes the mean plus half the spread. Nudge the mean up and the code moves up with it, one for one. The gradient has a clean path from the code back to the encoder, and the encoder can learn.

The idea was in the air. Kingma and Welling gave the deep latent-variable case its standard recipe and its name. But they credit Rezende, Mohamed and Wierstra with an independent, near-simultaneous derivation of the same connection between encoders, generative models and gradient-based inference.⁴ They also point to Salimans and Knowles a year earlier for a similar reparameterization used elsewhere.¹ The move was co-discovered, and the paper is more careful about that than its later reputation is.

From a CPU proof to Stable Diffusion's latent space

The experiments were small, and the paper is plain about it. The models were single-hidden-layer networks, run on handwritten digits from MNIST and on the Frey Face set of one person's changing expression.¹

Those 500- and 200-unit networks trained on a single CPU, roughly twenty to forty minutes per million samples.¹ What the paper reported was a likelihood. It gave the variational lower bound over training, against a competing method called wake-sleep, and a marginal-likelihood estimate sampled on the first thousand train and test points.¹ It reports no image-quality score anywhere. So the popular line that a VAE generates images runs ahead of what this paper measured: how probable the data was under the model, not how good its pictures looked.

The pictures tend to come out soft, and the reason is built into the reconstruction term. A plain VAE scores its rebuild pixel by pixel, under what Larsen and colleagues call an element-wise Gaussian model, and they argue such element-wise errors "are not adequate for images" with visual structure.⁵ When several outputs could each explain an input, minimizing pixel error settles on their average, and the average of many sharp possibilities is a smear.⁶ The famous blurriness is that averaging, not a verdict any one paper pronounced.

What survived was the latent space itself. The autoencoder in latent diffusion runs its whole diffusion process inside the compressed space of a pretrained encoder. One of the two encoders Rombach and colleagues train is regularized with "a slight KL-penalty towards a standard normal on the learned latent, similar to a VAE."⁷ The discrete branch of the family, the VQ-VAE, swaps the continuous blur for a learned codebook of discrete codes and a learned prior over them, which sidesteps posterior collapse.⁸ Scaled up, that line produced ImageNet samples the authors put at a quality "that rivals that of state of the art Generative Adversarial Networks."⁹ The fidelity came from a powerful prior over the latent grid. The autoencoder alone did not supply it.