Variational Inference

VI fundamentally reframes Bayesian inference. Instead of directly computing the intractable posterior distribution p(z|x) via integration (e.g., Markov Chain Monte Carlo), VI casts inference as an optimization problem. It seeks the best approximation q(z) from a chosen family of simpler distributions Q by minimizing their difference from the true posterior, typically measured by Kullback-Leibler (KL) divergence.

The core objective function in VI is the Evidence Lower Bound (ELBO), a tractable surrogate for the log marginal likelihood (evidence). Maximizing the ELBO is equivalent to minimizing the KL divergence between the variational approximation q(z) and the true posterior p(z|x).

The ELBO decomposes into two key terms:

• Reconstruction Term (Expected Log-Likelihood): Measures how well the model explains the data given the latent variables.
• Regularization Term (KL Divergence): Penalizes the approximation q(z) for straying too far from the chosen prior distribution p(z).

A key efficiency innovation in modern VI (e.g., in Variational Autoencoders). Instead of optimizing a separate set of variational parameters for each individual data point x_i, amortized inference trains a shared neural network (the encoder or inference network) to predict the parameters of q(z|x) for any input. This allows for rapid, single-pass inference on new data after training, scaling to massive datasets.

A common simplifying assumption to make the optimization tractable. The mean-field variational family assumes all latent variables z_j are mutually independent, so the joint variational distribution factorizes: q(z) = ∏_j q_j(z_j). This turns a complex high-dimensional optimization into a set of simpler, lower-dimensional problems, often solved via coordinate ascent. While powerful, it cannot capture posterior correlations between latent variables.

VI leverages stochastic gradient descent to scale to large datasets. The challenge is backpropagating gradients through the sampling of latent variables z ~ q(z).

The reparameterization trick solves this by expressing the random variable z as a deterministic, differentiable function z = g(φ, ε), where ε is an independent noise variable (e.g., ε ~ N(0,1)) and φ are the distribution parameters. This allows gradients to flow through g(·), enabling efficient end-to-end training of models like VAEs.

VI is often contrasted with Markov Chain Monte Carlo (MCMC) methods.

VI Advantages:

• Speed: Typically much faster convergence due to optimization.
• Scalability: Well-suited for large datasets via stochastic gradients.
• Amortization: Fast inference post-training.

VI Limitations:

• Approximation Bias: The choice of variational family Q imposes a fixed form, leading to an approximate, not exact, posterior.
• Underestimation of Uncertainty: Mean-field assumptions often lead to over-confident, compact posteriors.

MCMC is asymptotically exact but can be prohibitively slow for large-scale inference.

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

• Encourages the variational posterior to accurately reconstruct the observed data (via the reconstruction term).
• Regularizes the posterior by minimizing its Kullback-Leibler (KL) divergence from a chosen prior distribution. The ELBO decomposition is: ELBO = 𝔼[log p(x|z)] - KL(q(z|x) || p(z)), where the first term is the expected log-likelihood and the second is the KL regularization.

Kullback-Leibler (KL) Divergence is a non-symmetric measure of how one probability distribution diverges from a second, reference distribution. In variational inference, it acts as a regularizer within the ELBO:

• KL(q(z|x) || p(z)): This term measures the information loss when using the approximate posterior q to represent the true prior p. Minimizing it forces the learned latent distribution to conform to a simple prior (e.g., a standard Gaussian).
• It ensures the latent space is well-structured and prevents posterior collapse, where the model ignores the latent variables.
• The asymmetry is crucial; this reverse KL direction favors approximations that avoid assigning probability mass where the prior has none.

The reparameterization trick is a critical technique that enables gradient-based optimization through random sampling nodes in a computational graph. It transforms a stochastic node z ~ q_φ(z|x) (e.g., a Gaussian with mean μ and variance σ²) into a deterministic function z = g_φ(ε, x) where ε is an independent noise variable (e.g., ε ~ N(0, I)). Key Impact:

• It allows gradients (∇_φ) to flow from the reconstruction loss back through the sampling operation to the encoder parameters φ.
• Without it, stochastic sampling blocks gradient propagation, making VAEs untrainable with standard backpropagation.
• This trick is foundational for training any continuous latent variable model with stochastic gradient descent.

Amortized inference refers to the use of a neural network (the probabilistic encoder) to learn a global mapping from data points x to the parameters of their variational posterior q(z|x). Instead of optimizing latent variables z separately for each data point—a computationally prohibitive process—the encoder learns to perform fast, approximate inference for any input. Benefits:

• Drastically reduces computational cost from O(N) per-datapoint optimization to a single forward pass through the encoder.
• Enables generalization: the encoder can infer posteriors for new, unseen data points instantly.
• This is the "amortization" of the cost of inference across the dataset and the model parameters.

Mean-field approximation is a simplifying assumption used to make variational inference tractable for complex models. It assumes that the latent variables in the variational posterior q(z) are mutually independent, allowing the joint distribution to factorize into a product of individual distributions: q(z) = ∏_i q_i(z_i). Implications:

• It ignores correlations between latent variables, which is a major approximation but makes optimization feasible.
• Each factor q_i(z_i) is typically chosen from a tractable family (e.g., Gaussian).
• In variational autoencoders, this assumption is applied per datapoint: q(z|x) = ∏_i q(z_i|x), often with a diagonal covariance Gaussian, making computation efficient.

Stochastic Variational Inference (SVI) is a scalable optimization framework that combines variational inference with stochastic gradient descent. It enables training on massive datasets by using mini-batches of data to compute noisy, unbiased estimates of the ELBO gradient. Core Mechanism:

1. Sample a mini-batch of data x.
2. Use the amortized inference network to obtain variational parameters.
3. Compute a Monte Carlo estimate of the ELBO gradient using the reparameterization trick.
4. Update model parameters with the stochastic gradient. This methodology is what makes large-scale Bayesian deep learning, including the training of variational autoencoders on datasets like ImageNet, computationally practical.