Variational Autoencoders and the ELBO
How maximising an intractable data likelihood turns into maximising a tractable lower bound, why the reparameterisation trick is what makes that bound differentiable, and what posterior collapse costs you.
You want a generative model \(p_\theta(x)\) over images or molecules, and you believe the data was produced by sampling a low-dimensional latent \(z\) and then decoding it. Writing that belief down gives \(p_\theta(x) = \int p_\theta(x \mid z) p(z) \, dz\), and the integral is hopeless: for a 64-dimensional latent there is no closed form and no Monte Carlo estimator with usable variance. Every practical latent-variable deep generative model since 2013 works around this by optimising a bound instead of the thing itself (Kingma and Welling, 2013, Auto-Encoding Variational Bayes, arXiv:1312.6114; independently Rezende et al., 2014, arXiv:1401.4082).
Deriving the bound
Introduce an approximate posterior \(q_\phi(z \mid x)\), an encoder network that guesses which latents could have produced this \(x\). For any such \(q\), Jensen's inequality gives
The right-hand side is the evidence lower bound, the ELBO. The first term is reconstruction: decode the latent and score the original input. The second is a regulariser pulling the encoder's output distribution toward the prior, usually \(\mathcal{N}(0, I)\). The gap between the ELBO and \(\log p_\theta(x)\) is exactly \(D_{\mathrm{KL}}(q_\phi(z \mid x) \| p_\theta(z \mid x))\), the error in your posterior approximation. This is worth internalising: you are never told how loose the bound is, because measuring the gap requires the true posterior you could not compute in the first place.
Why the reparameterisation trick is the whole trick
The ELBO contains an expectation over a distribution whose parameters you are differentiating. Sampling \(z \sim \mathcal{N}(\mu_\phi(x), \sigma_\phi(x)^2)\) is not a differentiable operation, and the score-function (REINFORCE) estimator that handles it in general has variance high enough to stall training.
The fix is to move the randomness out of the path: sample \(\epsilon \sim \mathcal{N}(0, I)\) from a distribution with no parameters, then set \(z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon\). Now \(z\) is a deterministic, differentiable function of \(\phi\) and a fixed noise draw, and ordinary backpropagation flows through it. One sample per datapoint is usually enough. That single substitution is what turned variational inference from a per-model derivation exercise into something you can express in ten lines of a deep learning framework.
VAE versus autoencoder, and versus diffusion
A plain autoencoder minimises reconstruction error alone. Nothing forces its latent space to be dense or meaningful, so sampling a random point and decoding it produces noise. The KL term is what makes a VAE generative: it forces encoded data to occupy the same region the prior samples from, so decoding a fresh \(z \sim \mathcal{N}(0, I)\) lands somewhere plausible.
Diffusion models are best understood as VAEs with a fixed, hand-designed encoder (the forward noising process) and a hierarchy of latents, one per timestep. That is why the diffusion training objective is also a weighted ELBO, and why latent diffusion systems stack a diffusion model on top of a VAE's latent space rather than replacing it.
When it breaks
- Posterior collapse. If the decoder is powerful enough to model \(x\) on its own, notably an autoregressive decoder, the cheapest way to shrink the loss is to drive \(q_\phi(z \mid x)\) to the prior and ignore \(z\) entirely. The KL term goes to zero, reconstruction stays good, and the latent code carries no information (Bowman et al., 2015, arXiv:1511.06349). KL annealing, free bits, and weakening the decoder all attack this; none of them fixes it permanently.
- Blurry samples. With a Gaussian likelihood, maximising \(\log p_\theta(x \mid z)\) is mean-squared error, and the mean of several plausible reconstructions is a blur. This is a property of the loss, not of the architecture, which is why VQ-VAE replaced the continuous latent with a discrete codebook and a learned prior over codes (van den Oord et al., 2017, Neural Discrete Representation Learning, arXiv:1711.00937).
- The bound is not the likelihood. Two models with identical ELBOs can have different true likelihoods. Comparing VAEs to models trained by exact likelihood on the reported number is comparing a lower bound with the real quantity.
- KL weighting is a knob with no principled setting. \(\beta\)-VAE reweights the KL term to trade reconstruction fidelity for disentanglement, and the "right" \(\beta\) is found empirically, per dataset.
10 flashcards for this concept
Click a card to reveal the answer.