Variational Autoencoders

Variational Autoencoders represent one of the most elegant intersections of deep learning and probabilistic modeling. Unlike purely discriminative models that learn boundaries between classes, or simple autoencoders that merely compress and reconstruct, VAEs provide a principled probabilistic framework for learning the underlying distribution of data—enabling not just reconstruction, but genuine generation of new samples.

At the heart of VAEs lies a profound insight: we can frame generative modeling as a latent variable model and derive a tractable learning objective through variational inference. The result—the Evidence Lower Bound (ELBO)—gives us both a training signal and deep theoretical guarantees about what we're learning.

This page derives the VAE objective from first principles, ensuring you understand not just what the loss function is, but why it takes this particular form and what each term contributes to the model's behavior.

Before diving into the VAE objective, let's crystallize what we're trying to accomplish. In generative modeling, we have a dataset $\mathcal{D} = {\mathbf{x}^{(1)}, \mathbf{x}^{(2)}, ..., \mathbf{x}^{(N)}}$ of samples (images, text, audio, etc.), and our goal is to learn the true underlying data distribution $p_{\text{data}}(\mathbf{x})$.

The maximum likelihood principle tells us the optimal approach: find the model parameters $\boldsymbol{\theta}$ that maximize the probability of observing our data:

$$\boldsymbol{\theta}^* = \arg\max_{\boldsymbol{\theta}} \prod_{i=1}^{N} p_{\boldsymbol{\theta}}(\mathbf{x}^{(i)})$$

Equivalently, we maximize the log-likelihood:

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

This seems straightforward, but there's a fundamental challenge: for complex, high-dimensional data like images, directly defining and computing $p_{\boldsymbol{\theta}}(\mathbf{x})$ is extraordinarily difficult. How can we tractably model a distribution over all possible 256×256 RGB images?

The Latent Variable Solution

The key insight is to introduce latent variables $\mathbf{z}$ that capture the underlying factors of variation in our data. We hypothesize that each observed datapoint $\mathbf{x}$ was generated from some unobserved latent code $\mathbf{z}$ through a generative process:

1. Sample a latent code: $\mathbf{z} \sim p(\mathbf{z})$ (prior distribution)
2. Generate the observation: $\mathbf{x} \sim p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z})$ (conditional likelihood)

The marginal likelihood of data under this model is:

$$p_{\boldsymbol{\theta}}(\mathbf{x}) = \int p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}) p(\mathbf{z}) , d\mathbf{z}$$

This latent variable formulation is powerful because:

• The prior $p(\mathbf{z})$ can be simple (e.g., standard Gaussian)
• The conditional $p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z})$ can be modeled by a neural network
• Together, they implicitly define a complex distribution over $\mathbf{x}$

The latent variable formulation seems elegant, but it introduces a severe computational problem. To maximize $\log p_{\boldsymbol{\theta}}(\mathbf{x})$, we need to compute:

$$\log p_{\boldsymbol{\theta}}(\mathbf{x}) = \log \int p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}) p(\mathbf{z}) , d\mathbf{z}$$

This integral is intractable for two fundamental reasons:

1. High-Dimensional Integration

The integral is over the entire latent space, which may have hundreds of dimensions. No closed-form solution exists, and numerical integration is computationally impossible in high dimensions (curse of dimensionality).

2. Posterior Intractability

To efficiently estimate this integral via Monte Carlo methods, we would ideally sample from the posterior $p_{\boldsymbol{\theta}}(\mathbf{z}|\mathbf{x})$—the distribution of latent codes that could have generated a specific observation. But by Bayes' theorem:

$$p_{\boldsymbol{\theta}}(\mathbf{z}|\mathbf{x}) = \frac{p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}) p(\mathbf{z})}{p_{\boldsymbol{\theta}}(\mathbf{x})}$$

Computing this requires the very marginal likelihood $p_{\boldsymbol{\theta}}(\mathbf{x})$ we're trying to compute—a circular dependency.

Why Naive Monte Carlo Fails

One might attempt to estimate the marginal likelihood directly:

$$p_{\boldsymbol{\theta}}(\mathbf{x}) \approx \frac{1}{K} \sum_{k=1}^{K} p_{\boldsymbol{\theta}}(\mathbf{x} | \mathbf{z}^{(k)}), \quad \mathbf{z}^{(k)} \sim p(\mathbf{z})$$

This is called importance sampling with the prior as the proposal distribution. While unbiased, it has astronomically high variance for high-dimensional data.

Why it fails: Consider an image of a face. Most samples from the prior $p(\mathbf{z})$ will decode to noise or unrelated images. Only a tiny fraction of the latent space actually corresponds to faces. We would need an impossibly large number of samples to get even one that's relevant to our specific image.

Mathematically, if the posterior $p(\mathbf{z}|\mathbf{x})$ concentrates in a small region of latent space (which it does for meaningful data), samples from the prior $p(\mathbf{z})$ will almost never land in that region.

The solution: Instead of sampling from the prior and hoping, we need to sample from a distribution that actually concentrates in the relevant regions of latent space for each datapoint. This is what the variational approach provides.

The breakthrough of VAEs is to address the intractability through variational inference. Since we can't compute the true posterior $p_{\boldsymbol{\theta}}(\mathbf{z}|\mathbf{x})$, we'll introduce an approximate posterior $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ parameterized by $\boldsymbol{\phi}$.

This approximate posterior will be:

• A tractable distribution family (typically Gaussian)
• Conditioned on the observation $\mathbf{x}$ through a neural network (the encoder)
• Optimized to be as close as possible to the true posterior

The variational approach converts our intractable optimization problem (maximize log-likelihood) into a tractable one (maximize a lower bound while making the approximate posterior accurate).

The Amortized Inference Advantage

Traditional variational inference optimizes separate parameters for each datapoint. VAEs use amortized inference: a single neural network (the encoder) outputs the approximate posterior parameters for any input $\mathbf{x}$.

This amortization is crucial:

• Computational efficiency: One forward pass gives $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$
• Generalization: The encoder can process novel inputs at test time
• Learning: Encoder and decoder are trained jointly end-to-end

The Standard Choice: Diagonal Gaussian

The most common choice for the approximate posterior is a diagonal Gaussian:

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

Where:

• $\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x})$ is the mean, output by the encoder network
• $\boldsymbol{\sigma}^2_{\boldsymbol{\phi}}(\mathbf{x})$ is the variance (diagonal covariance), also from the encoder

This choice satisfies all our requirements:

• Tractable sampling: $\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}$, where $\boldsymbol{\epsilon} \sim \mathcal{N}(0, I)$
• Tractable density: Standard Gaussian density formula
• Closed-form KL: The KL divergence to a standard Gaussian has a closed form
• Differentiable: Smooth functions of $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$

Now we derive the central mathematical result: the Evidence Lower Bound (ELBO). This derivation is fundamental—understanding it deeply will inform your intuition about VAE behavior.

Starting Point: The Log-Likelihood

We want to maximize $\log p_{\boldsymbol{\theta}}(\mathbf{x})$. Let's introduce our approximate posterior by multiplying and dividing:

$$\log p_{\boldsymbol{\theta}}(\mathbf{x}) = \log \int p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}) , d\mathbf{z}$$

$$= \log \int p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}) \frac{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})}{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} , d\mathbf{z}$$

$$= \log \mathbb{E}{q{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \left[ \frac{p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})}{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \right]$$

Applying Jensen's Inequality

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

$$\log \mathbb{E}[X] \geq \mathbb{E}[\log X]$$

Applying this:

$$\log p_{\boldsymbol{\theta}}(\mathbf{x}) \geq \mathbb{E}{q{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \left[ \log \frac{p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})}{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \right]$$

The right-hand side is the Evidence Lower Bound (ELBO):

$$\mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x}) = \mathbb{E}{q{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \left[ \log \frac{p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})}{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \right]$$

So we have established:

$$\boxed{\log p_{\boldsymbol{\theta}}(\mathbf{x}) \geq \mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x})}$$

The ELBO can be decomposed in two illuminating ways. Each reveals different aspects of what the VAE is learning.

Decomposition 1: Reconstruction + Regularization

Using $p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}) = p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})p(\mathbf{z})$:

$$\mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x}) = \mathbb{E}{q{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \left[ \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) \right] - D_{\text{KL}}(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}))$$

This is the most common form, revealing two terms:

Term 1: Reconstruction Loss (Expected Log-Likelihood) $$\mathbb{E}{q{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \left[ \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) \right]$$

This measures how well the decoder reconstructs $\mathbf{x}$ from latent samples. Higher is better—we want the decoder to produce high probability for the original data.

Term 2: KL Regularization (Posterior Regularization) $$D_{\text{KL}}(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}))$$

This penalizes deviation of the approximate posterior from the prior. It encourages the encoder to produce latent codes that stay close to the prior distribution.

Decomposition 2: The Gap to True Log-Likelihood

There's an alternative decomposition that reveals exactly when ELBO equals the log-likelihood:

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

This equality tells us:

1. The gap between ELBO and true log-likelihood is exactly the KL divergence between approximate and true posteriors
2. ELBO is tight (equals log-likelihood) if and only if $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) = p_{\boldsymbol{\theta}}(\mathbf{z}|\mathbf{x})$
3. Better approximate posteriors → tighter bounds → better training signal

Since KL divergence is always non-negative, we always have $\mathcal{L} \leq \log p_{\boldsymbol{\theta}}(\mathbf{x})$, confirming it's a lower bound.

Let's see how to actually compute the ELBO for training. We'll work with the standard choices: Gaussian prior, Gaussian approximate posterior, and either Gaussian or Bernoulli likelihood.

The KL Term: Closed-Form for Gaussians

With a standard Gaussian prior $p(\mathbf{z}) = \mathcal{N}(0, I)$ and diagonal Gaussian posterior $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) = \mathcal{N}(\boldsymbol{\mu}, \text{diag}(\boldsymbol{\sigma}^2))$:

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

This has a closed-form solution! No sampling required. Each latent dimension contributes independently.

The Reconstruction Term: Monte Carlo Estimation

The reconstruction term $\mathbb{E}{q}[\log p{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})]$ is estimated via sampling:

$$\mathbb{E}{q}[\log p{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})] \approx \frac{1}{L}\sum_{l=1}^{L} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}^{(l)}), \quad \mathbf{z}^{(l)} \sim q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$$

In practice, $L=1$ (single sample) works remarkably well due to minibatch averaging.

For Binary Data (Bernoulli Likelihood): $$\log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) = \sum_i x_i \log \hat{x}_i + (1-x_i) \log(1-\hat{x}_i)$$

This is the binary cross-entropy loss! The decoder outputs probabilities $\hat{x}i = \text{sigmoid}(f{\boldsymbol{\theta}}(\mathbf{z}))$.

For Continuous Data (Gaussian Likelihood): $$\log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) \propto -\frac{1}{2\sigma^2}||\mathbf{x} - \hat{\mathbf{x}}||^2$$

With fixed variance, this reduces to mean squared error! The decoder outputs $\hat{\mathbf{x}} = f_{\boldsymbol{\theta}}(\mathbf{z})$.

Putting it all together, the VAE training objective is to maximize the ELBO, which we can equivalently express as minimizing the negative ELBO loss:

$$\mathcal{L}{\text{VAE}} = \mathbb{E}{\mathbf{z} \sim q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \left[-\log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})\right] + D_{\text{KL}}(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}))$$

$$= \underbrace{\text{Reconstruction Loss}}{\text{How well can we recover } \mathbf{x}?} + \underbrace{\text{KL Divergence}}{\text{How "normal" is the latent code?}}$$

In practice, a weighting factor β is often introduced:

$$\mathcal{L}{\beta\text{-VAE}} = \text{Reconstruction Loss} + \beta \cdot D{\text{KL}}$$

With $\beta = 1$, we have the standard VAE (maximizing ELBO). Different values of $\beta$ create different trade-offs and have spawned influential variants like β-VAE for disentangled representations.

Beyond the mathematics, it's essential to build intuition for what the VAE objective is achieving. Let's examine each component's role in shaping the learned representation.

The Tension That Drives Learning

The two terms of the ELBO create a productive tension:

Reconstruction wants specificity: To accurately reconstruct $\mathbf{x}$, the encoder should produce a narrow posterior centered exactly at the latent code that makes $\mathbf{x}$ most likely. Ideally, a delta function $q(\mathbf{z}|\mathbf{x}) = \delta(\mathbf{z} - \mathbf{z}^*)$.

KL regularization wants generality: The KL term penalizes posteriors that deviate from the prior. It wants the encoder to produce posteriors that are broad, close to $\mathcal{N}(0, I)$—essentially forgetting the input.

The compromise: The model finds a middle ground where posteriors encode enough information for reconstruction while maintaining a structured, prior-like distribution. This compromise is what gives VAEs their generative capabilities.

Why This Works for Generation

The VAE objective implicitly ensures good generative properties:

1. Coverage of the prior: The KL term forces $q(\mathbf{z}|\mathbf{x})$ to cover most of the prior support. If all posteriors collapse to a small region, KL would be high. This means random samples from the prior $p(\mathbf{z})$ will likely fall in regions where some training data mapped.

2. Smooth latent space: Overlapping posteriors create smooth transitions between encoded points. Moving in latent space corresponds to meaningful changes in generated outputs.

3. No holes: Unlike deterministic autoencoders where the decoder only learns to handle specific encoder outputs, VAE decoders must handle samples from broad Gaussian posteriors—making them robust to sampling from the prior.

The Information Theory Perspective

The ELBO has an elegant information-theoretic interpretation:

$$\mathcal{L} = \underbrace{\mathbb{E}q[\log p(\mathbf{x}|\mathbf{z})]}\text{Bits to reconstruct x given z} - \underbrace{D_{KL}(q(\mathbf{z}|\mathbf{x}) || p(\mathbf{z}))}_\text{Bits to transmit z}$$

The VAE learns an efficient encoding: minimize the bits needed to specify the latent code while keeping enough information to reconstruct the data. This is the rate-distortion trade-off from information theory!

We've built up a complete understanding of the VAE objective from first principles. Let's consolidate the key insights:

What's Next:

With the theoretical objective established, the next page explores the encoder and decoder networks—the neural network architectures that parameterize the approximate posterior and likelihood. We'll see how to design these networks for different data modalities and examine the architectural choices that impact VAE performance.