Variational inference#

Part 3 · Stage 10 · 🎛️ Modal & Variational Approximation · Lesson 087 of 144 · intermediate

◀ Previous · Example: hierarchical normal model (continued) · Next · Expectation propagation ▶ · ↑ Section

Important

✨ AI-generated content. This page was written with the assistance of an AI language model and is provided as a learning aid. Despite careful review, it may still contain mistakes, omissions, or out-of-date information. Whether you are new to the topic, a team lead, or a senior practitioner, treat it as a starting point rather than an authoritative reference: read it critically and independently verify anything you act on (code, commands, figures, and factual claims) against official documentation and primary sources before relying on it.

Turn inference into optimisation#

MCMC samples the posterior; variational inference approximates it, by choosing the member of a tractable family \(\mathcal{Q}\) that is closest to it. Inference becomes optimisation, which scales to models and datasets where sampling is hopeless.

The ELBO#

Closeness is measured by Kullback–Leibler divergence, and the decomposition is exact:

\[\log p(y) = \underbrace{\mathrm{E}_{q}\bigl[\log p(\theta, y)\bigr] - \mathrm{E}_{q}\bigl[\log q(\theta)\bigr]}_{\text{ELBO}(q)} \;+\; \mathrm{KL}\bigl(q(\theta) \,\|\, p(\theta \mid y)\bigr) .\]

Since \(\log p(y)\) is fixed and the KL term is non-negative, maximising the ELBO minimises the divergence — and the intractable evidence never has to be computed. The ELBO’s two pieces read naturally: fit the joint density well, but keep the approximation’s entropy high.

Families, and the mean-field default#

The commonest choice factorises the posterior across parameters — the mean-field family, \(q(\theta) = \prod_k q_k(\theta_k)\) — usually with Gaussian factors on an unconstrained scale. ADVI automates the whole thing: transform constrained parameters, fit a Gaussian, optimise the ELBO by stochastic gradients using automatic differentiation.

import pymc as pm
with model:
    approx = pm.fit(n=30_000, method="advi")     # optimise the ELBO
    idata = approx.sample(2000)                  # draws from q, not from the posterior

The direction of the KL#

Everything that goes wrong follows from which KL is minimised. \(\mathrm{KL}(q \| p)\) is mode-seeking: it heavily penalises putting \(q\) mass where \(p\) has none, and barely penalises the reverse. So \(q\) prefers to cover one mode and underestimate variance and tails. A mean-field family compounds this: forcing independence across correlated parameters shrinks the approximation’s spread further.

The consequences are systematic, not random. VI point estimates are often excellent; VI uncertainties are usually too narrow; and multimodal posteriors are typically represented by one mode. In a hierarchical model, VI will underestimate \(\tau\)’s uncertainty exactly where it matters.

Check it, always#

Unlike MCMC, VI has no \(\hat{R}\) — the optimiser converging says nothing about whether \(\mathcal{Q}\) contains anything close to the posterior. Diagnostics exist and should be used: PSIS importance-reweighting of the VI draws (whose \(\hat{k}\) flags a bad approximation), and simulation-based calibration. Use VI for exploration, for very large models, and as an initialiser — and never report a VI interval without checking it.

Hint