Introduction

Papers by Yang Song are always nerve-racking to read. In a deep-learning field already overflowing with equations, he makes even more spectacular use of mathematics and fills the appendices lavishly. I remember that “Score-Based Generative Modeling through Stochastic Differential Equations”, one of the foundational score-based diffusion papers reviewed just before this one, also had a formidable appendix. Anyone studying diffusion can struggle to decide where to begin. While reading presumed baselines such as DDPM and NCSN, there are moments when one wonders what foundation permits all these derivations. Since beginning to study diffusion, I have consulted countless foundational papers, blogs, videos, and equations throughout Bishop’s Pattern Recognition. What I learned over that time is that anything I skimmed without understanding never became truly mine, and later obstructed my understanding of subsequent papers. I also found that viewing the work as deriving a generative model from a mathematical model—rather than merely implementing something generative—made the equations easier to understand.

Rather than blindly following derivations, the most important thing is therefore to understand why diffusion can serve as a generative model in the first place, along with its limitations and their solutions.

Problems with diffusion models

Diffusion models still have not overcome slow generation. As image manipulation and large-language-model chatbots have recently become widely available services, users’ demand for speed and convenience has collided with providers’ need to minimize resources and costs. It is wonderful for AI to produce high-quality output, but what use is that if it takes forever? It feels like waiting after work for an express bus with a long headway.

Diffusion provides better modal stability than an implicit generative model such as a GAN, but guaranteeing it forces us to pay a trade-off in time and computation.

Designing faster sampling

Fast generation is easy to describe intuitively. Instead of diffusion’s artisanal process of removing noise one painstaking step at a time along time axis $t$, we want to take the express lane and generate a sample from noise in one shot. In other words, we want diffusion to support the fast $G(z)$ sampling available to GANs. Methods arising from this idea include

• DDIM: Define a non-Markovian forward process with the same marginal likelihood as a Markovian process and simplify the time steps in the sampling sequence.
• Diffusion-model distillation: Use a strong diffusion model such as DDPM to train a probability-flow ODE that can sample well in a single step.

Even with DDIM, however, accelerating probability-flow-ODE sampling inevitably lowers sample quality. The key is to minimize sampling steps without sacrificing performance, which conventional methods struggle to achieve. Distillation can transfer a good diffusion model’s performance into a probability-flow ODE, but dependence on diffusion-model generation creates a bottleneck in which training becomes remarkably slow.

Related work

The authors’ consistency-model outline does not depart from Yang Song’s earlier treatment of diffusion. Quoting the preceding paper, every diffusion process has an ODE with the same marginal likelihood $p(x_{0:T})$. For example, if the original diffusion SDE is

$$dx_t = \mu(x_t, t)dt + \sigma(t)dw_t$$

then an ODE with the same marginal likelihood is

$$dx_t = \left( \mu(x_t, t) - \frac{1}{2}\sigma(t)^2 \nabla_x \log p_t(x_t) \right) dt$$

The proof appears in the appendix of Yang Song’s SDE diffusion paper.

The ODE has an advantage over the SDE: it has no stochastic diffusion coefficient ($dw$). Given a starting point $x_0$ for the probability-flow ODE, every point $x_t$ through the endpoint $x_T$ lies on one trajectory solving the differential equation, as illustrated below. In an SDE, the drift term specifies only a direction while randomness controls the actual path. We know the start and end, but intermediate $x_t$ values cross and intertwine unpredictably, preventing a one-to-one mapping. In an ODE, by contrast, only the single time variable controls the trajectory. Given a model $s_\phi(x, t) \approx \nabla \log p_t(x)$ that predicts the score of data $x_t$ at a particular time, the equation takes the following form for perturbation kernel $p_t(x) = p_\text{data}(x) \otimes \mathcal{N}(0, t^2I)$:

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

This configuration follows “Elucidating the Design Space of Diffusion-Based Generative Models.” Papers design the drift and diffusion terms differently when defining a diffusion SDE. The equation above is the solution for $\mu = 0, \sigma = \sqrt{2t}$.

Thus, defining $x_T\sim \mathcal{N}(0,T^2I)$ and solving the corresponding probability-flow ODE $\frac{dx_t}{dt} = -ts_\phi (x_t, t)$ gives a single trajectory connecting $x_0$ and $x_T$. Numerical methods such as Euler and Heun solvers approximate the function’s shape to solve the differential equation.

This means predicting how the function changes under tiny changes in its variable when the differential equation has no analytic solution. As the figure shows, a numerical solution necessarily differs from the true solution, with variance growing as the time axis lengthens and sampling intervals widen.

$$\hat{x}_t,~t \in (0,~T)$$

To reduce numerical instability, the paper treats the solution at $t = \epsilon(0.002)$ as an approximation of the real data sample $x_0$ and uses $T = 80$ total time steps.

Slow sampling is ultimately diffusion’s greatest problem. When an ODE solver is used for sampling, the score prediction $s_\phi(x, t)$ again becomes the bottleneck. A numerical ODE solver can increase speed only by sacrificing quality or using distillation.

Despite these efforts, conventional ODE solvers cannot generate sufficiently high-quality data in one step. Distillation generally depends on a diffusion prior such as DDPM. Since DDPM must sample noisy data at each time step, it cannot eliminate the computational burden.

Progressive distillation is one attempt to solve this problem by gradually reducing the number of distilled time steps.

Learning a full trajectory from the outset requires correspondingly many samples for score prediction. Progressive distillation instead connects ODE score estimators trained over partial trajectories and reduces their number over time, achieving good performance without an enormous sample count from the beginning.

This paper does not use progressive distillation as shown above. It uses consistency distillation to make the ODE solver’s predictions consistent with prior predictions. Its main comparison is the progressive-distillation paper.

Consistency models

The paper opens by calling a consistency model a “new generative model.” It takes the probability-flow ODE derived from a diffusion SDE as its mathematical framework. If solving that ODE does not require distillation from a pretrained DDPM, it becomes the basis of a new generative model trainable from scratch. It may superficially resemble a normalizing flow, but it is fundamentally different because it starts from a diffusion stochastic differential equation.

Defining a consistency model

Consider the probability-flow ODE

$$dx_t = \left( \mu(x_t, t) - \frac{1}{2}\sigma(t)^2 \nabla_x \log p_t(x_t) \right) dt$$

and denote its solution trajectory by $\{x_t \}_{t \in [\epsilon, T]}$. A consistency function maps every point on this trajectory directly to $x_\epsilon$:

$$f : (x_t, t) \rightarrow x_\epsilon$$

The figure shows exactly this. The green line is the PF-ODE solution trajectory extending from $\epsilon$ to $T$ on the time axis. The main goal of the consistency model is to send every point on that axis back to the starting village.

Every point on the trajectory therefore has a consistent function output:

$$f(x, t) = f(x_{t^\prime},t^\prime),~\forall t,~t^\prime \in \{\tau \vert \epsilon \le \tau \le T\}$$

If the time argument is fixed—if a footprint remains on the time axis—the function is also invertible in the reverse direction.

A consistency model therefore predicts an ODE trajectory and sends the start of every remaining footprint to the same $x_\epsilon$.

Parameterization

The crucial requirement in deep learning is to parameterize the function so that it can be trained. Every consistency function has the following simple boundary constraint:

$$f(x_\epsilon, \epsilon) = x_\epsilon$$

In plain language, returning from the starting village ($x_\epsilon$) through $f(\cdot)$ must produce the starting village. It may seem obvious, but because the subject is a continuous function solving a differential equation, a correct constraint is essential. There are two ways to impose it. The first defines the function piecewise:

$$f_\theta(x, t) = \begin{cases}x,&t = \epsilon \\ F_\theta(x, t),& \epsilon <t<T \end{cases}$$

The second uses differentiable functions satisfying $c_\text{skip}(\epsilon) = 1$ and $c_\text{out}(\epsilon) = 0$ at the skip point $t = \epsilon$:

$$f_\theta(x, t) = c_\text{skip}(t)x + c_\text{out}(t) F_\theta(x, t)$$

The table comes from a related paper covering several consistency-like approaches. Most use the second method, as does this paper.

Sampling

Given a well-trained consistency model $f_\theta(\cdot, \cdot)$, sample from a known prior,

$$\hat{x}_T \sim \mathcal{N}(0, T^2I)$$

and pass the result directly through the function—the deep-learning model:

$$\hat{x}_\epsilon = f_\theta(\hat{x}_T, T)$$

This provides single-step generation. If for some reason we want multiple-step generation with a consistency model, as in a conventional diffusion model, we can simply repeat the process of returning to the starting village and adding noise again.

Zero-shot data editing

The defining property of a consistency model—converging to the original $x_0$ from any point on the trajectory corresponding to data under the prior—enables zero-shot image editing and manipulation. The simplest analogy is interpolation in a latent-variable model such as a GAN or VAE.

In a GAN or VAE, where the latent and generated sample are outputs of a parameterized implicit decoder, interpolating between latent $z_0$ that produces $x_0$ and latent $z_1$ that produces $x_1$ can generate an intermediate image $\text{Image}(x_0, x_1)$. This lets us manipulate feature vectors to modify generated images.

$$F_\Theta(\alpha \cdot z_0 + (1-\alpha) \cdot z_1) = \text{Image}(x_0, x_1)$$

Now consider the DDPM baseline, which directly follows the diffusion SDE. Its prior sample $x_T$ and generated sample $F_{\Theta_{1:T}}(x_T) = x_0$ do not form a one-to-one correspondence. Because a Markov process samples possible data-modality outputs $x_0^1, \cdots x_0^N$ for one latent sample $x_T$, latent interpolation does not yield meaningful image interpolation. Solving instead over a probability-flow-ODE solution, as a consistency model does, removes this one-to-many mapping from $x_T$ to the data modality. It inherits one of the advantages of GANs: easy image manipulation through latent manipulation.

The model can also be used zero-shot when a condition accompanies the sample modality. Diffusion-based models for high-quality inpainting, colorization, super-resolution, and similar tasks ordinarily require explicit task-specific training. A consistency model, however, learns to recover $x_\epsilon$ from any noise level and can therefore denoise across many levels.

Viewed differently, multiple-step generation from any input can find that input’s starting point. With a conditional input, we can presumably skip only the prior-sampling stage. If the input is a grayscale image, for example, we can apply the consistency model repeatedly—add noise, predict $x_0$, and repeat.

The fact that there is no need for separate conditional training makes the method naturally useful for inpainting, super-resolution, and SDEdit—painting to image—as shown below.

Training a consistency model

As the authors state, a consistency model can be trained either through distillation from a strong pretrained source model’s score or from scratch. The two methods are as follows.

Training through distillation

$$dx_t = \left( \mu(x_t, t) - \frac{1}{2}\sigma(t)^2 \nabla_x \log p_t(x_t) \right) dt,~~\frac{dx_t}{dt} = -ts_\phi (x_t, t)$$

Recall the PF-ODE equation above. The true data-distribution score in the left-hand equation is unknown; substituting a trained network’s score prediction gives the empirical PF-ODE on the right. Divide time interval $[\epsilon,~T]$ into $N-1$ subintervals. With boundary conditions $t_1 = \epsilon$ and $t_N = T$, define an increasing sequence $[t_1,~t_2,~\cdots,~t_N]$. The divisions could be arbitrary, but the paper follows a setting from related work:

$$t_i = \left(\epsilon^{1/\rho} + \frac{i-1}{N-1} (T^{1/\rho} - \epsilon^{1/\rho} )\right)^\rho,~\rho = 7$$

Since denser sampling brings a numerical ODE solver closer to the true solution, larger $N$ gives a more accurate prediction. Let $\hat{x}_{t_n}^\phi$ denote the solver’s predicted function value at a particular time:

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

This is the predicted next function value under the one-step ODE-solver update $\Phi(\cdots; \phi)$. Parameter $\phi$ appears because the score estimator solves an empirical PF-ODE using a pretrained score estimator. Applying the most common numerical solver, the Euler method, simplifies the equation to

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

The error introduced by converting the SDE to a PF-ODE corresponds to error in the true score-estimation function. To connect the two, we deliberately create a one-to-many mapping:

$$x \sim p_\text{data},~x = x+\eta \text{(Gaussian noise)}$$

Given data point $x$ defined this way, we can obtain adjacent points $(\hat{x}_{t_n}^\phi,~x_{t_{n+1}})$ on the PF ODE. Here $x_{t_{n+1}}$ is sampled from $\mathcal{N}(x,~t^2_{n+1}I)$ according to the SDE transition kernel, approximately as shown below.

The adjacent points then train the consistency network. The method is simple: pass both points—one sampled from the forward SDE, the other predicted from that sample by the score estimator and numerical ODE solver—through the network and make their outputs equal.

$$\mathcal{L}^N_{CD}(\theta, \theta^-;\phi) := \mathbb{E}(\lambda(t_n)d(f_\theta(x_{t_{n+1}}, t_{n+1}),~f_{\theta^-}(\hat{x}^\phi_{t_{n}},t_n)))$$

$\lambda(\cdot)$ weights the loss to account for changes in the kernel distribution over time, while $d$ is a distance between predictions. Student parameter $\theta$ receives gradient descent on the loss, and teacher parameter $\theta^-$ follows the student by EMA, just as in ordinary distillation. The distance can use many metrics; the paper experiments with MSE, L1, and LPIPS, which are common in image generation. It reports that simply fixing the weight $\lambda(\cdot)$ to $1$ performs well across tasks and datasets.

The figure above gives pseudocode for the training procedure. Assuming the numerical ODE has bounded error and consistency network $f$ satisfies a Lipschitz condition, one can prove that the supremum of the loss also converges. This demonstrates that the empirical PF ODE—the consistency model, or more precisely the consistency network—can converge with the numerical ODE along the true SDE trajectory from which it is distilled.

$$\begin{aligned} &\text{If local error uniformly bounded by }O((t_{n+1} - t_n)^{p+1}),\\ &\sup_{n, x} \parallel f_\theta(x, t_n) - f(x, t_n;\phi)\parallel_2 = O((\Delta t)^p) \end{aligned}$$

The proof appears in Appendix A.2 and follows straightforward induction, so it is omitted here.

The paper calls trainable $f_\theta$ the online network rather than student, and EMA-updated $f_{\theta^-}$ the target network rather than teacher. When consistency-distillation loss is trained over an unboundedly increasing number of time-step samples $N$, the target and online parameters can become equal. In other words, the distilled consistency network has inherited all information.

Training in isolation

The method above showed how score-network information and an ODE solver can make consistency loss converge. We now consider the training method that demonstrates the consistency model’s independence from conventional diffusion and establishes it as a new generative model built directly around the PF ODE.

Distillation requires a pretrained diffusion-process model and must use its score estimate $s_\phi(x, t)$ in the differential equation. Training a consistency model independently requires removing that dependency—finding $\nabla_x \log p_t(x_t)$ below:

$$dx_t = \left( \mu(x_t, t) - \frac{1}{2}\sigma(t)^2 \nabla_x \log p_t(x_t) \right) dt$$

It can be obtained without a score estimator as follows.

Projecting the desired score backward through the true data’s marginal distribution gives an integral:

$$\nabla \log p_t(x_t) = \nabla_{x_t} \log \int p_\text{data}(x) p(x_t \vert x) dx$$

The derivative of the logarithm has the closed form

$$\nabla \log p_t(x_t) = \frac{ \int p_\text{data}(x) \nabla_{x_t}p(x_t \vert x) dx}{\int p_\text{data}(x)p(x_t \vert x)dx}$$

The derivative of $p(x_t \vert x)$ can be expressed using the derivative of its log likelihood $\log (p(x_t \vert x))$:

$$\nabla \log p_t(x_t) = \frac{ \int p_\text{data}(x) p(x_t \vert x)\nabla_{x_t}\log p(x_t \vert x) dx}{\int p_\text{data}(x)p(x_t \vert x)dx}$$

Simplifying numerator and denominator gives

$$\nabla \log p_t(x_t) = \frac{ \int p_\text{data}(x) p(x_t \vert x)\nabla_{x_t}\log p(x_t \vert x) dx}{p_t(x_t)}$$

Because $x_t$ is independent of integration variable $x$, it can be treated as constant:

$$\nabla \log p_t(x_t) = \int\frac{ p_\text{data}(x) p(x_t \vert x)}{p_t(x_t)}\nabla_{x_t}\log p(x_t \vert x) dx$$

Bayes’ rule reverses the conditional in the leading factor:

$$\nabla \log p_t(x_t) = \int p(x \vert x_t) \nabla_{x_t} \log p(x_t \vert x) dx$$

This is the mean over $x$ under the distribution conditioned on $x_t$:

$$\nabla \log p_t(x_t) = \mathbb{E}(\nabla_{x_t} \log p(x_t \vert x) \vert x_t)$$

Since the diffusion process defines the conditional distribution with a Gaussian kernel, this can be approximated as

$$-\mathbb{E}\left(\frac{x_t - x}{t^2} \vert x_t\right)$$

Numerical error is inevitable because the expectation is computed over available samples. Nevertheless, this score permits sampling without a pretrained score-estimation network, and those samples can train the consistency network.

The derivation was long, but the procedure is simple: generate both adjacent samples from a predefined diffusion SDE and train the consistency model with them. For multivariate standard Gaussian $z \sim \mathcal{N}(0, I)$, the consistency-training loss becomes

$$\mathcal{L}^N_{CT}(\theta, \theta^-) := \mathbb{E}(\lambda(t_n)d(f_\theta(x + t_{n+1}z, t_{n+1}),~f_{\theta^-}(x + t_nz,t_n)))$$

One can likewise prove that minimizing this loss is ultimately equivalent to minimizing the distillation loss. This proof is genuinely important, but it extracts an $o(\Delta t)$ term through a lengthy Taylor expansion, so it too is omitted here. See Theorem 2 in the appendix.

Experiments

The experiments proceed intuitively, testing each factor in turn and selecting the best.

Panel (a) shows that LPIPS loss is the most effective distance metric $d(\cdot)$, so the next experiment, (b), fixes LPIPS as the metric. Panel (b) studies the solver; Heun, which uses a second-order approximation, performs better than first-order Euler.

Panel (c) tests a denser range of time-step sample counts $N$ to reduce bias. As expected—and as shown by the derivation—performance improves as $N$ grows. It also saturates beyond a certain $N$, perhaps reflecting the numerical ODE’s performance bottleneck.

Panel (d) covers training with CT. Its much poorer FID is an unavoidable limitation. Unlike CD, CT does not depend on the performance of a particular numerical ODE solver, because it samples for training from a predefined kernel, so it needs no solver. Since CT uses no distillation, the effect of $N$ is pronounced: a small $N$ converges quickly but samples poorly; a large $N$ converges more slowly but samples better. To combine both advantages, the paper gradually increases $N$ during training, as shown in purple, and raises EMA factor $\mu$ accordingly. The graph shows both rapid convergence and high sampling quality by FID.

Conclusion

The experimental results include few-step image generation, direct generation, and zero-shot image editing. By exploiting convergence of the empirical PF ODE, the consistency network is among the first approaches to train a network capable of sampling in one shot. It may open a new direction for fast sampling that conventional diffusion could not provide. If the method can combine with zero-shot text-to-image generation such as Stable Diffusion while preserving high sample quality, it may become a new business game changer.