Optimization as Inference: The VI Perspective

As we established in the chapter introduction, directly computing the posterior distribution $p(\mathbf{z}|\mathbf{x})$ in Bayesian models is often hindered by the intractable evidence term $p(\mathbf{x}) = \int p(\mathbf{x}, \mathbf{z}) d\mathbf{z}$. While MCMC methods provide a powerful simulation-based approach to approximate the posterior by generating samples, they can demand significant computational resources and time, particularly when dealing with large datasets or high-dimensional parameter spaces.

Variational Inference (VI) offers a fundamentally different approach. Instead of simulating samples from the posterior, VI transforms the inference problem into an optimization problem. The central idea is to select a family of probability distributions, $\mathcal{Q}$, over the latent variables $\mathbf{z}$, where each distribution $q(\mathbf{z}) \in \mathcal{Q}$ is designed to be tractable (easy to compute expectations, densities, etc.). We then seek the specific distribution $q^*(\mathbf{z})$ within this family that is "closest" to the true, but intractable, posterior $p(\mathbf{z}|\mathbf{x})$.

Measuring Closeness: The KL Divergence

How do we quantify the "closeness" between our approximation $q(\mathbf{z})$ and the true posterior $p(\mathbf{z}|\mathbf{x})$? The standard measure in VI is the Kullback-Leibler (KL) divergence, denoted $KL(q || p)$. For our specific case, it's defined as:

$$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}$$

The KL divergence is always non-negative ($KL(q || p) \ge 0$), and it equals zero if and only if $q(\mathbf{z})$ and $p(\mathbf{z}|\mathbf{x})$ are identical almost everywhere. Minimizing this KL divergence means finding the distribution $q(\mathbf{z})$ within our chosen family $\mathcal{Q}$ that best matches the true posterior.

The Challenge and the Solution: Introducing the ELBO

At first glance, minimizing $KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$ doesn't seem to solve our original problem. Calculating the KL divergence directly still requires computing the true posterior $p(\mathbf{z}|\mathbf{x})$, which involves the problematic evidence term $p(\mathbf{x})$.

However, we can algebraically rearrange the definition of the KL divergence. Let's expand the term inside the logarithm using the definition of conditional probability, $p(\mathbf{z}|\mathbf{x}) = \frac{p(\mathbf{x}, \mathbf{z})}{p(\mathbf{x})}$:

$$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 separate the terms in the logarithm:

$$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}$$

Recognizing the integrals as expectations with respect to $q(\mathbf{z})$, and noting that $\log p(\mathbf{x})$ is a constant with respect to $\mathbf{z}$ (so $\int q(\mathbf{z}) \log p(\mathbf{x}) d\mathbf{z} = \log p(\mathbf{x}) \int q(\mathbf{z}) d\mathbf{z} = \log p(\mathbf{x})$), we get:

$$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 gives us a significant relationship:

$$\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 known as the Evidence Lower Bound (ELBO). Since the KL divergence is always non-negative ($KL(q || p) \ge 0$), this equation tells us that the ELBO is always less than or equal to the log model evidence:

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

This relationship is the cornerstone of variational inference.

Maximizing the ELBO as an Inference Strategy

Consider the equation $\log p(\mathbf{x}) = \mathcal{L}(q) + KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$. The left side, $\log p(\mathbf{x})$, is the log evidence of our model given the data. For a fixed model and dataset, this value is constant, regardless of our choice of $q(\mathbf{z})$. Therefore, maximizing the ELBO, $\mathcal{L}(q)$, with respect to $q(\mathbf{z})$ must be equivalent to minimizing the KL divergence, $KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$.

Crucially, 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})]$, does not depend directly on the intractable evidence $p(\mathbf{x})$ or the true posterior $p(\mathbf{z}|\mathbf{x})$. It depends only on the joint distribution $p(\mathbf{x}, \mathbf{z})$ (which is typically defined by our model specification) and our approximating distribution $q(\mathbf{z})$.

This transforms the inference problem into an optimization problem: find the distribution $q^*(\mathbf{z})$ within the chosen family $\mathcal{Q}$ that maximizes the ELBO. The resulting $q^*(\mathbf{z})$ serves as our approximation to the true posterior $p(\mathbf{z}|\mathbf{x})$.

Variational Inference finds the distribution $q^*(\mathbf{z})$ within a tractable family $\mathcal{Q}$ that maximizes the Evidence Lower Bound (ELBO). This maximization is equivalent to minimizing the KL divergence between $q^*(\mathbf{z})$ and the true posterior $p(\mathbf{z}|\mathbf{x})$.

The practical challenge now lies in:

1. Choosing an appropriate family of distributions $\mathcal{Q}$ that balances tractability and flexibility.
2. Developing algorithms to perform the optimization required to maximize $\mathcal{L}(q)$.

The following sections will examine common choices for $\mathcal{Q}$, such as the mean-field approximation, and explore algorithms like Coordinate Ascent Variational Inference (CAVI) and Stochastic Variational Inference (SVI) designed to efficiently maximize the ELBO.