Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion
Authors: Dongjun Kim
Chieh-Hsin Lai
Wei-Hsiang Liao
Naoki Murata
Yuhta Takida
Toshimitsu Uesaka
Yutong He
Yuki Mitsufuji
Stefano Ermon
Venue: ICLR 2024 poster (2024)
Keywords: Diffusion Models; Score-based Models; Generative Models; Generative AI; Deep Generative Models; Distillation Models
Primary area: generative models
TL;DR: Our model integrates score-based models and distillation models by learning the entire ODE trajectory, resulting in SOTA FID for CIFAR-10 and ImageNet.

Abstract: Consistency Models (CM) (Song et al., 2023) accelerate score-based diffusion model sampling at the cost of sample quality but lack a natural way to trade-off quality for speed. To address this limitation, we propose Consistency Trajectory Model (CTM), a generalization encompassing CM and score-based models as special cases. CTM trains a single neural network that can -- in a single forward pass -- output scores (i.e., gradients of log-density) and enables unrestricted traversal between any initial and final time along the Probability Flow Ordinary Differential Equation (ODE) in a diffusion process. CTM enables the efficient combination of adversarial training and denoising score matching loss to enhance performance and achieves new state-of-the-art FIDs for single-step diffusion model sampling on CIFAR-10 (FID 1.73) and ImageNet at 64X64 resolution (FID 1.92). CTM also enables a new family of sampling schemes, both deterministic and stochastic, involving long jumps along the ODE solution trajectories. It consistently improves sample quality as computational budgets increase, avoiding the degradation seen in CM. Furthermore, unlike CM, CTM's access to the score function can streamline the adoption of established controllable/conditional generation methods from the diffusion community. This access also enables the computation of likelihood. The code is available at https://github.com/sony/ctm.

Published as a conference paper at ICLR 2024
CONSISTENCY TRAJECTORY MODELS: LEARNING PROBABILITY FLOW ODE TRAJECTORY OF DIFFUSION
Dongjun Kim*t & Chieh-Hsin Lai"
Sony AI
Tokyo, Japan
dongjannstaanfod.cco edu, chi chieh-hsin,.es watasiinnaiieee... wwwllllnllc.oooo com com
Wei-Hsiang Liao & Naoki Murata & Yuhta Takida & Toshimitsu Uesaka
Sony AI
Tokyo. Japan
Yutong Het
Yuki Mitsufuji
Stefano Ermon Stanford University CA, USA
Carnegie Mellon University PA, USA
Sony AI, Sony Group Corporation Tokyo, Japan
ABSTRACT
Consistency Models (CM) (Song al., 2023) accelerate score-based diffusion model sampling at the cost of sample quality but lack a natural way to trade-off quality for speed. To address this limitation, we propose Consistency Trajectory Model (CTM), generalization encompassing CM and score-based models special cases. CTM trains single neural network that can a single forward oass-u scores (i.e., gradients of log- log-density) and enables unrestricted traversal between any initial and final time along the Probability Flow Ordinary Differential Equation (ODE) in diffusion process. CTM enables the efficient combination of adversarial training and denoising score matching loss to enhance performance and achieves new state-of-the-arr FIDs for single-step diffusion model sampling on CIFAR-10 (FID 1.73) and ImageNet at 64 x 64 resolution (FID 1.92). CTM also enables a new family of sampling schemes, both deterministic and stochastic, involving long jumps along the ODE solution trajectories. It consistently improves sample quality computational budgets increase, avoiding the degradation seen CM. Furthermore, unlike CM. CTM's ' access to the score function can streamline the adoption of established controllablecconditiona generation methods from the diffusion community. This access also enables the computation of likelihood. The code is available at https: : / //gitttii tuu.. com/ sony/ / ctm.
INTRODUCTION
Deep generative models encounter distinct training and sampling challenges. Variational Autoencoder (VAE) (Kingma & Welling, 2013) can be trained casily but may suffer from posterior collapse, resulting blurry samples, while Generative Adversarial Network (GAN) (Goodfellow al., 2014) generates high-quality samples but faces training instability. Conversely, Diffussinn Model (DM) (Sohl- Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020b) addresses these issues by learning the score (i.e., gradient of log-density (Song & Ermon, 2019), which can generate high quality samples. However, compared to VAE and GAN excelling at fast sampling, DM involves a gradual denoising process that slows down sampling, requiring numerous model evaluations.
Score-based diffusion models synthesize data by solving the reverse-time (stochastic or deterministic) process corresponding to prescribed forward process that adds noise to the data (Song & Ermon, 2019; Song et al., 2020b). Although advanced numerical solvers (Lu et al., 2022b; Zhang & Chen,
Equal contribution
Work done during intersship SONY AT

---

Published as a conference paper at ICLR 2024
2022) of Stochastic Differential Equations (SDE) Ordinary Differential Equations (ODE) substan-
tially reduce the required Number of Function Evaluations (NFE), further improvements are challeng- ing due to the intrinsic discretization error present in all differential equation solvers (De Bortoli et al., 2021). Recent developments in sample efficiency thus focus on Distillation models (Salimans & Ho, 2021) (Figure 1) that directly estimates the integral along the Probability Flow (PF) ODE sample trajectory, amortizing the computational cost of numerical solvers, exemplified by the Consistency Model (CM) (Song al., 2023). However, their generation quality does not improve as NFE increases (the red curve of Figure 7). Theorem I (in this paper) explains this inherent absence of speed- -reed-quality trade-off CM's s multistep sampling by the overlapping time intervals between jumps. This persists as fundamental issue when training jumps solely to zero-time as in CM.
Score-based models
Distillation models
CTM
Training
Estánate the
Exthoate the culpocs of F2 008.
Lstimote the colrre FY UVE tialectrry
Scoee
NFE
Sampling
Discretisaton ecturs!
w/
NFE >
Add
Discrvcisntion recors!
Accmmnllttd crcoral
Figure Training and sampling comparisons of score-based and distillation models with CTM.
Score- -oooe models exhibit discretization errors during SDE/ODE solving, while distillation models can accumulate errors in multistep sampling. CTM mitigates these issues with 7-sampling (7 = 0).
This paper introduces the Consistency Trajectory Model (CTM) as a unified framework simulta- neously assessing both the integrand (score function) and the integral (jump) of the PF ODE, thus bridging
and distillation models. More specifically, CTM estimates jump. ranging both infnittssmallyy small jumps (score function) and long jumps (integral over any time horizon) along the PF ODE, providing increased flexibility at inference time. Particularly, our unique feature enables novel sampling method called y-sampling. which alternates forward and backward jumps along the solution trajectory, with governing the level of stochasticity.
peytime-to-anytinng
CTM's pertime-to-anytinne jump along the PF ODE greatly enhances its training flexibility as well. It allows the combination of the distillation loss and auxiliary losses, such as denoising score matching (DSM) and adversarial losses. These auxiliary losses measures statistical divergences between the data distribution and the sample distribution, which provides student high-quality training signal for better jump learning. Notably, leveraging these statistical divergences to student training enables train the student as good teacher, reaffirming the conventional belief established in the distillation community of classification tasks that auxiliary losses beyond distillation loss achieve the new Srate-Of-The-re (SOTA)
enhance student performance. In experiments, performance in both density estimation and image generation for CIFAR-10 (Krizhevsky, 2009) and ImageNet (Russakovsky et al., 2015) at resolution of 64 64.
'The DSM loss is closely linked to the KI. divergence (Song al., 2021; Kim et al., 2022c). Also, the adversarial GAN loss is proxy f-divergence (Nowozin al., 2016) IPMs (Arjovsky 2017).
2

---

Published as a conference paper at ICLR 2024
2 PRELIMINARY
In DM (Sohl-Dickstein et al., 2015; Song et al., 2020b), the encoder structure is formulated using set of continuous-tinn time random variables defined by a fixed forward diffusion process", dx, = 2 dwe, initialized by the data variable, X Pdam- A reverse-time process (Anderson, 1982) from T to 0 is established as dx, V2t log Pr (xt)dt /2t cw,, where w, is the standard Wiener process in reverse-time, and (x(x) is the marginal density of x; following the forward process. The solution of this reverse-time process aligns with that of the forward-time process marginally (in distribution) when the reverse-time process is initialized with XT ~ pr. The deterministic counterpart of the reverse-time process, called the PF ODE (Song et al., 2020b), given by
dx,
(x [x (x((x
VV (x.))
where (tu(x(x.) is the probability distribution of the solution of the reverse-timg stochastic process from time t to zero, initiated from x,, Here, (x))(x)x)) [x|x]] = + tV log pr (x,)) is the denoiser function (Efron, 2011), an alternative expression for the score function V log Pa (x²). For notational
simplicity, we omit Pro [x xx). ). a subscript in the expectation of the denoiser, throughout the paper.
In practice, the denoiser [ is approximated using a neural network Do. obtained by minimizing the DSM (Vincent, 2011; Song et al., 2020b) loss Ex (x(o)) |||x0- -D t) 1)]) where Poe. (x(x))
the transition probability from time initiated with Xo. With the approximated denoiser, the empirical PF ODE is given by
dx; dl
X2
(x(x,,t)
/
Sampling from DM involves solving the PF ODE, equivalent to computing the integral
Jr dx: dt JT
0 1
(x(x dt - Xu = xx + X, (x.. dt,
dt
(2)
t
where Xy sampled from prior distribution æ approximating Pr- Decoding strategies of DM primarily fall into two categories: score-based sampling with time-discretized numerical integral solvers, and distillation sampling where neural network directly estimates the integral.
Score-based Sampling Any off-the-shelf ODE solver, denoted as Sol Solver(xr, T. O: Φ) (with an initial value of Xy time and ending at time 0), can be directly applied solve Eq. (2) (Song et al, 2020b). For instance, DDIM (Song et al., 2020a) corresponds to Ist-order Euler solver, while EDM (Karras et al., 2022) introduces a 2nd-order Heun solver. Despite recent advancements in numerical solvers (Lu et al., 2022b; Zhang & Chen, 2022), further improvements may be challenging due to the inherent discretization eтror present in all solvers (De Bortoli et al., 2021), ultimately limiting the sample quality obtained with few NFES.
Distillation Sampling Distillation models (Salimans & Ho, 2021; Meng et al., 2023) successfully amortize the sampling cost by directly estimating the integral of Eq. (2) with single neural network evaluation. However, their multistep sampling approach (Song et al., 2023) exhibits degrading sample quality with increasing NFE, lacking clear trade-off between computational budget (NFE) and sample fidelity. Furthermore, multistep sampling is not deterministic, leading to uncontrollable sample variance. We refer to Appendix for a thorough literature review.
CTM: AN UNIFICATION OF SCORE-BASED AND DISTILLATION MODELS
To address the challenges in both score-based and distillation samplings, we introduce the Consistency Trajectory Model (CTM), which integrates both decoding strategies to sample from either SDE/ODE solving or direct seytime-to-anyting jump along the PF ODE trajectory.
3.1 DECODER REPRESENTATION OF CTM
"This paper be extended to VPSDE encoding (Song et al., 2020b) with re-scaling (Kim al., 2022a).
3

---

Published as a conference paper at ICLR 2024
CTM predicts both infinitesimally small step jump and long step jump the PF ODE trajectory. Specifically, we define G(x,, t, s) as the solution of the PF ODE from initial time t to final time S t:
CTM
Distillation
gix.,t, = 0)
Score- based g(x,, = f)
J.
X. [x du.
(6,4)
G(x,,4,8)= Xx
For stable training'. we express G as a mixture of xs and a function (inspired from the Euler solver"):
6( xx (1-3 (x(xxx)))))
Figure 2: Learning objectives of Score-based (t 8 line), distillation (s 0 line), and CTM (upper triangle).
where g(x,,4,8) =xp - Se Rixixul du. We predict
Go -xx + (1-7))0(xx,t.
(x,
as the neural jump, a combination of x, and a neural output go- This ensures the neural jump Go satisfies the initial condition Go t,t) = x2 for free. removes the necessity of enforcing the initial condition in neural network training, transforming the optimization problem from constrained unconstrained. Figure contrasts CTM's learning target with that of previous models.
A crucial characteristic of becomes evident when taking the limit as approaches t. From the definition, obtain
lim (x),,t,)) lim |. X E[x[xu] du E[x[x,].
Therefore, estimating g leads to the approximation of not only the t-to-s jump but also the infinitesimal t-to-t jump' (denoiser function). Indeed, from the Taylor expansion, we have
g(xr,t, X1 1
X - X du Xx + 1
(s
X
E[x[x,
+ +0((t-s)3) s))
[[[[[x + O(t
s|)-
Therefore, 9(x,, 8) (with general & & t) is interpreted as the denoiser function added with a residual term of the Taylor expansion.
trajectory
3.2 DISTILLATION Loss: SOFT CONSISTENCY Loss
Muttiple-step ODE solver
To achieve trajectory learning, CTM should match the neural jump Go to the true jump G by Go f, R G(x,,t, for any s S tt. We opt train Gg by comparing with a solution of the numerical solver, Solver(x,, 4, d), of the pre-trained PF ODE in Eq. (1):
X4
CTM matehes tropectors.
Gø(x,, 8) se Solver(x,, t, Φ)) v G(xi,t,
An illustration of CTM's two predictions at time
Figure 3:
With a perfect teacher d. Solver accurately reconstructs G(x,, t, s), and the optimal Gg- t, coincides with the with an initial value x,. ground truth G(x,,t,
1, s), given sufficient student network flexibility.
For more precise estimation of the entire solution trajectory, introduce soft consistency matching. As illustrated in Figure 3, soft consistency compares two s-predictions: one from the teacher and the other from the student. More precisely, the target prediction a mixture of teacher and student,
Directly learning x, + l x. xx-3(xx.1
X. [|xl]] with neural network easily lead divergence.
(t x -0x30xx Xx xx++(1 (1 [][[x(x,) Our scale choices of and thus, naturally derived from the Fuler representation by replacing TE [[ with g. refer to Appendix C.I for detailed analysis with this respect.
"Solving the PF ODE from i to with single-step Euler solver gives x
Buler
In contrast, G is unaware of this infinitesimal jump it collapses to the identity function when + and estimating the denoiser with finite differences along the time derivative imprecise due numerical issues.
4

---

Published as a conference paper at ICLR 2024
where solve the teacher PF ODE the -iiterrall and jump to using the stop- hop-gradient student. In summary, soft consistency compares
(x t, = ( (Solver(x,, u; Φ), u, s),
(3)
where random E (s, determines the amount of teacher information to distill, and sg is exponential moving average stop-gt -opggradient sg( <- stopgradduss{(( + n))).
By the choice of u, this soft matching spans two frameworks:
At u = s, Eq. (3) enforces global consistency, i.c., the student distills the teacher information the entire interval (s,
At
At, Eq. (3) local consistency, i.c., the student only distills the teacher information on single-step interval (t t), Moreover, it becomes CM's distillation target when s 0.
To quantify the dissimilarity between the student prediction Ga (x,,,,t, s) and the teacher prediction (Solo)(Solvee (Solver(x,, 4, u; Φ), u, s), we use a feature distance d in clean data space by transporting two s-time predictions to O-time using stop nop-gradient student (-) 0). More specifically. transported predictions become Xest (xs, t, s) = G Ga (X, t,,s,,
0) and Xtarget (Xe, ł, s) := Gegte) (Gzg(0) (Solver(xe. L, u; Φ), s). 8, 0). Summing altogether, the CTM loss is defined as
(0: Φ) :- E0
Exp 1x d(Xtargen (Xe, L U, 8), Xest (Xe, (,,
(Xta
which leads to the neural jump, optimum, match with the jump provided by solving the empirical PF ODE of Eq. (1), see Appendix B (Propositions and 5) for details.
3.3 AUXILIARY LOSSES FOR BETTER TRAINING OF STUDENT
In knowledge distillation for classification problems, it is widely acknowledged that the student classifier often performs as well as, or even outperforms, the teacher classifier. crucial factor contributing this success is the direct training signal derived from the data label. More precisely, the student loss classifier loss Laks which provides high-quality signal to the student with the data label.
student) + Les (data, student) combines a distillation loss Lassall and
However, in the context of generation tasks, distillation models tend to exhibit inferior sample quality compared the teacher. This is primarily because model optimization relies solely the distillation loss. In our approach, we extend the principles classification distillation to our model by introducing direct signals from both DSM and adversarial losses facilitate student learning.
First, we guide the student training with the DSM loss, given by
LDSM ( Ex (I 90 1
The optimal ge- obtained from the DSM loss becomes ge- (x,, f, t) E|x[x|] g(x,, t)t There- fore, the DSM loss improves jump precision when = f by acting as regularizer. We remark that the DSM loss mitigates the vanishing gradient problem of g learning when - t (because the scale factor - 0) and significantly improves the accuracy of small neural jumps.
Second, we employ adversarial training for enhanced student learning, utilizing the GAN loss
LGAN (0, m7) Ex, (log d (xx)] + Exc
xxo [log (1 dy ( (x, t,,,ttss))) s)))]
where d, is discriminator. This adversarial loss is motivated by VQGAN (Esser al., 2021), which shows that combination of reconstruction and adversarial losses is beneficial for generation quality.
In summary, CTM incorporates distillation, DSM and GAN losses as
(( m) == LeTM (0; x) + (DMM + AGAN LGAN (0, 17),
We opt to use the conventional DSM loss instead of the score matching loss of [R]]] (x. t) go 1)11 ensure the exact denoiser estimation the optimality.
5

---

Published as a conference paper at ICLR 2024
7
7
(a) y (Fully stochastic)
(b) 1 > >>0
(c) y = 0 (Deterministic)
Figure 5: Illustration of y-sampling with varying 7 value. It denoises with the network evaluation
and iteratively diffuses the sample in reverse by (t.. Denoise
v -1
Noisify
1n+11 )
into single and unified training framework"; and CTM solves the mini-max problem ming max, L(O n7) Following VQGAN, we employ adaptive weighting with ÀNSM Wa, (0:)) and AGAN VV, (cm(((b)) where Oz. is the last layer of the UNet output block. Iee, [0]]
TV, (o))
This adaptive weighting significantly stabilizes the training by balancing the gradient scale of each term. Algorithm 4 summarizes CTM's training algorithm.
4 SAMPLING WITH CTM
Score-baste Sumpling
Distilation Sumpling
CTM
EDM
CTM
CM
05
CTM enables exact score eval- uation through So S0 (), sup- porting standard score-based sam- pling with ODE/SDE solvers. In high-dimensional image synthesis, as shown in Figure 4's left two columns, CTM performs compara- bly to EDM using Heun's method as a PF ODE solver.
NFE 40
NFE 4
NFE
CTM additionally enables time traversal along the solution trajec- tory, allowing for the newly intro- duced 7-sampling method, refer to
Figure 4: Comparison of score-based models (EDM), distilla- tion models (CM), and CTM with various sampling methods and NFE trained on AFHQ-cat (Choi et al., 2020) 256 256.
Algorithm 2 and Figure 5. Suppose the sampling timesteps T to Xta R, where π is the prior distribution, y-sampling denoises Xt, to time v1-~22 with Go(xt, to, /1-y2t,) and perturb this neural sample with forward diffusion to the noise level at time t1. The s-sampling iterates this back-and-forth traversal until reaching to time tN = 0.
tv 0. With
Our sampling is a new distillation sampler that unifies previously proposed sampling techniques, including distillation sampling and score-based sampling.
Figure 5-(a): When y 1, it coincides to the multistep sampling introduced in CM, which is fully stochastic and results in semantic variation when NFE changes, e.g., compare samples of NFE 4 and 40 with the deterministic sample of NFE 1 in the third column of Figure 4. With a fixed XT, CTM reproduces CM's samples in the fourth column Figure 4.
Figure 5-(c): When & 0, it becomes the deterministic distillation sampling that estimates the solution of the PF ODE. A key distinction between the 7-sampling and score- based sampling is that CTM avoids the discretization errors, e.g., compare (score-based) samples in the lefumost column and (-- 0 distillation) samples in the rightmost column of Figure 4.
"The DSM loss is closely linked to the KL divergence (Song et al., 2021; Kim et et., al,, 2022b) and the GAN loss is proxy of f-divergenees (Nowozin et al. 2016) or IPMs (Arjovsky et al., 2017). Therefore, our comprehensive loss can he interpreted as Lataill (teacher, student) + DRL. (data]) (datal|]tuudnt)) + Dy (data, student), which combines the distillation loss (hetween teacher/student) with statistical divergences (hetween data/student).

---

Published as a conference paper at ICLR 2024
Reference
y y=1
y=0
Figure y controls sample variance in ttroke-based generation (see Appendix C.5).
Figure 5-(b): When 0 < < 1, it generalizes the EDM's 's stochastic sampler (Algorithm 1). Appendix B.4 shows that 7-sampling sample variances scale proportionally with y².
The optimal choice y depends on practical usage and empirical configuration (Karras et al., 2022;
Xu et al., 2023). Figure 6 demonstrates s-mmplingg in stroke-based generation (Meng et al., 2021), revealing that the sampler with Y = leads to significant semantic deviations from the reference stroke, while smaller Y values yield closer semantic alignment and maintain high fidelity. Moreover, Figure 7 showcases 7's impact on generation performance. In Figure 7-(a). 7 has less influence with small NFE, but the setup with 7 O is the only one that resembles the performance of the Heun's solver as NFE increases. Additionally, CM's multistep sampler ( = 1) significantly degrades sample quality as NFE increases. This quality deterioration concerning ? becomes more pronounced with higher NFEs, shown in Figure 7-(b), potentially attributed to the error accumulation during the iterative neural jump overlap to zero-time. We explain this phenomenon using a 2-steps y-sampling example in the following theorem, see Theorem 8 for generalized result for N steps.
Theorem 1 ((Infrrmal) 2-steps 7-samplnng). Let t E (O,T) and 7 E [0, 1]. Let pe-.2 denote as the density obtained from the y-sampler with the optimal CTM. following the transition sequence
T 72t + tt + 0, starting from pr. Then Drv (Dantas Po 2) =
=50== O(V v 7 + t).
When it becomes N steps, the 7-sampling with y = 1 iteratively conducts long jumps from fn to for each step n, which aggre- gates the error to O(VT+ + 1 + + (x). In contrast, such time overlap between jumps does not occur with = 0, elim- inating the error accumulation, resulting only ((VT) error, Appendix C.2. In summary, CTM addresses challenges asso- ciated with large NFE in distillation models with 7 = 0 and removes the discretization error in acore-based models.
NFE
NFF 35
Hus sdlaer (CTMS
0.
NFE
(a) FID by NFE
(b) Sensitivity to y
Figure 7: (a) CTM enables score-based sampling and distillation y-samplnng on CIFAR- (b) The FID de- grade highlights the importance of trajectory learning.
EXPERIMENTS
5.1 STUDENT (CTM) BEATS TEACHER (DM) QUANTITATIVE ANALYSIS
We evaluate CTM on CIFAR-10 and ImageNet 64 x 64, using the pre-trained diffusion checkpoints from EDM (CIFAR-10) and CM (ImageNet) as the teacher models. We adopt EDM's training config- uration for LDSM (0) and employ StyleGAN-XIL's (Sauer al., 2022) discriminator for LGAN n)) For student, we use EDM's s DDPM++ implementation on CIFAR- 10; and CM's ADM implementation on ImageNet. In addition these default architectures, we incorporate s-information via auxiliary temporal embedding with positional embedding ( (Vaswani et al., 2017), and add this embedding to the t-embedding. We minimally change the CM's design to comply the previous implementation, and we list important modifications in Table 4 in Appendix D.I: 1: 1) We find that a large y for aoogggradeent EMA significantly stabilizes the adversarial training; 2) We evaluate the model performance with student EMA rate of 0.999; 3) we reuse the skip connection and output scaling of the pre-trained diffusion model to our neural output modeling: ge 8x,,, t, s) Caki (t) X + Cour NN) (x,, t, s). where NN, is neural network. The selection of C.kip and Cout ensures that the initialized go (x,, t, t) closely
*The total variation of two densities and is defined as Drv g) = } [p(x) - q(x) dxx. Bold text indicates the best performance, while underlined text means the second hest performance.
7

---

Published as a conference paper at ICLR 2024
Table I: Performance comparisons on CIFAR-10 10
Table 2: Performance comparisons on Ima- geNet 64 x 64.
Model
NFF
Uncondtionl
Condttonsl
FIDI
NLL4
FIDJ
Model
NPE
FID
18
Rec1
Validaioo
64.10
Fltecck 2018)
8:1
StyloCi
2000)
292
ADM [Dharwal & Nichol, 2021)
250
2.07
0.63
SiyleAAN-DDD (Kang 20213
EDM (Kamas es al, 2022)
79
2,64
45.85
0.67
MtyleGANNNL
2003)
DIgGAN deep (Brock et al. 2015)
4,06
0.48
Sty.CAN-XL (Sancr et al, 2002)
2.09
82.35
0.52
Diffodon Models- - Score--oreeee chaee--sseeddee Sampling
DUPM (fto at atai)
1000
5.17
5:75
Distillation Sampling
DUM (Song et 20000)
100
10
13.35
4.16
(Salimuss Ho, 2021)
(Ga et 2023))
Score SDE (Song et 2000kg
2000
2:20
3.45
(ong 2023)
NUM (Kingna et 2021)
1000
7.41
CTM (oun)
20.38
LSGM (Vabdar atall 2021)
EDM4 (Kamas et 2002)
138
35
2.10
2,01
3.43
2.55
1,82
PD (Sainaas & Ho, 2021) CD USong atal, 2023)
4,70
8.95
0.65
0.64
Difftsion Model - Ditillation Sumplingg
CTM (ours)
1,73
64
64.29
0.57
KD) (lalmann&& & Lihoee, 2001)
9:36
DFNO (2heng clal. 2023)
3.76
2-Roctified Flow (1.aa ctall, 2022)
4.85
Clazaller relection sampling
NFE
(Salimes
9.12
CD (offical ropurt) (Song stall 2023)
3.55
Claailier guidance INFE 250)
NFE
CD intrainal)
10.51
CD GAN (l. al, 2003y
2,65
CTM
1.98
2.45
(Saliness 2021) 2003)
4.51
NFE
NPt
1.25
2:03
50
100
120
Models withoot Pre-trained DM- Direct Generation
15
8:50
Figure 8: FID-IS curve on ImageNet.
CTM (OnE))
2,39
(a) EDM (79 NFE)
(b) CTM w/o GAN (1 NFE) (c) CTM w/ GAN (1 NFE)
Figure 9: Samples generated by (a) EDM, (b) CTM without GAN (AGAN 0), and (c) CTM
with GAN (adaptive AGAN). More generated samples are demonstrated in Appendix E.
aligns with the pre-trained denoiser, with slight random noise introduced from the s-embedding. Reusing Caip and Cout directs the student network to focus on training long jumps while preserving the accuracy of small jumps (via LDSM) from the initial training phase. Consequently, achieving good performance requires only 100K iterations (10x faster) for CIFAR-10 and 30K iterations (20x faster) for ImageNet, compared to corresponding baselines.
CIFAR-10 CTM's NFE FID (1.98) excels not only CM (3.55) on unconditional generation, but CTM (1.73) outperforms the SOTA models, such as EDM (1.82 with 35 NFE) and StyleGAN-XL (1.85) on conditional generation. In addition, CTM achieves the SOTA FID (1.63) with 2 NFEs, surpassing all previous generative models. These results on CIFAR-10 10 are obtained upon the official py Torch code of CM, where retraining CM with their code yields FID of 10.53 (unconditional), significantly worse than the reported FID of 3.55. Additionally, CTM's ability to approximate scores using go (x,, t, f) enables evaluating Negative Log-Likelihood (NLL) (Song et al., 2021; Kim et al., 2022b), establishing a new SOTA NLL. This improvement can be attributed, in part, to CTM's reconstruction loss when
and improved alignment with the oracle process (Lai et al., 2023a).
ImageNet CTM surpasses any previous non- guided generative models in FID. Also, CTM most closely resembles the IS of validation data, which implics that StyleGAN-XL tends to generate sam- ples with a higher likelihood of being classified for a specific class, even surpassing the probabilities of world validation data, whereas CTM's s generation is statistically consistent in terms of the classifier likelihood. In sample diversity, CTM reports an intermediate level of recall, but the random samples in Figure 16 exhibits the actual samples are comparably diverse to those of EDM or CM. Furthermore, the high likelihood of CTM on CIFAR- 10 indirectly indicates that CTM has no issue on mode collapse. Lastly, we emphasize that all results in Tables I and 2 are achieved within 30K training iterations, requiring only 5% of the iterations needed to train CM and EDM.
Clasiifier-Reeection Sampling CTM's s fast sampling enables classificr-rojection sampling. In the evaluation, for each class, we select the top 50 samples out of 50 probability, where is the rejection ratio. This sampler, combined with NFE sampling, consumes an average of NFE, __y. In Figure 8, CTM shows a FID-IS trade-off comparable to classifier-guidee results (Ho & Salimans, 2021) achieved with high NFEs of 250 (see Figure 17 for samples).
1-r samples based on predicted class
8

---

Published as a conference paper at ICLR 2024
5.2 QUALITATIVE ANALYSIS
condneey rutcking
FID (NFE 1)
FID (NFE 16)
zonanirry nutcing
Cldbal coniioryy muschin
CTM Loss Figure 10 highlights that soft consistency outperforms local consistency and performs compara- ble to global consistency. Specifi- cally, local consistency distills only 1-step teacher, so the teacher of time interval [0, At] is not used to train the neural jump starting from XT. Rather, teacher on [t [t-2x مt, I] with E [0, T Ar) is distilled to student from neural jump starting from X,, not xy. The student, thus, has to extrapolate the learnt but scattered teacher across time intervals to estimate the jump from Xy, which could potentially lead to imprecise estimation. In contrast, the amount of teacher to be distilled in soft consistency is determined by random u, where 0 represents distilling teacher on the entire interval [0, T), Appendix C.3. Hence, soft matching serves as a computationally efficient and high-performing loss.
5001
T0Em 15000 2000 15033 Training iterations
300
1000 15000 Seeee 2000 Training iterations
(a) NFE 1
(b) NFE 18
Figure 10: Comparison of local, global, and the proposed soft consistency matching.
DSM Loss Figure 11 illustrates two benefits of incorporating fISM with LCTM- preserves sample quality for small NFE unless DSM scale out- weighs CTM. For large NFE sam- pling, it significantly improves sample quality due to accurate score estima- tion. Throughout the paper, we main- tain Ansm to be the adaptive weight (the case of CTM +1.0DSM), based on insights from Figure 11.
CTM
FID (NFE 1)
CIM 1
FLD (NFE 18)
6.IDSM
CTM 1OSM
55
(1oxne)
5.00SM
5300 DaBar to0M 2038) 20001 29181 Training iterations
9001 30000 19031 (RAR) 20060) 33000 Training iterations
(a) NFE 1
(b) NFE 18
Figure 11: The effect of DSM loss.
50
GAN Loss Figure 12 highlights the benefits of integrating the GAN loss for both small- and laEge-NFE sam- ple quality. In Figure 9, CTM shows superior sample production compared to the teacher, with GAN refining lo- cal details. Throughout the paper. we implement a GAN warm-up strat- egy: deactivating GAN training with ÂGAN 0 during warm-up iterations
CTM + DSM
4.5
CTM DSM
RID (NFE 1)
CTM DSM GAN
DSM GAN
30
CM (10xiiar)
4.0
15
GAN
355
30
10
35
20
10000
20081
3000
30000
20000
48000
Training iterations
Training iterations
(a) NFE 1
(b) NFE 18
Figure 12: The effect of GAN loss.
and subsequently activating GAN training with an adaptive AGAN, following the VQGAN's approach. Additional insights into the effects of GAN generated samples are discussed in Appendix C.4.
Training Without Pre-trained DM Leveraging our score learning capability, we replace the pre- trained score approximation, D t), with CTM's approximation, ge (x, t, t,. allowing us to obtain the corresponding empirical PF ODE dx, = xx=90(x....) Consequently, we can construct pretrained- free target, Ntarget == (a (G (Solver(x, u; (Q))) u. s), s, 0), to replace Xtarget in computing the CTM loss LCTM. When incorporated with DSM and GAN losses, it achieves NFE 1 FID of 2.39 unconditional CIFAR-10, 10, a performance on par with pre-trained DMs. Contrastive to CM, our CTM uses the identical form of loss from its score approximation capability.
6 CONCLUSION
CTM, a novel generative model, addresses issues in established models. With a unique training approach accessing intermediate PF ODE solutions, it enables unrestricted time traversal and seamless integration with prior models' training advantages. A universal framework for Consistency and Diffusion Models, CTM excels in both training and sampling. Remarkably, it surpasses its teacher model, achieving SOTA results in FID and likelihood for few-steps diffusion model sampling on CIFAR-10 10 and ImageNet 64 x 64, highlighting its versatility and process.
9

---

Published as a conference paper at ICLR 2024
ETHICS STATEMENT
CTM poses a risk for generating harmful or inappropriate content, including deepfake images, graphic violence, or offensive material. Mitigating these risks involves the implementation of strong content filtering and moderation mechanisms to prevent the creation of unethical or harmful media content.
REPRODUCIBILITY STATEMENT
The code is available at ht / //ggttitttt tuub. com/ www/sony/ ctt. Moreover, we outline our training and sampling procedures in Algorithms 4 and 2, and detailed implementation instructions for result reproducibility can be found in Appendix D.
ACKNOWLEDGEMENT
We sincerely acknowledge the support of everyone who made this research possible. Our heartfelt thanks go to Koichi Saito, Woosung Choi. Kin Wai Cheuk, and Yukara lkemiya for their assistance. Computational resource of AI Bridging Cloud Infrastructure (ABCI) provided by National Institute of Advanced Industrial Science and Technology (AIST) was used.
10

---

Published as a conference paper at ICLR 2024
REFERENCES
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Brian DO Anderson. Reverse-time diffusion equation models. Stochustic Processes and their Applications, 12(3):31-3266 1982.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap, Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbot, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.0m 0444488 2023.
Dennis D Boos. A converse scheffe's theorem. The Annals of Statistics, pp. 423 427, 1985.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2018.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference Learning Representations, 2022.
Kin Wai Cheuk, Ryosuke Sawata, Toshimitsu Uesaka, Naoki Murata, Naoya Takahashi, Shusuke Takahashi, Dorien Herremans, and Yuki Mitsufuji. Diffroll: Diffusion-basee generative music transcription with unsupervised pretraining capability. ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8188-8197, 2020.
Giannis Daras, Yuval Dagan, Alexandros G Dimakis, and Constantinos Daskalakis. Consis- tent diffusion models: Mitigating sampling drift by learning to be consistent. arXiv preprint arXiv:2302.09055 09057, 2023.
Valentin De Bortoli, James Thorton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising
diffusion for image restoration. Transactions on Machine Learning Research, 2023.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 101-1161,, 2011.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high- high-resolution image synthesis. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Ian Goodfellow, Jean Pouget-Abadics Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural information processing systems, 27, 2014.
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free distil- lation of denoising diffusion models with bootstrapping. In ICML 2023 Workshop on Structured Probabilistic Inference ((&) Generative Modeling, 2023.

---

Published as a conference paper at ICLR 2024
Carlos Hernandez-Olivan, Koichi Saito, Naoki Murata, Chieh-Hsin Lai, Marco A Martinez-Ramirez, Wei-Hssang Liao, and Yuki Mitsufuji. Vrdng: Vocal restoration via diffusion posterior sampling with multiple guidance. arXiv preprint Sat1:22309009934 2023.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840- 68860066681, 2020.
Minguk Kang, Woohyeon Shim, Minsu Cho, and Jaesik Park. Rebooting acgan: Auxiliary classifier gans with stable training. Advances in neural information processing systems, 34:23505-23518, 2021.
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in neural information processing systems, 33:1 2104-12114, 2020.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances Neural Information Processing Systems, 35:26565-26577. 2022.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593-23606 2022.
Dongjun Kim, Yeongmin Kim, Wanmo Kang, and Il-Chul Moon. Refining generative process with discriminator guidance in score- bose- diffusion models. arXiv preprint arXiv:2211.1709 2022a.
Dongjun Kim, Byeonghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, and I1-chul Moon. Maxi- mum likelihood training of implicit nonlinear diffusion model. Advances in Neural Information Processing Systems, 35:32270-32284 2022b.
Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and Il-Chul Moon. Soft truncation: universal training technique of score-based diffusion model for high precision score estimation. In International Conference on Machine Learning, pp. 11011-12228, PMLR, 2022c.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 14:21696-21707, 2021.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.611 14, 2013.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Tuomas Kynkäänniemi, Tero Karras, Miika Aittala, Timo Aila, and Jaakko Lehtinen. The role of imagenet classes in fréchet inception distance. Proc. ICLR, 2023.
Chich-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji. and Stefamo Ermon. FP-diffusion: Improving score-based diffusion models by enforcing the underlying score fokker- planck equation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 23-29 Jul 2023a. URL https: http://pcoccceen...c/e////. mll. malr.press press htt///22/////// add2dd...... html..
PMLR,
Chich-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka. Naoki Murata, Yuki Mitsufuji. and Stefano Ermon. On the equivalence of consistency-type models: Consistency models, consistent diffusion models, and ookkeer-planck regularization. arXiv preprint arXiv:2306.00 Ttl.:23066006677, 2023b.
Yangming L., Zhaozhi Qian, and Mihaela van der Schaar. Do diffusion models suffer error prop- agation? theoretical analysis and consistency regularization. arXiv preprint arXiv:2308.055: 05021, 2023.
Xingchao Liu, Chengyue Gong, et al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations, 2022.
12

---

Published as a conference paper at ICLR 2024
Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score- based diffusion odes by high order denoising score matching. In International Conference Machine Learning, pp. 14429-144600 PMLR, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022b.
Haoye Lu, Yiwei Lu, Dihong Jiang, Spencer Ryan Szabados, Sun Sun, and Yaoliang Yu. Cm- gan: Stabilizing gan training with consistency models. In ICMI. 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint wrXiv:2101.02388 2021.
Chenlin Meng, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. Sdedit: Image synthesis and editing with stochastic differential equations. arXiv preprint 10Xiv:2108.01073, 2021.
Chenlin Meng. Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Naoki Murata, Koichi Saito, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Gibbsddrm: partially collapsed gibbs sampler for solving blind inverse problems with denoising diffusion restoration. arXiv preprint arXiv:2301.12668 2023.
Scbastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. Advances in neural information processing systems, 29, 2016.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations, pp. 65-84. Springer, 2003.
W.T. Reid. Ordinary Differential Equations. Applied mathematics series. Wiley, 1971.
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, 10695, 2022.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-255 2015.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494 2022.
Koichi Saito, Naoki Murata, Toshimitsu Uesaka, Chieh-Hsin Lai, Yuhta Takida, Takao Fukui, and Yuki Mitsufuji. Unsupervised vocal dereverberation with diffusion- generative models. In ICASSP 2023-2023 IEEE International Conference on Acoustics. Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations, 2021.
Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xI: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings, pp. 1-10, 2022.
Shitong Shao, Xu Dai, Shouyi Yin, Lujun Li, Huanran Chen, and Yang Hu. Catch-up up distillation: You only need to train once for accelerating sampling. arXiv preprint arXiv:2305. 10769, 2023.
13

---

Published as a conference paper at ICLR 2024
Jascha
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning. pp. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2020a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems,
2020.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. bcore-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020b.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34:1415-1428, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
TJ Sweeting. On a converse to scheffé's theorem. The Annals of Statistics, 14(3):1252-1256, 1986.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference on machine learning, pp. 6105-61 14. PMLR, 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 14:11227-11302 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Pascal Vincent. A connection between matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. arXiv preprint arXiv:2306. 14878, 2023.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Kleventh International Conference on Learning Representations, 2022.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning, pp. 02390-42402 PMLR, 2023.
14

---

Published as a conference paper at ICLR 2024
CONTENTS
Introduction
2
Preliminary
3
CTM: An Unification of Score-based and Distillation Models
3.1
Decoder Representation of CTM
3.2
Distillation Loss: Soft Consistency Loss
3.3
Auxiliary Losses for Better Training of Student
4
Sampling with CTM
5
Experiments
5.1
Student (CTM) beats teacher (DM) Quantitative Analysis
5.2
Qualitative Analysis
6
Conclusion
A Related Works
17
B Theoretical Insights on CTM
17
B.1 Unification of score-based and distillation models
17
B.2 Convergence Analysis Distillation from Teacher Models
18
B.3 Non-Intersecting Trajectory of the Optimal CTM
19
B.4 Variance Bounds of y-sampling......
20
B.5
Accumulated Errors the General Form of samppiin
20
B.6 Transition Densities with the Optimal CTM
20
Algorithmic Details
21
C.1 Motivation of Parametrization
C.2 Characteristics of 7-sampling
22
Algorithmic Comparison of Local Consistency and Soft Consistency
C44 Comparison of GAN Effects in Generation
23
C.5 Trajectory Control with Guidance
24
Implementation Details
25
D.1 Training Details
25
D.2 Evaluation Details.
26
Additional Generated Samples
26
Theoretical Supports and Proofs
28
15

---

Published as a conference paper at ICLR 2024
F.1
Proof of Lemma
28
F2 Proof of Theorem
28
E3 Proof of Proposition 3.
29
E4
Proof of Proposition 5
30
F5
Proof of Proposition 6.
30
F6 Proof of Proposition 7
32
F.7 Proof of Proposition 9
33
16

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
Diffuusion Models DMs excel in high-fidelity synthetic image and audio generation (Dhariwal & Nichol, 2021; Saharia et al., 2022; Rombach al., 2022), as well in applications like media editing, restoration (Meng et al., 2021; ; Cheuk et al., 2023; Kawar et al., 2022; Saito et al., 2023; Memandezz-livaal
al., 2023; Murata et al., 2023). Recent research aims to enhance DMs in sample quality (Kim et al., 2022b;a), density estimation (Song et al., 2021; Lu et al., 2022a), and especially, sampling speed (Song et al., 2020a).
Fast Sampling of DMs The SDE framework underlying DMs (Song et al., 2020b) has driven research into various numerical methods for accelerating DM sampling, exemplified by works such as (Song et al., 2020a; Zhang & Chen, 2022; Lu et al., 2022b). Notably, (Lu et al., 2022b) reduced the ODE solver steps as few as 10-15. Other approaches involve learning the solution operator of ODEs (Zheng et al., 2023), discovering optimal transport paths for sampling (Liu et al., 2022), or employing distillation techniques (Luhman Luhman, 2021; Salimans Ho, 2021; Berthelot et al., 2023; Shao et al., 2023). However, previous distillation models may experience slow convergence or extended runtime. Gu ct al. (2023) introduced a bootstrapping approach for data-free distillation. Furthermore, Song et al. (2023) introduced CM which extracts DMs' PF ODE to establish direct mapping from noise to clean predictions, achieving one-step sampling while maintaining good sample quality. CM has been adapted to enhance the training stability of GANs, as (Lu et al., 2023). However, it's important to note that their focus does not revolve around achieving sampling acceleration for DMs, nor the results restricted simple datasets.
Consistency of DMs Score-based generative models rely on differential equation framework, employing neural networks trained on data to model the conversion between data and noise. These networks must satisfy specific consistency requirements due to the mathematical nature of the underlying equation. Early investigations, such as (Kim al., 2022c), identified discrepancies between learned scores and ground truth scores. Recent developments have introduced various consistency concepts, showing their ability to enhance sample quality (Daras et al., 2023; et al., 2023), accelerate sampling speed (Song et al., 2023), and improve density estimation in diffusion modeling (Lai ct al., 2023a). Notably, Lai et al. (2023b) established the theoretical equivalence of these consistency concepts, suggesting the potential for unified framework that can empirically leverage their advantages. CTM can be viewed as the first framework which achieves all the desired properties.
B THEORETICAL INSIGHTS ON CTM
In this section, we explore several theoretical aspects of CTM, encompassing convergence analysis (Section B.2), properties of well-trained CTM, variance bounds for y-sampling. and a more general form of accumulated errors induced by 7-sampling (cf. Theorem 1).
We first introduce and review some notions. Starting at time t with an initial value of x, and ending at time recall that G(x,, () represents the true solution of the PF ODE, and G(x,, 4, 8,8; Φ) is the solution function of the following empirical PF ODE.
dx, X,, D& (X.,
u E [0,
(4)
du
Here & denotes the teacher model's weights learned from DSM. Thus, G(xt, 8; () can be expressed as
8 x + (1 2 ))((x f, Φ).
G(x,, t, s;
t
where g(xe; 4, Φ)- X, similar formulation, albeit for different purposes.
Dx. du. (Delbracio & Milanfar, 2023) also derived
B. UNIFICATION OF SCORE-BASED AND DISTILLATION MODELS
The following lemma summarizes our dedicated G-expression using an auxiliary function g. allowing convenient access both the integral via G and the integrand via g. visualized in Figure 13.
17

---

Published as a conference paper at ICLR 2024
Lemma 2 (Unification of score-based and distillation models). Suppose that the score satisfies supx So | log Pr (x)ll2 du < oc. The solution G t, 8) can be expressed as:
1( x; (x(()) with g(x,,1,8) X, t
1
{" Xu
E| Xul du.
u
Here, g satisfies:
When 0, G(x,, t, g(x,, 0,0) is the solution of PF ODE = 0, initialized at X
As t, g t, s) E0 [x]]] Hence, g can be defined at t by its limit: g(x,, E[x
As outlined in Section 3.1, the g-expression for G is in- herently associated with the Taylor approximation to the integral:
Diffusion models
Consistency Trajectory Model (CTM)
Score-based models
Distillation models
(s 1) Xt [[x + 0 s
Sucs 50B: Siit et 2003A). Dora: (ile 2036. EDN (Komm sl, 2000)
xD (tatenan Labvon 3E31), 2thngg anDy
G(xi, 1, = Xx +
= - + (1- [|x|x,] + O (1t -
(x,.sttss)
Figure 13: Schematic illustration of our CTM.
To further elucidate why Taylor's expansion is the primary cause of discretization errors, we will provide an explana-
tion using the DDIM sampler with the oracle score function. The denoised sample with DDIM from t to Åt x = (1- - )x+ + At 1 [[ However, the Taylor expansion of the integration yields that the true trajectory sample is x = (1- Al )x,+ + At (E[x|x;] ((A)). Therefore, the DDIM trajectory differs from the true trajectory by (At),, which exactly represents the residual term of the Taylor expansion beyond the 2nd order. Consequently, the discretization error originates from the failure to estimate the residual term of the Taylor expansion.
B.2 CONVERGENCE ANALYSIS DISTILLATION FROM TEACHER MODELS
Convergence along Trajectory a Time Discretization Setup. CTM's practical implementation follows CM's one, utilizing discrete timesteps to 0 < 4 <... < IN T for training. Initially, we assume local consistency matching for simplicity, but this can be extended to soft matching. This transforms the continuous time CTM loss to the discrete time counterpart:
C (0; Φ) == Eng Exo. (x(xa) [d(xx (XSaagge (Xt,, tn tm): Xest (X,, tny mm
where .(;,.) is metric, and
Xest (Xt, In, tm) = Go Go(Xx,. tu, tm). tm.0)
Xtarget (xt,, tn: tn 1h-1,te tm) = Go Go tno tniin- tn tn f,
t 0).
In the following theorem, we demonstrate that irrespective of the initial time in and end time lmn. CTM Gol', In, lm; Φ), will eventually converge to its teacher model, G(,, Ins ().
Proposition 3. Define Axt == max {ltn+1 - t. Assume that Go is uniform Lipschitz in nE[L.1]
and that the ODE solver admits local truncation error bounded uniformly by O ((ANt))tt) with p Z 1. If there is Ox that LN x0;; Φ) 0, then for any n E [1, N] and E [1,n]
sup d(Gay (Gox (x, bny 1m), Lmm : 0), Go (G(x, L, km; Φ), bm 0)) = ((Axx())( - lm)- x
Similar argument applies, confirming convergence along the ODE trajectory, ensuring the local consistency with 0 replacing s@):
(,( Go(Solver(x Ad; (), Al.s)
18

---

Published as a conference paper at ICLR 2024
by enforcing the following loss
CCM (0; Φ) [n Bmm [u poo. (x/)x0) d( (xarggr (Xe,, Iin (nnntt)) tm)s
(Xe, Ena fm))
where
(xa( bny km) = Go :n, bm)
Xarget (x,,: In tn-1.tm) tm) == Go (Solver (Xe, tn. 1n: in tm))
Proposition 4. If there is a ON so that CCTM (6) = 0. then for any n E [1 and m E [1,
sup d(Gex (x, tn: tm); G(x, tn tmi Φ)) (((ANt)))(t -
x
Convergence of Densities. In Proposition 3, we demonstrated point-wise trajectory convergence, from which we infer that CTM may converge to its training target in terms density. More precisely, in Proposition 5, we establish that if CTM's target Xurget is derived from the teacher model (as defined above), then the data density induced by CTM will converge to that of the teacher model. Specifically, if the target Xtarget perfectly approximates the true G-function:
Xurget Xt,, a In In fm) = G (X bn, (m), for all E [1, E [0, N N.
(5)
Then the data density generated by CTM will ultimately learn the data distribution Pdata-
Simplifying, the l2 for the distance metric d and consider the prior distribution to be pr. which is the marginal distribution at time t = ' defined by the diffusion process:
dx, 2 dw.,
(6)
Proposition 5. Suppose that
(i) The uniform Lipschitzness of Ge (and G), sup [|Go(x,,,,)) Go(x', (, 8) L [||x x112, for all x,x E RR. RR,,,, [0,
(ii) The uniform boundedness in 0 of Go: there is L(x)= 2 so that sup (GGo(x,t,s)ll S L(x) < oc, for allall x RR. t, E [0. 0
If for any N, there is a ON such that (0 Φ) = Let po (.) denote the pushforward distribution of pr induced by Gey (-,T,), Then, N + 00, ||o (-) (o)l) O. Particularly,
if the condition in Eq. (5) is satisfied, then [|pay (-) Pdtto (-) + 0 as N + 0c.
B.3 NON-INTERSECTING TRAJECTORY OF THE OPTIMAL CTM
CTM learns distinct trajectories originating from various initial points xt and times t. In the following proposition, we demonstrate that the distinct trajectories derived by the optimal CTM, which effectively distills information from its teacher model (Go- (Go-(,t,s) = G(.,t, Φ) for any t, [0, T)), do not intersect.
Proposition 6. Suppose that a well-trained 0 such that Go. (-,t,)) G(-,t, s; for any t, [0. TJ, and that tt is Lipschitz, iee,, there is a constant L 0 so that for any x, y E R and t t [
[x Do t)]l S LG ||x -e--y12-
Then for any s E t), the mapping Gø- (,, .t, s): RD + RD is bi bi-Lisschitz. Namely, for any Xt.Yi RD
(a((-s)) ||x, yellz $ |Go (X,, Go (y,, s))l2 S cfall-o) |x+ -e--ell2-- This implies that x, # yt. Go- (x15 I, Go (yt; for all [0,
(7)
Specifically, the mapping from an initial value to its corresponding solution trajectory, denoted as X, + Ge- (xe, tt -), is injective. Conceptually, this ensures that if we use guidance at intermediate times to shift a point to another quided-targer trajectory, the guidance will continue to affect the
outcome / 0.
19

---

Published as a conference paper at ICLR 2024
B.4 VARIANCE BOUNDS OF 7--AMMPING
Suppose the sampling timesteps are T = to > 4i >.. IN = 0. In Proposition 7, we analyze the variance of
Xu+1 = bn V1 bn+1) + Zn:
resulting from n-step y-sampling, initiated
X1 == (x fo. v1-1?tt) + 72o. where Zn iid ~
Here, we assume an optimal CTM which precisely distills information from the teacher model Ga- (.) = G(,,4, 8; Φ) for all 1,8 E [O, for simplicity.
Proposition 7. We have
(X,,) + + n+1 I Var (Xa+1) S ((trs tn+1, 7) Var (X,) + 7 1+
where C(tns (nt1. 7) = exp 22 11-11--00 pt++)) and Lo is a Lipschitz constant of D D& t).
In line with our intuition, CM's multistep sampling (y= 1) yields a broader range of Var (X++1)
compared to y resulting in diverging semantic meaning with increasing sampling NFE.
B.5 ACCUMULATED ERRORS THE GENERAL FORM OF 7-SAMPLING.
We can extend Theorem for two steps 7-sampligg for the case of multisteps.
begin by clarifying concept "density transition by a function". a measurable mapping T : Ω + Ω and measure the measurable space Ω the notation Tov denotes the pushforward measure, indicating that a random vector X follows the distribution v, then T(X) follows the distribution Tiv.
Given a sampling timestep = > 1 > ... > IN = Let Pa-,N represent the density resulting from N-steps of y-sampling initiated at Pr. That is.
N-1
P
To
VI- thtt o T bn b -1--11))))
EpT.
n=0
N-1
Here, denotes the sequential composition. We assume an optimal CTM which precisely distills
information from the teacher model Gø- (-) G(, s; Φ) for all 1 E [0, T). Theorem 8 (Accumulated errors of N-steps 7-sampling). Let [0, 1].
DTV (Péara No) = O
Σ ==
V tn /1 7
Here, T + RD denotes the oracle transition mapping from to determined by Eq. (6). The pushforward density via Te-sa is denoted as Tr with similar notation applied to T where T denotes the transition mapping associated with the optimal CTM trained with LCTM-
B.6 TRANSITION DENSITIES WITH THE OPTIMAL CTM
In this section, for simplicity. we assume the optimal CTM. G = G with well-learned 0", which recovers the true G-function. We establish that the density propagated by this optimal CTM from any time subsequent time s aligns with the predefined density determined by the fixed forward process.
We now present the proposition ensuring alignment of the transited density.
Proposition 9. Let {p:TT= IT T be densities defined by the diffusion process Eq. (6), where Po Pdata- Denote (-) = G(,, t, s): RD + IRD for any t 2 s. Suppose that the score V log satisfies that there is function T(()) 2 0 that J T(()) dt oo and
20

---

Published as a conference paper at ICLR 2024
(i) Linear growth: [Vlogp(()]]2 (1 + [xx11)) all R
(ii) Lipschitz:
| (x Vlogg: - L(t) x foo all x, E RD.
Then for any [0, and [0.4. P. Tp=-
This theorem guarantees that by learning the optimal CTM, which possesses complete trajectory information, we can retrieve all true densities any time using CTM.
C ALGORITHMIC DETAILS
C.I MOTIVATION OF PARAMETRIZATION
Our parametrization of Ge is affected from the discretized ODE solvers. For instance, the one-step Euler solver has the solution of
E[x
t Xx + (1- E[x/xx))-
Faler
X (1
The one-step Heun solver is
Houn
Xx
E[x Xt]
x
Buler
[E[x [E[x|xtaber] buulerr
Xx
2
8
X1
x
Euler
8
[E(x|x1] +
R F[x/xxBuler)
Xt
+
2
t
2
t
xx +
E[x[xx]
+
TE(x x Eutor
2s
2s
Again, the solver scales x2 with ? and multiply 1 1-3 ? to the second term. Therefore, our G(x,, s)
[x (1 {)g(X,, t, s) is natural way to represent the ODE solution.
For future research, we establish conditions enabling access to both integral and integrand expressions. Consider continuous real-valued function a(t, s). We aim to identify necessary conditions on a(t, s) for the expression of G as:
t, t, = a(t, s)x, + - ((x,,,)
for a vector-valued function h(x,, t, s) and that h satisfies:
lim,- - 1xx,, exists;
it can be expressed algebraically with E[x[x,].
Starting with the definition of G, we can obtain
t( =x, + '." E[x[xu] du
/
X., Efx[Xx]]
= a(t, 8)x + (1 a(t, xt + 1- a(t,8)
du
1(xx....8)
Suppose that there is a continuous function c(t) so that
--t
lim
c(t))
+ 1 - s)t,s)
then
1
Xu. - E E[x [[x|Xx]] Xx]
lim h(x,, t) =x, + lim o-3
du
I at,, 8)
X1- E[x x
X; + lim
for some t" E [s,
alt, 8)
tz
21

---

Published as a conference paper at ICLR 2024
Xx 16/x Xx.
x, + c(t)) 11+ c(t)
c(t) B[x|xx]
Xt
The second equality follows from the value theorem (We omit the continuity argument details for Markov filtrations). Therefore, we obtain the desired property 2). We summarize the necessary conditions on a(s, t)
There some continuous function c(t) in so that lim
s-4 - aattt c(t).
(8)
We now explain the above observation with an example by considering EDM-type parametrization.
auia 2 and Cour (t, s) (1 {) Then G(x,, t, s) can be expressed
Consider (tkip(t,s) ==
( (t-o amm) 40
as
G(x 8) Cseip (1,8) Coul , 8) s))((x,,
where h is defined as
1
[11 Cakip )x, + f" x, E X
h(x,, t,
du
Caut
Then, we can verify that cakip (t, 8) satisfies the condition in Eq. (8) and that
E[x[x]] = g(xe,, t, + min + dana mmint (t mmin) )2 + o Jastn
The DSM loss with this Cstip(t, becomes
(D) E xo (90(x), + 2 + daan Tminn (t Tmin) 1 o data
=1|x 10
X,
However, empirically, we find that the parametrization of cakip (t, 8) and Cour(t, 8) other than the ODE solver-oriented one, i.e., (4kip(4, (4, = and Cakip (4, 8) = 1- 7. faces training instability. Therefore,
set G(x,,t, s) fxx + ()g(()e f, s) as default design and estimate g-function with the neural network.
C.2 CHARACTERISTICS OF SAMMLIIN
Connection with SDE When Go = G, a single step of y-sampling is expressed as:
X7 = =X,, + G( tn. V1 7 tn+1) + ytnt1E
u (
Inil
= =Xi,.
Pa (xx)) du +
w log (XCu) du + 7t++1E,
fuil
information
future information
where e NO, I). This formulation cannot be interpreted as a differential form (Øksendal, 2003) because it look-ahead future information (from 1x+1 to / - 1+ to generate the sample x7 at time fn+1- This suggests that there is no Itō's SDE that corresponds to our 7-sampler pathwisely, opening up possibilities for the development of new family of diffusion samplers.
Connection with EDM's stochastic sampler We conduct a direct comparison between EDM's stochastic sampler and CTM's sammpling. We denote Hewn t, s) Heun's solver initiated at time and point x, and ending at time s. It's s worth noting that EDM's sampler inherently experiences discretization errors stemming from the of Heun's solver, while CTM is immune such errors.
22

---

Published as a conference paper at ICLR 2024
Algorithm 1 EDM's sampler
Algorithm 2 CTM's 7-sampling
1: Start from Xto ~
1: Start from Xt ~ π
2: for 0 N do
2:
for n 0 to N 1 do
3: En - (1+
3: (n-1 f- / - 12
4: Diffuse X Xt, V f
4: Denoise Xin+1 x Go Ins bnt1)
5: Denoise _ f Heun(x,, Ens bxt1)
5: Diffuse Xt t X Xn11 + Autte
6: end for
6: end for
7:
Return XIN
7: Return Xty
The primary distinction between EDM's stochastic sampling in Algorithm I and CTM's 7-sampling in Algorithm is the order of the forward (diffuse) and backward (denoise) steps. However, through the iterative process of forward-backward time traveling, these two distinct samplers become indistin- guishable. Aside from the order of forward-backward steps, the algorithms essentially align if we opt to synchronize the CTM's time (1CTM, :n (CTM) to with the EDM's time (7HDM, (EDM) respectively, and their ys accordingly.
C.3 ALGORITHMIC COMPARISON OF LOCAL CONSISTENCY AND SOFT CONSISTENCY
In this subsection, we explain the algorithmic difference between the local consistency loss and the soft consistency loss focusing on how the neural jump Go(xT, T, 0) is trained.
Local Consistency (Implicit Information from Teacher) Let us assume that at some training iteration the maximum time T is sampled as random time t. Then CM matches the long jump provided by (XT, T, O) and =g (So (Solver(XT. T T At, 0). Hence, the neural jump (xr, 0) distills on the teacher information within the interval [T At, T" and may lack precision for the trajectory within [0, T Ar.. The transfer of teacher information for the interval [0, T Af] may occur in another iteration with a random time / S T - At. In this case, the student model Go (xx, f, 0) distills the teacher information solely within the interval [t At, t]. where the network is trained with X; as the input.
However, for 1-step generation, go (xr, T,0) 0) still lacks perfect knowledge of the teacher information within [L - At, t). This is because, when the student network input is Xy. . the teacher information for the interval [t At, t] with t < T At has not been explicitly provided, as the student was trained with the input x, to distill information within [t At, t). Given the non-overlapping intervals with distilled information from local consistency, the student neural network must extrapolate and attempt to connect the scattered teacher information. Consequently, this implicit signal provided by teacher results in slow convergence and inferior performance.
Soft Consisteney (Explicit Information from Teacher) At the opposite end of local consistency, there is glocal consistency, where the teacher prediction is constructed solely with an ODE solver to cover the entire interval [o, I'] (or [0, t] for random time t). In this case, the student model can explicitly extract information from the teacher. However, this approach is resource-intensin (3x slower than local consistency on CIFAR- 10) during training due to the ODE solving calls on the entire interval at each iteration.
In contrast, innovative loss, soft matching, constructs the teacher prediction by using ODE solver spanning from to a random u. Importantly, u is not limited to T At. but can take any value in the range of (0, T). As result, the teacher information has the opportunity to be distilled and transmitted over a broader range of [u, T). More precisely, if a random u is sampled with u S </ L- Av, the range [u. T'] contains the interval [t At, t). and the student network directly distills teacher information of the interval [t At, t). As is arbitrary, the student of CTM with input xy will ultimately receive the explicit information from the teacher for any intermediate timesteps. This renders CTM superior to CM for 1-step generation, as evidenced in Figure 10, while maintaining training efficiency (2x faster than global consistency on CIFAR-10).
C.4 COMPARISON OF GAN EFFECTS IN GENERATION
23

---

Published as a conference paper at ICLR 2024
(a) Teacher EDM (NFE 79)
(b) CTM without GAN (NFE 1)
(c) CTM with GAN (NFE 1)
Figure 14: Uncurated samples from (a) teacher, (b) CTM without GAN, and (c) CTM with GAN. For visualization purpose, upsample 64 64 samples to 224 224 resolution with bilinear upsampling technique. Best viewed with zoom-in.
This section investigates the effect of adversarial training with generated samples and statistics. In Figure 14, we compare the samples of (a) the teacher diffusion model, (b) CTM (NFE 1) without GAN, and (c) CTM (NFE 1) with GAN. shows that the samples with auxiliary GAN loss exhibit enhanced fine details, effectively addressing high-frequency alisasing artifacts. More- over, improvements in overall shapes Nutterflybbckkroundd and are evident in these sam-
Table 3: Effect of GAN Loss CIFAR-10. We use identical hyperparameters except the GAN loss for fair comparison.
Model
NFE
FID
features htigghhnness/cootassssssttrrtrer ples. Although existing literature (Kynkäännieni al., 2023) discusses the possibility that FID improvement may not neces- sarily correlate with an actual enhancement in human perceptual judgement, our observations in Figure 14 indicates that, in the case of CTM, the improvement achieved through GAN indeed perceptually discernible in human judgement.
CTM w/o GAN
5.19
1
CTM w/ GAN
2.28
CTM w/o GAN
3.00
18
CTM w/ GAN
2.23
compare FIDs of CTM trained with/without GAN Table 3. Consistent to the findings in previous research (Song & Ermon, 2020), the improved fine-details of GAN-augmented samples Figure 14-(c) results better FID than CTM without GAN, as indicated in the table. Moreover, Table 3 demonstrates that the use of adversarial loss is also beneficial the generation of large-NFE samples.
C.5 TRAJECTORY CONTROL WITH GUIDANCE
We could apply s-mmmiin for application tasks, such as image inpainting or colorization, using the (straightforwardll) generalized algorithm suggested in CM. In this section, however, we propose loss-based trajectory optimization algorithm in Algorithm 3 for potential application downstream tasks.
24

---

Published as a conference paper at ICLR 2024
Algorithm 3 Loss-based Trajectory Optimization
1: Xref is given
2: Diffuse Xto +- Xref + L06
3: for 1 to do
4: in + 11 11--2tn
5: Denoise x, f (x
6: for m to M do
7:
Sample e, c' (O.II)
Apply corrector Xi, Xi, + (Vlogrr, (xi,)-co -i, Vxx,, L(xr, Xref (ne)) vsc end for
8:
9:
10: Sample E ~ (0,1))
11: Diffuse Xt, < Xi, c
12:
end for
Algorithm uses the time traversal from tn-1 fn. and apply the loss-embedded corrector (Song et al. 2020b) algorithm to explore mmanfoll.. For instance, the loss could be a feature loss between Xr_ and Xter + E. With this corrector-based guidance, This loss-embedded corrector could also be interpreted as sampling from a posterior distribution.
could control the sample variance.
For Figure 6, choose N = 2 with (to. t1) ((omi (a 1/p o oidl) )0.4),, (o (a 1/p 1/0 )0.3))))... Cr, 1, and M = 10.
min
mmax
D IMPLEMENTATION DETAILS
D.1 TRAINING DETAILS
Algorithm 4 CTM Training
Following Karras al. (2022), we utilize the oan
1:
repeat
EDM's skip connection Cskip (t)
and
2:
Sample xo from data distribution
to
3:
Sample (N, /)
output scale (t) torm V1++oa
for ge modeling
4:
Sample E [0, T] . E [0, tl. u E [s,t)
as
5:
Calculate X, Xo + te
6: Calculate So Solve solveriae vollvr((x,, I, Φ)
go(x,, 1, (tki((t))x + Caul (N) 1,
7: Update < 0 ag L(B, n)
where NNg refers to the actual neural network output. The advantage of this EDM-style skip and output scalings are that if we copy the teacher model's parameters the student model's param- eters, except student model's s-embedding structure, go (x,,t,t) tt initialized with & would be close to the teacher denoiser D& [x, t). This good initialization partially explains the fast convergence speed.
8: Update + + on LOAN 0 17)
9: until converged
We use 4xV100 (16G) GPUs for CIFAR-10 experiments and 8xA100 (40G) GPUs for ImageNet experiments. We use the warm-up for AGAN hyperparameter. On CIFAR-10 10, deactivate GAN training with ÀGAN = until 50K training iterations (200K for training without pre-trained DM) and activate the generator training with the adversarial loss (added to CTM and DSM losses) by setting AGAN to be the adaptive weight. The minibatch per GPU is 16 in the CTM+DSM training phase, and 11 in the CTM+DSM+GAN training phase. On ImageNet, due the excessive training budget, we deactivate GAN only for 10k iterations and activate GAN training afterwards. We fix the minibatch be 11 throughout the CTM+DSM or the CTM+DSM+GAN training ImageNet.
We follow the training configuration mainly from CM. but for the discriminator training, we follow that of StyleGAN-XI. (Sauer et al., 2022). For LCTM calculation, we use LPIPS (Zhang et al., 2018) as feature extractor. We choose t and from the N-discretized timesteps to calculate LCTM. following CM. Across the training. we choose the maximum number of ODE steps to prevent a single iteration takes too long time. For CIFAR-10 10, we choose N 18 and the maximum number of ODE steps to be 17, i.e. we do nothing for CIFAR- 10 training. For ImageNet, we choose = 40 and the maximum number of ODE steps to be 20. We find the tendency that the training performance is improved by the number of ODE steps, so one could possibly improve our ImageNet result by choosing larger maximum ODE steps.
25

---

Published as a conference paper at ICLR 2024
Table 4: Experimental details on hyperparameters.
Hyperpurancter
CIFAR-10
ImageNNet 64x64
Unconditional
Conditional
Conditional
Training with
Training from Seratch
Training with
Training with
Learning rale
0.0004
0,0004
0.0004
0.000008
Discriminator leanning rale
0.002
0.002
0.002
Student's is stop grad EMA parameter
0,9999
01.999
0,999
0.999
ODE solver
Heun
Heun
Heun
Heun
Teacher
(x(
Se f,
(a(X,
Dg(X,.
Max. ODE steps
17
17
17
20
decay rate
0.999
0.999
0.999
Training iterations
100K
300K
100K
30K
Pixsed-rrecsiio (FP16)
True
True
True
True
Batch size
256
128
512
2048
Number of GPUs
For LOSM calculation, we select 50% of time sampling from EDM's original scheme of ~ (-2222, 1.22). For the other half time, we first draw sample from & ~ [0, 0.7] and transform 1/p using (a + f( min - m This specific time sampling blocks the neural network to forget the denoiser information for large time. For LGAN calculation, we two feature extractors to transform GAN input to the feature space: the EfficientNet (Tan & Le, 2019) and DeiT-base (Touvron et al., 2021). Before obtaining an input's feature, upscale the image 224x224 resolution with bilinear interpolation. After transforming to the feature space, we apply the cross-channel mixing and cross-scale mixing represent the input with abundant and non-overlapping features. The output of the cross-scale mixing is a feature pyramid consisting of four feature maps at different resolutions (Sauer et al., 2022). total, we use eight discriminators (four for EfficientNet features and the other four for DeiT-base features) for GAN training.
Following CM, we apply Exponential Moving Average (EMA) update sg(() by
sq) + stopgrad(uss(())
e)
However, unlike CM, we find that our model bestly works with μ = 0.999 or μ 0.9999, which largely remedy the subtle instability arise from GAN training. Except for the unconditional CIFAR-10 training with Φb, we set y to be 0.999 as default. Throughout the experiments, we use m'in 0.002, maax 80, / 7, and data 0.5.
D.2 EVALUATION DETAILS
For likelihood evaluation, we solve the PF ODE, following the practice suggested in Kim al. (2022b) with the RK45 (Dormand & Prince, 1980)
5.
FID (Lower the better)
GAN models
Score based models
ODE solver of tol 1e 3 and Imin 0.002.
as
CTM (ours)
DCPM
Throughout the paper, we choose & = 0 otherwise stated. In particular, for Tables and the sample quality metrics based on either the one- step sampling of CM the 7 sampling for NFE 2 case. For CIFAR-10 10, we calculate the FID score based on Karras et al. (2022) statisties. and Figure 15 summarizes the result. For ImageNet, we compute the metries following Dhariwal Nichol (2021) and their pre-caleulated statistics. For the StyleGAN-XL ImageNet result, we recaleu- lated the metrics based on the statistics released by Dhariwal & Nichol (2021), using StyleGAN-XL's official checkpoint
CDIFEE
25
report
CTM (MFF))
FRM
CTM NFF
13P
10
10?
NFE
(Lawer the faster)
Figure 15: SOTA on CIFAR-10. Closeness to the origin indicates better performance.
For large-NFE sampling, we follow the EDM's time discretization. Namely, if we draw n-NFE samples, we equi-divide [o, 1] with points and transform it (say {) to the time scale by (o (a 1/ a mux )/). However, we emphasize the time discretization for both training and sampling is modeler's S choice.
min
E ADDITIONAL GENERATED SAMPLES
26

---

Published as a conference paper at ICLR 2024
EDM
CTM
CM
EDM
CTM
CM
EDM
CTM
CM
EDM
CTM
CM
INFE 1I
D, INFE
(a) Tench
(b) Tree frog
Elephant
Kimono
Figure 16: Uncurated sample comparisons with identical starting points, generated by EDM (FID
2.44) with NFE 79, CTM (FID 2.19) with NFE 1, CTM (FID 1.90) with NFE and CM (FID 6.20) with NFE 1, on (a) tench (class id: 0), (b) tree frog (class id: 31), (c) elephant (class id: 386), and (d) kimono (class id: 614).
(a) W/o classifier-reeection sampling (NFE 1) (b) classifier classifier-rgeectionn sampling (avg. NFE
Figure 17: Random samples (Siberian Husky) (d) with and (c) without classifier-free sampling.
27

---

Published as a conference paper at ICLR 2024
F THEORETICAL SUPPORTS AND PROOFS
F.1 PROOF OF LEMMA 2
Proof of Lemma 2.
The reverse-time SDE is
dxy-r (T T) log Pr (x 9(T T)
where the forward-time SDE is given by
dx, = (
The reverse-time PF-ODE thus becomes
1 dxt 29 (T T)V log PT (x)
Therefore, by integrating from T too T (s t), obtain
T-s
T 2 2 (T - -) -o
x, - X
By change-of-variiee with T - T, the equation is derived to be
[ 2 (w) Pu
X₂ X
With g(u) 2u and the Tweedie's formula V (oggp((x.) paper:
]|a/ X=
derive Eq. (2) in
Xx E[x du. /-
X, G(x,,t, - Xx
u
Now, we derive the following equations:
8
8
Xu
E8 [x] x.
t) = 7x + (1
b |xx:
du
X1 + (1
,)(x +
8
{) /. x, TE. [x] xx]] due
(1
2
$
t
Xu xx [xu] dul 1-
X4 + (1
8 [ +
-
x + (1
8
t (x(, s),
where g(x,,t, 1,,) = X, + t-x Se x. -[xix=l du. As the score, V log (x)), is integrable, the Fundamental Theorem of Calculus applies, leading to
X.
E(xu X.]
lim g (X,, t, 8) = Xx + lim *-p
du
8-22 t
u
Xx E Exxol [x]]
Xt
E[xu||xi]
F.2 PROOF OF THEOREM
Proof of Theorem 1. Define Te-te as the oracle transition mapping from t to via the diffusion process Eq. (6). Let (.) represent the transition mapping from the optimal CTM, and T, (-) represent the transition mapping from the empirical probability flow ODE, Since all processes start at point T with initial probability distribution Pr and T (.) = T.,,.)), Theorem (Chen et al.,
2022) and Tr- Pt from Proposition 9 tell us that for L.
28

---

Published as a conference paper at ICLR 2024
Drv (Ti+ohPi, T(T.ApP) = Drv (T~+atPe, T(,atri) O(t - 8).
(9)
Drv
To-TT VI-
T T T- To-07 T
T
V 1- HpT
v
11--88--8
T
(a) DDrv (T-OT 1- T v 1-,2 PPr, T /1 T (r-v+-+0 | ,
+DTV T T V 1- - T T 1-121 LPT, T T
T
v 1-73 PT
( Drv TT T TI- + DrV
T
1 HpT, T 11 APT
@ Drv (T-odpPe T T(SotD)) + DTV TT-N/I-SH PP T
T - kpt
@ -O(V)) - v1-
Here (a) is obtained from the triangular inequality, (b) and (c) are due to T 31 T 1-11---
Ti and TIrr = Pr from Proposition 9, and (d) comes from Eq. (9).
F.3 PROOF OF PROPOSITION 3
Proof of Proposition
Consider a LPIPS- metric, denoted as :). determined by a feature extractor F of Pdnta- That is, d(x,y) y) |F(x) F(y)ll, for q 1. For simplicity of notation,
denote y as O. Since C (0: Φ) = 0, it implies that for any X n E [1, and m E [1,n]
Gool Xt 1, 0)) (GoolGG(( (x,,, i bas bm), tm 0))
(10)
Denote
tmi Φ), tm, 0)). Cn,m == FGoo In km, 0)) F(Go(G(X bni
Then due to Eq. (10) and G is an ODE-trajectory function that G(Xen111 tn+1 tm; Φ) =
G(Xt,, bny imi )b,, we have
ea+1, = F Goo tnt1, mm), , fm 0)) F(Go(G( In tm Φ), tm 0))
GGolGGo(x 1 ln 1, lm; 0)) F(Go(G(X),, by bms Φ): bm 0))
= =F(Go(Go(x) eb In 1m) tm. 0) ((Go(Go(xi.. In tm tim 0))
=F(Go(Go(Xe.. Eny bm)s tm. 0)) - FGoo(G(X tan tmi Φ), tm 0))
=F(Go(Go(xf) tyi tm), tm. 0)) F FGooGGo (Xi,, tux tm), tm 0)) + Cn,m
Therefore,
||a41,m l - F(Go(Go(x)) (antm),ttn 0)) F(Go (n; bm), tm 0)) |1 + +||en,all
L xx Xt, 1, + +||en.mlly
= =O((tn+1 - tn )rt ||en,mlll
Notice that since Go (Xim bms lm) = Xt.m G(xti bms bmi $), em.m = O.
So we can obtain via induction that
R-1
|en41, I 110mm
+ Σ O((t41 tz) km
29

---

Published as a conference paper at ICLR 2024
E O11x+11- - 1x) 1t)
- (((ANt)))(t - tmm)
Indeed, an analogue Proposition 3 holds for time-conditional feature extractors.
Let d.(..)) be LPIPS-like metric determined by time-conditional feature extractor F, That is,
dx(x. y) [x)( Fi(y)) l for q 2 1. We can similarly derive
sup dem (Go(x, tni bm), Gixx, bni tmi x6) 10000((AA) (((ANt)))(t, -mm. XERD
F.4
PROOF OF PROPOSITION 5
Proof of Proposition 5.
We first prove that for any E [0, and S t, as N oc,
sup ||Gox (Gay (x, t, s) 0), Gox (G(x, L,; 8, 0)]|,
(11)
XERD
We may assume (tm) so that im = s, in t, and m+1 - 8, bn+1 as A Nt - oc.
sup Supl/Gay (Go Xx L, 8), 8, 0), Ge [G
0)112
x
ssup |Gew (Gew x t, s). 0), Gow GG (x, t fm+ Φ), fm+ 0)1l
+ sup || (Gex (x, 1nt1, 1i )) 0), Ge, (G(x, (n 1, 4m+ Φ), m 0)
x
+ sup |Ge (G(x, 11 1m+1: Φ), 11 0), Ge. (G f, 8; s) s. 0)llz
Since both G and Gey are uniform continuous F [0, T [0, T). together with Proposition we obtain Eq. (11) as Aut + oc.
In particular, Eq. (11) implies that when N + sup Goo (Go (x, 0), - Ge. (G((x,T,,0 0), (,,
sup Gooy (x, T, 0) G(x. +))1 +
This implies that Pew (.), the pushforward distribution pr induced by Gev T)O,, converges in distribution to (-). Note that since {Ge bN is uniform Lipschitz
||Go(x,t,s) Go t, s)(l2 - L |xx - x1112 for all x, E RD t, [O,T], and O
{Gox lN is asymptotically uniformly equicontinuous. Moreover, {Gow JN is uniform bounded in ON. Therefore, the converse of Scheffé's theorem (Boos, 1985; Sweeting, 1986) implies 0 N
that (1) (.) p-)(-)
Similar argument be adapted to prove Iee (-) Pdata (-) + 0 as N + oo if the regression target Ps(.)) is replaced with Pdiaa (-).
F.5 PROOF OF PROPOSITION 6
Lemma 10. Let f: R [0, D be a function which satisfies the following conditions:
(a) [(,0) 1) is Lipschitz any E [0, T): there is a function L(() 2 0 so that for any I. E [0. 7] and x,y y E DRD
(x t) f(y,))l < TL(1) |X yll,
(b) Linear growth x: there is a L integrable function Mt) so that for any t [0, T] and x E R
|| 1) < M() ||x]]).
30

---

Published as a conference paper at ICLR 2024
Consider the following ODE
(T) = f(x(T), on [O,
(12)
Fix a E [0, the solution operator T of Eq. (12) with initial condition X2 is defined as
+xxi(()) = x I. (x)(::));TTdd 8 T.
(13)
Here x() denotes the solution time starting from the initial value Xt- Then T is an injective operator. Moreover, (s) RD + RD is bi- bi-Lipschitz; that is, for any X2, Xx= E IR
-((--t | x111 - [ - [x)l S (((-s)) ||x x
(14)
Here : Su L(t) cc. In particular, x, 7 & (sx, # (s) all [t, T
Proof of Lemma 10.
Assumptions (a) and (b) ensure the solution operator in Eq. (13) is well-
defined by applying Carathéodory-tye x(s;x,) x,) We need to prove that for any distinct initial values x, and &, starting from f. T[x,] # 7 [x T[x.[s0) For E [lo, 80
global existence theorem (Reid, 1971). We denote (s) as
Suppose on the contrary that there is so E [t, T] so that T [xxx](s0)
consider y(s; x.) := x(1 + 80 8; xx) and (88 == x(1o + 80 s; kt). Then both y (s; x,) and
&s; satisfy the following ODE
(y'(s) - f(y(s), $ E sol (y(t) = (x) -
(15)
Thus, the uniqueness theorem of solution Eq. (15) leads to y(so: y(sp: xt), which means X₁ = x. This contradicts to the assumption. Hence, T is injective.
Now we show that (s) R" + R is bi-Lipschitz for any s s [t, T. For any Xx, X, RR.
[ss) - [+)1( = ||x(s; xx) - x(s; 81)112
< ||xe xell2 + [.
||x &ell2 + L J. [x((:xx)) - x(T;
By applying Gröwnwall's lemma, we obtain
[7|x, - [s) = [xx(sxx,) - 8(s; 81) S cl(c-t) ||x: - -xell2
(16)
On the other hand, consider the reverse time ODE of Eq. (12) by setting T r(u) t u. (u) == x((++-u - and hy)(),, u) == - (yy u). then satisfies the following equation y' y'(u) h(y(w),, u), u E ().
(17)
Similarly, we define the solution operator to Eq. (17)
Sly.](): : [ h(y(u,y)): u) du.
(18)
Here y, denotes the initial value of Eq. (17) and y() is the solution starting from y., Due the Carathéodory global existence theorem, the operator Sl-](s) is well-defined and
s[x(s;x)](s (ss) Xx, {[x(s) x
For simplicity, let y: == x;) and yt:= x(s; Also, denote the solutions starting from initial values Ye and ýr as y(u and { respectively. Therefore, using similar argument, obtain ||xe &ell2 [|8(y=](s) -
{()11())11
< ||x(s:x))= -((s:x))l2 + J. [I(r(uyy)) u)
w)ll du
< |x((;xx) - x(s;xx)ll + L J." [y((uy))) - -5(u:))))l du.
31

---

Published as a conference paper at ICLR 2024
[s) 7 + J. Ily( S(u du.
By applying Gröwnwall's lemma, we obtain
Xx &2 22 S el(o-1) [T(x)](s) T[&:]())112
Therefore,
e- --((x-0) ||x xell2 S [7 ())1
Proof of Proposition
With the definition G(x;; t, s; Φ). we obtain
G(x,, t, s; Φ) = X + (1 ) t, Φ)
X+ {" Xu D& u) du.
Here, t,; Φ)) + - Se
X_ Do (x) du. Thus, the result follows by applying Lemma 10
the integral form of G(x², L, Φ).
F.6 PROOF OF PROPOSITION 7
Lemma 11. Let X be random vector on Ie. and h: IRD RR be bi-Lipschitz mapping with Lipschitz constant O; namely, for any x, RRD
L ||x yll2 I |1(xx) - h(y)ll2 - L yyll-
Then
L Var() < Var(h(X)) S VaV(
Proof of Lemma 11. Let Y be an i.i.d. copy of X. Then 1(X) and h(Y) are also independent.
Thus, cov(X,Y) = and cov(h(X), h(Y)) 0.
2ar (A(X)) = Var (h(X) h(Y))
= E (M(X) ( -(())) (E (N(X) h(X)1)?
(19)
Since h(X) and h(Y) are identically distributed, [(((X) b(Y) - 0.
Thus, by Lipschitzness of h
Var (h(X)) E (b(()) h(X))²]
(20)
< I.P. (X r)
Var (X).
The final equality follows the same reasoning as in Eq. (19), Likewise, we can apply the argument from Eq. (20) to show that
2aa (h(X)) EE (1(x [
2 I. 2B (X y)
22 22-2206 Var (X).
Therefore, T,- Ta- (X) & Var (X) Ta (X).
32

---

Published as a conference paper at ICLR 2024
Proof of Proposition 7. For any E N. since Gg- (Xnx In v1 72 and Zn+1 are indepen- dent,
Var (Xx+1) Var Go (Xn, tn: / tn42) + Var (Zx+1)
=aar Ge (Xnz In: v 711-))++1) + 22 (In+1).
(21)
Proposition 6 implies that Gø- (:tt: tn /1 y is bi- bi-Lipschitz and that for any x.
(tn (nt1: ||x yl12 < Go- (x, in / (2t+11 Go. (y, tn V 22t
S int1: 7) |x yl12:
(22)
where ((bn; (n+1, = exp 21 - - ((b+1)))))))) Proposition follows immediately from
the inequalities (21) and (22).
F.7 PROOF OF PROPOSITION 9
Proof of Proposition 9. {pr is known satisfy the Fokker-Planck equation (Oksendal, 2003) (under some technical regularity conditions). In addition, we can rewrite the Fokker-Planck equation of {pr l the following equation (see Eq. (37) in (Song et al., 2020b))
Öp: = div (Wep.), in (0, x R
(23)
ôt
where W == - log
Now consider the continuity equation for μr defined by W.
Öμ4 div (Wit) in (0, x RP.
(24)
a
Since the score V logg p. is of linear growth in and upper bounded by summable function in the vector field We == -tV log pe: [0. x R + R² satisfies that
sup W. (x)x)] + Lip(W.,, K) dt <000 xeR
for any compact set K C RR. Here Lip( Wt, K) denotes the Lipschitz constant W, on K.
Thus, Proposition 8.1.8 of (Ambrosio et al., 2005) implies that for pr-a.e. the following reverse time ODE (which is the Eq. (6)) admits a unique solution on [0, T
d X2(x) = W. (X)(
(25)
(x(()= x.
Moreover, Pt = Xxipr, for [0,7]. By applying the uniqueness for the continuity equation (Proposition 8.1.7 of (Ambrosio et al., 2005)) and the uniqueness of Eq. (25), we have Pt Mr
XitPr= = Tr- for E [0, Again, since the uniqueness theorem with the given pr, we obtain Pa Ep for any t [0. T and E [0.
33


Reviews:
[{"review_id": "GEDOTM0q59", "reviewer": "Reviewer_6usj", "summary": "The authors extend the recently introduced Consistency Models framework. Consistency Models were proposed as a distillation method to accelerate the sampling of diffusion models. They are trained to predict the output of the Probability Flow ODE trajectory. The authors of this work extend this method by training a model that can predict any intermediate point of the Probability Flow ODE trajectory. This has two benefits: i) the trained model can recover the conditional expectation (and the score) when the starting and the final time of the Probability Flow ODE are chosen to be close, ii) the framework provides a natural way to perform multistep sampling. The authors demonstrate successfully that their method works by achieving a new SOTA on CIFAR-10 and ImageNet.", "questions": "See weaknesses above. Also, why would the student outperform the teacher? And why this method gives better results that Consistency Models when s=0?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, I think this is a great submission. \n\n- The topic of the paper is timely. Consistency Models have been recently proposed and they offer a promising solution for one-step generation with diffusion models.\n\n- The proposed method is simple, intuitive, and novel. \n\n- This framework provides a natural way to perform multistep sampling. It mitigates a significant limitation of Consistency Models which is that the performance was not increasing significantly even if more computational budget was available.\n- The method has great empirical success -- it achieves a new SOTA for single-step generation in ImageNet and CIFAR-10.\n- The comparison with the baselines is thorough.\n- The presentation of the method is clear. I particularly liked Figure 2.", "weaknesses": "I don't see any major issues, but there are a couple of small concerns.\n- The presentation of the paper could be improved a little. Figure 1 is not very helpful I think -- I would advocate making Figure 2 the main Figure of the paper. There are some typos that could be fixed with more careful reading. For example, there is a parenthesis missing in the citation at the top of page 3.\n- In the definition of the $g$ function (Lemma 1), I think the integral bounds $s, t$ have to change order (or a minus sign needs to be added in front of the integral).\n- Some of the mathematical details are unnecessarily complicated or a little sloppy. For example, the fact that $G$ can be written in terms of $g$ is trivial and hence it should not be part of the lemma. What Lemma proves is that at the limit of $s\\to t$, $g$ becomes the conditional expectation. This can be easily seen by approximating the integrand as being constant in $[t, s]$ -- similar to how the first-order ODE solvers operate. \n- Following this point, I think being able to approximate the conditional expectation with the same network is not really important. For this method to work, we need to start from a pre-trained score model anyway. Additionally, the proposed method for learning this doesn't actually work because of numerical instabilities, and hence the authors resort to adding an additional loss term.\n- Section 3.2 is a little bit confusing. The notation added is too much and makes it hard to follow what's going on. I think (?) that the main point is that another network is used for the ODE solving and the student network learns these predictions. If that's the case, this needs to be a little bit more clear.", "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": "yQNngj9j6C", "reviewer": "Reviewer_csVy", "summary": "The work proposes a new work that is able to train generative models from scratch and distill a faster generative model from teacher models. At its core, CMT tries to learn an ODE transition function, which is achieved by minimizing a combination of soft-matching loss, and DSM loss. Moreover, the authors suggest additional GAN to treat numerical accuracy for perceptual quality. The work achieved SOTA performance with the help of a pre-train teacher model and additional gan discriminator.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- SOTA performance\n- Theorem 2 and related analysis and Classifier-Rejection Sampling is interesting and non-trivial. They are new and novel to me.", "weaknesses": "Many of the losses and concepts are not entirely new, and there are some existing works that share similar ideas. Since this work achieves state-of-the-art (SOTA) performance, I encourage the authors to include more experimental details (see below) and empirical analyses in the main paper.\n\nSome recent SOTA distillation or generative works are difficult to reproduce, such as the Guided Distilled Stable Diffusion (Meng et al.). The authors also encountered difficulties in reproducing CD, even when the source code is available. I would like the authors to discuss the challenges and unstable factors associated with reproducing CD, and whether similar factors will affect CMT. Have the authors encountered any stability issues?\n\nThe authors mentioned that the weights of GAN are tuned similarly to VQGAN. Can the authors provide concrete implementation details?", "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": "pNPgB3LlMh", "reviewer": "Reviewer_RDDt", "summary": "This paper proposes the Consistency Trajectory Model (CTM) that predicts any final time from any initial condition along the probability flow ODE of diffusion models. It can sample high-quality samples in one step while keeping the ability to access the score. The authors also introduce $\\gamma$-sampling method, which can consistently improve sample quality with more steps. The experiments show that CTM can achieve SOTA sampling results in one step.", "questions": "Listed in the weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well written, with valuable insights into the proposed method. It also clearly compares to related works, positioning it well in the literature. \n2. The idea is natural and nicely connects the score-based and distillation methods. \n3. SOTA FID results in one step while maintaining the flexibility of accessing the original score. \n4. The provided ablation studies are helpful in understanding the effect of each design component.", "weaknesses": "1. While adversarial training greatly improves FID, it might lead to poor mode coverage, which is a known issue of GAN. Can you measure the mode coverage of CTM on ImageNet64? For example, reporting Recall in Table 2. \n2. CTM loss brings a lot more additional training cost compared to other distillation methods. It takes at least 4 model forward passes of CTM, multiple steps of the ODE solver at each training step, the cost of computing DSM loss, and also the cost of training the discriminators. Can you provide a detailed description of how many model forward passes are used per training step? Can you also provide the total training time in Table 3 and compare it to the cost of training the teacher model?\n3. Training CTM seems to require many tricks like the warm-up scheme for GAN loss and a particular time sampling scheme as described in appendix D.1. But there is little ablation study about them. Can you elaborate more on how you design these training strategies and how important they are? \n4. Minors:\n\t1. Why is the retrained CM much worse than the official report? \n\t2. The official report of DFNO is 3.78. \n\t3. The FID of Rectified Flow 4.85 is from 2-Rectified Flow with distillation. It could be more concrete about which Rectified Flow model is compared in the table.", "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": "ZpW7I0sUQn", "reviewer": "Reviewer_DUTv", "summary": "This paper introduces the Consistency Trajectory Model (CTM), which extends the concept of Consistency Models (CM) and score-based diffusion models. CTM is a single neural network that can efficiently output scores in a single forward pass, allowing for flexible traversal between any initial and final time points in a diffusion process. CTM combines adversarial training and denoising score matching loss to achieve competitive performance in CIFAR-10 and ImageNet-64 datasets. Additionally, CTM accommodates various diffusion model inference techniques, including exact likelihood computation.", "questions": "1. It would be great if the authors could provide further motivation for using such complicated training objective functions with pretrained model over just training GANs.\n\n2. It would be helpful if the author could provide the wall-clock of the training procedure. It seems like the training will be even more expensive than the diffusion model or GAN.\n\n3. More comparisons with CMs are needed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a novel way to distill dynamical generative models such as Diffusion Model. To the best of my knowledge, the proposed method is novel and yields appealing results in different datasets. The experimental details are provided clearly and the paper is well-written.", "weaknesses": "I have some uncertainties regarding potential weaknesses, but I believe it is valuable to bring this matter to the forefront for discussion. \n\n1. The paper itself presents a novel concept, but the methodology employed appears to be excessively complex. The distillation process involves three networks, and various techniques, such as $\\lambda_{GAN}$ for ensuring stable training, are necessary to attain the results reported in Figure 15. Nevertheless, it is worth noting that this complexity also contributes to the uniqueness and innovation of the paper.\n\n2. One of the reasons for the popularity of diffusion models lies in their scalability with respect to data complexity. In Figure 15, it is evident that the distillation of diffusion models heavily depends on the underlying GAN training structure. This prompts the question of whether it might be more straightforward to solely employ GAN training for a single NFE generation. Essentially, this would involve passing through the computationally expensive pretrained Diffusion Model, followed by the resource-intensive distillation process (as CTM requires simulating trajectories), ultimately converging to the realization that the fundamental component driving performance improvement is the GAN training objective. Given the current state of the art where GANs can achieve superior results, it raises the question of why we necessitate a complex distillation methodology, notwithstanding the novelty of such an approach.\n\n3. Sorry for my ignorance, but I think some concepts need more elaboration. For example,  why the adversarial training can lead to better performance than the teacher model? Why Taylor's expansion is the main reason for the discretization error?\n\n4. To the fair comparison with extremely small NFE, the fair comparison should be Consistency Models instead of EDM.", "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)