Back

Mean-Field VI on a Banana

What you are seeing: the simplest variational inference setup, fit to the simplest distribution that breaks it. Target: a Rosenbrock-style "banana" logp(x,y)=((1x)2+100(yx2)2)/20\log p(x, y) = -((1 - x)^2 + 100 (y - x^2)^2)/20, a long curved valley. Variational family: a product q(x,y)=N(μx,σx2)N(μy,σy2)q(x, y) = N(\mu_x, \sigma_x^2)\,N(\mu_y, \sigma_y^2), axis-aligned Gaussian (mean-field). VI maximizes ELBO=Eq[logp]+H(q)\text{ELBO} = E_q[\log p] + H(q) by gradient ascent.

The contour lines are the target logp\log p. The orange ellipse is the current qq. The yellow trace below is the ELBO over training. The "best" Gaussian approximation cannot follow the banana's curvature; it settles into a compact ellipse at the bend.

Figure 1. Mean-field VI on a Rosenbrock banana, reparameterization gradient with K = 32 Monte Carlo samples per step.
lr0.005
K (MC)32
speed5

WHAT TO TRY

  • Watch the axis-aligned Gaussian try to fit the curved banana: mean-field VI cannot tilt, so it collapses onto one part of the valley and underestimates the variance. That is the classic failure mode.
  • Raise K (Monte Carlo samples): the ELBO gradient gets less noisy and the fit settles more smoothly, at higher cost per step. Lower K and it jitters.
  • Tune the learning rate: too high and the variational parameters oscillate around the valley, too low and they crawl. The ELBO trace shows convergence.