Generative Diffusion Models Explained (28): Understanding Consistency Models Step by Step
Continuing from where we left off: in Generative Diffusion Models Explained (27): Using Step Size as a Conditioning Input, we introduced the Shortcut model for accelerated sampling, one of whose comparison baselines was the "Consistency Models" approach. In fact, as far back as Generative Diffusion Models Explained (17): The General Procedure for Constructing an ODE (Part 2), when discussing ReFlow, some readers already brought up consistency models. But at the time I felt it looked more like a practical trick with somewhat thin theoretical grounding, so I didn't pay it much attention.
However, now that we've started following developments in accelerated sampling for diffusion models, consistency models are a piece of work we simply cannot skip. So, taking this opportunity, I'd like to share my own understanding of consistency models here.
A familiar recipe
Once again using a familiar recipe, our starting point is still ReFlow, since it's probably the simplest way to understand ODE-based diffusion. Let $\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)$ be a real sample from the target distribution, $\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)$ be random noise from the prior distribution, and $\boldsymbol{x}_t = (1-t)\boldsymbol{x}_0 + t\boldsymbol{x}_1$ be the noised sample. Then the training objective of ReFlow is: more
\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{t\sim U[0,1],\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\left[w(t)\Vert\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\Vert^2\right]\label{eq:loss}\end{equation}
where $w(t)$ is a tunable weight. Once training is complete, sampling can be carried out by solving $d\boldsymbol{x}_t/dt = \boldsymbol{v}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, t)$ to transform $\boldsymbol{x}_1$ into $\boldsymbol{x}_0$.
It should be pointed out that the noise schedule used in consistency models is $\boldsymbol{x}_t = \boldsymbol{x}_0 + t\boldsymbol{x}_1$ (when $t$ is large enough, $\boldsymbol{x}_t$ is likewise close to pure noise), which differs slightly from ReFlow. That said, the main purpose of this article is to try to gradually derive the same training idea and training objective as consistency models, step by step. I find ReFlow easier to understand, so I'll present things following ReFlow's framework, and readers are welcome to adjust the specific training details as needed.
Using $\boldsymbol{x}_t = (1-t)\boldsymbol{x}_0 + t\boldsymbol{x}_1$, we can eliminate $\boldsymbol{x}_1$ from the objective $\eqref{eq:loss}$:
\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{t\sim U[0,1],\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t)\Vert \underbrace{\boldsymbol{x}_t - t\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)}_{\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)} - \boldsymbol{x}_0\Vert^2\big]\label{eq:loss-2}\end{equation}
where $\tilde{w}(t) = w(t)/t^2$. Note that $\boldsymbol{x}_0$ is the real sample and $\boldsymbol{x}_t$ is the noised sample, so the ReFlow training objective is in fact also performing denoising. The model that predicts the clean sample is $\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)=\boldsymbol{x}_t - t\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$, and this function has an important property that always holds: $\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_0, 0)=\boldsymbol{x}_0$ — this is precisely one of the key constraints in consistency models.
Understanding it step by step
Next, let's deconstruct the ReFlow training process step by step, trying to find a better training objective along the way. First, we divide $[0,1]$ into $n$ equal parts, each of size $1/n$, and denote $t_k = k/n$; then $t$ only needs to be sampled uniformly from the finite set $\{0,t_1,t_2,\cdots,t_n\}$. Of course we could also choose a non-uniform discretization scheme, but these are non-essential details.
Since $t_0=0$ is trivial, we start from $t_1$, and the training objective for the first step is
\begin{equation}\boldsymbol{\theta}_1^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_1)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_1}, t_1) - \boldsymbol{x}_0\Vert^2\big]\end{equation}
Next, consider the training objective for the second step. If we followed the same pattern as $\eqref{eq:loss-2}$, it should be the expectation of $\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_2}, t_2) - \boldsymbol{x}_0\Vert^2$, but instead let's evaluate a new objective:
\begin{equation}\boldsymbol{\theta}_2^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_2)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_2}, t_2) - \boldsymbol{f}_{\boldsymbol{\theta}_1^*}(\boldsymbol{x}_{t_1}, t_1)\Vert^2\big]\end{equation}
That is, the prediction target is changed to $\boldsymbol{f}_{\boldsymbol{\theta}_1^*}(\boldsymbol{x}_{t_1}, t_1)$ instead of $\boldsymbol{x}_0$. Why make this change? Let's discuss it in terms of feasibility and necessity. On the feasibility side, $\boldsymbol{x}_{t_2}$ has more noise added compared to $\boldsymbol{x}_{t_1}$, so denoising it is harder — in other words, the degree of recovery achievable from $\boldsymbol{f}_{\boldsymbol{\theta}_2^*}(\boldsymbol{x}_{t_2}, t_2)$ is no better than that from $\boldsymbol{f}_{\boldsymbol{\theta}_1^*}(\boldsymbol{x}_{t_1}, t_1)$. So replacing $\boldsymbol{x}_0$ with $\boldsymbol{f}_{\boldsymbol{\theta}_1^*}(\boldsymbol{x}_{t_1}, t_1)$ as the training target for the second step is entirely feasible.
But even so, what's the necessity for making this switch? The answer is to reduce "trajectory crossing." Since $\boldsymbol{x}_{t_k} = (1-t_k)\boldsymbol{x}_0 + t_k\boldsymbol{x}_1$, as $k$ increases, the dependence of $\boldsymbol{x}_{t_k}$ on $\boldsymbol{x}_0$ becomes progressively weaker, to the point where two different values of $\boldsymbol{x}_0$ can yield very close values of $\boldsymbol{x}_{t_k}$. If we still used $\boldsymbol{x}_0$ as the prediction target in this case, we would run into the dilemma of "one input, multiple targets" — this is "trajectory crossing."
Faced with this dilemma, ReFlow's strategy is post-hoc distillation: once pretraining is finished, solving $d\boldsymbol{x}_t/dt = \boldsymbol{v}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, t)$ yields many pairs $(\boldsymbol{x}_0,\boldsymbol{x}_1)$, and using these paired $\boldsymbol{x}_0,\boldsymbol{x}_1$ to construct $\boldsymbol{x}_t$ avoids crossing. The idea behind consistency models is instead to change the prediction target to $\boldsymbol{f}_{\boldsymbol{\theta}_{k-1}^*}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})$, because for "the same $\boldsymbol{x}_1$, different $\boldsymbol{x}_0$," the discrepancy between different $\boldsymbol{f}_{\boldsymbol{\theta}_{k-1}^*}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})$ will be smaller than that between different $\boldsymbol{x}_0$, which likewise reduces the risk of crossing.
Put simply: predicting $\boldsymbol{f}_{\boldsymbol{\theta}_1^*}(\boldsymbol{x}_{t_1}, t_1)$ from $\boldsymbol{f}_{\boldsymbol{\theta}_2^*}(\boldsymbol{x}_{t_2}, t_2)$ is easier than predicting $\boldsymbol{x}_0$, and it achieves the desired effect all the same, so the prediction target is adjusted accordingly. Similarly, we can write
\begin{equation}\begin{gathered} \boldsymbol{\theta}_3^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_3)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_3}, t_3) - \boldsymbol{f}_{\boldsymbol{\theta}_2^*}(\boldsymbol{x}_{t_2}, t_2)\Vert^2\big] \\ \boldsymbol{\theta}_4^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_4)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_4}, t_4) - \boldsymbol{f}_{\boldsymbol{\theta}_3^*}(\boldsymbol{x}_{t_3}, t_3)\Vert^2\big] \\ \vdots \\[5pt] \boldsymbol{\theta}_n^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_n)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_n}, t_n) - \boldsymbol{f}_{\boldsymbol{\theta}_{n-1}^*}(\boldsymbol{x}_{t_{n-1}}, t_{n-1})\Vert^2\big] \end{gathered}\end{equation}
Consistency training
We have now finished deconstructing the ReFlow model and arrived at a new training objective that we believe is more reasonable — but at the cost of ending up with $n$ separate sets of parameters $\boldsymbol{\theta}_1^*,\boldsymbol{\theta}_2^*,\cdots,\boldsymbol{\theta}_n^*$, which is certainly not what we want; we only want a single model. So we posit that all the $\boldsymbol{\theta}_i^*$ can share the same set of parameters, giving us the training objective
\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{k\sim[n],\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_k)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_k}, t_k) - \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})\Vert^2\big]\label{eq:loss-3}\end{equation}
Here $k\sim[n]$ means that $k$ is sampled uniformly from $\{1,2,\cdots,n\}$. The problem with the expression above is that $\boldsymbol{\theta}^*$ is the very parameter we're trying to solve for, yet it also appears inside the objective function — which is clearly not sound (if we already knew $\boldsymbol{\theta}^*$, why would we still need to train?). So this objective must be modified to make it workable.
The meaning of $\boldsymbol{\theta}^*$ is the theoretically optimal solution. Considering that as training progresses, $\boldsymbol{\theta}$ gradually approaches $\boldsymbol{\theta}^*$, we can relax this condition in the objective function to a "look-ahead solution" — it just needs to be better than the current $\boldsymbol{\theta}$. How do we construct such a "look-ahead solution"? Consistency models' approach is to apply EMA (Exponential Moving Average) to the historical weights, which often gives a superior solution — a trick we used to rely on frequently back in the days of competing in machine learning contests.
So the final training objective of consistency models is:
\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{k\sim[n],\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_k)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_k}, t_k) - \boldsymbol{f}_{\bar{\boldsymbol{\theta}}}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})\Vert^2\big]\label{eq:loss-4}\end{equation}
where $\bar{\boldsymbol{\theta}}$ is the EMA of $\boldsymbol{\theta}$. This is exactly the "Consistency Training (CT)" from the original paper. In practice, we can also replace $\Vert\cdot - \cdot\Vert^2$ with a more general metric $d(\cdot, \cdot)$ to better fit the characteristics of the data.
Sampling analysis
Since we arrived here by "equivalently transforming" ReFlow step by step, one basic way to sample after training is still to solve the ODE just like in ReFlow:
\begin{equation}d\boldsymbol{x}_t/dt = \boldsymbol{v}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, t) = \frac{\boldsymbol{x}_t - \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, t)}{t}\label{eq:ode}\end{equation}
Of course, if after all that effort we end up with results identical to ReFlow, the whole exercise would have been pointless. Fortunately, models trained via consistency training have an important advantage: they support much larger sampling step sizes — even a step size of 1, which enables single-step generation:
\begin{equation}\boldsymbol{x}_0 = \boldsymbol{x}_1 - \boldsymbol{v}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1)\times 1 = \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1)\end{equation}
The reasoning is as follows:
\begin{equation}\begin{aligned} \Vert\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1) - \boldsymbol{x}_0\Vert =&\, \left\Vert\sum_{k=1}^n \Big[\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_k}, t_k) - \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})\Big]\right\Vert \\[5pt] \leq&\, \sum_{k=1}^n \Vert\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_k}, t_k) - \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})\Vert \\ \end{aligned}\label{eq:f-x1-x0}\end{equation}
We can see that consistency training is effectively optimizing an upper bound on $\Vert\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1) - \boldsymbol{x}_0\Vert$; when the loss is small enough, it means $\Vert\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1) - \boldsymbol{x}_0\Vert$ is also small enough, so one-step generation becomes possible.
But $\Vert\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1) - \boldsymbol{x}_0\Vert$ is precisely ReFlow's original training objective — why would directly optimizing it perform worse than optimizing its upper bound? This brings us back to the issue of "trajectory crossing." When training directly, the $\boldsymbol{x}_0,\boldsymbol{x}_1$ are randomly sampled without any one-to-one pairing, so it's impossible to directly train a one-step generative model this way. But when training the upper bound, the transitivity across multiple $\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_k}, t_k),\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})$ implicitly achieves the pairing of $\boldsymbol{x}_0,\boldsymbol{x}_1$.
If single-step generation quality isn't satisfactory, we can also increase the number of sampling steps to improve quality. There are two approaches here: (1) numerically solving $\eqref{eq:ode}$ with smaller step sizes; (2) converting it into an SDE-like stochastic iteration. The former is fairly conventional, so let's focus mainly on the latter.
First, note that in equation $\eqref{eq:f-x1-x0}$, replacing $\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1)$ with an arbitrary $\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, t)$ still yields a similar inequality relation, meaning that any $\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, t)$ is predicting $\boldsymbol{x}_0$. This means that, starting from $\boldsymbol{x}_1$, applying $\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1)$ gives us a preliminary $\boldsymbol{x}_0$, which might not be perfect. So we "mask" this imperfection by adding noise, obtaining $\boldsymbol{x}_{t_{n-1}}$, and then substituting into $\boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_{n-1}}, t_{n-1})$ gives a slightly better result, and so on:
\begin{equation}\begin{aligned} &\boldsymbol{x}_1\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I}) \\ &\boldsymbol{x}_0\leftarrow \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_1, 1) \\ &\text{for }k=n-1,n-2,\cdots,1: \\ &\qquad \boldsymbol{z} \sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I}) \\ &\qquad \boldsymbol{x}_{t_k} \leftarrow (1 - t_k)\boldsymbol{x}_0 + t_k\boldsymbol{z} \\ &\qquad \boldsymbol{x}_0\leftarrow \boldsymbol{f}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_{t_k}, t_k) \end{aligned}\end{equation}
For distillation
The training idea behind consistency models can equally be applied to distilling an existing diffusion model, producing what's called "Consistency Distillation (CD)." The method is to change the learning target for $\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_k}, t_k)$ in equation $\eqref{eq:loss-4}$ from $\boldsymbol{f}_{\bar{\boldsymbol{\theta}}}(\boldsymbol{x}_{t_{k-1}}, t_{k-1})$ to $\boldsymbol{f}_{\bar{\boldsymbol{\theta}}}(\hat{\boldsymbol{x}}_{t_{k-1}}^{\boldsymbol{\varphi}^*}, t_{k-1})$:
\begin{equation}\boldsymbol{\theta}^* = \mathop{\text{argmin}}_{\boldsymbol{\theta}} \mathbb{E}_{k\sim[n],\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\big[\tilde{w}(t_k)\Vert \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_k}, t_k) - \boldsymbol{f}_{\bar{\boldsymbol{\theta}}}(\hat{\boldsymbol{x}}_{t_{k-1}}^{\boldsymbol{\varphi}^*}, t_{k-1})\Vert^2\big]\label{eq:loss-5}\end{equation}
where $\hat{\boldsymbol{x}}_{t_{k-1}}^{\boldsymbol{\varphi}^*}$ is the $\boldsymbol{x}_{t_{k-1}}$ predicted by the teacher diffusion model starting from $\boldsymbol{x}_{t_k}$ as the initial value — for instance, with the simplest Euler solver, we have
\begin{equation}\hat{\boldsymbol{x}}_{t_{k-1}}^{\boldsymbol{\varphi}^*} \approx \boldsymbol{x}_{t_k} - (t_k - t_{k-1})\boldsymbol{v}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_{t_k}, t_k)\end{equation}
The reasoning for doing this is straightforward: if we already have a pretrained diffusion model, there's no need to hunt for a learning target along the straight line $\boldsymbol{x}_t = (1-t)\boldsymbol{x}_0 + t\boldsymbol{x}_1$, since that line is artificially defined and always carries some risk of crossing. Instead, we can let the pretrained diffusion model predict the trajectory. The learning target found this way might not be the "straightest" one, but it's guaranteed not to have crossings.
If cost is no object, we could also start from a randomly sampled $\boldsymbol{x}_1$, combine it with $\boldsymbol{x}_0$ solved by the pretrained diffusion model, and build the learning target from the paired $(\boldsymbol{x}_0,\boldsymbol{x}_1)$ — this is essentially ReFlow's distillation approach. Its downside is that it requires running the full sampling process on the teacher model, which is costly and slow. Consistency distillation, in contrast, only requires running the teacher model for a single step, making it much cheaper computationally.
That said, consistency distillation still requires real samples during the distillation process, which is also a drawback in certain scenarios. If you want a distillation process that neither runs the full teacher sampling process nor requires real data, one option is SiD, which we discussed previously — the cost being a considerably more complex derivation.
Summary
By progressively deconstructing and refining the ReFlow training procedure, this article has offered an intuitive path for gradually transitioning from ReFlow to Consistency Models.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.