Variational autoencoders

Problem

A plain autoencoder compresses well but cannot generate: decode a random code and you usually get garbage, because nothing decided what the space between training codes should mean. We want a model we can sample from. Pick a simple distribution p(z) for a hidden code, say a standard normal, and a decoder p_θ(x | z) that turns a code into data. Sampling is then easy: draw z, decode.

Training is the hard part. Maximum likelihood needs p_θ(x) = ∫ p_θ(x | z) p(z) dz, an integral over every possible code, which has no closed form once the decoder is a neural network. The posterior p_θ(z | x), which code explains this x, needs that same integral through Bayes' rule. Kingma and Welling's answer (and, independently, Rezende, Mohamed and Wierstra's) was to learn an encoder that approximates the posterior and to optimize a bound that only needs samples from it.

Intuition

The encoder no longer outputs one code but a small cloud: a mean and a spread. The decoder must reconstruct x from any point in the cloud, so nearby codes are forced to decode to similar data. A second term pulls every cloud toward the standard normal prior, so the clouds of all training examples overlap and tile the space with no gaps. Reconstruction wants the clouds small and far apart; the prior term wants them wide and centred. The balance is a smooth, densely used code space, and then a fresh draw from the prior lands somewhere the decoder knows.

The trick that makes this trainable is to move the randomness out of the way of the gradient: instead of sampling z from N(μ, σ²), sample ε from N(0, 1) and compute z = μ + σε. Now z is a differentiable function of μ and σ, and backpropagation flows straight through.

The noise is an input, not a node to differentiate through, so the sample is a deterministic function of what the encoder outputs.

Mechanics

The ELBO, derived

For any distribution q(z) over codes, multiply and divide inside the integral and apply Jensen's inequality (log of an average is at least the average of the log):

log p(x) = log 𝔼_q(z)[ p(x | z) p(z) / q(z) ] ≥ 𝔼_q(z)[ log p(x | z) ] − KL(q(z) ‖ p(z)) = ELBO

The first term is a reconstruction score: decode codes drawn from q and measure how likely x is. The second keeps q close to the prior, measured by the KL divergence. The bound is not loose by an unknown amount; expanding p(x | z) p(z) = p(z | x) p(x) gives an exact identity:

log p(x) = ELBO(q) + KL(q(z) ‖ p(z | x))

So raising the ELBO over q pushes q toward the true posterior, and the gap is exactly how far it still is. A VAE amortizes this: instead of fitting a separate q for every x, an encoder network outputs q_φ(z | x) = N(μ_φ(x), diag σ_φ²(x)), and the encoder φ and decoder θ are trained together to maximize the average ELBO over the data. Raising it either tightens the bound (better encoder) or raises log p(x) itself (better decoder).

The KL term in closed form

For two one-dimensional normals the KL has a closed form, which is what PyTorch's registered _kl_normal_normal computes:

KL(N(μ_p, σ_p²) ‖ N(μ_q, σ_q²)) = ½ [ (σ_p/σ_q)² + ((μ_p − μ_q)/σ_q)² − 1 − log (σ_p/σ_q)² ]

Against the standard normal prior and with the encoder predicting log σ² (so σ is always positive), it is ½ (σ² + μ² − 1 − log σ²) per dimension, summed over the independent latent dimensions. It is zero only when μ = 0 and σ = 1, that is, when the code says nothing about x. Predicting log σ² rather than σ is the usual choice because it is unconstrained and well-behaved near zero variance.

The reparameterization trick

The reconstruction term 𝔼_q[log p(x | z)] has no closed form for a neural decoder, so it is estimated with samples, usually one per example. We need its gradient with respect to the encoder, but a sample from N(μ, σ²) is not a differentiable function of μ and σ. Writing z = μ + σ ⊙ ε with ε ~ N(0, I) fixes that: ∂z/∂μ = 1 and ∂z/∂σ = ε. PyTorch makes the distinction explicit: Normal.rsample returns loc + eps * scale with a fresh standard normal ε and keeps the graph, while sample runs under torch.no_grad(). has_rsample = True advertises that a distribution supports it. The alternative, the score-function (REINFORCE) estimator f(z) ∇ log q(z), works for any distribution, including discrete ones, but its variance is typically far higher, as the worked example shows.

β-VAE, posterior collapse and latent diffusion

• β-VAE. Higgins et al. weight the KL term by β: loss = reconstruction + β · KL. With β > 1 the code is pushed harder toward the factorized prior, which tends to give more disentangled, individually interpretable latent dimensions at the cost of blurrier reconstructions. β < 1 does the reverse.
• Posterior collapse. If the decoder is powerful enough to model x without the code (an autoregressive decoder over text, say), the cheapest optimum sets q(z | x) = p(z) for every x: KL is zero and z is ignored. Individual dimensions can collapse too; they show up as units with μ ≈ 0, σ ≈ 1 and KL ≈ 0. KL annealing (start β near 0 and raise it) and free bits (no penalty below a per-dimension KL floor) are the usual remedies.
• Blurriness. A Gaussian decoder with fixed variance makes the reconstruction term a squared error, which favours averaged outputs. VAE samples are smooth rather than sharp.
• Latent diffusion. Stable Diffusion and its relatives first train an autoencoder with a small KL penalty (plus perceptual and adversarial losses) to map images to a smaller latent grid, then run diffusion in that latent space. The VAE does the compression; the diffusion model does the generating.

Worked example

The ELBO's gap is the KL between the dashed q and the true posterior; a VAE's encoder is trained to shrink it for every x at once.

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal, Independent, kl_divergence

class VAE(nn.Module):
    def __init__(self, d_in=784, hidden=400, latent=16):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(d_in, hidden), nn.ReLU())
        self.mu = nn.Linear(hidden, latent)
        self.logvar = nn.Linear(hidden, latent)
        self.dec = nn.Sequential(nn.Linear(latent, hidden), nn.ReLU(), nn.Linear(hidden, d_in))

    def forward(self, x):
        h = self.enc(x)
        mu, logvar = self.mu(h), self.logvar(h)
        q = Independent(Normal(mu, torch.exp(0.5 * logvar)), 1)   # event_shape (latent,)
        z = q.rsample()                                           # mu + sigma * eps, keeps the graph
        return self.dec(z), q

def loss_fn(logits, x, q, beta=1.0):
    recon = F.binary_cross_entropy_with_logits(logits, x, reduction="none").sum(-1)  # -log p(x|z)
    prior = Independent(Normal(torch.zeros_like(q.mean), torch.ones_like(q.mean)), 1)
    kl = kl_divergence(q, prior)                                  # closed form, summed over latent dims
    # same as 0.5 * (logvar.exp() + mu**2 - 1 - logvar).sum(-1)
    return (recon + beta * kl).mean()                             # negative ELBO per example

model = VAE()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for x, _ in loader:
    x = x.view(x.size(0), -1)                                     # pixels in [0, 1]
    logits, q = model(x)
    loss = loss_fn(logits, x, q)
    opt.zero_grad()
    loss.backward()
    opt.step()

with torch.no_grad():
    samples = torch.sigmoid(model.dec(torch.randn(64, 16)))       # generate from the prior

Sum the reconstruction over pixels and the KL over latent dimensions before averaging over the batch. Averaging the reconstruction over pixels (the default reduction="mean") silently shrinks it by a factor of 784 relative to the KL, which acts like a huge β and is a common cause of collapse. Independent(..., 1) does that summing for both log_prob and kl_divergence. To report log p(x) rather than the bound, use importance sampling with several draws from q.

Tradeoffs

• Likelihood and a smooth latent space versus sharpness. VAEs give a lower bound on likelihood, stable training and an encoder for free, but samples are blurrier than those of GANs and diffusion models, which is why modern image generators use a VAE only for compression.
• Amortization gap. One encoder serves every x, so even at the optimum its q can be worse than the best Gaussian for a particular x. More expressive posteriors (normalizing flows) or a few refinement steps per example shrink the gap.
• The Gaussian assumption. A diagonal Gaussian q cannot represent correlated or multimodal posteriors; the ELBO then underestimates log p(x) by at least that mismatch.
• β is a real knob. Higher β gives more structured, disentangled codes and worse reconstructions; lower β the reverse. The right setting depends on whether you want samples, features or compression.
• Discrete latents need other tools. Reparameterization needs a continuous z. Discrete codes use the score-function estimator, a continuous relaxation, or vector quantization with a straight-through gradient.

Connections

References

• Kingma and Welling. Auto-Encoding Variational Bayes. ICLR 2014. arXiv:1312.6114
• Rezende, Mohamed and Wierstra. Stochastic Backpropagation and Approximate Inference in Deep Generative Models. ICML 2014. arXiv:1401.4082
• Williams. Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. Machine Learning, 1992.
• Higgins et al. β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR 2017.
• Bowman et al. Generating Sentences from a Continuous Space. CoNLL 2016. arXiv:1511.06349
• Kingma et al. Improved Variational Inference with Inverse Autoregressive Flow. NeurIPS 2016. arXiv:1606.04934
• Cremer, Li and Duvenaud. Inference Suboptimality in Variational Autoencoders. ICML 2018. arXiv:1801.03558
• Kingma and Welling. An Introduction to Variational Autoencoders. Foundations and Trends in Machine Learning, 2019. arXiv:1906.02691
• Rombach, Blattmann, Lorenz, Esser and Ommer. High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022. arXiv:2112.10752
• PyTorch v2.15.0: torch/distributions/normal.py, torch/distributions/kl.py, torch/distributions/independent.py.