Deriving the Evidence Lower Bound (ELBO)

As introduced, Variational Inference (VI) seeks an approximation $q(\mathbf{z})$ to the true posterior $p(\mathbf{z}|\mathbf{x})$ by turning the inference problem into an optimization problem. The central quantity we optimize is the Evidence Lower Bound (ELBO). Let's derive this important objective function.

Our goal is to make $q(\mathbf{z})$ as close as possible to the true posterior $p(\mathbf{z}|\mathbf{x})$. A standard way to measure the dissimilarity between two probability distributions is the Kullback-Leibler (KL) divergence. We want to minimize the KL divergence from our approximation $q(\mathbf{z})$ to the true posterior $p(\mathbf{z}|\mathbf{x})$:

$$KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x})) = \int q(\mathbf{z}) \log \frac{q(\mathbf{z})}{p(\mathbf{z}|\mathbf{x})} d\mathbf{z}$$

Minimizing this KL divergence directly is difficult because it involves the true posterior $p(\mathbf{z}|\mathbf{x})$, which itself contains the intractable evidence term $p(\mathbf{x})$ since $p(\mathbf{z}|\mathbf{x}) = p(\mathbf{x}, \mathbf{z}) / p(\mathbf{x})$.

Let's expand the definition of the posterior within the KL divergence formula:

$$KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x})) = \int q(\mathbf{z}) \log \frac{q(\mathbf{z}) p(\mathbf{x})}{p(\mathbf{x}, \mathbf{z})} d\mathbf{z}$$

We can split the logarithm term:

$$KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x})) = \int q(\mathbf{z}) \log q(\mathbf{z}) d\mathbf{z} - \int q(\mathbf{z}) \log p(\mathbf{x}, \mathbf{z}) d\mathbf{z} + \int q(\mathbf{z}) \log p(\mathbf{x}) d\mathbf{z}$$

Now, let's analyze each term:

1. $\int q(\mathbf{z}) \log q(\mathbf{z}) d\mathbf{z}$: This is the definition of the negative entropy of the distribution $q(\mathbf{z})$, often written as $-\mathbb{H}(q)$. In expectation notation, it's $\mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$.
2. $\int q(\mathbf{z}) \log p(\mathbf{x}, \mathbf{z}) d\mathbf{z}$: This is the expectation of the log joint probability $p(\mathbf{x}, \mathbf{z})$ under the distribution $q(\mathbf{z})$, written as $\mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})]$.
3. $\int q(\mathbf{z}) \log p(\mathbf{x}) d\mathbf{z}$: Since $\log p(\mathbf{x})$ is a constant with respect to the integration variable $\mathbf{z}$, we can pull it out of the integral. The remaining integral $\int q(\mathbf{z}) d\mathbf{z}$ is simply 1, because $q(\mathbf{z})$ is a probability distribution. So, this term simplifies to $\log p(\mathbf{x})$.

Substituting these back into the equation for the KL divergence:

$$KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x})) = \mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})] - \mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] + \log p(\mathbf{x})$$

Rearranging this equation to isolate the log evidence $\log p(\mathbf{x})$ gives us a fundamental identity in variational inference:

$$\log p(\mathbf{x}) = \underbrace{\mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] - \mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]}_{\mathcal{L}(q)} + KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$$

The term denoted $\mathcal{L}(q)$ is the Evidence Lower Bound (ELBO):

$$\mathcal{L}(q) = \mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] - \mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$$

This can also be written using the entropy $\mathbb{H}(q) = -\mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$ as:

$$\mathcal{L}(q) = \mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] + \mathbb{H}(q)$$

Why is it a Lower Bound?

The KL divergence has the property that $KL(A || B) \ge 0$ for any distributions $A$ and $B$, with equality holding if and only if $A=B$. Therefore, $KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x})) \ge 0$.

Looking back at the rearranged equation:

$$\log p(\mathbf{x}) = \mathcal{L}(q) + KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$$

Since the KL term is non-negative, it immediately follows that:

$$\log p(\mathbf{x}) \ge \mathcal{L}(q)$$

This confirms that $\mathcal{L}(q)$ is indeed a lower bound on the logarithm of the model evidence.

The Optimization Objective

The identity $\log p(\mathbf{x}) = \mathcal{L}(q) + KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$ provides the key insight for variational inference. We want to find the distribution $q(\mathbf{z})$ within our chosen family that minimizes the KL divergence $KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$.

Notice that the log evidence $\log p(\mathbf{x})$ is a constant with respect to the parameters defining $q(\mathbf{z})$. Therefore, maximizing the ELBO $\mathcal{L}(q)$ is equivalent to minimizing the KL divergence $KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$.

Think of it like this: the total log evidence is fixed. It's composed of the ELBO and the KL divergence. If we increase the ELBO, the KL divergence must decrease by the same amount to keep the sum constant.

This is extremely useful because the ELBO, $\mathcal{L}(q) = \mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] - \mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$, depends only on the joint distribution $p(\mathbf{x}, \mathbf{z})$ (which is typically computable) and our chosen variational distribution $q(\mathbf{z})$. It crucially avoids the intractable evidence $p(\mathbf{x})$.

Variational inference thus proceeds by:

1. Choosing a family of distributions for $q(\mathbf{z})$, often parameterized by some variational parameters $\phi$. Let's write this as $q_{\phi}(\mathbf{z})$.
2. Finding the parameters $\phi$ that maximize the ELBO: $$\phi^* = \arg \max_{\phi} \mathcal{L}(q_{\phi})$$

The resulting distribution $q_{\phi^*}(\mathbf{z})$ serves as our approximation to the true posterior $p(\mathbf{z}|\mathbf{x})$. The value of the ELBO at the optimum, $\mathcal{L}(q_{\phi^*})$, also provides a lower bound on the log marginal likelihood, which can be useful for model comparison.

Understanding the ELBO and its relationship to the KL divergence and the model evidence is fundamental to grasping how variational inference methods operate. The next sections will explore specific families for $q(\mathbf{z})$ and algorithms for optimizing the ELBO.