Variational Inference

ELBO, KL divergence, mean field VI, CAVI, stochastic VI — when to use over MCMC

You want to fit a topic model (LDA) to 100,000 documents. The exact posterior P(topics | documents) requires summing over all possible topic assignments for every word in every document — exponential in the number of words. You cannot compute it. MCMC could sample from it, but running chains long enough to converge on 100,000 documents takes hours to days. You need a tractable alternative that scales with the data. This is the problem variational inference solves.

Variational inference reframes posterior computation as optimization. Instead of computing the true posterior p(z|x), you pick a simpler family of distributions q(z; φ) — typically factorized Gaussians or Dirichlet distributions — and find the member of that family closest to the true posterior. "Closest" is measured by KL divergence. You minimize KL(q ‖ p) by maximizing the ELBO (Evidence Lower Bound): ELBO = E_q[log p(x, z)] - E_q[log q(z)]. The ELBO is a lower bound on log p(x). Maximizing it pushes q as close to p as possible while keeping q tractable.

The mechanism: ELBO = log p(x) - KL[q(z) ‖ p(z|x)]. Since KL ≥ 0, ELBO ≤ log p(x) always. Maximizing ELBO is exactly equivalent to minimizing KL[q ‖ p(z|x)] — they are the same objective. Once you can compute and differentiate the ELBO, the posterior approximation reduces to gradient descent. Stochastic VI uses minibatch gradient estimates, reducing cost to O(batch_size) per update. This is how VI scales to millions of examples — making LDA at scale, SVGP, and VAE training all tractable.

NOT this. "Variational inference always gives accurate posteriors." The approximation quality depends entirely on the variational family. Mean-field VI — the fully factorized assumption q(z) = ∏ᵢ qᵢ(zᵢ) — assumes all latent variables are independent. Real posteriors almost always have correlations. Mean-field ignores all of them. The result: predictions are overconfident because the approximate posterior is too narrow. The forward KL[q ‖ p] is mode-seeking — it concentrates on one posterior mode and ignores others, systematically underestimating posterior variance. When posterior accuracy matters more than speed (clinical decision making, scientific inference), use richer variational families (normalizing flows) or MCMC. Mean-field VI is not a universal approximate Bayesian method; it is a fast approximation with known failure modes.

Key points

• The ELBO (Evidence Lower Bound): log p(x) = ELBO(q) + KL[q(z) ‖ p(z|x)]. Since KL ≥ 0, ELBO ≤ log p(x) always. Maximising the ELBO is equivalent to minimising KL[q(z) ‖ p(z|x)]. Written out: ELBO = E_q[log p(x|z)] + E_q[log p(z)] - E_q[log q(z)] — reconstruction quality plus how close q is to the prior. This is why VAEs maximise a lower bound rather than the true likelihood: the true likelihood requires integrating over z, which is the integral we're trying to avoid.
• Mean field VI assumes q(z) = ∏ᵢ qᵢ(zᵢ) — each latent variable is independent in the approximate posterior. Real posteriors almost always have correlations between latent variables. Mean field ignores all of them. The result: each marginal qᵢ can look correct individually while the joint approximation is completely wrong — a diagnostic you only catch by inspecting the joint samples, not the marginals.
• CAVI (Coordinate Ascent Variational Inference) updates each qᵢ(zᵢ) one at a time, holding all others fixed. The optimal update is qᵢ*(zᵢ) ∝ exp(E_{−i}[log p(z,x)]). With conjugate priors, these updates have closed forms that look like EM. Without conjugacy, you need gradient-based VI via the reparameterisation trick — which is what makes VAEs trainable.
• Why VI underestimates uncertainty: VI minimises KL[q ‖ p] = E_q[log q - log p]. This penalises q for putting mass where p is small, but never penalises q for failing to cover regions where p is large. The result: q concentrates on one mode of a multimodal posterior and ignores the rest. The reverse KL[p ‖ q] is mass-covering — it would spread q across all modes — but it's computationally harder. The choice of forward KL is what makes VI tractable and also what makes it systematically overconfident.
• Stochastic VI scales CAVI to large datasets. Standard CAVI sweeps all n points per update — O(n) per iteration. Stochastic VI uses minibatch ELBO gradient estimates with SGD/Adam, reducing cost to O(batch_size) per update. This is how VI runs on millions of examples — used in LDA at scale, SVGP, and the encoder-decoder training loop of VAEs. The price is noisier gradient estimates, but the variance is usually manageable with standard variance reduction.
• VI vs MCMC: MCMC is asymptotically exact — given infinite time, it converges to the true posterior. VI is biased — it converges to the best approximation within the chosen family, not the true posterior. Use VI when n is large and approximate uncertainty is acceptable. Use MCMC when exact uncertainty is the actual deliverable — scientific inference, clinical decision support — and you can afford the runtime. The choice is not about which is "better" but about whether you need exact uncertainty or approximate uncertainty quickly.
• Amortised VI trains an inference network to predict q(z|x) directly from x, rather than running iterative CAVI per datapoint. VAEs use this: the encoder outputs μ(x), σ(x) defining q(z|x) = N(μ(x), σ(x)²). Inference at test time is one forward pass. The cost is that the encoder is a single learned function shared across all datapoints — it trades per-datapoint accuracy for inference speed.
• The most common VI failure: the ELBO converges but posteriors are too narrow and predictions are overconfident, because Q is too restrictive for the true posterior. A diagonal Gaussian mean field cannot represent correlated latent variables. If the ELBO at convergence is much lower than the log evidence from a short MCMC run, the variational gap is large — you need a richer family such as normalising flows or hierarchical VI.

Recap

• VI reframes inference as optimisation: pick tractable $q(z;φ)$, minimise $KL[q | p]$ by maximising the ELBO.
• ELBO = log p(x) − KL[q ‖ p(z|x)] — a lower bound (KL ≥ 0); maximising ELBO ≡ minimising KL to the true posterior.
• Mean-field $q(z)=prod_i q_i(z_i)$ assumes independence — real posteriors are correlated; marginals can look right while the joint is wrong.
• Forward $KL[q | p]$ is mode-seeking — concentrates on one mode, systematically underestimates uncertainty (overconfident).
• Stochastic VI scales it: minibatch ELBO gradients, $O( ext{batch})$/step → runs on millions (LDA, SVGP, VAEs).
• Amortised VI: an inference network predicts $q(z|x)$ in one forward pass (the VAE encoder) — trades per-datapoint accuracy for speed.
• VI biased, MCMC asymptotically exact: VI when scale matters and approx uncertainty is fine; MCMC when exact uncertainty is the deliverable.

Check your understanding