Langevin Dynamics for Sampling

You have a score model: for any point, it tells you which direction increases log density. Turning that into samples looks like it should be gradient ascent. It is, plus one term, and that term is the entire difference between finding a mode and sampling a distribution.

The update

Unadjusted Langevin dynamics iterates

Drop the noise and this is gradient ascent, which converges to a mode and stays there. Keep it and the iterates converge in distribution to \(p\) itself as \(\epsilon \to 0\) and \(t \to \infty\). The drift pulls toward high density; the noise pushes back out, and the stationary distribution of that balance is exactly \(p\).

The coefficients are not arbitrary. The \(\epsilon/2\) on the drift against \(\sqrt{\epsilon}\) on the noise is what makes the stationary distribution correct; any other ratio gives a sampler for a different, tempered distribution. Sampling at a lower "temperature" by shrinking the noise term is a common trick for sharper samples and it is explicitly sampling from \(p^{1/T}\), not from \(p\).

Only the score appears, never the density, so a score model or an unnormalised energy is sufficient. This is what makes Langevin the natural partner for energy-based and score-based models.

Why it is not enough on its own

The convergence guarantee is asymptotic and the constants are terrible in high dimension.

Mode mixing. Moving between two modes separated by a low-density valley requires the noise to carry the chain across, and the probability of that per step falls exponentially with the depth of the valley. In image space the modes are effectively disconnected, and a chain started in one will not visit another within any feasible number of steps. A sampler that cannot mix between modes will reproduce whatever mode its initialisation landed in.

Discretisation bias. The correctness result holds as \(\epsilon \to 0\). At any usable step size the stationary distribution is biased. Metropolis-adjusted Langevin adds an accept-reject step that removes the bias exactly, at the cost of an acceptance rate that falls with dimension and a required density evaluation, which an unnormalised model cannot provide without the partition function.

Step size is a tightrope. Too large and the chain is unstable in high-curvature regions; too small and it takes an infeasible number of steps to move anywhere. There is generally no single step size that works across the whole space, because curvature varies enormously between the manifold and its surroundings.

Annealing is what makes it work

The practical fix is to not sample from \(p\) directly. Start with a heavily smoothed version of the distribution, where modes have merged into one broad basin and Langevin mixes easily, and gradually reduce the smoothing, running a few Langevin steps at each level. Early levels place the chain in the right global region; later ones refine detail.

This is annealed Langevin dynamics, and it is the sampler that made score-based generative models work. It is also, viewed from the other direction, exactly what a diffusion sampler does: the noise ladder is the annealing schedule, and the reverse SDE is Langevin dynamics with a time-dependent target. Modern diffusion samplers descend from this, with the correction step made explicit in predictor-corrector schemes.

When it breaks

The score is wrong exactly where the chain starts. Initialising from noise puts the chain in a region the score model never saw during training unless the noise ladder reaches that far. A truncated ladder produces chains that wander before finding the manifold, which shows up as wasted steps and occasional total failures.

Long chains do not fix short-chain problems. If the chain cannot cross between modes, running it ten times longer produces ten times more samples from the same mode. Diagnostics that check within-chain convergence will report success. Running multiple chains from dispersed initialisations and comparing their statistics is the check that catches it.

Temperature tricks are not free. Reducing the noise coefficient produces cleaner, more typical samples and quietly narrows the distribution. When a paper reports both good FID and a modified noise scale, those are related facts.

Gradient magnitude varies by orders of magnitude across the space. Far from the manifold the score is enormous, near it the score is small. A fixed step size therefore takes wildly different effective step lengths at different points, which is another argument for annealing, where each noise level has a well-matched scale.