Auto-Encoding Variational Bayes
Contributions
- Variational Autoencoder (VAE): introduces a deep generative model that combines a probabilistic encoder (recognition network) with a probabilistic decoder (generative network), jointly trained by maximizing a variational lower bound on the data log-likelihood.
- Reparameterization trick: solves the problem of backpropagating through a stochastic sampling operation by rewriting z ~ q(z|x) as z = μ + σ·ε with ε ~ N(0,I), enabling gradient-based optimization of the encoder parameters.
- Amortized variational inference: uses a neural network to map each data point to its approximate posterior parameters (μ, σ), amortizing the cost of inference across data points rather than optimizing per-data-point variational parameters separately.
- Scalable generative modeling: demonstrates a scalable algorithm that handles large datasets and continuous latent spaces, bridging deep learning and Bayesian latent variable models.
Methodology
Generative Model
The VAE defines a generative process:
- Sample latent variable: z ~ p(z) = N(0, I)
- Generate data: x ~ p_θ(x|z) parameterized by a decoder neural network
The marginal likelihood p_θ(x) = ∫ p_θ(x|z) p(z) dz is intractable due to the integral over z.
Evidence Lower Bound (ELBO)
The true posterior p_θ(z|x) is approximated by a recognition network q_φ(z|x) = N(μ_φ(x), diag(σ²_φ(x))). The training objective is the ELBO:
L(θ, φ; x) = E_{q_φ(z|x)}[log p_θ(x|z)] - KL(q_φ(z|x) || p(z))
- The first term is the reconstruction loss: how well the decoder reconstructs x from sampled z.
- The second term is the KL divergence: a regularizer encouraging the approximate posterior to stay close to the prior N(0,I). For Gaussian q and p, this KL has a closed-form expression.
Reparameterization Trick
The key technical challenge is that the expectation E_{q_φ(z|x)}[·] depends on φ through the sampling distribution. Backpropagation through sampling is not directly possible. The reparameterization trick rewrites:
z = μ_φ(x) + σ_φ(x) ⊙ ε, ε ~ N(0, I)
Now z is a deterministic function of φ and ε; gradients flow through μ and σ as usual, and ε is treated as fixed noise.
Architecture
The encoder and decoder are multi-layer perceptrons (MLPs) or convolutional networks. For image experiments, the encoder maps an image to (μ, log σ²) vectors; the decoder maps a sampled z back to a reconstructed image. The latent space is typically 2–20 dimensional in the original paper.
Results
Generative Modeling (MNIST, Frey Face)
- The VAE learns a smooth, interpolable latent space on MNIST: latent codes of different digits are organized meaningfully and interpolation between codes produces coherent digit transitions.
- On Frey Face dataset, the model captures pose and expression variation in a 2D latent space.
- Quantitative comparison via estimated marginal log-likelihood: VAE outperforms wake-sleep algorithm on MNIST (−86.47 vs. −89.0 nats).
Latent Space Structure
- Low-dimensional (2D) VAE on MNIST organizes digit classes into distinct clusters without class supervision, demonstrating that the latent space captures semantically meaningful variation.
Limitations
- The Gaussian prior and posterior are restrictive; more complex posteriors (normalizing flows, hierarchical VAEs) are needed for high-fidelity image generation.
- VAE-generated images tend to be blurry because the reconstruction loss (pixel-wise MSE or BCE) averages over all possible reconstructions of a given z.
- The ELBO is a lower bound, not the true log-likelihood; the tightness depends on the expressiveness of the recognition network.
- Later generative models (GANs, diffusion models) surpass VAEs in perceptual image quality, though VAEs remain competitive for representation learning and structured latent spaces.