Improved Techniques for Training Consistency Models
Authors: Yang Song
Prafulla Dhariwal
Venue: ICLR 2024 oral (2024)
Keywords: Consistency Models
Consistency Training
Diffusion Models
Score-Based Generative Models
Score-Based Diffusion Models
Distillation
Primary area: generative models
TL;DR: Consistency training works better than consistency distillation with improved techniques

Abstract: Consistency models are a nascent family of generative models that can sample high quality data in one step without the need for adversarial training. Current consistency models achieve optimal sample quality by distilling from pre-trained diffusion models and employing learned metrics such as LPIPS. However, distillation limits the quality of consistency models to that of the pre-trained diffusion model, and LPIPS causes undesirable bias in evaluation. To tackle these challenges, we present improved techniques for consistency training, where consistency models learn directly from data without distillation. We delve into the theory behind consistency training and identify a previously overlooked flaw, which we address by eliminating Exponential Moving Average from the teacher consistency model. To replace learned metrics like LPIPS, we adopt Pseudo-Huber losses from robust statistics. Additionally, we introduce a lognormal noise schedule for the consistency training objective, and propose to double total discretization steps every set number of training iterations. Combined with better hyperparameter tuning, these modifications enable consistency models to achieve FID scores of 2.51 and 3.25 on CIFAR-10 and ImageNet $64\times 64$ respectively in a single sampling step. These scores mark a 3.5$\times$ and 4$\times$ improvement compared to prior consistency training approaches. Through two-step sampling, we further reduce FID scores to 2.24 and 2.77 on these two datasets, surpassing those obtained via distillation in both one-step and two-step settings, while narrowing the gap between consistency models and other state-of-the-art generative models.

Published as a conference paper at ICLR 2024
IMPROVED TECHNIQUES FOR TRAINING CONSISTENCY MODELS
Yang Song & Prafulla Dhariwal
OpenAl (songyang, prafulla) a
@openai com
ABSTRACT
Consistency models are nascent family of generative models that can sample high quality data in one step without the need for adversarial training. Current consistency models achieve optimal sample quality by distilling from pre-trained diffusion models and employing learned metrics such as LPIPS. However, distil- lation limits the quality of consistency models to that of the pre-trained diffusion model, and LPIPS causes undesirable bias in evaluation. To tackle these challenges, we present improved techniques for consistency training, where consistency mod- els learn directly from data without distillation. We delve into the theory behind consistency training and identify a previously overlooked flaw, which we address by eliminating Exponential Moving Average from the teacher consistency model. To replace learned metrics like LPIPS, we adopt Pseudo-Huber losses from robust statistics. Additionally, introduce lognormal noise schedule for the consis- tency training objective, and propose to double total discretization steps every set number of training iterations. Combined with better hyperparameter tuning, these modifications enable consistency models to achieve FID scores of 2.51 and 3.25 on CIFAR- 10 and ImageNet 64 64 respectively a single sampling step. These scores mark 3.5x and 4x improvement compared to prior consistency training approaches. Through two-step sampling, we further reduce FID scores 2.24 and 2.77 these two datasets, surpassing those obtained via distillation in both one-step and two-step settings, while narrowing the gap between consistency models and other state-of-the-an generative models.
INTRODUCTION
Consistency models (Song et al., 2023) are an emerging family of generative models that produce high-quality samples using a single network evaluation. Unlike GANs (Goodfellow et al., 2014), consistency models are not trained with adversarial optimization and thus sidestep the associated training difficulty. Compared to score-based diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; 2020; Ho al., 2020; Song et al., 2021), consistency models do not require numerous sampling steps to generate high-quality samples. They are trained to generate samples in a single step, but still retain important advantages of diffusion models, such as the flexibility exchange compute for sample quality through multistep sampling, and the ability to perform zero-shot data editing.
We can train consistency models using either consistency distillation (CD) or consistency training (CT). The former requires pre- mre-raining diffusion model and distilling the knowledge therein into consistency model. The latter allows us train consistency models directly from data, establishing them as an independent family of generative models. Previous work (Song et al., 2023) demonstrates that CD significantly outperforms CT. However, CD adds computational overhead to the training process since it requires learning separate diffusion model. Additionally, distillation limits the sample quality of the consistency model to that the diffusion model. To avoid the downsides of CD and to position consistency models as independent family of generative models, aim to improve CT to either match or exceed the performance of CD.
For optimal sample quality, both CD and CT rely learned metrics like the Learned Perceptual Image Patch Similarity (LPIPS) (Zhang et al., 2018) in previous work (Song et al., 2023). However, depending LPIPS has primary downsides. Firstly, there could be potential bias in evaluation

---

Published as a conference paper at ICLR 2024
since the same ImageNet dataset (Deng et al., 2009) trains both LPIPS and the Inception network in Fréchet Inception Distance (FID) (Heusel et al., 2017), which is the predominant metric for image quality. As amalyzed in Kynkäänniemi et al. (2023), improvements of FIDs can come from accidental leakage of ImageNet features from LPIPS, causing inflated FID scores. Secondly, learned metrics require pre-training auxiliary networks for feature extraction. Training with these metrics requires backpropagating through extra neural networks, which increases the demand for compute.
To tackle these challenges, we introduce improved techniques for CT that not only surpass CD in sample quality but also eliminate the dependence on learned metrics like LPIPS. Our techniques are motivated from both theoretical analysis. and comprehensive experiments on the CIFAR- 10 dataset (Krizhevsky et al.. 2014). Specifically, we perform an in-depth study on the empirical impact of weighting functions, noise embeddings, and dropout in CT. Additionally, we identify an overlooked flaw in prior theoretical analysis for CT and propose simple fix by removing the Exponential Moving Average (EMA) from the teacher network. We adopt Pseudo- losses from robust statistics to replace LPIPS. Furthermore, we study how sample quality improves as the number of discretization steps increases, and utilize the insights to propose a simple but effective curriculum for total discretization steps. Finally, we propose new schedule for sampling noise levels in the CT objective based on lognormal distributions.
Taken together, these techniques allow CT to attain FID scores of 2.51 and 3.25 for CIFAR-10 10 and ImageNet 64 x 64 in one sampling step, respectively. These scores not only surpass CD but also represent improvements of 3.5x and 4x over previous CT methods. Furthermore, they significantly outperform the best few-step diffusion distillation techniques for diffusion models even without the need for distillation. By two-step generation, we achieve improved FID scores of 2.24 and 2.77 on CIFAR-10 10 and ImageNet 64 x 64, surpassing the scores from CD in both one-step and two-step settings. These results rival many top-tier diffusion models and GANs, showcasing the strong promise of consistency models as a new independent family of generative models.
CONSISTENCY MODELS
Central to the formulation of consistency models is the probability flow ordinary differential equation (ODE) from Song et al. (2021). Let us denote the data distribution by Pasta (x). When we add Gaussian noise with mean zero and standard deviation a to this data, the resulting perturbed distribution is given by (x(x)= Pdwaa (yaaay)).N [y)x(x v dy. The probability flow ODE, presented in Karras et al. (2022), takes the form of
dx
-o (x) a E [omins mmax]; do
where the term Vx log Pa (x) is known as the score function of Pa (x) (Song et al., 2019; Song & Ermon, 2019; 2020; Song et al., 2021). Here min is a small positive value such that Pa (x) & Pdma(x). introduced to avoid numerical issues in ODE solving. Meanwhile, mmax is sufficiently large so that Pa (x) x (0, 1). Following Karras et al. (2022); Song et al. (2023), we adopt Tmin 0.002, and Fmax 80 throughout the paper. Crucially, solving the probability flow ODE from noise level 01 to 02 allows us to transform a sample Xx, ~ Po, (x) into X ~ (x).
The ODE in Eq. (1) establishes a bijective mapping between noisy data sample x, (x( and Xom. ~ Pomm (x) 2 (x) This mapping. denoted as f* (xx, σ) + Xom. is termed the funcittiocy By its very definition, the consistency function satisfies the boundary condition f*(x, min)) = x. A consistency model, which we denote by fe(x, a), is neural network trained to approximate the consistency function f* (x, r). To meet the boundary condition, we follow Song et al. (2023) to parameterize the consistency model as
fo(x, σ) caip(o)x + Fo(x,
where (x, o) is a free-form neural network, while ca)i((o) and (o) are differentiable functions such that Cakip min) = 1 and Cour foun(-mi))= = 0.
To train the consistency model, we discretize the probability flow ODE using sequence of noise
levels Tmin = 01 < 02 < ... < NN = mmax. where we follow Karras et al. (2022); Song et al. (2023) in setting a; (a 1/1 + 1-1
N-1 (dii
1/p ) for [1, NT. and p 7, where [a, bl] denotes the set
2

---

Published as a conference paper at ICLR 2024
of integers {a, a + 1, .b). The model is trained by minimizing the following consistency matching (CM) loss over 0:
L (0, 0 - B (x 11 So- (Ka,, oz))];
where Xx. Xor-1 (o; 0++1)0;+1 Vx log Po..: (x))x_
Eq. (3), d(x y) is metric function comparing vectors æ: and y. and X(o) > 0 is a weighting function. Typical metric functions include the squared f2 metric d(xr. y) = I| - yll and the Learned Perceptual Image Patch Similar- ity (LPIPS) metric introduced in Zhang al. (2018). The expectation Eq. (3) is taken over the where wh N - !] represents the uniform distribution over 21. N - -],, and X ~ Po; , (x). Note that Xa; is derived from Xx, by solving the probability flow ODE in the reverse direction for a single step. In Eq. (3), fo and fo- are referred to as the student network and the teacher network, respectively. The teacher's parameter O obtained by applying Exponential Moving Average (EMA) to the student's parameter & during the course of training as follows:
following sampling process: i U[1,
0- stopgrad(w0 + (1 -0)0,,
with & J < representing the EMA decay rate. Here we explicitly employ the stopgrad operator to highlight that the teacher network remains fixed for each optimization step of the student network. However, in subsequent discussions, we will omit the slopgrad operator when its presence is clear and unambiguous. In practice, we also maintain EMA parameters for the student network to achieve better sample quality at inference time. is clear that as N increases, the consistency model optimized using Eq. (3) approaches the true consistency function. For faster training, Song al. (2023) propose a curriculum learning strategy where N is progressively increased and the EMA decay rate μ is adjusted accordingly. This currieulum for N and μ is denoted by N (k) and ((k), where E N is non-negative integer indicating the current training step.
Given that x.: relies on the unknown score function Vx log Pa:11 (x), directly optimizing the consis- tency matching objective in Eq. (3) is infeasible. To circumvent this challenge, Song et al. (2023) propose two training algorithms: consistency distillation (CD) and consistency training (CT). For consistency distillation, we first train diffusion model & (xx о) to estimate Vx Po po (x) via score matching (Hyvärinen, 2005; Vincent, 2011: Song et al., 2019; Song Ermon, 2019), then approximate Xz, with &g X (o; - dz+1 1 sistency training employs a different approximation method. Recall that X = =x+ 0+41% with Pdata (x) and (O, I). Using the same x and z, Song et al. (2023) define Xo. + 0, as an approximation to X which leads to the consistency training objective below:
0;+1). On the other hand, con-
(0, 0 E [X(o:)d + 0++12, 1)+1). fo- o,2, o))].
(5)
As analyzed in Song et al. (2023), this objective is asymptotically equivalent to consistency matching in the limit of N xc. We will revisit this analysis in Section 3.2.
After training a consistency model fe(x, a) through CD or CT, we can directly generate a sample x by starting with z NO o and computing x fe [z, o max). Notably, these models also enable multistep generation. For a sequence of indices 1 = ii < i2 <.. ik = N, we start by sampling Xx ~ ~N(0, 2 I) and then iteratively compute X < fo(x oi.
2 Z for - 1, K 2,... 1 where 2 ~ .N(,, I). The resulting sample x, approximates the distribution Pdata (x).. In experiments, setting K = 3 (two-step generation) often enhances the quality of one-step generation considerably, though increasing the number of sampling steps further provides diminishing benefits.
3 IMPROVED TECHNIQUES FOR CONSISTENCY TRAINING
Below we re-examine the design choices of CT in Song et al. (2023) and pinpoint modifications that improve its performance, which we summarize in Table 1. We focus on CT without learned metric functions. For our experiments, we employ the Score SDE architecture in Song et al. (2021) and train the consistency models for 400,000 iterations on the CIFAR-10 dataset (Krizhevsky et al., 2014) without class labels. While our primary focus remains CIFAR-10 in this section, observe similar improvements on other datasets, including ImageNet 64 x 64 (Deng et al., 2009). We measure sample quality using Fréchet Inception Distance (FID) (Heusel al., 2017).
3

---

Published as a conference paper at ICLR 2024
Table Comparing the design choices for CT in Song et al. (2023) versus our modifications.
Design choice in Song el al. (2023)
Our modifications
EMA decay rate for the teacher network
p((k) expl N(()
p(k)
Metric consistency loss
d(z. v) LPIPS((,
d(z, y) V -
Discretization curriculum
N(k)
[R) min( min(302 21
[V V 1 1)2 +)) % 1|11
where K'
logal /7
Noise schedule
0₁. where i (())
og where i ~ p(i), and p()) orf [ (oslere1)) Prem)
orf loogi Pun
Weighting function
Max
M(a,) O.L
9o 2, 91 150, po 0.9 on CIFAR-10
9p 10, 81 -280
Su = 2, 81 = 200, Mo 0.95 on ImageNet 64 x 64
c 0.00054 Vē, is data dimensionality
Parameters
Pazan 111,, Pus 2.0
As e [0, K]. where K is the total training itcrations
Up
Gaal o 2i))), where i e [1.N 7. 0.002, ams 80
min
3.1 WEIGHTING FUNCTIONS, NOISE EMBEDDINGS, AND DROPOUT
We start by exploring several hyperparameters that are known to be important for diffusion models, including the weighting function A(a). the embedding layer for noise levels, and dropout (Ho et al., 2020; Song et al., 2021; Dhariwal Nichol, 2021; Karras et al., 2022). We find that proper selection of these hyperparameters greatly improve CT when using the squared &2 metric.
The default weighting function in Song et al. (2023) uniform, i..,, M(a) = 1. This assigns equal weights to consistency losses at all noise levels, which we find to be suboptimal. We propose to modify the weighting function so that it reduces as noise levels increase. The rationale is that errors from minimizing consistency losses in smaller noise levels can influence larger ones and therefore should be weighted more heavily. Specifically, our weighting function (cf., Table 1) is defined as A(a,) a111-,"" The default choice for d, given in Section 2, ensures that A(a.) reduces monotonically as o; increases, thus assigning smaller weights higher noise levels. As shown in Fig. 1c, this refined weighting function notably improves the sample quality in CT with the squared f2 metric.
de 1:11 1-1,
In Song et al. (2023), Fourier embedding layers (Tancik et al., 2020) and positional embedding layers (Vaswani al., 2017) are used embed noise levels for CIFAR- 10 and ImageNet 64 64 respectively. It is essential that noise embeddings are sufficiently sensitive to minute differences to offer training signals, yet too much sensitivity can lead to training instability. As shown Fig. 1b, high sensitivity can lead to the divergence of continuous-time CT (Song et al., 2023). This is known challenge in Song et al. (2023), which they circumvent by initializing the consistency model with parameters from pre-trained diffusion model. In Fig. 1b, show continuous-timg CT on CIFAR-10 converges with random initial parameters, provided we use a less sensitive noise embedding layer with a reduced Fourier scale parameter, as visualized in Fig. Ia. For discrete-time CT, models are less affected by the sensitivity of the noise embedding layers, but as shown in Fig. Ic, reducing the scale parameter in Fourier embedding layers from the default value of 16.0 still leads to slight improvement of FIDs on CIFAR-10. For ImageNet models, we employ the default positional embedding, as it has similar sensitivity to Fourier embedding with scale 0.02 (see Fig. la).
smaller value of 0.02
Previous experiments with consistency models in Song et al. (2023) always employ zero dropout, motivated by the fact that consistency models generate samples in a single step, unlike diffusion models that do so in multiple steps. Therefore, it is intuitive that consistency models, facing more challenging task, would be less prone to overfitting and need less regularization than their diffusion counterparts. Contrary to our expectations, we discovered that using larger dropout than diffusion models improves the sample quality of consistency models. Specifically, as shown in Fig. 1c, a dropout rate of 0.3 for consistency models CIFAR-10 obtains better FID scores. For ImageNet 64 x 64, we find it beneficial to apply dropout of 0.2 to layers with resolution less than or equal to 16 x 16, following Hoogeboom et al. (2023). We additionally ensure that the random number
4

---

Published as a conference paper at ICLR 2024
400
Averngo aimorerc between
300
200
200
Rew aniigttnng
Bandiine Excuared
how waighting lequared
20
Fourler scale 0.2
100
Fourter 16.0 Fourior
Dasel the (LFPPS)
Iow rablo booding
Joquares
Rew ieiighting
40
doououl (UMPS)
embeedding
10
10
1000
10
20 Training steps (x:1
discretization
Training steps
(a) Sensitivity of noise embeddings.
(b) Continuous-time CT.
(c) Ablation study.
Figure (a) As the Fourier scale parameter decreases, Fourier noise embeddings become less sensitive to minute noise differences. This sensitivity is closest to that of positional embeddings when the Fourier scale is set to 0.02. (b) Continuous- sensitive minor noise differences. (c) An ablation study examines the effects of our selections for weighting function -,. ), noise embedding (Fourier scale = 0.02), and dropout (= 0.3) on CT using the squared l2 metric. Here baseline models for both metrics follow configurations in Song et al. (2023). All models are trained on CIFAR-10 without class labels.
CT diverges when noise embeddings are overly
generators for dropout share the same states across the student and teacher networks when optimizing the CT objective in Eq. (5).
By choosing the appropriate weighting function, noise embedding layers, and dropout, we signifi- cantly improve the sample quality of consistency models using the squared f2 metric, closing the gap with the original CT in Song et al. (2023) that relies on LPIPS (see Fig. 1c). Although our modifications do not immediately improve the sample quality of CT with LPIPS, combining with additional techniques in Section 3.2 will yield significant improvements for both metrics.
3.2 REMOVING EMA FOR THE TEACHER NETWORK
When training consistency models, we minimize the discrepancy between models evaluated at adjacent noise levels. Recall from Section that the model with the lower noise level is termed the teacher network, and its counterpart the student network. While Song et al. (2023) maintains EMA parameters for both networks with potentially varying decay rates, we present a theoretical argument indicating that the EMA decay rate for the teacher network should always be zero for CT, although it can be nonzero for CD. We revisit the theoretical analysis in Song et al. (2023) to support our assertion and provide empirical evidence that omitting EMA from the teacher network in CT notably improves the sample quality of consistency models.
To support the of CT, Song et al. (2023) present two theoretical arguments linking the CT and CM objectives as N - oc. The first line of reasoning, which we call Argument (i), draws upon Theorem 2 from Song et al. (2023) to show that under certain regularity conditions, C2 00,00'00 =
L-V (0, 0 ) + of a(Ao). That is, when N - cc, we have Ao O and hence 2 O, 0-) converges to LN (0, 0 asymptotically. The second argument, called Argument (ii), is grounded in Theorem 6 from Song et al. (2023) which asserts that when 0 = 0, both 1im (N - V) (0,00) and lim (N 1) Y 0-) are well-defined and identical. This suggests that after scaling by N - 1, gradients of the CT and CM objectives match in the limit of N - oc, leading to equivalent training dynamics. Unlike Argument (i), Argument (ii) is valid only when 0 0, which can be enforced by setting the EMA decay rate for the teacher network to zero in Eq. (4).
We show this inconsistency in requirements for Argument (i) and (ii) to hold is caused by flawed
theoretical analysis of the former. Specifically, Argument (i) fails if lim LN (0, 0 ) is not valid objective for learning consistency models, which we show can happen when 0- * 0. To give concrete example, consider data distribution (I) = 8(1 E), which leads to Px) = N(x; E, o²) and a ground truth consistency function f* 0) = Suis r + (1 Smin E. Let us define Aem
the consistency model as fo(x, = Aen x + (1 -1m=mm 0. In addition, let o; = Omin + N-I (amax
i-1
mmin) for E [1 be the noise levels, where we have Ao N -ann Given 2 A 1) and
Toat E++ 2, it is straightforward to show that Co.; Tots1 o++1 (o+ G V. log Po (3artt)
5

---

Published as a conference paper at ICLR 2024
200
200
200
Basetine Esquared E:1
6
improved isquared 6s)
Inquared wil
100
LPPS
Pseudd llubrr ic- 0.0011
100
LIPS
Poeuso-Hluber ic 0
Psoudo Huber (00000
Pscudo- 0.0033
50
Penude huber 0.03)
50
Pleuuo huber 0.01)
Impreved ILPIPS,
Pseudo 0.03)
Huber 0.03)
20
20
20
10 20
30
60
10
20
3D
60 0
10
20 Training steps (x101)
40
Training steps (x101)
Training steps (x10*)
(a) LPIPS & squared &g metrics.
(b) 50 2,
150
(c) 80 10, 81
1280
Figure 2: (a) Removing EMA in the teacher network leads to significant improvement in FIDs. (b, c) Escudo-Huber metrics significantly improve the sample quality of squared lz metric, and catches up with LPIPS when using overall larger N (k), where the Pseudo-Hube metric with 0.03 is the optimal. All training runs here employ the improved techniques from Sections 3.1 and 3.2.
simplifies to L., E + d.z. As a result, the objectives for CM and CT align perfectly this toy example. Building on top of this analysis, the following result proves that limy ( is not amenable for learning consistency models whenever 0 * 8.
(0, 0,,0) here
Proposition 1. Given the notations introduced earlier, and using the uniform weighting function
((r) = 1 along with the squared lg metric, we have
lim LZ (0,00
= lim CC 0
=E
Omin
2 (0
0
if #
(6)
-0
d BB
1
Oan a,
(0 - E)2
0 0
lim 1 dC~((..0))
dB
of
N Do
de
400,
0 < 0
-00,
0 > 0
Proof. See Appendix A.
Recall that typically 6- * θ when * 0. In this case, Eq. (6) shows that the CM/CT objective independent of &, thus providing no signal of the data distribution and are therefore impossible to train correct consistency models. This directly refutes Argument (i). In contrast, when we set u = 0 to ensure f- 8, Eq. (7) indicates that the gradient of the CM/CT objective, when scaled by 1/Ar,, converges to the gradient of the mean squared error between & and &. Optimizing this gradient consequently yields 0 = ç, accurately learning the ground truth consistency function. This analysis is consistent with Argument (ii).
As illustrated in Fig. 2a, discarding EMA from the teacher network notably improves sample quality for CT across both LPIPS and squared fg metrics. The curves labeled "Improved" correspond to CT using the improved design outlined in Section 3.1. Setting p(k) 0 for all training iteration k effectively counters the sample quality degradation of LPIPS caused by the modifications in Section 3.1. Combining the strategies from Section 3.1 with zero EMA for the teacher, we are able match the sample quality of the original CT in Song al. (2023) that necessitates LPIPS, by using simple squared &g metrics.
3.3 PSEUDO- HUBER METRIC FUNCTIONS
Using the methods from Sections 3.1 and 3.2, we are able to improve CT with squared fz metric. matching the original CT in Song et al. (2023) that utilizes LPIPS. Yet, as shown in Fig. 2a, LPIPS still maintains a significant advantage over traditional metric functions when the same improved techniques are in effect for all. To address this disparity, we adopt the Reeudo-Huber metric family (Charbonnier et al., 1997), defined as
a y) = VII - 212 + c c,
(8)

---

Published as a conference paper at ICLR 2024
where is adjustable constant. As depicted in Fig. 5a, Pseudo- Huber metrics smoothly bridge the { and squared lg metrics, with c determining the breadth of the parabolic section. In contrast to common metrics like En, l, and Ex Pseudo-Huber metrics are continuously twice differentiable, and hence meet the theoretical requirement for CT outlined in Song et al. (2023).
Compared to the squared & metric, the Seeudo-Huber metric is more robust outliers it imposes a smaller penalty for large errors than the squared &g metric does, yet behaves similarly for smaller errors. We posit that this added robustness can reduce variance during training. To validate this hypothesis, we examine the l2 norms of parameter updates obtained from the Adam optimizer during the course of training for both squared lg and Pseudo- Huber metric functions, and summarize results in Fig. 5b. Our observations confirm that the Geoudo-Huber metric results in reduced variance relative to the squared fz metric, aligning with our hypothesis.
We evaluate the effectiveness of Pseudo- Huber metrics by training several consistency models with varying e values on CIFAR-10 and comparing their sample quality with models trained using LPIPS and squared f2 metrics. We incorporate improved techniques from Sections 3.1 and 3.2 for all metrics. Fig. reveals that Pseudo-Huber metrics yield notably better sample quality than the squared &2 metric. By increasing the overall size of (Kk)-adjusting 5o and from the standard values of 2 and 150 in Song et al. (2023) to our new values of 10 and 1280 (more in Section 3.4)-we we for the first time surpass the performance of CT with LPIPS equal footing using traditional metric function that does not rely on learned feature representations. Furthermore, Fig. 2c indicates that 0.03 is optimal for CIFAR-10 images. We suggest that should scale linearly with ||æ: ul12- and propose a heuristic of e = 0.00054y for images with d dimensions. Empirically, we find this recommendation to work well on both CIFAR-10 10 and ImageNet 64 x 64 datasets.
3.4 IMPROVED CURRICULUM FOR TOTAL DISCRETIZATION STEPS
As mentioned in Section 3.2, CT's theoretical foundation holds asymptotically as - co. In practice, we have to select finite N for training consistency models, potentially introducing bias into the learning process. To understand the influence of N on sample quality, we train a consistency model with improved techniques from Sections 3.1 to 3.3. Unlike Song et al. (2023), we use an exponentially increasing curriculum for the total discretization steps .N, doubling N after a set number of training iterations. Specifically, the curriculum is described by
(1) = min (sp2) #1, s1) + 11, K' $1gg|s//am So
K
(9)
and its shape is labelled "Exp" in Fig. 3b. Here So and 81 control the minimum and maximum number of discretization steps, and K" is the total number of training iterations.
As revealed in Fig. 3a, the sample quality of consistency models improves predictably as N increases. Importantly, FID scores relative to N adhere to precise power law until reaching saturation, after which further increases in N yield diminishing benefits. As noted by Song al. (2023), while larger N can reduce bias in CT, they might increase variance. On the contrary, smaller N reduces variance the cost of higher bias. Based on Fig. 3a, we cap N at 1281 in N(k),, which we empirically find to strike a good balance between bias and variance. In our experiments, we set 8p and 81 in discretization curriculums from their default values 2 and 150 in Song et al. (2023) to and 1280 respectively.
Aside from the exponential curriculum above, we also explore various shapes for N(k) with the same 80 10 and 8 1280, including a constant function, the square root function from Song et al. (2023), a linear function, square function, and a cosine function. The shapes of various curriculums illustrated in Fig. 3b. As Fig. 3c demonstrates, the exponential curriculum yields the best sample quality for consistency models. Consequently, we adopt the exponential curriculum in Eq. (9) as our standard for setting (kk) going forward.
3.5 IMPROVED NOISE SCHEDULES
Song et al. (2023) propose to sample a random from 21 N 1] and select o; and 0 to compute the CT objective. Given that o; (o 1/p 1-1 Lill σ H/2)) this corresponds sampling from the distribution p(log o) pla 1 xip as N cc. As shown in Fig. 4a, this distribution exhibits
N-1
7

---

Published as a conference paper at ICLR 2024
100
000
Bashee
1350
870 i 1283
1200
20
100
F10
10
E
200
Bosches
Ulwer
260
Cosine
Square
10
20 40 100
230
400
1000
Number of discretization stops IN)
Training stops (x10P)
Trainnn stops (x10
(a) FID scores vs. N
(b) Various curriculums for (k).
(c) FIDs vs. (k) curriculums.
Figure 3: (a) FID scores improve predictably the number of discretization steps N grows. (b) The shapes of various curriculums for total discretization steps N (k). (c) The FID curves of various curriculums for discretization. All models are trained with improved techniques from Sections 3.1 to 3.3 with the only difference in discretization curriculums.
higher probability density for larger values of log a. This is odds with the intuition that consistency losses at lower noise levels influence subsequent ones and cause error accumulation, so losses at lower noise levels should be given greater emphasis. Inspired by Karras et al. (2022), we address this by adopting lognormal distribution to sample noise levels, setting mean of -1.1 and standard deviation of 2.0. As illustrated in Fig. 4a, this lognormal distribution assigns significantly less weight to high noise levels. Moreover, it also moderates the emphasis smaller noise levels. This is helpful because learning is easier at smaller noise levels due to the inductive bias in our parameterization of the consistency model to meet the boundary condition.
For practical implementation, we sample noise levels in the set (0,,00,"" ,ON} according to a discretized lognormal distribution defined
p(o,) erf (oo((o++))- - mean erf (log(o;) - mean
(10)
V2Paas
/2Psau
where Pincan 1.1 and Pad 2.0. As depicted in Fig. 4b, this lognormal noise schedule significantly improves the sample quality of consistency models.
4 PUTTING IT TOGETHER
Combining all the improved techniques from Sections 3.1 to 3.5, we employ CT to train several consistency models on CIFAR-10 and ImageNet 64 x 64 and benchmark their performance with competing methods in the literature. We evaluate sample quality using FID (Heusel al., 2017), Inception score (Salimans et al., 2016), and Precision/Recall (Kynkäänniemi et al., 2019). For best performance, We
larger batch size and an increased EMA decay rate for the student network in CT across all models. The model architectures are based on Score SDE (Song et al., 2021) for CIFAR-10 10 and ADM (Dhariwal & Nichol, 2021) for ImageNet 64 x 64. We also explore deeper variants of these architectures by doubling the model depth. We call our method iCT which stands for "improved consistency training". and the deeper variants iCT-deep. We summarize our results in Tables 2 and 3 and provide uncurated samples from both iCT and iCTT-deep in Figs. 6 to 9, More experimental details and results are provided in Appendix B.
It is important to note that we exclude methods based on FastGAN (Liu et al., 2020; Sauer et al., 2021) or StyleGAN-XI (Sauer al., 2022) from our comparison, because both utilize ImageNet pre-trained feature extractors in their discriminators. As noted by Kynkäänniemi et al. (2023), this can skew FIDs and lead to inflated sample quality. Methods based on LPIPS suffer from similar issues, LPIPS is also pre-trained on ImageNet. We include these methods in Tables 2 and 3 for completeness, but we do not consider them as direct competitors to iCT or iCT-deep methods.
Several key observations emerge from Tables 2 and 3. First, iCT methods surpass previous diffusion distillation approaches in both one-step and two-step generation on CIFAR- 10 and ImageNet 64 x 64, all while circumventing the need for training diffusion models. Secondly, iCT models demonstrate
8

---

Published as a conference paper at ICLR 2024
Table Comparing the quality of unconditional samples on CIFAR-10.
Table 3:
Comparing the quality of class-
conditional samples on ImageNet 64 x 64.
METHOD
NFE (J) FID (J) 1S (t)
METHOD
(F) FID (1) Prec. (1) Rec. (1)
Pust samplers & distillation for difftasion models
Fust samplers & distillation for diffusion models
DDIM (Song et al., 2020)
10
13.36
DDIM (Song al., 2020)
13.7
0.65
0.56
DPM-salva--fat (La al. 2022)
10
4,70
10
18.3
0.60
0.49
3-DEIS (Zhang& & Chen, 2022)
10
4.17
DPM solver (Lu al., 2022)
10
7.93
UniPC (Zhao ot sl., 2023)
10
3.87
3.42
Knowledge Distillation (Luhman Luhman, 2021)
DEIS (Zhang & Chen, 2022)
10
6.65
DENO (1 (LPS)) (Zheng al., 2022)
3.78
20
3.10
2 Reciified Flow (rdistill) (Liu et al., 2022)
4.85
9.01
DFNO (LPIPS) (Zheng et al., 2022)
7.83
0.61
TRACT (Berthelot al., 2023)
TRACT (Berthelot et 2023)
7.43
2
3.32
4.97
Dnff-Instruct (Luo et al., 2023)
4.53
9.89
BOOT (Gu et 2023)
16.3
0.68
0.36
P1 (Salimans 2023)
8:34
8.69
Diff-Instrect (Lao et al., 2023)
5.57
2
5.58
9.05
PD* (Salimans Ho, 2022)
15.39
0.59
0.62
CD (1 (LPISS) (Song et al., 2023)
3.55
9.48
2
2.93
9.75
8:95
0.63
0.65
Direct Generation
6.77
0.66
0.65
Score SDE (Song al., 2021)
2000
2.38
9.83
(1.PIPS) (Song al., 2023)
7.88
0.66
0.63
Scare SDE (dopp) (Song 2021)
2000
5.74
0.67
0.65
DDPM (Ho al. 2020)
1000
3.17
9.46
4.92
0.68
0.65
LSGM (Valdatt et al., 2021)
147
2.10
CD (1PPIS) (Song et al., 2023)
6.20
0.68
0.63
PRGM al., 2022)
110
2.35
9.68
4:70
0.69
0.64
EDM* (Karras et all. 2022)
35
2.04
9.84
4.32
0.70
0.64
EDM-G-+ (Kim et 2023)
36
Direct Generation
IGEBM (Du Mordatch, 2019)
60
40.6
6.02
RIN (Jabri et al., 2023)
1000
1.23
NVAE (Vabdat & Kautz, 2020)
I
23.5
7.18
DDPM (Hl et al., 2020)
250
11.0
0.67
0.58
Glow (Kingma Dhariwal, 2018)
3.92
iDDMM (Nichol Dbariwal, 2021)
2.92
0.74
0.62
Residual Flow (Chen 2019)
46.4
ADM (Dhuriwal & Nichol, 2021)
250
2,07
0.74
0.63
BigGAN (Bmak et al., 2019)
14.7
9.22
EDM (Karras et al., 2022)
511
1,36
StyleGAN2 (Karras al., 2020b)
EDM" (Heun) (Karaa et al., 2022)
79
2.44
0.71
0.67
StyleGAN2-ADA (Kamas et al., 2020aa
2.92
9.83
BigGAN deep (Brock 2019)
4.06
0.79
0.48
CT (LPIPS) (Song al., 2023)
8.70
8.49
CT (LPIPS) (Song al., 2023)
13.0
0,71
0.47
5,83
8.85
11.1
0.69
0.56
iCT (ours)
1
2,83
9.54
ICT (ours)
4.02
0.70
0.63
9.50
3.20
0.73
0.63
iCT-deep (ours)
1
2.51
9.76
ICT-deep (ours)
3.25
0.72
0.63
2
2,24
9.89
2.77
0.74
0.62
Most results for existing methods are taken from a previous paper. except for those marked with *, which are from our own se-implementation.
sample quality comparable to many leading generative models, including diffusion models and GANs. For instance, with one-step generation, iCT-deep obtains FIDs of 2.51 and 3.25 for CIFAR-10 and ImageNet respectively, whereas DDPMs (Ho al., 2020) necessitate thousands of sampling steps to reach FIDs of 3.17 and 1 11. (result taken from Gu et al. (2023)) on both datasets. The one-step FID for iCT already exceeds that of StyleGAN-ADA (Karras et al., 2020b) on CIFAR-10, and that of BigGAN-deep (Brock et al., 2019) on ImageNet 64 x 64, let alone iCT-deep models. For two-step generation, iCT-deep records an FID of 2.24, matching Score SDE in Song et al. (2021), a diffusion model with an identical architecture but demands 2000 sampling steps for an FID of 2.20. Lastly, iCT methods show improved recall than CT (LPIPS) in Song et al. (2023) and BigGAN indicating better diversity and superior mode coverage.
5 CONCLUSION
Our improved techniques for CT have successfully addressed its previous limitations, surpassing the performance of CD in generating high-quality samples without relying on LPIPS. We examined the impact of weighting functions, noise embeddings, and dropout. By removing EMA for teacher networks, adopting Pseudo- Huber losses in lieu of LPIPS, combined with a new curriculum for discretization and noise sampling schedule, we have achieved unprecedented FID scores for consis- tency models on both CIFAR- 10 and ImageNet 64 x 64 datasets. Remarkably, these results outpace previous CT methods by a considerable margin, surpass previous few-step diffusion distillation techniques, and challenge the sample quality of leading diffusion models and GANs.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We would like to thank Alex Nichol, Allan Jabri, Ishaan Gulrajani, Jakub Pachocki, Mark Chen and Ilya Sutskever for discussions and support throughout this project.
REFERENCES
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbot, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.04248 2023.
Andrew Brock. Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019. URL
net/ :ttp://wwwee//euttteemem
Pierre Charbonnier, Laure Blanc-Féraud, Gilles Aubert, and Michel Barlaud. Deterministic edge- preserving regularization in computed imaging. IEEE Transactions image processing. 6(2): 298-311, 1997.
Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Jömn-Henrik Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, pp. 9916-9926, 2019.
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-255 leee, 2009.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. arXiv preprint arXiv:2105.0522 2021.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based mod- els. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019, URL httphpppss//rrrcccccng.occo 2019/file/ 17.11le/3//000 (ts99//2222aacbcccccccceeeeaaae2
ec wwtppppp/applaeeeeeeeeeeeeeeeepe fellffiess/eeeees fil.es/pppee/////////p pdf.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural informa- tion processing systems. pp. 2672-2680, 2014.
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free distil- lation of denoising diffusion models with bootstrapping. In ICML 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. Advances in Neural Information Processing Systems, 33, 2020.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301.1109. 2023.
Aapo Hyvärinen. Estimation of Non-Normali Statistical Models by Score Matching. Journal of Machine Learning Research. 6(Apr):695-709 2005.
Allan Jabri, David J. Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
10

---

Published as a conference paper at ICLR 2024
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in neural information processing systems, 33:12104-12114 2020a.
Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. 2020b.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Proc. NeurIPS, 2022.
Dongjun Kim, Yeongmin Kim, Se Jung Kwon, Wanmo Kang, and Il-Chul Moon. Refining gener- ative process with discriminator guidance in score-based diffusion models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 16567-16598, PMLR, 23-29 Jul 2023. URL http:::/pprceeeiii
mmlr wae.sppres//ve /vlsss/v2222 kim23i........ html.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. In Advances in Neural Information Processing Systems, pp. 10215-10224, 2018.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. The CIFAR-10 Dataset. online: http:///ww.. cs. toronto. edu/kriz/cfar. huml, 55, 2014.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019.
Tuomas Kynkäänniemi, Tero Karras, Miika Aittala, Timo Aila, and Jaakko Lehtinen. The role of im- agenet classes in fréchet inception distance. In The Eleventh International Conference on Learning Representations, 2023. URL https: /ttp:////eppeeeeeeeeeeeeeeeeeeee wwt.:/iwwwneeeeeeeeemmmmeeeeeeee n forum 2tw.rmmmm00440000mmmm
Bingchen Liu, Yizhe Zhu, Kunpeng Song, and Ahmed Elgammal. Towards faster and stabilized gan training for high-fidelity few-shot image synthesis. In International Conference on Learning Representations, 2020.
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: 1908. 2019.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint rrXiv:2209.03033. 2022.
Cheng Lu, Yuhao Zhou, Fan Bao. Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.023883 2021.
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. arXiv preprint arXiv:2305.184. 2023.
Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2102.09672. 2021.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022, URL https: https:t/opeeeeeeeeeeeee /pppreevie..... net/ d-t
11

---

Published as a conference paper at ICLR 2024
Tim Salimans, Ian J. Goodfellow, Wojeiech Zaremba, Vicki Cheung. Alec Radford, and Xi Chen. Improved techniques for training GANs. In Daniel D. Lee, Masashi Sugiyama, Ulrkee von Luxburg. Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 2226-2234, 2016. URL https: http///oooccccoeecoocooocooec ings bag.... nww. ps. cc/ 7tt.aa////8333aaa22222222228777 01t.ae/0099000200666666000000000 Abbtaasasataaaaaaaaaattt .tt.raattttmm.....
Axel Sauer, Kashyap Chitta, Jens Müller, and Andreas Geiger. Projected gans converge faster. Advances in Neural Information Processing Systems, 11417780-17492 2021.
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.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep Unsupervised Learning Using Nonequilibrium Thermodynamics. In International Conference Machine Learning, pp. 2256-2265, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances Neural Information Processing Systems, pp. 11918-1 11930, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score- score-based generative models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced Score Matching: A Scalable Approach to Density and Score Estimation. In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence, pp. 204, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021. URL https: http:////eeeeeeeee.eeeee net forum!
id=PxTIG12RRHS.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 102211332252 PMLR, 23-29 Jul 2023. URL https: /ttp///peceee..cccpeeppppo mll. p.r..ppresss... / /202 song html.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich- Keil, 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.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 33:1 1967-999 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 14:1 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Pascal Vincent. A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7):1661-1674 2011.
12

---

Published as a conference paper at ICLR 2024
Yilun Xu, Ziming Liu, Max Tegmark, and Tommi S. Jaakkola. Poisson flow generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL dd=woV_TRqchh.
http://wwwww....eo........ forum?
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204.13900 2022.
Richard Zhang. Phillip Isola, Alexei Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features perceptual metric. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 586-595, 2018.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. arXiv preprint arXiv:2302.04680
2023.
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar, Fast sampling of diffusion models via operator learning. arXiv preprint arX14:2211.13444, 2022.
13

---

Published as a conference paper at ICLR 2024
A PROOFS
Proposition 1. Given the notations introduced in Section 3.2, and using the uniform weighting funcctio Mo) along with the squared metric, have
lim N
L (0,0 = lim
(0r(0,0)
mmie 2 (0 - )²
if0 * 0
(11)
N-I
0i
man
(0 - -3)3 ;:
B 0
1 dC
do
of
N Ao d0
lim
+00,
B 0
(12)
-00,
0 > 0
Proof. Since \(o) = 1 and d(x,)) = (Ir - y/), we can write down the CM and CT objectives
L (0, 0- ([) So (#a, )²] and (0, 0-) [ 0--1 ( 0,))²1 respectively. Since Paata (z) = 5(x E), we have pa(x) N(x | {, and therefore V Vlog
5=S According to the definition of oo, and Zo = E 0i+12, we have
to. = Tou41 - (a; 0:41)0++1 V Vlog P(oo
Tous + (o, - 0:+1)0 To E
of
Taai + (o; 0;+1)=
= E+0+1 + (o; - 01+1)=
+
As a result, the CM and CT objectives are exactly the same, that is, L C~100, = L.Y (0,00
Recall that the consistency model fa(x, o) is defined as ()) Jum + (1 -mis) 0, so we have fo o) = min + Onin E + (1- omnn 0. Now, let us focus on the CM objective
L (0,0
0241) fo- (#a. 0;))²1
E((fo(Foos. 0441) fo (ăo, {3
2
E
Tmin {
+
miin
0
mmin c
min
0
0it1
0i+1
di
27
E
Omin { +
1
min
M'in
1
Omin
)
0
0 + Ao
+
Ao
where Ao Cuie. because σ; Omin
N
i-1 N
mmmx
minn).
By taking the limit N 00, we
have Ao + 0. and therefore
lim (O,00
2
lim E
m
E +
1
T Fiin
0
mmin
min
0
40-0
0.
Ad
0 + Ao
0
= lim E
mmin
Ao
+
T
Tmin
g
Tmin
0-) )
27
+ +((oo)
40
Fi
+ Ao
Oi
lim E
min Ao
++
o min
mmin
0
2
+ o(Ao)
20
o²
oi + Ao
ai
lim
E
min Ao
{
min
Ao
0
Tmin
))
0-
27
ol o(Ao).
0i
σi
σi
14

---

Published as a conference paper at ICLR 2024
Suppose 0 * 8, we have
lim L (0,0-
lim E
Omin Ao
+ 1- min (1- _o
σi ))
11-
min
0- )
27
o (A0)
o?
0i
03
lim E (1
min 10 (0-07) -0) o(Ao)
20-0
0i
-E (1
mmin
(0 - 0 )²)
which proves our first statement in the proposition.
Now, let's consider VoL* (0,0-). It has the following form
VoL (0,07) = 2E
min
< + (1-
m
Omin
1
Omin 0- )
mmin
0i+1
0i+1
0i+1
As N + 00 and Ao we have
lim VoL (0,07
N
mmin Ao
min
Ao
2 θ- [(1
min
lim A0
2E
+
o
lim
2E
E
Jnn Ao
An
amn
8
of
21E (1
Auia
2(0 0-)
0-
x/
Amind (0- - [)] (1 Uan
lim
2
(13)
2E
min a,
)²(8 0
#
Now becomes obvious from Eq. (13) that when 0
0, we have
lim 1
VoL
VoL (0,0-)
lim 2r
2
min (0
{)
min
N
Ao
o²
σ,
2
o min
(0-
{)]
m
o
d
E
min
1
m
(0 --
de
Moreover, we can deduce from Eq. (13) that
lim 1
+00,
8 0
VoL
N Ao
-00,
0 < 0
which concludes the proof.
B ADDITIONAL EXPERIMENTAL DETAILS AND RESULTS
Model architecture Unless otherwise noted, we use the NCSN+ architecture (Song al., 2021) on CIFAR-10, and the ADM architecture (Dhariwal & Nichol, 2021) on ImageNet 64 x 64. For iCT-deep models in Tables 2 and 3, we double the depth of base architectures by increasing the number of residual blocks per resolution from 4 and 3 to 8 and 6 for CIFAR- and ImageNet 64 64 respectively. We use a dropout rate 0.3 for all consistency models on CIFAR-1O. For ImageNet 64 64, we use dropout rate of 0.2, but only apply them to convolutional layers whose the feature map resolution is smaller or equal to 16 x 16, following the configuration in Hoogeboom et al. (2023). We also found that AdaGN introduced in Dhariwal & Nichol (2021) 1) hurts consistency training and opt to remove it for our ImageNet 64 64 experiments. All models on CIFAR-10 are unconditional, and all models ImageNet 64 64 conditioned on class labels.
15

---

Published as a conference paper at ICLR 2024
0.200
400
Baseline notse schedute
Lognor mal noisa schodule
0.150
0:125
2 0.000
0.055
10
0.025
Baselie
Lognormal schedue
0.000
10n2
10t1
10"
13
10
10
20
40
loge
Training staps (x100)
(a) PDF log
(b) Lognormal vs, default schedules.
Figure 4: The PDF of log a indicates that the default noise schedule in Song et al. (2023) assigns more weight to larger values of log o, corrected by our lognormal schedule. We compare the FID scores of CT using both the lognormal noise schedule and the original one, where both models incorporate the improved techniques in Sections 3.1 to 3.4.
6.0000
0.0.003
0.30
c.0.03
0-0.01
norm 3f Adam uptates
3750
0.25
c.1.1
3500
3250
6
0.30
750
0.05
500
-0.4
-0.2
0.2
0.4
500
1000
1500
2000
2500
3000
Training
(a) d(0,, æ) as a function of æ.
(b) Comparing Adam updates.
Figure 5: (a) The shapes of various metric functions. (b) The le norms parameter updates Adam optimizer. Curves are rescaled to have the same mean. The Pseudo-Huber metric has lower variance compared to the squared & metric.
Training We train all models with the RAdam optimizer (Liu al., 2019) using learning rate 0,0001. All CIFAR-I0 models are trained for 400,000 iterations, whereas ImageNet 64 x 64 models are trained for 800,000 iterations. For CIFAR-10 models Section 3, we use batch size 512 and EMA decay rate 0.9999 for the student network. For iCT and iCTT-depp models in Table 2, we use batch size 1024 and EMA decay rate of 0.99993 for CIFAR-10 100 models, and batch size 4096 and EMA decay rate 0.99997 for ImageNet 64 64 models. All models are trained on cluster of Nvidia A100 GPUS.
Eseudo-Huber losses and variance reduction In Fig. 5, we provide additional analysis for the Pseudo-Huber metric proposed in Section 3.3. We show the shapes of squared fg metric, as well Pseudo-Huber losses with various values of in Fig. 5a, illustrating that Pseudo-Huber losses smoothly interpolates between the & and squared & metrics. In Fig. 5b, we plot the & norms of parameter updates retrieved from the Adam optimizer for models trained with squared Eg and Pseudo-Hube metrics. We observe that the Pseudo Hubuer metric has lower variance compared to the squared &2 metric, which consistent with our hypothesis Section 3.3.
Samples We provide additional uncurated samples from iCT and iCT-deep models on both CIFAR- 10 and ImageNet 64 x 64. See Figs. 6 to 9. For two-step sampling, the intermediate noise level Ø,, is 0.821 for CIFAR- 10 and 1.526 for ImageNet 64 64 when using iCT. When employing iCT-deep. Ti is 0.661 for CIFAR-10 10 and 0.973 for ImageNet 64 x 64.
16

---

Published as a conference paper at ICLR 2024
(a) One-step samples from the iCT model on CIFAR-10 (FID= 2.83).
(b) Two-step samples from the iCT model CIFAR- 10 (FID 2.46).
Figure 6: Uncurated samples from iCT models on CIFAR- 10. All corresponding samples use the same initial noise.
17

---

Published as a conference paper at ICLR 2024
(a) One-step samples from the iCT- model on CIFAR-10 (FID 2.51).
(b) Two-step samples from the iCT-ddep model CIFAR-10 (FID 2.24).
Figure 7: Uncurated samples from iCT-deep models on CIFAR-10. All corresponding samples use the same initial noise.
18

---

Published as a conference paper at ICLR 2024
One-step step samples from the iCT model ImageNet 64 64 (FID= 4.02).
Two-step samples from the iCT model ImageNet 64 64 (FID 3.20)..
Figure 8: Uncurated samples from iCT models on ImageNet 64 x 64. All corresponding samples use the same initial noise.
19

---

Published as a conference paper at ICLR 2024
One-step samples from the iCT-deep model ImageNet 64 64 (FID 3.25).
(b) Two-step samples from the iCT-deep model ImageNet 64 64 (FID 2.77).
Figure 9: Uncurated samples from iCT-deep models on ImageNet 64 x 64. All corresponding samples use the same initial noise.
20


Reviews:
[{"review_id": "H0DwvgkMhg", "reviewer": "Reviewer_ceYg", "summary": "This paper presents improved techniques for training consistency models, which include 1) adopting better weighting function, noise embeddings, and dropout 2) eliminating Exponential Moving Average from the teacher model; 3) utilizing Pseudo-Huber loss instead of $l_2$ metric and learned metrics like LPIPS; 4) introducing a new noise schedule and proposing a new curriculum for discretization steps. These modifications enable much better consistency model training.", "questions": "1. Do the authors have any insight as to why the Pseudo-Huber loss could bring this much improvement? I do not recall it utilized in any popular diffusion models. Could it be that consistency models, because of their unique features and training dynamics, benefit more from \"bounded\" gradients for all $t$?\n2. Do the authors have insights on why insensitive time embeddings are helpful or in this case crucial? There are similar design choices made in EDM: $1/4\\log(\\sigma)$ has a rather small range between $\\sigma_{\\textrm{min}}$ and $\\sigma_{\\textrm{max}}$.\n3. When examining the effects of different time embedding choices, positional embedding is mentioned. Given Fig.1 on the average $l_1$ distance, I assume the authors directly feed in $\\sigma$, rather than EDM's $1/4\\log(\\sigma)$? I think this needs to be clarified.\n4. In the paper, there are no explicit definitions of $s_0$ and $s_1$, which denote the start and end discretization steps.\n5. In Sec.3.2, third paragraph, the ground truth consistency function should have $\\xi$ rather than $\\mu$.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "This paper is really well-written and easy to follow. Each component of the design space is carefully explained and well-presented. Most of the choices are accompanied by intuitive demonstrations, which provide insights that could translate to other methods. Most importantly, the empirical results are competitive against the SOTA methods, establishing consistency models as an attractive family of generative models.", "weaknesses": "If I remember correctly, in the original consistency model, the teacher set to have EMA is more of an empirical decision. Using the same network as the student (with STOPGRAD operation) is definitely reasonable and even more intuitive (we want the model to be consistent with itself). Thus, I am not against dropping the EMA component, if it means better empirical performances.\n\nHowever, I do not feel the theoretical analysis presented in Sec.3.2 is justified. Note since I believe the main contribution of the paper is an empirical one, my complaint here does not change my assessment of the paper significantly. Still, I would appreciate it if the authors could clarify my concerns here.\n\n1. I think Eq.6 applies no matter the relation between $\\theta$ and $\\theta^-$, no?\n2. In Eq.7, why are we looking at the gradient scaled by $1/\\Delta\\sigma$? If we look at just the pure gradient, it should be 0 if $\\theta=\\theta^-$, and some finite value associated with their difference if not, and furthermore, the difference between $\\theta$ and $\\xi$ disappears. It may look like the loss function has nothing to do with learning the correct parameter $\\xi$ like the authors suggest in the second last paragraph in Sec.3.2. However, to me, this is not a surprise and totally expected. The self-consistency loss itself does not really force the network to learn anything correctly, i.e. the model could just predict a constant no matter the input and still be considered consistent with itself. In my understanding, it is really the boundary condition enforced through parameterizations that makes the model work, which is not a fact used in this toy example.\n3. In the current version, I am curious as to what the authors have in mind about \"if the consistency loss either does not exist or is unsuitable\". This notion does not seem to be explained well.", "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": "1BjwXopXgp", "reviewer": "Reviewer_YvKM", "summary": "\"The paper introduces a comprehensive set of techniques aimed at enhancing the efficacy of consistency training within few-shot diffusion models for image generation. The contributions are prioritized as follows:\n\n1. The elimination of the Exponential Moving Average (EMA) from the student model, a change supported by both theoretical arguments and empirical evidence, results in markedly improved consistency training performance.\n\n2. An improved framework that incorporates weighted timestep selection (inspired by EDM), time-step adaptive loss weighting, and the pseudo-Huber loss function to refine the training process.\n\n3. A refined schedule for discretization steps involved in computing the consistency loss, accompanied by a new heuristics curriculum that adjusts these steps throughout the training phases.\n\n4. Additional modifications to the network, such as the integration of dropout strategies and better time embeddings, to enhance model robustness and adaptability.\n\nCollectively, these advancements makes the consistency training model highly competitive, achieving performance on par with leading generative models on benchmarks like CIFAR-10 and ImageNet.", "questions": "I have several questions for the authors that could enrich the broader community’s understanding:\n\nQ1: Does the current consistency training outperform consistency distillation when applying the enhancements presented in this paper to the latter?\n\nQ2: Given the apparent benefits of scaling up the network architecture for larger datasets (e.g., ImageNet as shown in Table 3), is there a risk that one(few)-step methods will be inherently limited by network capacity? In other words, can these methods directly approximate the ODE solution with a network of the same capacity as the original diffusion model, or would alternative approaches like RectFlow [1] provide a more viable solution?\n\n[1] Liu, Xingchao, Chengyue Gong, and Qiang Liu. \"Flow straight and fast: Learning to generate and transfer data with rectified flow.\" arXiv preprint arXiv:2209.03003 (2022).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1: this model positions consistency models as competitive generative models. When the first consistency model was introduced, its performance was somewhat mediocre, especially on large datasets like ImageNet. The set of improvements proposed in this paper has made consistency training competitive, and it may become a leading approach in the future due to desirable properties such as stable training and accurate approximation of the probability flow ODE.\n\nS2: the technical contribution is robust. The theoretical analysis supporting the removal of the EMA is well-justified and is corroborated by experimental results.\n\nS3: Various other enhancements have also proven useful and are validated experimentally.\n\nS4: the paper is well-written and easy to follow", "weaknesses": "W1: The choice to use the pseudo-Huber loss is not entirely convincing, as its effectiveness appears sensitive to the chosen constant value, and it offers negligible improvements over the tuning-free LPIPS loss, which is known for its alignment with human perception and computational efficiency, especially in pixel-space models.\n\nW2: Apart from the significant change of removing EMA, the other enhancements resemble engineering optimizations rather than foundational advances. Their relevance to broader applications is questionable, especially since the empirical evaluation is limited to smaller datasets like CIFAR-10 and ImageNet-64x64. Expanding experiments to include higher-resolution images on ImageNet or more varied datasets such as COCO or LAION could substantiate the model's versatility.\n\nW3: The training speed for the consistency model is comparatively slow, especially when compared with models like GANs. It would be insightful if the authors could address this aspect and discuss its implications for scaling to larger datasets. \n\nW4: The applicability of consistency training appears confined to training unguided models, which significantly underperform in applications such as text-to-image generation.", "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": "fCOpnfJTwV", "reviewer": "Reviewer_3jcG", "summary": "In this work, the authors enhanced the training of consistency model, initially introduced by Song et al. [1], by implementing improved weighting functions, a refined discretization curriculum, eliminating Exponential Moving Average (EMA), introducing a new loss function, and improving the noise schedule. These advancements enable the model to achieve state-of-the-art results in both one and two-step generation processes without relying on pre-trained diffusion models or learned metrics. Overall, the paper is well-written and organized, and the improving techniques are supported by convincing discussions.\n\n\n\n\n\n\n[1] Song, Y., Dhariwal, P., Chen, M. and Sutskever, I., 2023. Consistency models.", "questions": "Could the author elaborate more on why higher discretization steps can reduce bias but increase variance?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper presents compelling discussions on improved training schemes, which notably facilitate the development of a consistency model. This model impressively achieves state-of-the-art results in both one and two-step generation processes, uniquely independent of pre-trained diffusion models or learned metrics.\n\n2. The empirical results on CIFAR-10 and ImageNet 64×64 datasets affirm the efficacy of these enhanced training methodologies.\n\n3. The research sets a new benchmark, providing a robust and successful framework for training consistency models. This paradigm holds potential for broader application in various one or two-step generalization models.", "weaknesses": "1. This paper functions primarily as a technical exploration, compiling successful practices in training diffusion models. While it includes numerous ablation studies, it lacks sufficient theoretical backing to fully support these practices.\n\n2. The paper’s improved training schemes do allow the consistency model to avoid relying on a pre-trained diffusion model. However, its theoretical basis still seems anchored in diffusion model principles. It would be beneficial if the authors explored the broader potential of the consistency training scheme. Specifically, whether this training would be effective when the generation process is not the reverse of a diffusion process but a more general corruption process, such as those described in [1] and [2].\n\n[1] Bansal, A., Borgnia, E., Chu, H.M., Li, J.S., Kazemi, H., Huang, F., Goldblum, M., Geiping, J. and Goldstein, T., 2022. Cold diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.09392.\n[2] Xu, Y., Liu, Z., Tegmark, M. and Jaakkola, T., 2022. Poisson flow generative models. Advances in Neural Information Processing Systems, 35, pp.16782-16795.", "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": "nn6FvPo3PU", "reviewer": "Reviewer_m7Dp", "summary": "Recently, consistency models were proposed. They can be used to either distill diffusion models into single- or few-step samplers (consistency distillation), or they also enable direct training of single- or few-step consistency models (consistency training). Previous work showed strong results primarily for the distillation setting. In this work the authors focus on the consistency training setting and propose multiple techniques to improve the performance of consistency training. For instance, the paper avoids the questionable use of LPIPS-based consistency matching, and employs a more general Huber loss. The authors also remove the exponential moving average and point out errors in previous theoretical analyses. Moreover, a new loss weighting function is proposed, the noise level embeddings are carefully analyzed and improved, the use of dropout is studied, and the discretization step curriculum is improved. The paper thoroughly ablates and analyzes all improvements and then evaluates the model on popular benchmarks, CIFAR10 and ImageNet64. It achieves very strong performance, significantly outperforming both previous consistency models and diffusion model distillation methods. The improved consistency models perform almost as good as state-of-the-art diffusion models and GANs, while being trained directly and only requiring a single or two synthesis steps.", "questions": "I only have some minor questions:\n- When the paper defines the ground truth consistency function in Section 3.2, is there a typo in the ground truth consistency function? Should it be $\\xi$ instead of $\\mu$ in the ground truth consistency function $f^{*}$ for the single data point data distribution?\n- What are $s_0$ and $s_1$ in Section 3.4? They are not properly explained when first introduced above Eq. (9).\n- Also in Section 3.4, the authors use an exponential discretization curriculum, which leads to the visualization in Fig. 3(a). Why is this exponential curriculum not included in the visualizations and ablations in Figs. 3(b) and 3(c)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper has several strengths:\n- The paper is well-written and easy to follow (good **clarity**).\n- The paper does not present a fundamentally new method (moderate **novelty**). However, it does significantly improve consistency models, a very novel and very promising class of generative models that enables single- or few-step sampling, in contrast to, for instance, diffusion models. I expect the insights provided by the paper to be used in follow-up work and I think that consistency models will find wide usage, too. This makes this work very **significant**\n- The paper presents thorough ablations and analyses of the proposed tricks and innovations. It is generally of high **quality**.\n- The final experimental results obtained when combining all modifications are very strong, and thoroughly compared to many and appropriate baselines.", "weaknesses": "I think the paper has no major weaknesses. However, there would be opportunities to further improve the work:\n- I am wondering how scalable consistency models are with the proposed modifications. Can one train, for instance, text-to-image consistency models? Or how about training on higher-resolution images?\n- The new consistency models do not require LPIPS losses anymore and are thereby more general in that they are not limited to image synthesis anymore. It would be interesting to validate that consistency models can also be successfully trained on non-image data (e.g. audio, graphs, 3D, video, etc).\n\nConclusion: Overall, I think this is a good and solid paper. It significantly improves consistency models, a very promising new class of generative models, and the paper is overall well-written and of high quality. Hence, I recommend acceptance.", "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)