Variational Autoencoders: Why the Latent Space Has to Be a Distribution

A plain autoencoder can learn to compress and reconstruct — but its latent space has holes, and interpolating between two points produces noise. The VAE fixes this by encoding distributions rather than points, regularising the latent space to be smooth and continuous. The math behind this is the ELBO — Evidence Lower Bound — and understanding it makes the design choices obvious.

An autoencoder is simple: an encoder compresses input x to a latent vector z, a decoder reconstructs x from z, and training minimises reconstruction loss. The encoder and decoder are neural networks; the bottleneck forces the representation to be compact. Autoencoders work well for compression, anomaly detection, and representation learning.

The problem: the latent space is not smooth. Each training example maps to a specific point z. Points between training examples in the latent space correspond to nothing the decoder has been trained on — they decode to garbage. You cannot generate new samples by sampling random points from the latent space. The space has no meaningful geometry for anything it has not memorised.

The VAE insight: encode distributions, not points

The Variational Autoencoder (Kingma & Welling, 2013) replaces the encoder's single point output with two outputs: a mean vector μ(x) and a log-variance vector log σ^2(x). These parameterise a Gaussian distribution: z ~ N(μ(x), σ^2(x)). During training, z is sampled from this distribution and decoded. The reconstruction is evaluated against the original input.

This forces adjacent regions of the latent space to decode similarly. If the encoder maps a single input to a distribution rather than a point, the decoder must learn to produce good reconstructions from any sample in that distribution — which means nearby z values must produce similar outputs.

The reparameterisation trick

Sampling is not differentiable — you cannot backpropagate through a random sampling operation. The reparameterisation trick rewrites z = μ + σ * ε where ε ~ N(0, 1). ε is sampled independently; the randomness is removed from the computation graph. Gradients can now flow through μ and σ normally.

The ELBO: reconstruction + regularisation

Training a VAE maximises the Evidence Lower Bound (ELBO): ELBO = E[log p(x|z)] - KL(q(z|x) || p(z)). The first term is the reconstruction quality — how well does the decoder reproduce the input from the sampled z? The second term is the KL divergence between the encoder's distribution and the prior p(z) = N(0, 1).

The KL term regularises the latent space: it penalises the encoder for learning a distribution that deviates from N(0, 1). This prevents the encoder from collapsing to very narrow distributions (equivalent to a plain autoencoder) or spreading to arbitrary shapes. The pressure toward N(0, 1) ensures the latent space is filled and continuous — random samples from N(0, 1) decode to meaningful outputs.

The tension: reconstruction vs regularisation

The two terms in the ELBO are in tension. A perfect reconstruction requires precise encoding — map each input to a tight distribution, reducing uncertainty. This pushes toward narrow distributions that violate N(0, 1). Perfect regularisation requires the encoder to map every input to exactly N(0, 1), making the encoding non-informative. Training finds the trade-off: distributions wide enough to regularise, tight enough to reconstruct.

A hyperparameter β (β-VAE) scales the KL term: ELBO = E[log p(x|z)] - β * KL(q(z|x) || p(z)). Larger β enforces more disentanglement — different dimensions of z capture independent generative factors. β-VAE representations are more interpretable (one dimension controls face rotation, another controls smile) at the cost of reconstruction quality.

Generation and interpolation

Once trained, generation is simple: sample z ~ N(0, 1), pass through the decoder. Because the KL term regularises the latent space toward N(0, 1), most samples decode to plausible outputs. Interpolation between two data points x1 and x2: encode both to μ1 and μ2, interpolate z = α*μ1 + (1-α)*μ2, decode for each α. Because the latent space is smooth, interpolated points decode to smooth transitions (e.g., between two faces) rather than noise.

VAEs vs GANs vs Diffusion Models

VAEs produce blurry generations because the reconstruction loss (typically MSE or BCE) averages over possible outputs. GANs produce sharper images by training a discriminator to detect fakes, but training is unstable (mode collapse, GAN training tricks). Diffusion models produce the highest-quality images by learning to denoise across many steps, at the cost of slow sampling. VAEs are valuable less for generation quality and more for structured, interpretable latent spaces useful in downstream tasks.

Try on Colab: train a VAE on MNIST or CelebA. After training, take a 2D slice of the latent space (fix all dimensions except two) and decode a grid of points across that plane — you should see smooth transitions between digit shapes or facial features. Then sample 100 random points from N(0,1) and decode them — compare the output quality to a plain autoencoder trained on the same data.