Variational inference: turning Bayes into optimization

Replace integration with optimization

Variational inference picks a tractable family of distributions q(w; φ) — say, independent Gaussians on every weight — and then searches within it for the member closest to the true posterior, closeness measured by KL divergence. We have turned an impossible integral into a familiar optimisation over the variational parameters φ. The catch: KL to the true posterior still contains the intractable marginal likelihood. The fix is the central object of the whole approach.

Rearranging that KL gives the evidence lower bound (ELBO): the expected log-likelihood under q minus the KL from q to the prior. Maximising the ELBO simultaneously minimises the KL to the true posterior, without ever touching the marginal likelihood. Read the two terms as a tug-of-war: the first pulls q toward weights that fit the data; the second is a complexity penalty pulling q back toward the prior. The ELBO is where Bayes and regularised optimisation become the same sentence.

\mathcal{L}(\phi)=\mathbb{E}_{q(w;\phi)}\!\left[\log p(\mathcal{D}\mid w)\right]-\mathrm{KL}\!\left(q(w;\phi)\,\|\,p(w)\right)

The ELBO: maximizing the expected log-likelihood under q while staying close to the prior is exactly minimizing the KL to the true posterior.

The reparameterization trick

To optimise the ELBO with SGD you need its gradient with respect to φ, but φ sits inside an expectation over the random weights — and you cannot push a gradient through a sampling step. The reparameterization trick is the unlock: instead of sampling w directly from q, sample fixed noise ε from a parameter-free distribution and compute w as a deterministic, differentiable function of φ and ε. The randomness is shoved into ε, which carries no gradient, so backprop flows cleanly through the rest.

\nabla_{\phi}\,\mathbb{E}_{q(w;\phi)}\!\left[f(w)\right]=\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I)}\!\left[\nabla_{\phi}\,f(\mu+\sigma\odot\epsilon)\right]

The reparameterization trick: write the random weight as mu + sigma odot epsilon so the gradient slides inside the expectation and flows to phi.

# Gaussian weight w ~ N(mu, sigma^2), reparameterized
eps = randn_like(mu)          # noise, no gradient
sigma = softplus(rho)         # keep sigma > 0
w = mu + sigma * eps          # differentiable in (mu, rho)
loss = -log_likelihood(w, batch) + kl(mu, sigma)  # one MC sample of the ELBO

One Monte-Carlo sample of the negative ELBO; gradients flow to mu and rho through w.

This is the very same trick that powers the variational autoencoder — there it reparameterizes a latent code, here a weight — which is why VAEs are the cleanest place to first meet it. A single noise sample per step gives a noisy but unbiased gradient of the ELBO; in practice one sample is usually enough because SGD is already averaging over minibatches.

Mean-field, and where it hurts

The simplest q treats every weight as independent — the mean-field approximation. It is cheap (two numbers, a mean and a variance, per weight) and it is the default in most libraries. But independence is a strong lie. Weights in a network are richly correlated, and forcing q to factorise makes it systematically underestimate variance: KL(q‖p) penalises q for putting mass where the posterior has none, so q tucks itself inside one mode and reports false confidence.

Scaling and amortizing

Computing the ELBO over the whole dataset every step is impossible at scale. Stochastic variational inference subsamples a minibatch, scales its likelihood term up to estimate the full sum, and runs plain stochastic optimisation on the result — the same idea that lets SGD train deep nets, now driving Bayesian inference. This is what makes variational Bayes feasible on datasets and models of modern size at all.

There is a second kind of scaling. When you must infer a per-input latent (as in a VAE) you would otherwise re-run an optimisation for every data point. Amortized variational inference instead trains a single inference network that maps any input straight to its variational parameters, paying the optimisation cost once and reusing it forever. You trade a little accuracy — the 'amortization gap' — for the ability to do approximate inference in one forward pass at test time.

Beyond Gaussians: flow posteriors

If the bottleneck is the shape of q, enrich it. Normalizing-flow posteriors start from a simple base distribution and push it through a chain of invertible, differentiable maps; the change-of-variables formula tracks the density exactly, so you can still evaluate the ELBO. The result is a variational family that can curve, skew, and correlate — capturing structure mean-field throws away — at the price of more parameters and more careful engineering.

w=f_K\circ\cdots\circ f_1(z_0),\qquad \log q_K(w)=\log q_0(z_0)-\sum_{k=1}^{K}\log\left|\det\frac{\partial f_k}{\partial z_{k-1}}\right|

A normalizing-flow posterior pushes a simple base through invertible maps; the log-density picks up a log-determinant-Jacobian term at every step.