Fokker–Planck, Time Reversal and Diffusion Models

6D.3 Stochastic Differential Equations and Feynman–Kac derived the Fokker–Planck equation $\partial_tp = L^*p$ for the density of a diffusion, found the Gibbs measure as the stationary density of Langevin dynamics, and sketched, in its opening box, how modern image generators run a diffusion backwards in time. This chapter develops the three ideas behind that sketch in full. First, convergence to equilibrium: the relative entropy of the density with respect to the stationary law decreases at a rate given by a relative Fisher information, and a logarithmic Sobolev inequality turns that into exponential convergence. Second, time reversal: a diffusion run backwards is again a diffusion, whose drift is corrected by the score $\nabla\log p_t$ of its own density. Third, learning the score: the score can be estimated from samples alone, by a regression that needs no knowledge of the density, and with it one can sample from the distribution by integrating the reverse dynamics.

Each idea has a twin in Perelman's work, and the chapter ends with the dictionary. Perelman's conjugate heat equation is a Fokker–Planck equation on a manifold moving by the Ricci flow, solved backward in time (6D.3 Stochastic Differential Equations and Feynman–Kac, 9B.7 The Heat Equation on a Manifold); his function $f$, defined by $u = (4\pi\tau)^{-n/2}e^{-f}$, has gradient $-\nabla f = \nabla\log u$, the score of the conjugate heat kernel; the term $\int|\nabla f|^2u$ in his entropy $\mathcal W$ is a Fisher information; and $\mathcal W$ is the derivative of a Nash entropy, in the same way that the Fisher information is the derivative of the Shannon entropy along the heat flow (6A.12 Entropy, Information and Diffusion). The mathematics of how a diffusion forgets its starting point and how it can be made to remember it is the mathematics of Perelman's monotone quantities.

By the end of this chapter you will be able to:

• write the Fokker–Planck equation of a reversible diffusion in gradient form and verify detailed balance;
• prove that the relative entropy $H(p_t\,|\,\pi)$ decreases at the rate of the relative Fisher information, and deduce exponential convergence from a logarithmic Sobolev inequality, with the Bakry–Émery constant for uniformly convex potentials;
• compare the $L^2$ (Poincaré) and entropy (log-Sobolev) routes to equilibrium, and describe hypocoercivity for kinetic Langevin dynamics;
• derive the time-reversal formula for a diffusion with additive noise from the Fokker–Planck equation, and the probability-flow ordinary differential equation with the same marginals;
• prove Tweedie's formula and the equivalence of explicit, implicit and denoising score matching;
• describe score-based diffusion models and their sampling schemes, and what is and is not proved about them;
• relate the Schrödinger bridge problem to entropic optimal transport;
• identify the score, the Fisher information and the Nash entropy in Perelman's $\mathcal W$.

4.1Reversible diffusions and detailed balance

Consider the overdamped Langevin equation on $\mathbb{R}^n$ of 6D.3 Stochastic Differential Equations and Feynman–Kac,

with $Z = \int e^{-V} < \infty$ and Gibbs measure $\pi = Z^{-1}e^{-V}$. Its generator is $L = \Delta - \nabla V\cdot\nabla$ and the Fokker–Planck equation for the density $p_t$ is

The last form (Exercise 4.4) is the key to everything in this section. Writing $h_t = \frac{p_t}{\pi}$, the relative density, the equation becomes $\partial_th = Lh$: the relative density solves the backward Kolmogorov equation. And $L$ is symmetric in $L^2(\pi)$,

a weighted integration by parts. Symmetry of the generator in $L^2(\pi)$ is reversibility: if $X_0 \sim \pi$, then the process run backwards in time has the same law as the process run forwards. In terms of transition densities it is detailed balance, $\pi(x)p_t(x, y) = \pi(y)p_t(y, x)$: at equilibrium, the flux of probability from $x$ to $y$ equals the flux from $y$ to $x$. Physical systems in thermal equilibrium without magnetic fields or external driving are reversible; molecular motors, driven chemical networks and anything powered by an external source are not.

The general diffusion $dX = b\,dt + \sigma\,dB$ is reversible with respect to some $\pi$ exactly when its drift has the form $b = \frac12\operatorname{div}a + \frac12a\nabla\log\pi$, with $a = \sigma\sigma^T$ (Kolmogorov's criterion in the diffusion setting): the drift is a gradient with respect to the geometry defined by the noise. Langevin dynamics is the case $a = 2I$, $\log\pi = -V$.

4.2Entropy and the approach to equilibrium

How fast does $p_t$ approach $\pi$? There are two natural measures of the distance, and two corresponding inequalities.

$L^2$ and the Poincaré inequality. Since $\partial_th = Lh$ with $L$ symmetric and nonpositive in $L^2(\pi)$, the variance $\operatorname{Var}_\pi(h_t) = \int(h_t - 1)^2\,d\pi$ satisfies

If $\pi$ satisfies a Poincaré inequality $\operatorname{Var}_\pi(f) \leq \frac1\kappa\int|\nabla f|^2\,d\pi$, then $\operatorname{Var}_\pi(h_t) \leq e^{-2\kappa t}\operatorname{Var}_\pi(h_0)$: exponential convergence at the rate of the spectral gap $\kappa$ of $-L$. For the Gaussian $\pi = N(0, I)$ ($V = \frac{|x|^2}{2}$) the gap is $\kappa = 1$, the Gaussian Poincaré inequality of 6A.13 Self-Similar Variables and Monotonicity in this normalisation.

Entropy and the logarithmic Sobolev inequality. The $L^2$ distance is finite only for densities with $\int\frac{p^2}{\pi} < \infty$, a strong condition in high dimensions. The relative entropy

is finite much more often, is nonnegative with equality only for $p = \pi$ (Jensen's inequality, 3A.7 Lᵖ Spaces and Jensen’s Inequality), and is the natural quantity in statistics (the Kullback–Leibler divergence) and in statistical physics (free energy above equilibrium).

If $\pi$ satisfies a logarithmic Sobolev inequality with constant $\kappa$,

then $\frac{d}{dt}H \leq -2\kappa H$, and

For the standard Gaussian this is Gross's inequality (6A.12 Entropy, Information and Diffusion) with $\kappa = 1$. When is it true in general? Dominique Bakry and Michel Émery (1985) found a curvature condition.

The proof differentiates the Fisher information along the flow and finds $\frac{d}{dt}I \leq -2\kappa I$, using the Bochner-type identity $\frac12L|\nabla f|^2 - \nabla f\cdot\nabla Lf = |\nabla^2f|^2 + D^2V(\nabla f, \nabla f)$ (6D.5 Diffusions on Manifolds and Curvature develops this calculus); integrating $I$ from $t$ to $\infty$ then gives the inequality. Uniform convexity of the potential is the flat-space form of a lower bound on the Bakry–Émery Ricci curvature $\operatorname{Ric} + \nabla^2V$ (6C.11 Optimal Transport and Ricci Curvature): on a manifold with measure $e^{-V}dV$ the criterion is $\operatorname{Ric} + \nabla^2V \geq \kappa g$. Holley and Stroock (1987) showed that bounded perturbations preserve the inequality, with a worse constant, so log-Sobolev holds for many non-convex potentials, but with constants that can be exponentially bad when $V$ has deep separated wells: a particle must climb over a barrier to equilibrate, and the time to do so grows like $e^{\Delta V}$ (the Eyring–Kramers law of chemical reaction rates).

Why entropy rather than $L^2$. In dimension $n$, a product measure $\pi^{\otimes n}$ satisfies log-Sobolev with the same constant as $\pi$ (the inequality tensorises), and the relative entropy of a product is the sum of the relative entropies. So for systems with many degrees of freedom, the entropy method gives convergence times that do not grow with the dimension, while $L^2$ distances of natural initial data can be exponentially large in $n$. This is why entropy methods dominate the analysis of high-dimensional sampling algorithms and of the kinetic equations of statistical physics.

Hypocoercivity. The underdamped (kinetic) Langevin equation, $dX = V'\,dt$ in position and $dV' = -\nabla U(X)\,dt - \gamma V'\,dt + \sqrt{2\gamma}\,dB$ in velocity, is not reversible in the above sense and has noise only in the velocity, so its generator is degenerate: the entropy dissipation controls only velocity gradients. Cédric Villani's theory of hypocoercivity (2009) showed that the interaction between the transport term and the noise nevertheless produces exponential convergence, by constructing a modified entropy that includes mixed position–velocity derivatives. Kinetic Langevin dynamics, with its momentum, often mixes faster than the overdamped version, and is the basis of the Hamiltonian Monte Carlo methods of computational statistics.

4.3Running a diffusion backwards

Let $X_t$, $0 \leq t \leq T$, solve $dX_t = b(t, X_t)\,dt + \sigma(t)\,dB_t$ with noise $\sigma(t)$ independent of position, and let $p_t$ be the density of $X_t$. Define the reversed process $\bar X_s = X_{T-s}$. Is it a diffusion, and with what coefficients?

The correction $a\nabla\log p$ is the score of the density, multiplied by the diffusion matrix. Its meaning is visible in the simplest example: for Brownian motion started at a point $x_0$, $p_t$ is a Gaussian centred at $x_0$ and the score $\nabla\log p_t = -\frac{x - x_0}{t}$ points back to $x_0$, so the reversed process is a Brownian bridge pulled back to the starting point. For a stationary reversible diffusion, $\bar b = b$ (Exercise 4.6): reversibility in the sense of the previous section is exactly the statement that time reversal changes nothing.

The probability-flow ODE. There is also a deterministic way to move along the same family of densities. The ordinary differential equation

transports the density $p_0$ into exactly $p_t$ at each time (Exercise 4.7), because its continuity equation $\partial_tp = -\operatorname{div}\big((b - \frac12a\nabla\log p)p\big)$ is the Fokker–Planck equation rewritten. It is the deterministic counterpart of the noisy dynamics, with the diffusion turned into a drift down the gradient of $\log p$. For the heat flow ($b = 0$, $a = 2I$) the velocity is $-\nabla\log p$, which is the Wasserstein gradient of entropy of 6C.11 Optimal Transport and Ricci Curvature: the probability-flow ODE of the heat equation is the Otto gradient flow, particle by particle.

4.4Learning the score

To run a diffusion backwards one needs the score $\nabla\log p_t$ at every time. In applications $p_0$ is known only through samples (a collection of photographs, say), and $p_t$ is the law of a noised sample, known only through noised samples. The remarkable fact is that the score can nevertheless be learned by least squares.

Tweedie's formula. Suppose $\tilde X = X + \sigma Z$, with $Z \sim N(0, I)$ independent of $X$, and let $p_\sigma$ be the density of $\tilde X$. Then

(Exercise 4.8). The posterior mean of the clean signal, given a noisy observation, is the observation plus $\sigma^2$ times the score. Herbert Robbins (1956) attributed it to Maurice Tweedie; Bradley Efron (2011) made it popular for empirical Bayes estimation. It says that denoising and score estimation are the same problem: the best mean-square denoiser determines the score, and vice versa.

Score matching. Aapo Hyvärinen (2005) observed that the squared error $\mathbb{E}|s(X) - \nabla\log p(X)|^2$ of a candidate score $s$ can be computed, up to a constant independent of $s$, without knowing $p$:

by integrating the cross term by parts (Exercise 4.9). Pascal Vincent (2011) showed that for noised data one can do better: minimising the denoising score matching loss

over $s$ gives the same minimiser as $\mathbb{E}|s(\tilde X) - \nabla\log p_\sigma(\tilde X)|^2$, because the two losses differ by a constant (Exercise 4.10). The target $-\frac{\tilde X - X}{\sigma^2}$ is computable from a clean sample and its noised version, so the score of the noised distribution is learned by regression on pairs (clean, noisy), which is Tweedie's formula turned into an algorithm.

4.5The Schrödinger bridge

Erwin Schrödinger asked in 1931 a question that connects time reversal with optimal transport. Suppose a cloud of independent Brownian particles is observed with distribution $\mu_0$ at time $0$ and, unexpectedly, $\mu_1$ at time $1$ (not the heat-flow evolution of $\mu_0$). What is the most likely way it got there? By large deviations, the answer is the law on paths that minimises the relative entropy with respect to Brownian motion, among laws with marginals $\mu_0$ and $\mu_1$: the Schrödinger bridge. Its density is a product $\varphi(t, x)\hat\varphi(t, x)$ of a solution of the backward heat equation and a solution of the forward heat equation, the Schrödinger system, and its drift is $\nabla\log\varphi$, a score-like correction, exactly as in time reversal.

Christian Léonard showed (2012–14) that the static version of Schrödinger's problem, the minimisation of $\int c\,d\pi + \varepsilon H(\pi\,|\,\mu_0\otimes\mu_1)$ over plans $\pi$, is entropic optimal transport, the problem solved by Sinkhorn's algorithm (6C.10 Fully Nonlinear Equations and Monge–Ampère), and that as the noise level $\varepsilon \to 0$ the Schrödinger bridge converges to the displacement interpolation of 6C.11 Optimal Transport and Ricci Curvature. So the three objects of this chapter and the last, time-reversed diffusions, entropic optimal transport and Wasserstein geodesics, are aspects of one variational problem, with the noise level as the parameter. Diffusion "bridges" between two data distributions (rather than between data and noise) are now trained in machine learning by exactly this connection.

4.6Perelman's entropy as a Fokker–Planck entropy

Let $g(t)$ be a Ricci flow on a closed manifold and $\tau = T - t$. Perelman's conjugate heat equation for a density $u$ (with respect to the moving volume $dV_{g(\tau)}$),

is the Fokker–Planck equation of Brownian motion on the evolving manifold run backward in time (6D.3 Stochastic Differential Equations and Feynman–Kac); it preserves $\int u\,dV$, the term $-Ru$ compensating the change of volume $\partial_\tau dV = R\,dV$. Write

Then:

• $-\nabla f$ is the score. $\nabla\log u = -\nabla f$. On flat space, for the backward heat kernel based at $x_0$, $f = \frac{|x - x_0|^2}{4\tau}$ and the score $-\frac{x - x_0}{2\tau}$ points to the base point, as in the Brownian bridge.
• $\int|\nabla f|^2u$ is a Fisher information. It is $\int u|\nabla\log u|^2\,dV$, the Fisher information of the density $u$ (6A.12 Entropy, Information and Diffusion). For the Euclidean heat kernel it equals $\frac{n}{2\tau}$ (Exercise 4.11).
• $\int fu$ is an entropy. $\int fu\,dV = -\int u\log u\,dV - \frac n2\log(4\pi\tau)$: minus the Shannon entropy of $u$, shifted by the entropy of the Gaussian of the same scale. Set $\mathcal N(\tau) = \int fu\,dV - \frac n2$, the Nash entropy; it vanishes for the Euclidean heat kernel.
• $\mathcal W$ is the derivative of the Nash entropy. Perelman's entropy is

a combination of the Fisher information (plus the curvature term $\int Ru$), the entropy, and constants. Along the coupled flow (Ricci flow and conjugate heat equation), $\mathcal W = \frac{d}{d\tau}\big(\tau\mathcal N(\tau)\big)$, equivalently $\mathcal N(\tau) = \frac1\tau\int_0^\tau\mathcal W(\sigma)\,d\sigma$ for the conjugate heat kernel; Lei Ni (2004) proved the analogous identity for the linear heat equation on a fixed manifold, and Bamler (2020) made the Nash entropy the basic quantity of his structure theory.

So Perelman's monotone quantity is built from exactly the ingredients of this chapter: a density evolving by a Fokker–Planck equation, its entropy, its score and its Fisher information. The monotonicity of $\mathcal W$ (12A.3 The 𝓦-Entropy) is the geometric analogue of the de Bruijn identity and the convexity of entropy along the heat flow of 6A.12 Entropy, Information and Diffusion, with curvature terms that the Ricci flow is designed to make vanish exactly on shrinking solitons.

4.7History

Fokker (1914) and Planck (1917) wrote the forward equation, and Kolmogorov (1931) both equations in general. Schrödinger posed his bridge problem in 1931–32. Robbins's empirical Bayes paper, with Tweedie's formula, appeared in 1956. Nelson studied time reversal and osmotic velocities in his stochastic mechanics (1966), Anderson proved the reverse-time diffusion formula in 1982, and Haussmann and Pardoux gave the general theorem in 1986. Bakry and Émery introduced their criterion in 1985, and Holley and Stroock their perturbation lemma in 1987. Hyvärinen introduced score matching in 2005, Vincent denoising score matching in 2011, and Villani hypocoercivity in 2009. Léonard connected the Schrödinger problem to optimal transport in 2012–14. Sohl-Dickstein and collaborators proposed diffusion-based generative models in 2015, Ho, Jain and Abbeel's DDPM appeared in 2020 and Song and collaborators' SDE formulation in 2021. Ni's entropy formula for the linear heat equation appeared in 2004, and Bamler's Nash entropy theory from 2020.

4.8Exercises

Solution

$\pi\nabla\frac p\pi = \nabla p - p\nabla\log\pi = \nabla p + p\nabla V$. For $h$: $\partial_tp = \pi\partial_th$ and $\operatorname{div}(\pi\nabla h) = \pi\Delta h + \nabla\pi\cdot\nabla h = \pi(\Delta h - \nabla V\cdot\nabla h)$; divide by $\pi$.

Solution

(a) By 6D.3 Stochastic Differential Equations and Feynman–Kac with $\theta = 1$, $\sigma = \sqrt2$: $X_t \sim N(e^{-t}x_0, (1 - e^{-2t})I)$. (b) With $s = 1 - e^{-2t}$ and $m = e^{-t}x_0$: $H = \frac12\big(n(-e^{-2t} - \log(1 - e^{-2t})) + e^{-2t}|x_0|^2\big)$. (c) The mean term is exactly $\frac12e^{-2t}|x_0|^2$, decaying at the rate $e^{-2t}$. For the variance term, $\phi(s) = s - 1 - \log s$ with $1 - s = \varepsilon = e^{-2t}$ is $-\varepsilon - \log(1 - \varepsilon) = \frac{\varepsilon^2}{2} + \frac{\varepsilon^3}{3} + \cdots$, which decays like $e^{-4t}$, faster than required; that $\phi(1 - \varepsilon)/\varepsilon$ is increasing in $\varepsilon$ gives $\phi(1 - e^{-2t}) \leq e^{-2(t - t_0)}\phi(1 - e^{-2t_0})$. The slowest mode is the mean, the first Hermite mode, as in 6A.13 Self-Similar Variables and Monotonicity.

Solution

Here $b = -x$, $a = 2I$, and $p_t = N(0, I)$ for all $t$, so $\nabla\log p = -x$ and $\bar b = x + 2(-x) = -x = b$. Started at $x_0$, $p_t = N(e^{-t}x_0, (1 - e^{-2t})I)$, $\nabla\log p_t = -\frac{x - e^{-t}x_0}{1 - e^{-2t}}$, and with $t = T - s$, $\bar b(s, x) = x - \frac{2(x - e^{-t}x_0)}{1 - e^{-2t}}$. As $s \to T$ ($t \to 0$) the second term is $\approx -\frac{x - x_0}{t}$, a strong pull toward $x_0$: the reversed process is a bridge that must end at the starting point.

Solution

$\operatorname{div}(pv) = \operatorname{div}(bp) - \frac12\operatorname{div}(a\nabla p)$, since $p\nabla\log p = \nabla p$; this is $-\partial_tp$. For the heat equation $\partial_tp = \Delta p$ from $N(0, s_0I)$: $p_t = N(0, (s_0 + 2t)I)$, $\nabla\log p_t = -\frac{x}{s_0 + 2t}$, and $v = \frac{x}{s_0 + 2t}$. The ODE $\dot x = \frac{x}{s_0 + 2t}$ gives $x(t) = x(0)\sqrt{\frac{s_0 + 2t}{s_0}}$: every particle moves radially outward, scaling the Gaussian, with no randomness at all.

Solution

$\nabla_y\phi_\sigma(y - x) = -\frac{y - x}{\sigma^2}\phi_\sigma(y - x)$, so differentiating under the integral gives the first formula. By Bayes, the conditional law of $X$ given $\tilde X = y$ is $\frac{\phi_\sigma(y - x)\,d\mu(x)}{p_\sigma(y)}$, so $\mathbb{E}[X - y\,|\,\tilde X = y] = \frac{1}{p_\sigma(y)}\int(x - y)\phi_\sigma(y - x)\,d\mu(x) = \sigma^2\frac{\nabla p_\sigma(y)}{p_\sigma(y)}$.

Solution

$\mathbb{E}_p\langle s, \nabla\log p\rangle = \int\langle s, \nabla p\rangle = -\int p\operatorname{div}s$. Expanding the square, $\mathbb{E}|s|^2 - 2\mathbb{E}\langle s, \nabla\log p\rangle + \mathbb{E}|\nabla\log p|^2 = \mathbb{E}|s|^2 + 2\mathbb{E}\operatorname{div}s + \mathbb{E}|\nabla\log p|^2$. The last term does not depend on $s$, so minimising the first two, which need only samples from $p$, minimises the error.

Solution

Write $w = -\frac{\tilde X - X}{\sigma^2}$. By Tweedie, $\mathbb{E}[w\,|\,\tilde X] = \frac{\mathbb{E}[X\,|\,\tilde X] - \tilde X}{\sigma^2} = \nabla\log p_\sigma(\tilde X)$. So $w = \nabla\log p_\sigma(\tilde X) + \eta$ with $\mathbb{E}[\eta\,|\,\tilde X] = 0$. Then $\mathbb{E}|s(\tilde X) - w|^2 = \mathbb{E}|s(\tilde X) - \nabla\log p_\sigma(\tilde X)|^2 + \mathbb{E}|\eta|^2$, the cross term vanishing after conditioning on $\tilde X$. The constant is $C = \mathbb{E}|\eta|^2$.

Solution

(a) $f = \frac{|x|^2}{4\tau}$ and $\nabla\log u = -\frac{x}{2\tau}$. (b) $|\nabla f|^2 = \frac{|x|^2}{4\tau^2}$ and $\int|x|^2u = 2n\tau$ (variance $2\tau$ per coordinate), so $\int|\nabla f|^2u = \frac{2n\tau}{4\tau^2} = \frac{n}{2\tau}$; $\int fu = \frac{2n\tau}{4\tau} = \frac n2$, so $\mathcal N = \frac n2 - \frac n2 = 0$. (c) $\mathcal W = \tau\cdot\frac{n}{2\tau} + \frac n2 - n = 0$, and $\frac{d}{d\tau}(\tau\cdot0) = 0$. The Gaussian is the equality case of every inequality in sight: log-Sobolev, the monotonicity of $\mathcal W$, and the vanishing of the Nash entropy.