Consistency Model

A consistency model is a generative model that maps any noisy image on a trajectory of a diffusion model’s probability-flow ODE directly to the clean image at that trajectory’s origin. Because pure noise lies on the same trajectories, one network evaluation turns noise into an image, and a few more evaluations, alternating with fresh noise, can improve it. Song, Dhariwal, Chen and Sutskever (OpenAI) introduced consistency models at ICML 2023, as a way to keep what iterative diffusion sampling offers, a trade between compute and quality and editing without task-specific training, while generating in a single step. A consistency model can be trained by consistency distillation (CD) from a pretrained diffusion model or by consistency training (CT) from data alone. The network is the diffusion model’s own architecture, a convolutional U-Net in the paper, with a modified skip connection; later work applied them to Stable Diffusion models as latent consistency models (LCM).

Problem Addressed

A diffusion model generates by solving its reverse SDE or probability-flow ODE in many small steps, each a pass through a large network; the samplers in Diffusion Sampling Methods reduce the number of these network function evaluations (NFE) but, in Song et al.’s summary, still needed more than 10 for competitive samples (§2). Their introduction puts the cost of iterative generation at typically 10 to 2,000 times that of single-step generators such as GANs, variational autoencoders and normalizing flows. Distillation trains a faster student instead. Song et al. note that most distillation methods, such as the knowledge distillation of Luhman and Luhman (2021) and DFNO, first build a large dataset of samples by running the diffusion model, and that progressive distillation (Salimans & Ho, 2022), which the Diffusion Sampling article summarizes, was the only one they knew of that avoids this (§2, §6.2). They wanted a model that samples in one step, can still spend more steps for better samples, and keeps the zero-shot editing of diffusion models, without adversarial training (§1).

Architecture

The Consistency Function

Song et al. work in the setting of Karras et al. (2022), the EDM formulation that Diffusion Sampling uses: the forward process adds Gaussian noise of standard deviation equal to the time $t$, so the noisy marginal is $p_t = p_{\text{data}} \otimes \mathcal{N}(0, t^2 I)$, with pixel values in $[-1, 1]$, $T = 80$ and a smallest time $\epsilon = 0.002$ (§2). A score model $s_\phi(x, t) \approx \nabla_x \log p_t(x)$, trained by score matching, gives the empirical probability-flow ODE

$$\frac{dx_t}{dt} = -t\, s_\phi(x_t, t)$$

(Eq. 3). In the denoiser form of the Diffusion Sampling article this is $dx/dt = (x - D(x; t))/t$, since the optimal denoiser satisfies $s = (D - x)/t^2$. Its solutions run from data at $t = \epsilon$ to approximately $\mathcal{N}(0, T^2 I)$ at $t = T$.

For a solution trajectory $\{x_t\}_{t \in [\epsilon, T]}$, the consistency function is $f: (x_t, t) \mapsto x_\epsilon$. Its defining property, self-consistency, is that $f(x_t, t) = f(x_{t'}, t')$ for all $t, t'$ on the same trajectory (§3). A consistency model $f_\theta$ is a network trained to have this property. The authors point out that neural flows (Biloš et al., 2021) use a similar definition, but consistency models are not required to be invertible.

Boundary Condition and Parameterization

At the start of a trajectory, $f(x_\epsilon, \epsilon) = x_\epsilon$: the consistency function is the identity at $t = \epsilon$. Song et al. call this the boundary condition, call it the most confining architectural constraint, and say it is crucial for successful training (§3). They give two ways to satisfy it with a free-form network $F_\theta(x, t)$ whose output has the shape of $x$:

$$f_\theta(x, t) = \begin{cases} x & t = \epsilon \\ F_\theta(x, t) & t \in (\epsilon, T] \end{cases} \qquad \text{or} \qquad f_\theta(x, t) = c_{\text{skip}}(t)\, x + c_{\text{out}}(t)\, F_\theta(x, t)$$

(Eqs. 4 and 5), the second with differentiable $c_{\text{skip}}(\epsilon) = 1$ and $c_{\text{out}}(\epsilon) = 0$. The second keeps $f_\theta$ differentiable at $\epsilon$, which the continuous-time losses need, and resembles the preconditioning of EDM, whose denoiser is $D_\theta(x; \sigma) = c_{\text{skip}}(\sigma)\, x + c_{\text{out}}(\sigma)\, F_\theta\big(c_{\text{in}}(\sigma)\, x;\ c_{\text{noise}}(\sigma)\big)$ (Karras et al., 2022, Eq. 7). All experiments use it. EDM’s $c_{\text{skip}}$ and $c_{\text{out}}$ give 1 and 0 only at $t = 0$, so the paper shifts them to $\epsilon$:

$$c_{\text{skip}}(t) = \frac{\sigma_{\text{data}}^2}{(t - \epsilon)^2 + \sigma_{\text{data}}^2}, \qquad c_{\text{out}}(t) = \frac{\sigma_{\text{data}}\,(t - \epsilon)}{\sqrt{\sigma_{\text{data}}^2 + t^2}}, \qquad \sigma_{\text{data}} = 0.5$$

(Appendix C; the official code’s get_scalings_for_boundary_condition computes the same, with EDM’s input scaling $c_{\text{in}} = 1/\sqrt{\sigma_{\text{data}}^2 + t^2}$). The boundary condition also rules out the trivial solution $f_\theta \equiv 0$, which would otherwise satisfy self-consistency perfectly (§4).

The network $F_\theta$ is the architecture of the diffusion model being distilled: NCSN++ (Song et al., 2021) for CIFAR-10, and the architectures of Dhariwal and Nichol (2021) for ImageNet $64 \times 64$ and LSUN Bedroom and Cat $256 \times 256$ (Appendix C). The ImageNet models are class-conditional (official launch.sh).

Training

Both methods train $f_\theta$ by making its outputs agree at two adjacent times on one trajectory. The time horizon is split into $N - 1$ intervals with boundaries $t_1 = \epsilon < t_2 < \dots < t_N = T$ placed by EDM’s rule,

$$t_i = \Big(\epsilon^{1/\rho} + \frac{i - 1}{N - 1}\big(T^{1/\rho} - \epsilon^{1/\rho}\big)\Big)^{\rho}, \qquad \rho = 7$$

(§4), the step schedule Diffusion Sampling describes.

Consistency Distillation

Given a data image $x$, CD draws $n$ uniformly from $1, \dots, N - 1$, noises the image to $x_{t_{n+1}} \sim \mathcal{N}(x, t_{n+1}^2 I)$, and takes one step of an ODE solver $\Phi$ with the pretrained score model to estimate the point one level lower on the same trajectory,

$$\hat x^{\phi}_{t_n} = x_{t_{n+1}} + (t_n - t_{n+1})\, \Phi(x_{t_{n+1}}, t_{n+1}; \phi)$$

(Eq. 6; for Euler, $\Phi = -t\, s_\phi(x, t)$). The consistency distillation loss compares the model’s outputs at the two points:

$$\mathcal{L}^N_{\text{CD}}(\theta, \theta^-; \phi) = \mathbb{E}\Big[\lambda(t_n)\, d\big(f_\theta(x_{t_{n+1}}, t_{n+1}),\ f_{\theta^-}(\hat x^{\phi}_{t_n}, t_n)\big)\Big]$$

(Eq. 7), where $d$ is a metric ($d(x, y) \ge 0$, zero only when $x = y$) and $\lambda > 0$ a weight. The paper tried squared $\ell_2$, $\ell_1$ and LPIPS, and found $\lambda \equiv 1$ worked well across tasks and datasets. The parameters $\theta^-$ of the target network are an exponential moving average of the online network’s,

$$\theta^- \leftarrow \operatorname{stopgrad}\big(\mu\, \theta^- + (1 - \mu)\, \theta\big)$$

(Eq. 8), and only $\theta$ receives gradients (Algorithm 2). The authors relate this to the target networks of deep Q-learning and to momentum-based contrastive learning (BYOL, MoCo), and found that the EMA and stop-gradient greatly stabilized training compared with $\theta^- = \theta$ (§4). Training starts from the pretrained diffusion model’s weights (Appendix C).

Theorem 1 justifies the loss: if $f_\theta$ is Lipschitz, the solver has local error $O((t_{n+1} - t_n)^{p+1})$ and the loss with $\theta^- = \theta$ is zero, then $f_\theta$ differs from the consistency function of the empirical ODE by $O((\Delta t)^p)$, where $\Delta t$ is the largest interval. The proof is the one for global error bounds of ODE solvers: the error is zero at $t_1$ by the boundary condition, and each interval adds at most one local solver error (Appendix A.2). A higher-order solver therefore gives a smaller error at the same $N$. In our reading, each target sits one solver step closer to the data than its input, so correct outputs enter training at the boundary and spread to higher noise levels.

Consistency Training

CT removes the pretrained model. For $x \sim p_{\text{data}}$ and $x_t \sim \mathcal{N}(x, t^2 I)$, the score has the unbiased estimate $\nabla \log p_t(x_t) = -\mathbb{E}[(x_t - x)/t^2 \mid x_t]$ (Lemma 1). Putting $-(x_t - x)/t^2$ in place of the score in one Euler step moves $x + t_{n+1} z$ exactly to $x + t_n z$ with the same noise $z$, which gives the consistency training loss

$$\mathcal{L}^N_{\text{CT}}(\theta, \theta^-) = \mathbb{E}\Big[\lambda(t_n)\, d\big(f_\theta(x + t_{n+1} z, t_{n+1}),\ f_{\theta^-}(x + t_n z, t_n)\big)\Big], \qquad z \sim \mathcal{N}(0, I)$$

(Eq. 10, Algorithm 3). Theorem 2 states that, with the Euler solver and a score model equal to the true score, $\mathcal{L}^N_{\text{CD}} = \mathcal{L}^N_{\text{CT}} + o(\Delta t)$, under smoothness assumptions on $d$ and $f_{\theta^-}$. Small $N$ gives a loss with less variance but more bias relative to CD, which speeds early training; large $N$ gives less bias and more variance, which suits the end of training (§5). CT therefore grows $N$ and the EMA rate $\mu$ during training with the schedules

$$N(k) = \left\lceil \sqrt{\frac{k}{K}\big((s_1 + 1)^2 - s_0^2\big) + s_0^2} - 1 \right\rceil + 1, \qquad \mu(k) = \exp\Big(\frac{s_0 \log \mu_0}{N(k)}\Big)$$

for iteration $k$ of $K$ (Appendix C). The CT models were initialized randomly.

Continuous-Time Losses

The CD loss extends to $N \to \infty$ when $\theta^- = \theta$ or $\theta^- = \operatorname{stopgrad}(\theta)$, and the CT loss when $\theta^- = \operatorname{stopgrad}(\theta)$ (Theorems 3 to 6, Appendix B). They need no time grid, but involve Jacobian–vector products and need forward-mode automatic differentiation. In the CIFAR-10 comparisons of Appendix B.3, discrete-time CD outperformed continuous-time CD, possibly, the authors suggest, because of the higher variance of the continuous objectives and because discrete CD can use a higher-order solver; continuous-time CT beat discrete-time CT with LPIPS in a comparison where both were initialized from a pretrained EDM, which the continuous version needed for stable training.

Choices and Settings

On CIFAR-10 (§6.1, Figure 3), LPIPS beat $\ell_1$ and $\ell_2$ by a large margin as the CD metric over all training iterations; Heun’s second-order solver beat Euler at each $N$ compared in Figure 3b (50, 80 and 120), which the authors connect to Theorem 1; and $N = 18$ was best among 9 to 120, with performance insensitive to $N$ once it was large enough. CD therefore used LPIPS and Heun, and CT used LPIPS too; CT needs no solver. CT’s convergence was highly sensitive to $N$, and the adaptive schedules improved both speed and final quality. For LPIPS on $32 \times 32$ and $64 \times 64$ images, both images were upsampled bilinearly to $224 \times 224$ first (Appendix C).

Table 3 lists the settings. CIFAR-10 used batch 512, learning rate $4 \times 10^{-4}$, 800,000 iterations on 8 GPUs, $\mu = 0$ for CD (the target network is a stop-gradient copy of the online one) and $s_0 = 2$, $s_1 = 150$, $\mu_0 = 0.9$ for CT; ImageNet $64 \times 64$ and LSUN used batch 2,048, mixed precision, 64 GPUs, $N = 40$ and $\mu = 0.95$ for CD, and 600,000 to 1,000,000 iterations. All models were trained with Rectified Adam, without learning-rate decay, warm-up or weight decay, with horizontal flips, and with an EMA of the online weights as well (rate 0.9999 on CIFAR-10). Every distilled model, CD and progressive distillation alike, came from EDM models the authors trained themselves.

Inference

One step. Draw $\hat x_T \sim \mathcal{N}(0, T^2 I)$ and return $f_\theta(\hat x_T, T)$: one forward pass of the network.

Multistep sampling. Algorithm 1 alternates denoising and noise injection: after the first step, for each of a decreasing sequence of times $\tau_1 > \tau_2 > \dots > \tau_{N-1}$ it adds noise to the current image, $\hat x_{\tau_n} = x + \sqrt{\tau_n^2 - \epsilon^2}\, z$, and maps it back, $x \leftarrow f_\theta(\hat x_{\tau_n}, \tau_n)$. Each step costs one evaluation. The authors chose the times greedily, one at a time, by ternary search on the FID of the samples, assuming that FID is unimodal in the next time point, which they found to hold (§3). The official launch.sh gives the resulting two-step settings as indices into the training grid: steps 0, 22 and 39 of 40 for CD on ImageNet $64 \times 64$, and 0, 106 and 200 of 201 for CT. Unlike a deterministic ODE sampler, multistep consistency sampling is stochastic: the image depends on the injected noise as well as on $\hat x_T$.

Zero-Shot Editing

Consistency models inherit the editing that diffusion models do without task-specific training (§3, §6.3, Appendix D). With a reference image $y$, a binary mask $\Omega$ and an invertible linear transform $A$ into a space where the known information is separate from the unknown, Algorithm 4 runs multistep sampling and after every step overwrites the known part of $A x$ with that of $A y$:

$$x \leftarrow A^{-1}\big[(A y) \odot (1 - \Omega) + (A x) \odot \Omega\big].$$

The time points follow EDM’s schedule, with $N = 40$ for LSUN Bedroom. The tasks differ only in $A$ and $\Omega$:

• Inpainting: $A$ is the identity and $\Omega$ marks the missing pixels.
• Colorization: $A$ multiplies each pixel’s colour by a $3 \times 3$ orthogonal matrix whose first column is proportional to the luminance weights $(0.2989, 0.5870, 0.1140)$, and $\Omega$ marks the other two channels, so colorization becomes inpainting in a colour space whose first channel is the grey level.
• Super-resolution: a $p^2 \times p^2$ orthogonal matrix whose first column is constant separates each $p \times p$ patch’s mean, which the low-resolution image fixes, from the rest.
• Stroke-guided generation: as in SDEdit (Meng et al., 2022), which the Image Editing article describes, $A = I$ and $\Omega$ is all ones, so, in our reading, the stroke painting enters only as the starting image; the paper used two steps at $t = 5.38$ and $2.24$.

Denoising needs no algorithm: an image with Gaussian noise of standard deviation $\sigma \in [\epsilon, T]$ is denoised by $f_\theta(x, \sigma)$. Interpolation feeds spherical interpolations of two noise vectors through $f_\theta(\cdot, T)$. All editing results in the paper are images from a CD model on LSUN Bedroom $256 \times 256$ (Figures 6 and 8 to 13), with no quantitative evaluation.

Performance

The paper reports FID and Inception Score on CIFAR-10 (Table 1) and FID with the precision and recall of Kynkäänniemi et al. (2019) on the others (Table 2).

Distillation. CD with LPIPS reached one-step FIDs of 3.55 on CIFAR-10, 6.20 on ImageNet $64 \times 64$, 7.80 on LSUN Bedroom and 11.0 on LSUN Cat, and two-step FIDs of 2.93, 4.70, 5.22 and 8.84. Progressive distillation (PD) of the same EDM teachers, with its original $\ell_2$ loss, scored 8.34, 15.39, 16.92 and 29.6 in one step and 5.58, 8.95, 8.47 and 15.5 in two. CD also beat distillation methods that build a synthetic dataset first: knowledge distillation (9.36), DFNO (4.12) and rectified flow with distillation (4.85 at best) on CIFAR-10. Figure 4 compares PD and CD with both metrics and up to several steps; CD was better throughout except for one-step generation on LSUN Bedroom, where CD with $\ell_2$ was slightly worse than PD with $\ell_2$. Using LPIPS improved PD too.

Direct generation. CT reached one-step FIDs of 8.70 on CIFAR-10, 13.0 on ImageNet $64 \times 64$, 16.0 on LSUN Bedroom and 20.7 on LSUN Cat, and two-step FIDs of 5.83, 11.1, 7.85 and 11.7. On CIFAR-10 this beat the single-step non-adversarial models in Table 1, such as Glow (48.9), DenseFlow (34.9) and DC-VAE (17.9), and several GANs, but not StyleGAN2-ADA (2.92) or StyleGAN-XL (1.85); the conclusion accordingly says that CT beats existing single-step models except GANs. The teachers needed many more evaluations for their quality: EDM scored 2.04 at 35 NFE on CIFAR-10 and 2.44 at 79 NFE on ImageNet $64 \times 64$. Samples of CT and of EDM from the same initial noise share much of their structure (Figure 5), which the authors read as a sign that CT is less likely to suffer from mode collapse.

Strengths

• One network evaluation per image, with the option to spend a few more for better samples (Algorithm 1).
• Training needs neither adversarial losses nor a dataset of teacher samples; CD needs one teacher solver step per training example (§6.2).
• Any diffusion architecture can be reused, since the boundary condition costs only a change of skip-connection scaling.
• Editing, denoising and interpolation work zero-shot, as with diffusion models, and the map from noise to image is deterministic for one-step sampling.

Limitations

• Quality gap. The best one-step and two-step FIDs remained above those of the diffusion teachers at tens of NFE, and CT trailed CD at equal steps (Tables 1 and 2).
• LPIPS. The best results use LPIPS, a VGG network trained on ImageNet. The model card warns that this can overstate quality, because FID and Inception Score also use a network trained on ImageNet, so some knowledge may leak between the training metric and the evaluation.
• Tuned sampling. Multistep time points were chosen by searching on FID for each model, and settings such as $N$, $\mu$ and the CT schedules were tuned per resolution (§6.1).
• Training cost. The reported models trained for 600,000 to 1,000,000 iterations on 8 to 64 GPUs (Table 3); continuous-time CT needed a pretrained initialization to train stably (Appendix B.3).
• Stated by the model card: the released models sometimes produce highly unrealistic images, particularly of human faces, and are meant for research only.

Historical Importance

Before consistency models, fast diffusion sampling came either from better ODE solvers or from distillation into a student that imitates a multi-step sampler. Consistency models recast the problem as learning the solution map of the probability-flow ODE, defined pointwise by self-consistency and anchored by the boundary condition, which gave one training rule for distillation and for training from scratch, and a sampler that can still trade steps for quality. The 2023 paper reported new state-of-the-art one-step FIDs for distilled diffusion models on CIFAR-10 and ImageNet $64 \times 64$, and later work, from improved consistency training to latent consistency models for text-to-image generation, builds on its definitions.

Variants

Improved consistency training (iCT). Song and Dhariwal (ICLR 2024) revised CT; only their poster was read here. It lists a weighting $\lambda(t_n) = 1/(t_{n+1} - t_n)$ that weights small noise more, a pseudo-Huber metric $d(x, y) = \sqrt{\lVert x - y \rVert_2^2 + c^2} - c$ in place of LPIPS, which outperformed both the squared loss and LPIPS in CT, no EMA in the teacher network ($\theta^- = \operatorname{stopgrad}(\theta)$), time steps drawn from a discretized lognormal distribution, a schedule that doubles the number of time steps every fixed number of iterations, less sensitive time embeddings and larger dropout. The poster reports one-step FIDs of 2.51 on CIFAR-10 and, on ImageNet $64 \times 64$, 4.02 for iCT and 3.25 for the model labelled iCT-deep, with two-step FIDs of 3.20 and 2.77, against 6.20 and 4.70 for CD with LPIPS.

Latent consistency models (LCM). Luo et al. (2023) introduced latent consistency models; the paper was not read here. According to the official repository, whose training scripts perform consistency distillation for Stable Diffusion models, a family of latent diffusion models, LCM distills classifier-free guidance into the model’s input and generates in very few steps, with 1 to 8 recommended; the released models include an LCM distilled from Dreamshaper v7, full-parameter LCMs for SDXL and SSD-1B, and LCM-LoRA adapters for SDXL, SSD-1B and Stable Diffusion 1.5, a low-rank module for accelerating Stable Diffusion models, and Diffusers has an LCM pipeline and scheduler (README, as of November 2026). In Sauer et al.’s user studies, one step of their adversarially distilled SDXL was preferred over four steps of LCM-XL (see the Stable Diffusion article).

Consistency trajectory models (CTM). Kim et al. (ICLR 2024) titled their paper after learning the probability-flow ODE trajectory of diffusion; the paper was not read here. Their official repository reports single-step FIDs of 1.73 on CIFAR-10 and 1.92 on ImageNet $64 \times 64$, and training commands labelled CTM+DSM and CTM+DSM+GAN.

Continuous-time consistency models (sCM). Lu and Song (ICLR 2025) published Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models; the paper was not read here.

Consistency distillation is one of several ways of distilling diffusion models. Progressive distillation is summarized in Diffusion Sampling Methods, guided distillation in Classifier-Free Guidance, and adversarial diffusion distillation, the method behind the Turbo models, in Stable Diffusion.

Implementations

• Official code. The openai/consistency_models repository (MIT licence) implements CD, CT, one-step and multistep sampling, the ternary search for time points and the editing algorithms in PyTorch for ImageNet $64 \times 64$ and LSUN $256 \times 256$; the CIFAR-10 experiments are in a separate JAX repository, openai/consistency_models_cifar10 (README). As of November 2026 the README links twelve checkpoints, an EDM teacher, CD with $\ell_2$, CD with LPIPS and CT for each of the three datasets, and a model card. The example commands in launch.sh do not all match Table 3: LSUN CD and CT run with a global batch of 256 rather than 2,048, ImageNet CT with a learning rate of $10^{-4}$ rather than $8 \times 10^{-6}$, and LSUN CT with $5 \times 10^{-5}$ rather than $10^{-5}$.
• Diffusers. The README shows the released models running through the ConsistencyModelPipeline of Hugging Face’s Diffusers library.

Example: Distilling an Exact Probability-Flow ODE

The example distils a consistency model from a teacher with no error of its own, so that every remaining error belongs to the consistency model. The data are the Gaussian mixture of Diffusion Sampling’s example: four $16 \times 16$ shapes (disk, square, plus sign, diagonal bar) with pixel values $\pm 1$, weights 0.4, 0.3, 0.2 and 0.1, and pixel noise of standard deviation 0.1. In the paper’s convention $x_t = x_0 + t z$, every $p_t$ is again a mixture, so the optimal denoiser $D(x, t) = \mathbb{E}[x_0 \mid x_t = x]$ is exact, and the ODE endpoint $x_\epsilon$ of any noisy image can be computed to high accuracy with 2,000 Heun steps. Diffusion Sampling’s example compares samplers by the shapes they produce; this one scores the learned map instead, by its distance to the exact endpoint.

Two consistency models are distilled with Algorithm 2 and the paper’s CIFAR-10 settings where they apply ($N = 18$, $\mu = 0$, $\lambda \equiv 1$, Rectified Adam with learning rate $4 \times 10^{-4}$, batch 512): one whose teacher step is Heun’s method, the paper’s choice, and one with Euler’s, the comparison of the paper’s Figure 3b, which used $N = 50$, 80 and 120 rather than 18. Three things differ from the paper. The metric is squared $\ell_2$, because LPIPS needs a VGG network trained on ImageNet; $F_\theta$ is a small multilayer perceptron with a sinusoidal embedding of $c_{\text{noise}} = \tfrac{1}{4}\ln t$, the EDM choice, which the consistency-model paper does not state (the official code feeds $1000 \cdot \tfrac{1}{4}\ln t$ to its time embedding); and training runs 4,000 iterations, not 800,000, and is evaluated with the online weights, without the additional EMA of those weights that the paper kept. The parameterization is that of Appendix C. Each image is a vector of 256 values; distances are root-mean-square (RMS) per pixel. The script prints, for 1,000 test noise images, the distance of $f_\theta(x_t, t)$ to the endpoint at points $x_t$ along the exact trajectories, taken at the levels of the 2,000-step solve just above 20, 5, 1, 0.2 and 0.05 (with the exact denoiser’s own estimate $D(x_t, t)$ for comparison), and the distance of two-step samples from Algorithm 1 run with the exact consistency function on the same noise. The two-step time $\tau$, which the paper finds by searching on FID, is not fixed by the sources, so three values are shown.

The script needs NumPy and PyTorch (output from PyTorch 2.14.1 and NumPy 2.5.4, four CPU threads); it ran in three to five minutes on a laptop CPU shared with other jobs, most of it training.

import sys
import numpy as np
import torch
import torch.nn as nn

seed = int(sys.argv[1]) if len(sys.argv) > 1 else 0
iters = int(sys.argv[2]) if len(sys.argv) > 2 else 4000
torch.set_num_threads(4)

# Data of the Diffusion Sampling article: four 16x16 shapes in [-1, 1], x0 = mu_k + 0.1 * noise.
yy, xx = np.mgrid[:16, :16]
shapes = [(yy - 7.5) ** 2 + (xx - 7.5) ** 2 < 5.5 ** 2,          # disk
          (abs(yy - 7.5) < 5) & (abs(xx - 7.5) < 5),               # square
          (abs(yy - 7.5) < 2) | (abs(xx - 7.5) < 2),               # plus
          abs(yy - xx) < 2.5]                                      # diagonal bar
mu = np.array([np.where(m, 1.0, -1.0).ravel() for m in shapes])   # (4, 256)
w = np.array([0.4, 0.3, 0.2, 0.1])
tau2 = 0.1 ** 2
EPS, T, RHO, SD = 0.002, 80.0, 7.0, 0.5                           # paper's eps, T, rho; sigma_data

def denoise(x, t):
    """Exact E[x0 | x_t = x] for x_t = x0 + t z (EDM convention, Song et al. 2023, Sec. 2)."""
    t = np.broadcast_to(np.asarray(t, float).reshape(-1), (len(x),))[:, None]
    v = tau2 + t ** 2
    d2 = (x ** 2).sum(1, keepdims=True) - 2 * x @ mu.T + (mu ** 2).sum(1)
    lg = np.log(w) - d2 / (2 * v)
    r = np.exp(lg - lg.max(1, keepdims=True))
    r /= r.sum(1, keepdims=True)
    m = r @ mu
    return m + tau2 / v * (x - m)

def karras(n, lo=EPS, hi=T):                    # n levels from hi down to lo (rho = 7)
    i = np.arange(n) / (n - 1)
    return (hi ** (1 / RHO) + i * (lo ** (1 / RHO) - hi ** (1 / RHO))) ** RHO

def heun(x, t, t_next):                         # one Heun step of dx/dt = (x - D(x, t)) / t
    d = (x - denoise(x, t)) / t
    x_e = x + (t_next - t) * d
    return x + (t_next - t) * 0.5 * (d + (x_e - denoise(x_e, t_next)) / t_next)

def ode_endpoint(x, t_start, steps=2000, keep=None):
    """Integrate the exact PF ODE from t_start down to eps; optionally record x at levels in keep."""
    ts, rec = karras(steps, EPS, t_start), {}
    for a, b in zip(ts[:-1], ts[1:]):
        if keep is not None:
            for k in keep:
                if a >= k > b and k not in rec:  # record the state at the nearest level above k
                    rec[k] = (x.copy(), a)
        x = heun(x, a, b)
    return (x, rec) if keep is not None else x

# ---------- consistency model: f(x, t) = c_skip x + c_out F(c_in x, c_noise) (App. C) ----------
class F(nn.Module):
    def __init__(s, d=256, h=512):
        super().__init__()
        s.emb = nn.Linear(1, 64)
        s.net = nn.Sequential(nn.Linear(d + 64, h), nn.SiLU(), nn.Linear(h, h), nn.SiLU(),
                              nn.Linear(h, h), nn.SiLU(), nn.Linear(h, d))
    def forward(s, x, c_noise):
        return s.net(torch.cat([x, torch.sin(s.emb(c_noise[:, None]))], 1))

def f_theta(net, x, t):
    t = t[:, None]
    c_skip = SD ** 2 / ((t - EPS) ** 2 + SD ** 2)
    c_out = SD * (t - EPS) / torch.sqrt(SD ** 2 + t ** 2)
    c_in = 1 / torch.sqrt(SD ** 2 + t ** 2)
    return c_skip * x + c_out * net(c_in * x, torch.log(t[:, 0]) / 4)

def sample_data(rng, n):
    return mu[rng.choice(4, n, p=w)] + 0.1 * rng.standard_normal((n, 256))

def train(mode, iters, seed, batch=512, lr=4e-4):
    """Consistency distillation (Algorithm 2) with N = 18 and mu = 0, as for CIFAR-10 (Table 3);
    mode 'cd-heun' or 'cd-euler' picks the solver step taken with the exact denoiser."""
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    net, tgt = F(), F()
    tgt.load_state_dict(net.state_dict())
    opt = torch.optim.RAdam(net.parameters(), lr=lr)
    for k in range(iters):
        N, ema = 18, 0.0
        ts = karras(N)[::-1]                     # t_1 = eps < ... < t_N = T
        n = rng.integers(0, N - 1, batch)        # index of t_n; t_{n+1} = ts[n + 1]
        t_n, t_n1 = ts[n], ts[n + 1]
        x = sample_data(rng, batch)
        z = rng.standard_normal((batch, 256))
        x_n1 = x + t_n1[:, None] * z
        if mode == "cd-euler":
            x_n = x_n1 + (t_n - t_n1)[:, None] * (x_n1 - denoise(x_n1, t_n1)) / t_n1[:, None]
        else:
            x_n = heun(x_n1, t_n1[:, None], t_n[:, None])
        f = lambda m, a, b: f_theta(m, torch.tensor(a, dtype=torch.float32),
                                    torch.tensor(b, dtype=torch.float32))
        with torch.no_grad():
            target = f(tgt, x_n, t_n)
        loss = ((f(net, x_n1, t_n1) - target) ** 2).sum(1).mean()   # squared l2, lambda = 1
        opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():                    # theta- <- stopgrad(mu theta- + (1 - mu) theta)
            for p_t, p in zip(tgt.parameters(), net.parameters()):
                p_t.mul_(ema).add_((1 - ema) * p)
    return net

def cm(net, x, t):
    with torch.no_grad():
        tt = torch.full((len(x),), float(t))
        return f_theta(net, torch.tensor(x, dtype=torch.float32), tt).double().numpy()

def rms(a, b):
    return np.sqrt(((a - b) ** 2).mean(1))

# ---------- test set: 1,000 noise images and their exact ODE endpoints ----------
rng = np.random.default_rng(1000 + seed)
xT = T * rng.standard_normal((1000, 256))
levels = [20.0, 5.0, 1.0, 0.2, 0.05]
x_end, rec = ode_endpoint(xT, T, keep=levels)
print(f"seed {seed}, {iters} iterations per model")
print(f"closest pair of shape means: RMS distance {min(rms(mu[i:i+1], mu[j:j+1])[0] for i in range(4) for j in range(i+1, 4)):.3f}")
print(f"ODE endpoints, 2000 vs 4000 Heun steps: max RMS difference {rms(x_end, ode_endpoint(xT, T, 4000)).max():.1e}")

z2 = rng.standard_normal((1000, 256))
taus = [0.5, 2.0, 8.0]
ref2 = {tau: ode_endpoint(x_end + np.sqrt(tau ** 2 - EPS ** 2) * z2, tau) for tau in taus}

nets = {m: train(m, iters, seed) for m in ["cd-heun", "cd-euler"]}

print("\nmedian RMS distance to the exact ODE endpoint of the same x_t (90th percentile in brackets)")
head = f"{'t':>8s}{'exact D(x_t,t)':>18s}" + "".join(f"{m:>18s}" for m in nets)
print(head)
pts = [(T, xT)] + [(rec[k][1], rec[k][0]) for k in levels]
for t, xt in pts:
    cols = [denoise(xt, t)] + [cm(nets[m], xt, t) for m in nets]
    print(f"{t:8.3f}" + "".join(f"{np.median(rms(c, x_end)):10.4f} [{np.percentile(rms(c, x_end), 90):.3f}]"
                                for c in cols))
print(f"boundary t = eps: max |f(x, eps) - x| = "
      f"{max(np.abs(cm(nets[m], x_end, EPS) - x_end).max() for m in nets):.1e}")

print("\ntwo-step sampling (Algorithm 1): median RMS distance to Algorithm 1 run with the exact")
print("consistency function on the same x_T and z (one-step row: distance to the ODE endpoint)")
print(f"{'':>12s}" + "".join(f"{m:>12s}" for m in nets))
print(f"{'one step':>12s}" + "".join(f"{np.median(rms(cm(nets[m], xT, T), x_end)):12.4f}" for m in nets))
for tau in taus:
    row = []
    for m in nets:
        x1 = cm(nets[m], xT, T)
        x2 = cm(nets[m], x1 + np.sqrt(tau ** 2 - EPS ** 2) * z2, tau)
        row.append(np.median(rms(x2, ref2[tau])))
    print(f"{'tau = ' + str(tau):>12s}" + "".join(f"{v:12.4f}" for v in row))

Output:

seed 0, 4000 iterations per model
closest pair of shape means: RMS distance 0.433
ODE endpoints, 2000 vs 4000 Heun steps: max RMS difference 7.9e-07

median RMS distance to the exact ODE endpoint of the same x_t (90th percentile in brackets)
       t    exact D(x_t,t)           cd-heun          cd-euler
  80.000    0.3708 [0.954]    0.2343 [0.318]    0.1417 [0.291]
  20.033    0.3311 [0.731]    0.2300 [0.313]    0.1337 [0.281]
   5.005    0.1333 [0.231]    0.1931 [0.294]    0.1036 [0.226]
   1.003    0.0899 [0.095]    0.1800 [0.249]    0.1129 [0.216]
   0.200    0.0552 [0.058]    0.1510 [0.224]    0.1278 [0.250]
   0.050    0.0107 [0.011]    0.0544 [0.081]    0.0468 [0.088]
boundary t = eps: max |f(x, eps) - x| = 6.0e-08

two-step sampling (Algorithm 1): median RMS distance to Algorithm 1 run with the exact
consistency function on the same x_T and z (one-step row: distance to the ODE endpoint)
                 cd-heun    cd-euler
    one step      0.2343      0.1417
   tau = 0.5      0.2271      0.1795
   tau = 2.0      0.1570      0.0821
   tau = 8.0      0.2188      0.1233

The boundary condition holds to float32 rounding, as the parameterization guarantees. At the highest noise levels both distilled models are much closer to the true endpoint than the exact denoiser: $D(x_T, T)$ is the posterior mean of the clean image, a blend of shapes when the noise has nearly erased which shape it came from, and its 90th-percentile distance (0.954) is more than twice the distance between the two most similar shapes, whereas a consistency function must commit to the single image at the end of the trajectory. From $t \approx 1$ down the comparison reverses for both models, and from $t \approx 5$ for the Heun model: there the posterior mean is already close to the endpoint, while both networks keep an error of about 0.1 to 0.2 down to $t = 0.2$, which falls only near the boundary, where $c_{\text{skip}}$ approaches 1. With an exact teacher, what remains is the network’s imperfect fit after 4,000 iterations together with the discretization error of the 18-level grid that Theorem 1 bounds; the example does not separate the two. The Heun model’s error grows steadily with $t$, as errors accumulating from the boundary would, but the Euler model’s does not.

Against the paper’s Figure 3b, the Euler-distilled model has the smaller median distance at every level and in every two-step setting, although its 90th percentile is larger at $t = 0.2$ and $0.05$. The two runs share the seed, the initial weights and every data and noise draw, since the teacher step consumes no random numbers, so the difference comes from the teacher step alone. In our reading, the networks are far from the zero-loss regime in which Theorem 1 favours the more accurate solver, and the ordering reflects how the two sets of targets shape optimization in a short run; the example does not show which, and the paper’s comparison used LPIPS, NCSN++, $N$ of 50 to 120 and up to 300,000 iterations.

Two-step sampling helps both models at $\tau = 2$ and $\tau = 8$, most at $\tau = 2$, where each lands closer to the exact two-step result than its one-step sample is to the exact endpoint; $\tau = 0.5$ helps the Heun model only slightly and makes the Euler model worse. The second evaluation starts from the model’s own first sample with noise added, whereas the reference starts from the exact endpoint with the same noise, so the first step’s error carries into the second. In our reading, noise of standard deviation 2 masks much of that error and moves the input to levels where the table shows both models’ errors are smaller than at $T$, while noise of 0.5 masks little; the dependence on $\tau$ is consistent with the paper’s choice to search for the times model by model.