Advanced Variational Families

While the mean-field approximation, where we assume $q(\mathbf{z}) = \prod_i q_i(z_i)$, simplifies the optimization considerably and makes Coordinate Ascent Variational Inference (CAVI) straightforward, it comes at a cost. This factorization enforces independence among the latent variables $z_i$ within the approximating distribution $q$. However, the true posterior $p(\mathbf{z}|\mathbf{x})$ often exhibits strong correlations between variables.

Imagine a posterior distribution over two variables, $z_1$ and $z_2$, that are highly correlated, perhaps resembling a tilted ellipse in a 2D plot. A mean-field approximation, restricted to axis-aligned ellipses (since $q(z_1, z_2) = q_1(z_1)q_2(z_2)$ implies independence), will fundamentally fail to capture this correlation structure.

The mean-field approximation (yellow, axis-aligned) cannot capture the correlation present in the true posterior (blue, tilted).

This mismatch can lead to several issues:

1. Underestimation of Variance: The marginal variances of $q_i(z_i)$ might be forced to be smaller than the true posterior marginal variances to fit within the restricted family.
2. Biased Estimates: While the mean might be captured reasonably well in some cases, the inability to model dependencies can distort the overall shape, potentially biasing predictions or parameter interpretations.
3. Looser ELBO: A less accurate $q$ distribution leads to a larger KL divergence $KL(q(\mathbf{z}) || p(\mathbf{z}|\mathbf{x}))$, meaning the Evidence Lower Bound $\mathcal{L}(q)$ is further from the true log model evidence $\log p(\mathbf{x})$.

To address these limitations, researchers have developed more expressive variational families that relax the strict independence assumption of mean-field VI.

Structured Mean-Field

A relatively simple extension is Structured Mean-Field (also known as Block Mean-Field or Variational Message Passing in some contexts). Instead of assuming full factorization, we partition the latent variables $\mathbf{z}$ into disjoint sets $\mathbf{z}_1, \dots, \mathbf{z}_M$ and assume factorization between these sets, but allow dependencies within each set: $q(\mathbf{z}) = \prod_{k=1}^M q_k(\mathbf{z}_k)$ This allows the approximation $q_k(\mathbf{z}_k)$ to capture correlations among the variables within the $k$-th group. The choice of partitioning is important and often guided by the structure of the probabilistic model itself (e.g., grouping variables that appear together in factors of the joint distribution $p(\mathbf{x}, \mathbf{z})$).

While more flexible than standard mean-field, optimizing structured mean-field approximations can be more complex. The CAVI updates involve expectations with respect to the joint distributions $q_k(\mathbf{z}_k)$ within each block, which might not have simple closed-form solutions unless the block structures are chosen carefully.

Normalizing Flows

A powerful and popular approach for constructing highly flexible variational distributions is using Normalizing Flows. The core idea is to start with a simple base distribution $q_0(\mathbf{z}_0)$ (e.g., a standard multivariate Gaussian) for which we can easily compute densities and draw samples. Then, we transform this simple distribution through a sequence of invertible functions $f_1, \dots, f_K$: $\mathbf{z}_0 \sim q_0(\mathbf{z}_0)$ $\mathbf{z}_1 = f_1(\mathbf{z}_0)$ $\mathbf{z}_2 = f_2(\mathbf{z}_1)$ $\dots$ $\mathbf{z}_K = f_K(\mathbf{z}_{K-1})$ The final variable $\mathbf{z} = \mathbf{z}_K$ follows a potentially much more complex distribution $q_K(\mathbf{z})$. Because each transformation $f_k$ is invertible, we can recover the initial noise $\mathbf{z}_0$ from the final sample $\mathbf{z}$: $\mathbf{z}_0 = f_1^{-1}(f_2^{-1}(\dots f_K^{-1}(\mathbf{z})))$.

Crucially, if the transformations $f_k$ are chosen such that the determinant of their Jacobian matrix, $\det J_{f_k}$, is computationally tractable, we can compute the density of the final distribution $q_K(\mathbf{z})$ using the change of variables formula from probability theory: $\log q_K(\mathbf{z}) = \log q_0(\mathbf{z}_0) - \sum_{k=1}^K \log |\det J_{f_k}(\mathbf{z}_{k-1})|$ Here, $\mathbf{z}_0$ and $\mathbf{z}_{k-1}$ are obtained by inverting the transformations starting from $\mathbf{z}$. The functions $f_k$ are typically parameterized (e.g., using neural networks), and these parameters are optimized to maximize the ELBO.

Examples of flow transformations include:

• Planar Flows: Simple affine transformations followed by element-wise nonlinearities.
• Radial Flows: Transformations that contract or expand space around a reference point.
• Autoregressive Flows (e.g., MAF, IAF): Use autoregressive neural networks to define transformations with triangular Jacobians, making the determinant easy to compute.
• Coupling Flows (e.g., RealNVP, NICE, Glow): Split variables into two parts and transform one part based on the other, leading to easily invertible functions and tractable Jacobians.

Normalizing flows allow $q(\mathbf{z})$ to approximate arbitrarily complex posterior distributions, provided the flow is sufficiently deep and expressive. They significantly improve the flexibility over mean-field approximations, often resulting in tighter ELBOs and better capturing of posterior dependencies.

Implicit Variational Distributions

Another class of advanced families involves Implicit Distributions. These are distributions $q(\mathbf{z})$ from which it is easy to sample, but difficult or impossible to evaluate the probability density function $q(\mathbf{z})$ itself. Often, these are defined through a sampling process: $\mathbf{z} = g(\boldsymbol{\epsilon}, \phi)$ where $\boldsymbol{\epsilon}$ is drawn from a simple noise distribution (e.g., Gaussian) and $g$ is a complex, non-invertible function (like a deep neural network) parameterized by $\phi$.

The challenge with implicit distributions is the entropy term $\mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$ in the ELBO, which requires evaluating the density $q(\mathbf{z})$. Various techniques exist to work around this:

• Density Ratio Estimation: Estimating the ratio $p(\mathbf{x}, \mathbf{z}) / q(\mathbf{z})$ directly, often using adversarial training techniques similar to Generative Adversarial Networks (GANs).
• Alternative Divergences: Optimizing divergences other than the KL divergence used in the standard ELBO, which might not require explicit density evaluation.

Implicit distributions are particularly useful when the primary goal is to generate samples that resemble the posterior, even if density evaluation is not strictly necessary for the downstream task.

Autoregressive Models

Related to some types of normalizing flows, Autoregressive Models explicitly factorize the variational distribution using the chain rule of probability, without necessarily requiring each step to be easily invertible: $q(\mathbf{z}) = \prod_{i=1}^D q(z_i | z_{1}, \dots, z_{i-1})$ Each conditional distribution $q(z_i | z_{<i})$ can be modeled using a flexible function, often a neural network that takes the preceding variables $z_{<i}$ as input and outputs the parameters of the distribution for $z_i$ (e.g., mean and variance if $q$ is Gaussian). Examples include models like MADE (Masked Autoencoder for Distribution Estimation) or PixelCNN/RNN applied to latent variables.

These models can capture arbitrary dependencies, as they do not impose conditional independence assumptions by the chain rule ordering. Sampling requires sequential generation, and density evaluation is typically straightforward by computing each conditional term.

Trade-offs and Considerations

Moving from mean-field VI introduces a trade-off between approximation accuracy and computational cost/complexity:

• Accuracy: Structured mean-field, normalizing flows, and autoregressive models can provide significantly more accurate approximations to complex posteriors compared to standard mean-field, leading to tighter ELBOs and better uncertainty estimates.
• Computational Cost:
  • Structured mean-field requires careful block selection and potentially more complex update steps.
  • Normalizing flows involve computing Jacobians and their determinants, which can be costly, especially for deep flows or high-dimensional latent spaces. The complexity depends heavily on the type of flow transformation used.
  • Implicit models might require adversarial training or complex density ratio estimation.
  • Autoregressive models require sequential computation for sampling or density evaluation (though density evaluation can sometimes be parallelized).
• Optimization: Optimizing the parameters of these more complex families (e.g., the parameters of the transformations in normalizing flows) can be more challenging than optimizing simple mean-field parameters. The optimization might be non-convex with many local optima.

The choice of variational family is an important modeling decision. While mean-field VI provides a scalable and often effective baseline, understanding and utilizing these more advanced families is essential when posterior dependencies are strong or when higher fidelity posterior approximations are required for accurate uncertainty quantification and model performance. Model checking and comparing the tightness of the ELBO across different families can help guide this choice.