Variational Autoencoders

Variational Autoencoders (VAEs) represent one of the most elegant intersections of probabilistic reasoning and deep learning. Introduced by Kingma and Welling in 2013, VAEs solved a fundamental problem: how do we learn meaningful latent representations while simultaneously building generative models that can synthesize new data?

Before VAEs, autoencoders could compress data but lacked probabilistic structure. Generative models like GMMs could model distributions but struggled with high-dimensional data like images. VAEs unified these worlds—creating a principled framework where neural networks learn to encode observations into structured latent spaces, and then decode from those spaces to generate new samples.

This page develops the complete mathematical formulation of VAEs from first principles. We'll see how the Evidence Lower Bound (ELBO) emerges naturally from variational inference, why it provides a tractable objective for deep learning, and how this formulation enables both generation and representation learning.

At its core, generative modeling asks: given a dataset of observations ${x^{(1)}, x^{(2)}, ..., x^{(N)}}$, can we learn the underlying distribution $p(x)$ that generated them?

If we could model $p(x)$ accurately, we could:

1. Generate new samples: Draw $x_{\text{new}} \sim p(x)$ to create realistic synthetic data
2. Evaluate likelihood: Compute $p(x_{\text{query}})$ for anomaly detection or density estimation
3. Learn representations: Understand the factors of variation that explain the data

The challenge is that real-world data—images, audio, text—lives in extremely high-dimensional spaces with complex dependencies. A 256×256 color image has 196,608 dimensions. Modeling the joint distribution over all pixels directly is intractable.

The Latent Variable Model

VAEs embrace this intuition by introducing latent variables $z$ that capture the underlying factors of variation. The generative process becomes:

1. Sample latent code: $z \sim p(z)$ (typically a simple prior like $\mathcal{N}(0, I)$)
2. Generate observation: $x \sim p_\theta(x|z)$ (a complex conditional distribution)

This defines the generative model:

$$p_\theta(x) = \int p(z) \cdot p_\theta(x|z) , dz$$

Here:

• $p(z)$ is the prior over latent variables—our belief about $z$ before seeing data
• $p_\theta(x|z)$ is the likelihood or decoder—how $z$ generates $x$
• $\theta$ are learnable parameters (e.g., neural network weights)

The elegance of this setup: we use a simple prior $p(z)$ (easy to sample) combined with a powerful, learnable decoder $p_\theta(x|z)$ (neural network) to model arbitrarily complex $p_\theta(x)$.

Given a dataset $\mathcal{D} = {x^{(1)}, ..., x^{(N)}}$, the natural objective is maximum likelihood estimation:

$$\theta^* = \arg\max_\theta \sum_{i=1}^{N} \log p_\theta(x^{(i)})$$

Maximizing log-likelihood finds parameters $\theta$ that make the observed data most probable under our model. This is principled, asymptotically consistent, and well-understood.

The problem: Computing $\log p_\theta(x)$ requires:

$$\log p_\theta(x) = \log \int p(z) \cdot p_\theta(x|z) , dz$$

This integral has no closed form when $p_\theta(x|z)$ is a deep neural network. We cannot evaluate the objective, let alone optimize it.

The Posterior Problem

Suppose we had access to the true posterior $p_\theta(z|x)$—the distribution over latent codes that could have generated observation $x$. Then we could:

1. For each $x$, sample $z \sim p_\theta(z|x)$ (codes likely to explain $x$)
2. These importance-weighted samples would give low-variance estimates

But the posterior is also intractable:

$$p_\theta(z|x) = \frac{p(z) \cdot p_\theta(x|z)}{p_\theta(x)} = \frac{p(z) \cdot p_\theta(x|z)}{\int p(z') \cdot p_\theta(x|z') , dz'}$$

The denominator is exactly the intractable marginal likelihood. We're stuck in a loop: we need the posterior to estimate the marginal, but we need the marginal to compute the posterior.

This is precisely where variational inference comes to the rescue.

Variational inference sidesteps intractable posteriors by introducing a tractable approximation. Since we cannot compute $p_\theta(z|x)$, we'll learn a simpler distribution $q_\phi(z|x)$ that approximates it.

This inference network or encoder takes an observation $x$ and outputs parameters of a distribution over latent codes. Typically:

$$q_\phi(z|x) = \mathcal{N}(z; \mu_\phi(x), \text{diag}(\sigma^2_\phi(x)))$$

Where $\mu_\phi(x)$ and $\sigma_\phi(x)$ are neural network outputs, and $\phi$ are the network parameters.

The key insight: Instead of computing the intractable posterior, we optimize $q_\phi(z|x)$ to be as close as possible to $p_\theta(z|x)$. Closeness is measured by KL divergence:

$$D_{\text{KL}}(q_\phi(z|x) | p_\theta(z|x)) = \mathbb{E}{q\phi(z|x)}\left[\log \frac{q_\phi(z|x)}{p_\theta(z|x)}\right]$$

Deriving the Evidence Lower Bound (ELBO)

Let's manipulate the log-marginal likelihood to reveal a tractable objective. Starting with any observation $x$:

$$\log p_\theta(x) = \log \int p(z) \cdot p_\theta(x|z) , dz$$

We introduce $q_\phi(z|x)$ via an importance-weighting identity:

$$\log p_\theta(x) = \log \int q_\phi(z|x) \cdot \frac{p(z) \cdot p_\theta(x|z)}{q_\phi(z|x)} , dz$$

Since $\log$ is concave, Jensen's inequality gives:

$$\log p_\theta(x) \geq \int q_\phi(z|x) \log \frac{p(z) \cdot p_\theta(x|z)}{q_\phi(z|x)} , dz$$

$$= \mathbb{E}{q\phi(z|x)}\left[\log p_\theta(x|z) + \log p(z) - \log q_\phi(z|x)\right]$$

This is the Evidence Lower Bound (ELBO):

$$\text{ELBO}(\theta, \phi; x) = \mathbb{E}{q\phi(z|x)}[\log p_\theta(x|z)] - D_{\text{KL}}(q_\phi(z|x) | p(z))$$

The ELBO admits multiple equivalent decompositions, each illuminating different aspects of what VAE training accomplishes.

Decomposition 1: Reconstruction + Regularization

$$\text{ELBO} = \underbrace{\mathbb{E}{q\phi(z|x)}[\log p_\theta(x|z)]}{\text{Reconstruction term}} - \underbrace{D{\text{KL}}(q_\phi(z|x) | p(z))}_{\text{KL regularizer}}$$

Reconstruction term: The expected log-probability of observation $x$ under the decoder, given latent codes sampled from the encoder. Maximizing this trains the decoder to accurately reconstruct inputs from their latent codes.

KL regularizer: How much the approximate posterior $q_\phi(z|x)$ deviates from the prior $p(z)$. Minimizing this prevents the encoder from learning arbitrarily complex posteriors—it must stay "close" to the simple prior.

Decomposition 2: Log-Evidence Minus Gap

Alternatively, we can show:

$$\log p_\theta(x) = \text{ELBO}(\theta, \phi; x) + D_{\text{KL}}(q_\phi(z|x) | p_\theta(z|x))$$

Since KL divergence is non-negative, ELBO ≤ log p(x), confirming it's a lower bound. The gap between ELBO and true log-likelihood is exactly the KL divergence between approximate and true posteriors.

Crucial insight: Maximizing ELBO simultaneously:

1. Increases log p(x) (maximizes marginal likelihood—what we really want)
2. Decreases the gap (makes approximate posterior closer to true posterior)

VAEs have a clean interpretation as probabilistic graphical models, connecting them to the broader literature on Bayesian modeling.

The Generative Model (Prior × Likelihood)

The generative process defines a directed graphical model:

$$z \rightarrow x$$

Where:

• $z \sim p(z) = \mathcal{N}(0, I)$ — latent variable drawn from prior
• $x|z \sim p_\theta(x|z)$ — observation drawn from conditional likelihood

This factorizes the joint:

$$p_\theta(x, z) = p(z) \cdot p_\theta(x|z)$$

The Inference Model (Approximate Posterior)

The encoder $q_\phi(z|x)$ represents amortized inference—a single neural network that outputs approximate posteriors for any input, rather than optimizing per-datapoint variational parameters.

This "runs the generative model in reverse":

$$x \rightarrow z$$

Mapping observations back to their likely latent codes.

Why Amortized Inference Matters

Traditional variational inference optimizes separate variational parameters for each datapoint—computationally expensive and doesn't generalize to new data.

Amortization uses a single encoder network that learns to perform inference for any input:

• Efficiency: O(N) vs O(N × optimization steps per point)
• Generalization: Can infer latent codes for unseen data
• Scalability: Enables mini-batch training on large datasets

This is sometimes called variational autoencoding because the encoder "encodes" data into approximate posterior distributions over latent codes.

The VAE formulation is agnostic to specific distributional choices, but standard implementations use carefully chosen forms that balance expressiveness with tractability.

The Standard Gaussian Prior

$$p(z) = \mathcal{N}(z; 0, I) = \prod_{j=1}^{d} \mathcal{N}(z_j; 0, 1)$$

Why this choice?

1. Tractable KL: KL divergence between two Gaussians has closed form
2. Easy sampling: Generating $z \sim \mathcal{N}(0, I)$ is trivial
3. Symmetry: No dimension is privileged—the prior is isotropic
4. Universal approximation: With flexible decoder, can still model complex data

Likelihood Choices

The decoder $p_\theta(x|z)$ should match the data type:

Continuous data (images, audio): $$p_\theta(x|z) = \mathcal{N}(x; \mu_\theta(z), \sigma^2 I)$$

The decoder network outputs the mean $\mu_\theta(z)$, with fixed or learned variance. This yields a mean squared error reconstruction loss when $\sigma^2 = 1$:

$$\log p_\theta(x|z) \propto -|x - \mu_\theta(z)|^2$$

Binary data: $$p_\theta(x|z) = \prod_i \text{Bernoulli}(x_i; \pi_\theta(z)_i)$$

The decoder outputs probabilities $\pi_\theta(z) \in [0,1]^D$ (via sigmoid). This yields binary cross-entropy loss:

$$\log p_\theta(x|z) = \sum_i [x_i \log \pi_i + (1-x_i) \log(1-\pi_i)]$$

For the standard VAE with Gaussian encoder and prior, the KL divergence has a beautiful closed-form solution. Let:

• Prior: $p(z) = \mathcal{N}(0, I)$
• Approximate posterior: $q_\phi(z|x) = \mathcal{N}(\mu, \text{diag}(\sigma^2))$

Where $\mu = \mu_\phi(x)$ and $\sigma = \sigma_\phi(x)$ are encoder outputs.

Closed-Form KL Divergence

$$D_{\text{KL}}(q_\phi(z|x) | p(z)) = \frac{1}{2} \sum_{j=1}^{d} \left( \mu_j^2 + \sigma_j^2 - \log \sigma_j^2 - 1 \right)$$

This decomposes into per-dimension contributions:

• $\mu_j^2$: Penalizes encoder means far from zero
• $\sigma_j^2$: Penalizes variance larger than 1
• $-\log \sigma_j^2$: Penalizes variance smaller than 1
• $-1$: Constant offset

Minimizing KL divergence pushes each latent dimension toward $\mathcal{N}(0, 1)$.

Interpreting the KL Penalty

The KL term acts as a regularizer on the latent space:

1. Prevents memorization: Without KL, encoder could assign unique, non-overlapping regions to each training point—no generalization
2. Enables generation: Forces latent space to resemble prior, so sampling $z \sim p(z)$ produces meaningful inputs to decoder
3. Smooth interpolation: Overlapping posteriors mean nearby latent codes decode to similar outputs

The tension: Reconstruction wants informative, distinctive codes; KL wants codes near the prior. This tension is fundamental to VAE learning—balancing data fidelity with latent space structure.

We've derived the complete mathematical foundation of Variational Autoencoders from first principles. Let's consolidate the key ideas:

What's next:

We've established what VAEs optimize (the ELBO) and why (variational inference). The next page explores how: the encoder-decoder neural network architecture that implements these distributions, translating mathematical abstraction into trainable models.