Variational Inference: Approximating the Impossible

In deep learning, we often want to know the probability of some "hidden" state $z$ given our observed data $x$. This is $P(z|x)$. Unfortunately, for complex data like images, this is usually intractable—it would take an infinite amount of time to calculate.

Variational Inference (VI) is our trick to turn an impossible integration problem into a solvable optimization problem.

---

1. The Core Idea: Distribution Matching

Since we can't find the true distribution $P(z|x)$, we choose a simpler "proxy" distribution $Q(z)$ (usually a Gaussian) and try to make it as similar to the true distribution as possible.

We measure this similarity using the Kullback-Leibler (KL) Divergence:

$D_{KL}(Q(z) || P(z|x)) = \mathbb{E}_{z \sim Q} \left[ \log \frac{Q(z)}{P(z|x)} \right]$

Our goal is to minimize this divergence.

---

2. The ELBO (Evidence Lower Bound)

Since we don't know $P(z|x)$, we can't minimize the KL Divergence directly. However, using some clever math (Bayes' Rule), we can derive the ELBO:

$\log P(x) \geq \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{Reconstruction Term}} - \underbrace{D_{KL}(q(z|x) || p(z))}_{\text{Prior Matching Term}}$

• Reconstruction Term: Does our hidden state $z$ actually represent the image $x$?
• Prior Matching: Is our distribution of hidden states "well-behaved" (e.g., similar to a standard Gaussian)?

In Diffusion models, the ELBO becomes our loss function. We maximize it to ensure our model generates realistic data.

---

3. Why VI for Diffusion?

Diffusion models are essentially deep hierarchical VAEs (Variational Autoencoders) with thousands of layers.

1. Forward Path: We know the distribution $q(x_t | x_{t-1})$—it's just adding noise.
2. Reverse Path: We don't know $p(x_{t-1} | x_t)$—that's what the neural network has to learn.

By using the ELBO, we provide the network with a mathematically sound objective to learn the denoising process.

---

Mathematical Insight

Think of Variational Inference as "Curve Fitting" for probability distributions. Instead of finding the exact curve, we pick a flexible curve (Gaussian) and stretch it until it fits the data as best as it can.