Truncated Consistency Models

## Abstract

Abstract Consistency models have recently been introduced to accelerate sampling from diffusion models by directly predicting the solution (i.e., data) of the probability flow ODE (PF ODE) from initial noise.
However, the training of consistency models requires learning to map all intermediate points along PF ODE trajectories to their corresponding endpoints. This task is much more challenging than the ultimate objective of one-step generation, which only concerns the PF ODE’s noise-to-data mapping.
We empirically find that this training paradigm limits the one-step generation performance of consistency models.
To address this issue, we generalize consistency training to the truncated time range, which allows the model to ignore denoising tasks at earlier time steps and focus its capacity on generation.
We propose a new parameterization of the consistency function and a two-stage training procedure that prevents the truncated-time training from collapsing to a trivial solution.
Experiments on CIFAR-10 and ImageNet 64 × 64 64 64 64\times 64 64 × 64 datasets show that our method achieves better one-step and two-step FIDs than the state-of-the-art consistency models such as iCT-deep,
using more than 2 × \times × smaller networks. Project page: https://truncated-cm.github.io/

## 1 Introduction

Diffusion models (Ho et al., ; Song et al., ) have demonstrated remarkable capabilities in generating high-quality continuous data such as images, videos, or audio (Ramesh et al., ; Ho et al., ; Huang et al., ).
Their generation process gradually transforms a simple Gaussian prior into data distribution through a probability flow ordinary differential equation (PF ODE). Although diffusion models can capture complex data distributions, they require longer generation time due to the iterative nature of solving the PF ODE.

Consistency models (Song et al., ) were recently proposed to expedite the generation speed of diffusion models by learning to directly predict the solution of the PF ODE from the initial noise in a single step. To circumvent the need for simulating a large number of noise-data pairs to learn this mapping, as employed in prior works (Liu et al., ; Luhman & Luhman, ), consistency models learn to minimize the discrepancy between the model’s outputs at two neighboring points along the ODE trajectory. The boundary condition at $t=0$ serves as an anchor, grounding these outputs to the real data. Through simulation-free training, the model gradually refines its mapping at different times, propagating the boundary condition from $t=0$ to the initial $t=T$.

However, the advantage of simulation-free training comes with trade-offs. Consistency models must learn to map any point along the PF ODE trajectory to its corresponding data endpoint, as shown in Fig. 1a. This requires the learning of both *denoising* at smaller times on the PF ODE, where the data are only partially corrupted, and *generation* towards $t=T$, where most of the original data information has been erased. This dual task necessitates larger network capacity, and it is challenging for a single model to excel at both tasks. Our empirical observations in Fig. 2 demonstrate the model would gradually sacrifice its denoising capability at smaller times to trade for generation quality as training proceeds. While this behavior is desirable as the end goal is generation rather than denoising, we argue for explicit control over this trade-off, rather than allowing the model to allocate capacity uncontrollably across times. This raises a key question: Can we explicitly reduce the network capacity dedicated to the denoising task in order to improve generation?

In this paper, we propose a new training algorithm, termed Truncated Consistency Models (TCM), to de-emphasize denoising at smaller times while still preserving the consistency mapping for larger times. TCM relaxes the original consistency objective, which requires learning across the entire time range $[0,T]$ of PF ODE trajectories, to a new objective that focuses on a truncated time range $[t^{\prime},T]$, where $t^{\prime}$ serves as the dividing time between denoising and generation tasks. This allows the model to dedicate its capacity primarily to generation, freeing it from the denoising task at earlier times $[0,t^{\prime})$. Crucially, we show that a proper boundary condition at $t^{\prime}$ is necessary to ensure the new model adheres to the original consistent mapping. To achieve this, we propose a two-stage training procedure (see Fig. 1a): The first stage involves pretraining a standard consistency model over the whole time range. This pretrained model then acts as the boundary condition at $t^{\prime}$ for the subsequent truncated consistency training stage of the TCM.

Experimentally, TCM improves both the sample quality and the training stability of consistency models across
different datasets and sampling steps.
On CIFAR-10 and ImageNet $64\times 64$ datasets, TCM outperforms the iCT (Song & Dhariwal, ), the previous best consistency model, in both one-step and two-step generation using similar network size.
TCM even outperforms iCT-deep that uses a $2\times$ larger network across datasets and sampling steps. By using our largest network, we achieve a one-step FID of $2.20$ on ImageNet $64\times 64$, which is competitive with the current state-of-the-art. In addition, the divergence observed during standard consistency training is not present in TCM. We show through extensive ablation experiments why the various design choices of truncated consistency models (including the strength of mandating boundary conditions, two-stage training, etc.) are necessary to obtain these results.

Contributions. (i) We identify an underlying trade-off between denoising and generation within consistency models, which negatively impacts both stability and generation performance.
(ii) Building on these insights, we introduce Truncated Consistency Models, a novel two-stage training framework that explicitly allocates network capacity towards generation while preserving consistency mapping. (iii) Extensive validation of TCM demonstrates significant improvements in both one-step and two-step generation, achieving state-of-the-art results within the consistency models family on multiple image datasets. Additionally, TCM exhibit improved training stability. (iv) We provide in-depth analyses, along with ablation and design choices that demonstrate the unique advantages of the two-stage training in TCM.

## 2 Preliminaries

Figure: Figure 1: (a) Two-stage training of TCM. In Stage 1, a standard consistency model is trained to provide both the boundary condition and initialization for TCM training in Stage 2. TCM focuses on learning in the $[t^{\prime},T]$ range, discarding denoising tasks at earlier times and allocating network capacity toward generation-like tasks at later times. (b) Sample quality (FID, lower is better) of the two training stages. TCM (Stage 2) improves over standard consistency training (Stage 1) across datasets. Additionally, standard consistency training shows instability on challenging datasets like ImageNet 64x64, where the model could diverge during training.
Refer to caption: x1.png

Figure: Figure 2: Evolution of the denoising FID (dFIDt) during standard consistency training for different $t$, where $0<t\leq 80$ follows the EDM noise schedule (Karras et al., ). The model gradually sacrifices its denoising capability at smaller times ($t<1.0$) to trade for the improved generation quality at $t=80$ as training proceeds.
Refer to caption: x2.png

### 2.1 Diffusion models

Diffusion models are a class of generative models that synthesize data by reversing a forward process in which the data distribution $p_{\text{data}}$ is gradually transformed into a tractable Gaussian distribution.
In this paper, we use the formulation proposed in Karras et al. (), where the forward process is defined by the following stochastic differential equation (SDE):

$$ $\displaystyle d\mathbf{x}_{t}=\sqrt{2t}d\mathbf{w}_{t},$ (1) $$

where $t\in[0,T]$ and $\mathbf{w}_{t}$ is the standard Brownian motion from $t=0$ to $t=T$.
Here, we define $p_{t}$ as the marginal distribution of $\mathbf{x}_{t}$ along the forward process, where $p_{0}=p_{\text{data}}$.
In this case, $p_{t}$ is a perturbed data distribution with the noise from $\mathcal{N}(0,t^{2}{\mathbf{I}})$.
In diffusion models, $T$ is set to be large enough so that $p_{T}$ is approximately equal to a tractable Gaussian distribution $\mathcal{N}(0,T^{2}{\mathbf{I}})$.

Diffusion models come with the reverse probability flow ODE (PF ODE) that starts from $t=T$ to $t=0$ and yields the same marginal distribution $p_{t}$ as the forward process in Eq. (1) (Song et al., ):

$$ $\displaystyle d\mathbf{x}_{t}=-t\mathbf{s}_{t}(\mathbf{x}_{t})dt,$ (2) $$

where $\mathbf{s}_{t}(\mathbf{x}_{t}):=\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ is the score function at time $t\in[0,T]$. To draw samples from the data distribution $p_{\text{data}}$, we first train a neural network to learn $\mathbf{s}_{t}(\mathbf{x})$ using the denoising score matching (Vincent, ), initialize $\mathbf{x}_{T}$ with a sample from $\mathcal{N}(0,T^{2}\mathbf{I})$,
and solve the PF ODE backward in time: $\mathbf{x}_{0}=\mathbf{x}_{T}+\int_{T}^{0}(-t\mathbf{s}_{t}(\mathbf{x}_{t}))dt$.
However, numerically solving the PF ODE requires multiple forward passes of the neural score function estimator, which is computationally expensive.

### 2.2 Consistency models

Consistency models instead aim to directly map from noise to data, by learning a consistency function that outputs the solution of PF ODE starting from any $t\in[0,T]$. The desired consistency function $\mathbf{f}$ should satisfy the following two properties (Song et al., ): (i) $\mathbf{f}(\mathbf{x}_{0},0)=\mathbf{x}_{0}$, and (ii) $\mathbf{f}(\mathbf{x}_{t},t)=\mathbf{f}(\mathbf{x}_{s},s),\;\forall(s,t)\in[0,
T]^{2}$.
The first condition can be satisfied by the reparameterization

$$ $\displaystyle\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x},t):=c_{\text{out}}( t)\mathbf{F}_{{\boldsymbol{\theta}}}(\mathbf{x},t)+c_{\text{skip}}(t)\mathbf{x},$ (3) $$

where ${\boldsymbol{\theta}}$ is the parameter of the free-form neural network $\mathbf{F}_{{\boldsymbol{\theta}}}:\mathbb{R}^{d}\times\mathbb{R}\to\mathbb{R}
^{d}$, and $c_{\text{out}}(0)=0$, $c_{\text{skip}}(0)=1$ following the similar design of Karras et al. ().
Here, instead of training ${\mathbf{f}}_{\boldsymbol{\theta}}$ directly, we train a surrogate neural network $\mathbf{F}_{\boldsymbol{\theta}}$ under the above reparameterization.
The second condition can be learned by optimizing the following consistency training objective:

$$ $\displaystyle\mathcal{L}_{\text{CT}}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{ f}_{\boldsymbol{\theta}}^{-}):=\mathbb{E}_{t\sim\psi_{t},\mathbf{x}\sim p_{ \text{data}},{\boldsymbol{\epsilon}}\sim\mathcal{N}(0,{\mathbf{I}})}[\frac{ \omega(t)}{\Delta_{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{ \boldsymbol{\epsilon}},t),\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t -\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})],$ (4) $$

where ${\boldsymbol{\theta}}^{-}=\text{stopgrad}({\boldsymbol{\theta}})$, $\psi_{t}$ denotes the probability of sampling time $t$ that also represents the noise scale, ${\boldsymbol{\epsilon}}$ denotes the standard Gaussian noise, $\omega(t)$ is a weighting function, $d(\cdot,\cdot)$ is a distance function defined in Sec. B.1.2, and $\Delta_{t}$ represents the nonnegative difference between two consecutive time steps that is usually set to a monotonically increasing function of $t$.

The gradient of $\mathcal{L}_{\text{CT}}$ with respect to ${\boldsymbol{\theta}}$ is an approximation of the underlying consistency distillation loss with a $O(\max_{t}\Delta_{t})$ error (See Appendix B).
Song et al. () empirically suggests that $\Delta_{t}$ should be large at the beginning of training, which incurs biased gradients but allows for stable training, and should be annealed in the later stages, which reduces the error term but increases variance.

##### Denoising FID

By definition, consistency models can both generate data from pure Gaussian noise as well as noisy data sampled from $p_{t}$ where $0<t<T$.
To understand how consistency models propagate end solutions through diffusion time, we need to empirically measure their denoising capability across different time steps. To this end, we define denoising FID at time step $t$, termed dFIDt, as the Fréchet inception distance (FID) (Heusel et al., ) between the original data $p_{\text{data}}$ and the denoised data by consistency models with inputs sampled from $p_{t}$.
When computing dFIDt, we first add Gaussian noise from $\mathcal{N}(0,t^{2}{\mathbf{I}})$ to 50K clean samples
and then denoise them using consistency models.
Hence, dFID0 is close to zero, and dFIDT is the standard FID.

## 3 Truncated consistency model

Standard consistency models pose a higher challenge in training than many other generative models:
instead of simply mapping noise to data, consistency models must learn the mapping from *any* point along the PF ODE trajectory to its data endpoint.
Hence, a consistency model must divide its capacity between denoising tasks (i.e., mapping samples from intermediate times to data) and generation (i.e., mapping from pure noise to data).
This challenge mainly contributes to consistency models’ underperformance relative to other generative models with similar network capacities (see Table 1).

Interestingly, standard consistency models navigate the trade-off between denoising and generation tasks implicitly.
We observe that during standard consistency training, the model gradually loses its denoising capabilities at the low $t$. Specifically,
Fig. 2 shows a trade-off in which, after some training iterations, denoising FIDs at lower $t$ ($t<1$) *increase* while the denoising FIDs at larger $t$ ($t>1$) (including the generation FID at the largest $t=80$) continue to decrease.
This suggests that the model struggles to learn to denoise and generate simultaneously, and sacrifices one for the other.

Truncated consistency models (TCM) aim to explicitly control this tradeoff by forcing the consistency training to ignore the denoising task for small values of $t$, thus improving its capacity usage for generation.
We thus *generalize* the consistency model objective in Eq. (4) and apply it only in the truncated time range $[t^{\prime},T]$ where the dividing time $t^{\prime}$ lies within $(0,T)$. The time probability $\psi_{t}$ in TCM only has support in $[t^{\prime},T]$ as a result.

##### Naive solution

A straightforward approach is to directly train a consistency model on the truncated time range.
However, the model outputs can collapse to an arbitrary constant because a constant function (i.e., $\mathbf{f}_{\boldsymbol{\theta}}(\mathbf{x},t)=\text{const}$) is a minimizer of the consistency training objective (Eq. (4)).
In standard consistency models, the boundary condition $\mathbf{f}(\mathbf{x}_{0},0)=\mathbf{x}_{0}$ prevents collapse, but in this naive example, there is no such meaningful boundary condition.
For example, if the free-form neural network $\mathbf{F}_{{\boldsymbol{\theta}}}(\mathbf{x},t)=-c_{\text{skip}}(t)\mathbf{x}
/c_{\text{out}}(t)$ for all $t\in[t^{\prime},T]$, $\mathbf{f}_{\boldsymbol{\theta}}(\mathbf{x},t)$ is $\mathbf{0}$, and thus Eq. (4) becomes zero.
To handle this, we propose a two-stage training procedure and design a new parameterization with a proper boundary condition, as outlined below.

##### Proposed Solution

Truncated consistency models conduct training in two stages:

- 1.
Stage 1 (Standard consistency training): We pretrain a consistency model to convergence in the usual fashion, with the training objective in Eq. (4); we denote the pre-trained model as $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$.
- 2.
Stage 2 (Truncated consistency training): We initialize a new consistency model $\mathbf{f}_{\boldsymbol{\theta}}$ with the first-stage pretrained weights $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$, and train over a truncated time range $[t^{\prime},T]$. The boundary condition at time $t^{\prime}$ is provided by the pretrained $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$. This stage is explained further below.

To explain the details of TCM, we first introduce the following parameterization:

$$ $\displaystyle\mathbf{f}_{{\boldsymbol{\theta}},{\boldsymbol{\theta}}_{0}^{-}}^ {\text{trunc}}(\mathbf{x},t)={\mathbf{f}}_{{\boldsymbol{\theta}}}(\mathbf{x},t )\cdot\mathds{1}\{t\geq t^{\prime}\}+\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-} }(\mathbf{x},t)\cdot\mathds{1}\{t<t^{\prime}\},$ (5) $$

where $\mathds{1}\{\cdot\}$ is the indicator function, and similarly, ${\boldsymbol{\theta}}_{0}^{-}=\text{stopgrad}({\boldsymbol{\theta}}_{0})$.
Intuitively, we only use our final model $\mathbf{f}_{{\boldsymbol{\theta}}}$ when $t\geq t^{\prime}$, and we inquire the pre-trained $\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}$ otherwise.
This approach ensures that (1) $\mathbf{f}_{{\boldsymbol{\theta}}}$ does not waste its capacity learning in the $[0,t^{\prime})$ range, and (2) if $\mathbf{f}_{{\boldsymbol{\theta}}}$ is trained well, it will learn to generate data by mimicking the pre-trained model $\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}$ at the boundary.
When $t^{\prime}=0$, we recover the standard consistency model parameterization Eq. (3).
During sampling, as $\mathbf{f}_{{\boldsymbol{\theta}},{\boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}
}=\mathbf{f}_{{\boldsymbol{\theta}}}$ for all $t\in[t^{\prime},T]$, we can discard this parameterization and just use $\mathbf{f}_{{\boldsymbol{\theta}}}$ for generating samples.

To describe the boundary condition, we then partition the support of the time sampling distribution $\psi_{t}$, i.e., $[t^{\prime},T]$ into two time ranges: (i) the boundary time region $S_{t^{\prime}}:=\{t\in\mathbb{R}:t^{\prime}\leq t\leq t^{\prime}+\Delta_{t}\}$, and (ii) the consistency training time region $S_{t^{\prime}}^{-}\triangleq[t^{\prime},T]\setminus S_{t^{\prime}}=\{t\in
\mathbb{R}:t^{\prime}+\Delta_{t}<t\leq T\}$.
To effectively enforce the boundary condition using the first-stage pre-trained model $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$, a nonnegligible amount of $t$’s, sampled from $\psi_{t}$, must fall within the interval $S_{t^{\prime}}$. Otherwise the consecutive time steps $t$ and $t-\Delta_{t}$ in consistency training will mostly be larger or equal to $t^{\prime}$, limiting the influence of the pre-trained model.

With this time partitioning and our new parameterization, Eq. (4) can be decomposed as follows:

$$ $\displaystyle\begin{split}\mathcal{L}_{\text{CT}}(\mathbf{f}_{{\boldsymbol{ \theta}},{\boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}},\mathbf{f}_{{ \boldsymbol{\theta}}^{-},{\boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}})&= \underbrace{\int_{t\in S_{t^{\prime}}}\psi_{t}(t)\frac{\omega(t)}{\Delta_{t}}d (\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t), \mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+(t-\Delta_{t}){ \boldsymbol{\epsilon}},t-\Delta_{t})dt}_{\text{Boundary loss}}\\ &+\underbrace{\int_{s\in S_{t^{\prime}}^{-}}\psi_{t}(t)\frac{\omega(t)}{\Delta _{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}}, t),\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t-\Delta_{t}){ \boldsymbol{\epsilon}},t-\Delta_{t})dt}_{\text{Consistency loss}},\end{split}$ (6) $$

where we apply our parameterization in Eq. (5) in the above two time partitions separately, and we drop the expectation over $\mathbf{x}\sim p_{\text{data}},{\boldsymbol{\epsilon}}\sim\mathcal{N}(0,{
\mathbf{I}})$ for notation simplicity.
Unlike standard consistency training, TCM have two terms: the boundary loss and consistency loss.
The boundary loss allows the model to learn from the pre-trained model, preventing collapse to a constant.

Training on the objective (6) can still collapse to a constant if we do not utilize the boundary condition sufficiently by not sampling enough time $t$’s in $S_{t^{\prime}}$.
In particular, this can happen for $\Delta_{t}$ close to zero when consistency training is near convergence (Song & Dhariwal, ; Geng et al., ).
To prevent this, we design $\psi_{t}$ to satisfy $\int_{t\in S_{t^{\prime}}}\psi_{t}(t)dt>0$.
In other words, we have a strictly positive probability of sampling a point in $S_{t^{\prime}}$, even when $\Delta_{t}$ is close to zero.
Specifically, we define $\psi_{t}$ as a mixture of the Dirac delta function $\delta(\cdot)$ at point $t^{\prime}$ and another distribution $\bar{\psi}_{t}$:

$$ $\displaystyle\psi_{t}(t)=\lambda_{b}\delta(t-t^{\prime})+(1-\lambda_{b})\bar{ \psi}_{t}(t),$ (7) $$

where the weighting coefficient $\lambda_{b}\in(0,1)$.
$\bar{\psi}_{t}$ has the support $(t^{\prime},T]$ and can be instantiated in different ways (e.g., log-normal or log-Student-$t$ distributions); the effect of different $\bar{\psi}_{t}$ choices is explored in Section 4.4.

By definition, we can see that $\int_{t\in S_{t^{\prime}}}\psi_{t}(t)dt\geq\lambda_{b}$, and $\lambda_{b}$ controls how significantly we emphasize the boundary condition.
Assume that the first-stage consistency model is perfectly trained in $[0,t^{\prime}]$, i.e., $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}(\mathbf{x}_{t},t)=\mathbf{x}_{0}$ for all $t\in[0,t^{\prime}]$.
If $\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t^{\prime}},t^{\prime})\neq
\mathbf{f}_{{\boldsymbol{\theta}}_{0}}(\mathbf{x}_{t^{\prime}},t^{\prime})$,
$\mathbf{f}_{{\boldsymbol{\theta}}}$ will be penalized by the boundary loss.
Minimizing the boundary loss enforces the boundary condition in second-stage model $\mathbf{f}_{{\boldsymbol{\theta}}}$ (i.e., $\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t^{\prime}},t^{\prime})=\mathbf
{f}_{{\boldsymbol{\theta}}_{0}}(\mathbf{x}_{t^{\prime}},t^{\prime})=\mathbf{x}
_{0}$), while minimizing the consistency loss propagates the boundary condition to the end time (i.e., $\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{T},T)=\mathbf{f}_{{\boldsymbol{
\theta}}}(\mathbf{x}_{t^{\prime}},t^{\prime})$). Consequently, the loss in Eq. (6) effectively guides the model towards the desired solution $\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{T},T)=\mathbf{x}_{0}$.
With the time distribution $\psi_{t}$ defined in Eq. (7), our training objective becomes

$$ $\displaystyle\mathcal{L}_{\text{CT}}(\mathbf{f}_{{\boldsymbol{\theta}},{ \boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}},\mathbf{f}_{{\boldsymbol{\theta}} ^{-},{\boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}})\approx{\lambda}_{b} \underbrace{\frac{\omega(t^{\prime})}{\Delta_{t^{\prime}}}d(\mathbf{f}_{{ \boldsymbol{\theta}}}(\mathbf{x}+t^{\prime}{\boldsymbol{\epsilon}},t^{\prime}) ,\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{ \prime}}){\boldsymbol{\epsilon}},t^{\prime}-\Delta_{t^{\prime}}))}_{\text{ Boundary loss}:=\mathcal{L}_{B}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{f}_{{ \boldsymbol{\theta}}_{0}^{-}})}$ (8) $\displaystyle+(1-\lambda_{b})\underbrace{\mathbb{E}_{\bar{\psi}_{t}}[\frac{ \omega(t)}{\Delta_{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{ \boldsymbol{\epsilon}},t),\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t -\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})]}_{\text{Consistency loss}:= \mathcal{L}_{C}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{f}_{{\boldsymbol{ \theta}}^{-}})}.$ (9) $$

where
the approximation in Eq. (8) holds when $\Delta_{t}$ is sufficiently small (which is true for the truncated training stage).
Please see Appendix C for the detailed derivation.
For simplicity of notation,
we relax the above objective by absorbing the $(1-\lambda_{b})$ factor into ${\lambda}_{b}$ and express our final training loss as:

$$ $\displaystyle\mathcal{L}_{\text{TCM}}:=w_{b}\mathcal{L}_{B}(\mathbf{f}_{ \boldsymbol{\theta}},\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}})+\mathcal{L}_{ C}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{f}_{{\boldsymbol{\theta}}^{-}}),$ (10) $$

where $w_{b}=\lambda_{b}/(1-\lambda_{b})$ is a tunable hyperparameter that
controls the weighting of the boundary loss. To estimate the two losses, we partition each mini-batch of size $B$ into two subsets. The boundary loss $L_{B}$ is estimated using $N_{B}$ = $\lfloor B\rho\rfloor$ samples, where $\rho\in(0,1)$ is a hyperparameter controlling the allocation of samples. The consistency loss $L_{C}$ is estimated with the remaining $B-N_{B}$ samples. Increasing $\rho$ reduces the variance of the boundary loss gradient estimator but increases the variance of the consistency loss gradient estimator, and vice versa. The final mini-batch loss is as follows:

$$ $\displaystyle\mathcal{L}_{\text{TCM}}$ $\displaystyle\approx\frac{w_{b}}{N_{B}}\sum_{i=1}^{N_{B}}\nabla_{{\boldsymbol{ \theta}}}({\mathcal{L}_{B}})_{i}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{f}_{ {\boldsymbol{\theta}}_{0}^{-}})+\frac{1}{B-N_{B}}\sum_{j=N_{B}+1}^{B}\nabla_{{ \boldsymbol{\theta}}}({\mathcal{L}_{C}})_{j}(\mathbf{f}_{\boldsymbol{\theta}}, \mathbf{f}_{{\boldsymbol{\theta}}^{-}}),$ (11) $$

where $({\mathcal{L}_{B}})_{i}$ and $({\mathcal{L}_{C}})_{j}$ are the boundary loss and the consistency loss at the $i$-th sample from $\delta(t-t^{\prime})$ and the $j$-th sample from $\bar{\psi}_{t}$, respectively.
We provide the training algorithm in Algorithm 1.

Figure: Algorithm 1 Truncated Consistency Training

## 4 Experiments

In this section, we evaluate TCM on standard image generation benchmarks and compare it against state-of-the-art generative models. We begin by detailing the experimental setup in Sec. 4.1. We then study the behavior of denoising FID and its impact on generation FID in Sec. 4.2. We benchmark TCM against a variety of existing methods in Sec. 4.3, and provide detailed analysis on various design choices in Sec. 4.4.

### 4.1 Setup

We evaluate TCM on the CIFAR-10 (Krizhevsky et al., ) and ImageNet $64{\times}64$ (Deng et al., ) datasets.
We consider the unconditional generation task on CIFAR-10 and class-conditional generation on ImageNet $64{\times}64$.
We measure sample quality with Fréchet Inception Distance (FID) (Heusel et al., ) (lower is better), as is standard in the literature.

For consistency training in TCM, we mostly follow the hyperparameters in ECT (Geng et al., ), including the discretization curriculum and continuous-time training schedule. For all experiments, we choose a dividing time $t^{\prime}=1$ and set $\bar{\psi}_{t}$ to the log-Student-$t$ distribution. We use $w_{b}=0.1$ and $\rho=0.25$ for the boundary loss. We discuss these choices in Sec. 4.4.
In line with Geng et al. (), we initialize the model with the pre-trained EDM (Karras et al., ) / EDM2 (Karras et al., ) for CIFAR-10 / ImageNet $64\times 64$, respectively. On CIFAR-10, we use a batch size of 512 and 1024 for the first and the second stage, respectively.
On ImageNet with EDM2-S architecture, we use a batch size of 2048 and 1024 for the first and the second stage, respectively.
For EDM2-XL, to save compute, we initialize the truncated training stage with the pre-trained checkpoint from the ECM work (Geng et al., ) that performs the standard consistency training, and conduct the second-stage training with a batch size of 1024. Please see Appendix D for more training details.

### 4.2 Truncated training allocates capacity toward generation

Figure: Figure 3: Denoising FID (dFID) for continuation of standard consistency training at later iterations (Stage 1) and TCM model (Stage 2) at various $t$s on CIFAR-10 during the course of training. For TCM, we set the dividing time $t^{\prime}=1$. We can see, in the second stage, the dFID exhibits a dramatic increase at times below the dividing time $t^{\prime}$, while the dFID at times above $t^{\prime}$ and FID at $t=T$ continue to improve. Notably, the rate of dFID in the truncated stage increase at earlier times is significantly faster compared to standard consistency training, suggesting a more efficient “forgetting” of the denoising tasks.
Refer to caption: x3.png

**Table 1: FID, NFE and # param. on CIFAR-10. Bold indicates the best result for each category and NFE.**
| Method | NFE | FID | # param. (M) |
| --- | --- | --- | --- |
| Diffusion models |  |  |  |
| EDM (Karras et al., ) | 35 | 1.97 | 55.7 |
| PFGM++ (Xu et al., ) | 35 | 1.91 | 55.7 |
| DDPM (Ho et al., ) | 1000 | 3.17 | 35.7 |
| LSGM (Vahdat et al., ) | 147 | 2.10 | 475 |
| Consistency models |  |  |  |
| 1-step |  |  |  |
| iCT (Song & Dhariwal, ) | 1 | 2.83 | 56.4 |
| iCT-deep (Song & Dhariwal, ) | 1 | 2.51 | 112 |
| CTM (Kim et al., ) (w/o GAN) | 1 | 5.19 | 55.7 |
| ECM (Geng et al., ) | 1 | 3.60 | 55.7 |
| TCM (ours) | 1 | 2.46 | 55.7 |
| 2-step |  |  |  |
| iCT (Song & Dhariwal, ) | 2 | 2.46 | 56.4 |
| iCT-deep (Song & Dhariwal, ) | 2 | 2.24 | 112 |
| ECM (Geng et al., ) | 2 | 2.11 | 55.7 |
| TCM (ours) | 2 | 2.05 | 55.7 |
| Variational score distillation |  |  |  |
| DMD (Yin et al., ) | 1 | 3.77 | 55.7 |
| Diff-Instruct (Luo et al., ) | 1 | 4.53 | 55.7 |
| SiD (Zhou et al., ) | 1 | 1.92 | 55.7 |
| Knowledge distillation |  |  |  |
| KD (Luhman & Luhman, ) | 1 | 9.36 | 35.7 |
| DSNO (Zheng et al., ) | 1 | 3.78 | 65.8 |
| TRACT (Berthelot et al., ) | 1 | 3.78 | 55.7 |
|  | 2 | 3.32 | 55.7 |
| PD (Salimans & Ho, ) | 1 | 9.12 | 60.0 |
|  | 2 | 4.51 | 60.0 |

Our proposed TCM aims to explictly reallocate network capacity towards generation by de-emphasizing denoising tasks at smaller $t$’s. Empirical analysis in Fig. 3 further characterizes this behavior, showing a rapid increase in dFIDs at smaller $t$’s below the threshold $t^{\prime}$ during the truncated training stage. Conversely, dFIDs continue to decrease at larger $t$’s. In addition, TCM exhibit a more pronounced “forgetting” of the denoising task compared to consistency training (Fig. 2) at earlier times. For instance, dFID at $t=0.2$ increases up to $3.5$ in the truncated training, whereas it remains below $1$ in the standard consistency training. TCM also significantly accelerate the process of forgetting the denoising tasks at these earlier times, achieving a substantially improved generation FID. This suggests that by explicitly controlling the training time range, the neural network can effectively shift its capacity towards generation.

Fig. 1(b) demonstrates how this reallocation of network capacity directly translates to improved sample quality and training stability. For CIFAR-10 / ImageNet $64\times 64$, the truncated training stage (Stage 2) is initialized from the Stage 1 model at 250K / 150K iterations, respectively.
We can see that the truncated training improves FID over the consistency training on the two datasets.
Moreover, we find that the truncated training is more stable than the original consistency training, as their ImageNet FID blows up after 150K iterations, while TCM continues to improve FID from $2.83$ to $2.46$, showcasing its robustness (See Figure 6 for more analysis).

**Table 2: FID, NFE and # param. on ImageNet $64\times 64$. Dotted lines separate results by # param. Bold indicates the best result for each category and NFE.**
| Method | NFE | FID | # param. (M) |
| --- | --- | --- | --- |
| Diffusion models |  |  |  |
| EDM2-S (Karras et al., ) | 63 | 1.58 | 280 |
| EDM2-XL (Karras et al., ) | 63 | 1.33 | 1119 |
| Consistency models |  |  |  |
| 1-step |  |  |  |
| iCT (Song & Dhariwal, ) | 1 | 4.02 | 296 |
| iCT-deep (Song & Dhariwal, ) | 1 | 3.25 | 592 |
| ECM (Geng et al., ) (EDM2-S) | 1 | 4.05 | 280 |
| TCM (ours; EDM2-S) | 1 | 2.88 | 280 |
| \cdashline1-4<br>MultiStep-CD (Heek et al., ) | 1 | 3.20 | 1200 |
| ECM (Geng et al., ) (EDM2-XL) | 1 | 2.49 | 1119 |
| TCM (ours; EDM2-XL) | 1 | 2.20 | 1119 |
| 2-step |  |  |  |
| iCT (Song & Dhariwal, ) | 2 | 3.20 | 296 |
| iCT-deep (Song & Dhariwal, ) | 2 | 2.77 | 592 |
| ECM (Geng et al., ) (EDM2-S) | 2 | 2.79 | 280 |
| TCM (ours; EDM2-S) | 2 | 2.31 | 280 |
| \cdashline1-4<br>MultiStep-CD (Heek et al., ) | 2 | 1.90 | 1200 |
| ECM (Geng et al., ) (EDM2-XL) | 2 | 1.67 | 1119 |
| TCM (ours ; EDM2-XL) | 2 | 1.62 | 1119 |
| Variational score distillation |  |  |  |
| DMD2 w/o GAN (Yin et al., ) | 1 | 2.60 | 296 |
| Diff-Instruct (Luo et al., ) | 1 | 5.57 | 296 |
| EMD-16 (Xie et al., ) | 1 | 2.20 | 296 |
| Moment Matching (Salimans et al., ) | 1 | 3.00 | 400 |
|  | 2 | 3.86 | 400 |
| SiD (Zhou et al., ) | 1 | 1.52 | 296 |
| Knowledge distillation |  |  |  |
| DSNO (Zheng et al., ) | 1 | 7.83 | 329 |
| TRACT (Berthelot et al., ) | 1 | 7.43 | 296 |
|  | 2 | 4.97 | 296 |
| PD (Salimans & Ho, ) | 1 | 15.4 | 296 |
|  | 2 | 8.95 | 296 |

### 4.3 TCM improves the sample quality of consistency models

To demonstrate the effectiveness of TCM, we compare our method with three lines of works that distill diffusion models to one or two steps: (i) consistency models (Song & Dhariwal, ; Kim et al., ; Geng et al., ) that distills the PF ODE mapping in a simulation-free manner; (ii) variational score distillation (Yin et al., ; Luo et al., ; Zhou et al., ) that performs distributional matching by utilizing the score of pre-trained diffusion models; (iii) knowledge distillation (Luhman & Luhman, ; Zheng et al., ; Berthelot et al., ; Salimans & Ho, ) that distill the PF ODE through off-line or on-line simulation using the pre-trained diffusion models. We exclude the methods that additionally use the GAN loss, which causes more training difficulties, for fair comparison.

Results.
In Table 1 and Table 2, we report the sample quality measured by FID and the sampling speed measured by the number of function evaluations (NFE), on CIFAR-10 and ImageNet-64$\times$64, respectively. We also include the number of model parameters. Our main findings are: (1) TCM significantly outperforms improved Consistency Training (iCT) (Song & Dhariwal, ), the state-of-the-art consistency model, across datasets, number of steps and network sizes. For example, TCM improves the one-step FID from $2.83$ / $4.02$ in iCT to $2.46$ / $2.88$, on CIFAR-10 / ImageNet.
Further, TCM’s one-step FID even rivals iCT’s two-step FID on both datasets. When using EDM2-S model, TCM also surpasses iCT-deep, which uses *$2\times$* deeper networks, in both one-step ($2.88$ vs $3.25$) and two-step FIDs ($2.31$ vs $2.77$) on ImageNet.
(2) TCM beats all the knowledge distillation methods and performs competitively to variational score distillation methods. Note that TCM do not need to train additional neural networks as in VSD methods, or to run simulation as in knowledge distillation methods.
(3) Two-step TCM performs comparably to the multi-step EDM (Karras et al., ), the state-of-the-art diffusion model. For example, when both using the same EDM network, two-step TCM obtains a FID of $2.05$ on CIFAR-10, which is close to $1.97$ in EDM with 35 sampling steps. We further provide the uncurated one-step and two-step generated samples in Fig. 5. Please see Appendix E for more samples.

### 4.4 Analyses of design choices

Time sampling distribution $\bar{\psi}_{t}$. We explore various time sampling distributions $\bar{\psi}_{t}$ supported on $[t^{\prime},T]$, and find that the truncated log-Student-$t$ distribution works best (i.e., $\ln(t)$ follows Student-$t$ distribution). The Student-$t$ distribution, being heavier-tailed than the Gaussian distribution employed in previous consistency training (Song & Dhariwal, ; Geng et al., ), inherently allocates more probability mass towards larger $t$’s. This aligns with the motivation of TCM, which emphasizes enhancing generation capabilities at later times. The degree of freedom $\nu$ effectively controls the thickness of the tail, with the Student-$t$ distribution converging to a Gaussian distribution as $\nu\to\infty$. Figure 4(a) shows the shape of $\bar{\psi}_{t}$ with varying standard deviation $\sigma$ and the degree of freedom $\nu$ in three cases:
(1) heavy-tailed and a low probability mass around small $t$’s ($\sigma=2,\nu=10000)$,
(2) heavy-tailed and a high probability mass around small $t$’s ($\sigma=0.2,\nu=0.01)$,
(2) light-tailed ($\sigma=0.2,\nu=2)$.
From Fig. 4(b), we observe that the log-Student-$t$ distribution with $\sigma=0.2,\nu=0.01$ is the best among the three.
Hence we use $\sigma=0.2,\nu=0.01$ in all the experiments.

Figure: (a)
Refer to caption: x4.png

**Table 3: CIFAR-10 FID when varying the dividing time $t^{\prime}$.**
| $t^{\prime}$ value | 0.17 | 0.8 | 1.0 | 1.5 |
| --- | --- | --- | --- | --- |
| FID | 2.70 | 2.69 | 2.56 | 2.79 |

Strength for boundary loss. Figure 4(c) shows the effect of two key hyper-parameters that control the strength of imposing boundary condition in the TCM objective (Eq. 10).
We observe that the FID is relatively stable with a wide range of $\rho$ and $w_{b}$ (from $0.1$ to $0.5$ for $\rho$ and from $0.1$ to $1$ for $w_{b}$).
However, when using a very small weight for the boundary loss ($w_{b}=0.01$), FID explodes as the model fails to maintain the boundary condition.
Thus, we use $\rho=0.25,w_{b}=0.1$ in all the experiments.

Dividing time $t^{\prime}$.
The boundary $t^{\prime}$ ideally represents the point where the task in the PF ODE transitions from denoising to generation. However, this transition is gradual, and there is no single definitive point. Fig. 2b suggests this transition occurs roughly between $t^{\prime}=0.8$ and $t^{\prime}=1.5$, where we observe a change in dFID behavior: it primarily deteriorates during training before this range, but then stabilizes afterwards (more indicative of a generation task). Based on this analysis, we experimented with multiple $t^{\prime}$ values around this range. Table 4 shows that $t^{\prime}=1$ provides the best results among the choices. Note that here we use a batch size of 128, while it is 1024 in our default setting.

Are two stages enough?
A natural question is whether we can extend our two-stage training procedures to three or more stages by gradually increasing $t^{\prime}$.
However, recall that our methodology was motivated by the fact that in the first-stage training (standard consistency training), we observe increasing dFIDs at smaller $t$ values of the training range, as seen in Fig. 2. This trade-off is notably absent in the second stage over the time range $[t^{\prime},T]$, as seen in Fig. 3. This suggests that during the second stage, the model tackles tasks that are more or less similar to generation, and introducing another truncated training stage may not yield further gains.
In support of this hypothesis, we implement the third stage where the dividing time is $t^{\prime}=4$, but do not observe improvement, as shown in Table 4.
We also consider adding an intermediate training stage between stage 1 and stage 2 that finetunes $\mathbf{f}_{\theta_{0}}$ in the time range $(0,t^{\prime})$ but it produces a slightly worse performance, which we discuss in Appendix A.

## 5 Related work

Consistency models.
Song et al. () first proposed consistency models as a new class of generative models that synthesize samples with a single network evaluation. Later, Song & Dhariwal (); Geng et al. () presented a set of improved techniques to train consistency models for better sample quality. Luo et al. () introduced latent consistency models (LCM) to accelerate the sampling of latent diffusion models.
Kim et al. () proposed consistency trajectory models (CTM) that generalize consistency models by enabling the prediction between any two intermediate points on the same PF ODE trajectory.
The training objective in CTM becomes more challenging than standard consistency models that only care about the mapping from intermediate points to the data endpoints.
Heek et al. () proposed multistep consistency models that divide the PF ODE trajectory into multiple segments to simplify the consistency training objective. They train the consistency models in each segment separately, and need multiple steps to generate a sample.
Ren et al. () combined CTM with progressive distillation (Salimans & Ho, ), by performing segment-wise consistency distillation where the number of ODE trajectory segments progressively reduces to one. Similar to CTM, it relies on the adversarial loss (Goodfellow et al., ) to achieve good performance.

Figure: Figure 5: Uncurated one-step (top) and two-step (bottom) generated samples from TCM (EDM) on CIFAR-10 and TCM (EDM2-XL) on ImageNet 64$\times$64, respectively.
Refer to caption: x7.png

Fast sampling of diffusion models. While a line of work aims to accelerate diffusion models via fast numerical solvers for the PF-ODE (Lu et al., ; Karras et al., ; Liu et al., ; Xu et al., ), they usually still require more than $10$ steps. To achieve low-step or even one-step generation, besides consistency models, other training-based methods have been proposed from three main perspectives: (i) *Knowledge distillation*, which
first used the pre-trained diffusion model to generate a dataset of noise and image pairs, and then applied it to train a single-step generator (Luhman & Luhman, ; Zheng et al., ).
Progressive distillation (Salimans & Ho, ; Meng et al., ) iteratively halves the number of sampling steps required, without needing an offline dataset.
(ii) *Variational score distillation*, which aims to match the distribution of the student and teacher output via an approximate (reverse) KL divergence (Yin et al., ; ; Xie et al., ), implicit score matching (Zhou et al., ) or moment matching (Salimans et al., ). (iii) *Adversarial distillation*, which leverages the adversarial training to fine-tune pre-trained diffusion models into a few-step generator (Sauer et al., ; ; Lin et al., ; Xu et al., ). Compared with these training-based diffusion acceleration methods, our method is most memory and computation efficient.

## 6 Conclusion

We have introduced a truncated consistency training method that significantly enhances the sample quality of consistency models.
To generalize consistency models to the truncated time range, we have proposed a new parameterization of the consistency function and a two-stage training process that explicitly allocates network capacity towards generation.
We also discussed about our design choices arising from the new training paradigm.
Our approach achieves superior performance compared to state-of-the-art consistency models, as evidenced by improved one-step and two-step FID scores across different datasets and network sizes.
Notably, these improvements are achieved while utilizing similar or even smaller network architectures than baselines.

Limitation.
TCM introduces a slight increase in computational cost due to the additional boundary loss in Eq. (10), beyond the standard consistency training loss. Standard consistency training necessitates two forward passes per training iteration, while our parameterization (Eq. 5) requires three. Also, our method incurs a minor additional memory cost as we need to maintain a pre-trained consistency model (in evaluation mode) for the boundary loss. We observe that on ImageNet 64$\times$64 with EDM2-S, TCMs have an 18% increase in training time and an 15% increase in memory cost.
Moreover, TCM still has a considerate performance gap from the diffusion models with a large NFE, but we believe this is an important step towards closing the gap.

## 7 Reproducibility statement

We provide sufficient details for reproducing our method in the main paper and also in the Appendix. D,
including a pseudo code of the training algorithm, model initialization and architecture, model parameterization, learning rate schedules, time step sampling procedures, and other training details. We also specify hyperparameter choices like the dividing time $t^{\prime}$, boundary loss weight $w_{b}$, and boundary ratio $\rho$. Additionally, we discuss the computational costs of our method compared to standard consistency training.
For evaluation, we describe our sampling procedure for both one-step and two-step generation.

## 8 Ethics statement

This paper raises similar ethical concerns to other papers on deep generative models. Namely, such models can be (and have been) used to generate harmful content, such as disinformation and violent imagery. We advocate for the responsible deployment of such models in practice, including guardrails to reduce the risk of producing harmful content. The design of these protections is orthogonal to our work.
Other ethical concerns may arise regarding the significant resource costs required to train and use deep generative models, including energy and water usage. This work increases the training cost of consistency models, but it also enables the models to be run with only 1 NFE and requires smaller neural network architectures, both of may which reduce inference-time costs relative to other diffusion-based models. Nonetheless, the environmental impact of training and deploying deep generative models remains an important limitation.

## Acknowledgments

We thank Zhengyang Geng for helpful feedback on reproducing ECM. This work was supported in part by the National Science Foundation
through RINGS
grant 2148359. GF also acknowledges the generous support
of the Sloan Foundation, Intel, Bosch, and Cisco.

## Appendix A Additional experiments

Fig. 6 shows that the gradient spikes during the first stage training while the second stage training is relatively smooth.
We hypothesize that the truncated training is more stable because it is less affected by the biased gradient norms across different $t$.

Figure: Figure 6: Gradient norm evolution during the first and second stage training on ImageNet $64\times 64$ (corresponding to Fig. 1(b)). The red circles indicate where the gradient norms are larger than 100. Stage 1 training blows up after the last few gradient spikes. It shows that the truncated consistency training is more stable than the standard consistency training.
Refer to caption: x8.png

Fig. 7 shows the dFIDt evolution during the standard consistency training.
We see that dFIDs at larger $t$’s start from larger values and converges more slowly.

Figure: Figure 7: Evolution of the denoising FIDs (dFIDt) at different times $t$’s during standard consistency training for different iterations. For $t\in(1,10)$, dFIDt has different convergence speeds while in both small times ($t<1$ ) and large ’s ($t>10$), dFIDt converges with a more similar speed.
Refer to caption: x9.png

##### Adding an intermedate training stage

In our parameterization Eq. (5), we only use the pre-trained model $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$ in $[0,t^{\prime})$.
One may wonder if we can fine-tune $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$ on the truncated time range $[0,t^{\prime})$ to provide a better boundary condition for the truncated training.
We find that although doing so improved the dFID${}_{t^{\prime}}$ of $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$ from 1.51 to 1.43, it led to a worse final FID of $>$2.7 for the truncated consistency model, regardless of whether we initialized $\mathbf{f}_{\boldsymbol{\theta}}$ with the pre-trained model or the fine-tuned model. In contrast, our proposed method achieved an FID of 2.61 with the same hyperparameters.
We hypothesize that fine-tuning the pre-trained model on the truncated time range $[0,t^{\prime})$
makes the model $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$ forget about how the early mappings properly propagate to the later mappings in the range of $[t^{\prime},T]$.
This may hinder
the learnability of its mapping at the boundary time, making it harder for $\mathbf{f}_{{\boldsymbol{\theta}}}$ to transfer the knowledge learned in $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$ to its generation capability.

## Appendix B Background on consistency models

Most of this part has been introduced by previous works (Song et al., ; Song & Dhariwal, ). Here, we introduce the background of consistency models, in particular the relationship between consistency training and consistency distillation, for completeness.

### B.1 Definition of Consistency Function

#### B.1.1 Probability flow ODE

The probability flow ODE (PF ODE) of Karras et al. () is as follows:

$$ $\displaystyle d\mathbf{x}_{t}=-t\mathbf{s}_{t}(\mathbf{x}_{t})dt,$ (12) $$

where $\mathbf{s}_{t}(\mathbf{x}_{t})$ is the score function at time $t\in[0,T]$. To draw samples from the data distribution $p_{\text{data}}$, we initialize $\mathbf{x}_{T}$ with a sample from $\mathcal{N}(0,T^{2}\mathbf{I})$
and solve the PF ODE backward in time. The solution $\mathbf{x}_{0}=\mathbf{x}_{T}+\int_{T}^{0}(-t\mathbf{s}_{t}(\mathbf{x}_{t}))dt$ is distributed according to $p_{\text{data}}$.

#### B.1.2 Consistency function

Integrating the PF ODE using numerical solvers is computationally expensive. Consistency function instead directly outputs the solution of the PF ODE starting from any $t\in[0,T]$. The consistency function $\mathbf{f}$ satisfies the following two properties:

- 1.
$\mathbf{f}(\mathbf{x}_{0},0)=\mathbf{x}_{0}$.
- 2.
$\mathbf{f}(\mathbf{x}_{t},t)=\mathbf{f}(\mathbf{x}_{s},s)\ \forall(s,t)\in[0,T
]^{2}$.

The first condition can be trivially satisfied by setting $\mathbf{f}(\mathbf{x},t)=c_{\text{out}}(t)\mathbf{F}(\mathbf{x},t)+c_{\text{
skip}}(t)\mathbf{x}$ where $c_{\text{out}}(0)=0$ and $c_{\text{skip}}(0)=1$ following EDM (Karras et al., ).
The second condition can be satisfied by optimizing the following objective:

$$ $\displaystyle\min_{\mathbf{f}}\mathbb{E}_{s,t,\mathbf{x}_{t}}[d(\mathbf{f}( \mathbf{x}_{t},t),\mathbf{f}(\mathbf{x}_{s},s))],$ (13) $$

where $d$ is a function satisfying:

- 1.
$d(\mathbf{x},\mathbf{y})=0\iff\mathbf{x}=\mathbf{y}$.
- 2.
$d(\mathbf{x},\mathbf{y})\geq 0$.
- 3.
$\frac{\partial d(\mathbf{x},\mathbf{y})}{\partial\mathbf{y}}|_{\mathbf{y}=
\mathbf{x}}=0$
- 4.
$\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}$ and $\frac{\partial d}{\partial\mathbf{y}^{2}}$ are well-defined and bounded.

### B.2 Consistency distillation

#### B.2.1 Objective

In practice, Song et al. () consider the following objective instead:

$$ $\displaystyle\min_{{\boldsymbol{\theta}}}\ \mathbb{E}_{t,\mathbf{x}_{t}}[d( \mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t),\mathbf{f}_{{\boldsymbol{ \theta}}^{-}}(\mathbf{x}_{t-\Delta_{t}},t-\Delta_{t}))],$ (14) $$

where we parameterize the consistency function $\mathbf{f}$ with a neural network $\mathbf{f}_{{\boldsymbol{\theta}}}$, and $0<\Delta_{t}<t$. Here, $\mathbf{f}_{{\boldsymbol{\theta}}^{-}}$ is the identical network with stop gradients applied and is called teacher.
Since $\Delta_{t}>0$, the teacher always receives the less noisy input, and the student $\mathbf{f}_{{\boldsymbol{\theta}}}$ is trained to mimic the teacher.
Optimizing Eq. (14) requires computing $\mathbf{x}_{t-\Delta_{t}}$, which we can be approximated using one step of Euler’s solver:

$$ $\displaystyle\mathbf{x}_{t-\Delta_{t}}=\mathbf{x}_{t}+\int_{t}^{t-\Delta_{t}}( -u\mathbf{s}_{u}(\mathbf{x}_{u}))du\approx\mathbf{x}_{t}+t\mathbf{s}_{t}( \mathbf{x}_{t})\Delta_{t}.$ (15) $$

When $\mathbf{s}_{t}(\mathbf{x}_{t})$ is approximated by a pre-trained score network, Eq. (14) becomes the consistency distillation objective in Song et al. ().
If $\Delta_{t}$ is sufficieintly small, the approximation in Eq. (15) is quite accurate, making $\mathcal{L}_{\text{CD}}$ a good approximation of Eq. (14).
The precision of the approximation depends on $\Delta_{t}$ and also the trajectory curvature of the PF ODE.

#### B.2.2 Gradient when Δ t → 0 → subscript Δ 𝑡 0 \Delta_{t}\to 0 roman_Δ start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT → 0

Let us rewrite Eq. (14) as follows:

$$ $\displaystyle\mathbb{E}_{t,\mathbf{x}_{t}}[d(\mathbf{f}_{{\boldsymbol{\theta}} }(\mathbf{x}_{t},t),\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t- \Delta_{t}},t-\Delta_{t}))]$ (16) $\displaystyle=\mathbb{E}_{t,\mathbf{x}_{t}}[d(\mathbf{f}_{{\boldsymbol{\theta} }}(\mathbf{x}_{t},t),\underbrace{\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x} _{t},t)}_{\mathbf{y}}+\underbrace{\mathbf{f}_{{\boldsymbol{\theta}}^{-}}( \mathbf{x}_{t-\Delta_{t}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}}( \mathbf{x}_{t},t)}_{\Delta\mathbf{y}})]$ (17) $\displaystyle=\mathbb{E}_{t,\mathbf{x}_{t}}[d(\mathbf{f}_{{\boldsymbol{\theta} }}(\mathbf{x}_{t},t),\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t))+ \frac{\partial d}{\partial\mathbf{y}}\Delta\mathbf{y}+\frac{1}{2}(\Delta \mathbf{y})^{T}\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}\Delta\mathbf{y}+O( ||\Delta\mathbf{y}||^{3})]$ (18) $\displaystyle=\frac{1}{2}\mathbb{E}_{t,\mathbf{x}_{t}}[(\Delta\mathbf{y})^{T} \frac{\partial^{2}d}{\partial\mathbf{y}^{2}}\Delta\mathbf{y}+O(||\Delta\mathbf {y}||^{3})]$ (19) $$

, where we define $\Delta\mathbf{y}=\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t-\Delta_{
t}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t)$.

Let’s take the derivative with respect to ${\boldsymbol{\theta}}$:

$$ $\displaystyle\frac{1}{2}\frac{\partial}{\partial{\boldsymbol{\theta}}}\mathbb{ E}_{t,\mathbf{x}_{t}}[(\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t- \Delta_{t}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t) )^{T}\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}(\mathbf{f}_{{\boldsymbol{ \theta}}^{-}}(\mathbf{x}_{t-\Delta_{t}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol {\theta}}}(\mathbf{x}_{t},t))+O(||\Delta\mathbf{y}||^{3})]$ (20) $\displaystyle=\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t-\Delta_{t }},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t))\frac{ \partial\mathbf{f}_{{\boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}+O(| |\Delta\mathbf{y}||^{3})].$ (21) $$

As

$$ $\displaystyle\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t-\Delta_{t}}, t-\Delta_{t})=\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t},t)-\frac{ \partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial\mathbf{x}_{t}}\frac{ \partial\mathbf{x}_{t}}{\partial t}\Delta_{t}-\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}^{-}}}{\partial t}\Delta_{t}+O(\Delta_{t}^{2}),$ (22) $$

we have

$$ $\displaystyle\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t-\Delta_{t}}, t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t},t)=-(\frac {\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial\mathbf{x}_{t}}\frac{ \partial\mathbf{x}_{t}}{\partial t}+\frac{\partial\mathbf{f}_{{\boldsymbol{ \theta}}^{-}}}{\partial t})\Delta_{t}+O(\Delta_{t}^{2}).$ (23) $$

Since $\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t)$ has the same value as $\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t},t)$, we can plug this into Eq. (21):

$$ $\displaystyle\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}_{t-\Delta_{t }},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t))\frac{ \partial\mathbf{f}_{{\boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}+O(| |\Delta\mathbf{y}||^{3})]$ (24) $\displaystyle=-\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}\frac{\partial\mathbf{x}_{t}}{\partial t}+\frac{\partial\mathbf {f}_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}\Delta_{t}+O(\Delta_{t}^{ 2})+O(||\Delta\mathbf{y}||^{3})]$ (25) $\displaystyle=-\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}\frac{\partial\mathbf{x}_{t}}{\partial t}+\frac{\partial\mathbf {f}_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}\Delta_{t}+O(\Delta_{t}^{ 2})].$ (26) $$

As the gradient is $O(\Delta_{t})$, it becomes zero when $\Delta_{t}\to 0$, so it cannot be used for training.
To make the gradient non-zero, Song et al. () divide the by $\Delta_{t}$. Then, we have

$$ $\displaystyle\frac{\partial}{\partial{\boldsymbol{\theta}}}\mathcal{L}_{\text{ CD}}({\boldsymbol{\theta}},{\boldsymbol{\theta}}^{-})=\frac{\partial}{\partial {\boldsymbol{\theta}}}\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{1}{\Delta_{t}}d( \mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}_{t},t),\mathbf{f}_{{\boldsymbol{ \theta}}^{-}}(\mathbf{x}_{t-\Delta_{t}},t-\Delta_{t}))]$ (27) $\displaystyle=-\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}\frac{\partial\mathbf{x}_{t}}{\partial t}+\frac{\partial\mathbf {f}_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}+O(\Delta_{t})]$ (28) $\displaystyle=-\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}\frac{\partial\mathbf{x}_{t}}{\partial t}+\frac{\partial\mathbf {f}_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}]$ (29) $\displaystyle=-\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}(-t\cdot\mathbf{s}_{t}(\mathbf{x}_{t}))+\frac{\partial\mathbf{f }_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}]$ (30) $\displaystyle=\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}(t\cdot\mathbf{s}_{t}(\mathbf{x}_{t}))-\frac{\partial\mathbf{f} _{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}]$ (31) $$

in the limit of $\Delta_{t}\to 0$.

##### Hessian of d 𝑑 d italic_d .

Here, we provide the Hessians of the L2 squared loss and the Pseudo-Huber loss.

- 1.
If $d(\mathbf{x},\mathbf{y})=||\mathbf{x}-\mathbf{y}||_{2}^{2}$, $\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}|_{\mathbf{y}=\mathbf{x}}=2{
\mathbf{I}}$.
- 2.
If $d(\mathbf{x},\mathbf{y})=\sqrt{||\mathbf{x}-\mathbf{y}||_{2}^{2}+c^{2}}-c$, $\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}|_{\mathbf{y}=\mathbf{x}}=\frac{1}
{c}{\mathbf{I}}$.

### B.3 Consistency training

Song et al. () show that Eq. (31) can be estimated without a pre-trained score network. From Tweedie’s formula, we express the score function as

$$ $\displaystyle\mathbf{s}_{t}(\mathbf{x}_{t})=\frac{\mathbb{E}_{p(\mathbf{x}| \mathbf{x}_{t})}[\mathbf{x}]-\mathbf{x}_{t}}{t^{2}}.$ (32) $$

Plugging this into Eq. (31), we have

$$ $\displaystyle\mathbb{E}_{t,\mathbf{x}_{t}}[\frac{\partial^{2}d}{\partial \mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial \mathbf{x}_{t}}(t\cdot\mathbf{s}_{t}(\mathbf{x}_{t}))-\frac{\partial\mathbf{f} _{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}]=\mathbb{E}_{t,\mathbf{x }_{t}}[\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{ {\boldsymbol{\theta}}^{-}}}{\partial\mathbf{x}_{t}}\frac{\mathbb{E}_{p(\mathbf {x}|\mathbf{x}_{t})-\mathbf{x}_{t}}[\mathbf{x}]}{t}-\frac{\partial\mathbf{f}_{ {\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{\boldsymbol {\theta}}}}{\partial{\boldsymbol{\theta}}}]$ (33) $\displaystyle=\mathbb{E}_{t,\mathbf{x}_{t}}[\mathbb{E}_{p(\mathbf{x}|\mathbf{x }_{t})}[\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}(\frac{\partial\mathbf{f}_ {{\boldsymbol{\theta}}^{-}}}{\partial\mathbf{x}_{t}}\frac{\mathbf{x}-\mathbf{x }_{t}}{t}-\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial t}) \frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta} }}]]$ (34) $\displaystyle=\mathbb{E}_{t,\mathbf{x},\mathbf{x}_{t}}[\frac{\partial^{2}d}{ \partial\mathbf{y}^{2}}(\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{ \partial\mathbf{x}_{t}}\frac{\mathbf{x}-\mathbf{x}_{t}}{t}-\frac{\partial \mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{ {\boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}],$ (35) $$

where we now have the expectation over three random variables $t,\mathbf{x},\mathbf{x}_{t}$ and do not require a score function.
In the next section, we will reverse-engineer an objective such that its gradient matches Eq. (35).

#### B.3.1 Objective

It turns out that the following objective is the one we are looking for:

$$ $\displaystyle\mathcal{L}_{\text{CT}}({\boldsymbol{\theta}},{\boldsymbol{\theta }}^{-})=\mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[\frac{1}{\Delta_{t}} d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t), \mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t-\Delta_{t}){\boldsymbol{ \epsilon}},t-\Delta_{t}))],$ (36) $$

where ${\boldsymbol{\epsilon}}\sim\mathcal{N}(0,\mathbf{I})$ is a random noise vector. The objective in Eq. (36) is called the consistency training objective.
We can show that the gradient of $\mathcal{L}_{\text{CT}}$ indeed matches Eq. (35) in the limit of $\Delta_{t}\to 0$.
First, we apply the Taylor expansion to the unweighted loss in Eq. (36):

$$ $\displaystyle\mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[d(\mathbf{f}_{{ \boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t),\underbrace{ \mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t)}_{ \mathbf{y}}+\underbrace{\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t- \Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{ \theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t)}_{\Delta\mathbf{y}})]$ (37) $\displaystyle=\mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[d(\mathbf{f}_{ {\boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t),\mathbf{f}_{{ \boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t))+\frac{\partial d }{\partial\mathbf{y}}\Delta\mathbf{y}+(\Delta\mathbf{y})^{T}\frac{\partial^{2} d}{\partial\mathbf{y}^{2}}\Delta\mathbf{y}+O(||\Delta\mathbf{y}||^{3})]$ (38) $\displaystyle=\mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[(\Delta\mathbf {y})^{T}\frac{\partial^{2}d}{\partial\mathbf{y}^{2}}\Delta\mathbf{y}+O(|| \Delta\mathbf{y}||^{3})]$ (39) $$

where we define $\Delta\mathbf{y}$ as $\Delta\mathbf{y}=\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t-\Delta_{
t}){\boldsymbol{\epsilon}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}}(
\mathbf{x}+t{\boldsymbol{\epsilon}},t)$.

Let’s take the derivative with respect to ${\boldsymbol{\theta}}$:

$$ $\displaystyle\frac{\partial}{\partial{\boldsymbol{\theta}}}\mathbb{E}_{t}[ \mathcal{L}_{\text{CT}}({\boldsymbol{\theta}},{\boldsymbol{\theta}}^{-})]= \mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[\frac{\partial^{2}d}{ \partial\mathbf{y}^{2}}(\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t- \Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{ \theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t))\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}+O(||\Delta\mathbf{y}||^{ 3})].$ (40) $$

Using the Taylor expansion, we have

$$ $\displaystyle\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t-\Delta_{t}){ \boldsymbol{\epsilon}},t-\Delta_{t})=\mathbf{f}_{{\boldsymbol{\theta}}^{-}}( \mathbf{x}+t{\boldsymbol{\epsilon}},t)-\frac{\partial\mathbf{f}_{{\boldsymbol{ \theta}}^{-}}}{\partial\mathbf{x}}{\boldsymbol{\epsilon}}\Delta_{t}-\frac{ \partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial t}\Delta_{t}+O(\Delta _{t}^{2}),$ (41) $\displaystyle\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t-\Delta_{t}){ \boldsymbol{\epsilon}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol{\theta}}^{-}}( \mathbf{x}+t{\boldsymbol{\epsilon}},t)=-\frac{\partial\mathbf{f}_{{\boldsymbol {\theta}}^{-}}}{\partial\mathbf{x}}{\boldsymbol{\epsilon}}\Delta_{t}-\frac{ \partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}}{\partial t}\Delta_{t}+O(\Delta _{t}^{2}).$ (42) $$

Since $\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t)$ has the same value as $\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t)$, we can plug this into Eq. (40):

$$ $\displaystyle\mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[\frac{\partial^ {2}d}{\partial\mathbf{y}^{2}}(\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x }+(t-\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})-\mathbf{f}_{{\boldsymbol {\theta}}}(\mathbf{x}+t{\boldsymbol{\epsilon}},t))\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}+O(||\Delta\mathbf{y}||^{ 3})]$ (43) $\displaystyle=\mathbb{E}_{t,\mathbf{x},{\boldsymbol{\epsilon}}}[\frac{\partial ^{2}d}{\partial\mathbf{y}^{2}}(-\frac{\partial\mathbf{f}_{{\boldsymbol{\theta} }^{-}}}{\partial\mathbf{x}}{\boldsymbol{\epsilon}}-\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{\boldsymbol{ \theta}}}}{\partial{\boldsymbol{\theta}}}\Delta_{t}+O(\Delta_{t}^{2})$ (44) $\displaystyle=\mathbb{E}_{t,\mathbf{x},\mathbf{x}_{t}}[\frac{\partial^{2}d}{ \partial\mathbf{y}^{2}}(-\frac{\partial\mathbf{f}_{{\boldsymbol{\theta}}^{-}}} {\partial\mathbf{x}}\frac{\mathbf{x}_{t}-\mathbf{x}}{t}-\frac{\partial\mathbf{ f}_{{\boldsymbol{\theta}}^{-}}}{\partial t})\frac{\partial\mathbf{f}_{{ \boldsymbol{\theta}}}}{\partial{\boldsymbol{\theta}}}\Delta_{t}+O(\Delta_{t}^{ 2}),$ (45) $$

where in Eq. (45), we use the reparametrization trick $\mathbf{x}+t{\boldsymbol{\epsilon}}=\mathbf{x}_{t}$ and ${\boldsymbol{\epsilon}}=\frac{\mathbf{x}_{t}-\mathbf{x}}{t}$.
Finally, we can show that Eq. (45) matches Eq. (35) in the limit of $\Delta_{t}\to 0$ and after dividing by $\Delta_{t}$:

$$ $\displaystyle\lim_{\Delta_{t}\to 0}\frac{\partial}{\partial{\boldsymbol{\theta }}}\mathcal{L}_{\text{CD}}({\boldsymbol{\theta}},{\boldsymbol{\theta}}^{-})= \lim_{\Delta_{t}\to 0}\frac{\partial}{\partial{\boldsymbol{\theta}}}\mathcal{L }_{\text{CT}}({\boldsymbol{\theta}},{\boldsymbol{\theta}}^{-})$ (46) $$

We can add a weighting function $\omega(t)$ without affecting this equality, leading to the objective in Eq. (4).

## Appendix C Training objective of TCMs

By substituting Eq. (7) into Eq. (6), we have:

$$ $\displaystyle\mathcal{L}_{\text{CT}}(\mathbf{f}_{{\boldsymbol{\theta}},{ \boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}},\mathbf{f}_{{\boldsymbol{\theta}} ^{-},{\boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}})=\lambda_{b}\frac{\omega(t^ {\prime})}{\Delta_{t^{\prime}}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x} +t^{\prime}{\boldsymbol{\epsilon}},t^{\prime}),\mathbf{f}_{{\boldsymbol{\theta }}_{0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{\prime}}){\boldsymbol{\epsilon}} ,t^{\prime}-\Delta_{t^{\prime}})$ (47) $\displaystyle+(1-\lambda_{b})\int_{t\in S_{t^{\prime}}}\bar{\psi}_{t}(t)\frac{ \omega(t)}{\Delta_{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{ \boldsymbol{\epsilon}},t),\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x }+(t-\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})dt$ (48) $\displaystyle+(1-\lambda_{b})\int_{t\in S_{t^{\prime}}^{-}}\bar{\psi}_{t}(t) \frac{\omega(t)}{\Delta_{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{ \boldsymbol{\epsilon}},t),\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t -\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})dt$ (49) $$

The first two terms in RHS represent the boundary loss, and the last term is the consistency loss.
Let us define $\Delta_{t}=(1+8\cdot\text{sigmoid}(-t))(1-r)t$.
We define $t_{m}$ to be the smallest point such that $t_{m}-\Delta_{t_{m}}=t^{\prime}$.
Then the volume of the set $S_{t^{\prime}}:=\{t\in\mathbb{R}:t^{\prime}\leq t\leq t^{\prime}+\Delta_{t}\}$ is $\Delta_{t_{m}}$.
We assume $\bar{\psi}_{t}$ is properly designed to be upper bounded by a finite value (see Sec. 4.4).
In the limit of $r\to 1$, we simplify the the second term in the above:

$$ $\displaystyle(1-\lambda_{b})\int_{t\in S_{t^{\prime}}}\bar{\psi}_{t}(t)\frac{ \omega(t)}{\Delta_{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{ \boldsymbol{\epsilon}},t),\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x }+(t-\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})dt$ (50) $\displaystyle=(1-\lambda_{b})\int_{t\in S_{t^{\prime}}}\bar{\psi}_{t}(t^{ \prime})\frac{\omega(t^{\prime})}{\Delta_{t^{\prime}}}d(\mathbf{f}_{{ \boldsymbol{\theta}}}(\mathbf{x}+t^{\prime}{\boldsymbol{\epsilon}},t^{\prime}) ,\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{ \prime}}){\boldsymbol{\epsilon}},t^{\prime}-\Delta_{t^{\prime}})+O(1)dt$ (51) $\displaystyle=(1-\lambda_{b})\bar{\psi}_{t}(t^{\prime})\frac{\omega(t^{\prime} )}{\Delta_{t^{\prime}}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t^{ \prime}{\boldsymbol{\epsilon}},t^{\prime}),\mathbf{f}_{{\boldsymbol{\theta}}_{ 0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{\prime}}){\boldsymbol{\epsilon}},t^{ \prime}-\Delta_{t^{\prime}})Vol(S_{t^{\prime}})+\int_{t\in S_{t^{\prime}}}O(1)dt$ (52) $\displaystyle\approx(1-\lambda_{b})\bar{\psi}_{t}(t^{\prime})\frac{\omega(t^{ \prime})}{\Delta_{t^{\prime}}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+ t^{\prime}{\boldsymbol{\epsilon}},t^{\prime}),\mathbf{f}_{{\boldsymbol{\theta} }_{0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{\prime}}){\boldsymbol{\epsilon}}, t^{\prime}-\Delta_{t^{\prime}})\Delta_{t^{\prime}}+\int_{t\in S_{t^{\prime}}}O (1)dt$ (53) $\displaystyle\approx(1-\lambda_{b})\bar{\psi}_{t}(t^{\prime})\omega(t^{\prime} )d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t^{\prime}{\boldsymbol{ \epsilon}},t^{\prime}),\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+( t^{\prime}-\Delta_{t^{\prime}}){\boldsymbol{\epsilon}},t^{\prime}-\Delta_{t^{ \prime}}),$ (54) $$

where in Eq. (51) we apply the Taylor expansion to the integrand, and in Eq. (53), we can see that $\Delta_{t_{m}}=(1+8\cdot\text{sigmoid}(-{t_{m}}))(1-r)t_{m}$ goes to zero as $r\to 1$.
Thus, $t_{m}\to t^{\prime}$ and then $Vol(S_{t^{\prime}})/\Delta_{t^{\prime}}=\Delta_{t_{m}}/\Delta_{t^{\prime}}=
\frac{(1+8\cdot\text{sigmoid}(-t_{m}))t_{m}(1-r)}{(1+8\cdot\text{sigmoid}(-t^{
\prime}))t^{\prime}(1-r)}=1$ in the limit.
Hence, the boundary loss is

$$ $\displaystyle(\lambda_{b}\frac{\omega(t^{\prime})}{\Delta_{t^{\prime}}}+(1- \lambda_{b})\bar{\psi}_{t}(t^{\prime})\omega(t^{\prime}))d(\mathbf{f}_{{ \boldsymbol{\theta}}}(\mathbf{x}+t^{\prime}{\boldsymbol{\epsilon}},t^{\prime}) ,\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{ \prime}}){\boldsymbol{\epsilon}},t^{\prime}-\Delta_{t^{\prime}})$ (55) $\displaystyle\approx\lambda_{b}\frac{\omega(t^{\prime})}{\Delta_{t^{\prime}}}d (\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t^{\prime}{\boldsymbol{\epsilon }},t^{\prime}),\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+(t^{ \prime}-\Delta_{t^{\prime}}){\boldsymbol{\epsilon}},t^{\prime}-\Delta_{t^{ \prime}}),$ (56) $$

where the first term ($O(1/\Delta_{t})$) dominates the second term ($O(1)$).
We hence arrive at Eq. (9):

$$ $\displaystyle\mathcal{L}_{\text{CT}}(\mathbf{f}_{{\boldsymbol{\theta}},{ \boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}},\mathbf{f}_{{\boldsymbol{\theta}} ^{-},{\boldsymbol{\theta}}_{0}^{-}}^{\text{trunc}})\approx{\lambda}_{b} \underbrace{\frac{\omega(t^{\prime})}{\Delta_{t^{\prime}}}d(\mathbf{f}_{{ \boldsymbol{\theta}}}(\mathbf{x}+t^{\prime}{\boldsymbol{\epsilon}},t^{\prime}) ,\mathbf{f}_{{\boldsymbol{\theta}}_{0}^{-}}(\mathbf{x}+(t^{\prime}-\Delta_{t^{ \prime}}){\boldsymbol{\epsilon}},t^{\prime}-\Delta_{t^{\prime}}))}_{\text{ Boundary loss}:=\mathcal{L}_{B}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{f}_{{ \boldsymbol{\theta}}_{0}^{-}})}$ (57) $\displaystyle+(1-\lambda_{b})\underbrace{\mathbb{E}_{\bar{\psi}_{t}}[\frac{ \omega(t)}{\Delta_{t}}d(\mathbf{f}_{{\boldsymbol{\theta}}}(\mathbf{x}+t{ \boldsymbol{\epsilon}},t),\mathbf{f}_{{\boldsymbol{\theta}}^{-}}(\mathbf{x}+(t -\Delta_{t}){\boldsymbol{\epsilon}},t-\Delta_{t})]}_{\text{Consistency loss}:= \mathcal{L}_{C}(\mathbf{f}_{\boldsymbol{\theta}},\mathbf{f}_{{\boldsymbol{ \theta}}^{-}})},$ (58) $$

where because $\Delta_{t}\to 0$, we can approximate $S_{t^{\prime}}^{-}\approx(t^{\prime},T]$, leading to Eq. (58).

In practice, we set $r$ close to one during the truncated training stage (see Section D).
Note that here, the boundary loss can dominate the consistency loss when $\mathbf{f}_{\boldsymbol{\theta}}$ and $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}^{-}$ are sufficiently different around $t=t^{\prime}$.
However, in practice, as we set $\Delta_{t}$ to be small enough but not all the way to zero, and as we use Pseudo-Huber loss (see Section D) with small $c$ value that normalizes the effect of the loss magnitude on the gradient norm, we can balance the training.

## Appendix D Implementation details

We provide detailed information about our implementation in the following.

##### Model initialization and architecture:

All stage 1 models are initialized from pre-trained EDM or EDM2 checkpoints as suggested by Geng et al. ().
For CIFAR-10, we use EDM’s DDPM++ architecture, which is slightly smaller than iCT’s NCSN++. For ImageNet $64\times 64$, we employ EDM2-S (280M parameters) and EDM2-XL (approximately 1.1B parameters) architectures. EDM2-S is slightly smaller than iCT’s ADM architecture (296M parameters).

##### Model parameterization:

Following EDM, we parameterize consistency models $\mathbf{f}_{{\boldsymbol{\theta}}}$ as $\mathbf{f}_{{\boldsymbol{\theta}}}=c_{\text{out}}(t)\mathbf{F}_{{\boldsymbol{
\theta}}}(\mathbf{x},t)+c_{\text{skip}}(t)\mathbf{x}$,
where $c_{\text{out}}(t)=\frac{t\sigma_{\text{data}}}{\sqrt{\sigma_{\text{data}}^{2}+
t^{2}}}$, $c_{\text{skip}}(t)=\frac{\sigma_{\text{data}}^{2}}{\sigma_{\text{data}}^{2}+t^
{2}}$, and $\sigma_{\text{data}}=0.5$.

##### Training details:

We set $\Delta_{t}=(1+8\cdot\text{sigmoid}(-t))(1-r)t$, where $r=\max\{1-1/2^{\lceil i/25000\rceil},0.999\}$ for CIFAR-10 and $\max\{1-1/4^{\lceil i/25000\rceil},0.9961\}$ for ImageNet $64\times 64$, with $i$ being the training iteration.
For the second stage, we start with the maximum $r$ values (i.e., 0.999 or 0.9961) and do not change them.
The weighting function $\omega(t)$ is set to 1 for CIFAR-10 and $\Delta_{t}/c_{\text{out}}(t)^{2}$ for ImageNet $64\times 64$.
As suggested by Song & Dhariwal (); Geng et al. (), we use the Pseudo-Huber loss function $d(\mathbf{x},\mathbf{y})=\sqrt{||\mathbf{x}-\mathbf{y}||_{2}^{2}+c^{2}}-c$, with $c=1e-8$ for CIFAR-10 and $c=0.06$ for ImageNet $64\times 64$.
This is especially crucial for our method as the boundary loss can dominate the consistency loss.
The boundary loss compares the outputs from the different model $\mathbf{f}_{\boldsymbol{\theta}}$ and $\mathbf{f}_{{\boldsymbol{\theta}}_{0}}$, it tends to be larger than the consistency loss, but Pseudo-Huber loss effectively normalize the effect of the loss magnitude on the gradient norm.
For ImageNet $64\times 64$, we employ mixed-precision training with dynamic loss scaling and use power function EMA (Karras et al., ) with $\gamma=6.94$ (without post-hoc EMA search).

##### Learning rate schedules:

EDM2 (Karras et al., ) architectures require a manual decay of the learning rate.
Karras et al. () suggest using the inverse square root schedule $\frac{\alpha_{\mathrm{ref}}}{\sqrt{\max\left(t/t_{\mathrm{ref}},1\right)}}$.
For the first stage training of EDM2-S on ImageNet, we use $t_{\text{ref}}=2000$ and $\alpha_{\text{ref}}=1e-3$ following Geng et al. ().
For the second stage training of EDM2-S, we use $t_{\text{ref}}=8000$ and $\alpha_{\text{ref}}=5e-4$.

Second stage training of EDM2-XL is initialized with the ECM2-XL checkpoint from Geng et al. ().
During the second stage, we use $t_{\text{ref}}=8000$ and $\alpha_{\text{ref}}=1e-4$ for EDM2-XL.

##### Time step sampling:

For the first stage training, we use a log-normal distribution for $\bar{\psi}_{t}$.
For CIFAR-10, we use a mean of -1.1 and a standard deviation of 2.0 following Song & Dhariwal ().
For ImageNet, we use a mean of -0.8 and a standard deviation of 1.6 following Geng et al. ().

For EDM2-XL, we also explore $t^{\prime}=1.5$ for truncated training, adjusting $\nu$ to 2 to ensure $\bar{p}_{t}$ has high probability mass around $t^{\prime}=1.5$ and also has a long tail as discussed in Sec. 4.4.
This way, we get the FID of 2.15, which is slightly better than the result in Table 2.

During two-step generation, we evaluate the model at $t=80,1$ on CIFAR-10 and $t=80,1.526$ for ImageNet.

## Appendix E Uncurated generated samples

We provide the uncurated generated samples in Fig. 8-10.

Figure: (a) One-step samples.
Refer to caption: extracted/5938731/figs/cifar-n1_8_20.png

Figure: (a) One-step samples.
Refer to caption: extracted/5938731/figs/s-n1-4_10.png

Figure: (a) One-step samples.
Refer to caption: extracted/5938731/figs/xl-n1_4_10.png

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: - Training consistency models by the original loss (i.e. the method used in stage-1) with the improved proposal distribution by focusing more on [0, t']. The optimal proposal distribution may be different from the stage-2 so it should be tuned a little bit for such task. Will the optimal result be comparable to the 2-stage training method?
- Consistency models include both consistency training (CT) and consistency distillation (CD), but this paper only shows the results of CT. Note that the original CT has the same gradient limit as CD as $\Delta t \to 0$. However, with truncation training, this conclusion will not hold anymore for t near to t'. What is the performance of CD by the truncated method?

(Please refer to Weakness section for detailed explanations)
soundness: 2
strengths: - The paper is well-written and is easy to follow. The math derivations are technically correct.
- The idea of truncating is novel, and the proposed method can further address the training instability of consistency models which are important to the community.
- The empirical results are strong, showing the effectiveness of the method.
summary: This work proposes a novel training technique for consistency models by truncating the time intervals into two splits: one is the boundary/denoising part near $t=0$, and the other is the consistency/generation part near $t=T$. By training with a weighted sum of the two objectives and introducing a dedicated proposal distribution for sampling the time, the proposed method beats previous 1-step and 2-step models on CIFAR10 and ImageNet64.
weaknesses: It is unclear whether the effectiveness is from the truncation or is just from the changing of proposal distribution by focusing more on the boundary parts. Specifically, the major techniques in this paper include two parts:

1. two-stage truncated training to avoid the overtraining of denoising tasks;
2. changing proposal distribution to focus more on the boundary conditions.

However, with only the first method, the training still diverges (L241-L242), which shows that the part 2 seems to be more important. Intuitively, as the supervision signals of consistency training only come from the boundary condition, the model will suffer from error accumulations as t increases from 0 to T. Therefore, focusing more on the boundary parts seem to be a natural idea to address such error accumulations and may be the reason why the proposed method is more stable than previous methods. Given such understanding, it is unclear whether the first part (truncation) is even necessary.

Therefore, this paper does not fully sound to me. It would be great if the author could conduct more rigorous ablations for the part 2, such as:

- Training consistency models by the original loss (i.e. the method used in stage-1) with the improved proposal distribution by focusing more on [0, t']. The optimal proposal distribution may be different from the stage-2 so it should be tuned a little bit for such task. Will the optimal result be comparable to the 2-stage training method?

Besides, there is another minor weakness but I think it does not affect the conclusions in this paper:

- Consistency models include both consistency training (CT) and consistency distillation (CD), but this paper only shows the results of CT. Note that the original CT has the same gradient limit as CD as $\Delta t \to 0$. However, with truncation training, this conclusion will not hold anymore for t near to t'. What is the performance of CD by the truncated method?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Could the authors discuss the relation with PCM [1]?
- What are the total training iterations and time for the first and second stages, respectively? Is the first stage exactly the same as previous consistency models? If so, why the authors only initialize the truncated training stage with ECT for EDM2-XL?
- I think a limit case for the truncated training is to only learn the 1-step mapping $T\rightarrow 0$, i.e., limiting the truncated region to a single point $T$. The model can be directly trained on noise-data pairs. Can this method achieve better 1-step FID?


[1] Phased Consistency Model
soundness: 3
strengths: - The motivation is the difficulty of learning the consistency function across the whole time zone, which is sound. The second stage model only needs to map noised data to clean data in a limited time interval, which is an easier task, and it is no surprise to bring better fitting for the ODE trajectory.
- The phenomenon that the consistency training gradually weakens the model's denoising capabilities at small $t$ is well illustrated in Figure 2.
- The method achieves better FID on standard image datasets than previous consistency models, especially in 1-step generation.
- Ablation studies for the dividing time and the number of stages are comprehensive.
summary: The paper divides the training of consistency models into two stages, where the first stage is the standard one, and the second stage only learns the consistency function in a truncated time region $[t',T]$. The training loss for the second stage is a mixture of the consistency loss and an extra boundary loss in $[t',t'+\Delta t]$. After careful tuning of the weighing functions and timestep schedules, the second stage model can achieve better 1-step and 2-step FID on CIFAR-10 and ImageNet 64x64 than standard consistency models.
weaknesses: - The overall idea and motivation are actually not new. ECT already points out the trade-off between the denoising capacity and the consistency capacity, suggesting the initialization of a consistency model with a diffusion model. The authors' work, in my opinion, is to use a two-stage method to replace the dedicated iteration-dependent training schedule in ECT.
- The 2-step performance improvement is relatively marginal compared to ECT. As suggested by ECT, it is better to use a 2-step generation with a small model instead of a 1-step generation with a large one. Therefore, the actual benefits provided by truncated training are limited, especially considering that it requires careful tuning of the weighing functions and timestep schedules.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Several studies, such as [1, 2], have proposed truncating the time interval in diffusion model training to address issues like the unbounded score problem or high Lipschitz continuity when $t\approx0$. In principle, similar phenomena could arise in consistency training, as it aims to learn a noise-to-data mapping that may also exhibit large Lipschitz continuity. How does the proposed method relate to this existing literature?


2. How can one ensure that the trained CT model, using the proposed parameterization in Eq. (5), will ultimately be continuous at the splitting time $t'$? If one performs multi-step sampling and selects (some) sampling timesteps close to $t'$, this could potentially cause issues?



[1] Kim, D., Shin, S., Song, K., Kang, W., & Moon, I. C. (2021). Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation.

[2] Yang, Z., Feng, R., Zhang, H., Shen, Y., Zhu, K., Huang, L., ... & Cheng, F. (2023). Eliminating lipschitz singularities in diffusion models.
soundness: 2
strengths: The paper is well-structured and clear. I appreciate the extensive experiments illustrating the trade-off between denoising and generation in consistency training. Additionally, the empirical results seem competitive to baselines.
summary: This work proposes a two-stage consistency training (CT) approach by dividing the time interval into two segments. Stage 1 employs standard CT, serving as a boundary anchor for Stage 2, which focuses on larger time consistency training. This method, the authors claim, emphasizes generation over denoising.
weaknesses: 1. The approach appears a bit over-engineered, with numerous handcrafted designs and hyperparameters, such as the weighting function $\psi_t(t)$, $\lambda_b$, $N_B$, $\Delta_t$, $\Delta_{t'}$, and the interval division $t'$.

2. How should one determine the terminate point of Stage 1 training? Monitoring its progress may introduce additional complexity. The number of training iterations for Stage 1 may also represent a crucial hyperparameter that may require further ablation studies.

3. The paper lacks ablation studies for certain hyperparameters relevant to the proposed method, such as $N_B$ (or $\rho$), $\Delta_t$, $\Delta_{t'}$.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Besides the weaknesses listed above, I also have some questions.
- I am not an expert on consistency models. But there are many other works on consistency distillation, such as CTM / TCD. Why not comparing against them?
- SiD seems to get better FID with significantly fewer parameters. Why variational score distillation methods are listed in separate categories and not compared against consistency models?
- How does the proposed method apply to guided sampling / text-to-image applications? How does guidance scale impact the proposed method? (Including more guided sampling results would improve the submission.)
soundness: 3
strengths: - The paper largely focus on an orthogonal problem of network capacity, which can be possibly combined with many consistency distillation works.
- Train two separate models with a boundary condition to enforce consistency is empirically challenging (though theoretically straightforward), but the paper shows that it is possible. This might be informative for subsequent works that require different models across time ranges.
summary: This paper proposes an interesting way to improve the training of consistency models. It points out, simutaneously learn to denoise and generate can require more capacity. Based on the belief, the paper proposes to separately use two models for denoising and generation, with a two-stage training procedure that first train a standard consistency model, and then train a generation model. The challenging part is the training loss and weighting function need to be carefully designed to enforce the boundary condition between two time ranges. The proposed approach seems to successfully train stably with continuously decreased FID, and improves upon existing work ECM.
weaknesses: - While the proposed method improves considerably on small models and few time steps, the improvement on the highest quality 2-step EDM2-XL is negible. The applicability of the proposed method on large-scale models is questionable.
- The proposed method seems to still perform worse than SiD.

---

confidence: 4
contribution: 2
initial_rating: 8
presentation: 4
questions: I have no specific question -- cf. the above comments. Overall, I recommend an "accept" but am waiting for the authors' response to my comments and look forward to discussing with the other reviewers on the matter.
soundness: 3
strengths: This paper is **very well written**: it is easy to read and understand. All ideas are presented and articulated clearly. The method's description and results are presented in a **simple and efficient manner**, making it easy to reimplement. Given its generality, I believe that Truncated Consistency Models may **become a widespread trick** to improve consistency models' performance.

The model is particularly **well motivated**, and its design **well supported** by sound experiments. The latter suscessfully highlight the denoising-generation tradeoff of standard consistency models. The efficiency of the proposed method to improve the generation performance at later training steps is demonstrated in the experiments. **Additional insights on design choices** (cutoff time, relevance of more training stages, etc.) are appreciated.

To my knowledge, both the presented insights and model are novel, even though the truncation principle has become a standard technique in diffusion (Balaji et al., 2022; Zheng et al., 2023 -- that could be cited in the related work) and consistency (Heek et al., 2024) models. Overall, the paper's contribution has a **good significance potential** as the presented insights are valuable for all practitioners, the experimental results are appealing and the method is generally applicable. The denoising-generation duality also echoes more fundamental work unveiling phase transition phenomena in diffusion models (Ambrogioni, 2023; Biroli et al., 2024; Sclocchi et al., 2024).

Ambrogioni. The statistical thermodynamics of generative diffusion models: Phase transitions, symmetry breaking and critical instability. arXiv:2310.17467, 2023.\
Balaji et al. eDiff-I: Text-to-Image Diffusion Models with Ensemble of Expert Denoisers. arXiv:2211.01324, 2022.\
Zheng et al. Truncated Diffusion Probabilistic Models and Diffusion-based Adversarial Auto-Encoders. ICLR 2023.\
Biroli et al. Dynamical Regimes of Diffusion Models. arXiv:2402.18491, 2024.\
Heek et al. Multistep Consistency Models. arXiv:2403.06807, 2024.\
Sclocchi et al. A Phase Transition in Diffusion Models Reveals the Hierarchical Nature of Data. arXiv:2402.16991, 2024.
summary: This paper identifies a shortcoming of consistency models: their joint temporal parameterization via a single network which limits their learnability. By noticing that early denoising timesteps, irrelevant for the final generation task, are neglected at the end of the training, the authors propose to train a consistency model on large noise levels only, with a pretrained consistency model serving as initial condition at the cutoff point. The relevance and the advantage of this approach are empirically demonstrated as the resulting model obtains state-of-the-art performance with limited number of parameters and reduced instability.
weaknesses: While they are not critical and may not significantly change my assessment, I would like the authors to tackle the following weaknesses.

### Unnecessary decomposition of the loss

I do not understand the necessity of the loss decomposition in Eq. (6). To my understanding, it amounts to giving a specific sampling weight for timesteps close to $t'$ in $\psi$. The parameterization of Eq. (5) should then handle the corner case. Additionally, it appears that $S_{t'} = \{t'\}$, so the notations may be simplified in this part of the paper.

### Additional insights & reproducibility

The following discussions / content would benefit the paper.
- The code used by the authors for the presented experiments should be included in the supplementary material for further reproducibility.
- One can notice in Figure 3 that the denoising FID increases at the beginning of stage 2. I would be interested in the authors' opinion on this phenomenon.
- Still in Figure 3, can the authors explain why the denoising FIDs do not start from the same value between stage 1 and stage 2 (assuming they start from the same checkpoint)?
- While the authors acknowledge an additional training and memory cost in Section 6, I am not sure they include the additional cost induced by the supplemental training iterations to complement the pretrained consistency model (up to 200k additional optimization steps based on the different plots).

### Minor issues

- Section 4.2 insists on the forgetting by the truncated model on the lower noise levels. This is not an advantage or a confirmation of the model's abilities, but a byproduct of the truncated training loss.
- The number of optimization steps is not specified in the hyperparameters listed in Section D.
- In Figure 3, the timestep $t = 1$ should be placed in the second row to be consistent with the rest of the paper.
- The paper should specify where the baseline results from Tables 1 and 2 are reported from.

Meta-review: This paper proposes a method to improve the training of consistency models by addressing the capacity challenge of simultaneously learning denoising and generation. It introduces a two-stage approach: first training a standard consistency model for denoising, then training a separate generation model. A key challenge is designing the loss and weighting function to enforce boundary conditions between time ranges. The proposed method achieves stable training with progressively decreasing FID and outperforms existing work.

All reviewers find the paper well-written, interesting, and with strong empirical results.

Decision: Accept (Poster)