Langevin dynamics, by hand

Companion to the foundations page Langevin dynamics: gradient descent plus noise. Back to the notes.

Every widget below runs the same one-line update, $w\leftarrow w-\eta\,\nabla U(w)+\sqrt{2D\eta}\;\xi$ with $\xi$ a fresh standard Gaussian: gradient descent plus a kick. What changes is the potential $U$, the temperature $D$, and what we measure.

1. Particles settle onto the Gibbs law

1500 particles move in the potential you pick. The bars are where they are now (a running histogram); the curve is the Gibbs law $e^{-U/D}/Z$ for the current temperature. Set $D$ to zero for plain gradient descent: each particle stops in the valley it started in. Click the potential to drop all particles at one point.

temperature $D$ step $\eta$

The potential $U(w)$ and the particles on it (dots). Click to drop them all at one spot.

particles (recent)Gibbs $e^{-U/D}/Z$

Try Narrow vs wide and click the narrow (deeper) well to drop every particle there. At $D=0.3$ they stay put: the barrier is $2.96$ above them and the mean escape time is about $5{,}400$ time units, although at equilibrium only $36\%$ of the mass belongs there. Raise $D$ to $0.8$: they leave within about $17$ time units and settle with roughly a quarter in the narrow well. Two separate questions, two separate answers: where the mass ends up is the Gibbs law, in which width counts by the factor $\sqrt{25}=5$; how soon is Kramers' law, set by the barrier. Then push $\eta$ up on Quadratic and watch the histogram come out wider than the curve (§9's step-size bias).

2. The density flows: Fokker–Planck

Instead of particles, follow the density $p(w,t)$ itself, solved on a grid with the Fokker–Planck equation. Time runs on a log scale, so you can see both clocks: the fast one that fits the density to its valley, and the slow one that moves mass over the barrier. Click the plot to restart from a spike at that point.

3. Sampling: the step-size bias and MALA

Here $U=-\log p$ for a target density $p$ and $D=1$, so Langevin should sample $p$. The unadjusted algorithm (ULA, orange) takes the step as it is. MALA (blue) treats it as a proposal and accepts it with the Metropolis–Hastings probability. 3000 chains each; histograms after a burn-in.

4. The shape of the noise

A quadratic loss with one stiff and one flat direction (curvatures $10s$ and $0.1s$, where $s$ is the sharpness slider). With isotropic noise the cloud is the Gibbs law, wide along the flat direction, and the average excess loss is $kD/2$ whatever the curvature. With noise shaped like the Hessian (roughly SGD's), the cloud is round and the excess loss is $c\operatorname{tr}H/2$, which grows with sharpness.

5. Score-based sampling, plain and annealed

The target is a mixture of three Gaussian blobs with weights $0.6$, $0.3$ and $0.1$. Langevin only needs the score $\nabla\log p$ (arrows). Plain Langevin, started uniformly, leaves each particle in whichever basin it began in, so the blob weights come out wrong. Annealed Langevin first samples a version smoothed by noise of size $\sigma$, where the blobs have merged, and lowers $\sigma$ step by step to $0$.

Everything runs in your browser. The same computations, with every number quoted in the notes, are in code/langevin.py.