Variational Autoencoders (VAEs)

Variational Autoencoders (VAEs)

Working notes on the mathematical foundations of VAEs

Problem Setup

Given: Samples $x$ from a distribution of interest

Goal: Model the true distribution $p(x)$, then generate new samples from our approximate distribution

Related approaches to explore: Autoregressive models, normalizing flows, variational autoencoders, energy-based models, score-based generative models

The Evidence Lower Bound (ELBO)

Why We Need It

Let $p(x, z)$ be the joint distribution of data and latent variables.

In a likelihood-based approach, we’d want to learn a model to maximize $p(x)$:

The difficulties:

• The integral over $z$ is difficult to compute

• We need access to the ground truth latent encoder $p(z

  x)$

Deriving the ELBO

Approach: Derive the Evidence Lower Bound (ELBO)

Starting from the log-likelihood:

Multiply and divide by the variational distribution:

Apply Jensen’s inequality (since log is concave):

This is the ELBO — a lower bound on the log-probability (the “evidence”).

Decomposing the ELBO

We can show that:

Since KL divergence is always $\geq 0$, we have:

The KL term represents the difference between the true encoder and the learned encoder. It’s non-negative, so we drop it from the objective.

Alternative: Maximize the ELBO. Since $\log p(x) = \text{ELBO} + D_{KL}$, and $\log p(x)$ is a constant, minimizing $D_{KL}$ is equivalent to maximizing the ELBO.

VAE Architecture

        Encoder                     Decoder
   ┌─────────────┐              ┌─────────────┐
   │             │    μ(x)      │             │
x ─┤    E_φ     ├──────┬───z───┤    D_θ     ├── x̂
   │             │      │       │             │
   └─────────────┘   log σ²(x)  └─────────────┘
                        │
                   z = μ_φ(x) + σ(x) ⊙ ε̂
                   
                   where ε̂ ~ N(0, I)

Expanding the ELBO for VAEs

This gives us:

• Reconstruction term: How well does the decoder reconstruct $x$ from samples of $z$?
• Prior matching term: How well does the learned encoder match the latent prior?

Evaluating the VAE Objective

Assumptions

• $x \in \mathbb{R}^{D_x}$, $z \in \mathbb{R}^{D_z}$
• Encoder: $\mathcal{E}_\phi : \mathbb{R}^{D_x} \to \mathbb{R}^{2D_z}$ (outputs mean and variance)
• Decoder: $\mathcal{D}_\theta : \mathbb{R}^{D_z} \to \mathbb{R}^{D_x}$
• Prior: $p(z) \sim \mathcal{N}(z; 0, I)$ — standard multivariate Gaussian

• Encoder distribution: $g_\phi(z

  x) \sim \mathcal{N}(z; \mu_\phi(x), \Sigma_\phi(x))$ — multivariate Gaussian with diagonal covariance

• Decoder distribution: $p_\theta(x

  z) \sim \mathcal{N}(x; \mu_\theta(z), \sigma^2 I)$, where $\sigma^2 = 1$

The Objective

Using Monte Carlo estimation with $L$ samples:

The KL term can be computed exactly for two multivariate Gaussians.

The Reparameterization Trick

Problem: Each latent $z^{(\ell)}$ is sampled, which is non-differentiable.

Solution: Use the reparameterization trick.

If $x \sim \mathcal{N}(x; \mu, \sigma^2)$, we can write: \(x = \mu + \sigma \epsilon, \quad \epsilon \sim \mathcal{N}(\epsilon; 0, 1)\)

In the VAE: \(z = \mu_\phi(x) + \sigma_\phi(x) \odot \hat{\epsilon}, \quad \hat{\epsilon} \sim \mathcal{N}(\hat{\epsilon}; 0, I)\)

Now gradients can be computed with respect to $\phi$!

After training: Generate new data by sampling directly from the latent prior $p(z)$ and running it through the decoder.

The Reconstruction Loss (Detailed Derivation)

This maps each $z$ to the log-probability of the fixed $x$.

For Gaussian decoder:

Setting $\sigma^2 = 1$:

Therefore:

where $\epsilon \sim \mathcal{N}(0, I)$.

The Prior Matching Loss (Detailed Derivation)

KL Divergence Background

For discrete distributions: \(KL(p \| g) = \sum_j p(y) \log \frac{p(y)}{g(y)} = -H(p) + H(p, g)\)

where $H(p)$ is entropy and $H(p, g)$ is cross-entropy.

For continuous distributions: \(KL(g \| p) = \int g(z) \log \frac{g(z)}{p(z)} \, dz = \text{Cross-entropy}(g, p) - \text{Entropy}(g)\)

KL Divergence for Univariate Gaussians

For $p_1(x) = \mathcal{N}(x; \mu_1, \sigma_1^2)$ and $p_2(x) = \mathcal{N}(x; \mu_2, \sigma_2^2)$:

If $p_2(x) = \mathcal{N}(x; 0, 1)$:

KL Divergence for Multivariate Gaussians

For Diagonal Covariance (VAE Case)

When both distributions have diagonal covariance, the KL divergence decomposes:

For the VAE prior matching term:

Final VAE Loss

Expanding:

Note: The $\beta$ term (when $\beta > 1$) gives the $\beta$-VAE, which encourages more disentangled representations.

Implementation Notes

In code, the encoder outputs:

• mu: $\mu_\phi(x)$
• logvar: $\log \sigma^2(x; \phi)$

Then: \(\sigma(x) = \exp(0.5 \cdot \text{logvar})\)

And the reconstruction: \(\hat{x} = \mathcal{D}_\theta(\mu_\phi(x) + \sigma(x) \odot \hat{\epsilon})\)

Bernoulli Decoder Variant

For binary data (e.g., MNIST), assume each pixel $x_i$ represents the expected value of a binary random variable $x_i’ \in {0, 1}$.

Each pixel is generated by a Bernoulli process:

• $x_i’ \sim \text{Bernoulli}(x_i)$ (latent variable)
• $x_i = \mathbb{E}_{x_i’ \sim \text{Ber}(x_i)}[x_i’]$ (observation)

The log-likelihood becomes: \(\log p_\theta(x|z) = \sum_{i=1}^{D} x_i \log \hat{x}_i + (1 - x_i) \log(1 - \hat{x}_i)\)

This is the binary cross-entropy loss.

Conditional VAE

The same architecture, but the encoder and decoder are conditioned on some additional information (e.g., class labels).

Further Reading

• Kingma & Welling (2014): Auto-Encoding Variational Bayes
• Murphy, PML-1: Chapter 6.2.1 (KL divergence), pages 105, 174, 205, 308
• Matrix Cookbook: Section 8.2.2 (equation 380)