A variational autoencoder (VAE) is an artificial neural network architecture introduced by Diederik P. Kingma and Max Welling in 2013. It is part of the families of probabilistic graphical models and variational Bayesian methods. In addition to being seen as an autoencoder neural network architecture, variational autoencoders can also be studied within the mathematical formulation of variational Bayesian methods, connecting a neural encoder network to its decoder through a probabilistic latent space (for example, as a multivariate Gaussian distribution) that corresponds to the parameters of a variational distribution.
The encoder maps each point (such as an image) from a large complex dataset into a distribution within the latent space, rather than to a single point in that space. The decoder has the opposite function, which is to map from the latent space to the input space, again according to a distribution (although in practice, noise is rarely added during the decoding stage). By mapping a point to a distribution instead of a single point, the network can avoid overfitting the training data. Both networks are typically trained together with the usage of the reparameterization trick, although the variance of the noise model can be learned separately. Although this type of model was initially designed for unsupervised learning, its effectiveness has been proven for semi-supervised learning and supervised learning.
Overview of architecture and operation
A variational autoencoder is a generative model with a prior and noise distribution respectively. Usually such models are trained using the expectation-maximization meta-algorithm (e.g., probabilistic PCA, sparse coding). Such a scheme optimizes a lower bound of the data likelihood, which is usually computationally intractable, and in doing so requires the discovery of q-distributions, or variational posteriors. These q-distributions are normally parameterized for each individual data point in a separate optimization process. However, variational autoencoders use a neural network as an amortized approach to jointly optimize across data points. In that way, the same parameters are reused for multiple data points, which can result in massive memory savings. The first neural network takes as input the data points themselves, and outputs parameters for the variational distribution. As it maps from a known input space to the low-dimensional latent space, it is called the encoder.
The decoder is the second neural network of this model. It is a function that maps from the latent space to the input space, e.g., as the means of the noise distribution. It is possible to use another neural network that maps to the variance, however this can be omitted for simplicity. In such a case, the variance can be optimized with gradient descent.
To optimize this model, one needs to know two terms: the "reconstruction error", and the Kullback–Leibler divergence (KL-D). Both terms are derived from the free energy expression of the probabilistic model, and therefore differ depending on the noise distribution and the assumed prior of the data, here referred to as p-distribution. For example, a standard VAE task such as IMAGENET is typically assumed to have a gaussianly distributed noise; however, tasks such as binarized MNIST require a Bernoulli noise. The KL-D from the free energy expression maximizes the probability mass of the q-distribution that overlaps with the p-distribution, which unfortunately can result in mode-seeking behaviour. The "reconstruction" term is the remainder of the free energy expression, and requires a sampling approximation to compute its expectation value.
More recent approaches replace Kullback–Leibler divergence (KL-D) with various statistical distances, see "Statistical distance VAE variants" below.
Formulation
From the point of view of probabilistic modeling, one wants to maximize the likelihood of the data \(x\) by their chosen parameterized probability distribution \(p_\theta(x)=p(x|\theta)\). This distribution is usually chosen to be a Gaussian \(N(x|\mu,\sigma)\) which is parameterized by \(\mu\) and \(\sigma\) respectively, and as a member of the exponential family it is easy to work with as a noise distribution. Simple distributions are easy enough to maximize, however distributions where a prior is assumed over the latents \(z\) results in intractable integrals. Let us find \(p_\theta(x)\) via marginalizing over \(z\).
\(p_\theta(x) = \int_z p_\theta({x,z}) \, dz,\)
where \(p_\theta({x,z})\) represents the joint distribution under \(p_\theta\) of the observable data \(x\) and its latent representation or encoding \(z\). According to the chain rule, the equation can be rewritten as
\(p_\theta(x) = \int_z p_\theta({x|z}) p_\theta(z) \, dz.\)
In practice, the integral over \(z\) is intractable because the latent space is high-dimensional and the true posterior \(p_\theta(z|x)\) is unknown. The VAE introduces an encoder network \(q_\phi(z|x)\) to approximate the true posterior, and a decoder network \(p_\theta(x|z)\) to reconstruct the data. The training objective is the evidence lower bound (ELBO), which is the sum of the reconstruction term and the KL divergence term. The reparameterization trick allows backpropagation through the sampling process by expressing \(z\) as \(z = \mu + \sigma \odot \epsilon\), where \(\epsilon\) is sampled from a standard normal distribution.
Training and optimization
The VAE is trained by maximizing the ELBO with respect to the parameters \(\theta\) and \(\phi\). The reconstruction term encourages the decoder to accurately reconstruct the input data, while the KL divergence term regularizes the latent space to be close to the prior (typically a standard normal distribution). This balance helps prevent overfitting and enables the generation of new data points by sampling from the prior and decoding.
The reparameterization trick is crucial for efficient training using stochastic gradient descent. By moving the stochasticity to the noise variable \(\epsilon\), the gradients can flow through the deterministic parts of the network. The variance of the noise model can be learned separately, and the reconstruction error is often computed using mean squared error for continuous data or binary cross-entropy for binary data.
Applications and impact
The VAE has become a foundational model in Generative AI, enabling tasks such as image generation, anomaly detection, and representation learning. It has been extended to various domains, including Deep learning and Machine learning, and has influenced the development of other generative models like Large language models, though those typically use different architectures. VAEs are also used in semi-supervised and supervised learning settings, demonstrating their versatility beyond unsupervised learning.
The original paper by Kingma and Welling has been highly influential, and the VAE framework has been adapted in numerous ways, such as using different statistical distances instead of KL divergence to address mode collapse. The model's ability to learn smooth latent spaces has made it useful for interpolation and manipulation of data attributes.
Variants and extensions
Several variants of VAEs have been proposed to improve their performance. For example, the beta-VAE introduces a weighting factor on the KL divergence term to encourage more disentangled representations. Other variants use different priors, such as Gaussian mixture models, or employ adversarial training to improve the realism of generated samples. The reparameterization trick has also been generalized to other distributions, such as categorical variables, using the Gumbel-softmax approximation.
In the context of Artificial intelligence, VAEs are often compared with generative adversarial networks (GANs), which are another class of generative models. While GANs can produce sharper images, VAEs tend to be more stable to train and provide a probabilistic framework that is useful for uncertainty estimation. As of the early 2020s, VAEs remain an active area of research, with applications in Neural network design and beyond.