A code that is a distribution
The encoder outputs a mean and a variance rather than a point, so one input maps to a cloud in the latent space.
Sampling from that cloud during training means nearby codes must decode to plausible outputs, which is what fills in the gaps a plain autoencoder leaves.
A prior is imposed on the whole latent space, so a vector drawn from a standard normal decodes to something valid, and the model can generate.
The ELBO
The likelihood of the data is intractable, so training maximises a lower bound on it, which is what the E, L and B stand for.
The first term is reconstruction: decode the sampled code and score the original against it.
The second is a regulariser pulling each input's cloud towards the prior, and it is the price paid for a latent space that is continuous.
For two Gaussians that divergence has a closed form, which is the practical reason the Gaussian assumption is made at all.
The reparameterisation trick
A sampling step has no gradient, so backpropagation cannot pass through it as written.
Moving the randomness into an input the parameters do not depend on makes the sampled code a deterministic function of the encoder's outputs.
The gradient then flows through it normally.
That one substitution is what made variational inference trainable by gradient descent, and it is the contribution of the 2013 paper.
One training step, in order
- Encode the input to a mean and a log variance.
- Draw one and form by the reparameterisation.
- Decode and compute the reconstruction term.
- Compute the KL term in closed form from the mean and log variance.
- Sum them and step, on both encoder and decoder at once.
Log variance is predicted rather than variance, so the output can be any real number and the exponential keeps it positive.
Posterior collapse, and the usual dial
If the decoder is powerful enough to do well without the code, the cheapest way to lower the KL term is to set equal to the prior and ignore entirely.
The model then generates fine and learns nothing in its latent space, which is a silent failure since the loss looks healthy.
Annealing from zero upward during training is the standard defence, and it is why almost every implementation has a warmup on the KL term.
The blurriness of VAE samples is the other known cost. Averaging over the posterior under a Gaussian likelihood is exactly an averaging operation, and averaged images are soft.