Reverse-Time SDEs and the Probability-Flow ODE

The discrete diffusion chain of the previous articles is one member of a family. Taking the limit of infinitely many infinitesimal steps turns the chain into a stochastic differential equation, and the family includes processes that cannot be written as a finite chain at all. The continuous formulation is worth learning because it separates two things that the discrete chain conflates: the distribution of samples at each noise level, and the particular random path taken to reach it.

That separation produces two samplers from one trained model. The reverse-time SDE is stochastic and its sample paths are different every run. The probability-flow ODE is deterministic and produces the same output for the same starting noise. Both have the same marginal distribution at every time, so both generate valid samples, and the choice between them is a trade between sample quality, step count, and the ability to invert a sample back to its noise.

From a finite chain to a continuous process

In the variance-preserving chain, each step scales the state by √(1 − β) and adds noise of variance β. For small β, √(1 − β) ≈ 1 − β/2, so the drift per step is −(β/2) x, and the diffusion coefficient per step is β. Writing β as a continuous rate and taking the limit gives the forward equation

dx = −½ β(t) x dt + √(β(t)) dW

where W is a standard Wiener process. The drift pulls the state toward zero, and the diffusion injects noise at a rate that grows with β(t). The solution has the same closed form as the discrete chain: x(t) = √(A(t)) x(0) + √(1 − A(t)) ε, with A(t) now the continuous cumulative product, an exponential of the integral of β.

The alternative convention, used in the variance-exploding formulation, has no drift and a diffusion coefficient that grows with time. The state is the data plus noise of increasing variance, and the process is a pure Brownian motion starting from the data distribution. It has no free parameters beyond the noise scale, which makes it easier to analyse, and it is numerically worse at small noise because the signal is not scaled down and the dynamic range of the state is the dynamic range of the data plus the largest noise scale.

Anderson’s reverse-time result

A forward diffusion process that runs from data to noise can be run backward, and the reverse process is also a diffusion, but its drift depends on the score of the marginal distribution at each time. For the variance-preserving forward equation, the reverse equation is

dx = [ −½ β(t) x − β(t) ∇_x log p_t(x) ] dt + √(β(t)) dW̄

where W̄ is a Wiener process running backward in time, and p_t is the marginal distribution of the forward process at time t. The score term is the only unknown. Everything else is fixed by the schedule.

The derivation is Anderson’s 1982 result, and its content is that the reverse process’s drift is the forward drift minus the diffusion coefficient times the score. The score enters with a negative sign because the reverse process must undo the spreading that the forward process performed: where the density is increasing in x, the forward process pushed samples outward, so the reverse process must push them back.

This is the sense in which a diffusion model is a score model. Training estimates ∇_x log p_t(x) at every t, using the denoising objective from score matching, and sampling solves the reverse equation with that estimate in place of the true score. The substitution is exact only if the score estimate is exact; the error in the estimate is the error in the generated distribution.

The probability-flow ODE and why it shares marginals

Given a diffusion process with marginal densities p_t, there is a deterministic ordinary differential equation whose solution curves have exactly those marginals. For the variance-preserving process it is

dx/dt = −½ β(t) [ x + ∇_x log p_t(x) ]

The ODE is the forward drift plus the score term, without the noise. Its trajectories are the paths along which probability mass flows; the Fokker-Planck equation for the SDE and the continuity equation for the ODE are the same equation when the ODE’s velocity field is chosen this way. Every marginal of the ODE at time t equals the marginal of the SDE at time t, even though individual trajectories differ completely.

This is a strong statement and it is easy to misread. It does not say the ODE and SDE produce the same sample from the same starting point; they do not, because the SDE adds noise that the ODE does not. It says that if the starting point is drawn from the same distribution, the endpoint is distributed identically. Sampling therefore starts from a draw of the terminal distribution, which is approximately standard normal, and integrates the ODE backward to t = 0.

The practical value is that the ODE can be solved with a high-order adaptive solver at far fewer steps than the SDE, because it has no stochastic term to average over. A solver that takes 50 to 200 steps can produce samples comparable to an SDE that needs a thousand, and the gap widens as the score estimate improves. The cost is that the ODE’s discretisation error is systematic rather than random: a coarse solver produces a consistent bias, often visible as a slight blur or a repeated artefact, while the SDE’s error averages out across runs.

import numpy as np

def euler_ode(x, t0, t1, n_steps, beta, score_fn):
    """Integrate dx/dt = -0.5*beta(t)*(x + score(x)) from t0 to t1."""
    dt = (t1 - t0) / n_steps
    for i in range(n_steps):
        t = t0 + i * dt
        b = beta(t)
        drift = -0.5 * b * (x + score_fn(x, t))
        x = x + drift * dt        # explicit Euler; first-order accurate
    return x

def euler_maruyama(x, t0, t1, n_steps, beta, score_fn, rng):
    """The same step with the diffusion term retained."""
    dt = (t1 - t0) / n_steps
    for i in range(n_steps):
        t = t0 + i * dt
        b = beta(t)
        drift = -0.5 * b * (x + score_fn(x, t))
        x = x + drift * dt + np.sqrt(b * abs(dt)) * rng.standard_normal(x.shape)
    return x

The stochastic version is Euler-Maruyama, and its strong order is one half rather than one: the error in a single trajectory scales as the square root of the step size, while the error in the ODE trajectory scales linearly. This is the formal reason the SDE needs more steps for the same accuracy, and it is also why the SDE’s error is less visible: it appears as noise in the output rather than as a systematic deviation.

Choosing between them

The SDE has one property the ODE lacks: it can correct its own errors. Because the noise term continuously re-injects randomness, a trajectory that drifts into a low-density region is pushed back toward the marginal by the score term, and the final sample is a genuine draw from the model’s distribution even if the score estimate is imperfect. The ODE has no such correction; once a trajectory leaves the high-density region, nothing brings it back, and the endpoint can be an artefact of the solver.

The ODE has the property that it is invertible. Integrating forward from a sample to its noise gives a unique noise vector for each sample, which enables editing by noise inversion, interpolation between samples in noise space, and exact likelihood computation through the change-of-variables formula. The likelihood computation requires the divergence of the velocity field, which for a network with d outputs costs d backward passes per evaluation, or can be estimated stochastically with Hutchinson’s trace estimator at the cost of variance.

A practical hybrid runs most of the trajectory with the ODE and applies stochasticity in a middle window. The reason is that the score estimate is most reliable at intermediate noise levels, where the denoising objective has the most signal, and least reliable at both extremes. Stochastic steps in the reliable region correct accumulated error; deterministic steps elsewhere avoid injecting noise where the model cannot compensate.

Numerical behaviour at the endpoints

The reverse process starts at t = T with a sample from the terminal distribution and integrates to t = 0. Both endpoints are numerically delicate for different reasons.

At t = T, the true terminal distribution is the data convolved with noise at the largest scale, which is close to but not exactly standard normal. Starting from an exact standard normal introduces a small mismatch that the first few steps must correct, and the correction appears as a slight shift in the global statistics of the output. Some implementations add a small amount of extra noise at the start to widen the terminal distribution toward the true one.

At t = 0, the score of the data distribution is concentrated near the data manifold and the ODE’s velocity field changes rapidly over short distances. A fixed-step solver that is accurate in the middle of the trajectory can overshoot at the end, producing samples that are slightly too sharp or that carry high-frequency artefacts. The standard remedies are a smaller final step, a solver with an adaptive step size, or stopping the integration slightly before t = 0 and treating the last step analytically.

The diffusion coefficient’s behaviour near the endpoints determines which solver is appropriate. A schedule whose β(t) is large at both ends makes the forward process aggressive there and the reverse process correspondingly stiff; a schedule that ramps β smoothly avoids the stiffness at the cost of needing more steps in the middle. This is the same design question that governs the discrete chain, and it is taken up in signal-to-noise ratio and schedule shape.

What the continuous view makes obvious

Two facts become clear only in the continuous formulation. The first is that the model’s only learned component is the score, and every difference between diffusion variants is a difference in how the score is parameterised, scheduled, and integrated. The second is that guidance modifies the score field rather than the sampler, so a guided model’s ODE and SDE are guided in the same way, and neither corresponds to the score of a normalised distribution. That is why the guidance scale changes the distribution being sampled rather than merely the path taken to it.