Variational Inference · demonstration

Which KL you minimise decides what you get wrong

Watch for this: both fits below are “optimal”. They minimise the same divergence between the same two distributions — only the argument order differs. One comes out too narrow. The other covers ground where the target has no mass at all. Neither is a bug.

What you are looking at

The target p(z) is a distribution over two variables that are coupled — knowing z₁ tells you something about z₂. The approximation q is forced to factorise, q(z₁)q(z₂), so it can never represent that coupling. This is the mean-field assumption, and it is the price of admission for tractable variational inference.

Given that q cannot be right, the question becomes: in which direction should it be wrong? That is decided entirely by which argument of the KL divergence you put it in.

Reverse KL — the one you are actually minimising

Maximising the ELBO is exactly minimising KL(q‖p). The expectation is taken under q, so the penalty is only ever felt where q puts mass. Wherever p ≈ 0 and q > 0, the integrand blows up — so q learns to stay strictly inside the target's support. It is mode-seeking, and on a correlated Gaussian it comes out visibly too narrow.

That narrowness is not approximate folklore; it is exact. For a Gaussian target the optimal factorised q has variance 1/Λᵢᵢ, the reciprocal of the precision diagonal, not the covariance diagonal. For the two-variable case that means the standard deviation shrinks by exactly √(1−ρ²). Set ρ = 0.9 above and the readout will report a width ratio of 0.436 — check it against √(1−0.81).

The practical consequence: a variational posterior reports error bars that are too small, and it does so confidently. If you use VI for uncertainty quantification and take the credible intervals at face value, you will be overconfident by a factor that grows with how correlated your parameters are — and parameters in real models are correlated.

Forward KL — the one you usually cannot compute

Reverse the arguments and the expectation is taken under p. Now the penalty falls wherever p has mass and q does not, so q is punished for leaving anything uncovered. It becomes mass-covering: on the two-mode target it stretches across both modes and places its own peak in the valley between them, where the target has almost no probability at all.

For a Gaussian family this fit has a closed form — match the mean and variance of p, which is why it is called moment matching. It recovers the marginal widths exactly. So why is it not the default? Because computing those moments requires expectations under p, and p is the intractable posterior you were trying to approximate in the first place. Forward KL is the better objective and the one you generally cannot have.

A coincidence worth noticing

On the correlated Gaussian, both directions report the same divergence — 0.8304 at ρ = 0.9. That is not a display error. For a Gaussian both minima equal −½ ln(1−ρ²), which is also the mutual information between z₁ and z₂. Identical cost in nats; two completely different approximations.

So the divergence tells you how much the factorisation throws away. It tells you nothing whatever about which way you will be wrong — and that is the part that decides whether your credible intervals are too small or your posterior has a mode you never saw.

Try this

Set the target to two modes and switch between the two directions. Reverse KL will commit to one mode and ignore the other entirely — a posterior that is confidently, precisely wrong about half the answer. Forward KL spans both and is confidently wrong about the middle. Then ask yourself which failure you would rather explain to a referee.

← Back to teaching