Mean-Field Approximation Details

Variational Inference (VI) transforms the challenge of computing the posterior $p(\mathbf{z}|\mathbf{x})$ into an optimization problem. We aim to find a distribution $q(\mathbf{z})$ within a chosen family $\mathcal{Q}$ that best approximates the true posterior, typically by maximizing 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})]$$

The main step is selecting the family $\mathcal{Q}$. If we allowed $\mathcal{Q}$ to contain all possible distributions, the optimal $q^*(\mathbf{z})$ would be the true posterior $p(\mathbf{z}|\mathbf{x})$ itself, bringing us back to the original intractable problem. Therefore, the essence of practical VI lies in choosing a restricted, tractable family for $q(\mathbf{z})$.

The Mean-Field Assumption: Factorizing Complexity

One of the most widely used strategies is the mean-field variational family. This approach introduces a significant simplification by assuming that the latent variables $\mathbf{z} = (z_1, z_2, ..., z_M)$ in the approximating distribution $q(\mathbf{z})$ are mutually independent. Mathematically, we enforce a fully factorized structure:

$$q(\mathbf{z}) = \prod_{j=1}^{M} q_j(z_j; \lambda_j)$$

Here, the joint variational distribution $q(\mathbf{z})$ is broken down into a product of $M$ independent factors, where each factor $q_j(z_j; \lambda_j)$ governs only a single latent variable $z_j$ (or sometimes a block of variables, though full factorization is common). Each factor $q_j$ is itself a probability distribution, often belonging to a simple parametric family (like Gaussian or Dirichlet) and parameterized by its own set of variational parameters $\lambda_j$. The goal of VI then becomes finding the optimal parameters $\{\lambda_j\}_{j=1}^M$ that maximize the ELBO.

Simplifying the ELBO

This factorization assumption has a direct impact on the ELBO terms. The entropy term $\mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$ becomes particularly manageable:

$$\mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})] = \mathbb{E}_{q(\mathbf{z})} \left[ \sum_{j=1}^{M} \log q_j(z_j) \right] = \sum_{j=1}^{M} \mathbb{E}_{q_j(z_j)} [\log q_j(z_j)]$$

The expectation of the sum becomes the sum of expectations, and because each $q_j$ only depends on $z_j$, the expectation $\mathbb{E}_{q(\mathbf{z})}$ simplifies to $\mathbb{E}_{q_j(z_j)}$ for the $j$-th term. This means the total entropy is just the sum of the individual entropies of the factors. If the factors $q_j$ belong to standard exponential families, their entropies often have closed-form expressions, making this term easy to compute.

The first term, $\mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})]$, also simplifies. The expectation is taken over the factorized distribution $q(\mathbf{z})$, meaning we integrate or sum over each $z_j$ according to its factor $q_j(z_j)$. The specific calculation depends heavily on the structure of the model's joint probability $p(\mathbf{x}, \mathbf{z})$.

Deriving the Optimal Factors

How do we find the optimal form for each factor $q_j(z_j)$? We can optimize the ELBO with respect to one factor $q_j$ while holding the others ($q_k$ for $k \neq j$) fixed. This is the core idea behind algorithms like Coordinate Ascent Variational Inference (CAVI), which we'll discuss in the next section.

Let's isolate the terms in the ELBO that depend on a specific factor $q_j(z_j)$. We can rewrite the ELBO as:

$$\mathcal{L}(q) = \int q_j(z_j) \left( \int \prod_{k \neq j} q_k(z_k) \log p(\mathbf{x}, \mathbf{z}) dz_{-j} \right) dz_j - \int q_j(z_j) \log q_j(z_j) dz_j + \text{terms not depending on } q_j$$

where $dz_{-j}$ denotes integration over all variables $\mathbf{z}$ except $z_j$. The term inside the parentheses in the first integral is an expectation of the log joint probability taken with respect to all factors except $q_j$:

$$\mathbb{E}_{q_{-j}} [\log p(\mathbf{x}, \mathbf{z})] \triangleq \int \log p(\mathbf{x}, \mathbf{z}) \prod_{k \neq j} q_k(z_k) dz_{-j}$$

This expectation $\mathbb{E}_{q_{-j}} [\log p(\mathbf{x}, \mathbf{z})]$ yields a function that depends on $z_j$ (and $\mathbf{x}$, and the parameters of the fixed $q_k$'s). Let's denote $\log \tilde{p}_j(z_j) = \mathbb{E}_{q_{-j}} [\log p(\mathbf{x}, \mathbf{z})]$. Then the terms in the ELBO depending on $q_j(z_j)$ are:

$$\mathcal{L}_j = \int q_j(z_j) \log \tilde{p}_j(z_j) dz_j - \int q_j(z_j) \log q_j(z_j) dz_j$$

This expression looks familiar. It is equal to the negative Kullback-Leibler (KL) divergence between $q_j(z_j)$ and an unnormalized distribution proportional to $\tilde{p}_j(z_j)$, plus a constant (the log normalizer of $\tilde{p}_j(z_j)$):

$$\mathcal{L}_j = -KL(q_j(z_j) || \frac{\tilde{p}_j(z_j)}{C_j}) + \log C_j \quad \text{where } C_j = \int \tilde{p}_j(z_j) dz_j$$

Maximizing $\mathcal{L}_j$ (and thus the full ELBO with respect to $q_j$, holding others fixed) is equivalent to minimizing the KL divergence $KL(q_j(z_j) || \frac{\tilde{p}_j(z_j)}{C_j})$. The minimum KL divergence value of zero is achieved when $q_j(z_j)$ is exactly equal to the normalized distribution $\frac{\tilde{p}_j(z_j)}{C_j}$. Therefore, the optimal solution for the factor $q_j^*(z_j)$ satisfies:

$$q_j^*(z_j) \propto \exp(\mathbb{E}_{q_{-j}} [\log p(\mathbf{x}, \mathbf{z})])$$

Or, equivalently:

$$\log q_j^*(z_j) = \mathbb{E}_{q_{-j}} [\log p(\mathbf{x}, \mathbf{z})] + \text{constant}$$

This important result provides a recipe for finding the optimal form of each factor $q_j^*(z_j)$, assuming all other factors $q_k(z_k)$ ($k \neq j$) are fixed. It states that the optimal log-density for $z_j$ is obtained by taking the log of the model's joint probability $p(\mathbf{x}, \mathbf{z})$ and then averaging over all other variables $z_k$ ($k \neq j$) according to their current variational distributions $q_k(z_k)$. This forms the basis for iterative update schemes like CAVI, where we cycle through the variables, updating each $q_j$ based on the current estimates of the others.

Implications and Limitations of Mean-Field

The primary strength of the mean-field approximation is computational tractability. By breaking the dependency structure among latent variables within the variational approximation $q$, we convert a potentially complex optimization problem over a high-dimensional distribution into a series of potentially simpler optimizations over lower-dimensional factors $q_j$. If the model structure and choice of $q_j$ families are compatible (e.g., using conjugate priors in the model often leads to recognizable forms for $q_j^*$), these updates can sometimes be derived analytically.

However, this simplification comes at a cost. The central assumption is that the variational posterior factors are independent: $q(\mathbf{z}) = \prod_j q_j(z_j)$. If the true posterior $p(\mathbf{z}|\mathbf{x})$ exhibits significant correlations between the latent variables $z_j$, the mean-field approximation $q(\mathbf{z})$ will fail to capture these dependencies by definition.

Consider a simple 2D Gaussian example. If the true posterior shows strong negative correlation between $z_1$ and $z_2$, the best mean-field approximation $q(z_1, z_2) = q_1(z_1) q_2(z_2)$ will be a Gaussian with zero correlation (its contours will be aligned with the axes), even if it manages to center itself correctly and approximate the marginal variances.

Illustration comparing a true posterior $p(z_1, z_2|\mathbf{x})$ with correlation (dashed orange contours) and its best mean-field approximation $q(z_1)q(z_2)$ (solid blue contours). The approximation centers correctly but enforces independence, thus missing the correlation structure present in the true posterior.

This inability to capture posterior correlations is a well-known characteristic of mean-field VI. Because the KL divergence $KL(q || p)$ penalizes placing probability mass with $q$ where $p$ has none, the resulting $q(\mathbf{z})$ tends to be more "compact" or concentrated around the mode than the true posterior $p(\mathbf{z}|\mathbf{x})$. This often leads to an underestimation of posterior variances and potentially overly confident uncertainty estimates.

Despite these limitations, the mean-field assumption is foundational in variational inference. Its computational advantages are significant, making Bayesian inference feasible for many complex models and large datasets where MCMC methods might struggle with convergence time or memory requirements. Understanding the mean-field approximation, its derivation via the optimal factor updates, and its inherent limitations regarding posterior dependencies and variance estimation is essential for applying VI effectively and interpreting its results critically. We will now explore the Coordinate Ascent Variational Inference (CAVI) algorithm, which directly implements the iterative updates derived from the mean-field assumption.