A variational autoencoder doesn't predict the data. It predicts a distribution over a latent (a mean and a variance), and then trains against a number it can never actually see. The likelihood is out of reach, so the whole method is a clever way around that one problem: bound what you can't compute, then close the gap. It's a beautiful little maneuver, and this post walks the whole chain, from what the network outputs to the loss we minimize.
So what does the network actually output?
A VAE has three pieces: a prior \(p(z) = \mathcal{N}(0, I)\) over a latent variable \(z\), a decoder \(p(x \mid z)\) that turns a latent back into data, and an encoder \(q(z \mid x)\) that maps data to a distribution over latents. The encoder is the neural network, and the first thing to be precise about is what it emits.
Given an input \(x\), the encoder does not output a latent vector. It outputs the parameters of a Gaussian over \(z\): a mean \(\mu(x)\) and a variance \(\sigma^2(x)\) (in practice a log-variance, so it stays unconstrained). That is what "the VAE predicts", mean and variance:
To actually get a latent we have to sample from that Gaussian, and sampling is not differentiable, which would cut the gradient off from \(\mu\) and \(\sigma\). The fix is the reparameterization trick: draw the randomness from a fixed standard Gaussian and push it through the predicted parameters,
Now \(z\) is a deterministic, differentiable function of \(\mu\) and \(\sigma\), with all the noise quarantined in \(\varepsilon\). The full derivation of this "standard Gaussian to latent" map (and the closed form of the KL it induces) is in the side note at the end.
The likelihood is out of reach (here's why)
We would like to train by maximum likelihood: maximize \(\log p(x)\). By the marginal,
which is an integral over the entire latent space. For any decoder expressive enough to be worth using, that integral has no closed form and cannot be enumerated: \(z\) is continuous and high-dimensional. And since Bayes' rule writes the posterior as
the intractable \(p(x)\) sits in its denominator, so the true posterior is out of reach too. That is exactly why we introduced \(q(z \mid x)\) to approximate it, and why we now need a way to measure how good that approximation is.
Two ways to measure the gap, only one is usable
To measure the gap between \(q(z \mid x)\) and \(p(z \mid x)\) we use a KL divergence, but the order matters, because KL is not symmetric. There are two choices:
- Forward, \(D_{\mathrm{KL}}\big(p(z \mid x) \,\|\, q(z \mid x)\big)\): an expectation over the true posterior \(p(z \mid x)\), the thing we just said we can't sample from. A non-starter.
- Reverse, \(D_{\mathrm{KL}}\big(q(z \mid x) \,\|\, p(z \mid x)\big)\): an expectation over our encoder \(q(z \mid x)\), which we built and can sample from cheaply via the reparameterization trick.
So we minimize the reverse KL, not only because it is the tractable direction, but because it is the one that turns into a quantity we can optimize. (Reverse KL is also mode-seeking: it would rather concentrate \(q\) on one mode of the posterior than smear probability where the posterior has none, which is the well-behaved failure mode for a sampler.)
Where KL even comes from
Before using KL it is worth seeing where it comes from. The information content of an outcome under a density is \(I(x) = -\log p(x)\): rarer outcomes carry more bits. If \(q\) is the distribution we are using and \(p\) is the true one, the excess information we pay per sample for using the wrong distribution is
and averaging that surprise gap over the true distribution is the Kullback-Leibler divergence,
zero only when the two distributions agree almost everywhere. That single non-negativity is the entire engine of what follows.
Now the derivation
Write the reverse KL between our encoder and the true posterior, and substitute Bayes' rule for \(p(z \mid x)\):
Replacing \(p(z \mid x) = \dfrac{p(x \mid z)\, p(z)}{p(x)}\) turns the ratio inside the log into \(\dfrac{q(z \mid x)\, p(x)}{p(x \mid z)\, p(z)}\), so
Split the logarithm and pull \(\log p(x)\) out of the integral, it does not depend on \(z\), and \(\int q(z \mid x)\, dz = 1\):
What the ELBO is, and why pushing it up helps twice
Call that integral \(\mathcal{L}\). Rearranging the line above is an exact identity, nothing has been approximated:
The KL term is non-negative, so \(\mathcal{L} \le \log p(x)\): \(\mathcal{L}\) is a lower bound on the log evidence, the evidence lower bound. Expanding it into its two pieces,
The ELBO
\( \mathcal{L} = \mathbb{E}_{q(z \mid x)}\!\big[\log p(x \mid z)\big] - D_{\mathrm{KL}}\!\big(q(z \mid x) \,\|\, p(z)\big) \le \log p(x). \)
We maximize \(\mathcal{L}\) because \(\log p(x)\) (the thing we truly want) is fixed and unreachable, while the gap to it is precisely \(D_{\mathrm{KL}}(q \,\|\, p(z \mid x)) \ge 0\). Pushing \(\mathcal{L}\) up therefore does double duty: it drives a proxy for the likelihood upward and squeezes that gap shut, pulling \(q(z \mid x)\) toward the true posterior. The bound is tight exactly when the encoder recovers \(p(z \mid x)\).
Flipping it into a loss
Optimizers minimize, so the training loss is just the negative ELBO:
The two terms are exactly the two halves of the autoencoder. The first is reconstruction: sample \(z\) with the reparameterization trick, decode it, and penalize a decoder that makes the original \(x\) unlikely: for a Gaussian decoder this is mean-squared error, for a Bernoulli decoder it is binary cross-entropy. The second is regularization: it pulls the predicted \(\mu, \sigma\) toward the standard-normal prior so the latent space stays smooth and samplable. Minimizing this loss maximizes the ELBO, which maximizes a lower bound on \(\log p(x)\).
Side note: moving the Gaussian into the latents
Two loose ends from above, both about the Gaussian the encoder predicts.
The reparameterization trick (the pretty part)
Here's the move I think is genuinely beautiful. You can't backprop through a coin flip, so you don't even try. You pull the randomness out into a fixed standard Gaussian, sample that, and push it through your predicted mean and variance. The noise stays quarantined, the gradients flow, and nothing about the network had to go deterministic. Simple, and a little bit magic.
We need samples \(z \sim \mathcal{N}(\mu, \sigma^2)\), but we also need gradients to flow back into \(\mu\) and \(\sigma\). A location- scale Gaussian factors cleanly: if \(\varepsilon \sim \mathcal{N}(0, I)\), then
The stochastic part is the fixed, parameter-free \(\varepsilon\); \(\mu\) and \(\sigma\) enter only through deterministic, smooth operations. So \(\partial z / \partial \mu = 1\) and \(\partial z / \partial \sigma = \varepsilon\), and the expectation \(\mathbb{E}_{q}[\,\cdot\,]\) becomes \(\mathbb{E}_{\varepsilon \sim \mathcal{N}(0,I)}[\,\cdot\,]\), which we estimate by Monte Carlo and differentiate term by term. The Gaussian noise is, quite literally, transported into the latent.
Closed form of the KL term
Because both the encoder \(q = \mathcal{N}(\mu, \sigma^2)\) and the prior \(p = \mathcal{N}(0, 1)\) are Gaussian, the regularizer needs no sampling. In one dimension,
Using \(\mathbb{E}_q[(z-\mu)^2] = \sigma^2\) and \(\mathbb{E}_q[z^2] = \mu^2 + \sigma^2\),
Summed over a diagonal latent of dimension \(J\), this is the term that actually appears in the loss,
a closed-form function of exactly the mean and variance the network predicts, which is where this note began.
It took me over ten tries to actually derive this from memory.