5. Samplers
DDIM and the design space Draft
DDIM
Training gives us a denoiser; sampling gives us choices. The same network can produce samples through different paths, solvers, and step counts, with large differences in cost. DDIM, introduced by Song, Meng and Ermon, showed how to shorten sampling without retraining the model.
We will derive its update and show that it is Euler's method for the probability-flow ODE in a different coordinate system. That connection leads to the design approach of Karras, Aittala, Aila and Laine: choose the clock, grid, integrator, and amount of noise separately.
The DDIM construction keeps the marginals of DDPM's forward chain and changes its joint law.
Think of the forward process as a collection of snapshots connected into movies. Denoising training observes a clean example and one noisy snapshot at a chosen time. It does not observe how the random displacements at two different times were related. DDIM exploits that freedom: it changes the connections between snapshots while preserving the training pairs at each time. The first figure of Section 5.1 draws the snapshots and three ways of joining them into movies.
Definition 5.1.1 (the non-Markovian forward family). Fix the schedule \(\bar\alpha_1, \dots, \bar\alpha_T\) of Section 3.1 and numbers \(\tilde\sigma_i \ge 0\) with \(\tilde\sigma_i^2 \le 1 - \bar\alpha_{i-1}\). The family \(q_{\tilde\sigma}\) is defined by \(q_{\tilde\sigma}(x_T \mid x_0) = \mathcal N(\sqrt{\bar\alpha_T}\, x_0, (1 - \bar\alpha_T) I)\) and, for \(i \ge 2\),
\[\begin{aligned} q_{\tilde\sigma}(x_{i-1} \mid x_i, x_0) &= \mathcal N\bigl(x_{i-1};\; m_i(x_i, x_0),\; \tilde\sigma_i^2 I\bigr), \\ m_i(x_i, x_0) &= \sqrt{\bar\alpha_{i-1}}\, x_0 + \sqrt{1 - \bar\alpha_{i-1} - \tilde\sigma_i^2}\; \frac{x_i - \sqrt{\bar\alpha_i}\, x_0}{\sqrt{1 - \bar\alpha_i}} . \end{aligned}\]Proposition 5.1.2 (the marginals are DDPM's). For every choice of \(\tilde\sigma\) and every \(i\), \(q_{\tilde\sigma}(x_i \mid x_0) = \mathcal N(\sqrt{\bar\alpha_i}\, x_0, (1 - \bar\alpha_i) I)\).
Proof. By backward induction from \(i = T\). If \(x_i \mid x_0\) has the stated law, write \(x_i = \sqrt{\bar\alpha_i}\, x_0 + \sqrt{1 - \bar\alpha_i}\, \varepsilon\); then the mean in Definition 5.1.1 is \(\sqrt{\bar\alpha_{i-1}}\, x_0 + \sqrt{1 - \bar\alpha_{i-1} - \tilde\sigma_i^2}\, \varepsilon\), and adding independent noise of variance \(\tilde\sigma_i^2\) gives a Gaussian with mean \(\sqrt{\bar\alpha_{i-1}}\, x_0\) and variance \((1 - \bar\alpha_{i-1} - \tilde\sigma_i^2) + \tilde\sigma_i^2 = 1 - \bar\alpha_{i-1}\). ∎
DDPM's denoising objective uses the marginals \(q(x_i \mid x_0)\). Since those marginals are unchanged, a network trained by \(L_{\mathrm{simple}}\) can serve every member of the DDIM family. Song, Meng and Ermon also show that the new variational bounds differ from DDPM's only in the weights of their terms. We can choose \(\tilde\sigma\) after training.J. Song, C. Meng and S. Ermon, "Denoising diffusion implicit models", ICLR 2021 (arXiv 2010.02502), Sections 3 and 4, Theorem 1 and Appendix B; "DDIMs can produce high quality samples 10× to 50× faster in terms of wall-clock time compared to DDPMs". The parameter \(\eta\) is defined in their Section 4.3.
chart the standardized displacement (x_i - sqrt(abar_i) x_0)/sqrt(1 - abar_i) of 48 of 2,048 movies of R2
on DDPM's chain of 1,000 steps against the index i, one movie followed: with sigma~ = 0 flat
lines, with DDPM's sigma~ a walk, with the kernel's white noise; at the probe i = 250 the
histogram of all 2,048 displacements against the N(0, 1) profile; or the movies in x through
the exact snapshots p_i
take the snapshots do not change, the movies do: at i = 250 the 2,048 frames are 0.028 from the exact
snapshot with sigma~ = 0, 0.032 with DDPM's and 0.025 with the kernel's, where 2,048 exact draws
sit 0.034 away; the correlation of e_250 with e_750 is 1.000, 0.064 and 0.003
(the live figure needs JavaScript)
To sample, replace \(x_0\) in Definition 5.1.1 with \(\hat x_0 = D_\theta(x_i)\) and replace \((x_i - \sqrt{\bar\alpha_i}\hat x_0)/\sqrt{1 - \bar\alpha_i}\) with \(\varepsilon_\theta(x_i)\). The paper controls the amount of fresh noise through
\[\tilde\sigma_i(\eta) = \eta\, \sqrt{\frac{1 - \bar\alpha_{i-1}}{1 - \bar\alpha_i}}\, \sqrt{1 - \frac{\bar\alpha_i}{\bar\alpha_{i-1}}}, \qquad \eta \in [0, 1],\]so that \(\eta = 1\) gives \(\tilde\sigma_i^2 = \tilde\beta_i\), the posterior variance of Proposition 3.2.1, and the sampler is DDPM's Algorithm 3.3.2, while \(\eta = 0\) gives a deterministic update.
Proposition 5.1.3 (the DDIM step on any schedule). In the notation \((\alpha_t, \sigma_t)\) of Section 1, for times \(s < t\), the DDIM step is
\[\begin{aligned} x_s &= \alpha_s\, D_\theta(x_t) + \sqrt{\sigma_s^2 - \tilde\sigma^2}\; \varepsilon_\theta(x_t) + \tilde\sigma\, \xi, \qquad \xi \sim \mathcal N(0, I), \\ \tilde\sigma &= \eta\, \sigma_{s \mid t} = \eta\, \sigma_s \sqrt{1 - \Bigl(\frac{\alpha_t \sigma_s}{\alpha_s \sigma_t}\Bigr)^2}, \end{aligned} \tag{5.1}\]where \(\sigma_{s \mid t}\) is the posterior standard deviation of Proposition 3.2.1. With \(\eta = 1\) it is the ancestral step of DDPM; with \(\eta = 0\) it is \(x_s = \alpha_s D_\theta(x_t) + \sigma_s \varepsilon_\theta(x_t)\). For a variance-preserving schedule \(\tilde\sigma\) reduces to \(\tilde\sigma_i(\eta)\) above, and for \(\alpha \equiv 1\) to \(\eta\, \sigma_s \sqrt{1 - \sigma_s^2/\sigma_t^2}\).
Proof. The posterior variance of Proposition 3.2.1 is \(\sigma_{s \mid t}^2 = \sigma_{t \mid s}^2\sigma_s^2/\sigma_t^2 = \sigma_s^2(1 - \alpha_{t \mid s}^2\sigma_s^2/\sigma_t^2)\), which is the stated \(\tilde\sigma^2/\eta^2\). For \(\eta = 1\) the mean of (5.1) is \(\alpha_s \hat x_0 + \sqrt{\sigma_s^2 - \sigma_{s \mid t}^2}\, \hat\varepsilon = \alpha_s \hat x_0 + (\alpha_{t \mid s}\sigma_s^2/\sigma_t)\, \hat\varepsilon\), and substituting \(x_t = \alpha_t \hat x_0 + \sigma_t \hat\varepsilon\) into the posterior mean of Proposition 3.2.1 gives the same expression, so the step is the DDPM step. The reduction to \(\tilde\sigma_i(\eta)\) uses \(\alpha^2 + \sigma^2 = 1\): \(\sigma_s^2(1 - \alpha_t^2\sigma_s^2/(\alpha_s^2\sigma_t^2)) = (\sigma_s^2/\sigma_t^2)(1 - \alpha_t^2/\alpha_s^2)\) after expanding \(\sigma_t^2 - \alpha_t^2\sigma_s^2/\alpha_s^2 = 1 - \alpha_t^2/\alpha_s^2\). ∎
The update has a simple interpretation. First estimate the clean point and noise from \(x_t\). Then rebuild a less noisy point, \(\alpha_s\hat x_0 + \sigma_s\,(\text{noise})\), using the old noise estimate, a fresh draw, or a mixture with the required variance.
The deterministic version reuses the noise estimate. Its update can connect any two times, allowing the sampler to skip training steps. Song, Meng and Ermon obtain useful samples with a subsequence of 20 or 50 of the 1000 training times. Running deterministic updates in the other direction approximately maps a sample back to noise. This is DDIM inversion, used in editing; the approximation reflects the discrete solver steps.
Skipping steps makes a numerical assumption: the current clean and noise estimates remain useful over the whole jump. Along an exact flow they generally change as the point moves. A smaller step refreshes them sooner; a larger step commits for longer to an estimate made at its starting point. The following example separates having an exact denoiser from integrating its flow exactly.
left one DDIM step from sigma = 3 to 0.2 at x_t = 2 on N(0.3, 0.5^2), the chord against the exact flow
line: the chord lands at 0.456 and the flow at 0.601, 0.145 apart; eight steps land 0.021 away,
and on a point mass (width 0) every step lands where the flow lands, at 0.413
right the noise estimate along the step: the value the step holds, 0.551, against eps-hat along the
exact flow, which falls to 0.208 at sigma = 0.2; the area between them is the step's error
(the live figure needs JavaScript)
Example 5.1.4 (DDIM is exact on a point mass, and only there). If the data law is a single point \(m\), then \(D_t \equiv m\) and \(\hat\varepsilon_t(x) = (x - \alpha_t m)/\sigma_t\) exactly, so the deterministic step gives \(x_s = \alpha_s m + (\sigma_s/\sigma_t)(x_t - \alpha_t m)\), which is exactly the flow of the probability-flow ODE of a point mass from \(t\) to \(s\): one step of any length lands where the continuous flow lands. On a Gaussian \(\mathcal N(m, s^2)\) it does not: the flow on the clock \(\sigma\) is \(x_s - m = \sqrt{(s^2 + \sigma_s^2)/(s^2 + \sigma_t^2)}\,(x_t - m)\), while one DDIM step from \(\sigma_t = 3\) to \(\sigma_s = 0.2\) at \(x_t = 2\), with \(m = 0.3\) and \(s = 0.5\), lands at \(0.456\) against the flow's \(0.601\). The program at the end of Section 5.2 prints both. The reason is Theorem 5.2.1: DDIM is Euler's method, exact when the quantity it holds constant along the step, the noise estimate, is constant along the flow, and that is the case only for a point mass. ◁
DDIM is Euler's method in the right coordinates
DDIM looks different from a standard ODE update because both the signal scale and the noise level change during a step. Divide out the signal scale and use the noise-to-signal ratio as the clock. In those coordinates, its deterministic update becomes ordinary Euler integration.
The coordinate change has a practical purpose. Part of the apparent motion comes from the known shrinking or expanding signal scale. Removing that scale leaves only the change associated with noise. An exact trajectory can be described in either coordinate system, but a finite Euler step freezes a different quantity in each. The choice of coordinates therefore affects numerical error. The figure of Section 5.2 draws one motion in both planes.
Theorem 5.2.1 (Song, Meng and Ermon). Let \(\bar x = x/\alpha_t\) and let \(\rho_t = \sigma_t/\alpha_t\) be the clock. Along a solution of the probability-flow ODE (4.3),
\[\frac{d\bar x}{d\rho} \;=\; \hat\varepsilon_t(x), \tag{5.2}\]and the deterministic DDIM step (5.1) with \(\eta = 0\) from time \(t\) to time \(s\) is the Euler step of (5.2) from \(\rho_t\) to \(\rho_s\), that is, \(\bar x_s = \bar x_t + (\rho_s - \rho_t)\, \hat\varepsilon_t(x_t)\).
Proof. From (4.3) and Proposition 4.3.2,
\[\begin{aligned} \frac{d}{dt}\Bigl(\frac{x}{\alpha_t}\Bigr) &= \frac{u_t - f_t x}{\alpha_t} = -\frac{g_t^2\, s_t}{2\alpha_t}, \\ \frac{d\rho_t}{dt} &= \frac{\dot\sigma_t\alpha_t - \sigma_t\dot\alpha_t}{\alpha_t^2} = \frac{\dot\sigma_t - \sigma_t\dot\alpha_t/\alpha_t}{\alpha_t} = \frac{g_t^2}{2\sigma_t\alpha_t}, \end{aligned}\]so \(d\bar x/d\rho = -\sigma_t s_t = \hat\varepsilon_t\). The Euler step gives
\[\begin{aligned} x_s &= \alpha_s\Bigl(\frac{x_t}{\alpha_t} + \Bigl(\frac{\sigma_s}{\alpha_s} - \frac{\sigma_t}{\alpha_t}\Bigr)\hat\varepsilon\Bigr) = \frac{\alpha_s}{\alpha_t}\,(x_t - \sigma_t\hat\varepsilon) + \sigma_s\hat\varepsilon \\ &= \alpha_s D_t(x_t) + \sigma_s \hat\varepsilon_t(x_t), \end{aligned}\]which is (5.1) at \(\eta = 0\). ∎Song, Meng and Ermon (2021), Section 4.3, "Relevance to neural ODEs", where (5.2) appears as \(d\bar x(t) = \varepsilon_\theta(x(t))\, d\sigma(t)\) in their notation, and where the deterministic sampler is identified with the Euler discretization and with the probability-flow ODE of Song et al. (2021).
After removing the magnification and changing the clock, the velocity in (5.2) is simply the noise estimate. Euler's method holds that velocity constant during a step, which is exactly what DDIM does when it reuses the current noise estimate.
left R1 under the variance-preserving schedule in the (t, x) plane, t from 0.4 at the left to 0.05 at
the right: sixteen exact flow lines, the line through the start x = 1.2 at t = 0.3, and from that
start one Euler step on dx/dt = u_t (straight here) and one DDIM step (bent here) to t = 0.05;
the flow lands at 1.106, DDIM at 0.983 and Euler at 1.274
right the same motion in (rho, xbar) = (sigma_t/alpha_t, x/alpha_t), where DDIM's step is the straight
chord and Euler's the bent one; on a point mass at +1 the flow lines are straight here and DDIM's
chord lands on them exactly, at 1.111 from this start, while Euler's misses by +0.282
(the live figure needs JavaScript)
On the EDM clock, \(\alpha \equiv 1\) and \(\rho = \sigma\). The equation becomes \(dx/d\sigma = \hat\varepsilon = (x - D(x; \sigma))/\sigma\), the form integrated by Karras et al.
We can go further by integrating the known linear dynamics analytically and approximating only the part involving the network. This spends the numerical approximation on the quantity we cannot write in closed form. Higher-order solvers estimate how that network prediction changes over the step, using extra evaluations or information retained from previous ones.
Proposition 5.2.2 (the exponential-integrator view). With \(\lambda = \log(\alpha_t^2/\sigma_t^2)\) as the clock, \(\rho = e^{-\lambda/2}\) and (5.2) integrates to
\[x_t \;=\; \frac{\alpha_t}{\alpha_s}\, x_s \;-\; \frac{\alpha_t}{2}\int_{\lambda_s}^{\lambda_t} e^{-\lambda/2}\, \hat\varepsilon_\lambda\bigl(x_\lambda\bigr)\, d\lambda \qquad (s < t \text{ or } s > t).\]Replacing \(\hat\varepsilon_\lambda\) under the integral by its value at the start of the step gives the DDIM step; replacing it by its Taylor polynomial of degree \(k - 1\) in \(\lambda\), with derivatives estimated from earlier evaluations, gives a solver of order \(k\). The linear part of the equation, the factor \(\alpha_t/\alpha_s\), is integrated exactly in every case.
Proof. \(d\rho = -\tfrac12 e^{-\lambda/2}\, d\lambda\), and \(\bar x = x/\alpha\). ∎
This is the design of Lu, Zhou, Bao, Chen, Li and Zhu's DPM-Solver, which treats the probability-flow ODE as a semi-linear equation, solves its linear part in closed form and applies a Taylor expansion to the nonlinear part in the clock \(\lambda\); their second- and third-order solvers reach an FID of \(4.70\) on CIFAR-10 in ten network evaluations and \(2.87\) in twenty, where DDIM needs several times as many.C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li and J. Zhu, "DPM-Solver: a fast ODE solver for diffusion probabilistic model sampling in around 10 steps", NeurIPS 2022 (arXiv 2206.00927), Proposition 3.1, with \(\lambda_t = \log(\alpha_t/\sigma_t)\), half the \(\lambda\) of this series, and Section 3.2 for the orders; "4.70 FID in 10 function evaluations and 2.87 FID in 20 function evaluations on the CIFAR10 dataset". The first-order member of their family is DDIM, which the paper notes. The figure of Section 5.3 runs the second-order single-step member of this family beside Heun.
The design space of EDM
EDM separates choices that earlier methods often bundled together. Karras, Aittala, Aila and Laine use the noise level as time, with \(\alpha \equiv 1\), so the ODE is \(dx/d\sigma = (x - D(x; \sigma))/\sigma\). They then choose a grid of noise levels, a numerical integrator, a network parameterization (Remark 2.4.3), and an optional injection of fresh noise. Studying these choices separately makes it easier to see which improvement comes from which part.
A grid allocates a fixed budget of network calls along the path. Long steps work well where the velocity changes slowly along a trajectory. Where it turns sharply, holding it fixed becomes inaccurate. The grid below puts more calls near small noise levels, where the examples develop narrow features and the denoiser changes rapidly. Its usefulness is a property of the field being integrated, not just of the number of steps.
Definition 5.3.1 (the EDM grid). For \(N\) steps, \(\sigma_{\min} < \sigma_{\max}\) and an exponent \(\rho \ge 1\),
\[\sigma_i = \Bigl(\sigma_{\max}^{1/\rho} + \frac{i}{N - 1}\bigl(\sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho}\bigr)\Bigr)^{\rho}, \quad i = 0, \dots, N - 1, \qquad \sigma_N = 0 .\]Karras et al. use \(\rho = 7\), \(\sigma_{\min} = 0.002\) and \(\sigma_{\max} = 80\); \(\rho = 1\) is a uniform grid in \(\sigma\), and larger \(\rho\) crowds the steps toward \(\sigma_{\min}\).
Euler uses the slope at the beginning of the step. Heun first makes that Euler prediction, asks the network for a slope at the predicted endpoint, and then advances using the average of the two slopes. The provisional endpoint is a place to measure the field; the corrected endpoint is the one we keep. This extra measurement lets the update follow some of the trajectory's bending.
Algorithm 5.3.2 EDM-HEUN
(deterministic; Karras, Aittala, Aila and Laine 2022)
Picture: one step from sigma_i to sigma_{i+1}. Euler
predicts with the slope at the start; the slope at the
predicted point is averaged in, and the step is retaken
with the mean slope (the trapezoidal rule).
x_i o------ slope d_i ------> x~ (Euler)
\ |
\ mean of d_i and d~ | slope d~ at x~
'----------------------> x_{i+1}
Input a denoiser D_theta(x; sigma); the grid of
Definition 5.3.1.
Output a sample.
1. x ~ N(0, sigma_0^2 I)
2. for i = 0, ..., N-1:
3. d <- (x - D_theta(x; sigma_i)) / sigma_i
4. x~ <- x + (sigma_{i+1} - sigma_i) d
5. if sigma_{i+1} > 0:
6. d~ <- (x~ - D_theta(x~; sigma_{i+1})) / sigma_{i+1}
7. x <- x + (sigma_{i+1} - sigma_i) (d + d~) / 2
8. else x <- x~ (the last step is Euler)
9. return x
Heun's method is second order: its error per step is cubic in the step length where Euler's is quadratic, at the price of two network evaluations per step, so \(N\) steps cost \(2N - 1\) evaluations. Karras et al. report an FID of \(1.97\) on unconditional CIFAR-10 and \(1.79\) class-conditional with 35 evaluations, that is, 18 steps of this algorithm, against the hundreds that DDPM's ancestral sampler needs; the opening figure of the introduction runs exactly it on the 64-centre law with 48 steps.Karras et al. (2022), Algorithm 1, Equation 5 for the grid, Table 2 for "1.79 (conditional) and 1.97 (unconditional) at NFE 35" with the deterministic Heun sampler; the abstract states "much faster sampling (35 network evaluations per image) than prior designs".
The stochastic member of their design injects noise before each step: raise the noise level from \(\sigma_i\) to \(\hat\sigma = (1 + \gamma_i)\sigma_i\) by adding fresh Gaussian noise of variance \(\hat\sigma^2 - \sigma_i^2\), then take the Heun step from \(\hat\sigma\) to \(\sigma_{i+1}\). This is a discretization of the family (4.2) with \(\lambda > 0\), since adding noise and then denoising more is a Langevin correction at a fixed noise level, and Remark 4.2.5 is the reason to do it.
squares Euler, Heun, DPM-Solver-2 and the chosen stochastic sampler (Euler-Maruyama, the ancestral
step or DDIM with eta = 0.5) on R3 with N steps of the EDM grid, 64 paths each over the exact
flow lines (RK4, 400 steps), the same noise in every square and one path followed; or the
pairs, each endpoint joined to the exact flow's endpoint from the same noise
charts the sliced W_1 distance of 2,048 endpoints to the data law against the number of score
evaluations, for all six samplers, and the error of Euler and Heun along the path
pairs the RMS distance between a sampler's endpoints and the exact flow's from the same noise, beside
the distance of its cloud to the data law: at N = 8 Heun 0.049 and 0.021, the ancestral step
1.032 and 0.116, Euler-Maruyama 1.353 and 0.031, DDIM with eta = 0.5 0.577 and 0.084, where
two independent draws of R3 are 1.42 apart
(the live figure needs JavaScript)
There are two ways to judge a sampler here. Comparing endpoints from the same starting noise measures how accurately a deterministic solver follows a particular flow map. Comparing the clouds of endpoints measures the output distributions, which also makes sense for stochastic samplers whose individual paths differ. The finite-sample reference in the plot is a scale for comparison; it is not a universal lower bound on the measured distance.
Example 5.3.3 (the samplers on R3). In the figure's default view, eight steps of Heun, fifteen evaluations, put the endpoints within a root-mean-square distance of \(0.049\) of the reference flow's endpoints from the same noise; eight Euler steps overshoot the bends of the flow lines and land \(0.164\) away. Against the data law, the reference is the exact flow's own endpoints, whose sliced Wasserstein distance to 8,192 fresh draws is the dotted floor of the chart; Heun comes within \(20\%\) of that floor at \(N = 16\) and Euler only at \(N = 128\), the chart's last tick.
The stochastic samplers are measured against the same floor, which they can cross, since their endpoints are different random variables: Euler–Maruyama with \(\lambda = 1\) reaches it at \(N = 64\), and the ancestral step at \(N = 64\). Setting \(\rho = 1\) instead of \(7\) spreads the steps uniformly in \(\sigma\) and visibly worsens every sampler at small \(N\): the paths bend most at small \(\sigma\), where the denoiser changes fastest, and a grid that does not crowd its steps there wastes them on the nearly straight part of the path. The second-order exponential integrator of Proposition 5.2.2 lands \(0.035\) from the reference at \(N = 8\) with the same evaluations as Heun, and the fitted orders over \(N \ge 8\) are \(0.79\), \(2.05\) and \(2.07\) for Euler, Heun and the exponential integrator, the order column of the table of Section 5.4 measured. The second view shows where Euler's error is made: half of its final error at \(N = 8\) is accumulated below \(\sigma = 0.95\), where the lines bend toward the moons. ◁
Remark 5.3.4 (trick: the last step). Every grid ends at \(\sigma = 0\), where the velocity \((x - D)/\sigma\) is a quotient of two vanishing quantities for a network and the exact velocity of a Gaussian mixture of positive width is finite but that of a point-mass law is not. Algorithm 5.3.2 takes the last step by Euler for this reason, and the figures of this series do the same; it is also why the flow-map figure of Section 8.1 reports the same error for Heun and Euler at \(N = 1\), where the only step is the last one. A last step that evaluates the denoiser at the smallest positive \(\sigma\) and returns \(D\) itself is the same thing written differently.
The program below checks the two statements of Example 5.1.4 and Theorem 5.2.1 numerically.
# DDIM with eta = 0 is Euler's method for dx/dsigma = (x - D(x; sigma))/sigma,
# and both are exact when the data law is a single point.
#
# DDIM (alpha = 1): x_next = D(x) + (sigma_next / sigma) (x - D(x))
# Euler: x_next = x + (sigma_next - sigma) (x - D(x)) / sigma
#
# Picture: one step holds the noise estimate fixed.
#
# x ----(x - D)/sigma held constant----> x_next
# D(x) o.....................o D(x) again (DDIM)
# sigma sigma_next
import math
W, M, S = (0.45, 0.35, 0.20), (-1.4, 0.2, 1.6), (0.25, 0.15, 0.30)
def denoise(x, s, comps, a=1.0): # E[x_0 | x_t = x] for a mixture of Gaussians
logw, means, gains = [], [], []
for w, m, sd in comps:
v = a * a * sd * sd + s * s
logw.append(math.log(w) - 0.5 * (x - a * m) ** 2 / v - 0.5 * math.log(2 * math.pi * v))
means.append(m + (a * sd * sd / v) * (x - a * m))
gains.append(v)
top = max(logw)
r = [math.exp(l - top) for l in logw]
z = sum(r)
return sum(ri / z * mi for ri, mi in zip(r, means))
comps = list(zip(W, M, S))
worst = 0.0
for x in (-2.0, -0.4, 0.7, 2.3):
for s0, s1 in ((3.0, 1.68), (1.2, 0.9), (0.44, 0.2)):
D = denoise(x, s0, comps)
ddim = D + (s1 / s0) * (x - D)
euler = x + (s1 - s0) * (x - D) / s0
worst = max(worst, abs(ddim - euler))
print(f"R2: largest |DDIM - Euler| over 12 steps = {worst:.1e}")
# A single Gaussian N(m, sd^2): the exact flow from sigma_0 to sigma_1 is
# x_1 = m + sqrt((sd^2 + sigma_1^2) / (sd^2 + sigma_0^2)) (x_0 - m).
# One DDIM step agrees with it only when the denoiser is constant along the
# step, which is the case of a point mass (sd = 0) and not otherwise.
m, x0, s0, s1 = 0.3, 2.0, 3.0, 0.2
for sd in (0.5, 0.0):
D = denoise(x0, s0, [(1.0, m, sd)])
exact = m + math.sqrt((sd * sd + s1 * s1) / (sd * sd + s0 * s0)) * (x0 - m)
ddim = D + (s1 / s0) * (x0 - D)
print(f"N({m}, {sd}^2), one step from sigma 3.0 to 0.2 at x = 2.0: exact flow {exact:.4f}, DDIM {ddim:.4f}")
R2: largest |DDIM - Euler| over 12 steps = 0.0e+00
N(0.3, 0.5^2), one step from sigma 3.0 to 0.2 at x = 2.0: exact flow 0.6010, DDIM 0.4562
N(0.3, 0.0^2), one step from sigma 3.0 to 0.2 at x = 2.0: exact flow 0.4133, DDIM 0.4133
Sources of error, and the samplers side by side
Nakkiran et al. distinguish errors made during training from errors made during sampling. Training estimates the conditional expectation; sampling approximates a continuous process with finitely many steps. The two errors can interact in ways that are still incompletely understood.Nakkiran, Bradley, Zhou and Advani (2024), Section 5, "Sources of error"; the same section notes that on small training sets these errors "act as a kind of regularization which prevents the diffusion model from just memorizing the train samples", a point Section 10.3 returns to.
The R3 comparison in this part uses the exact denoiser, isolating the solver's error. Under the smoothness conditions discussed in Section 10.1, a deterministic solver of order \(k\) has endpoint error of order \(N^{-k}\). The convergence results stated there also compare rates in distribution of order \(N^{-1/2}\) for the ancestral sampler and \(N^{-1}\) for the flow, under their respective assumptions. These are different kinds of guarantees. Stochastic samplers also have a reason to remain useful beyond their step-count rates: they can be more tolerant of score errors.
| sampler | what it discretizes | evaluations per step | order | noise |
|---|---|---|---|---|
| DDPM, ancestral (Algorithm 3.3.2) | the reverse SDE, \(\lambda = 1\), on the training grid | 1 | \(\tfrac12\) in the step, in law | \(\tilde\beta_i\) or \(\beta_i\) |
| DDIM, \(\eta \in (0, 1)\) (5.1) | the family (4.2) with \(\lambda = \eta\) in effect, on any subgrid | 1 | 1 | \(\eta\, \sigma_{s \mid t}\) |
| DDIM, \(\eta = 0\) | the flow (5.2), Euler in \((\bar x, \rho)\) | 1 | 1 | none |
| DPM-Solver-\(k\) (Proposition 5.2.2) | the flow, exponential integrator in \(\lambda\) | 1 (multistep) or \(k\) | \(k\) | none |
| EDM Heun (Algorithm 5.3.2) | the flow on the clock \(\sigma\) | 2 | 2 | none, or churn \(\gamma_i\) |
| predictor–corrector (Song et al.) | the reverse SDE plus Langevin steps at fixed \(t\) | 2 or more | \(\tfrac12\) | yes |
Example 5.4.2 (the samplers on the ε-model). We compare samplers for the \(\varepsilon\)-model of Section 11, widths (64, 64, 128), trained on FashionMNIST at \(28 \times 28\) on DDPM's chain of 1,000 steps, \(\beta_i\) linear from \(10^{-4}\) to \(0.02\) (Example 3.1.3), and the same architecture trained on the 150-step chain. We sample each model by DDIM with \(\eta = 0\) and by the ancestral chain re-spaced to \(K\) steps, at \(K\) from \(1\) to the full chain, from the same sixteen noises for the tiles and from 1,024 fixed draws of \(\mathcal N(0, I)\) for measuring the Fréchet distance FD_clf of the samples to 1,024 held-out images in the features of the small classifier of Example 3.1.5.
The deterministic sampler's images barely move once \(K\) exceeds ten, because the map from the same noise converges; the ancestral chain's are different draws at every \(K\), by Remark 4.3.5.
On DDPM's chain of 1,000 steps DDIM reaches \(363.1\) at ten steps, \(251.0\) at twenty-five and \(217.7\) at a hundred and fifty, with the ancestral chain at \(155.7\) there and \(208.8\) at twenty-five, on a scale where two halves of the held-out set are \(11.7\) apart and the full thousand steps of DDIM reach \(212.7\): twenty-five steps of DDIM come within a fifth of the full chain's distance and a hundred and fifty within five units of it. This illustrates Song, Meng and Ermon's finding that tens of steps can do much of the work of a thousand, here on a model of two and a half million parameters. The ancestral chain's fresh noise does better still on this model from twenty-five steps on, and keeps improving where the deterministic map has converged.
The 150-step chain shows the starting mismatch from Example 3.1.5: DDIM reaches \(749.2\) at three steps, \(723.5\) at ten and \(731.3\) at the full chain, and the ancestral chain \(641.1\) and \(586.7\), with \(727.9\) and \(599.9\) at twenty-five; neither sampler gets below \(580\), while the same twenty-five steps from the true start \(q(x_{149} \mid x_0)\) give \(20.2\) and \(22.3\), so the error is the start and not the steps.
The figure's last block runs the browser-sized model, the (16, 32, 64) member of the same architecture trained on DDPM's chain of 1,000 steps, in the page: its weights load when you press run, its samples at ten DDIM steps are garments, and the reader can watch the clean estimate settle. ◁
strips sixteen samples of the epsilon-model trained on DDPM's chain of 1,000 steps, or on the
150-step chain, by DDIM (eta = 0) and by the re-spaced ancestral chain (eta = 1) at K = 1 to
the full chain from the same noise, and their difference DDIM - ancestral as a change map:
the deterministic sampler's images barely move with K, the stochastic one's are different
draws at every K
chart beside them, the Frechet distance FD_clf against K for both samplers on both chains: at
K = 25 on DDPM's chain of 1,000 steps DDIM 251.0 and the ancestral chain 208.8, against
the held-out floor 11.7; on the 150-step chain 727.9 and 599.9 (Example 3.1.5 says why)
live the browser-sized (16, 32, 64) model trained on DDPM's chain of 1,000 steps,
576,897 parameters, run in a Web Worker in this page: four samples by DDIM in ten steps,
painted step by step, beside the film of the followed sample's clean estimate
(the live figure needs JavaScript and loads 0.65 MB of data, and about 1.2 MB of weights when you
press run)
Remark 5.4.1 (takeaway: the sampler is chosen after training). Example 5.4.2 uses one ε-model with two samplers and many step counts. There is no retraining between runs. More generally, a denoiser trained by Algorithm 2.3.7 can supply the score needed by DDPM's ancestral chain, every \(\eta\) of DDIM, Heun on the EDM grid, and the exponential integrators. These methods discretize members of the same family (4.2).
Better numerical methods reduced the required steps from thousands to tens. Reaching one or a few steps asks more of each update: it must account for how the velocity changes over a substantial part of the path. Section 8 studies methods that learn this larger move directly.