Variational Autoencoders

Generative model, ELBO, reparameterisation trick, posterior collapse, β-VAE, representation learning

Standard autoencoders compress data into a latent code and reconstruct it — but the latent space is irregular. Interpolating between two codes often passes through empty regions that decode into nonsense, because there is no constraint on what the latent space looks like globally. VAEs solve this by placing a probabilistic structure on the latent space: instead of mapping each input to a single point, the encoder maps it to a distribution q(z|x), and the KL regularisation term forces these distributions to stay close to a standard Gaussian prior.

The result is a continuous, densely populated latent space where interpolation makes sense and unseen points decode coherently. The reparameterisation trick — writing z = μ + σ·ε where ε ~ N(0,I) — is what makes gradients flow through the sampling step. The central failure mode is posterior collapse: an expressive decoder learns to model p(x) without using z at all, the encoder degenerates to the prior, KL drops to zero, and you have trained a very expensive unconditional generator.

Key points

• The generative model: p_θ(x,z) = p_θ(x|z)p(z) where p(z) = N(0,I). The decoder p_θ(x|z) maps a latent code to a distribution over x. The problem is fitting this model: to learn θ via maximum likelihood you need p_θ(x) = ∫ p_θ(x|z)p(z)dz, which requires marginalising over all possible latent codes — intractable. VAEs avoid this by optimising a lower bound (the ELBO) instead.
• The encoder q_φ(z|x) = N(μ_φ(x), diag(σ_φ(x)²)) is a neural network mapping input x to the parameters of an approximate posterior over z. This is amortised VI: rather than running a separate optimisation per input to find q(z|x), one encoder network handles all inputs at once. The cost is that the encoder is an approximation — it learns the best single function from inputs to posteriors, not the exact posterior for each input.
• VAE ELBO: L(θ,φ;x) = E_{q_φ(z|x)}[log p_θ(x|z)] - KL[q_φ(z|x) ‖ p(z)]. The reconstruction term rewards the decoder for explaining the data given latent codes sampled from the encoder. The KL term penalises the encoder for drifting from the prior N(0,I). For diagonal Gaussian q, the KL is closed form: -½ Σⱼ(1 + log σⱼ² - μⱼ² - σⱼ²) — so the only stochastic step that requires a gradient estimator is the expectation over q in the reconstruction term.
• The reparameterisation trick: you cannot backpropagate through z ~ N(μ, σ²) because the sampling step is stochastic and has no gradient. Fix: write z = μ_φ(x) + σ_φ(x)⊙ε where ε ~ N(0,I). Now z is a deterministic function of the encoder parameters and a fixed noise draw. Gradients flow through μ and σ back to φ. Without this trick, VAE training requires high-variance REINFORCE-style gradient estimates — end-to-end training becomes impractical. The trick breaks down precisely when z is discrete (Bernoulli, categorical), where no differentiable reparameterisation exists.
• Posterior collapse is the central VAE failure mode. When the decoder is sufficiently expressive (PixelCNN, autoregressive Transformer), it can model p(x) without using z at all. The encoder then learns q(z|x) ≈ N(0,I) — identical to the prior regardless of x. The KL term drops to near zero, reconstruction loss stays low, and training happily converges to a model where the latent space carries no information. Symptom: KL ≈ 0 after training. Cause: decoder power exceeds the bottleneck created by the KL penalty.
• Fixes for posterior collapse: KL annealing — start with β=0 (pure reconstruction), linearly ramp β to 1 over the first 30% of training. This forces the decoder to first learn to use z before the KL regularisation becomes active. Free bits — floor the KL per latent dimension at δ bits, so the optimiser cannot collapse dimensions to zero without incurring a penalty. Both interventions make the decoder see informative z before it has a chance to learn to ignore z.
• β-VAE multiplies the KL term by β: L = E[log p(x|z)] - β·KL[q(z|x) ‖ p(z)]. β > 1 over-penalises KL, forcing the encoder to compress information into fewer, more independent latent dimensions — each dimension learns to control one factor of variation. β < 1 relaxes regularisation to combat posterior collapse. The tradeoff is explicit: higher β gives better disentanglement and worse reconstruction quality. β = 1 is the standard VAE.
• Standard autoencoders have irregular latent spaces — interpolating between two encoded points often passes through low-density regions that decode into garbage. VAE latent spaces avoid this because the KL regularisation forces encoder outputs to stay close to N(0,I), which is dense everywhere. Spherical interpolation between two VAE codes z₁ and z₂ produces semantically coherent intermediate samples because the path stays in the high-density region of the prior.
• Production uses for VAEs: anomaly detection (low ELBO = poor reconstruction or high KL → flag as out-of-distribution), data imputation (infer z from observed dimensions, decode to fill missing values), molecule generation (VAE latent spaces over molecular graphs enable gradient-based optimisation of chemical properties). VAE outputs are blurrier than GAN or diffusion outputs — this is a direct mathematical consequence of optimising expected MSE under an approximate posterior, which averages over plausible reconstructions rather than sampling one.

Recap

• VAE = probabilistic latent space: encoder maps x to $q(z|x)$; KL to $N(0,I)$ keeps it dense so interpolation stays coherent.
• ELBO = reconstruction − KL[q(z|x) ‖ p(z)] — maximise a lower bound because the true likelihood needs the intractable $int p(x|z)p(z)dz$.
• Reparameterisation trick: $z = μ + σ⊙ε$, $ε sim N(0,I)$ — makes sampling differentiable so gradients flow to φ.
• Trick breaks for discrete z (Bernoulli, categorical) — no differentiable reparameterisation; use Gumbel-Softmax.
• Posterior collapse (KL → 0): expressive decoder models p(x) ignoring z; encoder degenerates to the prior.
• Fixes: KL annealing (β: 0 → 1) and free bits — force the decoder to use z before the KL penalty bites.
• β-VAE: β > 1 → disentanglement + worse reconstruction; VAE blurriness is intrinsic (expected MSE averages plausible reconstructions).

Check your understanding