Mean-Field Approximation

The core assumption of mean-field variational inference is that the latent variables z₁, z₂, ..., zₘ are mutually independent. This allows the complex, intractable true posterior p(z | x) to be approximated by a variational distribution q(z) that factorizes completely:

q(z) = ∏ᵢ qᵢ(zᵢ)

• Each qᵢ(zᵢ) is a distribution over a single latent variable or a small group.
• This factorization turns a high-dimensional inference problem into a set of simpler, lower-dimensional problems that can be optimized iteratively.

Due to the factorization, the Evidence Lower Bound (ELBO) can be optimized using coordinate ascent. The optimal form for each factor qⱼ(zⱼ) is derived by taking the expectation of the joint log-probability over all other variables z₋ⱼ:

log qⱼ(zⱼ) = 𝔼_{q₋ⱼ}[log p(x, z)] + const.*

• This results in an iterative algorithm: update each factor qⱼ in turn while holding the others fixed.
• The update often yields a recognizable distributional form (e.g., Gaussian), making the procedure tractable.

The primary motivation for mean-field approximation is to make Bayesian inference computationally feasible for complex models.

• It transforms integration (required for the true posterior) into optimization.
• The factorized form drastically reduces the number of variational parameters that need to be estimated.
• This enables inference in large-scale models, such as Latent Dirichlet Allocation (LDA) for topic modeling, where exact inference is impossible.

Mean-field is an approximation that introduces bias in exchange for scalability. The independence assumption is its main limitation.

• Pros: Enables fast, scalable inference on massive datasets.
• Cons: It cannot capture posterior dependencies between latent variables. If the true posterior has strong correlations, the mean-field approximation will be poor.
• The approximation error is the KL divergence between the factorized q(z) and the true posterior p(z | x).

In a Variational Autoencoder (VAE), the probabilistic encoder performs amortized mean-field inference. For a data point x, it outputs the parameters (e.g., mean and variance) of a factorized Gaussian variational distribution q(z | x).

• The encoder neural network amortizes the cost of inference by learning a mapping from data to variational parameters.
• The standard VAE uses a fully factorized Gaussian as its variational posterior: q(z | x) = ∏ᵢ 𝒩(zᵢ; μᵢ(x), σᵢ²(x)).
• This is a direct application of the mean-field assumption within a deep learning framework.

Structured mean-field is a generalization that relaxes the fully factorized assumption. Instead of assuming all variables are independent, it allows dependencies within predefined groups of variables while assuming independence between groups.

• Example: q(z) = q(z_A, z_B) q(z_C), where z_A and z_B are allowed to be dependent.
• This provides a more flexible and accurate approximation when the model's conditional dependencies are known, while still maintaining computational benefits over full inference.

The mean-field assumption posits that the latent variables in a model are mutually independent. This allows the complex, joint variational posterior distribution q(z) to be factorized into a product of simpler, independent distributions: q(z) = ∏ᵢ qᵢ(zᵢ).

• Key Benefit: Transforms an intractable inference problem over many correlated variables into a set of simpler, independent optimization problems.
• Mathematical Consequence: The Evidence Lower Bound (ELBO) decomposes into terms that can be optimized coordinate-wise, often via coordinate ascent variational inference (CAVI).

Mean-field approximation is the engine behind scalable variational inference for large datasets and models.

• Amortized Inference: In models like Variational Autoencoders (VAEs), a neural network encoder learns the parameters (mean, variance) of the factorized q distributions for any input, avoiding per-datapoint optimization.
• Stochastic Optimization: The factorized ELBO enables the use of stochastic gradient descent on mini-batches, making it feasible to train on massive datasets. This combination is known as Stochastic Variational Inference (SVI).

In a standard Variational Autoencoder (VAE), the mean-field approximation is applied to the latent space. The encoder outputs parameters for a diagonal Gaussian distribution: q(z|x) = N(z; μ(x), diag(σ²(x))).

• Diagonal Covariance: The "diag" notation explicitly enforces the mean-field assumption—no off-diagonal terms, meaning latent dimensions are independent given the data.
• KL Divergence Calculation: This factorization allows the KL divergence between q(z|x) and the prior p(z) to be computed analytically as a sum over latent dimensions, which is crucial for efficient training.

The simplification introduces a well-known trade-off between computational tractability and model fidelity.

• Limitation: It ignores posterior dependencies between latent variables. In models where variables are inherently correlated, this can lead to underfitting or poor uncertainty quantification.
• Posterior Collapse: In VAEs, a powerful decoder combined with a mean-field posterior can sometimes lead to posterior collapse, where the latent variables are ignored (KL divergence → 0).
• Mitigations: Techniques like hierarchical VAEs, β-VAE, or more structured variational families are used to capture necessary dependencies while retaining some tractability.

Structured mean-field variational inference relaxes the fully factorized assumption by allowing dependencies between subsets of variables, offering a middle ground.

• Application: Used in models like the Variational Graph Autoencoder (VGAE), where a factorized Gaussian is used for node embeddings, but the graph structure implicitly induces relationships.
• Time Series: In Variational RNNs (VRNNs), a mean-field Gaussian is often used at each timestep, but temporal dependencies are captured through the RNN's hidden state.

The term originates from statistical physics, particularly the study of Ising models for magnetic systems. The mean-field approximation there replaces interactions of a spin with all its neighbors with an interaction with an average "mean field."

• Conceptual Parallel: In machine learning, the complex interaction between a latent variable and all others is replaced by an interaction with a simplified, averaged effect.
• Historical Context: This cross-disciplinary borrowing highlights how variational methods in ML are grounded in well-established approximation techniques from physics and statistics.

Variational inference is a Bayesian approximation technique that transforms complex posterior inference into an optimization problem. Instead of computing the exact posterior distribution, it finds the best approximation from a tractable family of distributions by maximizing the Evidence Lower Bound (ELBO).

• Core Idea: Approximate an intractable posterior p(z|x) with a simpler distribution q(z).
• Optimization: Minimizes the Kullback-Leibler (KL) divergence between q(z) and the true posterior.
• Application: Enables scalable probabilistic modeling with deep neural networks, as seen in Variational Autoencoders (VAEs).

The Evidence Lower Bound (ELBO) is the objective function maximized in variational inference. It provides a tractable lower bound to the log marginal likelihood (evidence) of the data.

• Decomposition: ELBO = Expected Log Likelihood - KL Divergence. It balances data reconstruction accuracy with regularization of the latent space.
• Role in VAEs: Directly optimized via stochastic gradient descent. The reconstruction loss corresponds to the expected log likelihood, and the KL divergence regularizes the latent distribution.
• Tightness: The gap between the ELBO and the true log evidence is equal to the KL divergence between the approximate and true posterior.

Kullback-Leibler Divergence is a non-symmetric measure of how one probability distribution diverges from a second, reference distribution. In variational inference, it quantifies the information loss when using the approximation q to represent the true posterior p.

• Formula: (D_{KL}(q || p) = \mathbb{E}_{q}[\log q(z) - \log p(z|x)]).
• In VAEs: The KL term in the ELBO pushes the variational posterior q_φ(z|x) towards the prior distribution p(z), typically a standard normal.
• Properties: Always non-negative, and zero only if the distributions are identical almost everywhere.

Amortized inference is a technique where a single model (like a neural network encoder) is trained to perform fast, approximate inference for any input, rather than optimizing latent variables separately for each data point.

• Efficiency: The probabilistic encoder in a VAE amortizes the cost of inference by learning a mapping from data x to the parameters of q(z|x).
• Contrast: Non-amortized methods would require running an optimization loop for each new x at test time.
• Connection to Mean-Field: In a mean-field variational autoencoder, the amortized inference network outputs the parameters for the factorized variational distribution.

A full-covariance approximation is a more expressive, but computationally heavier, alternative to the mean-field assumption in variational inference. It models dependencies between latent variables by using a multivariate Gaussian with a dense covariance matrix.

• Comparison to Mean-Field: Captures latent variable correlations that mean-field ignores by assuming independence.
• Computational Cost: Requires estimating O(d²) parameters for the covariance matrix for d latent dimensions, versus O(d) for mean-field.
• Use Cases: Employed when posterior correlations are known to be significant, but scalability is limited compared to the mean-field approximation.

Stochastic Variational Inference is a scalable optimization framework that uses stochastic gradient descent on mini-batches of data to perform variational inference on massive datasets.

• Scalability: Enables training on datasets too large to process all at once by using unbiased gradient estimates from data subsamples.
• Natural Gradients: Often employs natural gradients for more efficient optimization in probability space.
• Integration: SVI is the standard training methodology for mean-field and other variational approximations in modern deep learning, forming the backbone of libraries like Pyro and TensorFlow Probability.