A visual exploration of Stein variational inference

One target, many particles, and a vector field with two jobs: seek probability and refuse collapse.

01

The particle view

From a gradient to a population

A score points uphill. A kernel lets nearby particles share that information—and gives them room to remain different.

step 0 / 180 mean log density —
Raw target score
  • score
  • repulsion
  • particle
x₁ x₂

Ask the target

The score points uphill

The normalization constant disappears under a gradient, so an unnormalized posterior still tells every particle which direction increases density.

\[p(x)=\frac{\bar p(x)}{Z},\qquad s_p(x)=\nabla_x\log p(x)=\nabla_x\log\bar p(x)\]

Used alone, score ascent sends each particle toward a mode. It finds peaks, not a distribution.

02

The diagnostic underneath

Move q, break the balance

The Stein operator averages to zero under the target p. Move the approximation q away, and the leftover signal tells us that transport can still improve it.

black p is fixed · blue q is yours to move Stein signal —

A calibration test

At p, every valid field balances

The orange score contribution and blue divergence contribution must cancel when samples come from p. Shift or squeeze q, and this probe usually becomes nonzero; one balanced probe alone is not proof that q = p.

\[\mathcal A_p\phi(x)=s_p(x)^\top\phi(x)+\nabla_x\!\cdot\phi(x)\] \[\mathbb E_{x\sim p}[\mathcal A_p\phi(x)]=0\]

That imbalance is precisely the first-order KL descent available to a transport field:

\[\left.\frac{\mathrm d}{\mathrm d\epsilon} \operatorname{KL}((I+\epsilon\phi)_\#q\,\|\,p) \right|_{\epsilon=0} =-\mathbb E_q[\mathcal A_p\phi]\]
score term divergence term Stein signal
03

A density in the wild

How long until Old Faithful erupts?

Real observations arrive as dots, not a probability law. A kernel density estimate turns 272 geyser records into a smooth target; Stein particles then make that target usable.

step 0 / 160 loading 272 geyser records
eruption duration × next wait
  • observation
  • KDE
  • Stein particle

Drag the vertical line to ask: after an eruption of this duration, how long is the next wait?

Observed phenomenon

Two eruption regimes emerge

Short eruptions tend to precede shorter waits; long eruptions tend to precede longer waits. Density estimation discovers both regimes without class labels.

\[\widehat p_H(x)=\frac1N\sum_{n=1}^{N} \mathcal N(x;x_n,H)\] \[s_{\widehat p}(x)=\sum_{n=1}^{N} \rho_n(x)H^{-1}(x_n-x),\qquad \rho_n(x)=\frac{\mathcal N(x;x_n,H)}{\sum_m\mathcal N(x;x_m,H)}\]

The side profile is a conditional slice, \(\widehat p_H(\text{wait}\mid\text{duration})\). Move the query line and watch a two-dimensional density answer a practical one-dimensional question.