A variational autoencoder (VAE) is a type of artificial neural network used for generative modeling, introduced by Diederik P. Kingma and Max Welling in 2013. It belongs to the family of probabilistic graphical models and variational Bayesian methods. A VAE learns to encode high-dimensional data (such as images) into a low-dimensional latent space and then decode samples from that latent space back into data, effectively learning the underlying probability distribution of the training data.
Overview of architecture and operation
A VAE consists of two main neural network components: an encoder and a decoder. The encoder maps each input data point \(x\) to a distribution over latent variables \(z\), typically a multivariate Gaussian with mean \(\mu(x)\) and variance \(\sigma^2(x)\). Instead of encoding a single point, the encoder outputs parameters for a probability distribution, which allows the model to represent uncertainty and capture continuous variations in the data. The decoder takes a sample \(z\) from this latent distribution and maps it back to the original data space, producing a reconstruction \(\hat{x}\).
The training objective of a VAE is to maximize the evidence lower bound (ELBO) on the log-likelihood of the data. This involves two terms: the reconstruction loss, which measures how well the decoder reconstructs the input from the latent sample, and the Kullback–Leibler divergence (KL divergence) between the encoder's output distribution and a prior distribution over \(z\) (usually a standard normal). The KL term acts as a regularizer, encouraging the latent space to be smooth and continuous.
To enable backpropagation through the random sampling step, the VAE uses the reparameterization trick: instead of sampling \(z\) directly from \(q_\phi(z|x)\), the model samples an auxiliary noise variable \(\epsilon \sim \mathcal{N}(0, I)\) and computes \(z = \mu(x) + \sigma(x) \odot \epsilon\). This makes the sampling operation differentiable with respect to the encoder parameters.
Mathematical formulation
From a probabilistic perspective, the goal is to maximize the marginal likelihood of the observed data \(p_\theta(x) = \int p_\theta(x|z) p(z) dz\). However, this integral is often intractable. The VAE introduces an approximate posterior \(q_\phi(z|x)\) and derives the ELBO:
\[
\log p_\theta(x) \geq \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) \| p(z))
\]
The first term is the reconstruction likelihood, and the second term is the KL divergence. Both terms are differentiable with respect to the parameters \(\theta\) and \(\phi\) when using the reparameterization trick.
The choice of prior \(p(z)\) and likelihood \(p_\theta(x|z)\) depends on the data type. For continuous data, a Gaussian likelihood is common; for binary data (e.g., binarized MNIST), a Bernoulli likelihood is used. The latent prior is typically a standard normal distribution.
Training dynamics
The VAE is trained by stochastic gradient descent on the ELBO. The reconstruction loss encourages the decoder to produce outputs close to the inputs, while the KL divergence pushes the encoder's posterior toward the prior. This creates a trade-off: too much weight on reconstruction leads to overfitting and a fragmented latent space, while too much weight on KL leads to a phenomenon called posterior collapse, where the latent variable becomes uninformative and the model ignores \(z\). Various techniques, such as KL annealing or beta-VAE, adjust this balance.
Variants and extensions
Several variants of the VAE have been developed. The beta-VAE introduces a hyperparameter \(\beta\) to weight the KL term, enabling more disentangled representations. Conditional VAEs (CVAEs) incorporate class labels or other conditioning information into both encoder and decoder. Vector-quantized VAEs (VQ-VAEs) use discrete latent codes, which are particularly effective for high-fidelity image generation. Other extensions include hierarchical VAEs, normalizing flows, and adversarial autoencoders.
Applications
The VAE has been widely applied in various domains. In computer vision, it is used for image generation, denoising, and inpainting. In natural language processing, VAEs have been applied to text generation and sentence representation learning. In medical imaging, they are used for anomaly detection. VAEs also serve as building blocks for more complex generative models, such as those used in drug discovery and molecular design.
Relationship to other generative models
VAEs are one of the three main types of deep generative models, alongside generative adversarial networks (GANs) and autoregressive models. Unlike GANs, which train a discriminator to distinguish real from fake samples, VAEs optimize a likelihood-based objective, making them more stable to train but often producing blurrier outputs. Autoregressive models, such as PixelCNN, model the joint distribution as a product of conditionals, but they lack a low-dimensional latent space. VAEs offer a balance between tractable inference and expressive generation.
Limitations and current research
Despite their success, VAEs have known limitations. The generated samples are often less sharp than those from GANs due to the Gaussian likelihood assumption. The latent space may not be fully disentangled, and posterior collapse remains an active research area. Recent work focuses on improving the expressiveness of the prior and posterior, using more flexible distributions, and combining VAEs with other models. The field continues to evolve, with applications expanding to reinforcement learning and representation learning.
See also
- Autoencoder
- Generative adversarial network
- Bayesian inference
- Latent variable model
- Reparameterization trick
- Evidence lower bound
- Beta-VAE
- Conditional variational autoencoder
- Vector-quantized variational autoencoder
- Normalizing flows
- Probabilistic graphical model
- Deep learning
- Unsupervised learning
- Stochastic gradient descent
- Kullback–Leibler divergence
- Gaussian distribution
- Bernoulli distribution
- MNIST database
- Image generation
- Representation learning
References
- Kingma, D. P., & Welling, M. (2013). Auto-Encoding Variational Bayes. arXiv:1312.6114.
- Doersch, C. (2016). Tutorial on Variational Autoencoders. arXiv:1606.05908.
- Blei, D. M., Kucukelbir, A., & McAuliffe, J. D. (2017). Variational Inference: A Review for Statisticians. Journal of the American Statistical Association.
- Higgins, I., et al. (2017). beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR.