Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

Samplers and Schedulers: Diffusion Sampling Is Numerical Integration

6 minute read

Published:

TL;DR: Sampling from a diffusion model means integrating a differential equation whose right-hand side is the trained network. DDPM's ancestral loop is a first-order stochastic solver; DDIM is first-order deterministic. Higher-order solvers (Heun) and exponential integrators that solve the equation's linear part exactly (DPM-Solver) cut the step count by roughly an order of magnitude with no retraining. The remaining lever is where you place those steps along the noise schedule, which at low step counts matters as much as the solver itself.

The reverse process is an ODE

Song et al. showed the forward corruption is a stochastic differential equation \(d\mathbf{x} = f(\mathbf{x},t)\,dt + g(t)\,d\mathbf{w}\), and that it can be reversed in two ways. Either as an SDE,

\[ d\mathbf{x} = \bigl[f(\mathbf{x},t) - g(t)^2\,\nabla_{\mathbf{x}}\log p_t(\mathbf{x})\bigr]dt + g(t)\,d\bar{\mathbf{w}}, \]

or as a deterministic probability-flow ODE with the same marginals at every \(t\):

\[ \frac{d\mathbf{x}}{dt} = f(\mathbf{x},t) - \tfrac{1}{2}g(t)^2\,\nabla_{\mathbf{x}}\log p_t(\mathbf{x}). \]

The score \(\nabla_{\mathbf{x}}\log p_t\) is what the network supplies, up to the fixed rescaling \(\boldsymbol{\epsilon}_\theta = -\sqrt{1-\bar{\alpha}_t}\,\nabla_{\mathbf{x}}\log p_t\). So a sampler is a choice of integrator, and the number of function evaluations (NFE) is the number of times the U-Net runs. The “1000 steps” of vanilla DDPM is not a property of the model; it is Euler’s method with a very small step.

In Karras et al.’s cleaner variance-exploding coordinates, where \(\mathbf{x}_\sigma = \mathbf{x}_0 + \sigma\boldsymbol{\epsilon}\) and \(D_\theta(\mathbf{x};\sigma)\) is the denoiser’s estimate of the clean image, the whole ODE collapses to

\[ \frac{d\mathbf{x}}{d\sigma} = \frac{\mathbf{x} - D_\theta(\mathbf{x};\sigma)}{\sigma}. \]

Read it literally: at every noise level, move away from the current point in the direction of the denoised estimate, scaled by how noisy you currently are.

Euler, Heun, and why order helps

Euler takes one derivative evaluation and steps along the tangent: \(\mathbf{x}_{i+1} = \mathbf{x}_i + (\sigma_{i+1}-\sigma_i)\,\mathbf{d}_i\). Because the true trajectory curves, the tangent leaves it, and the error per step is \(O(h^2)\).

Heun’s method corrects this. Take the Euler step to a provisional point, evaluate the derivative there, and re-step using the average of the two slopes. Two network calls per step, local error \(O(h^3)\). That trade is usually worth it: at a fixed NFE budget, halving the number of steps to double their accuracy wins once the steps are large.

Euler versus Heun on one large step of a curving trajectoryA curving grey trajectory runs from a start point at the upper left down to a true endpoint at the lower right. A straight orange line from the start point follows the initial tangent and ends well above the true endpoint, labelled "Euler, one evaluation, overshoots". A teal line follows the average of the initial slope and the slope measured at the Euler point, and ends close to the true endpoint, labelled "Heun, two evaluations". A dashed grey line marks the remaining Euler error. σ = σᵢ (start) σ = σᵢ₊₁ (target) Euler: 1 evaluation, follows the initial tangent Heun: 2 evaluations, averages both slopes true trajectory error One step of the probability-flow ODE
Notice that the extra cost of Heun buys accuracy only because the trajectory is curved. On a nearly straight path — the goal of rectified flow — the two coincide and the second evaluation is wasted.

DPM-Solver: solve the easy part exactly

The diffusion ODE is not a generic black box. Written in terms of \(\alpha_t\) and \(\sigma_t\), it is semi-linear: a linear term in \(\mathbf{x}\) plus a nonlinear term that is entirely the network. Variation of constants handles the linear part in closed form, leaving one integral to approximate. With the half-log-SNR variable \(\lambda_t = \log(\alpha_t/\sigma_t)\),

\[ \mathbf{x}_t = \frac{\alpha_t}{\alpha_s}\,\mathbf{x}_s \;-\; \alpha_t\!\int_{\lambda_s}^{\lambda_t} e^{-\lambda}\,\hat{\boldsymbol{\epsilon}}_\theta(\hat{\mathbf{x}}_\lambda,\lambda)\,d\lambda . \]
Key Insight — where the speed-up comes from: a general-purpose solver like Euler discretises the whole right-hand side, so it accumulates error from the stiff linear term as well as from the network. The exponential-integrator view moves that term out of the discretisation entirely — it is integrated exactly — so the only approximation left is a low-order Taylor expansion of \(\hat{\boldsymbol{\epsilon}}_\theta\) in \(\lambda\), which is a smooth, slowly varying function. Truncating it at first order recovers DDIM exactly; at second and third order you get DPM-Solver-2 and -3.

Ancestral or deterministic

Integrating the SDE means adding fresh noise at every step; integrating the ODE means not. The stochastic variants (DDPM ancestral, “Euler a”, DPM++ SDE) have a genuine benefit Karras et al. identified: injected noise acts as error correction, dragging the sample back towards the correct marginal after solver drift. The cost is that the seed no longer determines the image — change the step count and you get a different picture — and at high step counts the churn can erase fine detail.

Deterministic samplers give a fixed, invertible-ish map from noise to image, which is what makes latent interpolation, inversion and reproducible seeds work at all.

SamplerOrderNFE per stepNoise injectedTypical step count
DDPM (ancestral)11yesseveral hundred
DDIM11no20–50
Euler11no20–50
Euler ancestral11yes20–40
Heun (EDM)22optional15–30 steps
DPM-Solver-222no15–25
DPM-Solver++ (2M, multistep)21no15–25
UniPC2–31no10–20

Step counts here are the ranges people reach for in practice, not measured equal-quality points; the honest comparison is always at matched NFE on your own model.

The schedule is half the sampler

A sampler needs a list of noise levels, and spacing them uniformly in \(t\) is a poor choice: it wastes evaluations at high \(\sigma\), where the trajectory is nearly straight, and starves the low-\(\sigma\) region where detail is decided. Karras et al. propose

\[ \sigma_i = \left(\sigma_{\max}^{1/\rho} + \frac{i}{N-1}\left(\sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho}\right)\right)^{\rho}, \qquad \rho = 7, \]

which concentrates steps at low noise. At 50 steps the choice barely registers; at 10 it can dominate the choice of solver.

Where it breaks: solver order assumes a smooth right-hand side, and high classifier-free guidance makes the effective field stiff — a second-order solver at 10 steps and guidance 12 can be worse than first order. Also, a model trained on 1000 discrete timesteps has only been asked about those inputs; a continuous-\(\sigma\) solver silently interpolates, which is fine near the training grid and not guaranteed away from it.

Solvers get you to roughly 10–20 evaluations. Below that, no integrator helps, because the trajectory genuinely is curved — you have to change the model, which is what distillation and consistency models do.

References

  1. Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021.
  2. Song, J., Meng, C., & Ermon, S. Denoising Diffusion Implicit Models. ICLR 2021.
  3. Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. NeurIPS 2022.
  4. Karras, T., Aittala, M., Aila, T., & Laine, S. Elucidating the Design Space of Diffusion-Based Generative Models. NeurIPS 2022.
  5. Zhao, W., Bai, L., Rao, Y., Zhou, J., & Lu, J. UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models. NeurIPS 2023.

You May Also Enjoy

Noether’s Theorem: Symmetry, Conservation, and Equivariant Networks

7 minute read

Published:

Energy is conserved because the laws of physics do not care what time it is. That single sentence is Noether’s theorem, and its machine learning descendant is the reason an equivariant network needs less data than one that must learn the symmetry from examples.

Entropy and Free Energy: From the Second Law to the ELBO

7 minute read

Published:

Thermodynamic entropy and Shannon entropy differ by a constant with units. Once you accept that, the ELBO stops being an inference trick and becomes a free energy — and diffusion models stop being a clever architecture and become a driven non-equilibrium process.

Statistical Mechanics: The Boltzmann Distribution and the Cost of Z

6 minute read

Published:

Counting microstates gives you the Boltzmann distribution, and the Boltzmann distribution gives you softmax, simulated annealing and energy-based models. The partition function is not a bookkeeping constant — it is the object that contains every thermodynamic quantity, and it is intractable for exactly that reason.