Variational Inference Primer

Learning Objectives

By the end of this section, you will be able to:

1. Understand latent variable models and why direct likelihood computation is often intractable
2. Derive the Evidence Lower Bound (ELBO) and understand its two components: reconstruction and regularization
3. Apply the reparameterization trick to enable gradient-based optimization through sampling operations
4. Connect variational inference to diffusion models and understand how the diffusion ELBO decomposes into per-timestep terms
5. Implement a simple VAE in PyTorch using the ELBO objective

The Big Picture: Inference as Optimization

Variational inference emerged from a profound insight: when exact probabilistic inference is computationally intractable, we can turn inference into an optimization problem. Instead of computing the true posterior distribution exactly, we search for the best approximation from a tractable family of distributions.

Historical Context: Variational methods have roots in physics (mean-field theory) and statistical mechanics. The application to machine learning was pioneered by researchers like Geoffrey Hinton and Michael Jordan in the 1990s. The modern era of variational inference began with Kingma and Welling's VAE paper (2013) and Rezende et al.'s work on stochastic variational inference.

The key insight is that we can convert an intractable integration problem into a tractable optimization problem. This is exactly what we need for diffusion models, where we want to learn a generative model without being able to compute exact likelihoods.

Latent Variable Models

Latent variable models assume that observed data $x$ is generated through unobserved (latent) variables $z$. The generative process is:

The marginal likelihood of the data requires integrating over all possible latent configurations:

Why Latent Variables?

Latent variables serve multiple purposes in generative modeling:

• Expressiveness: Simple distributions in latent space can map to complex distributions in data space
• Disentanglement: Latent dimensions can capture independent factors of variation (e.g., pose, lighting, identity)
• Compression: High-dimensional data can be encoded into lower-dimensional representations
• Generation: New samples can be created by sampling latent codes and decoding them

Real-World Example (Vision): Consider face images. The latent space might capture age, expression, pose, and lighting as separate dimensions. A VAE learns this without supervision!

The Intractability Problem

For learning and inference, we need the posterior distribution:

The denominator $p(x)$ (the evidence or marginal likelihood) requires integrating over all possible latent configurations. For continuous latent spaces and nonlinear relationships (like those modeled by neural networks), this integral is almost always intractable.

Why Is This Hard?

Consider a VAE with a 128-dimensional latent space. Computing$p(x)$ exactly would require integrating over all points in $\mathbb{R}^{128}$. Even with Monte Carlo methods, the variance of such estimates would be prohibitive.

This is where variational inference comes to the rescue: instead of computing the true posterior, we approximate it with a simpler distribution that we can work with.

Variational Inference: Core Idea

The core idea of variational inference is to introduce an approximate posterior $q_\phi(z|x)$ from a tractable family of distributions (parameterized by $\phi$) and optimize it to be as close as possible to the true posterior.

We measure "closeness" using KL divergence (which we covered in the information theory section):

Recall from the previous section that KL divergence is asymmetric. We use $D_{\text{KL}}(q \| p)$ (reverse KL) rather than $D_{\text{KL}}(p \| q)$ (forward KL) because it leads to a tractable objective.

Key Insight: The reverse KL tends to produce mode-seeking behavior: $q$ will concentrate on regions where $p$ has high probability, rather than trying to cover all of $p$'s support. This is important for understanding VAE behavior.

Deriving the ELBO

We want to minimize $D_{\text{KL}}(q_\phi(z|x) \| p(z|x))$, but this requires knowing $p(z|x)$, which in turn requires the intractable $p(x)$. Here's the elegant solution:

Start with the definition of KL divergence:

Use Bayes' rule to expand $p(z|x)$:

Since $p(x)$ doesn't depend on $z$, it comes out of the expectation:

Rearranging:

The second term is the Evidence Lower Bound (ELBO):

Since KL divergence is always non-negative, we have the fundamental inequality:

This is why it's called a "lower bound" on the evidence! By maximizing the ELBO, we simultaneously:

1. Maximize the log-likelihood $\log p(x)$ (as much as the bound allows)
2. Minimize $D_{\text{KL}}(q_\phi \| p)$, making our approximation better

ELBO Decomposition

The ELBO has a beautiful interpretation with two competing terms:

Reconstruction Term

$\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]$measures how well we can reconstruct the input from the latent code. This is the "quality" term that encourages the model to encode useful information in $z$.

• For continuous data with Gaussian likelihood: equivalent to negative MSE
• For binary data with Bernoulli likelihood: equivalent to binary cross-entropy
• Encourages the encoder to preserve information

Regularization (KL) Term

$D_{\text{KL}}(q_\phi(z|x) \| p(z))$ measures how much the approximate posterior deviates from the prior. This term:

• Prevents the encoder from memorizing each data point
• Encourages smooth, regular latent spaces
• Enables generation by sampling from the prior

The Trade-off: If we only maximize reconstruction, the model becomes an autoencoder without generative capability. If we only minimize KL, we get random noise. The ELBO balances these naturally!

Mean-Field Approximation

The most common choice for the variational family is the mean-field approximation, where we assume the latent dimensions are independent:

For VAEs, we typically use a multivariate Gaussian with diagonal covariance:

The encoder neural network outputs the mean $\mu$ and variance $\sigma^2$ for each data point. In practice, we output $\log \sigma^2$ for numerical stability.

Advantages of Mean-Field

• Closed-form KL: For Gaussian distributions with Gaussian prior, the KL term has an analytical solution
• Simple parameterization: Only need to output mean and variance vectors
• Efficient sampling: Can sample from each dimension independently

The Closed-Form KL for Gaussians

When $q = \mathcal{N}(\mu, \sigma^2)$ and$p = \mathcal{N}(0, 1)$ (standard normal prior), the KL divergence has a simple form:

This is the famous KL regularization term used in VAEs!

The Reparameterization Trick

There's one remaining challenge: the ELBO requires taking expectations over $q_\phi(z|x)$, which depends on the parameters $\phi$ we want to optimize. We need to estimate gradients through this expectation.

The problem: If we sample $z \sim q_\phi(z|x)$directly, this sampling operation is not differentiable! We can't compute $\partial z / \partial \phi$.

The solution: The reparameterization trick rewrites the sampling operation as a deterministic transformation of a parameter-free random variable:

Now the randomness is in $\epsilon$, which doesn't depend on $\phi$. The gradients flow through$\mu$ and $\sigma$!

Why This Works

Consider the gradient of the reconstruction term:

The expectation is now over $\epsilon$, which doesn't depend on $\phi$. We can:

1. Sample $\epsilon_1, \ldots, \epsilon_K$ from $\mathcal{N}(0, I)$
2. Compute $z_k = \mu + \sigma \cdot \epsilon_k$ for each sample
3. Estimate the gradient with Monte Carlo:$\frac{1}{K}\sum_k \nabla_\phi \log p(x|z_k)$

Key Insight for Diffusion: The reparameterization trick is exactly how we train diffusion models! When we sample$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$, we're using the same reparameterization to enable gradient flow.

Connection to Diffusion Models

Diffusion models are a special case of hierarchical variational autoencoders! The key insight is that diffusion models have a fixed, predetermined forward process (the encoder) and only learn the reverse process (the decoder).

The Diffusion ELBO

In diffusion models, the ELBO decomposes into a sum over timesteps:

The Simplified Loss

Ho et al. (2020) showed that this complex ELBO simplifies to a remarkably simple objective when parameterized correctly:

This is just predicting the noise that was added! The connection to variational inference explains why this works: we're still maximizing a lower bound on the log-likelihood, just in a computationally convenient form.

Deep Connection: The diffusion model training objective is derived from variational inference principles. Each denoising step learns to approximate the true reverse conditional$q(x_{t-1}|x_t, x_0)$, which has a closed form because the forward process is Gaussian!

Implementation in PyTorch

Let's implement a simple VAE to solidify our understanding of the ELBO and reparameterization trick:

1import torch.nn as nn
2import torch.nn.functional as F
3
4class SimpleVAE(nn.Module):
5    def __init__(self, input_dim=784, latent_dim=32):
6        super().__init__()
7        self.encoder = nn.Sequential(nn.Linear(input_dim, 256), nn.ReLU())
8        self.fc_mu = nn.Linear(256, latent_dim)
9        self.fc_logvar = nn.Linear(256, latent_dim)
10        self.decoder = nn.Sequential(nn.Linear(latent_dim, 256), nn.ReLU(), nn.Linear(256, input_dim))
11
12    def reparameterize(self, mu, logvar):
13        std = torch.exp(0.5 * logvar)  # Convert log-variance to std
14        eps = torch.randn_like(std)     # Sample noise: eps ~ N(0, I)
15        z = mu + eps * std              # Reparameterization: z = mu + std * eps
16        return z
17
18    def forward(self, x):
19        h = self.encoder(x.view(-1, 784))
20        mu, logvar = self.fc_mu(h), self.fc_logvar(h)
21        z = self.reparameterize(mu, logvar)
22        recon = torch.sigmoid(self.decoder(z))
23        return recon, mu, logvar
24
25def vae_loss(recon_x, x, mu, logvar):
26    # Reconstruction loss (binary cross-entropy)
27    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
28    # KL divergence: -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
29    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
30    return BCE + KLD  # ELBO = -BCE - KLD, so we minimize BCE + KLD

Now let's see how the same principles apply to diffusion model training:

1from diffusers import DDPMScheduler
2import torch
3
4# Setup noise schedule
5scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear")
6alphas_cumprod = scheduler.alphas_cumprod
7
8def diffusion_loss(model, x_0, device):
9    """Simplified diffusion ELBO loss - just predict the noise!"""
10    batch_size = x_0.shape[0]
11    # Sample random timesteps
12    t = torch.randint(0, 1000, (batch_size,), device=device)
13    noise = torch.randn_like(x_0)
14
15    # Create noisy image using reparameterization: x_t = sqrt(alpha_bar) * x_0 + sqrt(1 - alpha_bar) * eps
16    sqrt_alpha = alphas_cumprod[t].sqrt().view(-1, 1, 1, 1)
17    sqrt_one_minus_alpha = (1 - alphas_cumprod[t]).sqrt().view(-1, 1, 1, 1)
18    x_t = sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise
19
20    # Model predicts the noise
21    predicted_noise = model(x_t, t)
22    loss = F.mse_loss(predicted_noise, noise)  # Simple L2 loss on noise prediction
23    return loss

Summary

Variational inference provides the theoretical foundation for training generative models when exact likelihood computation is intractable. Here are the key takeaways:

1. The Intractability Problem: Computing$p(x) = \int p(x|z) p(z) dz$ is intractable for nonlinear latent variable models
2. The ELBO: We can maximize a lower bound on log-likelihood that decomposes into reconstruction and regularization terms
3. Reparameterization: By writing$z = \mu + \sigma \cdot \epsilon$, we enable gradient-based optimization through sampling
4. Mean-Field: Assuming independent latent dimensions gives tractable optimization with closed-form KL terms
5. Diffusion Connection: Diffusion models are hierarchical VAEs with fixed encoders and learned decoders, trained using the same variational principles

Looking Ahead: In the next section, we'll explore Markov chains and stochastic processes - the mathematical framework that describes how diffusion models progressively add and remove noise. Understanding these processes will complete our prerequisite toolkit for diving into diffusion models proper.