Simplifying, Stabilizing and Scaling Continuous-time Consistency Models
Authors: Cheng Lu
Yang Song
Venue: ICLR 2025 Oral (2025)
Keywords: continuous-time consistency models
diffusion models
fast sampling
Primary area: generative models
TL;DR: 2-step continuous-time consistency models reduce the gap to within 10\% in sample quality (FID) compared to best diffusion models

Abstract: Consistency models (CMs) are a powerful class of diffusion-based generative models optimized for fast sampling. Most existing CMs are trained using discretized timesteps, which introduce additional hyperparameters and are prone to discretization errors. While continuous-time formulations can mitigate these issues, their success has been limited by training instability. To address this, we propose a simplified theoretical framework that unifies previous parameterizations of diffusion models and CMs, identifying the root causes of instability. Based on this analysis, we introduce key improvements in diffusion process parameterization, network architecture, and training objectives. These changes enable us to train continuous-time CMs at an unprecedented scale, reaching 1.5B parameters on ImageNet 512×512. Our proposed training algorithm, using only two sampling steps, achieves FID scores of 2.06 on CIFAR-10, 1.48 on ImageNet 64×64, and 1.88 on ImageNet 512×512, narrowing the gap in FID scores with the best existing diffusion models to within 10\%.

Published as a conference paper at ICLR 2025
SIMPLIFYING, STABILIZING & SCALING CONTINUOUS- TIME CONSISTENCY MODELS
Cheng Lu & Yang Song OpenAl
ABSTRACT
Consistency models (CMs) are powerful class diffusion-basee generative mod- els optimized for fast sampling. Most existing CMs are trained using discretized timesteps, which introduce additional hyperparameters and are prone to discretiza- tion errors. While continuous- success has been limited by training instability. To address this, we propose a simplified theoretical framework that unifies previous parameterizations of diffu- sion models and CMs, identifying the root causes of instability. Based on this analysis, we introduce key improvements in diffusion process parameterization, network architecture, and training objectives. These changes enable us to train continuous- time CMs at an unprecedented scale, reaching 1.5B parameters on ImageNet 512x512 12. Our proposed training algorithm, using only two sampling steps, achieves FID scores of 2.06 on CIFAR- 10, 1.48 on ImageNet 64x64, and 1.88 on ImageNet 5 122x112, narrowing the gap in FID scores with the best existing diffusion models to within 10%.
formulations can mitigate these issues, their
INTRODUCTION
Diffusion models (Sohl-Dickstein et al., 2015; Song &
Fréchet Incoption Distance (FI)) (FID)
BlgGAN
Ermon, 2019; Ho et al., 2020; Song et al., 2021b) have revolutionized generative AI, achieving remarkable results
ADM-G
MasACITT
in image (Rombach et al., 2022; Ramesh et al., 2022; Ho
al., 2022), 3D (Poole al., 2022: Wang et al., 2024; Liu et al., 2023b), audio (Liu et al., 2023a; Evans et al., 2024),
and video generation (Blattmann al., 2023; Brooks et al.,
L-VIT--H
2024). Despite their success, significant drawback is their slow sampling speed, often requiring dozens hun- dreds of steps generate single sample. Various diffu- sion distillation techniques have been proposed, includ- ing direct distillation (Luhman & Luhman, 2021; Zheng et al., 2023b), adversarial distillation (Wang et al., 2022; Sauer et al., 2023), progressive distillation (Salimans & Ho, 2022), and variational score distillation (VSD) (Wang al., 2024: Yin al., 2024b;a; Luo et al., 2024; Xie al., 2024b; Salimans et al., 2024). However, these methods come with challenges: direct distillation incurs extensive computational cost due to the need for numerous diffu- sion model samples; adversarial distillation introduces complexities associated with GAN training; progressive distillation requires multiple training stages and is less effective for one two-step generation; and VSD can produce overly smooth samples with limited diversity and struggles at high guidance levels.
VAR-414
VDM-
sCM (ours) MAGIVTE-2
MAR
10
1000
1000
Effective Sumpling Compute
Figure 1: Sample quality vs. effective sam- pling compute (hillion parameters X number of function evaluations during sampling). We compare the sample quality of different mod- els ImageNet 512x measured by FID (J). Our 2-step sCM achieves sample quality comparable to the best previous generative models while using less than 10% of the cf- fective sampling compute.
Consistency models (CMs) (Song et al., 2023; Song & Dhariwal, 2023) offer significant advantages in addressing these issues. They eliminate the need for supervision from diffusion model samples, avoiding the computational cost of generating synthetic datasets. CMs also bypass adversarial training, sidestepping its inherent difficulties. Aside from distillation, CMs can be trained from scratch with consistency training (CT), without relying on pre-trained diffusion models. Previous work (Song & Dhariwal, 2023; Geng et al., 2024; Luo al., 2023; Xie al., 2024a) has demonstrated the

---

Published as a conference paper at ICLR 2025
Figure 2: Selected 2-step samples from continuous time consisteney model trained ImageNet 512x512.
effectiveness of CMs in few-step generation, especially one or two steps. However, these results are all based on discrete-time CMs, which introduces discretization errors and requires careful scheduling of the timestep grid, potentially leading to suboptimal sample quality. In contrast, continuous-time CMs avoid these issues but have faced challenges with training instability (Song et al., 2023; Song & Dhariwal, 2023; Geng et al., 2024).
In this work, introduce techniques to simplify, stabilize, and scale up the training of continuous- time CMs. Our first contribution is TrigFlow, new formulation that unifies EDM (Karras et al., 2022; 2024) and Flow Matching (Peluchetti, 2022; Lipman al., 2022; Liu al., 2022; Albergo al., 2023; Heitz et al., 2023), significantly simplifying the formulation of diffusion models, the associated probability flow ODE and CMs. Building on this foundation, we analyze the root causes of instability in CM training and propose complete recipe for mitigation. Our approach includes improved time-conditioning and adaptive group normalization within the network architecture. Additionally, we re-formulate the training objective for continuous time CMs, incorporating adaptive weighting and normalization of key terms, and progressive annealing for stable and scalable training.
With these improvements, we elevate the performance of consistency models in both consistency training and distillation, achieving comparable or better results compared previous discrete-time formulations. Our models, referred to as sCMs, demonstrate success across various datasets and model sizes. We train sCMs on CIFAR-10 10, ImageNet 64 x64, and ImageNet 512x512, reaching an unprecedented scale with 1.5 billion parameters-tte the largest CMs trained to date (samples in Figure 2). We show that sCMs scale effectively with increased compute, achieving better sample quality in predictable way. Moreover, when measured against ofate-of-the-arr diffusion models, which require significantly more sampling compute, sCMs narrow the FID gap to within 10% using two-step generation. In addition, provide rigorous justification for the advantages of continuous-
2

---

Published as a conference paper at ICLR 2025
time CMs over discrete-time variants by demonstrating that sample quality improves as the gap between adjacent timesteps narrows to approach the continuous- time limit. Furthermore, we examine the differences between sCMs and VSD, finding that sCMs produce more diverse samples and are more compatible with guidance, whereas VSD tends to struggle at higher guidance levels.
2
PRELIMINARIES
2.1 DIFFUSION MODELS
Given a training dataset, let Pa denote its underlying data distribution and aa its standard deviation. Diffusion models generate samples by learning to reverse noising process that progressively perturbs a data sample Zu ~ Pd into noisy version It = 0g + 022, where 2 (0, I) is standard Gaussian noise. This perturbation increases with E [0, TT. where larger i indicates greater noise.
We consider two recent formulations for diffusion models.
EDM (Karras et al., 2022; 2024). The noising process simply sets 1 and L, with the training objective given by Emo...t (t) || (x,, t) 0ol12] where wtt is a weighting function.
The diffusion model is parameterized as f 1) = (1):a + Com (1) Fo (Cin (D),,, Cuoise (4)), where Fg is neural network with parameters 0, and Cakip, Couts Cins and Craise manually designed coefficients that ensure the training objective has the unit variance across timesteps at initialization. For sampling, EDM solves the probability flow ODE (PF-ODE) (Song et al., 2021b), defined by c2 (a,,, ()] starting from æт (N T I) and stopping at æzo
dæ. de
Flow Matching. The noising process uses differentiable coefficients ar and ot, with time derivatives denoted by og and of (typically, oe
1. and a, (). The training objective is given by Ex,,,,t (t) [ t) (a/20 + x where wt) is weighting function and Fş is neural network parameterized by 9. The sampling procedure begins at t = 1 with x1 N(O I) and solves dæ.
the probability flow ODE (PF-ODE), defined by de (a t), from t = 1 to t = 0.
2.2 CONSISTENCY MODELS
A consistency model (CM) (Song et al., 2023; Song & Dhariwal, 2023) is neural network fo t) trained to map the noisy input a directly to the corresponding clean data 2o in one step, by following the sampling trajectory of the PF-ODE starting at Zz. A valid Sa must satisfy the boundary condition, fo (x: 0) = æ. One way meet this condition is to parameterize the consistency model as fo(æ,, () = Caki (4) Jx, + Com (4) (Gn (1)æ,, Cnoise (1)) with (0) = 1 and Cour (0) O. CMs are trained to have con- sistent outputs at adjacent time steps. Depending on how nearby time steps selected, there are two categories of consistency models, as described below.
O(A))
DEN P
Ri-51
Discrete-time CMs. The training objective is defined at two adjacent time steps with finite distance:
Ban, (()))((f(( fo (RI-as, A()). (1)
Figure 3: Discrete-time CMs (top & middle) vs. continuous-time CMs (bottom). Discrete- time CMs suffer from discretization errors from numerical ODE solvers, causing impre- cise predictions during training. contrast. contimuous-time CMs stay on the ODE trajec- tory by following its tangent direction with infinitesimal steps.
where 0- denotes st)pgrad(O),, w(()) is the weighting function, At >> 0 is the distance between adjacent time steps. and o(-,-) is metric function: common choices are E2 loss d(x, 3) = x y 12. Pseudo-Huber loss d(ær, y) / VI y - c for 0 (Song & Dhari- wal, 2023), and LPIPS loss (Zhang et al., 2018). Discrete- time CMs are sensitive to the choice of At, and therefore require manually designed annealing schedules (Song & Dhariwal, 2023; Geng et al., 2024) for fast convergence. The noisy sample H; At at the preceding time step At is often obtained from a
3

---

Published as a conference paper at ICLR 2025
by solving the PF-ODE with numerical ODE solvers using step size At, which cause additional discretization errors.
Continuous-time CMs. When using d(x,, y) = 1|x ylll and taking the limit At O, Song et al. (2023, Remark 10) show that the gradient of Eq. (1) with respect to converges to
dfo- (mc, dt
Volar, w wt)ff (xx,
(2)
where dfe- (dt..) de Vz.fo (zt, dr + a.fo- (xx, is the tangent of fo- at (x,, t) along the trajectory of the PF-ODE dæ, dr Notably, continuous-time CMs do not rely on ODE solvers, which avoids discretization errors and offers more accurate supervision signals during training. However, previous work (Song et al., 2023; Geng et al., 2024) found that training continuous-time CMs, or even discrete-time CMs with an extremely small At, suffers from severe instability in optimization. This greatly limits the empirical performance and adoption of continuous-time CMs.
Consistency Distillation and Consistency Training. Both discrete-time and continuous-time CMs can be trained using either consistency distillation (CD) or consistency training (CT). In consistency distillation, a CM trained by distilling knowledge from a pretrained diffusion model. This diffusion model provides the PF-ODE, which can be directly plugged into Eq. (2) for training continuous- time CMs. Furthermore, by numerically solving the PF-ODE to obtain Zt- from Et, one can also train discrete-time CMs via Eq. (I), Consistency training (CT), by contrast, trains CMs from scratch without the need for pretrained diffusion models, which establishes CMs as a standalone family of generative models in their
right. Specifically, CT approximates 2t-At in discrete-time CMs as Tt-Ar = 21 - reusing the same data xy and noise z when sampling V; = az + 2 In the continuous-tii time limit, as At 0. this approach yields unbiased estimate of the PF-ODE das + a,20 + ofz, 2, leading to unbiased estimate of Eq. (2) for training continuous-time CMs.
3 SIMPLIFYING CONTINUOUS -IMM CONSISTENCY MODELS
Previous consistency models (CMs) adopt the model parameterization and diffusion process for- mulation in EDM (Karras et al., 2022). Specifically, the CM parameterized as (a Cxip(t) (t)æz + Cou (t) Fo (t)xz, Caoise (t)), where F, is a neural network with parameters A. The coefficients (t) Cour (t), (n(t) are fixed to ensure that the variance of the diffusion objective is equalized across all time steps at initialization, and Cnnise()) is transformation of L for better time conditioning. Since EDM diffusion process variance-exploding (Song et al., 2021b), meaning that 2, = 3D + tz, we can derive that Caki((t) = d22((t + oz), (oun(1) = σa t/ /o + f2, and (t) = 1/ /t2 + + (see Appendix B.6 in Karras et al. (2022)). Although these coefficients are important for training efficiency, their complex arithmetic relationships with and ao complicate theoretical analyses of CMs.
To simplify EDM and subsequently CMs, we propose TrigFlow, a formulation of diffusion models that keep the EDM properties but satisfy Caki (t) = cos(t), (t = -a- sin(t), and cm = 1/ (proof Appendix B). TrigFlow is a special case of flow matching (also known as stochastic interpolants or rectified flows) and v-prediction parameterization (Salimans & Ho, 2022). It closely resembles the trigonometric interpolant proposed by Albergo & Vanden-Eiinden (2023); Albergo et al. (2023); Ma et al. (2024), but is modified to account for dd. the standard deviation of the data distribution pa- Since TrigFlow is a special case of flow matching and simultancously satisfies EDM principles, combines the advantages of both formulations while allowing the diffusion process, diffusion model parameterization, the PF-ODE, the diffusion training objective, and the CM parameterization to all have simple expressions, as provided below.
Diffusion Process. Given Co ~ pa(at) and z ~ (O, o the noisy sample is defined as 2;, cos(L)aza sin(()z for & [0, 5]- As special case, the prior sample æ (N,
Diffusion Models and PF-ODE. We parameterize the diffusion model as fo (x,,
Fg (2 Cnoise (t)), where F, is a neural network with parameters 6, and Croise (t) is a transfor- mation of to facilitate time conditioning. The corresponding PF-ODE is given by
dx dt
Dt Cooise (1) (1))
= Fg
(3)
0
4

---

Published as a conference paper at ICLR 2025
Diffusion Objective. In TrigFlow, the diffusion model is trained by minimizing
Dun (0) = Re od Fe
Cnoise (1) (1)-000
-00
(4)
where 0, = cos(t)= - sin(t) is the training target.
Consistency Models. As mentioned in Sec. 2.2, valid CM must satisfy the boundary condition fo(2,0) 0) = æ. To enforce this condition, we parameterize the CM as the single-step solution of the PF-ODE in Eq. (3) using the first-order ODE solver (see Appendix B.I for derivations). Specifically, CMs in TrigFlow take the form of
(t)) fo 1) cos(t)æz, sin(t)or (z), Cnoise (
(5)
where Croise (t) is time transformation for which defer the discussion Sec. 4.1.
STABILIZING CONTINUOUS -IME CONSISTENCY MODELS
Training continuous-time CMs has been highly unstable (Song et al., 2023; Geng al., 2024). As result, they perform significantly worse compared to discrete-time CMs in prior works. To address this issue, we build upon the TrigFlow framework and introduce several theoretically motivated im- provements to stabilize continuous- time CMs, with focus on parameterization, network architecture, and training objectives.
4.1 PARAMETERIZATIIN AND NETWORK ARCHITECTURE
Key to the training of continuous-time CMs is Eq. (2), which depends on the tangent function dfa (wi.) Under the TrigFlow formulation, this tangent function is given by
de
Ft. -= afe (x1, cos(t) oaFo dt
dx: di
dF
sin(t)
+ od
(6)
Od
dl
where represents the PF-ODE, which either estimated using a pretrained diffusion model in consistency distillation, or using an unbiased estimator calculated from noise and clean samples in consistency training.
To stabilize training, it necessary to ensure the tangent function in Eq. (6) is stable across different dt and the noisy sample æ; are all time steps. Empirically, we found that Fo the PF-ODE dæz
dF,, relatively stable. The only term left in the tangent function now sin(t) de sin(t) Vt) Fo dæ de
(in(t)a,FF Fp-. After further analysis, we found V Fz Fo- dæs is typically well-conditioned, so instability originates from the dime-derivative sin Fg-, which can be decomposed according to
di
sin(1)a,Fe =sin( Cenons (t)
Demb((noiiee)
F
Öt
DCnoise
Jembl Cnoise
where emb(-) refers the time embeddings. typically in the form of either positional embeddings (Ho et al., 2020; Vaswani, 2017) or Fourier embeddings (Song et al., 2021b; Tancik et al., 2020) in the literature of diffusion models and CMs.
Below we describe improvements to stabilize each component from Eq. (7) in turns.
Identity Time Transformation (Croise (t) t). Most existing CMs use the EDM formulation. which can be directly translated to the TrigFlow formulation as described in Appendix B.2. In particular, the time transformation becomes Cuoise (t) x log(a tan t). Straightforward derivation
shows that with this Guoise (t), sin(t) &cCnnise (t) = 1/ cos (t) blows up whenever t + x To mitigate numerical instability, propose to Cnoise (t) as the default time transformation.
Positional Time Embeddings. For general time embeddings in the form of emb(c) = sin(s 2n Φ), have 0.em((c) 2nw cs 2nw + Φ). With larger Fourier scale 8, this derivative
5

---

Published as a conference paper at ICLR 2025
250
dd
ad
a.f
125
4zf
OF
dk
vector nom
12.-
V,f
dx
Seee nomb
V
dx
10
0.50
vector
0.25
22
8.0
0.08
0
0.1
0.0
0.30
0.25
0.00
0.5
0.5
1.50
125
1.50
0.30
0.50
1.50
3.55
1550
EDM, Fourier scale 16.0
EDM, positional embedding
TrigFlow, positional embedding
Figure 4: Stability of different formulations. We show the norms both terms df,-
ddr V=fa-
0.fe- for diffusion models trained with the EDM (Cooss (t) log(a a tan(t))) and TrigFlow (Cmosse
formulations using different time embeddings. We observe that large Fourier scales in Fourier embeddings cause
instabilities. In addition, the EDM formulation suffers from numerical issues when & + 3. while TrigFlow (using positional emheddings) has stable partial derivatives for both 22, and f.
L-ERA sigra Begger
12
vagging
walgiig
Zern, un dindlrr
4.5-
RID
FID
255
24
30000
22000
20000
40520
20503
Number of training iterations
Number of training iterations
Number training iterationn
(a) Tangent Normalization
(b) Adaptive Weighting
(c) Discrete vs. Continuous
Figure 5: Comparing different training objectives for consistency distillation. The diffusion models are EDM2 (Karras et al., 2024) pretrained on ImageNet 512x512. (a) 1-step and 2-step sampling of continuous-
time CMs trained by using raw langents dfe= clipped tangents clip( 4fg- dr 1) and normalized tangents
di )/11 45. dt + 0.1). ). (b) Quality of 1-step and 2-step samples from continuous-tine time CMs trained w/ and w/o adaptive weighting. both are w/ tangent normalization. (c) Quality of 1-step samples from continuous-time time CMs discrete- time CMs using varying number of time steps (N), trained using all techniques Sec. 4.
has greater magnitudes and oscillates more vibrantly, causing worse instability. To avoid this, we use positional embeddings, which amounts to As 0.02 in Fourier embeddings. This analysis provides a principled explanation for the observations Song Dhariwal (2023).
Adaptive Double Normalization. Song & Dhariwal (2023) found that the AdaGN layer (Dhariwal & Nichol, 2021), defined as y = norm(x)) s(t) + b(t), negatively causes CM training diverge. Our modification is adaptive double normalization, defined as y norm (z) @epo)mppommmmmm [ (so))((porm((()()) where p-orm(-) denotes pixel normalization (Karras, 2017). Empirically we find retains the expressive power of AdaGN for diffusion training but removes its instability in CM training.
As shown in Figure we visualize how our techniques stabilize the time-derivates for CMs trained on CIFAR-10. Empirically, find that these improvements help stabilize the training dynamics of CMs without hurting diffusion model training (see Appendix G).
4.2 TRAINING OBJECTIVES
Using the TrigFlow formulation in Sec. 3 and techniques proposed in Sec. 4.1, the gradient of continuous-time CM training in Eq. (2) becomes
Øt dfo (2),,) di
Vg
wt) w(t)) F
od
Below we propose additional techniques to explicitly control this gradient for improved stability.
Tangent Normalization.
As discussed in Sec. 4.1, most gradient variance in CM training comes from the tangent function dfa- We propose to explicitly normalize the tangent function by replacing dfg-
di
dse- || c), where we empirically c 0.1. Alternatively, clip the de with dfp- /(1

---

Published as a conference paper at ICLR 2025
50g ACD
F10 Ratio
HD Rudlo
22
FID
Single Forward Flops
Single Forward Plops
1437
Single 2010 Hops
1030
Single Forward Flops
ImageNet 64 64
ImageNet 512x 512
ImageNet 64x 64
ImageNet 512x 512
(a) FID (L) as a function of single forward flops.
(b) FID ratio (L) as a function of single forward flops.
Figure 6: sCD scales commensurately with teacher diffusion models. We plot the (a) FID and (b) FID ratio against the teacher diffusion model (at the same model size) on ImageNet 64 64 and 512x sCD scales better than sCT, and has constant offset in the FID ratio across all model sizes, implying that sCD has the same scaling property of the teacher diffusion model. Furthermore, the offset diminishes with more sampling steps. tangent within [-1, 1], which also caps its variance. Our results in Figure 5(a) demonstrate that either normalization or clipping leads to substantial improvements for the training of continuous- C.Ms.
Adaptive Weighting. Previous works (Song & Dhariwal, 2023; Geng et al., 2024) design weighting functions w() manually for CM training. which can be suboptimal for different data distributions and network architectures. Following EDM2 (Karras et al., 2024), we propose to train an adaptive weighting function alongside the CM, which not only cases the burden of hyperparameter tuning but also outperforms manually designed weighting functions with better empirical performance and negligible training overhead. Key our approach is the observation that Vo F y) = 1 V [ Fa Fo- + v1l21 where y is arbitrary vector independent of 6 When training continuous- tittoouus--teeeeeo CMs using Eq. (2), we have y = (t)oa sin(t) dfa-
de This observation allows us to convert Eq. (2) into the same approach in Karras et al. (2024) to
the gradient of MSE objective. We can therefore train an adaptive weighting function that minimizes the variance of MSE losses across time steps (details in Appendix D). In practice, we find that integrating a prior weighting w(l) og tan(r) further reduces training variance. By incorporating the prior weighting, we train both the network Fb and the adaptive weighting function (t) by minimizing
L.cM(0, Be
(1) D
Fe
Pt.
Fg
cos(t) -t
-cos()) dfe (are, b) Od
us (t)
(8)
di
where D is the dimensionality of æo. and we sample tan(t) from a log-Normal proposal distribu- tion (Karras et al., 2022), that is, e tan(t) ~ NN(mmen, Paa) (details in Appendix G).
Diffusion Finetuning and Tangent Warmup. For consistency distillation, we find that finetuning the CM from a pretrained diffusion model can speed up convergence, which is consistent with df.= Song et al. (2023); Geng et al. (2024). Recall that in Eq. (6), the tangent di can be decomposed into two parts: the first term cos(t)(a - de:) ) is relatively stable, whereas the second term dF.. sin (a; + ad di may cause instability. We introduce an optional technique named as tangent warmup by replacing the coefficient sin(t) with sin(t), where linearly increases from 0 to 1 over the first 10k training iterations. We find that the tangent normalization does not affect sample quality but may reduce some gradient spikes during training.
With all techniques in place, the stability of both discrete-time and continuous time CM training substantially improves. We provide detailed algorithms for discrete-time CMs in Appendix E, and train continuous-time CMs and discrete-time CMs with the same setting. As demonstrated in Figure 5(c), increasing the number of discretization steps N in discrete-time CMs improves sample quality by reducing discretization errors, but degrades once N becomes too large (after > 1024) suffer from numerical precision issues. By contrast, continuous-time CMs significantly outperform discrete-time CMs across all N's which provides strong justification for choosing continuous-timg CMs over discrete- time counterparts. We call our model sCM (s for simple, stable, and scalable), and provide detailed pseudo-code for sCM training in Appendix A.
5 SCALING UP CONTINUOUS- TIME CONSISTENCY MODELS
Below we test all the improvements proposed in previous sections by training large-scale sCMs on a variety of challenging datasets.
7

---

Published as a conference paper at ICLR 2025
0.10
6.70
12
0.33
Precision
0.80
Rocall
6.40
sCD
0.2
0.59
0.78
2f
6.10
V50 (1:tts)
V52 mmao)
0.46
V5D 22D stoppp
v 63B dtep)
V52 <in strp)
VSD 1 PED (toppp
1.2
Guidance 16
2,0
1.2
Guidance
34
1.
20
Guidance Scale (c) FID score
14
29
(a) Precision score (t)
(b) Recall score (1)
Figure 7: sCD has higher diversity compared to VSD: Sample quality comparison the EDM2 (Karras et al., 2024) diffusion model, VSD (Wang et al., 2024; Yin et al., 2024b), sCD, and the combination of VSD and sCD, across varying guidance scales. All models are of EDM2-M size and trained on ImageNet 512.
Table I: Sample quality unconditional CIFAR-10 and class-conditional ImageNet 64 64.
Unconditional CIEAR-10
Conditionl ImageNet
METHOD
NFE (J)
FID
ME THOD
NFE
FID (L)
Diffusion models & Fast Samplers
Diffusion models Fast Samplers
Score SDE (deep) (Song et al., 2021b)
2000
2.20
ADM (Dhariwal & Nichol, 2021)
250
2.07
EDM (Karras al., 2022)
35
2.01
RIN (Jahri et al., 2022)
1000
1.23
Flow Matching (Lipman et al., 2022)
142
6.35
DPM-Solver
al., 2022ag
3.42
(Tong al., 2023)
EDM (Houn) (Karra ot al., 2022)
79
2,44
DPM-Solver (Lu al., 2022a)
10
4.70
(Heun) (Karras al., 2024)
DPM-Solvere+ (Lu et al., 20226)
2.91
Joint Training
DPM-Solver-v3 (Zhng al., 2023c)
Styl:GAN-XL (Saner et al., 2022)
1.52
Joint Training
Diff-Instract (Luo 2024)
Diffusion GAN (Xian al., 2022)
3.75
EMD (Xie et al., 2024b)
2.20
Diffusion StyleGAN (Wang et 2022)
(Yin al., 20244)
StyleGAN XI. (Saner et ol., 2022)
1.52
DMD2 (Yin al., 2024a)
1.28
CTM (Kim al., 2023)
1.87
SiD (Zhou et 2024)
1.52
Diff Instructt (Lan al., 2024)
4.53
CTM (Kim al., 2023)
1.92
DMD (Yin et al., 2024bb
3.77
1.73
SiD (Zhou 2024)
Moment Matching (Solimans 2024)
3,00
3.86
Diffusion Distillation
Diffusion Distillation
DENO (LPISS) (Zheng et al., 2023b)
3.78
DENO (LPIPS) (Zheng 2023h)
7.83
2-Roctified Flow (Liu al., 2022)
4.85
PID (1 PIPS) (Tee et 2024)
9.49
PTD (I PIPS) (Tee al., 2024)
3.92
TRACT (Berthooot et 2023)
Consistency-Pe PM (Yang et al., 2024)
5.34
(Salimans Ho, 2022)
4.97
PD (Salimans Ho, 2022) (rcimpl. Dom Hock al. (2024))
10.70
5.58
4,70
TRACT (Bemhelot et al., 2023)
3,78
CD (LPIPS) (Song et al., 2023)
6.20
3.32
4.70
CD (LPIPS) (Song et al., 2023)
3.55
MultiStep-CD (Hock et al., 2024)
3.20
cCD (ours)
3.66
sCD (ours)
2.44
2.52
Consistency Training
Consistency Training
ICT (Song Dhariwal, 2023)
ICT (Song & Dhariwal, 2023)
4.02
2.46
3.20
iCT-doep (Song Dhariwal, 2023)
iCT-deep (Song & Dhuriwal, 2023)
3.25
2.24
(Geng et al., 2024)
3.60
ECT (Geng al., 2024)
2,49
1.67
sCT (onrs)
2,85
cCT
2.04
1.48
5.1 TANGENT COMPUTATION IN LARGE-SCALE MODELS
The common setting for training large-scale diffusion models includes using half-precision (FP16) and Flash Attention (Dao et al., 2022; Dao, 2023). As training continuous-time CMs requires computing the tangent dt accurately, we need to improve numerical precision and also support memory-efficient attention computation, detailed below.
dfg-
JVP Rearrangement. Computing df.- involves calculating dF.- = V Fg. dF which can be efficiently obtained via the Jacobian-vector product (JVP) for F (z) with the input
de. + O
dt
de
8

---

Published as a conference paper at ICLR 2025
Table 2: Sample quality on class-conditional ImageNet 512x 512. 'Our reimplemented teacher diffusion model based on EDM2 (Karras et al., 2024) but with modifications in Sec. 4.1.
METHOD
NFE (
FID (
#Params
METHOD
NFE (1)
FID (
@Paraas
Diffusioon models
Diffusion
ADM-G (Dhariwal & Nichol, 2021)
250x2
7.72
559M
EDM2-S (Kamus al., 2024)
63x2
2.29
280M
2023)
EDM2-M M (Kamus et al., 2024)
63x2
2.00
498M
U-VTTH4 (Bao et al., 2023) 2023)
250x2
4.05
501M
EDM2-L (Karras at.,, 2024)
2024)
63x2
778M
DIT-XL/2 (Pocbles Xie, 2023)
3.04
EDM2-XL (Karras et al,, 2024)
63x2
1,80
1.1
SimDill (Hoogeboom el al., 2023)
512K2
3.02
2B
EDM2-XXI, (Karms al., 2024)
63x2
(Kingma & 2024)
512x2
DitiT (Haianizakh et ol., 2023)
250x2
2.67
361M
Consistency Training (sCT, ours)
DIMR-XI.3K
250x2
525M
CTT- (ours)
1
10:13
280M
Dm::SM-XLL (Yam et ol., 2024)
250x2
3.41
673M
DIM-H (Tlng 2004)
250x2
3.78
-CT-M (ours)
5.84
498M
U-DIT (Tian al.,, 2024b)
250
15.39
201M
2
5.53
498M
STT-XI (Ma 2024)
2.62
-CT-L (ours)
I
5.15
778M
Large-DiT (Alphn-VL..MM, 2024)
250x2
2.52
3B
2
4,65
778M
MaskDIT (/heng
2023a)
736M
(ours)
DIS-H/2 (Fei et al., 2024a)
DRWKV-H2 (Fei
20246)
250x2
250x2
2.95
2.88
779M
900M
2
3.73
IJIB
EDM2-S (Karms et al., 2024)
63x2
2.23
280M
sCT-XXL (ours)
4,29
1,58
IDM2-M
2024)
63x2
2,01
498M
2
3.76
1,5
EDM2-L (Karras et al., 2024)
63x2
1.88
778M
Consistency Distillation (sCD, ours)
EDM2-XL (Karras al. 2024)
63x2
1.85
L.B
EDM2-XXL (Karras et al., 202())
63x2
1.81
1.5B
SCD-S
3.00
280M
2
2.50
280M
GANs Musked Models
cCD-M
2.75
498M
BigGAN (Brock, 2018)
8.43
160M
2
2.26
498M
StyleGAN-XXL (Saver al.,
1xx2
2,41
CD--
1
2.55
778M
VQGAN (Esser et al., 2021)
1024
26.52
227M
MaskGFT (Chang nl. 2002))
7.32
227M
CD-XI
2.40
1,18
MAGVIT-22 et al., 2023)
64x2
1.91
307M
2
1.93
1.18
MAR (Li et al., 2024)
64x2
1.73
481M
CD-XXI
2.28
VAR-666s (Tian et al., 2024a)
(TR-a
10x2
2.63
2.3B
2 1.88
1,5B
vector (2t,, t) and the tangent vector (dæ: ,1. However, we empirically find that the tangent may overflow in intermediate layers when t near 0 g. To improve numerical precision, propose to rearrange the computation of the tangent. Specifically, since the objective in Eq. (8) contains can compute the JVP as:
cos(t) df0= and dfa= is proportional to sin(t) dF de
di
dF -
Fi (Vazz Fo-) (cos(t) dt
da:t ) + a Fg-
cos(t) sin(t) dt
which the JVP for (o-(..) with the input (=), t) and the tangent (cos(t) sin(t) dæs, cos(t) sin(t)oa). This rearrangement greatly alleviates the overflow issues in the intermediate layers, resulting in more stable training in FP16.
JVP of Flash Attention. Flash Attention (Dao et al., 2022; Dao, 2023) widely used for attention computation in large-scale model training, providing both GPU memory savings and faster training. However, Flash Attention does not compute the Jacobian- vector product (JVP). To fill this gap, we propose similar algorithm (detailed in Appendix F) that efficiently computes both softmax self- attention and its JVP in a single forward pass in the style of Flash Attention, significantly reducing GPU memory usage for JVP computation in attention layers.
5.2 EXPERIMENTS
To test our improvements, we employ both consistency training (referred to as sCT) and consistency distillation (referred to as sCD) to train and scale continuous-time CMs on CIFAR-10 (Krizhevsky, 2009), ImageNet 64 x4 and ImageNet 512x512 (Deng et al., 2009). We benchmark the sample quality using FID (Heusel et al., 2017). We follow the settings of Score SDE (Song et al., 2021b) on CIFARIO 10 and EDM2 (Karras et al., 2024) on both ImageNet 64 x4 and ImageNet 512x512, while changing the parameterization and architecture according to Section 4.1. We adopt the method proposed by Song et al. (2023) for two-step sampling of both sCT and sCD, using a fixed intermediate time step 1.1. For sCD models on ImageNet 512x512, since the teacher diffusion model relies on classifier-free guidance (CFG) (Ho Salimans, 2021), we incorporate an additional input s into the model F2 represent the guidance scale (Meng et al., 2023). We train the model with sCD

---

Published as a conference paper at ICLR 2025
by uniformly sampling [1, 2] and applying the corresponding CFG to the teacher model during distillation (more details are provided in Appendix G). For sCT models, we do not test CFG since it is incompatible with consistency training.
Training compute of sCM. We use the same batch size as the teacher diffusion model across all datasets. The effective compute per training iteration of sCD is approximately twice that of the teacher model. We observe that the quality of two-step samples from sCD converges rapidly, achieving results comparable to the teacher diffusion model using less than 20% of the teacher training compute. In practice, we can obtain high-quality samples after only 20k finetuning iterations with sCD.
Benchmarks. In Tables and we compare our results with previous methods by benchmarking the FIDs and the number of function evaluations (NFEs). First, sCM outperforms all previous few-step methods that do not rely on joint training with another network and is on par with, or even exceeds, the best results achieved with adversarial training. Notably, the 1 1-step FID of sCD-XXL on ImageNet 512x512 surpasses that of StyleGAN-XL (Sauer et al., 2022) and VAR (Tian et al., 2024a). Furthermore, the two-step FID of SCD-XXL outperforms all generative models except diffusion and comparable with the best diffusion models that require 63 sequential steps. Second, the two-step sCM model significantly narrows the FID gap with the teacher diffusion model to within 10%, achieving FIDs of 2.06 on CIFAR-10 10 (compared to the teacher FID of 2.01). 1.48 on ImageNet 64x64 (teacher FID of 1.33), and 1.88 on ImageNet 5 512x512 12 (teacher FID of 1.73). Additionally, we observe that sCT is more effective at smaller scales but suffers from increased variance at larger scales, while sCD) shows consistent performance across both small and large scales.
Scaling study. Based on our improved training techniques, we successfully scale continuous -iem CMs without training instability. We train various sizes of sCMs using EDM2 configurations (S, M, L, XL. XXL) on ImageNet 64x64 and 512x512, and evaluate FID under optimal guidance scales, as shown in Fig. 6. First, as model FLOPs increase, both sCT and sCD show improved sample quality, showing that both methods benefit from scaling, Second, compared to sCD, sCT is more compute efficient smaller resolutions but less efficient at larger resolutions. Third, sCD scales predictably for a given dataset, maintaining a consistent relative difference in FIDs across model sizes. This suggests that the FID of sCD decreases at the same rate as the teacher diffusion model, and therefore sCD is as scalable as the teacher diffusion model. As the FID of the teacher diffusion model decreases with scaling, the absolute difference in FID between sCD and the teacher model also diminishes. Finally, the relative difference in FIDs decreases with more sampling steps, and the sample quality of the two-step sCD becomes on par with that of the teacher diffusion model.
Comparison with VSD. Variational score distillation (VSD) (Wang et al., 2024; Yin et al., 2024b) and its multi-step generalization (Xie et al., 2024b; Salimans et al., 2024) represent another diffusion high-resolution images (Yin et al., 2024a).
distillation technique that has demonstrated scalability We apply one-step VSD from time T to 0 to finetune a teacher diffusion model using the EDM2-M configuration and tune both the weighting functions and proposal distributions for fair comparisons. As shown in Figure 7, we compare sCD, VSD, a combination of sCD and VSD (by simply adding the two losses), and the teacher diffusion model by sweeping over the guidance scale. We observe that VSD has artifacts similar to those from applying large guidance scales in diffusion models: it increases fidelity (as evidenced by higher precision scores) while decreasing diversity (as shown by lower recall scores). This effect becomes more pronounced with increased guidance scales, ultimately causing severe mode collapse. In contrast, the precision and recall scores from two-step sCD are comparable with those of the teacher diffusion model, resulting in better FID scores than VSD.
6 CONCLUSION
Our improved formulations, architectures, and training objectives have simplified and stabilized the training of continuous-time consistency models, enabling smooth scaling up to 1.5 billion parameters on ImageNet 512x512. We ablated the impact of TrigFlow formulation, tangent normalization, and adaptive weighting, confirming their effectiveness. Combining these improvements, our method demonstrated predictable scalability across datasets and model sizes, outperforming other few-step sampling approaches at large scales. Notably, we narrowed the FID gap with the teacher model to within 10% using two-step generation, compared to state-of- of-te-of-the-are tht-hhe-rtt art diffusion models that require significantly more sampling steps.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
Michael Samuel Albergo and Eric Vanden-Fijnden. Building normalizing flows with stochastic
interpolants. In The Eleventh International Conference on Learning Representations, 2023. URL https://<penneviie: net wttp:/ooummmiiiieeeeiaoiio-aa-ae
Alpha-VLLM.
Large-DiT-Inman
2024. Commit f7fe198.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pallern recognition, pp. 22669-226777 2023.
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.04248, 2023.
Andreas Blattmann, Tim Dockhom, 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.111127 2023.
Andrew Brock. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1809. 1096, 2018.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models world simulators. 2024. URL https://openaat com/ research/ video-generate ion- conn-mooelss--. aswlsssssswrrer simulators.
Huiwen Chang. Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-1133 11325, 2022.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning (2023). arXiv preprint arXiv:2307.0669 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16359, 2022.
Jia Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248 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-8794 8794, 2021.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederie Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning. 2024.
Zach Evans, CJ Carr, Josiah Taylor, Scott Hawley, and Jordi Pons. Fast timing ciming-conditioned latent audio diffusion. arXiv preprint 2rXiv:2402.0480 2024.

---

Published as a conference paper at ICLR 2025
Zhengeong Fei, Mingyuan Fan, Changqian Yu, and Junshi Huang. Scalable diffusion models with state space backbone. arXiv preprint arXiv:2402.05600 2024a.
Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Diffusion-rwkv: Scaling rwkv-like architectures for diffusion models. arXiv preprint arXiv:240.004477, 2024b.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. arXiv preprint arXiv:2406. 14548, 2024.
Ali Hatamizadeh, Jiaming Song. Guilin Liu, Jan Kautz, and Arash Vahdat. Diffit: Diffusion vision transformers for image generation. arXiv preprint arXiv:2312.02139, 2023.
Jonathan Heek, Emiel Hoogeboom, and Tim Salimans. Multistep consistency models. arXiv preprint 200x:22403.06807, 2024.
Eric Heitz, Laurent Belcour, and Thomas Chambon. Iterative a-(de) blending: A minimalist deterministic diffusion model. In ACM SIGGRAPH 2023 Conference Proceedings, pp. 1-8, 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to a local Nash equilibrium. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, pp. 6626-6633 6637, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NewrIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
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, al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02300 2022.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pp. 13213-13232. PMLR, 2023.
Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprint 22. 2022.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
Tero Karras. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint 171. 2017.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv preprint 2xXiv2206.003664 2022.
Tero Karras, Miika Aittala, Jaakko I Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CV Conference Computer Vision and Pattern Recognition, pp. 24174-24184, 2024.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Ue- saka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint 10.0022310.02279 2023.
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances Neural Information Processing Systems, 36, 2024.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Advances in Neural Information Processing Systems, 2021.
12

---

Published as a conference paper at ICLR 2025
Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Alexander Korotin. Optimal flow matching: Learning straight trajectories in just one step. arXiv preprint arXiv:2403.13117, 2024.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization. arXiv preprint arXiv:2406. 1838, 2024.
Shanchuan Lin, Bingchen Liu, Jiashi Li, and Xiao Yang. Common diffusion noise schedules and sample steps are flawed. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pp. 5404 5411 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hame Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Haohe Liu, Zehua Chen, Yi Yuan, Xinhao Mei, Xubo Liu, Danilo Mandic, Wenwu Wang, and Mark D Plumbley, Audioldm: Text-to-addio generation with latent diffusion models. arXiv preprint arXiv:2301.12500 2023a.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv: wwt.:/008000000 2019.
Qihao Liu, Zhanpeng Zeng, Ju He, Qihang Yu, Xiaohui Shen, and Liang-Chieh Chen. Alleviating dis- tortion in image generation via multi-resolution diffusion models. arXiv preprint arXiv:2406.09416, 2024.
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick. Zero- 1-to33 Zero-shot one image to 3d object. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 299899309, 2023b.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0388 03000,, 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu, Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models, arXiv preprint arXiv:2211.01095 2022b.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.0233 2021.
Simian Luo, Yigin Tan, Longbo Huang. Jian Li, and Hang Zhao. Latent consistency models:
Synthesizing high-resolution images with few-step inference. arXiv preprint ar.00:23100043378 2023.
Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diff- instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden- Eijnnen,, and Saining Xie. Sit: Exploring flow and diffusion- generative models with scalable interpolant transformers. arXiv preprint 240.::220100777400 2024.
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.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171, PMLR, 2021.
13

---

Published as a conference paper at ICLR 2025
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195 4205-420,440 2023.
Stefano Peluchetti. Non-denoising forward-time diffusions, 2022. URL htt.httss///ppereeeeei net/ peet/fornnt wediodoiiiii V
Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall, Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209. 14988, 2022.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06115, 2022.
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.
2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations, 2022.
Tim Salimans, Thomas Mensink, Jonathan Heek, and Emiel Hoogeboom. Multistep distillation of diffusion models via moment matching. arXiv preprint arXiv:2406.04100 04000000,,,, 2024.
Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings, pp. 1-10, 2022.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint arXiv:2311.17042 2023..
Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. Advances in neural information processing systems, 31, 2018.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, pp. 1256-2266 PMLR, 2015.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310. 111119910,,,189, 2023.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution.
In Advances in Neural Information Processing Systems, volume 32, pp. 11895-11907, 2019.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances Neural Information Processing Systems, volume 34, pp. 1418-1428, 2021a.
Yang Song. Jascha Sihl-Dickstten. Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference Learning Representations, 2021b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, pp. 3221 1-22522. PMLR, 2023.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Kan Nithin Raghavan, Utkarsh Singhal. Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537-7547, 2020.
Joshua Tian Jin Tee, Kang Zhang, Hee Suk Yoon, Dhananjaya Nagaraja Gowda, Chanwoo Kim, and Chang D Yoo. Physics informed distillation for diffusion models. Transactions on Machine Learning Research, 2024.
14

---

Published as a conference paper at ICLR 2025
Yao Teng, Yue Wu, Han Shi, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, and Xihui Liu. Dim: Diffusion mamba for efficient high-resolution image synthesis. arXiv preprint arXiv:2405.1422 2024.
Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. arXiv preprini arXiv:2404.00290. 2024a.
Yuchuan Tian, Zhijun Tu, Hanting Chen, Jie Hu, Chao Xu, and Yunhe Wang. U-dits: Downsample tokens in u shaped diffusion transformers. arXiv preprint arXiv:2405.0777 2024b.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.00425 2023.
Vaswani. Attention all you need. Advances in Neural Information Processing Systems, 2017.
Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-gan: Training gans with diffusion. arXiv preprint arXiv:2206.0226,, 2022.
Zhengyi Wang, Cheng Lu, Yikai Wang. Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Pro- lifiedreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36, 2024.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference on Learning Representations, 2022.
Qingsong Xie, Zhenyi Liao, Zhijie Deng, Shixiang Tang, Haonan Lu, et al. Mlem: Multistep consistency distillation of latent diffusion model. arXiv preprint wrXiv:2406.05760 2024a.
Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, and Ruiqi Gao. EM distillation for one-step diffusion models. arXiv preprint ar arviv:2405. 16852, 2024b.
Jing Nathan Yan, Jiatao Gu, and Alexander M Rush. Diffusion models without attention. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8239 8249, 2024.
Ling Yang, Zixiang Zhang, Zhilong Zhang, Xingchao Liu, Minkai Xu, Wentao Zhang, Chenlin Meng, Stefano Ermon, and Bin Cui. Consistency flow matching: Defining straight flows with velocity consistency. arXiv preprint arXiv:2407.02395 2024.
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. arXiv preprint 2rX. 2024a.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 6613-66132 2024b.
Lijun Yu, José Lezama, Nitesh Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Agrim Gupta, Xiuye Gu, Alexander G Hauptmann, al. Language model beats diffusion- tokenizer key to visual generation. arXiv preprint arXiv:2310.05737 2023.
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, and Anima Anandkumar. Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306.09300 2023a.
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. 42390-42402. PMLR, 2023b.
15

---

Published as a conference paper at ICLR 2025
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Dpm-solver-vB: Improved diffusion ode solver with empirical model statistics. Advances in Neural Information Processing Systems, 36: 55502-55542, 2023c.
Mingyuan Zhou, Huangjie Zheng, Zhendong Wang, Mingzhang Yin, and Hai Huang. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Forty-first International Conference on Machine Learning. 2024.
16

---

Published as a conference paper at ICLR 2025
DISCUSSIONS AND LIMITATIONS
sCT is less effective than sCD in latent spaces. As listed in Tables and 2, sCT consistently outperforms sCD on CIFAR-10 and ImageNet 64x64 but is less effective than sCD across different model scales on ImageNet 512x512. We believe the higher training variance of CT is the main issue, particularly in the complex latent spaces defined by the pretrained encoder. We hypothesize that the current encoder/decoder may not be optimal for consistency models. Theoretically. since the ground truth mapping in consistency models aims to transform a Gaussian distribution into multimodal data distribution with potentially disconnected supports, its tangent can become ill-conditioned at boundary points, resulting in worse optimization dynamics. If we could develop better encoder/decoder to create more well-conditioned ground truth mapping in the latent space, the training of consistency models would likely become significantly easier.
Computation costs of sCM. As Jacobian- vector product can be efficiently computed using forward- mode automatic differentiation, which requires the same memory and compute as standard forward pass without saving intermediate activations. This is significantly cheaper than backpropagation, which relies on reverse-mode automatic differentiation. Consequently, our continuous-time consis- tency models require similar compute and memory to train when compared to their discrete-time counterparts, which perform two forward passes at each iteration.
Limitations. Despite large improvements in FID scores, our method can still produce images with noticeable artifacts. These artifacts are commonly observed when training generative models on the ImageNet dataset with class labels, whereas training on larger datasets with caption conditions may significantly alleviate this issue. Furthermore, our 2-step sCM still shows a small gap compared to state-of-the-art diffusion models, which we believe may be further reduced by incorporating our proposed techniques into multi-step consistency models (Heek et al., 2024). Additionally, since FID scores do not capture all semantic details, further validation is needed to determine whether our method can scale effectively to image video generation tasks that require larger resolutions and fine details. Besides, ensuring the training stability of sCM requires several significant modifications of the network architecture, thus sCM may be not suitable for some architectures designed for diffusion models. Moreover, our best performing method sCD still highly relies on the performance of pretrained diffusion models, which restricts the architecture family and potentially limits the few-step generation performance. Addressing these quality issues might require new sampling strategies or enhanced architectures to maintain high fidelity even with limited sampling steps.
APPENDIX
We include additional derivations, experimental details, and results in the appendix. The detailed training algorithm for sCM, covering both sCT and sCD, is provided in Appendix A. We present comprehensive discussion of the TrigFlow framework in Appendix B. including detailed derivations (Appendix 1. and its connections with other parameterization (Appendix B.2). We introduce new algorithm called adaptive variational score distillation in Appendix C, which eliminates the general framework
need for manually designed training weighting. Furthermore, elaborate for adaptive training weighting in Appendix D, applicable to diffusion models, consistency models, and variational score distillation. As our improvements discussed in Sec, 4 are also applicable for discrete- time consistency models, we provide detailed derivations and the training algorithm for discrete-tee time consistency models in Appendix E, incorporating all the improved techniques of sCM. We also provide complete description of the lacobian-vector product algorithm for Flash Attention in Appendix F. Finally, all experimental settings and evaluation results are listed Appendix G, along with additional samples generated by our SCD-XXI. model trained on ImageNet at 512x512 resolution in Appendix H.
A TRAINING ALGORITHM OF SCM
We provide the detailed algorithm of sCM in Algorithm 1, where we refer to consistency training of sCM as sCT and consistency distillation of sCM as sCD.
17

---

Published as a conference paper at ICLR 2025
Algorithm 1 Simplified and Stabilized Continuous-time Consistency Models (sCM).
Input: dataset D with std. da pretrained diffusion model F pretrain with parameter Opeeminn model Fo, weighting was learning rate y, proposal (Peean; Paa), constant e, warmup iteration H. Init: &- pretrain: Iters < 0.
repeat
20 ~ D. = ~ (O, (.) ~ (P Pla) t f arctan ( Dt +- cos(t) + sin(t)= cos(t)= -tnn)))-s)()) if consistency training else dæz <- aa F pretrain
(in()) r f min(1, H)
dæ:
2t
dr
aa
>
Tangent warmup
f cos² (oaFo
dæy df - r cos(t) (t) + ad
dF. dt
b JVP rearrangement
g - g/[lg]] +c)
Tangent normalization
[(0, 6) ="a()
Fg R t) Fo g w
Adaptive weighting
(0 (0, ②)
V (00,60)
Iters < Iters +
until convergence
B TRIGFLOW: A SIMPLE FRAMEWORK UNIFYING EDM, FLOW MATCHING AND VELOCITY PREDICTION
B.1 DERIVATIONS
Denote the standard deviation of the data distribution Pa as dd- We consider general forward diffusion process at time t [0, with æ; ag + 0, for the data sample Vo Pa and the noise sample z ~ N NN0, (0,,0 o) (note that the variance of z is the same as that of the data wo) where o 0, are noise schedules such that 02/0, is monotonically decreasing w.r.t. t, with ag = 0, = 0. The general training loss for diffusion model can always be rewritten as
Cour( Eug,z, (w(t) Do(2,, zu
(9)
where different diffusion model formulation contains four different parts:
1. Parameterization of Dg, such score function (Song & Ermon, 2019; Song al., 2021b),
noise prediction model (Song & Ermon, 2019; Song et al., 2021b; Ho et al., 2020), data prediction model (Ho al., 2020; Kingma et al., 2021; Salimans & Ho, 2022), velocity pre- diction model (Salimans & Ho, 2022), EDM (Karras et al., 2022) and flow matching (Lipman et al., 2022; Liu al., 2022; Albergo et al., 2023).
2.
Noise schedule for de and 0,, such as variance preserving process (Ho et al., 2020; Song et al., 2021b), variance exploding process (Song et al., 2021b; Karras et al., 2022), cosine schedule (Nichol & Dhariwal, 2021), )), and conditional optimal transport path (Lipman et al., 2022).
3. Weighting function for (t). such as uniform weighting (Ho al., 2020; Nichol Dhariwal, 2021; Karras et al., 2022), weighting by functions of signal-to-noise--eeeee (SNR) (Salimans & Ho, 2022), monotonic weighting (Kingma & Gao, 2024) and adaptive weighting (Karras et al., 2024).
4. Proposal distribution for t, such uniform distribution within [0, T] (Ho et al., 2020: Song et al., 2021b), log-normal distribution (Karras et al., 2022), SNR sampler (Esser et al., 2024), and adaptive importance sampler (Song et al., 2021a; Kingma al., 2021).
Below show that, under the unit variance principle proposed in EDM (Karras et al., 2022), obtain a general but simple framework for all the above four parts, which can equivalently reproduce all previous diffusion models.
For any diffusion process with æ, al ofe where N/0,, I), can always equivalently convert sl to x; = a/2s + Pg (aze) and let z = age, 0r == oi, oz == does result in any loss of generality.
So the assumption for (N o
18

---

Published as a conference paper at ICLR 2025
Step 1: General EDM parameterization. principle in EDM (Karras et al., 2022) by
We consider the parameterization for Da as the same
Da t) Cskip (1)ar, Coun(t) (t) Fo ( Crois (t)),
(10)
and thus the training objective becomes
Loir = Ex w(t) (t) Fo Fo Cnoise (t)) (1
Cskip (t)ip()aa))20) Ckip(t)o
(11)
Cour (t)
To ensure the input data of F, has unit variance, we should ensure (ar((m(())a))) by letting
Gi Oayo +
(12)
To ensure the training target of F6 has unit variance, have
out (t) o3 (t)ap()) + oi
(13)
To reduce the error amplification from Fg to Dg. we should ensure Coun (t) to be as small possible, which means we should take (Lkp(() by letting
Bram dessi
O, which results in
Ct2 07 oz
Cakip (t)
Cour =
(14)
Va + +
Though equivalent, we choose Cour (1)
which can simplify some derivations below.
Val+o?
In summary, the parameterization and objective for the general diffusion noise schedule are
Do(x,, af + oi
Lt
Fg
Cnoise (1))
(15)
V a o7
JJV oj of
LDiff Exc.
w(t)
of af + of
JdFo
Cnoise (t)
0,2 0,
(16)
JJV of + oi
Step 2: All noise schedules can be equivalently transformed. One nice property of the unit variance principle is that the ag; a, in both the parameterization and the objective are homogenous, which means we can always assume af + o² = without loss of generality. To see this, we can apply and a;; Vaz+o?
simple change-of-varian of â; Vo?+ot +o?
δe Va?+o?
â,35 + o.2,
thus we have
x. Caoise (t))
Do(x,, t) = dg::: Gr
(17)
w(t)o}
+t Cnosse( -(à1x 3,20)
2
Coir = Exp.xx
JaFo
(18)
As for the sampling procedure, according DPM-Solver++ el, 2022b), 20222). the exact solution of diffusion ODE from time to time satisfies
Tt xx σ1 S.a² > D'Doxx. G
(19)
where X₁ log an the sampling procedure is also homogenous for Oeey Tr. To this, can use de the fact that 2t &; and d log & = de, thus the above equation is equivalent a; àn
ō. n. - oe (,a She e Do Åjad,
(20)
which is exactly the sampling procedure of the diffusion process 2,, which means noise schedules of diffusion models won't affect the performance of sampling. In other words, for any diffusion
19

---

Published as a conference paper at ICLR 2025
process (at, dz. 2t) time t, we always divide them by Val + oo to obtain the diffusion process (å,, oe, Bi) with àž + o = and the parameterization, training objective and sampling procedure can be equivalently transformed. The only difference is the corresponding training weighting 2 in Eq. (18), which we will discuss the next step.
A straightforward corollary is that the "opttiaal transport path" (Lipman al., 2022) in flow matching with 0e = 1, o = t can be equivalently converted to other noise schedules. The reason of its better empirical performance is essentially due the different weighting during training and the lack of advanced diffusion sampler such as DPM-Solver series (Lu et al., 2022a:b) during sampling, the "straight path" (Lipman et al., 2022) itself. By converting the diffusion process the space satisfying / a o the path connection 2% and has consistent variance which matches the unit-variance design principles of EDM.
Step 3: Unified framework by TrigFlow. As we showed in the previous step, we can always assume a? + of L. An equivalent change-of-variable of such constraint is to define
f = arctan
= arctan
Oe
(21)
so f E [0, 2 is monotonically increasing function of t [0, T). thus there exists a mapping between and i convert the proposal distribution p(t) to the distribution of denoted as (i). As cos
âe = cos (f) Ô. = sin (f), the training objective in Eq. (18) is equivalent to
Exng, l:² p (1) w (1) sin (1)
LDIA
FaFo
Cnoise (i)
(cos (1 -sin (1) xo
di
training weighting
independent from and
(22)
Therefore, we can always put the influence of different noise schedules into the training weighting of the integral for f from 0 to 2. while the & norm loss each f is independent from the choices of o and ot. As we equivalently convert the noise schedules by trigonometric functions, we name such framework for diffusion models as TrigFlow.
For simplicity and with slight abuse of notation, we omit the i and denote the whole training weighting as a single wr)), we summarize the diffusion process, parameterization, training objective and samplers of TrigFlow as follows.
Diffusion Process. In ~ >p((z)), ~~ ~NO,, cos(t)æn + sin(t) for E [0, 51.
Parameterization.
Do( 1) =con())):ssii(())) (#: Et
(noisee (
(23)
where Croise (t) is the conditioning input of the noise levels for Fe, which can be arbitrary one-one mapping of I. Moreover, the parameterized diffusion ODE is defined by
dæ dt
Fo
It od
Gnoise (1)
(24)
Training Objective.
(Di = E Exzo x Lo w(t)
Zt oaFo
Cnoise (t)
(cos(t) sin(t)xo))
dt
(25)
0
where w(t) is the training weighting, which will discuss in details in Appendix D. As for the sampling procedure, although we can directly solve the diffusion ODE in Eq. (24) by Fuler's or Heun's solvers in flow matching (Lipman al., 2022), the parameterization for aa
20

---

Published as a conference paper at ICLR 2025
may be not the optimal parameterization for reducing the discreteization errors. As proved in DPM- Solver-v3 (Zheng et al., 2023c), the optimal parameterization should cancel all the linearity of the ODE, and the data prediction model Dg is an effective approximation of such parameterization. Thus, we can also apply DDIM, DPM-Solver and DPM-Solver++ for TrigFlow by rewriting the coefficients into the TrigFlow notation, listed below.
Ist-order Sampler by DDIM.
Starting from 2, at time 8, the solution Dt at time / is
Ft = cos(s I) - sin(s 1)
Cnoise (8)
(26)
Od
One good property of TrigFlow is that the 1st-order sampler naturally support zero-SNR sam- pling (Lin et al., 2024) by letting s = without any numerical issues.
2nd-order Sampler by DPM-Solver.
Starting from Bg at time s. by reusing previous solution
x. at time s', the solution p; at time is
x, = cos(s - x))x,- sin(s - Fo am= Cruise(
sin(s - 1)
2r, cos(s) (Eo(Z,, - co( 8)) (27)
where Eg (a 1) = sin(t) cos(t)aaFo (ml, Croise (f) is the noise prediction model, and r's
log tan(a) ()) --) log tan(=') log tan(s) log tam(?)
2nd-order Sampler by DPM-Solver++. T at time s', the solution I; at time is
Starting from æ, at time s, by reusing a previous solution
Cuoise (s) + sin(s - t)
cos(s - t) sin Fe
(Do s') - D(x,, s)), (28)
2r, sin(s)
0d
where T.
log (au(())) Lan(s') log (=) tan(r)
B.2 RELATIONSHIP WITH OTHER PARAMETERIZATTIN
As previous diffusion models define the forward process with zy = a0 ag Ze (oae) da
for E ~ (O, I), we can obtain the relationship between 1' and TrigFlow time steps E [0, 5 by
arctan
VOjoz+O7 Tr.
(29)
Thus, we can always translate the notation from previous noise schedules to TrigFlow notations. Moreover, below we show that TrigFlow unifies different current frameworks for training diffusion models, including EDM, flow matching and velocity prediction.
EDM. As our derivations closely follow the unit variance principle proposed in EDM (Karras et al., 2022), our parameterization can be equivalently converted to EDM notations. Specifically, the transformation between TrogFlow (æ,, 1) and EDM (ær, o) is
= arctan
xz cos(t)
(30)
The reason why TrigFlow notation is much simpler than EDM is just because we define the end point of the diffusion process as z (O0 I) with the same variance as the data distribution. Thus, the unit variance principle can ensure that all the intermediate æ; does not need to multiply other coefficients as EDM.
Flow Matching. Flow matching (Lipman et al., 2022: Liu et al., 2022; Albergo al., 2023; Kornilov et al., 2024) defines a stochastic path between two samples To from data distribution and from tractable distribution which is usually some Gaussian distribution. For general path 2t = Cg + 02 with ag = 1,a7 = 0, do = 0, oz = 1, the conditional probability path is
da:
xu + do:
V2 =
z,
(31)
dt
dt
21

---

Published as a conference paper at ICLR 2025
and it learns parameterized model (a t) by minimizing
B w(t) (00 b) vell
(32)
and the final probability flow ODE is defined by
dæ: di
vy (0 6).
(33)
As TrigFlow uses a = =os((t) and o; = sin(t), it is easy to see that the training objective and the diffusion ODE of TrigFlow are also the same as flow matching with vo t) Fo Cuoise (t)).
To the best of our knowledge, TrigFlow is the first framework that unifies EDM and flow matching for training diffusion models.
Velocity Prediction. The velocity prediction parameterization (Salimans & Ho, 2022) trains parameterization network with the target a, o, As TrigFlow 0n (los()),0))) or sin sin((), is easy to see that the training target in TrigFlow is also the velocity.
Discussions on SNR. Another good property of TrigFlow is that it can define data-varianee invariant SNR. Specifically, previous diffusion models define the SNR at time as SNR(t) = c for 2t 0g T&E with ~ NO, I). However, such definition ignores the influence of the variance of æn: we rescale the data æp by constant, then such SNR doesn't ' get rescaled correspondingly, which is not reasonable in practice. Instead, in TrigFlow
define the SNR by
a of
1
(
(34)
tan (t))
which is data-variance-invariat and also simple.
C ADAPTIVE VARIATIONAL SCORE DISTILLATION IN TRIGFLOW FRAMEWORK
In this section, we propose the detailed derivation for variational score distillation (VSD) in TrigFlow framework and an improved objective with adaptive weighting.
C.I DERIVATIONS
Assume we have samples Vn E R.D from data distribution Pa with standard deviation das and define a corresponding forward diffusion process {p:)T= starting at po = Pa and ending at pr (/(0, 8 I), with Peo (z, ( = N |o420, T). Variational score distillation (VSD) (Wang et al., 2024; Yin al., 2024b;a) trains generator go RRD Rd aiming to map noise samples N' 6200 to the data distribution, by minimizing
min E, Dxt (qf pt) Ex,z,e [w(t) (log d. (augo(z) + o.e) log (oggo() +
where N'(, I). of is the diffused distribution time with the same forward diffusion process as P: while starting at o as the distribution go (z), w(() is an ad-hoc training weighting (Poole et al., 2022; Wang al., 2024; Yin et al., 2024b), and t follows proposal distribution such as
uniform distribution. It is proved that the optimum of of satisfies 9o Pa (Wang et al., 2024) and thus the distribution of the generator matches the data distribution.
22

---

Published as a conference paper at ICLR 2025
Moreover, by denoting 2. argo(z) + ore and taking the gradient w.r.t 8, we have
VoEe DKL (9 pi)]
Rt,, [u(t)Vo (log ol (a.) - pa
dxc
)]
Er,,, w(()
Og log of (2x1) (Va, log logg (2t) Va log (rgp(()))))
08
04 0r,go 00
-B (vo())0g log of [are]] +R1,2, w((() (Vn, log log (xre) Van log
Oga
-,, (tw(() (Va,, log of (xx) - Van pe (xx))
00
Therefore, we need to approximate the score functions Vz, log qi (xt) for the generator and Væ, log px(;)) for the data distribution. VSD trains a diffusion model for samples from ge ga(z) approximate we, log of (ax) and uses a pretrained diffusion model to approximate Va, log Po (xt).
In this work, we train the diffusion model in TrigFlow framework, with cos(t), og sin sin Fd. 2. Specifically, assume we have pretrained diffusion model F prctrain parameterized by TrigFlow, and we train another diffusion model F& to approximate the diffused generator distribution, by
min Eg 2,
w(t)
FaFs
where 24 cos(t) + sin(t) cos(t)= sin(t) z' (O, az Moreover, the relationship between the ground truth diffusion model Frar(ar, I) and the score function Va, p..:)( is
NNO. 12 To go(z') with
Fa t) B[u, [ = tan(t)
sin(t) En ie [zo]
xt (os(t))zg
cos(t)o (on(t Fair + sin(t) og sin(()
Væ, (xt)) Exg wx
of sin (l)
Thus, train the generator go by the following gradient w.r.t. 0:
Ex
(t) aa sin(t) w(t) Fpretrain
aa (111
F
0go
ad
30
which is equivalent to the gradient of the following objective:
cos (t)
sin(t) w(t) go((z') go- (z') + Fpretrain
at Fs
E,,,,''
od
where 9o- (z') is the same as (z())) but stops the gradient for 6. Note that the weighting functions sin²(1)
used in previous works (Wang et al., 2024; Yin et al., 2024b) is proportional to cos(t) thus the prior weighting is proportional to sin(t) cos(!), which has U-shape similar to the log-normal distribution used Karras et al. (2022). Thus, we can instead use log-normal proposal distribution and apply the adaptive weighting by training another weighting network we (t). We refer Appendix D for detailed discussions about the learnable adaptive weighting. Thus we can obtain the training objective, listed below.
C.2 TRAINING OBJECTIVE
Training Objective of Adaptive Variational Score Distillation (aVSD).
It 0
min Lpit = Ex,x'
w(t)
FoF
f
2
(())
min LVsD(0, v)) = Er 6.8
ga(x')
g0 (x')= - F pretrain
_
It
F
-
w.s
(1)
D
0.
od
23

---

Published as a conference paper at ICLR 2025
And also choose proportional distribution of for estimating (0, 15) by log(tan(t)oz)
N N (Pmean: P and tune these two hyperparameters (note that they may be different from the proposal distribution for training (ó), as detailed in Appendix G.
In addition, for consistency models fo t), we choose x' ~ (N0, 0 and go(z') = fo ;)
Fo oaa 2), and thus the corresponding objective is
z' od 2
z' Od 2'
Ftt a
min E 2,, 8,00
Fo
Fa
F
pretrain
F
w w
od
od
D ADAPTIVE WEIGHTING FOR DIFFUSION MODELS, CONSISTENCY MODELS AND VARIATIONAL SCORE DISTILLATION
We first list the objectives of diffusion models, consistency models and variational score distillation (VSD). For diffusion models, as shown in Eq. (25). the gradient of the objective is
V (iIrin( Vgl, 2.. -e.,, w(t [l ve v = V VelE [w(t)o&FF) (a.Fo- - -a)Fo--))]-
where Fg- is the same Fp but stops the gradient w.r.t. O. For VSD, the gradient of the objective is
VoLDan Vo [w)())o(=z) (Fpretrain
F$)]
And for continuous-time CMs parameterized by TrigFlow, the objective is dfo
(OCcm(O) = V Vo
(()) sin(
di
where fe- is the same fo but stops the gradient w.r.t. O. Interestingly, all these objectives can be rewritten into form inner product between a neural network and target function which has the same dimension (denoted D) as the output of the neural network. Specifically, assume the neural network is Fs parameterized by 0, we study the following objective:
min E. [Fg y].
where we do not compute the gradients w.r.t. for y. such case, the gradient will be equivalent to
VgE,
Fg Fo
212
where Fo- is the same Fg but stops the gradient w.r.t. 0. In such case, we can balance the gradient variance w.r.t. t by training an adaptive weighting network wy(t) to estimate the loss norm, i.c., minimizing
e (t)(
min E,
Fa
Fg- + 21l2 w.((t)
This is the adaptive weighting proposed by EDM2 (Karras et al., 2024), which balances the loss variance across different time steps, inspired by the uncertainty estimation of Sener & Koltun (2018). By taking the partial derivative w.r.t. w in the above equation, it is easy to verify that the optimal w" (l) satisfies
E E
Fg + + 11
Therefore, the adaptive weighting reduces the loss variance across different time steps. In such case, all we need to do is to choose
1. A prior weighting X(L) for y. which may be helpful for further reducing the variance of Then the objective becomes
(t)(())
min E, 8,0
F. 1 + X(t) (()yll-ws(tt)
D
e.g., for diffusion models and VSD, since the target either y ve or y Eorctrai Fo which are stable across different time steps, we can simply choose X(t) = 1; while for consistency models, the target sin(t) of may have huge variance, we choose X(t)
da (a) to reduce the variance of A(t)y, which empirically is critical for better performance.
24

---

Published as a conference paper at ICLR 2025
2. A proposal distribution for sampling the training t, which determines which part of should focus on more. For diffusion models, we generally need to focus on the intermediate time steps since both the clean data and pure noise cannot provide precise training signals. Thus, the common choice to choose a normal distribution over the log-SNR of time steps, which is proposed by Karras al. (2022) and also known log-normal distribution.
In this way, we do not need to manually choose the weighting functions, significantly reducing the uning complexity of training diffusion models, CMs and VSD.
E DISCRETE-TIME CONSISTENCY MODELS WITH IMPROVED TRAINING OBJECTIVES
Note that the improvements proposed in Sec. can also be applied discrete-time consistency models (CMs). In this section, we discuss the improved version of discrete-time CMs for consistency distillation.
E.1 PARAMETERIZATION AND TRAINING OBJECTIVE
Parameterization.
We also parameterize the CM by TrigFlow:
fe (z,, t) cos(t)==; - Fa sin(t())
And denote the pretrained teacher diffusion model as dæ: F
Pretrrin
= t).
Reference sample by DDIM. Assume we sample 2o ~ Pds z No of and a; cos(t)aro sin())z, need reference sample at time t' < é to guide the training of the CM, which can be obtained by one-step DDIM from to l':
Ict add"
2t = cos(t t()x, - ddd sis(t t') Ppretrain
Thus, the output of the consistency model at time t' is
att (aa
Ce t'
fg (')= = cos(t') cossttt')az--ao cosst" F pretrain
-oaii((())))
(35)
Original objective of discrete-time CMs. The consistency model at time / can be rewritten into
It [=
So ().,t) = cos(t)az sasiin()FF..
ad
(cos(L') cos(1 (') sin(t') sin(l 1'))
(36)
- aa(ssin( - (') (1') + cos(| (') (in) Fg-
od
Therefore, by computing the difference between Eq. (35) and Eq. (36), we define (x,,t,t') = fo- (xi, t) fo- (x0))'')
sin(t t')
= =ooos((') (oaFF)
ct 0g
Cc _=,1
pretrain
Od
Od
(37)
Od cocos(t - f
ite od
0 Fg
sii(tt)(() (x +
sin(t (t 1')
Rid d
25

---

Published as a conference paper at ICLR 2025
Comparing Eq. (6), it is easy to that limy-se Ag- (2 1,1') -f (2), Moreover, when using d( y) || y 212. and At 1 t', the training objective of discrete-time CMs in Eq. (1)
di
becomes
Exnt (t))lffo(() fo (xi- (c--anttAt)]) -t)1
which has the gradient of
Bu,t V (2t, t) (fo- (xc, - fo (Xi- At))] Volaa (())s(ss( - (t)()))))))) fo (0t, 1)Ag (a L, t')]
(38)
VolE,, [w(() sin Fg (z=,t)Ag- ( t, t,t")
Adaptive weighting for discrete-time CMs. Inspired by the continuous-time consistency models, we can also apply the adaptive weighting technique into discrete-time training objectives Eq. (38).
Specifically, since Ag (æt, t, t') is first-order approximation of d/g= (Iz, t), we can directly replace the tangent in Eq. (8) with Ag (az,, t, t'), and obtain the improved objective of discrete-time CMs by:
(t))))
CcM(B, 0) E
Fe
Fo
It
:)
cos(t) (z ws(()
od
Od
(39)
where (t) the adaptive weighting network.
Tangent normalization for discrete-time CMs. We apply the simliar tangent normalization method as continuous-time CMs by defining
(x,, (,, t, t') go - (a,, t,t cos(t) (Ie, t, t')l +
where > O is hyperparameter, and then the objective in Eq. (39) becomes
(1)
Lo (0 E
Fo
- Fe-
ge- (2 121
wo mo(()
D
od
add"
Tangent warmup for discrete-time CMs.
We replace the Ag- (at, t,t') with the warmup version:
Dt
Ag = - cos(t') oaFo
- oaE Epeerani
0d
Td cos(t F')
O2Fo
-r sin(t')
Et +
sin(t
t')
where r linearly increases from to over the first 10k training iterations.
We provide the detailed algorithm of discrete-time sCM (dsCM) in Algorithm 2, where we refer to consistency distillation of discrete-time sCM as dsCD.
Algorithm 2 Simplified and Stabilized Discrete-tin time Consistency Distillation (dsCD).
dataset D with std. σa. pretrained diffusion model Prearain with parameter Opettzins model Fo, weighting We learning rate proposal (Pmean: Psa), constant c, warmup iteration H. Init: +- Opeerrin. Iters +- repeat
Input:
2p ~ D. N,,,2(0, N(P meaan P2), ) 1 t- arctan (.) x. +- cos(!)a0o + sin(l)z
æx. cos(t ('),, aasin(t 1') ( Peretrain
min(1, HH)
Tangent warmup
f- cos(t)Ap t,t,t,t',r)
b JVP rearrangement
g/(lg]] c
Tangent normalization
C(0, 6) + m)(0
Fe 0. t) - Fg m1
gl1 (t(()
Adaptive weighting
(0, (0, ②)
V
C(0,b)
Iters f Iters +1
until convergence
26

---

Published as a conference paper at ICLR 2025
E.2 EXPERIMENTS OF DISCRETE-TIME SCM
We use the algorithm in Algorithm 2 to train discrete-time sCM, where we split [0, Σ into N intervals by EDM sampling spacing. Specifically, we first obtain the EDM time step by o; mo
1/p
1/p o min )) with p 7, min 0.002 and o mux = 80, and then obtain f, arctan(o,/or aa)) and set fg = 0. During training, we sample with a discrete categorical distribution that splits the log-normal proposal distribution used in continuous-time sCM, similar to Song & Dhariwal (2023).
M max
As demonstrated in Figure 5(c), increasing the number of discretization steps N in discrete-ting CMs improves sample quality by reducing discretization errors, but obviously degrades once N becomes too large (after N >> 1024) to suffer from numerical precision issues. By contrast, continuous-time CMs significantly outperform discrete- teereteetimeeeeee CMs across all N's which provides strong justification for choosing continuous-time CMs over discrete-timg counterparts.
E.3 COMPARISON WITH ECT
We compare the 1-step sampling FID scores different training iterations between ECT (Geng et al., 2024) and sCT on CIFAR-I0. As shown in Table 3, our proposed sCT significantly outperforms ECT during the training, demonstrating the effectiveness of the compute efficiency and faster convergence of sCT.
For fair comparison, we use the same network architecture with ECT on CIFAR-10, which is the DDPM++ network proposed by Ho et al. (2020) and does not have AdaGN layer, and use the same dropout rate of 0.20 as ECT, and use the same batch size (128) as ECT (which is different from our default setting of 512 in our reported results in Table 1). We choose Pmean = 1.0 and Pad 1.8 for SCT, and use TrigFlow parameterization (with Cnoise = t). All the other hyperparameters are the same as the experiments in Table 1.
Table 3: Sample quality measured by FID score (L) of ECT (Geng et al., 2024) and sCT at different training iterations on 10..
Training Iterations
100k
200k
400k
ECT
4.54
3.86
3.60
sCT (ours)
3.97
3.51
3.09
F RACOBIAN-VECTO PRODUCT OF FLASH ATTENTION
The attention operator (Vaswani. 2017) needs to compute y = softmax (x) where RIXL V RLXD y E IRIDD Flash Attention (Dao et al., 2022: Dao, 2023) computes the output by maintaining three variables m(æ) E R, ((x)) IR, and f(=) with the same dimension as æ. The computation done recursively: for each block, we have
mar) max max(e");
((x) Σ m(n)
f(x) et v
and for combining two blocks x = [x((1) (01], we merge their corresponding m, l, f by m(x) max (0))) (z()),x(00) 20(1)) ((x))= ml(al)) mm(()) + ) x)) (l (x (8))
f f(x) ( mm(a) f(x(")), e m(()))) mm(x) f(
f(x)
y
f(x)
However, to the best of knowledge, there does not exist an algorithm for computing the Jacobian- Vector product of the attention operator in the Flash Attention style for faster computation and memory saving. We propose a recursive algorithm for the JVP computation of Flash Attention below.
Denote p so)t Denote the tangent vector for x E RIXL, P E RIXL V E LLDD y E R as tz E IIX tp E IIxL tv E RLID by E RIxD correspondingly. The JVP for attention is computing (xx tm), (V, (v) + (y, by), which
by V + ply
where L, (p tx)
V
(pt
(pV)
softmax (x)
1XL
1x1
27

---

Published as a conference paper at ICLR 2025
Notably, the computation for both ptv and pV can be done by the standard Flash Attention with the value matrices V and tv. Thus, to compute ty. we only need to maintain a vector g == (pota ftg) V
and a scalar μ)(²) = pt= during the Flash Attention computation loop. Moreover, since we do not know p during the loop, each block,
reuse the intermediate m, e, f in Flash Attention. Specifically, for
g ((x) (c=-m(a) ④ tiz) V, paax) Σ - m())
and for combining two blocks x = [x((), (01),, we merge their corresponding g and μ by g(x) (om((ol)))mlal g(z(o)) (m(a")))mma g(x g(x(1))
m(x) = e mm()) ((x((1) + e m ( mmm() /1 (z)
and after obtaining m, m, f.f.fggt. 9. for the row vector the final result of tp V is
g(x)
p(x))
1p V
y.
(xx:
E(x)
Therefore, we can use a single loop to obtain both the output y and the JVP output by. which accesses the memory for the attention matrices only once and avoids saving the intermediate activations, thus saving the GPU memory.
G EXPERIMENT SETTINGS AND RESULTS
G.1 TRIGFLOW FOR DIFFUSION MODELS
We train the teacher diffusion models CIFAR-10, ImageNet 64x64 and ImageNet 512x512 with the proposed improvements of parameterization and architecture, including TrigFlow parameteriza- tion, positional time embedding and adaptive double normalization layer. We list the detailed settings below.
CIFAR-10. Our architecture based the Score SDE (Song al., 2021b) architecture (DDPM++). We use the same settings of EDM (Karras et al., 2022): dropout rate is 0.13, batch size is 512, number of training iterations is 400k, learning rate is 0.001, Adam e 10 ,8 0.9, .3 = 0.999. We use 2nd-order singlc-step DPM-Solver (Lu et al., 2022a) (DPM-Sollvr-2S)) with Heun's 's intermediate time step with 18 steps (NFE=35), which is exactly equivalent to EDM Heun's sampler. We obtain FID of 2.15 for the teacher model.
ImageNet 64x64. We preprocess the ImageNet dataset following Dhariwal & Nichol (2021) by
1. Resize the shorter width height to 64 x 64 resolution with bicubic interpolation.
2. Center crop the image.
3. Disable data augmentation such as horizontal flipping.
Except for the TrigFlow parameterization, positional time embedding and adaptive double normaliza- tion layer, we follow exactly the same setting in EDM2 config G (Karras al., 2024) train models with sizes of S, M, L, and XI. while the only difference is that we use Adam = 10-11.
ImageNet 512x512. We preprocess the ImageNet dataset following Dhariwal & Nichol (2021) and Karras et al. (2024) by
1. Resize the shorter width height to 512 x 512 resolution with bicubic interpolation.
2. Center crop the image.
3. Disable data augmentation such as horizontal flipping.
4. Encode the images into latents by stable diffusion VAE² (Rombach et al., 2022; Janner et al., 2022), and rescale the latents by channel mean μ., [1.56, -0.695, 0,483, 0.729] and channel std a. = [5.27, 5.91, 4.21, 4.31, We keep the aa = 0.5 as in EDM2 (Karras et al., 2024), for each latent we substract Me and multiply it by a2/O..
stabi tityii parri//d-vaae ft-mse
28

---

Published as a conference paper at ICLR 2025
When sampling from the model, we redo the scaling of the generated latents and then run the VAE decoder. Notably, our channel mean and channel std are different from those in EDM2 (Karras et al., 2024). It because when training the VAE, the images are normalized to [-1, before passing to the encoder. However, the channel mean and std used in HDM2 assumes the input images are in [0, 1] range, which mismatches the training phase of the VAE. We empirically find that it is hard to distinguish the reconstructed samples by human eyes of these two different normalization, while it has non-ignorable influence for training diffusion models evaluated by FID. After fixing this mismatch, our diffusion model slightly outperforms the results of EDM2 at larger scales (XL and XXL). More results are provided in Table 6.
Except for the TrigFlow paramcterization, positional time embedding and adaptive double normaliza- tion layer, we follow exactly the same setting in EDM2 config G (Karras et al., 2024) to train models with sizes of S, M. L,, XI. and XXL., while the only difference is that we use Adam = 10- 11. We enable label dropout with rate 0.1 to support classifier free guidance. We use 2nd-order single-step DPM-Solver (Lu et al., 2022a) (DPM-Solver-25) with Heun's intermediate time step with 32 steps (NFE=663, which is exactly equivalent to EDM Heun's 's sampler. We find that the optimal guidance scale for classifier-fre guidance and the optimal EMA rate also the same EDM2 for all model sizes.
G.2 CONTINUOUS-TING CONSISTENCY MODELS
In all experiments, we use c = 0.1 for tangent normalization, and use H = 10000 for tangent warmup. We always use the same batch size as the teacher diffusion training, which is different from Song & Dhariwal (2023). During sampling, we start at 4mxx = arctan Omn with maxx 80 such that it matches the starting time of EDM (Karras et al., 2022) and EDM2 (Karras et al., 2024). For 2-step sampling, we use the algorithm in Song et al. (2023) with an intermediate t - 1.1 for all the experiments. We always initialize the CM from the EMA parameters of the teacher diffusion model. For sCD, we always use the F Prettrai of the teacher diffusion model with its EMA parameters during distillation.
We empirically find that the proposal distribution should have small P mcaa- i.e. close to the clean data, to ensure the training stability and improve the final performance. Intuitively, this is because the training signal of CMs only come from the clean data, so we need to reduce the training error for t near to 0 to further reduce the accumulation errors.
CIFAR-10. For both sCT and sCD, we initialize from the teacher diffusion model trained with the settings in Appendix G.1, and use RAdam optimizer (Liu et al., 2019) with learning rate of 0.0001, 31 0.9, 822 0.99, 6 = 10-$, and without learning rate schedulers. proposal distribution of Pinca = - 1.0, Padd = 1.4. For the attention layers, we use the implementation in (Karras et al., 2022) which naturally supports JVP by PyTorch (Paszke et al., 2019) auto-grad. We use EMA half-life of 0.5 Ming (Karras et al., 2022). We use dropout rate of 0.20 for sCT and disable dropout for sCD.
ImageNet 64x64. We only enable dropout at the resolutions equal or less than 16, following Simple Diffusion (Hoogeboom et al., 2023) and iCT (Song & Dhariwal, 2023). We multiply the learning rate of the teacher diffusion model by 0.01 for both sCT and sCD. We train the model with half precision (FP16), and use the flash attention jvp proposed in Appendix F for computing the tangents of flash attention layers. Other training settings are the same as the teacher diffusion models. More details of training and sampling are provided in Table 4 and Table 8. During sampling, we always use EMA length ret = 0.05 for sampling from CMs.
ImageNet 512x512. We only enable dropout at the resolutions equal to or less than 16, following Simple Diffusion (Hoogeboom et al., 2023) and iCT (Song & Dhariwal, 2023). We multiply the learning rate of the teacher diffusion model by 0.01 for both sCT and sCD. We train the model with half precision (FP16), and use the flash attention jvp proposed in Appendix F for computing the tangents of flash attention layers. Other training settings are the same as the teacher diffusion models. More details of training and sampling are provided in Table 5 and Table 6. During sampling, we always use EMA length det 0.05 for sampling from CMs.
We add an additional input in F (0 t, s) where represents the CFG guidance scale of the teacher model, where s is embedded by positioinal embedding layer and an additional linear layer, and the embedding is added to the embedding of 1, similar to the label conditioning. During training,
29

---

Published as a conference paper at ICLR 2025
Table 4: Training settings of models training algorithms on ImageNet 64 64 dataset.
Model Size
S
L
XL
Model details
Batch size
2048
2048
2048
2048
Channel multiplier
192
256
320
384
Time embedding layer
positional
positional
positional t
positional
noise conditioning Guoise (1)
adaptive double normalization
/
V
v
Learning rate decay (fxz)
35000
35000
35000
35000
Adam B1
0.9
0.9
0.9
0.9
Adam B2
0.99
0.99
0.99
0,99
Adam
1.0e-11
1.0e-11
1.0e-11
1,0e-11
Model capacity (Mparams)
280.2
497.8
777.6
1119.4
Training details diffusion models (TrigFlow)
Training iterations
1048k
1486k
761k
540k
Learning rate max (Oner)
1.0c-2
9.0c-3
8.0c-3
7.0c-3
Dropout probability
0%
10%
10%
10%
Proposal Preea
-0.8
-0.8
-0.8
-0.8
Proposal Pas.
1.6
1.6
Shared details of consistency models
Learning rate max (Oieet
1.0c-4
9.0e-5
8.0e-5
7.0c-5
Proposal P monn
-1.0
-1.0
-1.0
-1.0
Proposal Pas.
1.6
1.6
1.6
1.6
Tangent normalization constant (c)
0.1
0.1
Tangent warm up iterations
10k
10k
10K
10k
EMA length (aaa) of pretrained diffusion
0.075
0.04
0.04
Training details of sCT
Training iterations
400k
400k
400k
400k
Dropout probability resolution 16
45%
45%
45%
45%
Training details of sCD
Training iterations
400k
400k
400k
400k
Dropout probability for resolution 16
0%
0%
0%
0%
uniformly sample [] and apply CFG with guidance scale the teacher diffusion model to
get Fpretrain-
VSD experiments. We do not use EMA for F6 in VSD, instead we always use the original model for F, for stabilizing the training. The learning rate of F& is the same as the learning rate CMs. More details and results are provided in Tables 5 to 7.
30

---

Published as a conference paper at ICLR 2025
Table 5: Training settings all models and training algorithms on ImageNet 512x512 datasct.
Model Size
XXL
Model details
Batch size
2048
2048
2048
2048
2048
Channel multiplier
256
320
384
Time embedding layer
positional
positional
positional t
positional
positional
noise conditioning (t)
f
adaptive double normalization
V
V
V
Learning decay (fet)
70000
70000
70000
70000
70000
Adam
0.9
0.9
0.9
0.9
Adam /22
0.99
0.99
0.99
0.99
0.99
Adam
1.0c-11
1.0c-11
1.0c-11
1.0c-11
1.0c-11
Model capacity (Mparams)
280.2
497.8
777.6
1119.4
1523.4
Training details of diffusion models (TrigFlow)
Training iterations
1048k
1048k
696k
598k
376k
Learning
(Otet)
1.0c-2
9.0c-3
8.0e-3
7.0c-3
6.5c-3
Dropout probubility
10%
10%
Proposal P Peean
-0.4
-0.4
-0.4
-0.4
-0.4
Proposal Pat.
1.0
1.0
1.0
1.0
1.0
Shared details of consistency models
Learning rate max (Over)
L.e-4
9.0e-5
8.0e-5
7.0e-5
6.5e-5
Proposal P Prean
-0.8
-0.8
-0.8
-0.8
-0,8
Proposal Pas-
1.6
1.6
1.6
1.6
1.6
Tangent normalization constant (c)
Tangent warm up iterations
10k
10k
10k
10k
10k
EMA length (o.a) of pretrained diffusion
0.025
0.03
0,015
0.02
0.015
Training details of sCT
Training itcrations
100k
100k
100k
Dropout probability for resolution <. 16
25%
35%
35%
35%
35%
Training details of sCD
Training itcrations
200k
200k
200k
Dropout probability for resolution & 16
0%
10%
10%
10%
10%
Maximum of CFG scale
2.0
2.0
2.0
2.0
2.0
Training details of sCD with adaptive VSD
Training iterations
20k
20k
20k
20k
20k
Learning rate max (aver) for F.
1.0e-4
9,0e-5
8.0e-5
7.0e-5
6.5e-5
Dropout probability for F,
0%
10%
10%
10%
10%
Proposal Precu for Crar (b)
-0.8
Proposal Pas. for Lnarl (Ø)
1.6
1.6
1.6
1.6
1.6
Number of updating of d per updating of
I
1
step sampling starting time fwns
arctan( 2
arctan( 80
arctan( 80
arctan( 80)
arctan( (2)
Proposal Premn for Lvsn(()
0.4
Proposal Pas. for (8)
2.0
2.0
2.0
2.0
2.0
Loss weighting Avsp for Lvsp
1.0
1.0
1.0
1.0
1.0
31

---

Published as a conference paper at ICLR 2025
Table 6: Evaluation of sample quality of different models 5 dataset. Results of EDM2 (Kar- ras et al., 2024) are with EDM parameterization the original AdaGN layer. 'The FDuioozin EDM2 are obtained by tuned EMA rate, which is different from our EMA rates that are tuned for FID scores.
Model Size
M
L
XL
XXL
Sampling by diffusion models (NFE= 126)
EMA length (ama)
0.025
0.030
0.015
0.020
0.015
Guidance scale for FID
1.4
1,2
1.2
1.2
1.2
'Guidance scale for FDDIN02
1.8
1.8
1.8
FID (TrigFlow)
2.29
2,00
1.87
1.80
1.73
FID (EDM2)
2.23
2,01
1.88
1.85
1.81
FDaiNova (TrigFlow)
52.08
43.33
39.23
36.73
35.93
"FDoinox2(EDMP with Fea for FDDINOv2
52.32
41.98
38.20
35.67
33.09
Sampling by consistency models trained with sCT
Intermediate time Eii in 2-step sampling
1.1
1.1
1.1
1.1
1.1
1-step FID
10.13
5,84
5.15
4.33
4,29
2-step FID
9,86
5,53
4.65
3.73
3.76
1-step FDDINOV2
278.35
192.13
169.98
147.06
146.31
2-step FDDINOv2
244.41
160.66
135.80
114.65
112.69
Sampling by consistency models trained with sCD
Intermediate time Emid in 2-step sampling
1.1
1.1
1.1
1.1
1.1
Guidance scale for FID, 1-step sampling
1.5
1.3
1.3
1.3
1.3
Guidance scale for FID, 2-step sampling
1.4
1.2
1.2
1.2
1.2
Guidance scale for FD 1-step sampling
2.0
2.0
2.0
2.0
Guidance scale for FDDINOX2, 2-step sampling
2.0
2.0
1.9
1.9
1.9
1-step FID
3.07
2.75
2.55
2.40
2.28
2-step FID
2.50
2.26
2.04
1.93
1.88
1-step FDDINOV2
104.22
83.78
76.10
70.30
67.80
2-step FDDINOV2
71.15
55.70
50.63
46.66
44.97
Sampling by consistency models trained with multistep sCD
Guidance scale for FID
1.4
1.2
1.2
1.15
1.15
Guidance scale for FDDINO12
2.0
2.0
2.0
1.9
1.9
FID, M = 2
2.79
2.51
2.32
2.29
2.16
FID, M 4
2.78
2.46
2.28
2.22
2.10
FID, M =8
2,49
2.24
2.04
2.02
1.90
FID,
2.18
1.99
1.90
1.82
FD Diioo M =
76.29
60.47
54.91
51.91
50.70
FDoIOo22
72.01
56.38
50.99
47.61
46.78
DDDNO-2, M 8
60.13
49.46
44.87
41.26
40.56
FDoINOKa, 16
55.89
46.94
42.55
39.30
38.55
Sampling by consistency models trained with sCD
adaptive VSD
Intermediate time Emia in 2-step sampling
1.1
1.1
1.1
1.1
1.1
Guidance scale for FID, 1-step sampling
1.0
1.0
Guidance scale for FID, 2-step sampling
1.2
1.0
1.0
1.0
1.0
Guidance scale for FDoINo2, 1-step sampling
1.5
1.5
Guidance scale for FDDINOx2, 2-step sampling
1.7
1.5
1.6
1.5
1.5
1-step FID
3.37
2.67
2.26
2.39
2.16
2-step FIDD
2.70
2.29
1.99
2.01
1.89
1-step FDDINOv2
72.12
54.81
50.46
48.11
45.54
2-step FDDINOv2
53.53
48.54
46.61
43.93
32

---

Published as a conference paper at ICLR 2025
Table 7: Ablation of adaptive VSD and sCD on ImageNet 512x512 dataset with model size M.
Method
VSD
sCD
sCD VSD
EMA length (dhdl)
0.05
0.05
0.05
Guidance scale for FID, I 1stepp sampling
1.1
1.3
1.0
Guidance scale for FID, 2-step sampling
1.2
1.0
Guidance scale for FDnINova, 1-step sampling
1.4
2.0
1.5
Guidance scale for FDDINOw,, 2-step sampling
2.0
1.5
1-step FID
3.02
2.75
2.67
2-step FID
2.26
2.29
1-step FDDINO
57.19
83.78
54.81
2-step FDDINO22
55.70
53.53
Table 8: Evaluation of sample quality of different models ImageNet 64 64 dataset.
Model Size
XL
Sampling by diffusion models (NFE=63)
EMA length (ara)
0.075
0.06
0.04
0.04
FID (TrigFlow)
1.70
1.55
1.44
1.38
Sampling by consistency models trained with sCT
Intermediate time faid in 2-step sampling
1.1
1.1
1.1
1.1
1-step FID
3.23
2.25
2.08
2.04
2-step FID
2.93
1.81
1.57
1.48
Sampling by consistency models trained with sCD
Intermediate time fiiit in 2-step sampling
1.1
1.1
1.1
1-step FID
2.97
2.79
2,43
2.44
2-step FID
2.07
1.66
33

---

Published as a conference paper at ICLR 2025
H
ADDITIONAL SAMPLES
Class (robin), guidance 1.9
Class 29 (axolot)), guidance 1.9
Class 33 (loggerhead), guidance 1.9
Class 88 (macaw), guidance 1.9
Figure 8: Uncurated I 1-step samples generated by our sCD-XXL trained on ImageNet 512x512.
34

---

Published as a conference paper at ICLR 2025
Class 15 (robin), guidance 1.9
Class 29 (axolot)), guidance 1.9
Class 33 (loggerhcad), guidance 1.9
Class 88 (macaw), guidance 1.9
Figure 9: Uncurated 2-step samples generated by our sCD-XXL trained on ImageNet 512x512.
35

---

Published as a conference paper at ICLR 2025
Class 127 (white stork), guidance 1.9
Class 323 (monarch), guidance 1.9
Class 388 (giant panda), guidance 1.9 Class 387 (lesser panda), guidance 1.9
Figure 10: Uncurated 1-step samples generated by our sCD-XXL. trained on ImageNet 512x512.
36

---

Published as a conference paper at ICLR 2025
Class 127 (white stork), guidance 1.9
Class 323 (monarch), guidance 1.9
Class 388 (giant panda), guidance 1.9 Class 387 (lesser panda), guidance 1.9
Figure 11: Uncurated 2-step samples generated by our sCD-XXL trained on ImageNet 512x512.
37

---

Published as a conference paper at ICLR 2025
Class 417 (balloon), guidance 1.9
Class 425 (bam), guidance 1.9
Class 973 (coral reef), guidance 1.9 Class 933 (chcesehurger) guidance 1.9
Figure 12: Uncurated 1-step samples generated by our sCD-XXL. trained on ImageNet 512x512.
38

---

Published as a conference paper at ICLR 2025
Class 417 (balloon), guidance 1.9
Class 425 (bar), guidance 1.9
Class 973 (coral reef), guidance 1.9 Class 933 (chceseburger) guidance 1.9
Figure 13: Uncurated 2-step samples generated by our sCD-XXL. trained on ImageNet 512x512.
39


Reviews:
[{"review_id": "0qzpk8E8Ds", "reviewer": "Reviewer_2i3a", "summary": "This paper investigates a fundamental topic in consistency models (CMs), specifically the challenges of discretization errors and the resulting training stability issue. Consistency Models can be trained in discrete or continuous time, either from scratch using a dataset or distilled from pretrained teacher scores. CMs' theoretical foundation elucidates the importance of controlling the discretization error and eventually achieving consistency in continuous time. While continuous-time CMs eliminate the discretization errors present in their discrete-time counterparts, they suffer from training instability, a problem that is not yet well understood in the research community.\n\nThis work conducts a comprehensive study into continuous-time CMs, covering forward process parameterization, network architecture, and training techniques. The authors first develop a simplified diffusion process formulation called TrigFlow, which unifies EDM and Flow Matching for the first time. Building upon this foundation, they analyze the gradient flow of continuous-time CMs, identify the root cause of training instability, and mitigate this issue through modifications to time embeddings and adaptive group normalization. Additional training techniques, such as adaptive weighting functions and annealing, further contribute to improved training stability and scalability.\n\nThe resulting method, sCT/sCD, allows continuous-time CMs to be trained at an unprecedented scale, scaling up to 1.5B parameters on ImageNet 512x512. These results significantly narrow the performance gap between CMs and state-of-the-art diffusion models to less than 10% in FID, while matching or even surpassing adversarial methods and discrete/continuous autoregressive models in both performance and efficiency.", "questions": "1. I appreciate the scaling study from ImageNet 64x64 to 512x512, where the former operates directly in the image space, while the latter relies on additional image compression models. The difference between sCT and sCD at the 512x512 resolution is intriguing, as even increasing model size and batch size cannot easily close the gap, while at the 64x64 resolution, scaling is sufficient to compensate for the variance induced by Monte Carlo estimation.\n    - Data complexity at 512x512 and the lack of effective mode decomposition definitely contribute to this discrepancy. However, I am curious about the authors' thoughts on the extent to which the increased variance could be caused by the pretrained image encoder/decoder. In other words, could data modes become *more dispersed* in the latent space, making it harder for sCT to learn effectively? Would it be better to train CMs directly in the image space (with some special architecture design and weighting schemes) or to find a latent space that is more suitable for CMs? In some sense, modern autoregressive models focus on tokenizer design for visual generation. Could latent space compression for CMs require properties distinct from those used in DMs? While this is more hypothetical, I would be happy to hear more thoughts.\n\n2. If I understand correctly, there are three weighting functions applied, namely learnable adaptive weighting, tangent normalization (per-sample basis weighting), a prior weighting $w(t) = \\frac{1}{\\sigma_d \\tan(t)} = \\frac{1}{\\sigma_t}$. To what extent does this prior weighting contribute to variance reduction? Is it helpful to stabilize the learnable adaptive weighting layer? I assume the time embeddings of learnable adaptive weighting could be either positional embedding or Fourier embedding since it is not directly involved in $\\frac{\\mathrm{d} \\boldsymbol{f}}{\\mathrm{d}t}$.\n\nAs an additional comment, could the authors consider conducting distillation experiments using the data-free formulation in [1]? I am curious how continuous-time CMs would scale and how stable they would be without using an extra dataset. No hurry to complete these experiments during the limited rebuttal period!\n\n[1] Consistency Models Made Easy\n\n### Minor\n\n1. In Line 213 and in Line 227, both diffusion models and consistency models are denoted as $f_\\theta(\\mathbf{x}_t, t)$ but with different equations.\n\n2. The Adaptive Double Normalization is less explained. Is it the same as local response normalization applied to the modulation layer?", "limitations": "", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This is a very strong paper in analysis, practical techniques, writing, and experiment results. \n\n- Novelty. This paper's novelty is evident in several aspects.\n    - First, it studies an important but less studied problem: consistency models in continuous time, together with the training stability and discretization error of consistency models. \n    - The proposed TrigFlow, as a novel unification of EDM and Flow Matching, substantially simplifies the analysis presented later and the practical techniques.\n    - The gradient analysis of continuous-time objective reveals the root cause of instability. To the best of my knowledge, this is the first paper to establish the gradient analysis for CMs.\n    - Model architecture modifications are original since existing works are mostly inherited from Diffusion Models' design and focus on the training techniques and formulations, leaving the architectural design underexplored.\n\n- Soundness. Its technical claims are well backed up by both theoretical analysis and empirical results. I particularly appreciate the in-depth investigation into the training dynamics and gradient analysis of continuous-time CMs.\n\n- Presentation. The logical flow of this paper is well structured and smooth. The problem statement is clearly defined, and the explanation of why discretization errors matter for CMs and the motivation toward continuous-time formulation is crystal clear. The gradient analysis into continuous-time CMs is thoughtfully motivated and carefully organized. Even the appendix is well-written, offering useful insights into the proposed techniques.\n\n- Experiments. Proposed techniques allow for training continuous-time Consistency Models (sCMs) at an unprecedented scale. Experiment results are impressive, matching/outperforming adversarial approaches, score distillation, and recent autoregressive models.\n    - Gradient variances have been carefully controlled via adaptive weighting and normalization techniques.\n    - Comprehensively studying the scaling behaviors of sCMs under continuous-time training. \n    - Comparisons with improved score distillation baseline using many methods developed in this work confirm the mode coverage of CMs.\n    - Additionally, the paper discusses efficient and stable implementation strategies for continuous-time CMs.\n\nGiven the potential impact of this paper, I strongly recommend acceptance with conference highlights. It was a great pleasure to read through the manuscript!", "weaknesses": "I did not find any apparent weaknesses in the analysis or experiments (including both ablation studies and performance evaluation). There are research questions worth further investigation, as discussed 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": "pKXVPOnfqj", "reviewer": "Reviewer_Fdt4", "summary": "The authors propose improvements to the consistency models generative paradigm and named their new method sCM. Specifically they -vastly- improve the FID for consistency models with the introduction of several new ideas to both stabilize and simplify continuous consistency models. \n\nMy understanding is the main claim of simplification for sCM comes from the simplification of EDM (Kerras et al.) normalizing design, resulting in $c_{in}=c_{skip}=1$ which in turns simplifies the continuous expression of consistency models. Another simplification is the combination of both EDM and Flow Matching concepts into their method which they call TrigFlow. Yet another simplification, not claimed as such by authors, is the use of vanilla L2 loss compared to Huber/LPIPS use in previous iterations of consistency models. This last simplification has the additional benefit to be more probabilistically grounded.\n\nThere are 3 main proposed ideas to stabilize the training of consistency models:\n1. Identity-time transformation as a replacement to the log-transformation from EDM\n2. Fourier embedding of the time dimension are replaced by positional embeddings\n3. AdaGN is modified to also normalize the conditioning inputs for scale and bias.\nMore ideas are also proposed in the training objective to stabilize training, namely: tangent normalization and tangent warmup. It is my understanding that the adaptative weighing is the same as in EDM.\n\nThe paper also provides ample ablations to demonstrate the effects and the reasoning motivating these 3 proposed improvements.", "questions": "These are mostly from the weakness section:\n1. Are there more limitations other than the 10% worse than diffusion SOTA?\n2. Is the method really fully stable?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is very well grounded mathematically and experimentally. The analysis is based on understanding the causes of training instabilities by decomposing the loss, validating each component experimentally and proposing changes to solve the root causes.\nThe mathematics while greatly simplified are still pretty complex and the paper shines in its clarity to make the logical reasoning easy to follow.\nThe experimental results are also outstanding resulting in very significant gains, essentially taking consistency models within 10% of the SOTA for diffusion models.", "weaknesses": "1. The limitations of the method are not very clear to me besides the fact that it's still 10% worse than SOTA for diffusion.\n2. The section on positional embeddings (line 269 and on) lacks details to be fully understandable without having to read another paper. Maybe beefing up that section would make the paper more self-contained.\n3. I found all figure very useful with the exception of Figure (3) which I felt did not add much value.\n\nI also noticed a typo (definitely not affecting my score), just leaving it there for authors to fix their manuscript:\n- Line 362: \"cause instability\" => \"causes instability\"", "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": "6n4mI5kMr5", "reviewer": "Reviewer_DGp7", "summary": "The paper presents a unified perspective on diffusion-based and flow-based generative models and introduces a comprehensive set of techniques aimed at improving the training stability and overall performance of continuous-time consistency models for large-scale image generation. The techniques include: 1) enhancing time transformation and embeddings, 2) replacing the AdaGN layer with Adaptive Double Normalization, 3) normalizing the tangent function and applying tangent warm-up, 4) implementing an adaptive weighting function in the training objective, and 5) optimizing forward-mode differentiation. These techniques mitigate the numerical instability issues in continuous-time consistency models and enable the model to achieve highly competitive performance in class-conditioned image generation.", "questions": "Q1 - In Section 4.1, the authors emphasize the importance of time transformation in mitigating numerical instability. Have the authors considered other potential candidates for time transformation?\n\nQ2 - In Figure 6(a) and Section 5.2, the authors mention that sCT is less effective at higher resolutions. Do the authors have any insights into why this might be?\n\nQ3 - In Figure 6(b), it appears that the performance of sCD-XL under the two-step sCD setting is better than that of sCD-XXL, which contradicts the results reported in Table 2. Could the authors clarify the specific settings used in Figure 6(b)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "S1 - The paper provides a comprehensive analysis and set of solutions addressing the numerical instability issues in continuous-time consistency models, significantly improving performance and enabling the model to achieve competitive results on selected benchmarks.\n\nS2 - Many of the enhancements are supported by detailed theoretical justification and experimental results.\n\nS3 - The unified perspective on previous diffusion and flow-matching parameterizations is thorough, complete, and well-grounded, offering novel insights that could benefit the community.\n\nS4 - The paper is well-structured and easy to follow.", "weaknesses": "W1 - Several design choices appear arbitrary and lack supporting evidence. \n\nFor example, in Section 4.1, the authors discuss the preference for Adaptive Double Normalization over AdaGN, but there is no experimental evidence supporting this choice. It would be more insightful to add a Figure similar to Figure 5 show experimental comparison between Adaptive Double Norm and AdaGN. \n\nSimilarly, in Section 4.2, the authors propose training with linear warm-up w.r.t the model's time derivative, yet no evidence is provided to demonstrate this choice’s effectiveness. Again, including an ablation study or comparative analysis showing the impact of the linear warm-up on training stability or performance metrics would provide more concrete evidence for the effectiveness of this specific choice\n\nFurthermore, Figure 5(b) suggests that incorporating adaptive weighting in a two-step setting may lead to worse performance, while in the one-step setting, it only yields marginal improvement. Have the authors considered alternative designs for the two step case?\n\nW2 - In Sections 4.1 and 5.2, the paper discusses the training compute of sCM. However, including a comparison of compute efficiency with other models (e.g., ECT [1]) would be more insightful, maybe a table or figure comparing the compute efficiency (e.g., FLOPs or training time) of sCM against ECT and other relevant baselines for a given performance level. \n\nAdditionally, given that the model is trained on a large-scale dataset (ImageNet 512) under latent setting, it would be beneficial to include discussions related to text-to-image generation.\n\n[1] Geng, Zhengyang, Ashwini Pokle, William Luo, Justin Lin, and J. Zico Kolter. \"Consistency Models Made Easy.\" arXiv preprint arXiv:2406.14548 (2024).", "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": "d6c0TeL69Z", "reviewer": "Reviewer_QC4D", "summary": "This paper tackles the problem of  instability in continuous consistency models and presents many contributions. \n\nThe first novelty is the simplification of model parametrization in the EDM/CM frameworks. It is shown that rescaling the mean ($\\alpha_t$ ) and noise ($\\sigma_t$) schedule of conditional probability flows with the norm $||(\\alpha_t, \\sigma_t)||$, produces a second conditional flow, such that the EDM model parametrization and loss in the first are identical to those in the second. Furthermore, the ODE sampling procedure remains unchanged. However, the connection between the new flow schedule parameters and the  EDM scaling parameters is greatly simplified after normalization, facilitating further theoretical analysis, while maintaining the benefits of the EDM model definition (unit variance of input, target and minimizing the scaling of the output). The formulation with $\\alpha_t=cos(t)$ and  $\\sigma_t=sin(t)$ is named *TrigFlow*.\n\nThen the training objective is studied and probed for causes of instability. Several factors are detected such as: an inappropriate $c_{noise}(t)$ which is fixed by being set to $t$; inappropriate Fourier scales in the time positional embeddings which are then reduced; the usage  'AdaGN layer', which is then replaced with 'Adaptive Double Normalization'; highly varying target norm alleviated by target normalization; inappropriate weighting mitigated by adaptive weighting and unstable terms solved by slowly introducing such terms into the loss with respect to the number of parameter updates. \n\nIn addition, methods for scaling such models to large sizes and datasets are proposed, namely JVP Rearrangement and JVP of Flash Attention.\n\nFinally, the proposed models are compared against the state of the art methods, and results show that continuous consistency models are excellent generators, that only require 1 or 2 sampling steps to generate new quality data from learned distributions with continuous support.", "questions": "How do the proposed models compare with [1} and [2]? Even tests on a small dataset using small models of equal size would be informative.. I completely understand however  if such comparisons cannot be made due to the short length of the discussion period. \n\nShouldn't $c_{skip}$ and $c_{out}$ be $\\sigma_t$ and  $ \\alpha_t \\sigma_d$, that is for $\\sigma_t=t$ and  $\\alpha_t=1$: $c_{skip}=t$  and $c_{out}=\\sigma_d$ in line 201/202?\n\nBased on the experiments performed, are there any indications that the continuous consistency models will still face instability issues for even larger scales, in particular as compared to diffusion/flow models? This question mostly relates to this part in the paper: \"Additionally,\nwe observe that consistency training is more effective at smaller scales but suffers from increased\nvariance at larger scales, while consistency distillation shows consistent performance across both\nsmall and large scales.\"\n\n[1]  Tong et al. 2024. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport\n\n[2]  Kornilov et al 2024. Optimal Flow Matching: Learning Straight Trajectories in Just One Step", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The contributions of this paper are novel, clear, significant and positioned correctly.\n\nTrigflow simplifies the theoretical analysis of flows by showing that  the mean ($\\alpha_t$ ) and noise ($\\sigma_t$) schedule in the definition of conditional flows can be normalized while preserving the model/loss formulation and the integrator-generated paths, while simplifying the connection between the scaling parameters of the model and those of the conditional flow. **(novelty, significance)**\n\nSeveral components of continuous consistency models are studied and probed for potential causes of instability. They include: an inappropriate $c_{noise}(t)$ which is fixed by being set to $t$; inappropriate Fourier scales in the time positional embeddings which are then reduced; the usage  'AdaGN layer', which is then replaced with 'Adaptive Double Normalization'; highly varying target norm alleviated by target normalization; inappropriate weighting mitigated by adaptive weighting and unstable terms solved by slowly introducing such terms into the loss with respect to the number of parameter updates.  **(novelty, significance)**\n\nStrategies for scaling such models to large sizes and datasets are proposed, namely JVP Rearrangement and JVP of Flash Attention. **(quality, significance)**\n\nThe proposed method shows competitiveness with the state of the art models, while only using 1 or 2 steps for generation, and outperforms all other tested methods that use 1 to 2 generation steps **(significance, quality)**", "weaknesses": "**While the contributions of the paper are numerous, the paper could be strengthened even further by:**\n\n1) Comparing the proposed model with more recent flow models such as [1] and [2] and adding the results for rectified flows with 2 generation steps.\n\n2) Placing the number of parameters and the number of parameter updates (or training time on identical hardware) for each model in Table 1 is very important as using an equal umber of parameters/compute is essential for ensuring a fair comparison.\n\n**Some smaller issues and suggestions:**\n\na) The paper would be improved if an intuitive explanation is given for the loss in Equation 2. Also it could save readers some time if it is mentioned that the loss is derived in Song et al 2023, *Remark 10*. \n\nb) The paper would be enriched by adding some generated images with one step.\n\nc) Shouldn't $c_{skip}$ and $c_{out}$ be $\\sigma_t$ and  $ \\alpha_t \\sigma_d$, that is for $\\sigma_t=t$ and  $\\alpha_t=1$: $c_{skip}=t$  and $c_{out}=\\sigma_d$ in line 201/202?\n\nd) In Equation (20) Appendix, shouldn't it be $\\hat{D}$ for $D_{\\theta}(x_t)=\\hat{D}_{\\theta}(\\hat{x}_t(x_t))$?\n\ne) The paragraph from line 924 to 938 (appendix) needs additional elaboration regarding the implications of having $||(\\alpha_t, \\sigma_t)||=1$ with respect to the invariance of the geometric set connecting $x_0$ and $z$.\n\n\n\n**Typos:**\n\ni) In line 126, a 2 is squared instead of the norm.\n\nii) in line 122, $z_t$ does not depend on time.\n\n\n**References**\n\n\n[1]  Tong et al. 2024. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport\n\n[2]  Kornilov et al 2024. Optimal Flow Matching: Learning Straight Trajectories in Just One Step", "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": "0NEBgKS68P", "reviewer": "Reviewer_pgsr", "summary": "This work proposed a set of improved training techniques to stabilize the training of continuous-time consistency models, including new consistency function formulations, new network architectures and new training objectives. With these new training techniques, the proposed method called sCMs outperformed all previous consistency models in terms of one-step and two-step FIDs.", "questions": "See the weaknesses in the above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is very well-written and easy to read. \n- This work proposed a new diffusion formulation, called TrigFlow, that unifies EDM and Flowing Matching, and also simplifies the analysis of continuous-time consistency models. \n- It provided a thorough analysis of the training stability of continuous-time consistency models, from the perspective of network architecture, training objective and diffusion process parameterization.\n- Experiments on CIFAR-10, ImageNet-64 and ImageNet-512 demonstrate the effectiveness of the proposed method and the scalability of continuous-time consistency models.", "weaknesses": "- Although I really like the improvements of continuous-time consistency models, which could fundamentally eliminate the discretization error in discrete-time consistency models, it comes with more time and memory costs related to JVP computation in the loss function. To this end, this work introduces JVP of Flash Attention to reduce the costs, which is great. Still, there may be a considerable gap between the continuous-time and discrete-time consistency models. I wonder if the paper can provide a more detailed comparison between sCMs and the previous discrete-time consistency models - ECMs, in terms of the training convergence and memory cost. \n- There is no explanation for the phenomenon that sCT performs better than sCD on CIFAR-10 and ImageNet-64, but sCTs performs worse on ImageNet-512. Any intuition of why sCT suffers from increased variance at larger scales? \n- There are no ablation study results on “Adaptive Double Normalization” except for claiming it “removes its instability in CM training”.\n- In Figure 5b, it looks like “w/o adaptive weighting” achieves better two-step FIDs than “w/ adaptive weighting” and very similar one-step FIDs to “w/ adaptive weighting”. Why do we need adaptive weighting?\n- In Figure 5c, do discrete-time CMs have a constant number of time steps $N$ or a timestep schedule up to the maximum number of steps $N$? If it is the former one, it seems to be a bit unfair to discrete-time CMs because the scheduling of time steps is very important to them. Does it make more sense to compare with the best-performing discrete-time CMs?\n- In Figure 7, does the paper apply TTUR proposed by DMD2 (Yin et al. 2024a)? From the DMD2 paper, TTUR improves the performance of VSD. Thus, a comparison with VSD + TTUR is more convincing.\n- In Figure 7, sCDs condition the consistency network on the guidance scale $s$. I wonder if VSD also condition the generator on the guidance scale, for a consistent evaluation setting? \n- A minor issue: In line 266, should it be $c_{\\text{noise}}(t) = \\frac{1}{4} \\log(\\sigma_d \\tan t) $?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)