6 Variational inference: Scaling to large datasets

The chapter introduces variational inference as a faster alternative to MCMC for Bayesian inference on complex models and large datasets. Rather than drawing many samples from a posterior until convergence, variational inference turns inference into an optimization problem: it searches over a simple family of candidate distributions and selects the one that best approximates the true posterior. This makes it especially useful in large-scale machine learning settings where speed, frequent retraining, and real-time prediction matter.

To make this work, the chapter explains how to compare the approximating distribution to the true posterior using KL divergence, and then how to avoid directly computing that hard quantity by maximizing the ELBO instead. The ELBO reformulation uses Bayes’ theorem to express the objective in a more tractable form, allowing optimization with gradient descent. Because sampling breaks differentiability, the reparameterization trick is introduced so samples from a normal approximator can be rewritten as a deterministic transformation of a standard normal draw, restoring gradient flow and enabling automatic differentiation.

The chapter then shows how variational inference is made practical at scale through minibatch training, where each update uses a small random subset of data to estimate the ELBO and its gradient. In PyMC, fitting a variational model looks similar to MCMC, but uses a fit step and then samples from the learned approximation. The chapter closes by emphasizing that VI is usually much faster than MCMC, but often less faithful: it can miss multiple posterior modes and underestimate variance, so its quality depends on the flexibility of the approximating family and the problem at hand.

Unlike MCMC, which generates samples to approximate an unknown posterior, variational inference (VI) fits a simple distribution like a normal distribution to the posterior and uses the fitted distribution as the approximation.

The KL divergence is large when the approximating distribution 𝑞 yields high density where 𝑝 is small. When the two distributions match, the KL divergence decreases to be close to 0.

The evolution of the approximating distribution 𝑞 and how it’s adjusted by gradient descent throughout optimization.

Approximate posteriors produced by MCMC (left) and variational inference (right) in the Bayesian linear regression problem. The two approximations are similar, having the same mean and 94% highest-density interval.

Optimization progress of the ELBO across 200,000 steps. The function’s value flattens out towards the end, indicating convergence to an optimum.

Demonstrations of pathological behaviors in variational inference. Left: mode-seeking causes the approximator to miss out on other modes. Right: variance underestimation results from the KL divergence formulation.

• MCMC is powerful but slow, especially for high-dimensional models or large datasets.
• Variational inference (VI) reframes inference as optimization, fitting a simple distribution to approximate the posterior.
• The KL divergence measures the gap between the approximating distribution and the true posterior; maximizing the ELBO minimizes this gap.
• Gradient descent with automatic differentiation makes optimization possible, but sampling blocks gradients—solved by the reparameterization trick.
• Minibatch VI allows scaling to very large datasets, similar to how deep learning trains.
• Pathological behaviors: VI tends to be mode-seeking and can underestimate variance due to KL’s asymmetry.
• Advanced VI methods (mean-field, structured VI, mixtures, amortized inference, normalizing flows) extend flexibility while balancing speed.
• MCMC vs. VI tradeoff: MCMC is slower but more faithful; VI is faster but approximate. Our choice is based on dataset size and inference needs.

FAQ