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.