This section is about the score of a probability distribution, the gradient of its log density with respect to the data. Picture an arrow at every point of space. It points in the direction in which the data become more likely, and its length says how fast the log density rises that way. Three facts make this field of arrows the center of diffusion models. A network can learn it from samples alone, samples can be drawn using nothing else, and the noise predictor that §33.2 trains is the score of the noised data multiplied by a known number.

Here is the problem the score solves. Suppose you want a model that can put probability on data of any shape. A flexible way is to take any function $E_\theta(\mathbf{x})$, for example a neural network with one output, call it the energy, and declare

Low energy means high probability. The number $Z_\theta$, the normalizing constant, makes the density integrate to one. Here $d$ is the number of values in one data point, not a model width. For a $256\times256$ color image $d = 196{,}608$, and $Z_\theta$ is an integral over every possible image, which nobody can compute. Maximum likelihood picks $\theta$ to make the data as probable as possible (§0.4). It needs $\log Z_\theta$ and its gradient, so it is stuck before it starts.

Now differentiate the log density with respect to the data $\mathbf{x}$ instead of the parameters $\theta$,

$Z_\theta$ is one number that does not change with $\mathbf{x}$, so its gradient in $\mathbf{x}$ is zero. The normalizer is gone.

The section follows that observation to the end. We define the score and compute it by hand. We learn it from data, first with a classical trick that needs no normalizer, then by adding noise, which turns learning into ordinary regression. A classical identity, Tweedie's formula, then shows that the best noise predictor, the score, and the best denoiser are one function written three ways. Next we draw samples with the score, using gradient ascent with added noise (Langevin dynamics) and a schedule of falling noise. Finally we let the noise be added continuously. Diffusion then becomes a stochastic differential equation, a differential equation with a random term. That equation has a deterministic twin, the probability flow ODE, with exactly the same distribution at every time. The fast samplers and flow matching of the rest of the chapter are built on that twin.

The score is a field of arrows that ignores the normalizer

Start with the standard normal in one dimension. Its log density is $-x^2/2$ plus a constant, so its derivative is $-x$. At $x = 2$ the score is $-2$, an arrow of length 2 pointing back toward 0. At $x = -1$ it is $+1$, again pointing toward 0. At the mean it is 0. The arrows always point at the mean, and they grow with the distance from it.

In plain words, the score is an arrow at every point of space. It points where the density increases fastest, and its length says how fast the log density changes. Following the arrows walks you uphill toward the data.

The name collides with one you may have met. The score-function estimator of §0.4 estimates the gradient of an expected value using $\nabla_\theta\log p_\theta(\mathbf{x})$, and the policy-gradient method of §15.3, called REINFORCE, is built on it. That is a gradient with respect to the parameters, with one entry per parameter. The score of this chapter is a gradient with respect to the data, with one entry per pixel. They share a name and nothing else.

Three cases build the intuition.

A Gaussian. For $\mathcal{N}(\mu, s^2)$ in one dimension, $\log p(x) = -(x-\mu)^2/(2s^2)$ plus a constant, so the score is $-(x - \mu)/s^2$. It points back toward the mean, and a narrow Gaussian pulls harder. In $d$ dimensions, $\mathcal{N}(\boldsymbol\mu, \boldsymbol\Sigma)$ has score $-\boldsymbol\Sigma^{-1}(\mathbf{x} - \boldsymbol\mu)$, a field that is linear in $\mathbf{x}$.

A mixture. Let $p(\mathbf{x}) = \sum_k \pi_k p_k(\mathbf{x})$, a weighted sum of component densities $p_k$ with weights $\pi_k \ge 0$ that add up to 1. Using $\nabla p_k = p_k\nabla\log p_k$,

The weights $r_k(\mathbf{x})$ are the responsibilities. By Bayes rule, $r_k(\mathbf{x})$ is the probability that the point $\mathbf{x}$ came from component $k$. In plain words, the score of a mixture is the average of the components' scores, weighted by how likely each component is to have produced the point.

An energy-based model. A model defined by an energy, as in the opening, is called an energy-based model. Its score is $-\nabla_{\mathbf{x}} E_\theta$, as we saw. The energy-based view of learning studies these models in their own right later in the book. Here they are only the reason the normalizer matters.

Mixtures of Gaussians are the workhorse of this chapter, because adding Gaussian noise keeps them mixtures of Gaussians. Recall the forward process of §33.1. It shrinks a clean point $\mathbf{x}_0$ by a signal scale $\alpha_t$ and adds Gaussian noise with a noise scale $\sigma_t$, so $\mathbf{x}_t = \alpha_t\mathbf{x}_0 + \sigma_t\boldsymbol\epsilon$ with $\boldsymbol\epsilon\sim\mathcal{N}(\mathbf{0}, \mathbf{I})$ independent of $\mathbf{x}_0$. We write $p_t$ for the density of $\mathbf{x}_t$. If $\mathbf{x}_0 \sim \sum_k \pi_k\mathcal{N}(\boldsymbol\mu_k, \boldsymbol\Sigma_k)$, then component $k$ becomes a Gaussian with mean $\alpha_t\boldsymbol\mu_k$ and covariance $\alpha_t^2\boldsymbol\Sigma_k + \sigma_t^2\mathbf{I}$, and the weights stay the same (the theorem on noised mixtures in §33.1). So $p_t$ and its score have closed forms at every $t$.

That gives us an oracle, a function that returns the true score of every noised distribution. Every algorithm in this section can be tested against it before any network is trained. The cell below implements the oracle for components with covariance $s_k^2\mathbf{I}$ and checks it against finite differences of $\log p_t$. Then it draws the arrows for the chapter's 8-Gaussian ring (8 modes on a circle of radius 2, each with standard deviation 0.1) at three noise levels, with $\alpha_t = 1$ so that only the noise changes.

The responsibilities need care. Far from every mode, each $\log \pi_k p_k$ is a large negative number, and exponentiating them directly gives zero divided by zero. The cell subtracts the largest of them before exponentiating, the stable log-sum-exp of §0.4.

The finite-difference check agrees to better than $10^{-9}$. The far point $(50, 0)$ gets a finite score of 184.6, which is $48/0.26$. That is its distance to the nearest noised mode, at $(2, 0)$, divided by that mode's variance $0.1^2 + 0.5^2$.

The arrows are drawn with squashed lengths and colored by size, and they tell the story of the whole section. At $\sigma = 0.05$ the field is enormous. The median arrow on the grid has length 76 and the largest 179, and every arrow points straight at the nearest mode, which says nothing about the ring as a whole. At $\sigma = 2$ the field is gentle (no arrow is longer than 0.7) and points at the center of the ring, as if the data were one blob. The middle level has both, a pull toward the ring and a pull toward each mode.

So small noise shows the local detail, and large noise shows the global shape. A sampler needs both, which is why diffusion uses many noise levels. An interviewer who asks why diffusion models need a whole range of noise levels wants this picture.

import numpy as np

def gmm_score(x, weights, means, covs, alpha, sigma):
    raise NotImplementedError

Learning a score you never see

Suppose we have samples $\mathbf{x}_1, \dots, \mathbf{x}_N$ from an unknown density $p$, and we want a network whose arrows match the true arrows, $\mathbf{s}_\theta \approx \nabla\log p$. The natural loss is the squared distance between the two sets of arrows, averaged over where the data are.

We cannot compute this loss, because its target $\nabla\log p$ is exactly what we do not know. There is a way around it, found by Hyvärinen (2005). Expand the square into three terms, $\tfrac12\lVert\mathbf{s}_\theta\rVert^2 - \mathbf{s}_\theta^\top\nabla\log p + \tfrac12\lVert\nabla\log p\rVert^2$. The last term has no $\theta$ in it, so it does not affect training. The first needs only the model. Only the middle term contains the unknown score. Integration by parts moves the derivative off $\log p$ and onto the model, and the unknown disappears.

The identity is easiest to believe on a model with one parameter, where both losses can be written down by hand.

The cell checks the example with 200,000 samples.

The difference is 0.50 at every $\theta$, up to sampling noise, and both columns are smallest at $\theta = 1$, as the theorem says.

Why the implicit form does not scale

The trace is the problem. Reverse-mode automatic differentiation gives one row of a Jacobian per backward pass (§0.3). The exact $\operatorname{tr}\nabla_{\mathbf{x}}\mathbf{s}_\theta$ needs every diagonal entry, so it costs $d$ backward passes. For a $256\times256\times3$ image that is 196,608 backward passes for every training example.

A random estimate of the trace is cheaper but noisy (the box below). And there is a deeper problem that no trick for the trace fixes.

Real data lie near a low-dimensional set inside $\R^d$. Natural images, for example, fill only a thin sliver of all possible pixel arrays. This belief is called the manifold hypothesis. Off that set the density is zero or nearly so. The log density is then undefined or extremely steep, and the score is not a smooth field.

Song and Ermon (2019) saw this directly. Sliced score matching on raw CIFAR-10 images, a standard set of $32\times32$ color images, gave a loss that fell and then fluctuated without converging. Adding Gaussian noise of variance $10^{-4}$, invisible on pixels scaled to $[0, 1]$, made it converge. Noise spreads the data over the whole space, so the density is positive everywhere and the score is well defined. So we stop fighting the noise and use it.

Denoising score matching

Add noise on purpose. Draw a clean point $\mathbf{x}_0 \sim p$ and noise $\boldsymbol\epsilon\sim\mathcal{N}(\mathbf{0}, \mathbf{I})$, and set $\mathbf{x}_t = \alpha\mathbf{x}_0 + \sigma\boldsymbol\epsilon$. This is the forward process at one fixed time. We drop the subscript $t$ on $\alpha$ and $\sigma$ until we need several noise levels. The noised data have density $p_t$, and we cannot compute its score. But if we know which clean point a noisy point came from, the score of that one Gaussian is easy.

A tiny example first. Take one dimension, $\alpha = 1$, $\sigma = 0.5$, a clean point $x_0 = 1$, and a noise draw $\epsilon = 0.4$. Then $x_t = 1.2$. Given $x_0$, the noisy point has density $\mathcal{N}(1, 0.25)$, whose score at 1.2 is $-(1.2 - 1)/0.25 = -0.8$. That is $-\epsilon/\sigma = -0.4/0.5$. The arrow points back at the clean point, and it is the added noise, reversed and divided by $\sigma$.

In general, write $q(\mathbf{x}_t\mid\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \alpha\mathbf{x}_0, \sigma^2\mathbf{I})$ for the density of a noisy point given its clean point. The noised density is the average of these over the data, $p_t(\mathbf{x}) = \int q(\mathbf{x}\mid\mathbf{x}_0)\,p(\mathbf{x}_0)\,\mathrm{d}\mathbf{x}_0$. The score of one of them is

This conditional score is a target we can compute for every training pair. Vincent (2011) proved that regressing on it is as good as regressing on the unknown score of $p_t$.

There is a second way to see the minimizer, using a fact you already know. $\mathcal{L}_{\text{DSM}}$ is a squared-error regression of the target $-\boldsymbol\epsilon/\sigma$ on the input $\mathbf{x}_t$. §32.1 proved that the best squared-error predictor is the conditional mean of the target given the input. So $\mathbf{s}^\star(\mathbf{x}_t) = -\mathbb{E}[\boldsymbol\epsilon\mid\mathbf{x}_t]/\sigma$. Tweedie's formula, proved below without any score matching, shows that this conditional mean is exactly $\nabla\log p_t$. Together the two facts give a second proof of the theorem.

Notice what was never needed. There is no trace, no Jacobian, and no density. Each training example costs one forward pass and has one regression target. The price is that we learn the score of $p_t$, not of $p$. That is a feature, because $p_t$ is smooth and positive everywhere.

The noisy targets are why DSM needs many samples. At $\sigma = 0.5$ each target $-\epsilon/\sigma$ has standard deviation 2 in every coordinate, while the field we want is smooth. Regression averages the noise away.

The next cell shows this on a 1D two-mode mixture that the rest of the section reuses, $p = 0.2\,\mathcal{N}(-2, 0.25) + 0.8\,\mathcal{N}(2, 0.25)$. The model is a weighted sum of 33 fixed bumps, $s_\theta(x) = \sum_j \theta_j\phi_j(x)$, where each $\phi_j$ is a Gaussian bump centered on a grid point. The model is linear in $\theta$, so DSM is ordinary least squares, solved in one line.

Single targets scatter with standard deviation 2.0. Yet the fitted field is within about 0.2 of the exact score of $p_t$ wherever the noised data have appreciable density, and within 3 percent on average. Where $p_t$ is tiny the fit is poor, because there are no training points there. This low-density problem will matter when we sample, because a sampler started from noise begins exactly where the data are rare.

In an interview, the one-line summary is that DSM turns an impossible target into a computable one, the added noise, and that regression on that noisy target still recovers the exact score of the noised data.

Many noise levels in one network

One noise level is a compromise. A small $\sigma$ gives a score close to the data's but badly behaved, and a large $\sigma$ gives a smooth but blurry one. Song and Ermon (2019) used many levels at once. They trained one network $\mathbf{s}_\theta(\mathbf{x}, \sigma)$ that takes the noise level as an input, and called it a noise conditional score network (NCSN). It is trained on a ladder of levels $\sigma_1 > \sigma_2 > \dots > \sigma_L$, with a weight $w(\sigma)$ for each level,

Their paper writes the weight as $\lambda(\sigma)$. We write $w$, because in this chapter $\lambda$ is the log signal-to-noise ratio of §33.1.

Which weight? The target $\boldsymbol\epsilon/\sigma$ has squared norm about $d/\sigma^2$. Without a weight, the smallest $\sigma$ dominates the sum. With levels from 1 down to 0.01 it outweighs the largest by a factor of $10^4$. Their choice is $w(\sigma) = \sigma^2$, which puts every level on the same scale. To see why, use the identity

and note that $\sigma\mathbf{s}^\star = -\mathbb{E}[\boldsymbol\epsilon\mid\mathbf{x}_t]$ is of order one in each coordinate at every level. Look at the right side once more. It regresses a network output on the noise $\boldsymbol\epsilon$. That is the loss of the denoising diffusion model of §33.2 (DDPM) in another form. Tweedie's formula, next, makes that exact.

import numpy as np

def dsm_loss(score_fn, x0, sigmas, rng):
    sig = np.full(len(x0), float(np.mean(sigmas)))          # placeholder: one level for everyone
    eps = rng.standard_normal(x0.shape)
    s = score_fn(x0 + sig[:, None] * eps, sig)
    return float(0.5 * np.mean(np.sum((s + eps / sig[:, None]) ** 2, axis=1)))

Tweedie's formula, and why noise, score, and denoiser are one function

A denoiser takes a noisy point and guesses the clean point it came from. Under squared error the best guess is the posterior mean $\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t]$, the average of all clean points that could have produced this noisy point, each weighted by how likely it is. That average looks like it needs the unknown data distribution. It needs only the score of the noisy data. Move the noisy point uphill by $\sigma^2$ times the score, then divide by $\alpha$ to undo the shrinking.

§32.1 checked this for Gaussian data with noise added and no shrinking, the case $\alpha = 1$. Here is the general case, for any data distribution and any signal scale $\alpha$.

Proof. Differentiate $p_t(\mathbf{x}) = \int \mathcal{N}(\mathbf{x}; \alpha\mathbf{x}_0, \sigma^2\mathbf{I})\,p(\mathbf{x}_0)\,\mathrm{d}\mathbf{x}_0$ under the integral. The gradient of the Gaussian in $\mathbf{x}$ is the Gaussian times $-(\mathbf{x} - \alpha\mathbf{x}_0)/\sigma^2$, so

By Bayes rule, the fraction inside the last integral is the posterior density of $\mathbf{x}_0$ given $\mathbf{x}_t = \mathbf{x}$, so the integral is $\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t = \mathbf{x}]$. Dividing by $p_t$ gives $\nabla\log p_t(\mathbf{x}) = -(\mathbf{x} - \alpha\,\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t = \mathbf{x}])/\sigma^2$. Solving for the expectation gives the formula. ∎

In plain words, to denoise, step uphill on the log density of noisy data by $\sigma^2$, then undo the shrinking by dividing by $\alpha$. The formula is named after the statistician Maurice Tweedie. Efron (2011) made it widely known in empirical Bayes, the branch of statistics that estimates a prior from the data themselves. It is exact for any data distribution, and it is the reason diffusion works. A network that knows the score of noisy data knows the best denoiser.

For the two-mode mixture the posterior mean is not linear. The next cell checks Tweedie there by brute force. It noises two million samples. Then, for a few chosen values of $x_t$, it averages the clean $x_0$ of the pairs whose $x_t$ lands within 0.02 of that value.

The two columns agree to about 0.01, the Monte Carlo error of averaging over a window of width 0.04. Look at the shape. Between $x_t = -1$ and $x_t = 0$ the posterior mean swings from $-1.5$ to $+1.05$. A point there could have come from either mode, and the denoiser hedges between them. That hedging is the blur of a one-step denoiser, the problem of §32.4. Diffusion avoids it by taking many small denoising steps instead of one big one.

Tweedie's formula also gives a dictionary between three things a network might output. Write the noise as $\boldsymbol\epsilon = (\mathbf{x}_t - \alpha\mathbf{x}_0)/\sigma$ and take conditional expectations given $\mathbf{x}_t$.

Proof. Taking $\mathbb{E}[\,\cdot\mid\mathbf{x}_t]$ of $\boldsymbol\epsilon = (\mathbf{x}_t - \alpha\mathbf{x}_0)/\sigma$ gives $\mathbb{E}[\boldsymbol\epsilon\mid\mathbf{x}_t] = (\mathbf{x}_t - \alpha\,\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t])/\sigma$, which is the second identity. Substituting Tweedie's $\alpha\,\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t] = \mathbf{x}_t + \sigma^2\nabla\log p_t$ gives $\mathbb{E}[\boldsymbol\epsilon\mid\mathbf{x}_t] = -\sigma\nabla\log p_t$, the first. ∎

Now look again at the denoising diffusion model of §33.2, DDPM. It trains a network $\boldsymbol\epsilon_\theta(\mathbf{x}_t, t)$ to predict the added noise with the simple loss $\mathbb{E}_{t,\mathbf{x}_0,\boldsymbol\epsilon}\lVert\boldsymbol\epsilon - \boldsymbol\epsilon_\theta(\mathbf{x}_t, t)\rVert^2$, with the step $t$ drawn uniformly. Write the network as a score, $\boldsymbol\epsilon_\theta = -\sigma_t\mathbf{s}_\theta$. Then

That is twice the DSM loss at level $t$ with weight $w = \sigma_t^2$, the same weight the noise conditional network chose. In plain words, DDPM is denoising score matching over all noise levels at once, with time drawn uniformly and the weight that makes all levels equally loud. Ho et al. (2020) pointed out this connection in the DDPM paper. Two lines of research, one from variational bounds and one from score matching, had built the same model. Interviewers ask for exactly this derivation.

import numpy as np

def tweedie_denoise(x_t, score, alpha, sigma):
    return x_t + sigma * score          # placeholder

def score_to_eps(score, sigma):
    return sigma * score                # placeholder

Tweedie's formula has a second-order version. It gives the posterior covariance of the clean data from the Hessian of $\log p_t$, so it says how uncertain the best denoiser is. §33.6 and Exercise 33.3f use it.

Langevin dynamics, sampling with nothing but the score

We can learn the score. Can we sample with it? Gradient ascent on $\log p$ (§0.6) climbs to a peak of the density and stops there. That gives the most likely point, not a sample. Add exactly the right amount of random noise at each step, and the climb becomes a random walk whose long-run distribution is $p$ itself.

One step by hand first. Take $p = \mathcal{N}(0, 1)$, whose score at $x$ is $-x$, a step size $\eta = 0.1$, and a current point $x = 2$. The uphill move is $\eta \cdot (-2) = -0.2$. The noise has standard deviation $\sqrt{2\eta} = 0.45$. So the next point is $1.8 + 0.45z$ for a fresh standard normal $z$. Far from the mean the pull dominates and the point drifts back. Near the mean the pull is weak, and the noise keeps the point moving.

Why $\sqrt{2\eta}$ and not some other amount? The uphill step contracts the cloud of points by a fraction of order $\eta$ per step. The noise spreads it by a variance of order $\eta$ per step. The factor 2 balances the two exactly when the cloud is distributed as $p$. We prove this with the Fokker-Planck equation later in the section. The Gaussian case can be done by hand now, and it also shows what a finite step costs.

The cell starts each chain exactly at the target and runs it at several step sizes.

The chain starts at the right answer and walks away from it, to exactly the biased variance the formula predicts. Score-based samplers live with this bias and keep $\eta$ small, for the reason in the box below. An interviewer who asks why diffusion samplers do not use a Metropolis correction wants that reason.

Why plain Langevin fails on real data

Run Langevin with the exact score of a mixture whose modes are far apart, and it gets the proportions wrong. Song and Ermon (2019) give the cleanest argument. Let $p = w\,p_1 + (1-w)\,p_2$, and suppose the two components never overlap. On the region where $p_1$ lives, $\log p = \log w + \log p_1$, so the score is $\nabla\log p_1$. The weight $w$ has disappeared.

A sampler that sees only the score therefore cannot tell how much mass each mode should get. When the modes overlap slightly, the weight shows up only in the thin valley between them, where the density is tiny. A chain has to cross that valley many times to learn the weight. The expected time to cross grows exponentially with the depth of the valley, measured in log density.

Three problems compound on real data.

1. The score is undefined off the data manifold.
2. Scores are poorly estimated where data are rare, as the least-squares fit above showed, and a chain started from noise begins exactly there.
3. Separated modes mix slowly, so the relative weights come out wrong.

Noise fixes all three at once. A large noise level fills the space. It makes the score well defined and well estimated everywhere, and it merges the modes so that the chain can move between them. A small noise level stays close to the data. So sample at a large level first, then lower it step by step.

Why $\eta_i \propto \sigma_i^2$? The score at level $\sigma$ has size of order $1/\sigma$ (the table "Score, noise, and denoiser, at a glance" above). So the uphill move $\eta\,\mathbf{s}$ has size $\eta/\sigma$, and the noise $\sqrt{2\eta}\,\mathbf{z}$ has size $\sqrt\eta$. Their ratio is $\sqrt\eta/\sigma$, which is the same at every level exactly when $\eta \propto \sigma^2$. With this rule, the balance between moving and diffusing does not depend on the noise level.

The next cell repeats Song and Ermon's toy experiment with the exact score. The data are $0.2\,\mathcal{N}((-5, -5), \mathbf{I}) + 0.8\,\mathcal{N}((5, 5), \mathbf{I})$, and 2,000 particles start uniformly in the square $[-8, 8]^2$. Plain Langevin runs 1,000 steps at $\eta = 0.05$. Annealed Langevin uses 10 geometric levels from $\sigma = 20$ to $\sigma = 1$, with $M = 100$ and $\eta_L = 0.05$. These are the paper's settings, converted to our $\eta$.

Plain Langevin ends with almost exactly half the particles in each mode, the split of the uniform start. Each particle slides into the basin it started in and never leaves. Annealed Langevin ends at about 0.79, close to the true 0.8 (the standard error with 2,000 particles is about 0.009). At $\sigma = 20$ the two modes are one blob, and the particles spread across it in the right proportions. As $\sigma$ falls the blob splits, and each part keeps its share.

The last level here is $\sigma = 1$. So these samples come from the data noised once more with unit variance, and they are visibly wider than the data. Real samplers end at a tiny $\sigma$. Many finish with one Tweedie denoising step, $\mathbf{x} \leftarrow \mathbf{x} + \sigma_L^2\,\mathbf{s}_\theta(\mathbf{x}, \sigma_L)$, which jumps to the posterior mean of the clean data (here $\alpha = 1$).

The widget runs both samplers on a smaller copy of this mixture, with the same weights 0.2 and 0.8. The arrows are the exact score at the current noise level, the dots are the particles, and the dashed circles mark the two data modes. Press Sample to run 1,500 Langevin steps from a uniform start, which takes five seconds. The annealed run spends its first 1,000 steps on 10 levels of 100 steps each, from $\sigma = 3$ down to the noise level you set, and the rest at that level. The big number is the share of particles in the left mode.

Once you have run both samplers, raise the step size $\varepsilon$ toward its maximum and watch the cluster std below the big number. It is measured across the line through the two modes, where both noised modes have the same spread. The widget sets $\eta = \varepsilon(s^2 + \sigma^2)$, where $s^2$ is the variance of one data mode. So the step is $\varepsilon$ in units of one noised mode's variance, which follows the rule $\eta\propto\sigma^2$. By the $2/(2-\eta)$ bias, each cluster's variance then grows by the factor $2/(2-\varepsilon)$. At the top of the slider that is about 1.5, so the measured std comes out about 1.24 times the exact one.

import numpy as np

def annealed_langevin(score_fn, sigmas, n_steps, step_scale, x_init, rng):
    x = np.array(x_init, dtype=float)
    s = sigmas[-1]                                   # placeholder: plain Langevin at the last level
    for _ in range(n_steps * len(sigmas)):
        x = x + step_scale * score_fn(x, s) + np.sqrt(2 * step_scale) * rng.standard_normal(x.shape)
    return x

Stochastic differential equations from first principles

DDPM uses $T = 1000$ discrete steps. The first noise conditional network used 10 levels with 100 steps each. Both numbers look arbitrary. Song et al. (2021) took the limit in which noise is added continuously. Then the forward process is a stochastic differential equation (SDE), a differential equation with a random term. Every sampler becomes a numerical method for solving an equation. The number of steps becomes a choice made at sampling time rather than a property of the model.

We need one new object, and we define it through the simulation anyone would write. Take a small time step $\Delta t$. At each step, move by a deterministic amount proportional to $\Delta t$, then add Gaussian noise whose variance is proportional to $\Delta t$. For example, take no deterministic move, a noise strength of 1, and $\Delta t = 0.01$. Each step adds noise of standard deviation $\sqrt{0.01} = 0.1$. After 100 steps, one unit of time, the variances add up to $100 \times 0.01 = 1$.

In plain words, each tiny step moves by the drift and then jiggles by a Gaussian whose variance is proportional to the step length. The $\sqrt{\Delta t}$ is the whole subtlety. Over a fixed time $t$ there are $t/\Delta t$ independent jiggles, so their variances add up to $g^2 t$, a finite amount. Had the jiggle been proportional to $\Delta t$, the total variance would be $g^2t\,\Delta t$, which goes to zero, and the noise would vanish in the limit. Had it been larger than $\sqrt{\Delta t}$, the variance would explode.

Brownian paths are continuous but nowhere differentiable, because their increments are of size $\sqrt{\Delta t}$, far larger than $\Delta t$. That is why we write $\mathrm{d}\mathbf{w}_t$ and never $\mathrm{d}\mathbf{w}_t/\mathrm{d}t$. The rigorous theory, Itô calculus, proves that the limit exists. Everything below can be checked at the level of the recursion.

DDPM and the noise conditional network are two SDEs

Take DDPM's chain, $\mathbf{x}_i = \sqrt{1-\beta_i}\,\mathbf{x}_{i-1} + \sqrt{\beta_i}\,\mathbf{z}_{i-1}$ for $i = 1, \dots, N$, where $\beta_i$ is the variance added at step $i$. Put the steps on a time grid $t_i = i/N$ with step $\Delta t = 1/N$. Write each step's variance as a rate times the step, $\beta_i = \beta(t_i)\,\Delta t$, for a continuous noise rate $\beta(t)$. For small $\Delta t$, $\sqrt{1 - \beta(t)\Delta t} = 1 - \tfrac12\beta(t)\Delta t + O(\Delta t^2)$, so

That is Euler-Maruyama for the SDE $\mathrm{d}\mathbf{x} = -\tfrac12\beta(t)\mathbf{x}\,\mathrm{d}t + \sqrt{\beta(t)}\,\mathrm{d}\mathbf{w}_t$. It is called the variance preserving (VP) SDE, because, as in §33.1, it keeps $\alpha_t^2 + \sigma_t^2 = 1$.

The noise conditional network's chain keeps the data unscaled and adds independent noise to reach each level, $\mathbf{x}_i = \mathbf{x}_{i-1} + \sqrt{\sigma_i^2 - \sigma_{i-1}^2}\,\mathbf{z}_{i-1}$. With $\sigma_i = \sigma(t_i)$, the added variance is $\sigma_i^2 - \sigma_{i-1}^2 \approx \frac{\mathrm{d}\sigma^2}{\mathrm{d}t}\Delta t$. That is Euler-Maruyama for an SDE with no drift, called the variance exploding (VE) SDE because the variance of $\mathbf{x}_t$ keeps growing instead of staying at 1.

Both rows are special cases of one rule that converts any Gaussian schedule $(\alpha_t, \sigma_t)$ into an SDE and back.

Check the rule on VP. With $f = -\tfrac12\beta$ we get $\alpha_t = e^{-\frac12\int_0^t\beta}$. The variance equation $v' = -\beta v + \beta$ has the solution $v = 1 - \alpha_t^2$, so $\alpha_t^2 + \sigma_t^2 = 1$ for all $t$. Going the other way, $\frac{\mathrm{d}\sigma_t^2}{\mathrm{d}t} = -\frac{\mathrm{d}\alpha_t^2}{\mathrm{d}t} = \beta\alpha_t^2$, and $g^2 = \beta\alpha_t^2 + \beta(1 - \alpha_t^2) = \beta$. For VE, $f = 0$ gives $g^2 = \mathrm{d}\sigma_t^2/\mathrm{d}t$.

The rule also builds new processes. Keep VP's drift but ask for the smaller noise variance $\sigma_t^2 = (1 - \alpha_t^2)^2$. The converse formula gives $g^2 = \beta(t)(1 - \alpha_t^4)$, which Song et al. (2021) call the sub-VP SDE. Its variance is below VP's at every time, and it gave their best likelihoods.

In practice this rule means that every schedule of §33.1 has an SDE, so every sampler in the rest of this section works for it. §33.4 and §33.6 use the two formulas for $f$ and $g$ directly.

The Fokker-Planck equation, how a density flows

An SDE moves individual points. We care about the whole cloud, the density $p_t$ of $\mathbf{x}_t$ when the start $\mathbf{x}_0$ is drawn from $p_0$. How does $p_t$ change in time? Two things happen. The drift carries probability along, like a current carrying dye. The noise spreads probability out, from places where it is dense to places where it is sparse.

Two pieces of notation come first. The divergence of a vector field $\mathbf{u}$ is $\nabla\cdot\mathbf{u} = \sum_i\partial u_i/\partial x_i$, the net rate at which the field flows out of a point. The Laplacian of a function is $\nabla\cdot\nabla p = \sum_i\partial^2p/\partial x_i^2$. It is negative at a peak of $p$ and positive where $p$ is lower than its surroundings.

In plain words, the drift moves probability like a current, and the noise spreads it from high density toward low density, at a rate set by the curvature of the density.

The proof watches averages. A test function $\varphi$ is a smooth function that is zero outside a bounded region. If we know the average $\int\varphi\,p_t$ for every test function, we know $p_t$. So we track how such averages change over one Euler-Maruyama step.

Langevin dynamics is a one-line application. Its drift is $\mathbf{f} = \nabla\log p$ and $g = \sqrt2$. Put $p_t = p$ into the right side and use $p\,\nabla\log p = \nabla p$,

So $p$ does not change under the Langevin diffusion, the stationarity we promised. The factor 2 in $\sqrt{2\eta}$ is exactly what makes the two Laplacians cancel.

That the chain also reaches $p$ from any start needs more than stationarity. Under mild conditions on a smooth positive $p$ it does, but the speed varies enormously. Convergence is fast when $p$ is log-concave, meaning $\log p$ is a concave function, a single bump with no valleys. How many steps it takes is the chain's mixing time, and the two-mode experiment showed that it can be enormous.

Running time backward

Two results turn diffusion into a generative model, and both come from the Fokker-Planck equation. The trick is to rewrite its noise term with the score. Since $p\,\nabla\log p = \nabla p$,

Substituting this into the Fokker-Planck equation leaves a single divergence,

That is a continuity equation, the density equation of an ODE with no noise at all.

Proof. The density of the ODE's state obeys the continuity equation with velocity $\mathbf{u} = \mathbf{f} - \tfrac12g^2\nabla\log p_t$. The rewriting above shows that the SDE's marginals obey that same equation. Both start from $p_0$, and a continuity equation with a smooth velocity field has a unique solution for a given start, so the two densities are equal for all $t$. The ODE can also be run backward, so starting from $p_1$ and integrating down to $t = 0$ produces $p_0$. ∎

In plain words, the SDE moves each point randomly, and the ODE moves each point deterministically along the current of probability. Someone who looks only at the cloud of points at each time cannot tell the two apart. The individual paths are completely different.

For the VE SDE the ODE has a clean form. There $\mathbf{f} = \mathbf{0}$ and $g^2 = \mathrm{d}\sigma^2/\mathrm{d}t = 2\sigma\,\mathrm{d}\sigma/\mathrm{d}t$, so $\mathrm{d}\mathbf{x}/\mathrm{d}t = -\sigma\,(\mathrm{d}\sigma/\mathrm{d}t)\,\nabla\log p_\sigma$. Written with the noise level $\sigma$ as the clock, it reads

whatever the schedule $\sigma(t)$.

The second result keeps the noise and reverses it.

The theorem is about whole paths, and its proof needs stochastic calculus. That the marginals come out right follows from the Fokker-Planck equation in a few lines.

Compare the two theorems. The reverse SDE's drift is the ODE's drift with a second $-\tfrac12g^2\nabla\log p_t$ added. So split one reverse step in two. One part is the probability flow ODE. The other part, written in the reverse time $\tau$, is $\tfrac12g^2\nabla\log p_t\,\mathrm{d}\tau + g\,\mathrm{d}\bar{\mathbf{w}}$. That is Langevin dynamics for $p_t$ with its clock sped up by $\tfrac12g^2$. (Substitute $\mathrm{d}\tau' = \tfrac12g^2\mathrm{d}\tau$ into $\nabla\log p_t\,\mathrm{d}\tau' + \sqrt2\,\mathrm{d}\mathbf{w}$.)

To sample, replace the unknown $\nabla\log p_t$ by a trained $\mathbf{s}_\theta(\mathbf{x}, t)$. Start from the prior $p_1 \approx \mathcal{N}(\mathbf{0}, \mathbf{I})$, which is accurate for VP since $\alpha_1^2 = 4.3\times10^{-5}$. Integrate either equation from $t = 1$ down to a small $t_{\min}$.

The next cell does this with the exact score of the 1D mixture $0.2\,\mathcal{N}(-2, 0.25) + 0.8\,\mathcal{N}(2, 0.25)$ and the VP SDE with $\beta(t) = 0.1 + 19.9\,t$. It runs Euler-Maruyama on the reverse SDE and Euler on the ODE, at three step counts, from the same starting noise. It reports the fraction of samples in the right mode, their mean, and their variance. The data have weight 0.8, mean 1.2, and variance $0.25 + 4 - 1.44 = 2.81$.

With 500 and with 50 steps both samplers reproduce the data within sampling error. With 4,000 samples the weight has a standard error of about 0.006. The variance is mostly the spread between the two modes, $0.25 + 16w(1 - w)$, so a weight off by 0.006 moves it by about 0.06. The SDE's variance of 2.97 at 50 steps follows mostly from its slightly low weight of 0.789, since $0.25 + 16 \times 0.789 \times 0.211 = 2.91$.

With 10 steps both samplers are off, and the way they are off is telling. The reverse SDE's variance jumps to about 3.3, too wide. Each big Euler-Maruyama step injects noise that the coarse drift does not fully remove. The ODE's error is smaller. Errors at few steps are the subject of §33.4.

The trajectories show what the intuition box described. The ODE paths are smooth curves that pick their mode early and then glide in, and each one is a fixed function of its starting noise. The SDE paths wander, and some change their minds late. The histograms at the end agree.

The widget runs the same experiment with 5 to 200 steps, from the same noise draws for both samplers. The left panel shows a dozen paths of each sampler over the shaded true density, and the right panel shows the histogram of all samples against the dashed true density. Press Sample to watch the paths grow and the histograms follow the truth.

Two more connections tie this back to the earlier sections. First, DDPM's sampler of §33.2 draws each step from the learned reverse Gaussian, which is called ancestral sampling. It is itself a discretization of the reverse VP SDE (Song et al. 2021, Appendix E). The 1,000 steps of DDPM were solving this equation all along.

Second, ODE paths cannot cross, because two solutions of the same smooth ODE that meet at one time coincide at all times. So the probability flow ODE defines a one-to-one map between noise and data, and it gives every image a unique noise code. It gives exact likelihoods through the change of variables for ODEs (§33.6). And an adaptive ODE solver picks its own step sizes. With one, Song et al. (2021) cut the number of network calls by over 90 percent without visible loss of quality.

In an interview, expect to write all three equations, the forward SDE, the reverse SDE with the full $g^2$, and the ODE with half of it, and to say why the last two share marginals.

Predictor-corrector samplers

We have a score at every time, so we can mix the two ideas. A predictor-corrector (PC) sampler alternates two moves. The predictor takes one step of a solver for the reverse SDE or the ODE, from $t$ to $t - \Delta t$. The corrector then runs a few Langevin steps at the new time. These pull the samples back toward $p_{t-\Delta t}$ without changing $t$. Annealed Langevin is the special case with no predictor, and DDPM's sampler is the special case with no corrector.

Song et al. (2021) choose the corrector's step as $\eta = 2(r\lVert\mathbf{z}\rVert/\lVert\mathbf{s}_\theta\rVert)^2$. Here $r$ is a target ratio of the uphill move to the injected noise, which they call a "signal-to-noise ratio" (0.16 for their VE models on CIFAR-10). This sets the balance between moving and diffusing directly, as the rule $\eta\propto\sigma^2$ did.