Elucidating the Preconditioning in Consistency Distillation
Authors: Kaiwen Zheng
Guande He
Jianfei Chen
Fan Bao
Jun Zhu
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion Models
Distillation
Consistency Trajectory Models
Primary area: generative models

Abstract: Consistency distillation is a prevalent way for accelerating diffusion models adopted in consistency (trajectory) models, in which a student model is trained to traverse backward on the probability flow (PF) ordinary differential equation (ODE) trajectory determined by the teacher model. Preconditioning is a vital technique for stabilizing consistency distillation, by linear combining the input data and the network output with pre-defined coefficients as the consistency function. It imposes the boundary condition of consistency functions without restricting the form and expressiveness of the neural network. However, previous preconditionings are hand-crafted and may be suboptimal choices. In this work, we offer the first theoretical insights into the preconditioning in consistency distillation, by elucidating its design criteria and the connection to the teacher ODE trajectory. Based on these analyses, we further propose a principled way dubbed \textit{Analytic-Precond} to analytically optimize the preconditioning according to the consistency gap (defined as the gap between the teacher denoiser and the optimal student denoiser) on a generalized teacher ODE. We demonstrate that Analytic-Precond can facilitate the learning of trajectory jumpers, enhance the alignment of the student trajectory with the teacher's, and achieve $2\times$ to $3\times$ training acceleration of consistency trajectory models in multi-step generation across various datasets.

Published as a conference paper at ICLR 2025
ELUCIDATING THE PRECONDITIONING IN CONSIS- TENCY DISTILLATION
Kaiwen Zhengil; Guande Hel!; Jianfei Chen', Fan Bao Jun Zhu 'Dept. of Comp. Sci. & Tech., Institute for AI, BNRist Center, THBI Lab 'Tsinghua-Bon -Bssch Joint MI. Center, Tsinghua University, Beijing, China
Shengshu Technology, Beijing Pazhou Lab (Huangpu), Guangzhou, China zkwthu@gmail com; guande waande.hel@@uut.lle 101177ootloo com;
fan. www.ssennngn.......... ai; (jianfeic, dcszj] (eststsinnguaaa edu. en
ABSTRACT
Consistency distillation is a prevalent way for accelerating diffusion models adopted in consistency (trajectory) models, in which student model is trained to traverse backward on the probability flow (PF) ordinary differential equation (ODE) trajectory determined by the teacher model. Precondationing is a vital tech- nique for stabilizing consistency distillation, by linear combining the input data and the network output with pre-defined coefficients as the consistency function. It imposes the boundary condition of consistency functions without restricting the form and expressiveness of the neural network. However, previous precondition- ings hand-crafted and may be suboptimal choices. In this work, we offer the first theoretical insights into the preconditioning in consistency distillation, by elu- cidating its design criteria and the connection to the teacher ODE trajectory. Based on these analyses, we further propose a principled way dubbed Analytic-Precond analytically optimize the preconditioning according to the consistency gap (de- fined as the gap between the teacher denoiser and the optimal student denoiser) on generalized teacher ODE. We demonstrate that Analytic-Precomt can facilitate the learning of trajectory jumpers, enhance the alignment of the student trajec- tory with the teacher's, s, and achieve 2x to 3x training acceleration of consistency trajectory models in multi-step generation across various datasets.
INTRODUCTION
Diffusion models are class of powerful deep generative models, showcasing cutting-edge per- formance in diverse domains including image synthesis (Dhariwal & Nichol, [022]] Karras et al.] 2022), speech and video generation (Chen et al.] 2021 Ho et al. (022)) controllable image ma- nipulation (Nichol et al. 2022. Ramesh et al. 2022 Rombach et al. 2022 Meng et al. 2022b),
density estimation (Song et al. 202 Kingma et al. 2021 Lu et al.
2022a Zheng al. 2023b)
and inverse problem solving (Chung et al. 2022 Kawar et al. (022). Compared to their generative counterparts like variational auto-encoders (VAEs) (Kingma & Welling 2014) and generative adver- sarial networks (GANs) (Goodfellow et al. 2014), diffusion models excel in high-quality generation while circumventing issues of mode collapse and training instability. Consequently, they serve as the cornerstone of next-generation generative systems like text-to-image (Rombach et al. (022) and
text-to-video (Gupta et al. 2023 Bao al. 2024) synthesis.
The primary bottleneck for integrating diffusion models into downstream tasks lies in their slow inference processes, which gradually remove noise from data with hundreds of network evaluations. The sampling process typically involves simulating the probability flow (PF) ordinary ODE back- ward in time, starting from noise (Song et al. 2021c). To accelerate diffusion sampling, various training- free samplers have been proposed as specialized solvers of the PF-ODE (Song et al. 2021a Zhang & Chen 2022 Lu al. 2022b), yet they still require over 10 steps to generate satisfactory samples due to the inherent discretization errors present in all numerical ODE solvers.
Work done during an internship Shengshu;
Equua contribution;
The corresponding author.

---

Published as a conference paper at ICLR 2025
Nata:
Teacher PF-ODE
(F,x,)
(6.x))
Naur Netwark
(3,4.)
Fo
folk..t.s) axFFolxe.t.)) Bazks
Agt
(0,*a)
Consistency Punction (Student)
Figure 1: Consistency distillation with preconditioning coefficients a,
single-step generation of diffusion models are concen- trated on distillation methods (Luhman & Luhman 2021 Salimans & Ho 2022 Meng et al. 2022a
Recent advancements in few-step
Song et al. 2023 Kim et al. 2023] Sauer et al. (023). Particularly, consistency models (CMs) (Song
et al. 2023) have emerged as a prominent method for diffusion distillation and successfully been ap-
plied to various data domains including latent space (Luo al. 2023), audio (Ye al. 2023 and
video (Wang et al. 2023). CMs consider training a student network to map arbitrary points on the PF-ODE trajectory to starting point, thereby enabling one-step generation that directly maps noise to data. A follow-up work named consistency trajectory models (CTMs) (Kim et al., (023) extends CMs by changing the mapping destination encompass not only the starting point but also intermediate ones, facilitating unconstrained backward jumps on the PF-ODE trajectory. This design enhances training flexibility and permits the incorporation auxiliary losses.
In both CMs and CTMs, the mapping function (referred to the consistency function) must adhere to certain constraints. For instance, in CMs, there exists boundary condition dictating that the linear
starting point maps itself. Consequently, the consistency functions are parameterized combination of the input data and the network output with pre-defined coefficients. This approach naturally satisfied without constraining the form expressive-
ensures that boundary conditions ness of the neural network. We term this parameterization technique preconditioning in consistency distillation (Figure T). aligning with the terminology in EDM (Karras al., 2002). The precondi- tionings in CMs and CTMs are intuitively crafted but may be suboptimal. Besides, despite efforts, CTMs have struggled identify any distinct preconditionings that outperform the original one.
In this work, we take the first step towards designing and enhancing preconditioning in consistency distillation learn better "trajectory jumpers". We elucidate the design criteria of precondition- ing by linking it to the discretization of the teacher ODE trajectory. We further convert the teacher PF-ODE into a generalized form involving free parameters, which induces a novel family precon- ditionings. Through theoretical analyses, we unveil the significance of the consistency gap (referring to the gap between the teacher denoiser and the optimal student denoiser) in achieving good initial- ization and facilitating learning. By minimizing a derived bound of the consistency gap, we can optimize the preconditioning within our proposed family. We name the optimal preconditioning under our principle as Analytic-Precond, it can be analytically computed according the teacher model without manual design or hyperparameter tuming. Moreover, the computation is efficient with less than % time cost of the training process.
We demonstrate the effectiveness of Analytic-Precmm by applying it to CMs and CTMS on standard benchmark datasets, including CIFAR-10 10, FFHQ 64x66 and ImageNet 64x64. While the vanilla preconditioning closely approximates Analytic-Precome and yields similar results in CMs, Analytic- Precond exhibits notable distinctions from its original counterpart in CTMs, particularly concerning intermediate jumps on the trajectory. Remarkably, Analytic-Precome achieves 2x to 3 training acceleration in CTMs in multi-step generation.
BACKGROUND
2.1 DIFFUSION MODELS
Diffusion models (Song al. 2021c Sohl-Dickstein et al. 2015 Ho al. 2020 transform d-
2

---

Published as a conference paper at ICLR 2025
dimensional data distribution qu( into Gaussian noise distribution through a forward stochastic differential equation (SDE) starting from 2 ~ qu:
dx. = (t) + g(t)dw
(1)
where t E [0, for some finite horizon T. f.g [0, + R is the scalar-valued drift and dif- fusion term, and wt E Rd is standard Wiener process. The forward SDE is accompanied by a series of marginal distributions {a=)= of (æx [ and f.g are properly designed so that the terminal distribution is approximately pure Gaussian, i.e., Or (T() *O of. An intriguing characteristic of this SDE lies in the presence of the probability flow (PF) ODE (Song et al. 202 dæ, - [f(() - (1)
log ai (2x) []]] whose solution trajectories at time when solved back- ward from time T to time 0, are distributed exactly as qt- The only unknown term Va, log (x is the score function and can be learned by denoising score matching (DSM) (Vincent 2011).
A prevalent noise schedule f 0,g /2t is proposed by EDM (Karras et al. 2022) and fol- lowed in recent text-to-image generation (Esser et al. (024)) video generation (Blattmann et al.) 2023), as well as consistency distillation. In this case, the forward transition kernel of the for- ward SDE (Egn. (I) owns a simple form (( (2o) = (Zv,, I'I). and the terminal distribution dr R (O, I"). Besides, the PF-ODE can be represented by the denoiser function (x(,,t))
dæ; de
xt Dz( t)
(2)
where the denoiser function is trained to predict æn given noisy data Φ, Do te, NO0,
any time f, i.e., minimizing E, B. (2t Jofe, loeg) (f) [ Do (x,, t) Zo [2 for some weighting w(t).
This denoising loss is equivalent to the DSM loss (Song al. 2021c In EDM, another key insight
is to employ preconditioning by parameterizing D& as D x. 1) Cakip (t)æ + Com (z
where
dana
Cakip (1) =
Caut (1)
data 1
Vo +
Fg is a free-form neural network and dama is the variance of the data distribution.
2.2 CONSISTENCY DISTILLATION
Denote 6 as the parameters of the teacher diffusion model, and 0 as the parameters of the student net- work. Given a trajectory (a,; Yf" with a fixed initial timestep e of teacher PF-ODE consistency models (CMs) (Song et al. 2023) aim to learn consistency function fo (xx, t) H. which maps the point Dy at any time t on the trajectory to the initial point x.. The consistency function is forced to satisfy the boundary condition fo(æ, x. To ensure unrestricted form and expressiveness of the neural network, fo is parameterized as
fo(x,t)= Jaa + (t c)2 + daata (t. - 5) Vodam f2
data
Fo t)
(4)
which naturally satisfies the boundary condition for any free-form network Fe(x,, We refer to this technique as preconditioning in consistency distillation, aligning with the terminology in EDM. The student 0 can be distilled from the teacher by the training objective:
(. (aa (wo))) (fo 1), (Solvers s)))
(5)
where w(()) is a positive weighting function, d(: is a distance metric, sg is the (exponential moving average) hoop-gradient and Solver& is any numerical solver for the teacher PF-ODE.
Consistency trajectory models (CTMs) (Kim al., 2023) extend CMs by changing the map- ping destination to not only the initial point but also any intermediate ones, enabling uncon- strained backward jumps the PF-ODE. Specifically, the consistency function is instead defined as fe (æ, t, s) + H., which maps the point 2; at time on the trajectory to the point T, at any
Moor precisely, (s t) = (tip(t) (f) + Coua(t) (o (G (f)an, Caoise Coose(t)): (f)). Since Gin (t) and Croise (t) take effects inside the network, we absorb them into the definition of Fg for simplicity. In EDM, the range of timesteps is typically chosen e 0.002, T 80.
3

---

Published as a conference paper at ICLR 2025
previous time < t. The boundary condition is fo( x, which is forced by the following preconditioning:
fola,t,s) 8 x+ 1-++111 (1-
D1-,))o(z,tt,) Do(z,t,s)
(6)
where (o 1, = Caap (f)= + Cou(t) Fo Fo t, is the student denoiser function, and Fo t, is a frre-ffomm network with an extra timestep input The student network is trained by minimizing
B.g.t:/T.. (x0))))) (#a)g((r=(aa)))
( (Festo)(Jo (Jo( L, 6,,, (ese
8), 8, s,e))
An important property of CTM's precondtioning is that when 8 - t, the optimal denoiser satisfies Dø- (æx, 1, (a,,,, t), i.e. the diffusion denoiser. Consequently, the DSM loss diffusion models can be incorporated to regularize the training of 0, which enhances the sample quality as the number sampling steps increases, enabling speed-quality trade-off.
METHOD
Beyond the hand-crafted preconditionings outlined in Eqn. 44) and Eqn. 6 we seek a general paradigm of preconditioning design in consistency distillation. We first analyze their key ingredients and relate them to the discretization of teacher ODE. Then we derive generalized ODE form, which can induce novel family of preconditionings. Finally, propose a principled way to analytically obtain optimized preconditioning by minimizing the consistency gap.
3.1 ANALYZING THE PRECONDITIONING IN CONSISTENCY DISTILLATION
We examine the form of consistency function fa I, 8) in CTMs, wherein it subsumes CMs special case by setting the jumping destination as the initial timestep Assume fe is parameterized as the following form of skip connection:
fa(æ, t, s) f(t, s)æ + g(t, Da(r,, t, s)
(8)
where Da s) Cti (t) + Ft) (x, t, s) represents the student denoiser function in align- ment with EDM. and ft,, gtt, coefficients that linearly combine and Dø. We identify two essential constraints on the coefficients f and g.
Boundary Condition For any free-form network F6 or De the consistency function Sa must adhere to fa(æ, æ (in-place jumping retains the original data point). Therefore, ʃ and should meet the conditions f(t, 1 and g(t, t) for any time t.
Alignment with the Denoiser Denote the optimal consistency function that precisely follows the teacher PF-ODE trajectory as fo- (x, t, 8), and the optimal denoiser as Dg- (,s) = Lo- (x..)))))) according to Eqn. (8). In CTMs, f and g are properly designed so that the limit ((t.8)
(s t). Thus, the student denoiser at t, i.c. Do(x, 1, ideally aligns with the teacher denoiser Dz- This alignment offers two advantages: (1) Dø(æ,, t) acts valid diffusion denoiser and is amenable to regularization with the DSM loss. (2) The teacher model D& serves as an effective initializer the student D, at s t, implying that Do solely < t is suboptimal and requires further optimization.
Dg. (x, lim,-,: [e=(
g(f..)
Precondionings satisfying these constraints can be derived by discretizing the teacher PF-ODE. Sup- pose the discretization from time t time is expressed as x, f(, s)æ, + g(t, s) (2,, (z t), then f. g naturally satisfy the conditions: the discretization from to f must be 2; = Φ;; as s - t, the discretization error tends O, and the optimal student for conducting infinitesimally small jumps is just D (2 1). For instance, applying Fuler method to the PF-ODE in Eqn. 2 yields:
C2 -
Tt - Dol 7
8 -
+1 (1- -o((2.))
(9)
which exactly matches the preconditioning used in CTMs by replacing D& (I,,,, t) with Do (æ,, t, s). Elucidating preconditioning ODE discretization also closely approximates CMs' choice in
4

---

Published as a conference paper at ICLR 2025
Eqn. (4) For t > have - t. therefore fo in Eqn. approximately equals the denoiser Da. On the other hand, as R 0, CTMs' choice in Equ. ( also indicates fo Do- Therefore, CMs' preconditioning is only distinct from ODE discretization when t is close to e, which is not the case in one-step or few-step generation.
3.2 INDUCED PRECONDITIONING BY GENERALIZED ODE
Based on the analyses above, the preconditioning can be induced from ODE discretization. Drawing inspirations from the dedicated ODE solvers in diffusion models (Lu et al. 2022b Zheng et al.
2023a), we consider a generalized representation of the teacher ODE Eqn. 2 which can give rise to alternative preconditionings that satisfy the restrictions.
Firstly, we modulate the ODE with a continuous function L, to transform it into an ODE with respect to L,æ, rather than Et- Leveraging the chain rule of derivatives, we obtain L. dm, + dla Cr, where dæy can be substituted by the original teacher ODE, resulting in de
d(f,m)))
de
d(Lrz:) Le [(1+
d doog Li, :) 2x -Doo))
dt i
(10)
dl
By changing the time variable from to d log 1, the ODE can be further simplified to
d(L,x))) = Loge(ze t),
go(x)t)) = Do( ((xttt)) t) - (1 11)xx
(11)
dde
where denote 4 dlog L1s, and tx c is the inverse function of A Moreover, L, can be represented by 4 as Le = el Aa lex dA
Secondly, instead using or A as the time variable in d(-) d(.) we can employ generalized time representation
the ODE (i.e., formulate the ODE as dt
d';
Y St, St,dA, dλ, where St is any positive continuous function. This transformation ensures that monotonically increases with respect to X, enabling one-to-one inverse mappings byp Aq. To align d St,, Using 7/2 as the new time
with Tap express S, as efay st,dA dA where denote 8; dA variable, we have dr = Lr, St, d', and the ODE in Eqn. (II) is further generalized to
d(L,z,) dm:
go(x,,
(12)
St
The final generalized ODE in Eqn.
(1) is theoretically equivalent to the original teacher PF-ODE Eqn. (2). albeit with a set of introduced free parameters (4 8t) Applying the Euler method
leads to different discretizations from Eqn. 9
go(x,,))
L873 Lix, = (n. n1)
(13)
St
which can be rearranged as
T. L.S, + (4 - 1)(77 170) 23 + Ne Tt Da(a,,t)
(14)
L.S,
LgS.
Hence, the induced preconditioning can be expressed by Eqn. with novel set of coefficients L.
st, 4:5 1)(m--9)) gg(t, s) -== Originating from the Euler discretization of an L.
equivalent teacher ODE, these coefficients adhere the constraints outlined in Section [.1 under any parameters {4,, 8 / thus opening avenues for further optimization. The induced preconditioning under specific selections 4 st
can also degenerate CTM's case ft, s) & gt.(t. s)
-1 for T..
3.3 PRINCIPLES FOR OPTIMIZING THE PRECONDITIONING
Derived from the generalized teacher ODE presented Eqn. a range of preconditionings is now at our disposal with coefficients f.g from Eqn. (14), governed by the free parameters {4, )7 Our aim is to establish guiding principles for discerning the optimal sets {4, st) thereby attaining superior preconditioning compared to the original one in Eqn. (6).
5

---

Published as a conference paper at ICLR 2025
Table 1: Comparison between different preconditionings used in consistency distillation.
Method
CM Song el al. 2023
BCM Li He 2024
CTM (Kim et al. 2023
Analytic-Precond (Ours)
Free-form
1)
Network
Denoiscr
Do(x,t) (t)(()) Sal (aln,
Do (t)p + Coul (1) Fo(ar, (,
Function
Consisteney
fo(x,)) t) e(t Fl(a,))
fo 8) =(tt, s)a glt, Fo (.,)),)))
fo(n,t,8) 8) f(t,s)a glt, (a(x,
Function
amm
2mm + 18 im +
Le9.
f(t,s)
a (1 s)2
f
L,S, + 1,)( 13))
damil a)
Canlt 8) dmn
g(t, a)
dan +
L,S. + ()((0- m9)
Firstly, drawing from the insights of Rosenbrock-tye sttpbbrock-typeeeee exponential integrators and their relevance in diffusion models (Hochbruck & Ostermann] 2010 Hochbruck et al. |0099] /Zeng et al. (023a),, it is suggested that the parameter Le chosen to restrict the gradient of Eqn. 12 right-hand side term with respect to dee- This choice ensures the robustness of the resulting ODE against errors in Zt. An analytical solution for 4 is derived as follows:
4 argmin Eyy,.)) [II a t) = = E (l( [tr Vv, (D
(15)
d
where d is the data dimensionality, || ||e denotes the Frobenius norm and tr(-) represents the trace of a matrix. Secondly, determine the optimal value of st, we dive deeper into the relationship between the teacher denoiser Da(a,,) 1) and the student denoiser D2 s). As elucidated in Section 3.1 the preconditioning is properly crafted to ensure that the optimal student denoiser sat- isfies Dø- tx, = D t). We further explore the scenario where < t by examining the gap ( t, (2 l) which we refer as the consistency gap. Minimizing this gap extends the alignment of D and Dø- to cases where < 1, ensuring that the teacher denoiser also serves as good trajectory jumper. In the subsequent proposition, we derive a bound depicting the asymptotic behavior of the consistency gap:
Proposition 3.1 (Bound for the Consistency Gap, proof in Appendix A. Suppose there ex- ists some constant C > O that the parameters 8t, ) are bounded by [el, Isel S C, then the optimal student denoiser function Dø- under the preconditioning f(,,8) (1s-+(()--immmm ,t(,,)) L.S.
De-3e satisfies
L.S.
DD (0,, (11 S (t/s)30 -
dg 8, ( ddr
max #STSt
(16)
3C
The proposition conforms to the constraint Dg- (x,, t,t) D (I,, Dø
when t. Moreover,
sidering in local neighborhood of i, by Taylor expansion we have (t/a) 30 20x
1 + Ollogt log s). Therefore, the consistency gap for s E (t - 5, t). when 5 is small, is roughly Minimizing this yields analytic solution for sçi
(=..T) s,g0 222.
5ax.5rst dx.
go( 1) (w.
(x,, sigo t) dd,
By(ar))
dA,
5, = argmin BamiinE,ceo (()))
(17)
Ey(m,)) (x [
We the resulting preconditioning as Analytic-Preconde as St are analytically determined by the teacher & using Eqn. (15) and Eqn. (7). Though 4, St are defined over continuous timesteps, we can compute them on hundreds of discretized ones, while obtaining reasonable estimations of their related terms L,, St, nr- The computation is highly efficient utilizing automatic differentiation in modern deep learning frameworks, requiring less than 1% of the total training time (Appendix B.
Despite the approximation (1/a)=C-1 Rz 1 hold-
Backward Euler Method for Training Stability
ing true in local neighborhoods of i, the coefficient [1/-)a0 in the bound exhibits exponential

---

Published as a conference paper at ICLR 2025
4.5
CIM
62-
44
16.9
(EM))
62
70- (CTM)
4.2
al. (CM)
0.4
(CTM)
34-
4.5
24-
4.4
7/
100
12.5
25
10.0
175
2
(eeuoees Cx100001
peruitor (x10000)
legt
(a) CM
(b) CTM
(c) Coefficients f(t, e) g(t,
Figure 2: Training curves for single-step generation, and visualization of preconditionings for single-
jump CIFAR-I0 (conditional).
behavior when >> 1. In practice, directly applying the preconditioning derived from Eqn. (4) may cause training instability, especially on long jumps with large step sizes. Drawing inspiration from the stability of the backward Euler method, known for its efficacy in handling stiff equa- tions without step size restrictions, we propose a backward rewriting of Eqn. (14) from to I as 2t f(s, ()s, + t(s,t)D&s where f.9 are the original coefficients from Eqn. (14). Rear-
ranging this equation yields V, F(a.1) De 2(a.t) D&, giving rise to the backward coefficients 0(s.t) [(,
[(x,6)
st, T(.s g(t, s)
We summarize different preconditionings in Table [ where we also included a concurrent work called bidirectional consisteney models (BCMs) (Li & Hc 2024) which proposed an alternative preconditioning to CTMs.
4 RELATED WORK
Fast Diffusion Sampling Fast sampling of diffusion models can be categorized into training-fiee and training-based methods. The former typically seek implicit sampling processes (Song et al.) 2021a Zheng et al.] 2024a b) or dedicated numerical solvers to the differential equations corre- sponding to diffusion generation, including Heun's methods (Karras et al. 2022), splitting numerical methods (Wizadwongsa & Suwajanakorn 2022), pscudo numerical methods (Liu et al. 2021) and exponential integrators (Zhang & Chen 2022 Lu et al. 2002b)) Zheng al. [023a] Gonzalez al.
2023). They typically require around 10 steps for high-quality generation. In contrast, training- based methods, particularly adversarial distillation (Sauer et al., (023) and consistency distillation (Song al. 2023 Kim et al.) 2023), have gained prominence for their ability achieve high-quality gen- eration with just one or two steps. While adversarial distillation proves its effectiveness one-step generation of text-to-image diffusion models (Sauer et al. 2024), it is thcoretically less transparent than consistency distillation due to its reliance on adversarial training. Diffusion models can also be
accelerated using quantized sparse attention (Zhang al. 2025a 2024 2025b).
Parameterization in Diffusion Models Parameterization is vital in efficient training and sampling of diffusion models. Initially, the practice involved parameterizing noise prediction network (Ho) et al. 22200) (Song et al. (202cc), which outperformed direct data prediction. A notable subsequent enhancement the introduction of "v" prediction (Salimans Ho 2022), which predicts the veloc- ity along the diffusion trajectory, and is proven effective in applications like text-to-image generation (Esser al. 2024) and density estimation (Zheng et al., (023b). EDM (Karras al. 2022) further advances the field by proposing preconditioning technique that expresses the denoiser function as a linear combination of data and network, yielding state-of-the-at sample quality alongside other techniques. However, the parameterization in consistency distillation remains unexplored.
EXPERIMENTS
In this section, demonstrate the impact of Analytic-Precond when applied to consistency distilla- tion. Our experiments encompass various image datasets, including CIFAR- 10 (Krizhevsky) 2009).
FFHQ (Karras al. 2019) 64x64, and ImageNet (Deng et al. 2009) 64 x44 under both uncon-
7

---

Published as a conference paper at ICLR 2025
4.0-
CTM
3.6-
CTX
8.0
CTM Ours
CTK
6.4
CTM
F10.00
4.4
FID-500
3.4-
FTD Solk
7.0-
4.2-
4.0
Uueedup
Spredap
2x speedui
6.0-
3.8
3.0
6.0-
5.8
10
15
10
Training Iterations [x] 00000
Training Lterations (x10000)
Training Iterations (<10000)
Training Iterations (x10000)
(a) CIFAR-10 (Unconditional)
(b) CIFAR-10 (Conditional)
(c) FFHQ x4 (Unconditional)
(d) ImageNet 64 x64 (Conditional)
Figure 3: Training curves for two-step generation.
ditional and conditionn settings. We deploy Analytic-Precond across two paradigms: con- sistency models (CMs) (Song et al.] 2023) and consistency trajectory models (CTMs) (Kim et al.) 2023), wherein we solely substitute the preconditioning while retaining other training procedures. For further experiment details, please refer to Appendix B
Our investigation aims to address two primary questions:
Can Analytic-Precond yield improvements over the original preconditioning of CMs and CTMs, across both single-step and multi-step generation?
How does Analytic-Precond differ from prior preconditioning across datasets, concerning the coefficients ftt, s) and s)t,
5.1 TRAINING ACCELERATION
Effects on CMs and Single-Step CTMs We first apply Analytic-Precond to CMs, where the consistency function fo t) is defined to map In on the teacher ODE trajectory to the starting point x, fixed time c. The models are trained with the consistency loss defined in Eqn. (5) and on the CIFAR-10 dataset, with class labels as conditions. As depicted in Figure 2 (a). we observe that Analytic-Precond yields training curves similar to original CM, measured by FID. Since multi- step consistency sampling in CMs only involves evaluating fe(æ,, t) multiple times, the results remain comparable even with an increase sampling steps. Similar phenomena emerge in CTMs with single-step generation, as illustrated in Figure 2 (b). The commonality between these two scenarios lies in the utilization of only the jumping destination at c. To investigate further, we plot the preconditioning coefficients f(t, E) and g(t, e) in CMs, CTMs and Analytic-Precome as a function of log L, as illustrated in Figure 2 (c). It is evident that across varying t, different preconditioning coefficients f and exhibit negligible discrepancies when is fixed to This elucidates the rationale behind the comparable performance, suggesting that the original preconditionings for + e are already quite optimal with minimal room for further optimization.
Effeets on Two-Step CTMs We further track sample quality during the training process on CTMs,
particularly focusing two-step generation where an intermediate jump is involved (T - to c)-
The models are training with both the consistency trajectory loss in Eqn. 7 and the denoising
score matching (DSM) loss EpEpaan(w00gee (mo) [wt] (* i, f) shown in Figure 3 across diverse datasets, Analytic- Precood enjoys superior initialization and up to 3x training acceleration compared CTM's preconditioning. This observation indicates the suboptimality of the original intermediate trajectory jumps t + s >>>> e. We provided the generated samples in Appendix C
wo l12). following CTM As
5.2 GENERATION WITH MORE STEPS
Apart from the superiority of CTMs over CMs in single-step generation (Figure 2) another notable advantage of CTMs is the regularization effects of the DSM loss. This ensures that D9(x,, t, t) functions as valid denoiser diffusion models, facilitating sample quality enhancement with ad- ditional sampling steps. To evaluate the effectiveness of Analytic-Precond with more steps, we
CCTMs also propose combine the GAN loss for further enhancing quality, which will discuss later.
8

---

Published as a conference paper at ICLR 2025
10
05
04
60
02
bet
bet
(a) CTM
(h) CIFAR-10
(c) FFHQ 64
(d) ImageNet 64 X64
Figure 4: Visualizations of the preconditioning coefficient g(t, s) for CTM, and for Analytic- Precond under different datasets.
Table 2:
FID results in multi-step generation with different number of function evaluations (NFEs).
NFE
FID
2
5
8
10
2
5
8
10
CIFAR-10 10 (Unconditional)
CIFAR- 10 (Conditional)
CTM
3.83
3.58
3.43
3.33
3.22
3.00
2.82
2.59
2.67
2.56
CTM + Ours
3.77
3.54
3.38
3.30
3.25
2.92
2.75
2.62
2.60
2.65
FFHQ 64 64 (Unconditional)
ImageNet 256
256
(Conditional)
CTM
5.96
5.80
5.53
5.39
5.23
5.95
6.16
5.43
5.44
5.98
CTM + Ours
5.71
5.56
5.47
5.31
5.12
5.73
5.67
5.34
5.43
5.70
employ the deterministic procedure in CTMs, which employs the consistency function to jump on consecutively decreasing timesteps from T' to E. As shown in Table [ Analytic-Preoom brings consistent improvement over CTMs as the number of steps increases, indicating better alignment between the consistency function and the denoiser function.
CTM (NFE-1)
5.3 ANALYSES AND DISCUSSIONS
INPRWID
6.0
Visualizations To intuitively understand the distinctions between Analytic-Precond and the original preconditioning in CTMs, we investigate the variations in coefficients f(t,,), g(,, s). We find that Analytic-Precond yields f(t, s) close to that of CTMs, denoted as fCTM (t, s). with FCM (t, s) - f(t, s)l < 0.03 across various f and s. However, gtt, s) produced by amalytic-Precone tends to be smaller, with disparities of up to 0.25 compared to g CTM (t. This distinction is visually demonstrated in Figure 4 where we depict g(4, 8) as a binary function of log l and log 8. Notably, the distinc- tion is pronounced for short jumps
4.5
4.
4.-
Figure 5: Effects of BCM's pre- conditioning on CTMs.
where s1/t is small.
Comparison to BCMs In a concurrent work called bidirectional consistency models (BCMs) & He 2024), novel preconditioning derived from EDM's first principle (specified in Table BCM's preconditioning also accommodates flexible transitions from to along the trajectory. However, as shown in Figure 5) replacing CTM's preconditioning with BCM's fails bring im- provements in both one-step and two-step generation.
Compatibility with GAN loss CTMs introduce GAN loss to fur- ther enhance the one-step generation quality, employing discrim- inator and adopting an alternative optimization approach akin GANs. As shown in Figure (6, when GAN loss is incorporated on CIFAR-10, Analytic-Prees Precond demonstrates comparable perfor- mance. However, in this scenario, the consistency function no longer faithfully adheres to the teacher ODE trajectory, and one- step generation is better than two-step, deviating from 2- theoretical foundations. Nevertheless, the utilization of Analytic- Precond does not lead to performance degradation.
GAN
3.-
CTM GAN Oun INTE-1)
CT54 GAN Ouzs (NPE-2)
2.5
244
Texnding Stredinna 24100002
Figure 6: Effects of Analytic- Precond with GAN loss.

---

Published as a conference paper at ICLR 2025
(a) Teacher
(b) CTM's
(c) Ours
Figure 7: Visualizations of the trajectory alignment, comparing teacher and 3-step student.
Enhancement of the Trajectory Alignment We observe that our method also leads to lower mean square error (MSE) in the multi-step generation of CTM, when compared to the teacher diffusion model under the same initial noise, indicating enhanced fidelity to the teacher's trajectory. To better in improving trajectory alignment, we adopt a toy example
illustrate the effect of Analytic-Precom where the data distribution is simple 1-D Gaussian mixture [- 1)+ (11, 0.25). In this case, we can analytically derive the optimal denoiser and visualize the ground-truth truth teacher trajectory. We initialize the consistency function with the optimal denoiser and apply different preconditionings. As shown in Figure 7 our preconditioning produces few-step trajectories that better align with the teacher's ' and yields more accurate final distribution.
CONCLUSION
In this work, we elucidate the design criteria of the preconditioning in consistency distillation for the first time and propose novel and principled preconditioning that accelerates the training of CTMs in multi-step generation by 2x 3x. The crux of our approach lies in our theoretical in- sights, connecting preconditioning to ODE discretization, and emphasizing the alignment between the consistency function and the denoiser function. Minimizing the consistency gap fosters coor- dination between the consistency loss and the denoising score-matching loss, thereby facilitating speed-quality trade-offs. Our method provides the first guidelines for designing improved trajectory jumpers on the diffusion ODE, with potential applications in other types of ODE trajectories such as the dynamics of control systems robotic path planning.
Limitations and Broader Impact Despite notable training acceleration in multi-step generation, the final FID improvement is relatively insignificant. Besides, Analytic-Precond fails to differ from previous preconditionings on long jumps, resulting in comparable performance in single-step gen- eration. Achieving accelerated distillation generative modeling may also raise concerns about the potential misuse for generating fake and malicious media content. Furthermore, it may amplify undesirable social bias that could already exist in the training dataset.
ACKNOWLEDGMENTS
This work supported by the National Natural Science Foundation of China (Nos. 62350080, 62106120, 92270001), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University; Zhu was also supported by the XPlorer Prize.
REFERENCES
Fan Bao, Chendong Xiang, Gang Yue, Guande He, Hongzhou Zhu, Kaiwen Zheng, Min Zhao, Shilong Liu, Yaole Wang, and Jun Zhu. Vidu: a highly consistent, dynamic and skilled text-to- video generator with diffusion models. arXiv preprint N0X::2405.00233, 2024.
Andreas Blattmann, Tim Dockhor,, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models large datasets. arXiv preprint arXiv:231..15127 2023.
10

---

Published as a conference paper at ICLR 2025
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. Wave- grad: Estimating gradients for waveform generation. In International Conference on Learning Representations, 2021.
Hyungjin Chung, Jeongsol Kim. Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh Interna- tional Conference on Learning Representations, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255. IEEE, 2009.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems, volume 34, pp. 8780 10000087944-- 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403. 032006 2024.
Martin Gonzalez, Nelson Fernandez. Thuy Tran, Elies Gherbi, Hatem Hajri, and Nader Masmoudi. Seeds: Exponential sde solvers for fast high-quality sampling from diffusion models. arXiv preprint arXiv:2305. /4267, 2023.
lan Goodfellow, Jean Pouget-Abadic Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, volume 27, pp. 2672-2680, 2014.
Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprini arXiv:2312.06665 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840 6851, 2020.
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303, 2022.
Marlis Hochbruck and Alexander Ostermann. Exponential integrators. Acta Numerica, 19:209-286, 2010.
Marlis Hochbruck, Alexander Ostermann, and Julia Schweitzer. Exponential rosenbrock- type meth- ods. SIAM Journal on Numerical Analysis, 47(1):786 803, 2009.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410 2019.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Advances in Neural Information Processing Systems, 2022.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems, 2022.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning prob- ability flow ode trajectory of diffusion. In The Twelfth International Conference on Learning Representations, 2023.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In
Advances in Neural Information Processing Systems, 2021.
11

---

Published as a conference paper at ICLR 2025
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Liangchen Li and Jiajun He. Bidirectional consistency models. arXiv preprint arXiv:2403.1800 2024.
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations, 2021.
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 on Machine Learning, pp. 14429-14460. 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. In Advances in Neural Information Processing Systems, 2022b.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 2021.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthe- sizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04378; 2023.
Chenlin Meng, Ruiqi Gao, Diederik P Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In NeurIPS 2022 Workshop on Score-Based Methods, 2022a.
Chenlin Meng, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. SDEdit: Image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2022b.
Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Megrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. In International Conference on Machine Learning, pp- 16784-16804. PMLR, 2022.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with CLIP latents. arXiv preprint 2XXXi:22004001255, 2022.
Robin Rombach, Andreas Blattmann, Dominík Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations, 2022.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. arXiv preprint arXiv:2311.17044 2023.
Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rom- bach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. arXiv preprint arXiv:2403.1200 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference Machine Learn- ing. pp. 2256-2266 PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021a.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. In Advances in Neural Information Processing Systems, volume 34, pp. 1415-1428, b021b.
12

---

Published as a conference paper at ICLR 2025
Yang Song, Jascha Sohl-Dickstein Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021c.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, pp. 32211-32252 PMLR, 2023.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compu- tation, 23 2011.
Xiang Wang, Shiwei Zhang, Han Zhang, Yu Liu, Yingya Zhang, Changxin Gao, and Nong Sang. Videolem: Video latent consistency model. arXiv preprint urXiv:2312.09109 2023.
Suttisak Wizadwongsa and Supasorn Suwajanakom. Accelerating guided diffusion sampling with splitting numerical methods. In The Eleventh International Conference on Learning Representa- tions, 2022.
Zhen Ye, Wei Xue, Tan, Jie Chen, Qifeng Liu, and Yike Guo. Comospeech: One-step speech and singing voice synthesis via consistency model. In Proceedings of the 31 ACM International Conference on Multimedia, pp. 1831-1839, 2023.
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageatten- tion2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization. arXiv preprint arXiv:2411. 1. 10958, 2024.
Jintao Zhang, Jia Wei, Pengle Zhang, Jun Zhu, and Jianfei Chen. Sageattention: Accurate 8-bit attention for plug-an and-play inference acceleration. International Conference on Learning Rep- resentations (ICLR), 2025a.
Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattn: Accurate sparse attention accelerating any model inference. arXiv preprint 2011::25022.8137, 2025b..
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2022.
Richard Zhang, Phillip Isola, Alexci 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
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Dpm-solver-v3: Improved diffusion ode solver with empirical model statistics. In Thirty-seventh Conference on Neural Information Pro- cessing Systems, 2023a.
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved techniques for maximum likelihood estimation for diffusion odes. In International Conference on Machine Learning, pp. 42363- 42389, PMLR, 2023b.
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint arXiv:2409.02900 02900, 2024a.
Kaiwen Zheng, Guande He, Jianfei Chen, Fan Bao, and Jun Zhu. Diffusion bridge implicit models. arXiv preprini 440iv:2405. /5885, 2024b.
13

---

Published as a conference paper at ICLR 2025
A PROOFS
A.1 PROOF OF PROPOSITION 3.1
Proof. Denote {a,)+== as data points on the same teacher ODE trajectory. The generalized ODE in Eqn. (12) can be reformulated as an integral:
/ ne
1 has (2 )nn
Lg
L =
(18)
where h.s (at, Io (
S. and 9o is defined by the teacher denoiser Do in Eqn. 11 On the other hand, by replacing the teacher denoiser Dø with the student denoiser Dy in the Euler discretization (Eqn. (13). the optimal student 8" should satisfy
L, Lix, (n= -nn) n (2),
(19)
where hgg go are defined similarly he. Io as
ho t),) go(
go(s,t, Do(x,,t, (1 -4)xx
(20)
S₁
Combining the above equations, have
Dg- (2r, - Da(at, t) = Si (at, t,s) - (2t, t))
I,g
L,
S
hs (2(()))t)
(21)
S. I"h he (2n) 1e t)dr
According to the mean value theorem, there exists some T E [t, (xp] satistying
cchp(ar:
dh
(ar (x,,) h (ar,, f) S (m - m/))
(22)
dn
Besides, the derivative dhg can be calculated as dn
dho(wr, dha((r,,7)
dnp
ddr dn//dd,
1 dgo(x,,
d dlog S, go(
(23)
St dd,
dd,
S
L,S,
dga(a,.
-
dA+
S.
where we have used d dog S.
8, and dy: dx.
L, Therefore,
dx.
DDr (D 1,8) Di 0)112 < S.
dgo(xr, T) s.go(arr: (7)
( 17 dn
max 1)5 T)t STSE
ddr
n1 S.
dg
(.g0((ET) R)
/
L Sz, Se
< max #ST
dd
ddr
1. S2
(24)
Since we assumed [4/. [se| I c. according to tx, we have
Lex
= Bir Lig dx < ccld-Ar),
Sta s (c()-2r),
S
S cc(d-Ae)
(25)
L,
S
Sr
Therefore,
Lt, St, St dp
< f ~ 30(2->e) 3c((20 -1
(1/s)
(26)
S
3c
3c
Substituting Eqn. 266 in Eqn.
24
completes the proof.
14

---

Published as a conference paper at ICLR 2025
Table 3: Experimental configurations.
Configuration
CIFAR-10
FFHQ 64x64
ImageNet 64
Uncond
Cond
Uncond
Cond
Learning rate
0.0004
0.0004
0.0004
0.0004
Student's stop-grad EMA parameter
0.999
0.999
0.999
0.999
NN
18
18
18
40
ODE solver
Heun
Heun
Heun
Heun
Max. ODE steps
17
17
17
20
EMA decay rate
0.999
0.999
0.999
0.999
Training iterations
200K
150K
150K
60K
Mixed-Precision (FP16)
True
True
True
True
Batch size
256
512
256
2048
Number of GPUs
4
8
32
Training Time (A800 Hours)
490
735
900
6400
B EXPERIMENT DETAILS
B.1 COFFFICIANTS COMPUTING
At every time f. the parameters l and can be directly computed according to Eqn. (5] and Eqn. (77. relying solely on the teacher denoiser model Dz. The computation of l, involves evaluat- ing tr V m, (zt, f)), which is the trace of a Jacobian matrix. Utilizing Hutchinson' 's trace estima- tor, can be unbiasedly estimated as Ex= w. VVV... Do (x t)v, where obeys d-dimensional distribution with zero mean and unit covariance. Thus, only the Jacobian-vector product (JVP) Do(2,, f)er required, achievable in (d) computational cost via automatic differentiation. Once 4 is obtained, the function 96 (x,,,) = ( () - (1 - 1)),, is determined. The computation of
(m which expands as follows:
8, involves evaluating
db.
dgo(wr,)) () dd:
dDo(cr..) dle
dx;
dd. dde tt - -1 -1--1, 11-1, (1)
dA.
(27)
dD dDo (x
dli
dd, axt ar (1)
where dD2(a..t) can also be calculated in (d) time by by automatic differentiation. dd
For the CIFAR-IO and FFHQ 64x64 datasets, we compute 4, and 5, across 120 discrete timesteps uniformly distributed in log space, with 4096 samples used to estimate the expectation E, For ImageNet 64x64, computations are performed across 160 discretized timesteps following EDM's
scheduling (/1/p + N (t2/ 1/2 ),, using 1024 samples to estimate the expectation Eg.,..) The total computation times for CIFAR-10, FFHQ 64x64, and ImageNet 64x64 8 NVIDIA A800 GPU cards are approximately 38 minutes, 54 minutes, and 38 minutes, respectively.
B.2 TRAINING DETAILS
Throughout the experiments, we follow the training procedures of CTMs. The teacher mod- els are the pretrained diffusion models the corresponding dataset, provided by EDM. The network architecture of the student models mirrors that their respective teachers, with the addition of time-conditioning variable minimizing the consistency loss outlined in Eqn. | and the denoising score matching loss E,E.pon ( (()))(( (2o) (we(() | (a () 20p ( For the consistency loss, use LPIPS (Zhang et al., 2018) as the distance metric d(.,-). which is also the choice of CMs. t and in the con- sistency loss are chosen from N discretized timesteps determined by EDM's scheduling (t1/e + N
input. Training of the student models involves
i (1 )P The Heun sampler in EDM is employed as the solver in Eqn. The number of sampling steps, determined by the gap between and 8, is restricted to avoid excessive training time. For CIFAR-10 10 and FFHQ 64x64, we select N 18 and the maximum number of sampling steps as 17, i.c., not restricting the range of jumping from t to s. For ImageNet 64x64, we set N 40 and the maximum number of sampling steps 20, that the jumping range most half
15

---

Published as a conference paper at ICLR 2025
of the trajectory length. sg(() in Equ. 7 an exponential moving average stop noppgradient version of 8, updated by
(g) (top-gracieet(xss(((( + (1-
(28)
We follow the hyperparameters used in EDM, setting Ginn = 0.002, max T = 80.0, data 0.5 and p = 7. The training configurations are summarized Table 3
We run the experiments on a cluster of NVIDIA A800 GPU cards. For CIFAR-10 (unconditional), we train the model with batch size of 256 for 200K iterations, which takes days on 4 GPU cards. For CIFAR-10 1-0 (conditional), which takes days on 8 GPU cards. For FFHQ 64 (unconditional), we train the model with a batch size of 256 for 150K iterations, which takes 5 days 8 GPU cards. For ImageNet 64 x 64 (conditional), we train the model with batch size of 2048 for 60K iterations, which takes 8 days on 32 GPU cards.
train the model with batch size of 512 for 150K iterations,
B.3 EVALUATION DETAILS
For single-step well as multi-step sampling of CTMs, we utilize their deterministic sampling
procedure by jumping along a set of discrete timesteps T to + l 7 ... 1.. - In with the consistency function, formulated as the updating rule Tt, = fo 1: 1-,, The timesteps {4:);= are distributed according to EDM's scheduling (U + & (1 ///2
(mmx)))) )) where Emin e, mmax T. We generate 50K random samples with the same seed and report the FID on them.
N
B.4 LICENSE
Table 4: The used datasets, codes and their licenses.
Name
URL
Citation
License
CIFAR-10
BLLps: p:0::///wa -oronLo w.d.nl/eeddddeeeee Car. han.hlni hLml
(Krzhevsky et al
2009
FFHQ
2019
CC.Y-NC--A
ImageNet
REtps:/ way. mage-nas .or.
Denz et al 201
EDM
/gthup..tgg.......... abbs/edm
2024
CC BCBYNC--AA
(LLos://qitinn com/ ouenss/cunt reyns www.r/ccleeeecceoooooooooo modey Ca
Sont ot al 2023
Apacho-2.0
CTM
(K
MIT
We list the used datasets, codes and their licenses in Table
C ADDITIONAL SAMPLES
16

---

Published as a conference paper at ICLR 2025
(a) CTM (CIFAR-10, Uncond)
(b) CTM + Ours (CIFAR-10, Uncond)
(c) CTM (CIFAR-10, 10, Cond)
(d) (CIFAR-II Cond)
(c) CTM (FFHQ 64 x 64, Uncond)
(0 CTM Ours (FFHQ 64, Uncond)
(g) CTM (ImageNet 64 x 64, Cond)
(h) CTM Ours (ImageNet 64 64, Cond)
Figure 8: Random samples produced by CTM and CTM + Analytic-Precond (Ours) with NFE=2.
17


Reviews:
[{"review_id": "VXbmZ4p8Wq", "reviewer": "Reviewer_ae52", "summary": "This paper proposed a general paradigm of preconditioning design in consistency distillation, which is a common technique used for accelerating the inference time of consistency models based on teacher-student training (i.e., knowledge distillation). Specifically, this paper focused on preconditioning, which is a vital technique for stabilizing consistency distillation. Compared to previous hand-crafted choices of preconditioning, this paper proposed a principled way called \"Analytic-Precond\" to analytically optimize the preconditioning based on the consistency gap associated with the teacher probability flow ODE. Numerical experiments on multiple datasets are included to justify the effectiveness of \"Analytic-Precond\".", "questions": "From line 265-266 the authors mentioned that the parameter $l_t$ in \"Analytic-Precond\" is chosen to be the minimizer of the expected gradient norm $E_{q(x_t)}[\\|\\nabla_{x_t}g_{\\phi}(x_t,t)\\|_F]$ based on earlier work [1]. Would it be possible for the authors to further expand on why such choice ensures the robustness of the resulting ODE again errors in $x_t$? Which section/part of [1] discussed the reason behind such choice?\n\nReferences:\n\n[1] Zheng, Kaiwen, Cheng Lu, Jianfei Chen, and Jun Zhu. \"Improved techniques for maximum likelihood estimation for diffusion odes.\" In International Conference on Machine Learning, pp. 42363-42389. PMLR, 2023.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Complete proofs are included for each proposition in the manuscript.\n2. Extensive numerical experiments are provided to validate the effectiveness of the proposed methodology.", "weaknesses": "Presentation of the manuscript can be further improved by rewriting certain phrases and expanding on some technical details. For instance, the phrase \"CMs aim to a consistency function\" on line 134 might be better rephrased as \"CMs aim to learn a consistency function\". For possible ways of explaining technical details in a better way, one may refer to the \"Questions\" section below.", "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": "dvyOaBMRaP", "reviewer": "Reviewer_6jeK", "summary": "The paper titled \"Elucidating the Preconditioning in Consistency Distillation\" examines consistency distillation techniques for diffusion models, where a student model learns to follow the probability flow trajectory set by a teacher model. This distillation accelerates generation by reducing the inference steps. The paper specifically explores preconditioning, a method that combines input data with network outputs to improve stability during training. Traditionally, preconditioning has been handcrafted, but this paper introduces a theoretically optimized method named \"Analytic-Precond.\" This new approach minimizes the gap between teacher and student denoisers, thereby improving training efficiency and trajectory alignment. Experimental results demonstrate that Analytic-Precond achieves up to 3x acceleration in training across various datasets, indicating its potential in enhancing consistency models for faster multi-step generation.", "questions": "See weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Theoretical Innovation in Preconditioning**: The paper introduces \"Analytic-Precond,\" a novel, analytically derived preconditioning method that theoretically optimizes the consistency distillation process. This goes beyond prior handcrafted preconditionings, offering a principled approach that minimizes the consistency gap between the teacher and student models. This theoretical grounding not only strengthens the methodology but also provides new insights into consistency distillation.\n\n**Significant Training Acceleration**: Experimental results show that Analytic-Precond achieves 2-3x faster training in multi-step generation tasks on standard datasets. This improvement in speed is impactful, especially for resource-intensive applications of diffusion models, as it directly addresses the bottleneck of slow inference that has historically limited diffusion models.", "weaknesses": "- This paper does not provide whether BCM is better than CTM+ Analytic-Precond in terms of FID.\n- Analytic-Precond does not perform better when GAN is incorporated into the CTM. Can the authors provide an explanation or intuition for this?", "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": "SqD0h0NcXW", "reviewer": "Reviewer_yEcg", "summary": "The paper works on _Preconditioning_, a technique used in the consistency distillation of diffusion models to obtain consistency functions that directly satisfy the boundary conditions required by the problem. Preconditioning consists in linearly linking the input to the output of a network. In the literature, the choice of linear coefficients is based on intuition. The paper introduces instead a new analytical method, called _Analytic-Precond_, for setting the coefficients. The method consists in applying a parametric discretization of the probability flow ODE, and then optimizing the parameters by minimizing the gap between the optimal student and the teacher, while keeping the discretization as robust as possible. Finally, some numerical proofs show that the derived result leads to a speed-up in the inference of diffusion models.", "questions": "- The paper focuses on the case where $f=0, g=\\sqrt{2t}$ because of the recent literature. Is the method still applicable for other choices of $f$ and $g$?\n- As far as I understand, the whole discussion depends on finding a good discretization of ODE (2). Both (9) and (13) use first order (Euler) methods. Can we get more insight if we try to use a better integrator?\n\nMinor:\n- Why is $q_T$ on line 118 a 0 mean Gaussian? Do we have some condition on $\\mathbb{E}[q_0]$?\n- The $\\lambda_t$ below equation (11) might be confused with $\\lambda(t)$ in equations (5), (7).\n- x$ and $x_t$ are used interchangeably in the RHS and LHS of the equations, better to be consistent. Example: line 182, line 276.\n- Is the code of the experiments released?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper presents strong mathematical arguments to support the choice of coefficients, including an explanation for the CMT choices that is not just based on intuition as previous methods.\n- The paper shows numerical proofs of the claims made, underlying when _Analytic-Precond_ offers no advantage (single step) and when it does (two or more steps).", "weaknesses": "- The paper might be a little hard to read for who is not familiar with distillation. I personally took a while to grasp the setting and all the notation. For example, $\\phi$ is used many times before definition. It could be worth having a brief discussion about some nomenclature like _teacher_ & _student_.", "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": "Gb38YU9umQ", "reviewer": "Reviewer_TK51", "summary": "Consistency (trajectory) distillation typically uses a network parameterized as $f(x, t, s) = \\alpha_{t, s} F_\\theta(x, t, s) + \\beta_{t, s} x$, with specific expressions for the coefficient (i.e. preconditioners) $\\alpha_{t, s}, \\beta_{t, s}$ so that the boundary conditions are automatically satisfied. This paper proposes an efficient method to find alternative preconditioners that yield improved performance and faster training. They derive the expressions for the preconditioners by re-expressing the underlying ODE in terms using certain additional variables and propose simple objectives whose minimizers can be found analytically to set the values for these variables.", "questions": "How does performance and the actual values for the preconditioners change as we change the number of samples used to estimate Eqs. 15 and 17 (expressions for the preconditioners)? Results in the paper are obtained estimating them for 120 different times, using 4096 samples to estimate expectations. This yields good results for 2 step sampling, at a modest computational cost. This is definitely not strictly necessary, but I think it could be interesting to see how the performance and preconditioner values change as we change the number of samples used. The estimators used have variance and bias, both of which decrease with the number of samples. How small can we make the number of samples used while still retaining good performance? Can we further boost performance using more samples?\n\nThis is all for distillation. Did you think about consistency training? These preconditioning parameters also show up in that case, but we don’t have access to the teacher, so it is unclear how to use the ideas in this work.\n\nWhile multi-step consistency models have been observed to underperform CTMs, it would be nice to have values in some of the results reported. I understand training is exactly the same, so I’m not expecting the new preconditioners to help there. But it would be nice having plain multi-step CMs in some results.\n\nWhat dataset is used for the GAN experiment?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Consistency (trajectory) distillation represents an extremely popular approach for fast generation, solving the main drawback of diffusion models. Finding approaches to improve distillation, either the final performance achieved or its training cost, is extremely relevant to improve the practicality and applicability of large models currently being trained in multiple domains.\n\nThis paper tackles this problem by changing the preconditioning parameters used for the neural network. To my knowledge the procedure described in the paper is novel, and empirical results show that, for CTMs, the proposed approach yields a training speedup of 2x for multiple datasets.\n\nThe final proposed method is quite simple to implement, with closed-form expressions for the preconditioners. (However, these expressions involve nonlinear functions intractable expectations which are estimated with samples, meaning that the actual values obtained have both bias and variance.)", "weaknesses": "The paper deals with the consistency (trajectory) distillation case. In practice, a very relevant alternative that is also widely used (and also uses similar preconditioners for the network) is consistency training. That is, directly training fast samplers without distilling a base model. The paper does not address this problem at all. While I understand this is not necessary, and distillation in itself is an important task, it would be interesting to see whether certain ideas from this work can be extended to consistency training. (Most of the expressions for the preconditioners rely on having a pre-trained model, so unclear how to generalize this approach, if possible.)\n\nThe approach seems to help when using CTM with 2 steps or more. For single step generation performance and training curves pretty much overlap with and without the proposed approach. 1 step generation plays an important role for real time generation, and it would be very interesting to develop methods that can improve training and final performance in that setting as well. Additionally, the fact that the derived preconditioners are essentially the same as the ones naively used by consistency (trajectory) distillation raises the question of whether there are other preconditioners that can be used that might help in this case as well. Again, not exploring this does not reduce the paper’s merit, but I think it is an interesting question too.\n\nThe method does not yield any benefits when used in concert with a GAN auxiliary loss. Using this loss has been observed to lead to improved performance, and indeed, the best results reported in the paper are using the GAN loss, if I understand correctly. The proposed approach does not yield benefits (but at least does not hurt) when using this auxiliary loss.", "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)