Variational Autoencoders

## The two terms A variational autoencoder trains an encoder and a decoder together against a single scalar loss with exactly two pieces: ``` loss(x) = reconstruction(x, decoder(z)) + KL( q(z|x) || p(z) ), z ~ q(z|x) ``` **Reconstruction.** The decoder maps a latent code `z` back to the data space and is scored by how much probability it puts on the actual input. The scoring function is fixed by the decoder's assumed output distribution: a Gaussian decoder with constant variance gives summed squared error, a Bernoulli decoder over binary pixels gives summed binary cross-entropy, a categorical decoder over tokens gives token-level cross-entropy. This term is an *expectation* over `z` drawn from the encoder's distribution, which in practice is estimated with a single draw per data point per step — noisy, but unbiased, and minibatch averaging smooths it out. **KL to the prior.** The encoder does not emit a single point. It emits the parameters of a distribution `q(z|x)` over codes — for the standard formulation, a diagonal Gaussian with a mean vector and a per-dimension variance. The second term measures how far that per-input Gaussian is from a fixed prior `p(z) = N(0, I)`. Because both sides are diagonal Gaussians, this term is available in closed form: ``` KL = 0.5 * sum_j ( mu_j^2 + sigma_j^2 - 1 - log sigma_j^2 ) ``` summed over latent dimensions, in nats. Nothing is sampled to compute it. It is zero exactly when every `mu_j` is 0 and every `sigma_j` is 1 — that is, when the encoder has stopped saying anything input-specific. ## Why both are needed The reconstruction term alone would push each input toward its own distant, tightly concentrated code: the easiest way to rebuild an input is to memorise a private address for it. That gives an accurate rebuild but a latent space with no shared coordinate system — you cannot pick a fresh `z` and expect the decoder to produce anything sensible. The KL term alone is trivially satisfied by ignoring the input: emit the prior for every `x` and pay nothing. That gives a perfectly samplable latent that carries no information. Training lands between the two. The KL term is the price, in nats, of every bit of input-specific information the encoder pushes through the code; the reconstruction term is what that information buys. This is why the KL value is often read as the *rate* of the code and the reconstruction error as the *distortion*: the loss is one point on a rate–distortion trade-off, not an absolute measure of quality. ## Reading the numbers during training Log the two terms separately, never just their sum. A falling total tells you nothing about which side is winning. Useful readings: - **KL near zero, reconstruction stuck high.** The code is carrying no information; the decoder is producing an input-independent average. - **KL very large, reconstruction excellent.** The per-input Gaussians are narrow and far apart. Reconstructions look fine, but codes drawn from the prior land in regions no training input ever occupied, so *generated* samples are poor even though the loss looks good. - **Per-dimension KL.** Report the KL split across latent dimensions. Dimensions sitting at zero are dead — the model has switched them off and is using a smaller effective code than you provisioned. ## Common misreadings The KL term is often described as "regularisation on the weights". It is not: it constrains the *output distribution of the encoder*, dimension by dimension, and is unrelated to weight decay. It also does not measure any distance between images — both of its arguments are distributions over latent codes. A second misreading is that a lower total loss always means better generation. The two terms are in different units of the same nat scale, but they are traded against one another, and the trade-off that minimises the sum is not necessarily the one that makes prior samples decode well. If your goal is generation, evaluate by decoding codes drawn from the prior, not by staring at the loss curve. A third is that the prior must be learned. In the standard formulation it is fixed at `N(0, I)` and never trained; the whole burden of matching it falls on the encoder. Learned or hierarchical priors exist as extensions, but the plain VAE's prior is a constant. ## Scaling caveat Whether the reconstruction term is *summed* over pixels or *averaged* over them silently rescales it relative to the KL, sometimes by thousands. Two implementations of "the same" VAE with different reduction conventions sit at completely different points on the trade-off. Fix the convention, and report the KL in nats per data point so the numbers mean something across runs.