Variational Autoencoder (Details)

A variational autoencoder (VAE) is a neural network architecture introduced in 2013 by Diederik P. Kingma and Max Welling that learns a probabilistic latent space for generative modeling, combining an encoder and decoder trained with variational Bayesian methods.

A variational autoencoder (VAE) is an artificial neural network architecture introduced by Diederik P. Kingma and Max Welling in 2013. It belongs to the families of probabilistic graphical models and variational Bayesian methods. VAEs are generative models that learn to represent data in a compressed latent space and generate new samples from that space. Unlike standard autoencoders, which map inputs to fixed points in a latent space, VAEs map inputs to a probability distribution, typically a multivariate Gaussian, which helps avoid overfitting and enables smooth interpolation and generation.

The architecture consists of two neural networks: an encoder and a decoder. The encoder maps each input data point (e.g., an image) to parameters of a variational distribution in the latent space, such as a mean and variance vector. The decoder maps from the latent space back to the input space, producing reconstructions or new samples. Both networks are trained jointly using the reparameterization trick, which allows backpropagation through stochastic sampling. VAEs were initially designed for unsupervised learning, but they have also proven effective for semi-supervised and supervised learning tasks.

Overview of architecture and operation

A variational autoencoder is a generative model with a prior distribution over latent variables and a noise distribution for the data. Traditional generative models, such as probabilistic PCA or sparse coding, are trained using the expectation-maximization (EM) meta-algorithm. EM optimizes a lower bound on the data likelihood, which is often computationally intractable, and requires discovering q-distributions (variational posteriors) for each data point in a separate optimization process. In contrast, VAEs use a neural network as an amortized approach to jointly optimize across all data points. This means the same network parameters are reused for multiple data points, resulting in significant memory savings.

The encoder network takes input data points and outputs parameters for the variational distribution. Since it maps from a known input space to a low-dimensional latent space, it is called the encoder. The decoder is the second neural network; it maps from the latent space to the input space, typically outputting the means of the noise distribution. While it is possible to use another network to output variance, this is often omitted for simplicity, and the variance can be optimized via gradient descent.

To optimize the model, two terms are needed: the reconstruction error and the Kullback–Leibler divergence (KL-D). Both terms are derived from the free energy expression of the probabilistic model, and they differ depending on the noise distribution and the assumed prior of the data (p-distribution). For example, a standard VAE task such as ImageNet typically assumes Gaussian noise, while tasks like binarized MNIST require Bernoulli noise. The KL-D term maximizes the probability mass of the q-distribution that overlaps with the p-distribution, which can lead to mode-seeking behavior. 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 KL-D with various statistical distances, as discussed in the section on statistical distance VAE variants.

Formulation

From a probabilistic modeling perspective, the goal is to maximize the likelihood of the data x under a chosen parameterized probability distribution p_θ(x) = p(x|θ). This distribution is often chosen to be a Gaussian N(x|μ, σ), parameterized by mean μ and variance σ, which is easy to work with as a member of the exponential family. Simple distributions are straightforward to maximize, but when a prior is assumed over latent variables z, the resulting integrals become intractable. To find p_θ(x), one marginalizes over z:

p_θ(x) = ∫ p_θ(x, z) dz,

where p_θ(x, z) is the joint distribution of observable data x and latent encoding z under p_θ. Using the chain rule, this can be rewritten as:

p_θ(x) = ∫ p_θ(x|z) p_θ(z) dz.

This integral is generally intractable because the latent space is high-dimensional and the likelihood p_θ(x|z) can be complex. Variational inference approximates the true posterior p_θ(z|x) with a simpler distribution q_φ(z|x), parameterized by the encoder network. The evidence lower bound (ELBO) is derived to optimize the model:

log p_θ(x) ≥ E_{q_φ(z|x)}[log p_θ(x|z)] - KL(q_φ(z|x) || p_θ(z)).

The first term is the reconstruction likelihood, and the second term is the KL divergence between the approximate posterior and the prior. The reparameterization trick allows sampling from q_φ(z|x) by expressing z as μ + σ * ε, where ε is drawn from a standard normal distribution, enabling gradient-based optimization.

Encoder and decoder networks

The encoder network, often denoted as q_φ(z|x), takes an input x and outputs the parameters of a variational distribution, typically the mean μ and log-variance log σ² of a Gaussian. This distribution represents the latent encoding of the input. The decoder network, denoted as p_θ(x|z), takes a latent sample z and outputs parameters for the data distribution, such as the mean of a Gaussian for continuous data or the probabilities for Bernoulli data.

Both networks are typically trained together using the reparameterization trick. The encoder and decoder are usually implemented as multilayer perceptrons (MLPs) or convolutional neural networks (CNNs), depending on the data type. For image data, CNNs are common, as seen in models like the convolutional VAE. The choice of network architecture affects the model's capacity and performance.

Training and the reparameterization trick

Training a VAE involves maximizing the ELBO, which is the sum of the reconstruction term and the KL divergence term. The reconstruction term encourages the decoder to accurately reconstruct the input from the latent sample, while the KL divergence term encourages the approximate posterior to be close to the prior, typically a standard normal distribution N(0, I). This balance prevents overfitting and ensures a smooth latent space.

The reparameterization trick is crucial for training because it allows gradients to flow through the stochastic sampling process. Instead of sampling z directly from q_φ(z|x), which is non-differentiable, the encoder outputs μ and σ, and z is computed as z = μ + σ * ε, where ε ~ N(0, I). This makes the sampling operation differentiable with respect to the parameters, enabling standard backpropagation.

During training, the variance of the noise model can be learned separately or fixed. In some implementations, the decoder outputs both mean and variance, while in others, the variance is a hyperparameter or optimized separately. The training process typically uses stochastic gradient descent (SGD) or variants like Adam, with learning rate schedules.

Loss function and KL divergence

The loss function for a VAE is the negative ELBO, which consists of two terms: the reconstruction loss and the KL divergence. The reconstruction loss measures how well the decoder reconstructs the input from the latent sample. For continuous data with Gaussian noise, this is often the mean squared error (MSE) between the input and the reconstructed output. For binary data, it is the binary cross-entropy.

The KL divergence term is computed between the approximate posterior q_φ(z|x) and the prior p(z). For Gaussian distributions, this has a closed-form expression:

KL(q_φ(z|x) || p(z)) = -0.5 * (1 + log σ² - μ² - σ²),

assuming the prior is N(0, I). This term penalizes the latent distribution for deviating from the prior, encouraging a compact and continuous latent space.

The KL divergence can lead to mode-seeking behavior, where the model focuses on a few modes of the data distribution. To address this, various alternatives have been proposed, such as using different statistical distances like the Wasserstein distance or maximum mean discrepancy (MMD).

Variants and extensions

Since the introduction of VAEs, numerous variants have been developed to improve their performance and address limitations. Some notable variants include:

β-VAE: Introduces a weighting factor β on the KL divergence term to encourage more disentangled representations. This is often used in unsupervised learning of interpretable factors.
Conditional VAE (CVAE): Conditions both the encoder and decoder on additional information, such as class labels, enabling controlled generation.
VQ-VAE: Uses vector quantization in the latent space, producing discrete latent codes, which are useful for tasks like image generation and representation learning.
Statistical distance VAE variants: Replace the KL divergence with other statistical distances, such as the Wasserstein distance, to improve training stability and sample quality.

These variants have been applied in various domains, including image generation, text generation, and anomaly detection.

Applications and significance

Variational autoencoders have a wide range of applications in machine learning and artificial intelligence. They are used for generative modeling, where they can produce new data samples similar to the training data, such as images, audio, and text. VAEs are also used for representation learning, where the latent space captures meaningful features of the data, enabling tasks like clustering and dimensionality reduction.

In semi-supervised learning, VAEs can leverage unlabeled data to learn useful representations, improving performance on tasks with limited labeled data. In supervised learning, VAEs can be used as feature extractors or for data augmentation. VAEs have also been applied in anomaly detection, where the reconstruction error or likelihood of a data point indicates whether it is normal or anomalous.

The significance of VAEs lies in their ability to learn probabilistic latent representations in an unsupervised manner, providing a principled framework for generative modeling. They have influenced many subsequent developments in deep learning, including the broader field of generative AI, and continue to be an active area of research.

Relationship to other generative models

VAEs are one of several prominent generative model families, alongside generative adversarial networks (GANs) and diffusion models. While GANs focus on generating realistic samples through an adversarial training process, VAEs emphasize probabilistic inference and a smooth latent space. Diffusion models, which have gained popularity recently, generate data by reversing a noising process. Each approach has its strengths and weaknesses; VAEs offer stable training and a well-defined latent space, but may produce blurrier samples compared to GANs. The choice of model depends on the specific application and requirements.

In the context of deep learning, VAEs are often used as building blocks in larger architectures, such as in variational autoencoder-based transformers or as components in multi-modal models. They are also relevant to the development of large language models, though those typically rely on autoregressive or transformer architectures rather than VAEs.