Use the gradient of log p to glide across high-dimensional distributions — the same physics that powers Stan, NUTS, and diffusion models.
You are fitting a Bayesian logistic-regression model with 200 weights. Your Metropolis sampler — the one you built in the MCMC lesson — runs for an hour, and the trace plot looks like a drunk crawling home: it keeps revisiting the same neighborhood and never reaches the far side of the posterior.
You did not break it. Random walks are simply this slow. The handicap is not a bug in your code — it is a property of randomness itself, and no amount of tuning removes it.
Here is the defining fact. A random walker that takes a step of size step in a random direction does not travel a net distance of step per move. Its displacement is a random walk, and after N moves the typical (root-mean-square) net distance is only step×√N — it grows like the square root of the number of steps, not linearly.
That √N is the entire story. To cross a distribution of width L, you need step×√N = L, which means N = (L/step)² steps. Double the distance and you quadruple the work. This is called diffusive exploration, and it is hopelessly slow.
The fix is to stop walking blindly. We have information the random walker ignores: we can differentiate log p. The gradient of log p — the score — points toward higher probability. Following it lets a sampler move ballistically, a fixed O(1) distance per step, crossing in √N-fewer moves.
That single idea — use the gradient of log p — is the engine inside Stan, NumPyro, and even the diffusion models that draw images. This lesson builds it from the ground up. First, let's race a walker against a glider.
Both start at the left edge of the same target. The red walker proposes tiny symmetric jumps (often rejected, bouncing back). The teal glider surfs along the density. Watch the walker's displacement track the diffusive prediction step·sqrt(N) while the glider's grows linearly.
Let the target spread over a distance L = 10, and let the random-walk proposal step be 0.5 per move. Ignore rejections for the moment (they only make it worse). The root-mean-square displacement after N steps is step×√N.
To traverse L = 10 we need step×√N = 10, so √N = 10 / 0.5 = 20, giving N = 20² = 400 steps — just to cross the distribution once.
Meanwhile a sampler that moves a fixed O(1) distance per (more expensive) step crosses in roughly 20 steps. And the walker is even worse than 400 in practice, because some of those proposals get rejected and waste a move entirely.
One more number to feel the curve: after only 100 steps the walker's RMS displacement is 0.5×√100 = 0.5×10 = 5.0 — it has covered half a crossing after a hundred steps. The other half costs another three hundred.
# The diffusive law, by simulation: net displacement ~ step*sqrt(N), NOT step*N import numpy as np step = 0.5; pos = -5.0; acc = 0 # target N(0,1), start far left for n in range(400): prop = pos + step*np.random.randn() # symmetric proposal logr = -0.5*prop**2 + 0.5*pos**2 # log p(prop) - log p(pos) if np.log(np.random.rand()) < logr: # Metropolis accept pos = prop; acc += 1 # after 400 steps you have, on average, crossed the target just ONCE
And the library baseline this lesson is reacting against — and the fix we build toward:
import pymc as pm # pm.sample(step=pm.Metropolis()) -> the slow diffusive baseline # pm.sample(step=pm.NUTS()) -> the gradient-driven fix we are building
Your 2-parameter model sampled fine. The same Metropolis code on a 200-parameter Bayesian neural net rejects 99.9% of proposals unless you shrink the step to a crawl. Nothing in the code changed — only the dimension d.
In MCMC you built the Metropolis proposal and its accept/reject rule; we will not re-derive detailed balance here. The new question is purely geometric: how must the step shrink as the dimension grows, and why?
The reason is the typical set. In high dimensions, almost all of a distribution's probability mass does not sit at the mode — it concentrates on a thin shell at a characteristic radius. A standard normal in d dimensions has essentially all its mass at radius ≈ √d, in a shell whose thickness barely grows.
Now fire an isotropic random proposal from a point on that shell. It points in a random direction in d-dimensional space, and almost surely that direction points off the shell — into the near-empty interior or the near-empty exterior. The proposal lands in a low-probability region and gets rejected.
To stay on the shell, each coordinate's perturbation must be tiny. If you perturb every coordinate by an independent step s, the squared jump length is about d·s² (one s² per dimension, summed). Holding the total jump magnitude fixed forces s to shrink like 1/√d.
That is the curse in one line: the affordable per-coordinate step shrinks as 1/√d, so per-step distance vanishes and the number of steps to explore grows with d. The famous result (Roberts, Gelman & Gilks, 1997) is that optimally tuned random-walk Metropolis accepts about 23.4% of proposals — and even at that sweet spot, you are still crawling.
Drag the dimension dial. The bar shows the fraction of fixed-length isotropic proposals from the center that land inside the high-probability shell (an acceptance proxy). As d rises the shell thins and proposals miss; the teal curve plots the step size s = 1/sqrt(d) that restores acceptance.
The target is a standard normal in d dimensions. A proposal perturbs every coordinate by an independent step s, so the squared jump length is about d·s². To keep the total jump magnitude fixed — say jump² = 1 — as d grows, we solve for s.
Set d·s² = 1, so s = 1/√d. At d = 1: s = 1/√1 = 1.0 — a healthy step. At d = 100: s = 1/√100 = 0.1 — ten times smaller, per coordinate.
So going from a 1-D toy to a 100-D model forces your per-coordinate step down by a factor of ten, before acceptance even has a chance. That is the multiplicative penalty no amount of "run it longer" can repair.
def rwm_accept_rate(d, step, N=2000): x = np.zeros(d); acc = 0 for _ in range(N): y = x + step*np.random.randn(d) logr = -0.5*(y@y) + 0.5*(x@x) # standard normal target if np.log(np.random.rand()) < logr: x = y; acc += 1 return acc/N # fixed step -> rate falls with d ; step = 1/sqrt(d) -> rate stays ~stable
# the optimal RWM scaling: ell ~ 2.38 gives the 0.234 acceptance sweet spot step = ell / np.sqrt(d)
A blind random walker has no idea which way the density rises. But you do: you can differentiate log p. The result is a vector that points toward higher probability, and it is the single quantity that turns blind walking into informed gliding.
Define the score as s(x) = ∇ log p(x) — the gradient of the log-density. For a Gaussian, that gradient is just an arrow pointing back toward the mean, longer the farther out you are. This same object is what diffusion models learn from data; we will meet that connection in Chapter 9.
It is convenient to talk in terms of energy. Define the potential energy U(x) = −log p(x). High-probability regions are low-energy valleys; low-probability regions are high-energy hills. The score and the energy gradient are negatives of each other: ∇U = −s.
So the score points downhill in energy, toward mass; the energy gradient ∇U points uphill, away from it. This sign bookkeeping is the only thing you must keep straight for the rest of the lesson.
Here is the property that makes the score practical: it does not depend on the normalizing constant. Write log p = log p̃ − log Z, where p̃ is the unnormalized density and Z the (usually intractable) normalizer. Because Z is a constant in x, ∇ log Z = 0.
Therefore the score of an unnormalized density is exactly the score of the normalized one. Gradient samplers never need Z — the very quantity that makes Bayesian posteriors hard to normalize simply vanishes under the gradient. That is why HMC works on raw, unnormalized posteriors.
A 2D Gaussian shown as filled contours. Drag the particle anywhere. The teal arrow is the score s=∇log p (toward the mode); the warm arrow is the energy gradient ∇U=−score (uphill). Toggle the view between probability and the energy bowl. "Follow score" slides the particle to the mode.
Take a univariate target p = N(μ=2, σ²=4). Then log p(x) = −(x−μ)² / (2σ²) + const, and differentiating gives the score s(x) = −(x−μ)/σ².
At x = 5: s(5) = −(5−2)/4 = −3/4 = −0.75. It is negative, pointing back toward the mean at 2 — exactly as a force toward mass should.
The potential is U(x) = −log p = (x−μ)² / (2σ²) + const, so its gradient at x=5 is ∇U(5) = (5−2)/4 = +0.75 = −s(5). The two arrows are equal and opposite, as promised.
A 2D sanity check: for a standard normal N(0, I), the score is simply s(x) = −x. At the point (3, −1) the score is (−3, +1) — pointing straight back to the origin. The mode is the unique place where the score is zero.
def score_gauss(x, mu, var): return -(x - mu) / var # grad log N(mu, var) print(score_gauss(5.0, 2.0, 4.0)) # -0.75 # general case: U = -log p ; gradU via autodiff def gradU(x, logp): return -jax.grad(logp)(x) # the force pointing uphill in energy
import jax score = jax.grad(logp) # autodiff gives grad log p for any differentiable model
In 1987, physicists simulating quarks on a lattice — Duane, Kennedy, Pendleton and Roweth — had exactly the high-dimensional sampling problem you have. Their fix came from physics class: give your sample point a mass and a velocity, then let it roll along the energy landscape like a frictionless puck.
The puck's momentum carries it far across the distribution in one shot — ballistically, not diffusively. That is the whole trick of Hamiltonian Monte Carlo (originally "Hybrid Monte Carlo"). Let's build the puck.
HMC augments the position q (your parameter of interest) with a fresh momentum p, drawn from a standard Gaussian. It assigns the momentum a kinetic energy K(p) = p²/2, the familiar ½mv² with unit mass.
The total energy is the Hamiltonian H(q,p) = U(q) + K(p), where U(q) = −log p(q) is the potential energy from Chapter 2. The joint distribution over (q,p) is proportional to exp(−H) = exp(−U)·exp(−K) = p(q)·N(p; 0, 1).
This factorization is the magic. The joint splits cleanly into your target times an independent Gaussian on the momentum. So if you sample the joint and then throw the momentum away, the marginal over q is exactly your target p(q) — untouched.
The momentum exists only to generate distant, high-acceptance moves. You invent it, ride it across the distribution, and discard it. Now watch a frictionless puck conserve energy while position and velocity trade back and forth.
A 1D energy bowl U(q)=q²/2 with a puck (position q) and a momentum arrow p. The stacked bar shows the energy budget: potential U(q) + kinetic K(p) = total H (the white line). "Kick" resamples p from N(0,1); "roll" runs the frictionless dynamics. Watch H stay flat while U and K trade.
Target N(0, 1), so U(q) = q²/2. Use unit-mass momentum, K(p) = p²/2. Start the puck at rest at the edge: q = 1, p = 0. The Hamiltonian is H(q,p) = (q² + p²)/2.
Initial energy: H = (1² + 0²)/2 = 0.5 — all potential, none kinetic. The puck is poised at the lip of the bowl.
For this quadratic U, the exact frictionless dynamics trace a circle in phase space: q(t) = cos t, p(t) = −sin t. At t = π/2 the puck has rolled to q = cos(π/2) = 0 (the mode) with p = −sin(π/2) = −1 (full speed).
Check the energy there: H = (0² + (−1)²)/2 = 0.5 — conserved, now entirely kinetic. The momentum carried the sample from the edge to the center in one smooth glide, with no rejections at all.
def U(q): return 0.5*q*q # -log N(0,1), the potential def K(p): return 0.5*p*p # kinetic, unit mass def H(q, p): return U(q) + K(p) # total energy q = 1.0; p = np.random.randn() # 'kick': fresh momentum from N(0,1) # exact flow for this U: q = cos t, p = -sin t (a circle; H conserved)
# In Stan / NumPyro the momentum (and its mass matrix) is handled internally. # You only supply U(q) = -log p(q). The puck is built for you.
The puck in Chapter 3 rolled along a perfect circle because the target was Gaussian and we could write the flow down by hand. For any real posterior you cannot solve the dynamics exactly — you must integrate them numerically, one small step at a time.
The obvious choice, plain Euler integration, is a disaster here. Euler updates position and momentum using the gradient at the start of the step, and that systematic staleness makes the puck spiral outward: energy leaks in, H grows without bound, and your acceptance dies.
Molecular dynamics solved this decades ago with a scheme so good it is almost magic: the leapfrog integrator. One leapfrog step is three sub-moves: a half-step momentum kick, a full-step position drift, then a second half-step momentum kick — with the second kick using the gradient at the new position.
The half-then-full-then-half pattern is what makes it symplectic: reversible and volume-preserving. Reversible means running the step backward returns you exactly where you started. Volume-preserving means a little cloud of phase-space points keeps its volume, which (with reversibility) is precisely what a valid Metropolis correction will demand in Chapter 5.
The structural payoff is bounded energy error. Euler's energy error drifts — it grows steadily no matter how small the step. Leapfrog's energy error oscillates around the true value and stays bounded forever. It is not just more accurate; it is a different kind of correct.
Let's compute one leapfrog step by hand and watch the energy barely move. The tiny residual we find is exactly the quantity the next chapter's accept/reject step will absorb.
The (q,p) plane with the exact energy contour (a circle for the Gaussian). The teal leapfrog dots hop half-kick / drift / half-kick and stay on the circle. Toggle the red naive-Euler trace to watch it spiral off. Sliders set step size eps and number of steps L.
Target N(0,1): U(q) = q²/2 so ∇U(q) = q. Unit mass. Step size eps = 0.1. Take one leapfrog step from (q,p) = (1.0, 1.0). The three moves are: p½ = p − (eps/2)·∇U(q); q′ = q + eps·p½; p′ = p½ − (eps/2)·∇U(q′).
First half-kick. The gradient at q=1.0 is 1.0, so p½ = 1.0 − (0.1/2)·1.0 = 1.0 − 0.05 = 0.95.
Drift. Move the position by a full step using the half-stepped momentum: q′ = 1.0 + 0.1·0.95 = 1.0 + 0.095 = 1.095.
Second half-kick, using the gradient at the new position ∇U(1.095) = 1.095: p′ = 0.95 − (0.1/2)·1.095 = 0.95 − 0.05475 = 0.89525.
Now check the energy. Before: H = (1² + 1²)/2 = 1.0. After: H = (1.095² + 0.89525²)/2 ≈ (1.199 + 0.8015)/2 ≈ 1.00025. The energy error is only about 0.0002 — tiny and bounded. That tiny residual is exactly what the Metropolis correction in the next chapter absorbs.
def gradU(q): return q # for N(0,1): U=q^2/2, dU/dq=q def leapfrog(q, p, eps, L): p = p - (eps/2)*gradU(q) # half kick for i in range(L): q = q + eps*p # full drift if i != L-1: p = p - eps*gradU(q) # full kick (interior) p = p - (eps/2)*gradU(q) # final half kick return q, p print(leapfrog(1.0, 1.0, 0.1, 1)) # (1.095, 0.89525)
# NumPyro / Stan run a tuned leapfrog with a learned mass matrix -- # you never write the integrator by hand.
You now have all the parts. The score gives the force, momentum gives the push, leapfrog integrates the roll. Assemble them and you get a sampler that proposes a point halfway across the posterior and accepts it 95%+ of the time.
One full HMC iteration is three moves. Resample the momentum from N(0,1) — a fresh draw, exactly the "kick" from Chapter 3. Roll: run L leapfrog steps to propose a distant point. Accept: keep the proposal with probability min(1, exp(−ΔH)), where ΔH = Hnew − Hold.
That acceptance rule is the entire Metropolis correction — and it is just the same Metropolis accept/reject you built in MCMC, but living in phase space. The target there is exp(−H), and because leapfrog is reversible and volume-preserving (Chapter 4), the Hastings ratio collapses to the clean energy-only form exp(−ΔH).
Here is why acceptance is near 1: leapfrog nearly conserves H, so ΔH ≈ 0 and exp(−ΔH) ≈ 1 — even for a move that crosses half the distribution. The expensive, far-reaching proposal is almost always accepted. That is the payoff for all the machinery.
But it is "near 1," not "1." The continuous dynamics conserve H exactly; the discrete leapfrog leaves a small residual ΔH, and that residual is precisely why the accept/reject step must exist. Skip it and your samples are biased by the integrator's leftover error.
And when the step size is too large, leapfrog goes unstable: ΔH explodes, exp(−ΔH) crashes, and the proposal is rejected. Practitioners call these blow-ups divergences, and they are the loudest warning sign that eps needs lowering. Let's run the loop and watch acceptance hold — then break it.
A correlated 2D target with contours. Each iteration: a momentum arrow appears (resample), a teal leapfrog arc sweeps to the proposal, then accept (dot stays teal) or reject (flashes red, snaps back). The side bar shows Hold vs Hnew and the accept probability. Try the "too-big eps" preset to trigger divergences.
Target N(0,1), U(q) = q²/2. Run a leapfrog trajectory of L = 10 steps with eps = 0.1 from (q,p) = (1.0, 1.0). The proposed end state is (qL, pL) = (1.3827, −0.3007). Sanity-check against the exact flow at t = eps·L = 1.0: q = cos1+sin1 = 1.3818, p = −sin1+cos1 = −0.3012 — leapfrog tracks it closely.
Energy before: Hold = (1² + 1²)/2 = 1.0. Energy after: Hnew = (1.3827² + 0.3007²)/2 = (1.9119 + 0.0904)/2 = 1.00114.
The change is ΔH = 1.00114 − 1.0 = 0.00114. Accept probability = min(1, exp(−0.00114)) = exp(−0.00114) = 0.9989 — essentially certain acceptance for a move of distance ≈ 0.38 across the target.
Now crank eps to 1.5. Leapfrog destabilizes, ΔH jumps to about 0.579, and accept = min(1, exp(−0.579)) = 0.5604. Same algorithm, same target — the only thing that changed was the integrator's stability. The energy error is the acceptance penalty: small for good eps, brutal once the integrator destabilizes.
def hmc_step(q, eps, L): p = np.random.randn() # resample momentum H0 = U(q) + 0.5*p*p qs, ps = leapfrog(q, p, eps, L) # roll L steps H1 = U(qs) + 0.5*ps*ps dH = H1 - H0 if np.log(np.random.rand()) < -dH: # accept w.p. exp(-dH) return qs return q # reject -> stay put
import numpyro mcmc = numpyro.infer.MCMC(numpyro.infer.HMC(model), num_warmup=500, num_samples=1000)
HMC has two dials: step size eps and trajectory length L. Step size you can tune to a target acceptance. But L is brutal. Get it wrong and the puck either barely moves or rolls all the way around and comes back to where it started — pure waste.
Picture the problem. Too short an L and the momentum barely carries you off the starting point — you have paid for a leapfrog trajectory and squandered it. Too long an L on a roughly Gaussian target and the puck does a full orbit, returning near the start: you computed dozens of gradients to move almost nowhere. The best L also depends on the local geometry, so a single fixed value is wrong almost everywhere.
Hoffman & Gelman's No-U-Turn Sampler (NUTS, 2014) removes this dial. It grows the trajectory automatically — doubling it forward and backward in time, building a balanced binary tree — and stops the instant the path starts to curve back on itself. It is why you can run Stan or NumPyro without ever setting L.
How does it detect "curving back"? With a U-turn criterion. While the trajectory still makes progress — while the vector from the earliest point to the latest point still points the same way as the momentum — keep extending. The instant that alignment flips negative, the leading momentum now points back toward where you started: a U-turn. Stop and sample.
Concretely the test is a dot product: keep going while (q+ − q−)·p+ > 0, stop when it goes negative. (The full algorithm checks both ends of the doubled tree, but the idea is this one sign.) When it stops, NUTS draws one state from the whole trajectory it built, in a way that keeps the chain valid.
The remaining noise comes from real-world systems: even after NUTS picks L, the step size eps must still be tuned. Stan and NumPyro do that during warmup by dual-averaging to hit a target acceptance (~0.8). NUTS made L automatic; it did not make eps free. Let's see the U-turn it watches for.
A 2D target with a leapfrog trajectory unrolling. The arrow from the earliest point q− to the latest q+ is drawn, along with the leading momentum. "Grow" doubles the path; "auto" runs until the U-turn statistic (q+−q−)·p+ goes negative, then freezes and samples one point (highlighted).
Take the 1D illustration of the criterion: keep extending while (q+ − q−)·p+ > 0; stop when it goes negative, meaning the leading momentum now points back toward the start. Suppose the trajectory has reached q− = −0.5, q+ = 1.3, with leading momentum p+ = −0.4.
First the span: q+ − q− = 1.3 − (−0.5) = 1.8. The trajectory currently spans 1.8 units — it has made real progress.
Now the U-turn statistic: (q+ − q−)·p+ = (1.8)·(−0.4) = −0.72. It is negative, so the leading momentum points back toward the start — the trajectory is about to retrace itself. NUTS stops here and samples a state from what it has built.
The sign is everything. Had p+ been +0.4 instead, the statistic would be (1.8)·(+0.4) = +0.72 — still moving outward — and NUTS would double again. That single dot product, evaluated each doubling, is what replaces the hand-tuned L.
def uturn(q_minus, q_plus, p_plus, p_minus): # stop when EITHER end starts curving back toward the other back_fwd = np.dot(q_plus - q_minus, p_plus) back_bwd = np.dot(q_plus - q_minus, p_minus) return (back_fwd < 0) or (back_bwd < 0) # NUTS doubles the trajectory until uturn(...) is True, then samples from it
import numpyro numpyro.infer.MCMC(numpyro.infer.NUTS(model), num_warmup=500, num_samples=1000).run(key, data) # L is tuned per draw; you never set the trajectory length
Full HMC trajectories are great, but each one costs L gradient evaluations. What if you take just one noisy gradient step?
You get Langevin dynamics: the simplest gradient-driven sampler. A Langevin step is a drift in the direction of the score plus a dose of Gaussian noise: y = x + (eps/2)·score(x) + √eps·z, with z ~ N(0,1). The drift pulls toward higher density; the noise keeps it exploring rather than collapsing onto the mode.
You can read this as the L = 1 limit of HMC: resample momentum, take a single leapfrog half-roll, and the momentum's Gaussian-ness becomes the √eps·z noise term. It is the bridge to diffusion models in the next chapter, where exactly this update samples images.
But one step breaks energy conservation, and that breaks HMC's tidy acceptance rule. A single Langevin step is an asymmetric proposal: the drift depends on where you start, so proposing y from x is not the mirror of proposing x from y. The forward and reverse drifts differ.
So the Metropolis-Adjusted Langevin Algorithm (MALA) must use the full Metropolis-Hastings ratio — the one with proposal densities — not HMC's energy-only form. The accept ratio is min(1, [p(y)·q(x|y)] / [p(x)·q(y|x)]), where q(·|·) is the Gaussian proposal density centered at the drifted mean.
That correction makes MALA exact. Drop it — never reject — and you get Unadjusted Langevin (ULA): faster, but biased. Chapter 8 measures that bias precisely. First, let's derive MALA's asymmetric acceptance by hand.
A 1D target. Each step is a drift (eps/2)·score plus a noise jiggle √eps·z. MALA computes the asymmetric accept ratio and may reject; ULA takes the same step but never rejects. "Run" overlays the two histograms on the true density — MALA matches, ULA drifts wider as eps grows.
Target N(0,1), score(x) = −x, step eps = 0.5. MALA proposes y ~ N(x + (eps/2)·score(x), eps), an asymmetric proposal. Take the current point x = 0.5 and an uphill proposed move to y = 1.6 (away from the mode, so acceptance is below 1).
Density term. log p(y) − log p(x) = −0.5·(1.6² − 0.5²) = −0.5·(2.56 − 0.25) = −0.5·2.31 = −1.155. The proposed point is less probable, which hurts acceptance.
Forward proposal q(y|x): the drifted mean from x=0.5 is 0.5 + (0.5/2)·(−0.5) = 0.5 − 0.125 = 0.375. So log q(1.6|0.5) = −0.5·ln(2π·0.5) − (1.6 − 0.375)²/(2·0.5) = −2.0730.
Reverse proposal q(x|y): the drifted mean from y=1.6 is 1.6 + (0.5/2)·(−1.6) = 1.6 − 0.4 = 1.2. So log q(0.5|1.6) = −0.5·ln(2π·0.5) − (0.5 − 1.2)²/(2·0.5) = −1.0624.
Assemble. log-accept ratio = (log p(y) + log q(x|y)) − (log p(x) + log q(y|x)) = (−1.155 + (−1.0624)) − (0 + (−2.0730)) = −2.2174 + 2.0730 = −0.1444. Accept = min(1, exp(−0.1444)) = 0.8656. ULA would take the same drifted mean 0.375 and always accept — fast, but biased.
def mala_step(x, eps): g = -x # score of N(0,1) y = x + 0.5*eps*g + np.sqrt(eps)*np.random.randn() def logq(to, frm): # Gaussian proposal density (drifted mean) m = frm + 0.5*eps*(-frm) return -((to-m)**2)/(2*eps) logr = (-0.5*y*y + logq(x,y)) - (-0.5*x*x + logq(y,x)) return y if np.log(np.random.rand()) < logr else x
# ULA = drop the accept step entirely (biased but cheap): x = x + 0.5*eps*score(x) + np.sqrt(eps)*np.random.randn()
When your dataset has a billion points, even one full-gradient evaluation per step is too expensive — let alone an accept/reject pass that needs the full likelihood. Welling & Teh's Stochastic Gradient Langevin Dynamics (SGLD, 2011) drops both: it uses a minibatch gradient and skips the correction, sampling Bayesian deep-net posteriors at SGD speed.
The engine underneath is Unadjusted Langevin (ULA): the Langevin step from Chapter 7 with the Metropolis correction thrown away. ULA never rejects, so it scales beautifully — but it pays a price, and the price is a bias in the distribution it actually samples.
Here is the crucial, counter-intuitive part: the bias is not a finite-sample error. ULA's stationary distribution — the one it converges to — is itself wrong by O(eps). Running longer only nails down the wrong distribution more precisely. Only shrinking eps (or adding the Metropolis correction, i.e. switching to MALA) reduces it.
We can compute the bias exactly for a Gaussian. For target N(0,1), the ULA update is the linear map x′ = (1 − eps/2)x + √eps·z. That is an AR(1) process: a contraction by a factor a = 1 − eps/2 plus noise of variance eps.
An AR(1) process x′ = a·x + (noise variance eps) has stationary variance eps/(1 − a²). Since the true variance is 1 but the ULA variance is eps/(1 − a²), the gap from 1 is exactly the discretization bias — and it scales with eps.
SGLD lives with this by decaying eps toward zero over training, so the asymptotic bias vanishes while early large steps still explore fast. That schedule — not a magic correction — is what makes minibatch Bayesian deep learning work. Let's measure the bias.
The teal ULA histogram over the true density (variance 1). At large eps the histogram is visibly too wide. Lower eps and it tightens. The gauge compares the empirical variance to the closed-form prediction eps/(1−(1−eps/2)²).
Target N(0,1) (true variance = 1). The ULA update is x′ = (1 − eps/2)x + √eps·z, an AR(1) with a = 1 − eps/2 and noise variance eps. Its stationary variance is eps/(1 − a²). Use eps = 0.1.
First the contraction: a = 1 − 0.1/2 = 0.95. Then a² = 0.9025, so 1 − a² = 0.0975.
The ULA stationary variance is eps/(1 − a²) = 0.1 / 0.0975 = 1.0256 — about 2.56% too large versus the true variance of 1. That over-dispersion is the discretization bias, plain as day.
Halve eps to 0.05 and the bias roughly halves too (it is O(eps)); take eps → 0 and ULA becomes exact but infinitely slow. For one concrete noisy step from x=2 with eps=0.1 and a fixed z=1: x′ = 2 + (0.1/2)(−2) + √0.1·1 = 2 − 0.1 + 0.3162 = 2.2162.
def ula(x0, eps, N): x = x0; xs = [] for _ in range(N): x = x + 0.5*eps*(-x) + np.sqrt(eps)*np.random.randn() xs.append(x) return np.array(xs) print(ula(0., 0.1, 100000).var()) # ~1.026, not 1.0 (the O(eps) bias)
# SGLD: replace -x with a minibatch -grad and DECAY eps over training eps_t = a / (b + t)**0.55 # Welling & Teh 2011 schedule
Everything you built — score, Langevin step, noise schedule — is the engine of image generators like Stable Diffusion. This is the flagship payoff of the lesson, and it is exact, not an analogy.
Song & Ermon (2019) showed that if a neural net learns the score ∇ log p(x) of your data, you can sample new data by running exactly the Langevin update from Chapter 7. The catch is that a single noise level fails, so they anneal across many. Let's make the link concrete.
Start with the key identity. Perturbing a clean data point x0 with Gaussian noise gives pσ(x) = N(x; x0, σ²). Its score is ∇ log pσ(x) = (x0 − x)/σ² — and notice that (x0 − x) is exactly the denoising direction, pointing from the noisy observation back toward the clean point.
So the score is the denoiser in disguise. A network trained by denoising score matching — given a noisy x, predict the clean direction — is learning precisely (x0 − x)/σ², the score of the perturbed data.
Why anneal? At a single low noise level the score is sharp but the modes are isolated, so Langevin gets stuck in one blob — the same multimodality trap that broke naive samplers. Annealing fixes it: start at high noise (where the smeared density is nearly unimodal and easy), then lower the noise level by level, letting samples condense onto the data manifold. The step size scales with σ² so high-noise levels take big strides and low-noise levels take fine ones.
That annealed-Langevin sweep — from large noise down to small, taking Langevin steps at each level — is the sampling loop of a score-based diffusion model. The denoising you see in an image generator is line-for-line the Langevin update from this lesson. Watch a cloud of noise condense into a swirl.
Particles start as pure noise and run annealed Langevin along the (analytically known) perturbed score. The noise dial sweeps high→low; particles de-noise from a blob into the swirl. Toggle "single level" to see Langevin get stuck (fails) vs annealed (succeeds).
Perturbing a data point x0 with Gaussian noise gives pσ(x) = N(x; x0, σ²), whose score is (x0 − x)/σ². Take a clean point x0 = 0, a noisy observation x = 3, and noise level σ = 2.
The score is (x0 − x)/σ² = (0 − 3)/2² = −3/4 = −0.75. It is negative — it points from the noisy x=3 back toward the clean x0=0, scaled by 1/σ². A network trained by denoising learns exactly this map: from a noisy x, predict (x0−x)/σ².
Sampling then runs the Langevin update x ← x + (α/2)·score + √α·z across a geometric noise schedule. For example, σ from 10 down to 0.01 in 10 levels has a constant ratio between consecutive levels.
That ratio is (10/0.01)1/9: the first-to-second ratio is np.geomspace(10, 0.01, 10)[0] / [1] = 2.1544. Each level multiplies the noise by about 0.464 (divides by 2.1544), with the step α scaled like σ² so high-noise levels take big strides and low-noise levels take fine ones. That annealed Langevin sweep is the sampling loop of a score-based diffusion model.
def perturbed_score(x, x0, sigma): return (x0 - x) / sigma**2 # denoising direction = score def annealed_langevin(x, score_net, sigmas, alpha0): for sig in sigmas: # high -> low noise a = alpha0*(sig/sigmas[-1])**2 # step scaled by sigma^2 for _ in range(10): x = x + 0.5*a*score_net(x, sig) + np.sqrt(a)*np.random.randn(*x.shape) return x
# Stable Diffusion / score-SDE: the sampler IS this annealed-Langevin # (or its continuous SDE) loop over a learned score.
Time to put it all together. Here is a nasty correlated 2D posterior — the kind that breaks naive samplers. We race three methods on it: random-walk Metropolis, MALA, and full HMC, all on the same target.
Watch not just the trail but the effective sample size (ESS) each one earns. This is the experiment that made Stan the default tool for applied Bayesian inference.
The honest scorecard is not iteration count — it is ESS. Correlated draws carry redundant information: consecutive samples that look alike count for less than one independent sample each. For an AR(1)-like chain with lag-1 autocorrelation ρ, the effective sample size is ESS ≈ N·(1 − ρ)/(1 + ρ).
Crank the correlation slider and watch the random walk's trail collapse into a tight diffusive scribble while HMC's leaps across the long axis of the correlation. The ESS gap between them widens dramatically as the target gets harder. The "same budget" toggle gives each method equal gradient/likelihood work so the race is fair.
This is the whole lesson in one arena: the gradient is the difference between a sampler that crawls and one that flies. Play with it — there is no quiz, the sim is the test.
Each pane runs its own chain over the same correlated 2D Gaussian; trails drawn in the contours. Below each: live acceptance, lag-1 autocorrelation, and effective sample size ESS = N·(1−rho)/(1+rho). Crank correlation to widen the gap.
Effective sample size for an AR(1)-like chain with lag-1 autocorrelation ρ is ESS ≈ N·(1 − ρ)/(1 + ρ). On a strongly correlated target, RWM has high autocorrelation ρ ≈ 0.9; a well-tuned HMC chain has ρ ≈ 0.1. Compare 10,000 raw draws each.
RWM efficiency factor: (1 − 0.9)/(1 + 0.9) = 0.1/1.9 = 0.05263. So 10,000 RWM draws are worth only 0.05263·10,000 ≈ 526 effective independent samples.
HMC efficiency factor: (1 − 0.1)/(1 + 0.1) = 0.9/1.1 = 0.81818. So 10,000 HMC draws are worth ≈ 8,182 effective samples — roughly 15× more information from the same number of iterations.
This is why iteration count lies and ESS tells the truth: the random walk's autocorrelation throws away about 95% of its draws. A chain that reports "10,000 samples" can be carrying the information of a few hundred — ESS is the only honest measure.
def ess(chain): x = chain - chain.mean() rho1 = (x[1:]*x[:-1]).sum() / (x*x).sum() # lag-1 autocorrelation return len(chain)*(1-rho1)/(1+rho1) # run rwm / mala / hmc on the same target, same budget; compare ess(...) and ess/sec
import arviz as az az.ess(idata) # and az.summary(idata) reports ESS + R-hat per parameter
Why does HMC sometimes still struggle? Not because of dimension per se, but because of shape. A posterior stretched 100× in one direction and squeezed in another forces a tiny step size in the tight direction and endless steps in the loose one.
The textbook horror is Neal's funnel: a varying-curvature trap where the right step size at the wide mouth is catastrophically wrong at the narrow neck. This lab lets you sculpt the geometry and watch HMC adapt or fail, revealing what Stan's warmup is really doing.
The single number that summarizes the difficulty is the condition number κ: the ratio of the widest scale to the narrowest. A circle has κ = 1; a needle-thin ridge has a huge κ. The cost of HMC scales with the square root of κ — you need on the order of √κ leapfrog steps to traverse the elongated direction.
The reason eps cannot simply grow to compensate: the step size is capped by the tight direction. Take a step big enough for the loose direction and you overshoot the tight one, the integrator destabilizes, and you get divergences. So you are stuck taking small steps, many of them.
The fix is geometry, not patience. A learned mass matrix M rescales the momentum to match the covariance — it whitens the posterior so the needle becomes a circle in the transformed space, dropping κeff toward 1 and the required steps toward 1. This is precisely what Stan's and NumPyro's warmup estimates.
Stretch the covariance below, watch divergences appear, then turn on the mass matrix and watch them vanish. There is no quiz — sculpting the geometry and seeing the chain rescued is the test.
Reshape the target from a circle to a needle with the scale sliders and rotation (condition number updates). An HMC chain samples it live. Toggle the mass matrix to whiten the geometry — the needle straightens into a circle in the transformed view and divergences drop.
Take a 2D Gaussian with covariance diag(1, 100): one direction has scale 1, the perpendicular direction has scale 10 (variance 100). The condition number is the ratio of largest to smallest variance, κ = 100/1 = 100.
The number of leapfrog steps to traverse the elongated direction scales like √κ. So √κ = √100 = 10: HMC needs on the order of 10 leapfrog steps per trajectory just to cross the long axis once.
And the step size eps is capped by the tight direction (scale 1), so it cannot simply take bigger steps to cover the long axis faster. You are forced into many small steps — the geometry, not the dimension, sets the cost.
A learned diagonal mass matrix M = diag(1, 1/100) rescales the momentum so the transformed target becomes isotropic (κ → 1), cutting the required steps back toward 1. This is exactly what Stan's and NumPyro's warmup estimates: a mass matrix that whitens the posterior so NUTS sees a sphere, not a needle.
import numpy as np Sigma = np.diag([1.0, 100.0]) kappa = np.linalg.cond(Sigma) # 100 n_leapfrog = int(np.sqrt(kappa)) # ~10 steps to traverse M = np.linalg.inv(Sigma) # mass matrix whitens momentum -> kappa_eff ~ 1
# Stan / NumPyro warmup auto-estimates the mass matrix. # For funnels, use a non-centered reparameterization instead.
You started with a random walk that never arrived and ended holding the engine behind Stan, NumPyro, and Stable Diffusion. The thread was always the same: use the gradient of log p.
That one idea — the score — unifies a surprising sweep of modern ML. NUTS uses it for Bayesian inference; SGLD uses it for Bayesian deep learning at scale; score-based diffusion uses it for generation. Click the map below to recall what each branch is and where it came from.
The map is also a quantitative cheat-sheet. Each method carries the number that defines it: HMC's near-1 acceptance and ~15× ESS advantage, ULA's O(eps) bias, the √κ leapfrog-step rule. Together they are a field guide.
And the field guide has a clean decision rule. The difficulty scales with the condition number κ, and the leapfrog cost scales like √κ. Pick NUTS when you can differentiate and the data fits; SGLD when the data is too big for full gradients; diffusion when you are generating rather than inferring.
But know the limit. Gradient samplers dominate only when the target is differentiable. Discrete parameters, combinatorial structures, and black-box likelihoods have no usable ∇ log p — there you fall back to the Metropolis-Hastings and Gibbs methods of the MCMC lesson, often hybridized.
Knowing when the gradient is unavailable is as important as knowing how to use it. The Monte Carlo family is a toolkit, not a single hammer. Here is the map, the comparison table, and the doors this idea opens.
Click a node to recall its method, its defining number, and its real-world system. The center is grad log p (the score); branches are HMC/NUTS, Langevin/MALA/ULA, SGLD, and score-based diffusion.
A practical decision rule that summarizes the whole lesson: difficulty scales with the target's condition number κ, and the leapfrog steps (hence cost) scale like √κ. Compare a well-conditioned posterior to an ill-conditioned one.
Well-conditioned, κ = 4: √4 = 2 leapfrog steps per trajectory — cheap, NUTS sails. Ill-conditioned, κ = 100: √100 = 10 steps — 5× costlier, and you should reach for a mass matrix or reparameterization first.
This single √κ heuristic, combined with the ESS-not-iterations rule from Chapter 10 and the O(eps) bias rule from Chapter 8, is the whole field guide in three numbers.
Pick NUTS when you can differentiate and the data fits; SGLD when it doesn't; diffusion when you are generating rather than inferring; and classical Metropolis-Hastings/Gibbs when the gradient simply does not exist.
| Method | Correction | Best when | Real system |
|---|---|---|---|
| HMC / NUTS | exp(−ΔH), auto-L | differentiable, data fits | Stan, NumPyro |
| MALA | full Hastings ratio | cheap single-step, exact | research / hybrids |
| ULA / SGLD | none (O(eps) bias) | huge data, minibatch | Bayesian deep nets |
| annealed Langevin | none (learned score) | generation | Stable Diffusion |
| Metropolis / Gibbs | symmetric / conditional | no gradient (discrete) | MCMC lesson |
# Decision guide as code if differentiable and data_fits_in_memory: use = 'NUTS (Stan/NumPyro)' elif differentiable and huge_data: use = 'SGLD (decaying eps, minibatch score)' elif generating_new_data: use = 'annealed Langevin / diffusion' else: use = 'Metropolis-Hastings / Gibbs (see mcmc)'