Truncated Consistency Models
Authors: Sangyun Lee
Yilun Xu
Tomas Geffner
Giulia Fanti
Karsten Kreis
Arash Vahdat
Weili Nie
Venue: ICLR 2025 Poster (2025)
Keywords: Consistency models
diffusion models
generative models
Primary area: generative models
TL;DR: We improve the one-step and two-step generation performance of consistency models by training them on a truncated time interval

Abstract: Consistency models have recently been introduced to accelerate the generation speed of diffusion models by directly predicting the solution (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 prevent the truncated-time training from collapsing to a trivial solution.
Experiments on CIFAR-10 and ImageNet $64\times64$ 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.

Published as a conference paper at ICLR 2025
TRUNCATED CONSISTENCY MODELS
Sangyun Lee" Carnegie Mellon University
Yilun Xu
Tomas Geffner NVIDIA
Giulia Fanti
NVIDIA
Camegie Mellon University
Karsten Kreis
Arash Vahdat
Weili Nie NVIDIA
NVIDIA
NVIDIA
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 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 carlier time steps and focus its capacity on generation. We propose a parameterization of the consistency function and two-stage training procedure that prevents the truncated- time training from collapsing to a trivial solution. Experiments CIFAR-10 and ImageNet 64 x 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 smaller networks. Project page: https://ruuncated-mm /
INTRODUCTION
Diffusion models (Ho et al., 2020 Song et al. 2020) have demonstrated remarkable capabilities in generating high-quality continuous data such as images, videos, or audio (Ramesh et al., [202]] Ho et al. 2022 Huang et al.) (023). Their generation process gradually transtorms 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. 2023) 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 single step. To circumvent the need for simulating large number of noise-data pairs to learn this mapping, as employed in prior works (Liu et al., 2022b Luhman & Luhman (00T), consistency models learn 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 = 0 to the initial t = T.
However, the advantage of simulation-fre 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. To. 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, where most of the original data information has been erased. This dual task necessitates larger network capacity, and it is challenging for single model to excel at both tasks. Our empirical observations in Fig. 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-offi off, rather than allowing the model to allocate capacity
Work mostly done while interning NVIDIA

---

Published as a conference paper at ICLR 2025
uncontrollably across times. This raises key question: Can explicitly reduce the network capacity dedicated to the denoising task in order to improve generation?
In this paper, we propose new training algorithm, termed Trancated 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, 7 of PF ODE trajectories, to a new objective that focuses on a truncated time range T] where e' 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'). Crucially, we show that proper boundary condition t' is necessary ensure the new model adheres to the original consistent mapping. To achieve this, we propose two-stage training procedure (see Fig. [a): The first stage involves pretraining standard consistency model over the whole time range. This pretrained model then acts as the boundary condition at t' 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-IO and ImageNet 64 x 64 datasets, TCM outperforms the iCT (Song & Dhariwal 2023), the previous best consistency model, in both one-step and two-step generation using similar network size. TCM even outperforms iCT--cep that uses 2 larger network across datasets and sampling steps. By using largest network. we achieve a one-step FID of 2.20 on ImageNet 64 x 64, which is competitive with the current otate-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
2.1 DIFFUSION MODELS
Diffusion models are class of generative models that synthesize data by reversing forward process in which the data distribution Pdata is gradually transformed into a tractable Gaussian distribution. In this paper, we use the formulation proposed in Karras al. (2022), where the forward process is defined by the following stochastic differential equation (SDE):
dx:- = V2tdwi
(I)
where t E [), and w, is the standard Brownian motion from t = 0 to t T. Here, we define Pe as the marginal distribution of x, along the forward process, where po = Paatn- In this case, pt is perturbed data distribution with the noise from (0, In diffusion models, 7 is set to be large enough so that pr is approximately equal to a tractable Gaussian distribution (0,
Diffusion models come with the reverse probability flow ODE (PF ODE) that starts from t T to t t= and yields the same marginal distribution P: as the forward process Eq. (T) (Song et al. 2020):
dx, = (x;
where (x,) = Vx log (x) is the score function at time 1. [0. 7]. To draw samples from the data distribution Pdata; we first train a neural network to learn sx() using the denoising score matching (Vincent, 2011), initialize xT with a sample from AO 11), and solve the PF ODE backward in time: Xn = Xy + [ (-ts (x;)) However, numerically solving the PF ODE requires multiple forward passes of the neural score function estimator, which is computationally expensive.
2

---

Published as a conference paper at ICLR 2025
Data
Probability ODE
Noise
2.5
iterations (x10*)
Stage I: Standard consistency training
CIFAR-10
6.0
Stage
5.3
Stage
5.0
245
blows up!
0
4.0
35
16
Itarations (x (x10%)
Stage Truncated consistency training
(a)
(b)
EmageeNet
Figure 1: (a) Two-stage training of TCM. In Stage I, standard consistency model is trained to provide both the boundary condition and initialization for TCM training in Stage TCM focuses on learning in the [t', 7"I range, discarding denoising tasks at earlier times and allocating network capacity toward generation-like tasks 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.
0.2
0.4
t= 0.8
3.5
2.50
1.2
2.25
3.0-
6 1.1-
1.75
2.0-
300002 200b00 300001 Iteration
100000
300063 Iteration 4,0
200000 300000 iteration
1.5
100002
Iterati
16.0
80.0
JFID-
0RR0 20
2-
300002 200bon
100000
2000n0
102002
Iteration
Figure 2: Evolution of the denoising FID (dFID,) during standard consistency training for different where 0 < < 80 follows the EDM noise schedule (Karras al., 2022). The model gradually sacrifices its demoising capability at smaller times (t < 1.0) to trade for the improved generation quality t 80 as training proceeds.
2.2 CONSISTENCY MODELS
Consistency models instead aim directly map from noise to data, by learning consistency
function that outputs the solution of PF ODE starting from any t = [0. T.. The desired consistency function f should satisfy the following two properties (Song et al. 2023): (i) f(xo. 0) = xo, and (ii)
f(x,,t) f f(x,, s), (s, t) [0, The first condition can be satisfied by the reparameterizatioo
fx(x, !) == ( t) Cskip(()).
(3)
where 0 is the parameter of the free-form neural network Fş RRd x R + gd and Cour(0) 0. Casi (0) = 1 following the similar design Karras et al. (2022). Here, instead of training fo directly, train surrogate neural network Fe under the above reparameterization. The second condition can be learned by optimizing the following consistency training objective:
(t) dffe(x + te, t), fo (x + (t (t-Ai -
(4) where = stopgrad(()) 1h denotes the probability of sampling time t that also represents the noise scale, denotes the standard Gaussian noise, w(t) is weighting function, d(...) is distance
(fr, fg) := Er xx..pm.ee..........mm... ((...) A,
3

---

Published as a conference paper at ICLR 2025
D.1.2 and A represents the nonnegative difference between two consecutive time steps that is usually set to a monotonically increasing function of t.
function defined in Sec. Sec.
The gradient of Let with respect to 0 is an approximation of the underlying consistency distillation loss with a (max, A)) error (See Appendix [D). Song et al. (023)) empirically suggests that A 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: 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 dFID,, as the Fréchet inception distance (FID) (Heusel et al., (017) between the original data Pdnt and the denoised data by consistency models with inputs sampled from ps. When computing dFID,, we first add Gaussian noise from (O0 121 to 50K clean samples and then denoise them using consistency models. Hence, dFID,, is close to zero, and dFIDr is the standard FID.
3 TRUNCATED CONSISTENCY MODEL
Standard consistency models pose 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 its data endpoint. Hence, 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 data). This challenge mainly contributes to consistency models' underperformance relative to other generative models with similar network capacities (see Table D
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 < L) increase while the denoising FIDs at larger t 1) (including the generation FID at the largest = 80) continue 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 explicitly control this tradeoff by forcing the consistency training to ignore the denoising task for small values of I, thus improving its capacity usage for generation. We thus generalize the consistency model objective in Eq. (4) and apply only in the truncated time range [t', T" where the dividing time E' lies within (0, T). The time probability 1 in TCM only has support in [t', T] as a result.
Naive solution A straightforward approach is directly train consistency model the truncated time range. However, the model outputs can collapse to an arbitrary constant because a constant function (i.e., (x, I) const) is minimizer of the consistency training objective (Eq. 4 In standard consistency models, the boundary condition f fxx, 0) = Xo prevents collapse, but in this naive example, there no such meaningful boundary condition. For example, if the free-form neural network F F f)= Cskip (f)x//cc (txx/C/uut)))) for all E T). fo (x, t) is O, and thus Eq. 44) becomes zero To handle this, we propose a two-stage training procedure and design 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 convergence in the usual fashion, with the training objective in Eq. 4 we denote the pre-trained model as fo.-
2. Stage 2 (Truncated consistency training): We initialize a new consistency model fo with the first-stage pretrained weights fe,, and train a truncated time range [t', T). The boundary condition at time i' is provided by the pretrained fø.. This stage is explained further below.
To explain the details of TCM, we first introduce the following parameterization:
girune (x, t) fo t) 1(t t'} + fo; (x.t) 1(t < t'}. 0.0
(5)
4

---

Published as a conference paper at ICLR 2025
where 1(-) is the indicator function, and similarly, 0 stopgrad(Ou). Intuitively, we only
final model fg when 2 t', and we inquire the pre-trained fo, otherwise. This approach ensures that (1) fo does not waste its capacity learning in the [0, f') range, and (2) if fo is trained well, it will
learn to generate data by mimicking the pre-trained model f0r at the boundary. When = 0, we
recover the standard consistency model parameterization Eq. (3) During sampling, firune = fe for 0,0,
all & [t', T]. can discard this paramcterization and just use fo for generating samples.
To describe the boundary condition, we then partition the support of the time sampling distribution th, i.e., [t', T] into two time ranges: (i) the boundary time region St {t E R :1' S l' Ar). and (ii) the consistency training time region Sp [t', I] Sy (t R it + A, <t<T). To effectively enforce the boundary condition using the first-stage pre-trained model fe.. nonnegligible amount of t's, sampled from the, must fall within the interval Sy Otherwise the consecutive time steps and t - A, in consistency training will mostly be larger or equal to t', limiting the influence of the pre-trained model.
With this time partitioning and our new parameterization, Eq. 4 can be decomposed as follows:
w(t) d(fo(x + te, t), foo (x + (t - De)E, - D.)dt
CC (frrunc 0.0," 0
pirune
.0
JIES
(e(())
Boundary loss
(6)
(t(t) (01 d(fo + t) (t-(tii))) - -i)dtt,
SES, A.
Consistency loss
where we apply our parameterization in Eq. in the above two time partitions separately, and we drop the expectation over x ~ Pdatn, E ~ (0, 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 the objective (6) can still collapse to constant if do not utilize the boundary condition sufficiently by not sampling enough time f's in Sp. In particular, this can happen for A, close to zero when consistency training is prevent this, we design th to satisfy JIES, Vhx (t) > O. In other words, we have a strictly positive probability of sampling a point in Sr, even when A, is close to zero. Specifically, we define v as mixture of the Dirac delta function 8()) at point and another distribution dh:
convergence (Song & Dhariwal 2023 Geng et al.) 2024). To
(1) Aso(t (') (1 (1->x) (t),
where the weighting coefficient A 1). &h has the support (t', T] and can be instantiated in different ways (e.g., log-normal or log-Student-t distributions); the effect of different ii choices is explored Section 4.4
By definition, we can see that SEes,, (x(d)) 2 As and A controls how significantly we emphasize the boundary condition. Assume that the first-stage consistency model is perfectly trained in [o, t'), i.e., (x,, t) Xa for all & E [0, l'). If ('( + fe. ('), fo will be penalized by the boundary loss. Minimizing the boundary loss enforces the boundary condition in second-stage model fo (i.e.. fo (x(x,,t')
fe, (xv,t') t') xo), while minimizing the consistency loss propagates the boundary condition to the end time (i.e., fo (XT, T) = fo (xx,,, t')). Consequently, the loss in Eq. (6) effectively guides the model towards the desired solution (x(,,, T) Xo. With the time distribution
Rhe defined in Eq. 7. our training objective becomes
Rd A w0t'),
CC petuun pinine
d(fo(x l'e, 1)), S0; (x (1' A:)E, Ar))
(8)
.0
Ar
Boundary loss: (Lo. (Co.
w(I) d d(fo(x + te, (x + (t-A-rp) Di)e,t Ad).
+(1 Ao) E
(9)
A
Consistency loss= Lo ffo
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Truncated Consistency Training
1: Standard consistency training
2: Ou <- arg ming Ler (fg. (a-)
Optimize consistency training loss for the regular model
3: Truncated training
Ne t -[pp]
Number boundary samples
5: for each training iteration do
6: X1, xB ~ Paata E1, EL,,.., ED 1)
7: Set te,...,, &Ng to f', and
NNt.1.... tB the
8:
Compute 2 (CB):
using Eq. ( with (xi, 6i, ti) for = L,..., NB
9:
Compute _j (Cc), (fo, using Eq.
(9) with (x), 65, for j = Na + I,... B
10: Compute CreM using Eq. [TT
Update 0 using the computed gradient
12:
end for
where the approximation in Eq. 8 holds when Ar is sufficiently small (which is true for the truncated training stage). Please see Appendix for the detailed derivation. For simplicity of notation, we relax the above objective by absorbing the (1 As)) factor into As and express our final training loss as:
LreM == W&CB fa;) + Cf fg-),
(10)
where w/ >s/(1 Aa)) is 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 LB is estimated using NB [B]] samples, where E (0, 1) is a hyperparameter controlling the allocation of samples. The consistency loss Le estimated with the remaining B -N samples. Increasing 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-bater loss is follows:
Na
B
LTCM
Whb
Σ
(fo. fo
Σ
Vo(Cc), fa-),
(11)
NB
B B-N B
j=No+1
where (CB)s and (Lc)) are the boundary loss and the consistency loss the i-th sample from 5(t t') and the j-th sample from Vh, respectively, We provide the training algorithm in Algorithm
EXPERIMENTS
In this section, we evaluate TCM on standard image generation benchmarks and compare it against state-of-the-arr generative models. We begin by detailing the experimental setup in Sec. |4. We then study the behavior of denoising FID and its impact 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. 2009) and ImageNet 64 (Deng et al.) 2009) datasets. We consider the unconditional generation task on CIFAR-10 and class-conditional generation on ImageNet 64x64. We measure sample quality with Fréchet Inception Distance (FID) (Heusel al. 2017) (lower is better), as is standard in the literature.
For consistency training in TCM, we mostly follow the hyperparameters in ECT (Geng et al. 2024). including the discretization curriculum and continuous-time training schedule. For all experiments. we choose a dividing time t' = 1 and set 3h to the log- Studetttttt-t distribution. We use wA = 0.1 and p = 0.25 for the boundary loss. We discuss these choices in Sec. 4.4 In line with Geng et al. (2024), we initialize the model with the pre-trained EDM (Karras et al.) 2022) / EDM2 (Karras et al.) 2024) for CIFAR-I0, ImageNet 64 x 64, respectively. On CIFAR-IO, 10, we use a batch size of 512 and 1024 for the first and the second stage, respectively. On ImageNet with EDM2-S architecture, use batch size of 2048 and 1024 for the first and the second stage, respectively. For EDM2-XL, to compute, we initialize the truncated training stage with the pre-trained checkpoint from the ECM work (Geng et al. 2024) that performs the standard consistency training, and conduct the second-stage training with batch size of 1024. Please see Appendix for more training details.

---

Published as a conference paper at ICLR 2025
0.4
3.5
1.5
1.7-
3.0
2.4
6
dF.D:
1.
Stape 1
1.0
12
300000 400000
300000
400000
300000
400000
300000 400800
Iteration
lteration
Iteration
Iteration
t- 1.5
= 4.0
t- 16.0
80.0
2.5
1.9
Stape
Stsee
1.8
Stage
Stage
Stage
Stage
2.8
-a- Stage
Shae?
2.3
2.5
300000 400000
300000
400000
300000
400000
300000
400000
Figure 3: Denoising FID (dFID) for continuation of standard consistency training at later iterations (Stage 1) and TCM model (Stage 2) at various ts on CIFAR- 10 during the course of training. For TCM, we set the dividing time 1' 1. We can see, in the second stage, dFID exhibits dramatic increase at times below the dividing time L', while the dFID at times above l' and FII) at earlier times is significantly faster compared to standard consistency training, suggesting more efficient "Forgetting" the denoising tasks.
continue improve. Notably, the rate of dFID the truncated stage increase
4.2 TRUNCATED TRAINING ALLOCATES CAPACITY TOWARD GENERATION
Our proposed TCM aims explictly reallocate network capacity towards generation by de-emphasizing denoising tasks smaller t's. Empirical analysis in Fig. 3g. further characterizes this behavior, showing rapid increase in dFIDs at smaller t's below the threshold t' during the truncated training stage. Conversely, dFIDs continue to decrease at larger ('s. In addition, TCM exhibit more pronounced "forgetting" of the denoising task compared to consistency 2 at earlier times. For instance, dFID at t 0.2 increases up to 3.5 in the truncated training, whereas it remains below 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.
Table 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 (Karaa et al
2022
35
1.97
55.7
PFGM++ X0 &
2025
55.7
DDPM
Ho cl
1000
3.17
35.7
LSGM Vanial
2.10
475
training (Fig.
Consistency models I-step
ICT (Son & Dhariwal 2023
2.83
56.4
iCT-deep SSogg & Dianwal 202.3
2.51
112
CTM Kini 2023
5.19
55.7
BCM Geng et 2024
3.60
55.7
TCM (Olr))
2-step
iCT (Song Dhariwal
56.4
iCT deep (Song & Dhanwal 2023
2.24
112
ECM (Geng et al 2024 TCM (ow)s)
2.11
$5.7
55.7
score distillation DMD Yin al. 2024b
Fig. [b) demonstrates how this reallocation of network capacity directly translates im- proved sample quality and training stability. For CIFAR-10/ ImageNet 64 x 64, the truncated training stage (Stage 2) is initialized from the Stage model at 250K 150K iterations, respec- tively. We
3.77
S5.7
Insf-Iii | J0n et al
1024
$5.7
SiD Zhou
1.92
55.7
Knowledge distillatonn
KD [ Luhmman Luhunu
202
9.36
35.7
DSN (2Ner al 41
3.78
65.8
TRACT [Brthhloo
3.78
55.7
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 improve FID from 2.83 2.46. showcasing its robustness (See Figure 7 for more analysis).
3.32
55.7
Salimans Ho
4.51
60.0
4.3 TCM IMPROVES THE SAMPLE QUALITY CONSISTENCY MODELS
7

---

Published as a conference paper at ICLR 2025
To demonstrate the effectiveness of TCM,
Table 2: FID, NFE and # param. on ImageNet 64 64. Dotted lines separate results by # param. Bold indicates the best result for cach category and NFE.
we compare our method with three lines of works that distill diffusion models to one
or two steps: (i) consistency models (Song
Melhod
NFE
FID
# param. (M)
& Dhariwal, 2023) Kim et al. 2023 Geng
et al. 2024) that distills the PF ODE mapping in a simulation-free manner, (ii) that performs distributional matching by
Diffusion models
EDM2-S (Kama et al
2024
63
1.58
280
variational score distillation (Yin et al.
EDM2 xL. Karaa al
1.33
1119
2024b Luo et al., 2024] Zhou et al., 2024)
Consistency models
I-step
utilizing the score of pre-trained diffusion
ICT Song & Dhariwal 2023
4.02
296
models; (iii) knowledge distillation (Luh-
iCT (Song Dharrwal 2023
3.25
592
BCM Geng el al 2024 (2022
4.05
280
man & Luhman 2021 Zheng et al. 2022a
(owrs RDMZ3)
2.88
280
Berthelot et al.) 2223 [Salimans & Ho 2022) 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.
MultiStep-CD (Heek et
2024
3.20
1200
ECM (Geng ot a. 2024 DDM--L)
2.49
1119
TCM (ours, EDM2-XLJ
2.20
1119
2-step
iCT (Song Dhariwal 2023
3.20
296
iCTFddep (Song de Dharaaal 2023
2
2.77
592
(Gnn et
2.79
TCM (owrs; EDM
2
2.31
280
Results. In Table D and Table 2 we re-
port the sample quality measured by FID and the sampling speed measured by the number of function evaluations (NFE), on CIFAR-10 10 and ImageNet-64 x46, respec- tively. We mostly borrow the baseline re- sults from the original papers. We also include the number model parameters. Our main findings are: (1) TCM signif- icantly outperforms improved Consis- tency Training (iCT) (Song & Dhariwal 2023), the state-of-the-art consistency model, across datasets, number of steps and network sizes. For example, TCM im- proves the one-step FID from 2.83/ 4.02 in iCT to 2.46 / 2.88, on CIFAR-10 ImageNet. Further, TCM's 's one-step FID even rivals iCT's two-step FID on both datasets. When using HDM2-S model, TCM also surpasses iCT-deep, which uses 2x deeper networks, in both one-step (2.88 vs 3.25) and two-step FIDs (2.31 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.) 2022), 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 CIFAR-10, which is close 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 G for more samples.
MultiSter CD (llekk
1.90
FCM (Geng in 2014 (2D122-XX) -X.)
2
1.67
1119
TCM (ours EDM2-X))
1.62
1119
Variational score distillation DMD2
et 2024g
2.60
296
Diff' Instract [Luo et i
5.57
296
EMD-16 (Xie ct 2024)
2.20
296
Moment Matching
2024
3.00
400
3.86
400
SiD Zhou al
2024
1.52
296
Knowledge distillation
DSNO (Zheng al 2022
7.83
329
TRACT Herthllx et
7.43
296
4.97
296
PD (Salimans & Ho
2022
1
15.4
296
296
4.4 ANALYSES OF DESIGN CHOICES
Time sampling distribution 5. We explore various time sampling distributions 3. supported on [t', 7), and find that the truncated log-Student-f distribution works best (i.e., In(t) follows Student-r distribution). The Student-f distribution, being heavier-tailed than the Gaussian distribution employed in previous consistency training (Song & Dhariwal 2023; Geng et al. 2024), 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 freedom & effectively controls the thickness of the tail, with the Student-t distribution converging Figure 4a shows the shape of eh with varying standard deviation o and the degree of freedom v in three cases: (1) heavy-tailed and low probability mass around small i's (a 2, 10000),
Gaussian distribution as v co.
8

---

Published as a conference paper at ICLR 2025
Table CIFAR-IO FID when varying the dividing time
Table 4: CIFAR-10 FID for different training stages.
l' value
0.17
0.8
1.0
1.5
Stage
Stage
Stage 3
FID
2.70
2.69
2.56
2.79
FID
2.77
2.46
2.46
(2) heavy-tailed and a high probability mass around small t's (a = 0.2, v = 0.01). (2) light-tailed (a 0.2, // 2). From Fig. 4bb observe that the log- Student-t distribution with a = 0.2, 0.01 is the best among the three. Hence we use o 0.2, v = 0.01 in all the experiments.
3.6
3500
0-0.2, u=0.01
0-02, V
0.0. Ro
Frequency
3000
2500
a=2, y=10000
a=0.2, p=2
3.4
a=0.2, v=0.01
d-2, p010000
3.4
3.2
p-025 a D1
p-025, 8 2.01
3.2
p-025, 411
2000
0 3.0
1500
3.0
2.8
1000
2.8
500
2.6
2.6
2.4
3
4
In(t)
Iterations (x10F)
Iterations (x109)
(a)
(b)
(c)
Figure 4: (a) Comparison of Student-I distributions with different standard deviations and degree
of freedom v. (b) FID evolution on CIFAR-10 for different a and v. Wp = -011, p = 0.25, t' = 1, and
batch size of 128 are used for all plots. (c) Effect of p and wy on the FID on CIFAR-10. We use batch size of 128. t' is set to 1.
Strength for boundary loss. Figure 4c 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 wide range of p and wy (from 0.1 to 0.5 for and from 0.1 to for wy). How- ever, when using a very small weight for the boundary loss (w,, = 0.01). FID explodes as the model fails to maintain the boundary condition. Thus, we use p 0.25, wh 0.1 in all the experiments.
Dividing time t'. The boundary t' ideally represents the point where the task in the PF ODE transitions from denoising generation. However, this transition is gradual, and there is no single definitive point. Fig. b suggests this transition occurs roughly between t' = 0.8 and t' = 1.5, where we observe 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' values around this range. Table 3 shows that t' 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? natural question whether we can extend our two-stage training pro- cedures to three or more stages by gradually increasing f'. However, recall that our methodology was motivated by the fact that in the first-stage training (standard consistency training), we observe increas- ing dFIDs at smaller i values of the training range, seen in Fig. 2 This trade-off is notably absent the second stage over the time range [t', T.. as seen in Fig. This suggests that during the second stage, the model tackles tasks that more or less similar to generation. and introducing another trun- cated training stage may not yield further gains. In support of this hypothesis, we implement the third stage where the dividing time is L' = 4, but do not observe improvement, as shown in Table 4 We also consider adding an intermediate training stage between stage and stage 2 that finetunes To, in the time range (0, t') but produces slightly worse performance, which we discuss in Appendix C
RELATED WORK
Consistency models. Song et al.) (2023) first proposed consistency models a new class of generative models that synthesize samples with a single network evaluation. Later, [Song & Dhariwal (2023): Geng et al. (2024) presented a set of improved techniques to train consistency models for better sample quality. Luo et al. (2023) introduced latent consistency models (LCM) to accelerate the sampling of latent diffusion models. Kim et al. (2023) 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

---

Published as a conference paper at ICLR 2025
1-step
2-step
CIFAR-10
ImageNet 64x64
Figure 5: Uncurated one-step (top) and two-step (bottom) generated samples from TCM (EDM) on CIFAR-10 10 and TCM (EDM2-XL.) on ImageNet 64x64, respectively.
than standard consistency models that only care about the mapping from intermediate points to the data endpoints. Heeek et al. (2024) 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. Similar direction is also explored in Wang al. (2024). |Ren et al. (2024) combined CTM with progressive distillation (Salimans & Ho] 2022), by performing segment- 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.) 2014) to achieve good performance.
Fast sampling of diffusion models. While line of work aims to accelerate diffusion models via fast numerical solvers for the PF-ODE (Lu et al., 2022 Karras et al. 2022; Liu et al. [022a]] Xu et al. 2023a), 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 dataset of noise and image pairs, and then applied it to train single-step generator (Luhman &
Luhman 2021 Zheng et al. 2022a). Progressive distillation (Salimans & Ho 2022 Meng et al. offline dataset.
2023) iteratively halves the number of sampling steps required, without needing
(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., 2024bb a Xie et al. 2024), implicit score matching (Zhou et al. 2024) or moment matching (Salimans et al. 2024). (iii) Adversarial distillation, which leverages the advessaiaal training fine-tune pre- pre-trained diffusion models into a
few-step generator (Sauer et al. 2023 2024 Lin et al. 2024 Xu et al. (024)). Compared with these
training-based diffusion acceleration methods, our method is most memory and computation efficient.
Truncated training of diffusion models. Balaji et al. (2022) propose to train different diffusion models for each time step range. Since consistency models solve a more difficult task (learning to integrate PF-ODE) than diffusion models (learning the drift of PF-ODE). they can benefit more from such a strategy but also require a specific parameterization (Eq. (5) to satisfy the boundary condition. Zheng et al. (2022b) use GANs to generate the noised data and use diffusion models to map them to clean data. Different from ours, they train diffusion models on the first half of the interval.
CONCLUSION
We have introduced 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 -tate-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 even smaller network architectures than baselines.
10

---

Published as a conference paper at ICLR 2025
7 REPRODUCIBILITY STATEMENT
We provide sufficient details for reproducing our method in the main paper and also in the Appendix, F 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', boundary loss weight wh. and boundary ratio p. 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.
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 work. Other ethical concerns may arise regarding the significant resource costs required 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 with only NFE and requires smaller neural network architectures, both of may which reduce inference- time costs relative to other iiffusion-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.
REFERENCES
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, et al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01333 2022.
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.044 04.00444044.... 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. leee, 2009.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. arXiv preprint arXiv:2406. 14548, 2024.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Jonathan Heek, Emiel Hoogehoom, and Tim Salimans. Multistep consistency models. arXiv preprint 20020:2403.06807 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. Classifier- free diffusion guidance. arXiv preprint arXiv:2207.12595 2022.
11

---

Published as a conference paper at ICLR 2025
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in 2020.
Neural Information Processing Systems,
Jonathan Ho. Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458, 2022.
Rongjie Huang, Jiawei Huang, Dongchao Yang, Yi Ren, Luping Liu, Mingze Li, Zhenhui Ye, Jinglin Liu, Xiang Yin, and Zhou Zhao. Make-an-audio: Text-to-audio generation with prompt-enhanced diffusion models. arXiv preprint arXiv:2301. /2661, 2023.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv preprint 221.:22000003664, 2022.
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174-24184, 2024.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Ue- saka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279, 2023.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Toronto, ON, Canada, 2009.
Shanchuan Lin, Anran Wang, and Xiao Yang. Saxl-lightning: Progressive adversarial diffusion distillation. arXiv preprint arXiv:2402.13929 2024.
Tsung- Yi Lin, Michael Maire, Serge Belongie. James Hays, Pictro Perona, Deva Ramanan, Piotr Dollár, and Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision- ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings. Part V 13, pp. 740-755. Springer, 2014.
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models manifolds. International Conference on Learning Representations, 2022a.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022b.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. arXiv preprint (:Xii:2309006880 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.00 00927, 2022.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.00 023388, 2021.
Simian Luo, Yigin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models:
Synthesizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04378, 2023.
Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diff- instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Chenlin Meng. Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation guided difffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.00 06125, 2022.
12

---

Published as a conference paper at ICLR 2025
Yuxi Ren, Xin Xia, Yanzuo Lu, Jiacheng Zhang, Jie Wu, Pan Xie, Xing Wang. and Xuefeng Xiao. Hyper-sd: Trajectory segmented consistency model for efficient image synthesis. arXiv preprint NrXi::2404113686 2024.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.0051:1 2022.
Tim Salimans, Thomas Mensink, Jonathan Heek, and Emiel Hoogeboom. Multistep distillation of diffusion models via moment matching. arXiv preprint wrXiv:2406.04103, 2024.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint arXiv:2311.10042 2023.
Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rombach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. arXiv preprint arXiv:2403. 12015, 2024.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint 10..v:2210.14188: 2023.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
Yang Song, Prafulla Dhariwal, Mark Chen, and llya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 14:11287-11348 11302, 2021.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- lion, 23(7):1661-1674, 2011.
Fu- Yun Wang. Zhaoyang Huang, Alexander William Bergman, Dazhong Shen, Peng Gao, Michael Lingelbach, Keqiang Sun, Weikang Bian, Guanglu Song, Yu Liu, et al. Phased consistency model. arXiv preprini 2405:22005.8400) 2024.
Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, and Ruiqi Gao. Em distillation for one-step diffusion models. arXiv preprint arXiv:2405. 16852, 2024.
Yanwu Xu, Yang Zhao, Zhisheng Xiao, and Tingbo Hou. Ufogen: You forward once large scale text-to-image generation via diffusion gans. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 8196-8206, 2024.
Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. Advances in Neural Information Processing Systems, 36.76806-76838, 2023a.
Yilun Xu, Ziming Liu, Yonglong Tian, Shangyuan Tong. Max Tegmark, and T. Jaakkola. Pfgm++: Unlocking the potential of physies-inspired generative models. In International Conference on Machine Learning, 2023b.
Tianwei Yin, Michael Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. arXiv preprint 4rXiv:240.114667 2024a.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEFF/CVF Conference Computer Vision and Pattern Recognition, pp. 6613-6623, 2024b.
13

---

Published as a conference paper at ICLR 2025
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. arXiv preprint arXiv:2211.13449, 2022a.
Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Truncated diffusion probabilistic models and diffusion-based adversarial auto-encoders. arXiv preprint arXiv:220.00667), 2022b,
Mingyuan Zhou, Huangjie Zheng. Zhendong Wang. Mingzhang Yin, and Hai Huang. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Forty-first International Conference on Machine Learning. 2024.
14

---

Published as a conference paper at ICLR 2025
Stage
Stage
Figure 6: One-step text-to-image generation results of the standard consistency model (stage 1) and TCM (stage 2).
A LIMITATION
TCM introduces an additional training stage on top of the standard consistency model training. Compared to the standard consistency training, the truncated training requires slight increase in per-iteration training time due to the additional boundary loss in Eq. (1O). Standard consistency training necessitates two forward passes per training iteration, while our parameterization (Eq, (5) requires three. Also, the truncated training 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 x 64 with EDM2-S, TCMs have an 18% increase in training time per iteration and an 15% increase in memory cost. Moreover, the one-step sample quality of TCM still has considerable performance gap from the diffusion models with a large NFE, but we believe this an important step toward closing the gap.
15

---

Published as a conference paper at ICLR 2025
B TEXT-TO-IMAGE RESULTS
Table 5: Zero-shot FID scores on MSCOCO dataset measured with 30k generated samples.
Stage
Stage
FID
18.32
15.58
To show the scalability of our method, we train TCM on COYO dataset using consistency distillation with fixed classifier-free guidance (Ho Salimans| 2022) scale We initialize our models with stable diffusion (Rombach et al. (022) 1.5. We use batch size of 512 for quick validation, though using a larger batch size (2 1, 024) is standard (Liu et al.) 2023 Yin al. 2024a) and would lead to better generative performance. For the first stage, train for 80,000 iterations (after which FID starts to increase), and in the second stage, we additionally train for another 200,000 iterations. We provide visual comparison between the standard consistency model and TCM in Fig. 6 Captions used are: "A photo of an astronaut riding a horse on Mars", "Robot serving dinner, metallic textures, futuristic atmosphere, high-tech kitchen, elegant plating, intricate details, high quality, misc-architectural style, warm and inviting lighting", and "A photo of a dog" for each row.
We also measure the FID on MSCOCO dataset (Lin al. 2014) in Table. achieves a better FID than the standard consistency model (the first stage).
We see that TCM
C ADDITIONAL EXPERIMENTS
Fig. 7 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 f.
200
175
Stage 2
Gradient Norm
Stage
125
100
75
50
25
50000
100000 150000 200000 250000 300000
Iterations
Figure 7: Gradient norm evolution during the first and second stage training on ImageNet 64 x 64 (corresponding to Fig. [b)) The red circles indicate where the gradient norms are larger than 100. Stage training blows up after the last few gradient spikes. It shows that the truncated consistency training more stable than the standard consistency training.
Fig. shows the dFID, evolution during the standard consistency training. We see that dFIDs at larger i's start from larger values and converges more slowly.
Adding an intermedate training stage In our parameterization Eq. (5) we only use the pre-trained model fe, in [0, t'). One may wonder if we fine-tune fo, the truncated time range [0, t') to provide better boundary condition for the truncated training. We find that although doing so improved the dFID,, of fo, from 1.51 to 1.43, it led to worse final FID of 227 for the truncated consistency model, regardless of whether we initialized fo 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') makes the model fo, forget about how the early mappings properly propagate to the later mappings the range of [t', T). This may hinder the learnability of its mapping at the boundary time, making it harder for fo to transfer the knowledge learned in fg, to its generation capability.
16

---

Published as a conference paper at ICLR 2025
30
Iter 80000
25
Iter 160000
Iter 200000
20
15
Iter 280000
Iter 320000
10
360000
10a
10
10
Figure 8: Evolution of the denoising FIDs (dFID,) at different times t's during standard consistency training for different iterations. For (1, 10), dFID, has different convergence speeds while in both small times (t < 1 and large 's (t > 10), dFID, converges with more similar speed.
D BACKGROUND ON CONSISTENCY MODELS
Most of this part has been introduced by previous works (Song al. 2023 Song & Dhariwal 2023). Here, we introduce the background consistency models, in particular the relationship between consistency training and consistency distillation, for completeness.
D.1 DERNITIIO OF CONSISTENCY FUNCTION
D.1.1 PROBABILITY FLOW ODE
The probability flow ODE (PF ODE) of Karras et al. (2022) is as follows:
dx, (x:
(12)
where (x(x,) the score function time f E [0, 7). To draw samples from the data distribution Pdata- initialize xT with sample from II) and solve the PF ODE backward in time. The solution Xg
Xy [ (-ts, (x,)) is distributed according to Pdata
D.1.2 CONSISTENCY FUNCTION
Integrating the PF ODE using numerical solvers computationally expensive. Consistency function instead directly outputs the solution of the PF ODE starting from any t E [0, The consistency function f satisfies the following two properties:
f(xo, 0) X0-
2. f(x,,t) = f(x., v(s,t) [0,7]3
The first condition be trivially satisfied by setting f(x, t) Cour (t) + (tkpp(t) where (0) and Cskip = following EDM (Karras et al. 2022). The second condition can be satisfied by optimizing the following objective:
min Ex,xx. [d(f(xe,,) t), f(x., s))]
(13)
f
where d a function satisfying:
1. y(x.y) = - x x
2. d(x.y) > O.
3. Od(x.y)
dy y
àd 4. afp and öy
well-defined and bounded.
17

---

Published as a conference paper at ICLR 2025
D.2
CONSISTENCY DISTILLATION
D.2.1 OBJECTIVE
In practice, Song et al. 2023 consider the following objective instead:
min [ [x( t), fo (x--2,, t -,)),,
(14)
0
where we parameterize the consistency function f with a neural network fa. and 0 < A, < L. Here, fo- is the identical network with stop gradients applied and is called teacher. Since A, 0, the teacher always receives the less noisy input, and the student fa is trained to mimic the teacher. Optimizing Eq. (14) requires computing Xt, which we can be approximated using one step of Euler's s solver:
X,- =x; + J. (-uss(xx))dd se X, + (x,)
(15)
When (xt) approximated by a pre-trained score network, Eq. (4 becomes the consistency distillation objective in Song et al. (2023). If A. is sufficieintly small, the approximation in Eq. (15)
quite accurate, making LCD a good approximation of Eq. (4).. The precision of the approximation depends on A and also the trajectory curvature of the PF ODE.
D.2.2 GRADIENT WHEN A. +0
Let us rewrite Eq. 14 as follows:
Exx, [x(fo(X,, t), fo- (x4-2,, A,))
(16)
= Exx, (x(f t), to (x,, t)
fo- -XI- t A.) fo fx(,,.)) t))]
(17)
Ay
Od 8y
32d
Bxx, [d( (), Co ()) +
Ay + (Ay))
Ay + 0(1lyy])))) [3]]
(18)
2
oy²
マ xx [( T 32d
1
Ay + 0(1lay))))
|3]
(19)
2
δy²
where we define Ay fg- (x:-a.: Ae) fo (xx, t).
Let's take the derivative with respect to 0:
a
Fr.xI[ffo (XI- L - Ae) - (x,, 1))² Odd
(fo (x:- t. A.) (xo(x ()) +
+(1)))))))))))))
2 00
0y²
(20)
jdd
(fg (XI-A.; - Ai) - fo t)) af.
E Exxe Oy²
+ +(1ly []]
30
(21)
As
Ofg 0x,
Ogg
fo- (x,- A,, - A,) = (x(,))
A,
A, 0(())
(22)
0x: Øt
Øt
we have
(x t- - A.) fg- (xx,)) = - OEg
0x:
Oig
+
)4, + 0(43).
(23)
0x,
Öt
at
Since fo (x,,t) has the same value as (x,, t), we can plug this into Eq. (1)
Odd (fo- (Xi- - A.) (x,, 1)) Hig
Ex,x,
+ 0(1lA)13)]
(24)
gy²
00
Odd Hr- Ox.
= = 0y² txi ōt
A, + 0(2 + 0(|lA)|) [1]]
(25)
Øt
00
Idd
dx.
0g-
of,
- 0y²
A + 0(A?))
(26)
Oxi
Øt
Øt
00
18

---

Published as a conference paper at ICLR 2025
As the gradient is O(A,), it becomes when A. - 0, so cannot be used for training. To make the gradient non-zero, Song et al. (2023) divide the by A. Then, we have
0 00 (c.
θ
1
E d(fo
((fo(x
t), fo (XI- Ad))
(27)
00
D
Od
Oc
8x.
H6-- Hh
-
+ 0(4.)
(28)
0y²
ox,
0t
Öt
00
0d
0x,
0fg Hh
E3 0y
(29)
0x,
0l
01
00
0dd Or.
O.g H6,
= =x,x
(-1 sx(x)))
(30)
0y² 8xt
Øt 00
Idd No
0.. Hio
= Ex,xx
(1 sx(x)))
(31)
0y² 0xt
Øt 00
in the limit of A, - O.
Hessian of d. Here, we provide the Hessians of the L2 squared loss and the Pseudo- Huber loss.
1. y) = || - y
ddd By
y====== 2I.
0.d 2. If - 2 c, 0y² y=x II
D.3 CONSISTENCY TRAINING
Song et al. (2023) show that Eq. (31 can be estimated without a pre-trained score network. From Tweedie's formula, we express the score function as
E [x(x(x,))
[x] Xx
8y (x.)
(32)
12
Plugging this into Eq.
31). we have
Fd Ho (t. sxx)))
σ6g Hfg
8dd
Ofa- Ep(x(x.) x.(xx [x]
Ofg Hfg
Ex,x, 0y² 0x.
Exxx, 0y²
Ot 00
0x.
t
0t
00
(33)
Ex [pxx)x)) )l
OPd Ofg
Xt.
0fg
Ofg
11
Oy2 0x:
t
00
(34)
Od Or.
Xt
0f
Ofg
x,x
0y2 ix;
t
00
(35)
where we now have the expectation over three random variables t, x, x, and do not require a score function. In the next section, we will reverse-engineer an objective such that its gradient matches Eq. (35)
D.3.1
OBJECTIVE
turns out that the following objective the one we are looking for:
1 A.
(e (0 Exx.
d(fo(x te, fo (x A.ee,t - A.)),,
(36)
where E ~ Ao, I) is a random noise vector. The objective in Eq. (36 is called the consistency
training objective. We show that the gradient of Let indeed matches Eq. (35) in the limit of
19

---

Published as a conference paper at ICLR 2025
A O. First, apply the Taylor expansion to the unweighted loss in Eq. (36):
Exx. _ex.e[((fo(x) [d(ff(x
+ te, t), fa(x x + te,t) + to (x+ (t A.)E,t A.) to + te,
(37)
Od
Ay + (Ay)² 7 0dd
= -xxx,el((fo(x + t), fx (x+ tc, + δy
Ay 0yl
(38)
δy²
= Exx (EIxx,((A9)) (4y) T Odd
Ay 0(11
(39)
dy²
where we define Ay as Ay = fo (x + (t Ad),, (1) - -fo(x t)
Let's take the derivative with respect to A:
i E, ( 0 Ex,, Od (fo (x + (t- Ai)c,t D.) - ff((x t)) Uf.
+ +0 Ayll)]
00
0y²
00 (40)
Using the Taylor expansion, we have
Hig
Ofg
(x + (t D.)E, De) fo- (x te,t)
cA,
A. + +044),
(41)
0x
Hfg
8f 0-
fg (x + (t De)e,t - A)) fg- (x + te,t)
cA.
A + 0(A))
(42)
0x
Ōt
Since fx te, t) has the same value as fo- (x + te, f), we can plug this into Eq. 440
Odd
Ofg
Bx,,,c Øy² (fo (x + (1 D,)E,t D.)
fo(x + te,
0(11Ay|)))
(43)
00
8dd
Ofg
Ofg
Ufg
A, - 0(4))
(44)
0y²
0x
00
E, x
8
O..
X,
0rg
0fg
A 00A)),
(45)
0y2
Ox
Øt
00
where in Eq. 455 we use the reparametrization trick x + te x, and x1=x Finally, we can show that Eq. 445 matches Eq. (35) in the limit of A, + 0 and after dividing by A,:
Cc 0
lim 2, 00 CCt (C,
(46)
d0
We can add weighting function w()) without affecting this equality, leading the objective in Eq. 44p.
E
TRAINING OBJECTIVE OF TCMS
By substituting Eq. into Eq. (6). we have:
w(l)
LCT (frune plune 0,0 -00 = >
d(fo(x + ('E, (x+ (t' Dr)e, - -Dr)
(47)
+(1 -d6 [ vi(11)
w(t)
(xfo(x (e, (x (1 -i)e,,, A.)dl
(48)
IES,
A.
+(1 Xb) / (i(t) w(t)
d(fo te, te, fo- (x
A.)c,t A,)dt
(49)
JEST
A.
The first two terms in RHS represent the boundary loss, and the last term is the consistency loss. Let us define A, = (1 + 8 (omnoid(-t))( r)t. We define tm to be the smallest point such that
I'm A.,, f'. Then the volume of the set Sp (1 IR &' < t l' +.) is A.... We assume
20

---

Published as a conference paper at ICLR 2025
de is properly designed to be upper bounded by finite value (see Sec. 4.4). In the limit of - 1,
we simplify the the second term the above:
(1 - An) / Va(1) w(t) d(fo + (e,1), for (x + -,)e, A.)dl
LES
A.
(50)
-11 - Ao)
d(fo ( ('',)) d(()))))((((((()/))))))) (x- Av) + O(1)dl
1ES
Dr
(51)
=1 xe)) ((t')
([o + ('e,,()) (x (1' Ar)E, -n)e,l(AA\Volsss + / O(1)d
JIES,
A
(52)
~ (1 - (1) w(t')
dl (x - (x= (1' (e)) - O
De
JES
(53)
= (1 (()
(x+ (t' - Ar),
(54)
where in Eq. 51 we apply the Taylor expansion to the integrand, and in Eq. 53 we can see that At... (1 sigmoid( (oad((mm)))) (m)) - fm goes to zero as - Thus, im - l' and then in the limit. Hence, the boundary loss is (148 (1)(111)
Vol(Sv) Dr A =
(1+8 (i+8-sigmoidd tm (1-r)
w(l') +1 1o) (( +'e, 50 (x (t' - -u))et - Av)
(Ab De
(55)
w(t Dr
No d(fo L'e, fo (x + +f' - Dc)ett' - Ar),
(56)
where the first term (0(1/At) dominates the second term (O(1)). We hence arrive at Eq. 99
w0(() d(fo(x + t,,,t)) (.), foo (x (t' Dr)c,t Av))
(f ptrone
0,0.
.0- Dr
d
(57)
Boundary loss: Los:=< (fo.Fo
w(t) d(fo + t), fo (x + Dr)e, - A.)..
+(1 - E,
(58)
A,
Consistency loss: (Fo.
where because A, + 0, we can approximate Sp = (t', TJ, leading to Eq. 58
In practice, we close to one during the truncated training stage (see Section Note that here, the boundary loss can dominate the consistency loss when fo and fø, are sufficiently different around t = t'. However, in practice, as we set A, to be small enough but not all the way to zero, and as we use Pseudo-Hull Huber loss (see Section F with small value that normalizes the effect of the loss magnitude on the gradient norm, we can balance the training.
F IMPLEMENTATION DETAILS
We provide detailed information about our implementation in the following.
Model initialization and architecture: All stage I models are initialized from pre-trained EDM
EDM2 checkpoints as suggested by Geng et al. (2024). For CIFAR-10, 10, we use EDM's DDPM+ architecture, which is slightly smaller than iCT's NCSN++. For ImageNet 64 64, employ EDM2- (280M parameters) and EDM2-XL (approximately 1.IB parameters) architectures. EDM2-S is slightly smaller than iCT's ADM architecture (296M parameters).
21

---

Published as a conference paper at ICLR 2025
Model parameterization:
Following EDM, we parameterize consistency models fo as fo and data 0.5.
Coun (f))F t) + Ca),(()),, where Cour (t)
taen
Cskip (t)
Gún +
Vo +
Training details: We set D. (1 + sigmoid(- (t))(L r)t, where r max
1 1/2ft/250001 0.999) for CIFAR-10 and max [1 1 /4T1/250001 .0.9961 for ImageNet 64 x 64, with i being the training iteration. For CIFAR-10. 10. we train for 250K iterations in Stage and 200K iterations Stage 2. For ImageNet 64 64, EDM2-S is trained for 150K iterations in Stage 1 and 120K iterations in Stage 2, while EDM2-XI. is trained for 40K iterations in Stage 2 only. See Fig. [ for the FID evolution during training. For the second stage, we start with the maximum values (i.e.. 0.999 or 0.9961) and do not change them. The weighting function (t) is set to for CIFAR- 10 and Au/Cout (t)² for ImageNet 64 x 64. As suggested by Song Dhariwal (2023); Geng et al. 2024). ). we
use the Pseudo-Huber loss function d(x,y) and c 0.06 for ImageNet 64 x 64. This is especially crucial for our method the boundary loss can dominate the consistency loss. The boundary loss compares the outputs from the different model fo and fe.. 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 x 64, we employ mixed-precision training with dynamic loss scaling and use power function EMA (Karras et al. 2024) with y 6.94 (without post-hoc EMA search).
V|Ix y11 +
with c= le 8 for CIFAR-10
Learning rate schedules: EDM2 (Karras et al. 2024) architectures require a manual decay of the learning rate. Karras et al. 2024)) suggest using the inverse square root schedule
the first stage training EDM2-S on ImageNet, we use fee 2000 and Ose /max(t/(.or.] For
Orae
3 following
Geng et al. 2024). For the second stage training of EDM2-S, we use ee = 8000 and o 5c 4.
Second stage training of EDM2-XL initialized with the ECM2-XL checkpoint from Geng et al. (2024). During the second stage, we use leef = 8000 and Orer le 4 for EDM2-XL.
Time step sampling: For the first stage training, we use a log-normal distribution for th. For CIFAR-10. use mean of -1.1 and a standard deviation of 2.0 following SSng & Dhariwal (2023). For ImageNet, we use a mean of -0.8 and a standard deviation 1.6 following Geng et al. (2024).
For EDM2-XL, we also explore t' = 1.5 for truncated training, adjusting 2 to ensure has high probability mass around = 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 f = 80, 1.526 for ImageNet.
G UNCURATED GENERATED SAMPLES
We provide the uncurated generated samples in Fig.
22

---

Published as a conference paper at ICLR 2025
(a) step samples.
(b) Two-step samples.
Figure 9: Uncurated one-step and two-step samples on CIFAR-10.
(a) One-ster samples.
(b) Two-step samples.
Figure 10: Uncurated one-step and two-step samples on ImageNet (EDM2-S).
23

---

Published as a conference paper at ICLR 2025
(a) One-step samples.
(b) Two-step samples.
Figure 11: Uncurated one-step and two-step samples on ImageNet (EDM2-XL).
24


Reviews:
[{"review_id": "RMZz7lgaAi", "reviewer": "Reviewer_uC7B", "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.", "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?\n- 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?\n\n(Please refer to Weakness section for detailed explanations)", "limitations": "", "rating": 8, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and is easy to follow. The math derivations are technically correct.\n- 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.\n- The empirical results are strong, showing the effectiveness of the method.", "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:\n\n1. two-stage truncated training to avoid the overtraining of denoising tasks;\n2. changing proposal distribution to focus more on the boundary conditions.\n\nHowever, 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.\n\nTherefore, 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:\n\n- 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?\n\nBesides, there is another minor weakness but I think it does not affect the conclusions in this paper:\n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "c2vvACrmEW", "reviewer": "Reviewer_ra9j", "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.", "questions": "- Could the authors discuss the relation with PCM [1]?\n- 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?\n- 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?\n\n\n[1] Phased Consistency Model", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n- The phenomenon that the consistency training gradually weakens the model's denoising capabilities at small $t$ is well illustrated in Figure 2.\n- The method achieves better FID on standard image datasets than previous consistency models, especially in 1-step generation.\n- Ablation studies for the dividing time and the number of stages are comprehensive.", "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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2CgW1MXIWK", "reviewer": "Reviewer_jsgK", "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.", "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?\n\n\n2. 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?\n\n\n\n[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.\n\n[2] Yang, Z., Feng, R., Zhang, H., Shen, Y., Zhu, K., Huang, L., ... & Cheng, F. (2023). Eliminating lipschitz singularities in diffusion models.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 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.", "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'$.\n\n2. 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.\n\n3. The paper lacks ablation studies for certain hyperparameters relevant to the proposed method, such as $N_B$ (or $\\rho$), $\\Delta_t$, $\\Delta_{t'}$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3F1CduXsVc", "reviewer": "Reviewer_2JKb", "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.", "questions": "Besides the weaknesses listed above, I also have some questions.\n- 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?\n- 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?\n- 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.)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper largely focus on an orthogonal problem of network capacity, which can be possibly combined with many consistency distillation works.\n- 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.", "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.\n- The proposed method seems to still perform worse than SiD.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z487BBy1E2", "reviewer": "Reviewer_NmLx", "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.", "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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "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.\n\nThe 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.\n\nTo 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).\n\nAmbrogioni. The statistical thermodynamics of generative diffusion models: Phase transitions, symmetry breaking and critical instability. arXiv:2310.17467, 2023.\\\nBalaji et al. eDiff-I: Text-to-Image Diffusion Models with Ensemble of Expert Denoisers. arXiv:2211.01324, 2022.\\\nZheng et al. Truncated Diffusion Probabilistic Models and Diffusion-based Adversarial Auto-Encoders. ICLR 2023.\\\nBiroli et al. Dynamical Regimes of Diffusion Models. arXiv:2402.18491, 2024.\\\nHeek et al. Multistep Consistency Models. arXiv:2403.06807, 2024.\\\nSclocchi et al. A Phase Transition in Diffusion Models Reveals the Hierarchical Nature of Data. arXiv:2402.16991, 2024.", "weaknesses": "While they are not critical and may not significantly change my assessment, I would like the authors to tackle the following weaknesses.\n\n### Unnecessary decomposition of the loss\n\nI 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.\n\n### Additional insights & reproducibility\n\nThe following discussions / content would benefit the paper.\n- The code used by the authors for the presented experiments should be included in the supplementary material for further reproducibility.\n- 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.\n- 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)?\n- 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).\n\n### Minor issues\n\n- 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.\n- The number of optimization steps is not specified in the hyperparameters listed in Section D.\n- In Figure 3, the timestep $t = 1$ should be placed in the second row to be consistent with the rest of the paper.\n- The paper should specify where the baseline results from Tables 1 and 2 are reported from.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)