Consistency Models

A diffusion model defines a probability-flow ODE whose trajectories carry each noise sample $x_T$ to a data sample as time runs from $T$ down to a small $\epsilon > 0$ (Song et al., 2021). A consistency function maps every point on such a trajectory to the trajectory's end point (Song et al., 2023):

Two properties follow:

• Self-consistency: $f(x_t, t) = f(x_{t'}, t')$ whenever $x_t$ and $x_{t'}$ lie on the same trajectory.
• Boundary condition: $f(x_\epsilon, \epsilon) = x_\epsilon$. Song et al. build this into the network as $f_\theta(x, t) = c_{\text{skip}}(t)\, x + c_{\text{out}}(t)\, F_\theta(x, t)$ with $c_{\text{skip}}(\epsilon) = 1$ and $c_{\text{out}}(\epsilon) = 0$.

Training enforces self-consistency between neighbouring times $t_n < t_{n+1}$ of a discretized trajectory:

where $d$ is a distance between images and $\theta^-$ is a copy of $\theta$ that receives no gradient (an exponential moving average in the original paper). In consistency distillation, $\hat{x}_{t_n}$ comes from one ODE-solver step of a pretrained diffusion model. In consistency training, both points are made from the same data sample and the same noise, so no pretrained model is needed.

Sampling takes one network evaluation, $x_\epsilon = f_\theta(x_T, T)$ with $x_T \sim \mathcal{N}(0, T^2 I)$, or alternates denoising and re-noising for a few more steps to trade compute for quality.