Diffusion Models

The general forward SDE. Song et al. (2021) use the "variance-preserving" SDE $dX = -\tfrac12\beta(t)X\,dt + \sqrt{\beta(t)}\,dW$ with a noise schedule $\beta(t)$. This page is the constant case $\beta = 2$. Its solution is $X_t = a_t X_0 + \sqrt{1 - a_t^2}\,Z$ with $a_t = e^{-\frac12\int_0^t \beta(s)\,ds}$.

The discrete version (DDPM). Ho, Jain & Abbeel (2020) use steps $x_k = \sqrt{1 - \beta_k}\,x_{k-1} + \sqrt{\beta_k}\,z_k$, a discretisation of the same SDE (compare Euler–Maruyama, section 5.1). The idea of learning to reverse a noising process goes back to Sohl-Dickstein et al. (2015).

References: J. Sohl-Dickstein et al., ICML (2015). J. Ho, A. Jain & P. Abbeel, NeurIPS (2020). Y. Song et al., "Score-Based Generative Modeling through Stochastic Differential Equations", ICLR (2021).

import numpy as np
rng = np.random.default_rng(20)
w, m, s = np.array([0.35, 0.65]), np.array([-2.0, 1.5]), np.array([0.3, 0.5])   # the "data": two bumps
k = rng.choice(2, size=100_000, p=w)
x0 = rng.normal(m[k], s[k])
for t in (0.0, 0.5, 1.0, 3.0):
    xt = np.exp(-t) * x0 + np.sqrt(1 - np.exp(-2 * t)) * rng.normal(size=x0.size)   # exact noising to time t
    print(f"t = {t}: mean {xt.mean():+.3f}, spread {xt.std():.3f}, share left of the valley {np.mean(xt < -0.25):.3f}")
# pure noise N(0, 1): mean 0, spread 1, share left of -0.25 = 0.401

Output when we ran it:

t = 0.0: mean +0.281, spread 1.724, share left of the valley 0.349
t = 0.5: mean +0.166, spread 1.317, share left of the valley 0.363
t = 1.0: mean +0.099, spread 1.124, share left of the valley 0.375
t = 3.0: mean +0.014, spread 1.003, share left of the valley 0.397

Tweedie's formula: the score is a denoiser. For $X_t = a X_0 + \sqrt{1 - a^2}\,Z$ with $a = e^{-t}$: $\;\mathbb E[X_0 \mid X_t = x] = \big(x + (1 - a^2)\,s(t, x)\big)/a$. The best guess of the clean sample is the noisy one, nudged along the score and scaled back up.

Denoising score matching. A network $s_\theta(t, x)$ is trained to minimise $\mathbb E\,\big\lvert s_\theta(t, X_t) + \tfrac{X_t - aX_0}{1 - a^2} \big\rvert^2$, averaged over data $X_0$, times $t$ and noise. The target uses only the known noise that was added, yet the minimiser is the true score $s(t, x)$ (Vincent, 2011), a consequence of Tweedie's formula. This is how every diffusion model is trained.

References: A. Hyvärinen, J. Machine Learning Research 6 (2005), 695–709. P. Vincent, Neural Computation 23 (2011), 1661–1674. B. Efron, "Tweedie's formula and selection bias", J. Amer. Statist. Assoc. 106 (2011), 1602–1614.

import numpy as np
w, m, s = np.array([0.35, 0.65]), np.array([-2.0, 1.5]), np.array([0.3, 0.5])

def score(t, x):                                     # d/dx log p(t, x) for the noised two-bump data
    mean, var = m * np.exp(-t), s ** 2 * np.exp(-2 * t) + 1 - np.exp(-2 * t)
    dens = w * np.exp(-(x - mean) ** 2 / (2 * var)) / np.sqrt(2 * np.pi * var)
    return np.sum(dens * -(x - mean) / var) / np.sum(dens)

for t in (0.05, 1.0, 5.0):
    print(f"t = {t}: " + "  ".join(f"s({x:+.0f}) = {score(t, x):+.2f}" for x in (-3, 0, 3)))
print("pure noise N(0, 1) has score -x: +3, 0, -3")

Output when we ran it:

t = 0.05: s(-3) = +6.21  s(+0) = +4.43  s(+3) = -4.90
t = 1.0: s(-3) = +2.62  s(+0) = +0.15  s(+3) = -2.73
t = 5.0: s(-3) = +3.00  s(+0) = +0.00  s(+3) = -3.00
pure noise N(0, 1) has score -x: +3, 0, -3

Anderson's theorem (1982). If $dX = f(t, X)\,dt + g(t)\,dW$ has densities $p_t$, then the same process run backward in time solves $dX = \big[f(t, X) - g(t)^2\,\partial_x \log p_t(X)\big]dt + g(t)\,d\overline W$, where $\overline W$ is a Brownian motion running backward. Here $f = -x$ and $g = \sqrt2$, which gives the equation above.

Probability-flow ODE. The ODE $\dot x = f - \tfrac12 g^2\,\partial_x \log p_t$ has the same densities $p_t$ at every time. Substituting it into the Fokker–Planck equation (section 5.3) shows its drift term reproduces the diffusion term exactly. It enables fast deterministic sampling and exact likelihoods.

In practice. The score is a trained network. Sampling uses a few dozen to a thousand steps of an SDE or ODE solver, often in a compressed "latent" space (Stable Diffusion) and steered by a text prompt ("guidance").

References: B. Anderson, "Reverse-time diffusion equation models", Stoch. Proc. Appl. 12 (1982), 313–326. Song et al., ICLR (2021). R. Rombach et al., "High-Resolution Image Synthesis with Latent Diffusion Models", CVPR (2022).

import numpy as np
rng = np.random.default_rng(21)
w, m, s = np.array([0.35, 0.65]), np.array([-2.0, 1.5]), np.array([0.3, 0.5])

def score(t, x):                                     # exact score of the noised data, for many x at once
    mean, var = m * np.exp(-t), s ** 2 * np.exp(-2 * t) + 1 - np.exp(-2 * t)
    dens = w * np.exp(-(x[:, None] - mean) ** 2 / (2 * var)) / np.sqrt(2 * np.pi * var)
    return np.sum(dens * -(x[:, None] - mean) / var, axis=1) / np.sum(dens, axis=1)

T, steps = 3.0, 600
h = T / steps
x = rng.normal(size=20_000)                          # start from pure noise
for i in range(steps):                               # run time backward from T to 0
    x = x + h * (x + 2 * score(T - i * h, x)) + np.sqrt(2 * h) * rng.normal(size=x.size)
left = x < -0.25
print(f"left bump {left.mean():.3f} (data 0.35); centres {x[left].mean():.2f}, {x[~left].mean():.2f} (data -2, 1.5)")
print(f"widths {x[left].std():.2f}, {x[~left].std():.2f} (data 0.30, 0.50)")

Output when we ran it:

left bump 0.352 (data 0.35); centres -2.00, 1.50 (data -2, 1.5)
widths 0.31, 0.50 (data 0.30, 0.50)

Why $\pi$ is the resting distribution. For $dX = s(X)\,dt + \sqrt2\,dW$ the Fokker–Planck equation can be written $\partial_t p = \partial_x\big(p\,\partial_x \log(p/\pi)\big)$, which is zero exactly when $p = \pi$. Under mild conditions $p_t \to \pi$ from any start. The speed is governed by constants such as the spectral gap or a log-Sobolev constant.

Metastability. With $U = -\log\pi$, the expected time to cross from one bump to the other grows like $e^{\Delta U}$, where $\Delta U$ is the height of the barrier between them (the Eyring–Kramers law). A deep valley means an exponentially long wait.

In practice. The discrete update $x \leftarrow x + h\,s(x) + \sqrt{2h}\,z$ (the "unadjusted Langevin algorithm") is slightly biased for $h > 0$. Adding a Metropolis accept/reject step removes the bias (MALA). Song & Ermon (2019) generated images with Langevin steps at a decreasing sequence of noise levels, the idea behind section 7.3.

References: G. Roberts & R. Tweedie, "Exponential convergence of Langevin distributions and their discrete approximations", Bernoulli 2 (1996), 341–363. A. Bovier et al., "Metastability in reversible diffusion processes I", J. Eur. Math. Soc. 6 (2004), 399–424. Y. Song & S. Ermon, "Generative Modeling by Estimating Gradients of the Data Distribution", NeurIPS (2019).

import numpy as np
rng = np.random.default_rng(22)
w, m, s = np.array([0.35, 0.65]), np.array([-2.0, 1.5]), np.array([0.3, 0.5])   # target: the two bumps

def score(x):                                        # d/dx log pi(x), for many x at once
    dens = w * np.exp(-(x[:, None] - m) ** 2 / (2 * s ** 2)) / s
    return np.sum(dens * -(x[:, None] - m) / s ** 2, axis=1) / np.sum(dens, axis=1)

dt, steps = 0.005, 2000                              # run to t = 10
for start, x in (("all in the right bump", 1.5 + 0.5 * rng.normal(size=5000)),
                 ("spread over [-4, 4]", rng.uniform(-4, 4, 5000))):
    for _ in range(steps):
        x = x + dt * score(x) + np.sqrt(2 * dt) * rng.normal(size=x.size)   # Langevin step
    print(f"start {start}: share in left bump at t = 10 is {np.mean(x < -0.25):.3f} (target 0.35)")

Output when we ran it:

start all in the right bump: share in left bump at t = 10 is 0.004 (target 0.35)
start spread over [-4, 4]: share in left bump at t = 10 is 0.414 (target 0.35)