Progressive Distillation for Fast Sampling of Diffusion Models
Authors: Tim Salimans
Jonathan Ho
Venue: ICLR 2022 Spotlight (2022)
Keywords: Diffusion Models
Generative Models
fast sampling

Abstract: Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation. A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations. Here we make two contributions to help eliminate this downside: First, we present new parameterizations of diffusion models that provide increased stability when using few sampling steps, compared to models in the literature. Second, we present a method to distill a trained deterministic diffusion sampler, using many steps, into a new diffusion model that takes half as many sampling steps. We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time. On standard image generation benchmarks like CIFAR-10, ImageNet, and LSUN, we start out with (near) state-of-the-art samplers taking 1024 or 8192 steps, and are able to distill down to models taking as little as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps. Finally, we show that the full progressive distillation procedure does not take more time than it takes to train the original model, thus representing an efficient solution for generative modeling using diffusion at both train and test time.

Published as a conference paper at ICLR 2022
PROGRESSIVE DISTILLATION FOR FAST SAMPLING OF DIFFUSION MODELS
Tim Salimans & Jonathan Ho Google Research, Brain team
(salimans, jona mmathanho] wow.nnhe.@oogg....oooooo com
ABSTRACT
Diffusion models have recently shown great promise for generative modeling, out- performing GANs on perceptual quality and autoregressive models at density es- timation. A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations. Here we make two contributions to help eliminate this downside: First, we present new parameterizations of diffusion models that provide increased stability when using few sampling steps. Second, we present a method to distill a trained deterministic diffusion sampler, using many steps, into many sampling steps. We then keep progressively applying this distillation proce- dure to our model, halving the number of required sampling steps each time. On standard image generation benchmarks like CIFAR-10, 10, ImageNet, and LSUN, we start out with state-of-the-ee samplers taking as many as 8192 steps, and are able to distill down to models taking as few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 10 in 4 steps. Finally, we show that the full progressive distillation procedure does not take more time than takes to train the original model, thus representing an efficient solution for generative modeling using diffusion at both train and test time.
diffusion model that takes half
INTRODUCTION
Diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) are an emerg- ing class of generative models that has recently delivered impressive results on many standard gen- erative modeling benchmarks. These models have achieved ImageNet generation results outper- forming BigGAN-deep and VQ-VAE-2 in terms of FID score and classification accuracy score (Ho et al., 2021; Dhariwal & Nichol, 2021), and they have achieved likelihoods outperforming autore- gressive image models (Kingma et al., 2021; Song et al., 2021b). They have also succeeded in image super-resolution (Saharia et al., 2021; Li et al., 2021) and image inpainting (Song al., 2021c), and there have been promising results in shape generation (Cai et al., 2020), graph generation (Niu et al., 2020), and text generation (Hoogeboom et al., 2021; Austin al., 2021).
major barrier remains to practical adoption of diffusion models: sampling speed. While sam- pling can be accomplished in relatively few steps in strongly conditioned settings, such text-to- speech (Chen et al., 2021) and image super-resolution (Saharia et al., 2021), or when guiding the sampler using an auxiliary classifier (Dhariwal & Nichol, 2021), the situation substantially differ- ent in settings in which there is less conditioning information available. Examples of such settings are unconditional and standard class-conditional image generation, which currently require hundreds or thousands of steps using network evaluations that are not amenable to the caching optimizations of other types of generative models (Ramachandran et al., 2017).
In this paper, we reduce the sampling time of diffusion models by orders of magnitude in uncondi- tional and class-conditional image generation, which represent the setting in which diffusion models have been slowest in previous work. We present procedure to distill the behavior of a N Nsstep DDIM sampler (Song et al., 2021a) for pretrained diffusion model into a new model with N/2 steps, with little degradation in sample quality. In what we call progressive distillation, repeat this distilla- tion procedure to produce models that generate in as few as 4 steps, still maintaining sample quality competitive with state-of-the-art models using thousands of steps.

---

Published as a conference paper at ICLR 2022
==
e
e
e
23/4 = f(zi; n)
Distillation
2/2 0//2= ((33/45
Distillation
X = f(z1;
114 = f(21/231)
Distillation
x = f(21/15n)
0=
Figure 1: A visualization of two iterations of our proposed progressive distillation algorithm. A sampler f(z; n). mapping random noise to samples x in deterministic steps, distilled into new sampler f(z; 0) taking only a single step. The original sampler is derived by approximately integrating the probability flow ODE for learned diffusion model, and distillation can thus be understood as learning to integrate in fewer steps, amortizing this integration into the new sampler.
2 BACKGROUND ON DIFFUSION MODELS
We consider diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) specified in continuous time (Tzen & Raginsky, 2019a; Song al., 2021c; Chen et et., 2021; Kingma et al., 2021). We use x ~ p(x) to denote training data. A diffusion model has latent variables z = (z: [0, 1]} and is specified by a noise schedule comprising differentiable functions a, at decreases monotonically with t.
such that A = logl a /ot). the log signal-to-noise-ratie
These ingredients define the forward process q(z) x), a Gaussian process satisfying the following Markovian structure:
((2,1x) N(z;;,,,,0 g(z,22,) = (ar/aa og.
where 0 S <
$ 1 and o (1-e21-2-)0?
The role of function approximation in the diffusion model to denoise z, ~ q(z, (x) into an estimate & (z.) (the function approximator also receives Å an input, but omit this to keep our notation clean). We train this denoising model Xo using weighted mean squared error loss
Be, [m<(Ac]] & (z.) x
over uniformly sampled times & E [0. 1]. This loss can be justified as a weighted variational lower bound on the data log likelihood under the diffusion model (Kingma et al., 2021) or form of denoising score matching (Vincent, 201 Song & Ermon, 2019). We will discuss particular choices of weighting function w(,) later
Sampling from a trained model can be performed in several ways. The most straightforward way is discrete time ancestral sampling (Ho et al., 2020). To define this sampler, first note that the forward process can be described in reverse as (22 a Zz x) = (2,5 Poje x), a|c (noting < t). where
(Ze, x) ~ (a./a.)
02. (1 d e )o
(3)
We use this reversed description of the forward process to define the ancestral sampler. Starting at 2i NO, I), the ancestral sampler follows the rule
z, = P.,e &o (Z+)) + / (6
(4)
= (c-s>((as/a))) +(1 e )a->)))a + (0?,, (2
(5)
2

---

Published as a conference paper at ICLR 2022
where is standard Gaussian noise, and is a hyperparameter that controls how much noise is added during sampling, following Nichol & Dhariwal (2021).
Alternatively, Song al. (2021c) show that our denoising model X xe(z,) can be used to determinis-
tically map noise 21 ~ (0, 1) to samples x by numerically solving the probability flow ODE:
(t)V: dz, [f(z,,t) (tV. log lo
(6)
where V, log (zi) = a (1)0(2,))
Following Kingma et al. (2021), we have f(z,,t) d'log Zt dt
of
do? 2 22 at of. Since 8o(z+) is parameterized by neural network, this equation and (t) di
dt
is special case of neural ODE (Chen et al., 2018), also called continuous normalizing flow (Grathwohl et al., 2018).
Solving the ODE in Equation 6 numerically can be done with standard methods like the Euler rule or the Runge-Kutta method. The DDIM sampler proposed by Song et al. (2021a) can also be understood as an integration rule for this ODE, as show in Appendix B, even though it was originally proposed with a different motivation. The update rule specified by DDIM is
2, = (zz))(Zz) - o. 2, (x)
Tt
= (((ii-a)/((ssaaa.)) z + (1 - (Ac-)))11 ja, (x()).
(8)
and in practice this rule performs better than the aforementioned standard ODE integration rules in our case, as we show in Appendix C.
If (z,) satisfies mild smoothness conditions, the error introduced by numerical integration of the probability flow ODE is guaranteed to vanish as the number of integration steps grows infinitely large, i.e. N - co. This leads to a trade-off in practice between the accuracy of the numerical integration, and hence the quality of the produced samples from our model, and the time needed to produce these samples. So far, most models in the literature have needed hundreds or thousands of integration steps to produce their highest quality samples, which is prohibitive for many practical applications of generative modeling. Here, we therefore propose a method to distill these accurate, but slow, ODE integrators into much faster models that are still very accurate. This idea is visualized in Figure 1, and described in detail in the next section.
3 PROGRESSIVE DISTILLATION
To make diffusion models more efficient sampling time, we propose progressive distillation: an algorithm that iteratively halves the number of required sampling steps by distilling a slow teacher diffusion model into a faster student model. Our implementation of progressive distillation stays very close to the implementation for training the original diffusion model, as described by e.g. Ho al. (2020). Algorithm 1 and Algorithm 2 present diffusion model training and progressive distillation side-by-side, with the relative changes in progressive distillation highlighted in green.
We start the progressive distillation procedure with a teacher diffusion model that is obtained by training in the standard way. At every iteration of progressive distillation, we then initialize the student model with a copy of the teacher, using both the same parameters and same model definition. Like in standard training, we then sample data from the training set and add noise to it, before forming the training loss by applying the student denoising model to this noisy data z,, The main difference in progressive distillation is in how we set the target for the denoising model: instead of the original data x, we have the student model denoise towards a target x that makes a single student DDIM step match 2 teacher DDIM steps. We calculate this target value by running 2 DDIM sampling steps using the teacher, starting from Z, and ending at Z, then calculate the value the student model would need to predict in order to move from 24 to 2,-1/N in a single step, as we show in detail in Appendix G. The resulting target value x(z;) is fully determined given the teacher model and starting point 2₁, which allows the student model to make a sharp prediction when evaluated at z₁ In contrast, the original data point x is not fully determined given 2,, since multiple different data points can produce the same noisy data 2; this means that the original denoising model is
student sampling steps. By inverting a single step of DDIM, 1,-1/N- with being the number of
3

---

Published as a conference paper at ICLR 2022
predicting weighted average of possible values, which produces blurry prediction. By making sharper predictions, preepittooiiiiinnn the student model can make faster progress during sampling.
After rumning distillation to leam student model taking N sampling steps, repeat the pro- cedure with N/2 steps: The student model then becomes the new teacher, and a new student model is initialized by making copy of this model.
Unlike our procedure for training the original model, we always run progressive distillation in dis- crete time: we sample this discrete time such that the highest time index corresoonds to signal-to- noise ratio of zero, i.e. on = 0, which exactly matches the distribution of input noise 2 ~ NN, I) that is used at test time. We found this work slightly better than starting from non-zero signal- to-noise ratio as used by e.g. Ho et al. (2020), both for training the original model as well as when performing progressive distillation.
Algorithm Standard diffusion training
Algorithm Progressive distillation
Require: Model (g(2;) to be trained
Require: Trained teacher model (2,)
Require: Data set D
Require: Data set 'D
Require: Loss weight function w(()
Require: Loss weight function w()
Require: Student sampling steps N
for K iterations do
7
b Init student from teacher
while not converged do
while not converged do
X ~ D
Sample data
D
t t~ U10,
b. Sample time
: = i/N, in Cat L....... N
<~
N(0, I)
>> Sample noise
E N(0,1)
Z: axx + oee Add noise data
2: 0iX OGE
steps of DDIN with teacher
E' = t- 0.5/N, t" = 1/N
2, = ar xn (2:) + dr H (Zr ar Xxy (Ze)
x,
(Zt a2 (z)
Zr" ar" Xn (2(2(r)
X > Clean data is target for & b log-SNR
X 2,"' (0.1//o.22
0,6 (ave lauer
Teacher & target
A = log[a? /o?
A = log [a [o?]
La w/ {1 x (o(2))
D Loss
Lg w/ ((( || X (ze))
0 + 0 VoLo
> Optimization
0 - 0 VaLo
end while
end while
4- 0
Student becomes next teacher
N N/2
Halve number of sampling steps
end for
4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS
In this section, discuss how to parameterize the denoising model Xg. and how to specify the reconstruction loss weight w()e). We assume a standard variance- preserving diffusion process for which of 1 - af. This is without loss of generalization, appendix G): different specifications of the diffusion process, such as the variance-exploding spec- ification, can be considered equivalent to this specification, up rescaling of the noisy latents 2,- We use a cosine schedule a = co.5t.)) similar to that introduced by Nichol & Dhariwal (2021).
shown by (Kingma et al., 2021,
Ho et al. (2020) and much of the following work choose to parameterize the denoising model through directly predicting with a neural network éØ (z.), which implicitly sets {g(2;) _ [Z (Zz-0z60(z.)
In this case, the training loss is also usually defined as mean squared error in the c-space:
1 - (z. axx)
1 - (z)))
Lg ||e En (z)))))
(a()))
(9)
which can thus equivalently be seen as a weighted reconstruction loss x space, where the weight- ing function is given by w²(A.) exp():), for log signal-to-noise ratio X, loglo? / of))-
4

---

Published as a conference paper at ICLR 2022
Although this standard specification works well for training the original model, it is well suited for distillation: when training the original diffusion model, and at the start of progressive distillation, the model is evaluated at wide range of signal-to-noise, ratios a /of. but as distillation progresses we increasingly evaluate at lower and lower signal-to-noise ratios. As the signal-to-noise ratio goes to zero, the effect of small changes in the neural network output éo (2z) on the implied prediction in x-space is increasingly amplified, since (z,) = de (Z: (Zz)(Z))) divides by at + O. This is not much of a problem when taking many steps, since the effect of early missteps is limited by clipping of the 2; iterates, and later updates can correct any mistakes, but becomes increasingly important as we decrease the number of sampling steps. Eventually, if we distill all the way down to a single sggaalttonoiiss ratio of zero, i.c. a, = 0, = 1. At this extreme, the link between e-prediction and x- x-prediction breaks down completely: observed data 2₁ is no longer informative of x and predictions é (Ze) no longer implicitly predict x. Examining our reconstruction loss (equation 9), we see that the weighting function w (>t)) gives zero weight to the reconstruction loss at this signal-to-noise ratio.
sampling step, the input to the model only pure noise E, which corresponds
For distillation to work, we thus need to parameterize the diffusion model in way for which the implied prediction xz(2;) remains stable as A = Log[a//of varies. We tried the following options, and found all work well with progressive distillation:
Predicting x directly.
Predicting both x and e, via separate output channels (Nx(2z), (Z:) of the neural net- work, and then merging the predictions via X = of (zt) + az (z: ozE6(2))). thus smoothly interpolating between predicting directly and predicting via e.
Predicting v = CzCC o,x, which gives = a: o V.vo(Zz), as we show in Appendix D.
In Section 5.1 we test all three parameterizations on training an original diffusion model (no distil- lation), and find them to work well there also.
In addition determining appropriate parameterization, tion loss weighting wAr)). The setup of Ho et al. (2020) weights the reconstruction loss by the signal-to-noise ratio, implicitly gives weight of zero to data with zero SNR, and is therefore not suitable choice for distillation. We consider two alternative training loss weightings:
also need to decide on reconstruc-
Lg max x R 11 ||e = max c 1)[[x & 'truncated SNR' weighting.
. Lg = IV, 1.113 = (1 - ||x 11: 'SNR+1' weighting.
We examine both choices in our ablation study in Section 5.1, and find both to be good choices for training diffusion models. In practice, the choice of loss weighting also has to take into account how dr, oz are sampled during training, this sampling distribution strongly determines the weight the expected loss gives to each signal-to-noise ratio. Our results are for a cosine schedule og = c0.5tt)) where time is sampled uniformly from [0, 1]. In Figure 2 we visualize the resulting loss weightings, both including and excluding the effect of the cosine schedule.
5 EXPERIMENTS
In this section we empirically validate the progressive distillation algorithm proposed in Section 3, as well as the parameterizations and loss weightings considered in Section 4. We consider various image generation benchmarks, with resolution varying from 32 32 128 128. All experiments use the cosine schedule a; cos(0.5mr)) and all models use a U-Net architecture similar to that introduced by Ho et al. (2020), but with BigGAN-style up- and downsampling (Brock et al., 2019), as used in the diffusion modeling setting by Nichol & Dhariwal (2021): Song al. (2021c). Our training setup closely matches the open source code by Ho et al. (2020). Exact details are given in Appendix E.
5.1 MODEL PARAMETERIZATION AND TRAINING LOSS
As explained in Section 4, the standard method of having our model predict e, and minimizing mean squared error in the e-space (Ho et al., 2020), is not appropriate for use with progressive distillation.
5

---

Published as a conference paper at ICLR 2022
SNR weight
SNR
log weight (excluding schedule)
truncated SNR
0.4
truncated SNR
weggn (inc Inclluding schedule)
SNR+1 weight
SNR+1
0.3
0.2
0.1
-6
-4
2
4
6
6
-4
2
0
4
log SNR
log SNR
Figure 2: Left: Log weight assigned reconstruction loss |
function of the log-SNR = logla² ]o]] for each of our considered training loss weightings, excluding the influence of the at, ae schedule. Right: Weights assigned to the reconstruction loss including the effect of the cosine schedule a, = co.Tt.)) with ~ U[0,, 1]. The weights are only defined up to a constant, and we have adjusted these constants to fit this graph.
8x112 -
Network Output
Loss Weighting Stochastic sampler
DDIM sampler
(x, c) combined
SNR Truncated SNR
2.54/9.88
2.78/9.56
2.47/9.85 85
2.76/9.49 .49
SNR+1
2.52/9.79
2.87/9.45
SNR
2.65/9.80
2.75/9.56
Truncated SNR SNR+1
2.53/9.92
2.51/9.58
2.56/9.84
2.65/9.52
SNR
2.59/9.84 884
2.91/9.52
Truncated SNR SNR+1
N/A
N/A
2.56/9.77
3.27/9.41
SNR
2.65/9.86
3.05/9.56
Truncated SNR
2.45/9.80
2.75/9.52
SNR+1
2.49/9.77
2.87/9.43
Table 1: Generated sample quality as measured by FID and Inception Score (FID/IS) on uncondi- tional CIFAR-10 10, training the original model (no distillation), and comparing different parameteri- zations and loss weightings discussed in Section 4. All reported results are averages over 3 random seeds of the best metrics obtained over 2 million training steps; nevertheless we find results are still +0.11 due to the noise inherent in training our models. Taking the neural network output to represent a prediction of in combination with the Truncated SNR loss weighting leads divergence.
We therefore proposed various alternative parameterizations of the denoising diffusion model that are stable under the progressive distillation procedure, as well as various weighting functions for the reconstruction eтror x-space. Here, we perform complete ablation experiment of all pa- rameterizations and loss weightings considered in Section 4. For computational efficiency, and for comparisons to established methods in the literature, we use unconditional CIFAR-10 as the bench- mark. We measure performance of undistilled models trained from scratch, to avoid introducing too many factors of variation into our analysis.
Table lists the results of the ablation study. Overall results are fairly close across different pa- rameterizations and loss weights. All proposed stable model specifications achieve excellent perfor- mance, with the exception of the combination of outputting with the neural network and weighting the loss with the truncated SNR, which we find to be unstable. Both predicting x directly, as well predicting v, the combination (c, x), could thus be recommended for specification of diffusion

---

Published as a conference paper at ICLR 2022
models. Here, predicting is the most stable option, as it has the unique property of making DDIM step-sizes independent of the SNR (see Appendix D), but predicting x gives slightly better empirical results in this ablation study.
5.2 PROGRESSIVE DISTILLATION
We evaluate our proposed progressive distillation algorithm on 4 data sets: CIFAR-IO, 64 x 64 downsampled ImageNet, 128 128 LSUN bedrooms, and 128 x 128 LSUN Church-Outdoon For each data set we start by training a baseline model, after which we start the progressive distillation procedure. For CIFAR-10 10 we start progressive distillation from a teacher model taking 8192 steps. For the bigger data sets we start at 1024 steps. At every iteration of distillation we train for 50 thousand parameter updates, except for the distillation to and sampling steps, for which we use 100 thousand updates. We report FID results obtained after each iteration of the algorithm. Using these settings, the computational cost progressive distillation to 4 sampling steps is comparable or less than for training the original model. In Appendix I we show that this computational cost can be reduce much further still, at a small cost in performance.
In Figure 4 we plot the resulting FID scores (Heusel et al., 2017) obtained for each number of sampling steps. We compare against the undistilled DDIM sampler, as well as to a highly optimized stochastic baseline sampler. For all four data sets, progressive distillation produces near optimal results up to 4 or 8 sampling steps. 2t or sampling steps, the sample quality degrades relatively more quickly. In contrast, the quality of the DDIM and stochastic samplers degrades very sharply after reducing the number of sampling steps below 128. Overall, we conclude that progressive distillation is thus an attractive solution for computational budgets that allow less than equal to 128 sampling steps. Although our distillation procedure is designed for use with the DDIM sampler, the resulting distilled models can in principle also be used with stochastic sampling: we investigate this in Appendix F, and find that it achieves performance that falls in between the distilled DDIM sampler and the undistilled stochastic sampler.
Table 2 shows some of our results on CIFAR-10, and compares against other fast sampling methods in the literature: Our method compares favorably and attains higher sampling quality in fewer steps than most of the alternative methods. Figure 3 shows some random samples from our model obtained at different phases of the distillation process. Additional samples are provided in Appendix H.
(a) 256 sampling steps
(b) sampling steps
(c) sampling step
Figure 3: Random samples from our distilled 64 x 64 ImageNet models, conditioned on the 'mala- mute' class, for fixed random seed and for varying number of sampling steps. The mapping from input noise to output image is well preserved the number of sampling steps is reduced.
6 RELATED WORK ON FAST SAMPLING
Our proposed method is closest to the work of Luhman & Luhman (2021), who perform distillation of DDIM teacher models into one-step student models. A possible downside of their method is
7

---

Published as a conference paper at ICLR 2022
CIFAR-I0
64x64 ImageNet
20
20
Distilled
Distilled
DDIM
DDIM
Stochastic
Stochastic
10
10
9
FID
FID
6
4
4
2
2
2
4
8
16
32
64
128
256
512
2
4
8
16
32
64
128
256 512
sampling steps
sampling steps
128x128 LSUN Bedrooms
128x128 LSUN Church-Outdoor
20
20
Distilled
Distilled
DDIM
DDIM
Stochastic
Stochastie
10
10
9
FID
5
5
3
3
2
4
8
16
32
64
128 256 512
2
4
8
16
32
64
128 256 512
sampling steps
sampling steps
Figure 4: Sample quality results measured by FID for distilled model unconditional CIFAR-10, class-conditional 64x64 ImageNet, 128x128 LSUN bedrooms, and 128x128 LSUN church-outdoor. We compare against the DDIM sampler and against optimized stochastic sam- pler, each evaluated using the same models that were used to initialize the progressive distillation procedure. For CIFAR-10 report average only use a single run because of their computational demand. For the stochastic sampler we set the variance a log-scale interpolation between an upper and lower bound the variance, follow- ing Nichol & Dhariwal (2021), but we use single interpolation coefficient rather than a learned coefficient. We then tune this interpolation coefficient separately for each number of sampling steps and report only the best result for that number of steps: this way we obtained better results than with the learned interpolation.
4 random seeds. For the other data sets
that requires constructing a large data set by running the original model its full number of sampling steps: their cost of distillation thus scales linearly with this number of steps, which can be prohibitive. In contrast, our method never needs to run the original model the full number of sampling steps: every iteration of progressive distillation, the number of model evaluations independent of the number of teacher sampling steps, allowing our method to scale up to large numbers of teacher steps logarithmic cost in total distillation time.
DDIM (Song et al., 2021a) was originally shown to be effective for few-step sampling, as was the probability flow sampler (Song et al., 2021c). Delicoeur-Martinee et al. (2021) study fast SDE integrators for reverse diffusion processes, and Tzen Raginsky (2019b) study unbiased samplers which may be useful for fast, high quality sampling as well.
8

---

Published as a conference paper at ICLR 2022
Other work fast sampling be viewed as manual automated methods to adjust samplers diffusion processes for fast generation. Nichol & Dhariwal (2021); Kong & Ping (2021) describe methods adjust discrete time diffusion model trained many timesteps into models that can sample few timesteps. Watson et al. (2021) describe a dynamic programming algorithm to reduce the number of timesteps for diffusion model in way that is optimal for log likelihood. Chen al. (2021); Saharia et al. 12; Ho et al. (2021) train diffusion models over continuous noise levels and tune samplers post training by adjusting the noise levels of few-step discrete time reverse diffusion process. Their method is effective in highly conditioned settings such as text-to- text-to-speech and image super-resolution.
al. (2021) train a new network estimate the noise level of noisy data and show how to use this estimate speed up sampling.
A Alternative specifications of the diffusion model can also lend themselves fast sampling, such as modified forward and reverse processes (Nachmani et al., 2021; Lam et al., 2021) and training diffusion models in latent space (Vahdat et al., 2021).
Method
Model evaluations
FID
Progressive Distillation (ours)
9.12
2
4.51
4
3.00
8
2.57
Knowledge distillation (Luhman Luhman, 2021)
1
9.36
DDIM (Song et al., 2021a)
10
13.36
20
6.84
50
100
4.16
Dynamic step-size extrapolation + VP-deep Oolicoeur-Martiness et al., 2021)
82.42
151
2.73
180
2.44
274
2,60
330
2.56
FastDPM (Kong & Ping, 2021)
10
9,90
5.05
50
3,20
100
2.86
Improved DDPM respacing
25
7.53
(Nichol & Dhariwal, 2021), our reimplementation
50
4.99
LSGM (Vahdat et al., 2021)
138
2.10
Table 2: Comparison of fast sampling results CIFAR-10 for diffusion models in the literature.
DISCUSSION
We have presented progressive distillation, a method drastically reduce the number of sampling steps required for high quality generation of images, and potentially other data, using diffusion models with deterministic samplers like DDIM (Song et al., 2020). By making these models cheaper to run at test time, we hope to increase their usefulness for practical applications, for which running time and computational requirements often represent important constraints.
In the current work we limited ourselves to setups where the student model has the same architecture and number of parameters as the teacher model: future work we hope to relax this constraint and explore settings where the student model is smaller, potentially enabling further gains in test time computational requirements. In addition, we hope to move past the generation images and also explore progressive distillation of diffusion models for different data modalities such e.g. audio (Chen et al., 2021).
In addition the proposed distillation procedure, some of our progress realized through differ- ent parameterizations of the diffusion model and its training loss. We expect to see more progress in this direction the community further explores this model class.

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
We provide full details on model architectures, training procedures, and hyperparameters in Ap- pendix E, in addition our discussion in Section In Algorithm 2 provide fairly detailed pseu- docode that closely matches our actual implementation, which is available in open source at https: //gi thub. wt..tttuhooooooooe.eeoooommm.. cww.google-eeeooooooooooooooooe /eterrcceeoeeeeeeeeeerre//ee///e ffusiio stillation..
research/trreemmstteeee
ETHICS STATEMENT
In general. generative models can have unethical uses, such as fake content generation, and they can suffer from bias if applied to data sets that are not carefully curated. The focus of this paper specifically is on speeding up generative models at test time order reduce their computational demands; do not have specific concerns with regards to this contribution.
REFERENCES
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne den Berg. Structured denoising diffusion models discrete state-spaces. CoRR, abs/2107.03006 2021.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019.
Ruojin Cai, Guandao Yang, Hadar Averbuch-Elot Zekun Hao, Serge Belongie, Noah Snavely, and Bharath Hariharan. Learning gradient fields for shape generation. arXiv preprint arXiv:2008.0653 2020.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan. Wave- Grad: Estimating gradients for waveform generation. International Conference on Learning Rep- resentations, 2021.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differ- ential equations. In Advances in Neural Information Processing Systems, pp. 6571-6583, 2018.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. arXiv preprint arXiv:2/05.055333 2021.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367, 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to local Nash equilibrium. In Advances in Neural Information Processing Systems, pp. 6626 6637, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, pp. 6840-6851, 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. arXiv preprint 2tt.:2//0.112822. 2021.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Foréé, and Max Welling. Argmax flows and multinomial diffusion: Towards non-autoregressive language models. arXiv preprint arXiv:2/02.053399 2021.
Alexia
Ke Li, Rémi Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2105. 10.22005.14000 14080, 2021.
Tal Kachman, and Ioannis Mitliagkas.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint arXiv:2 107.00630, 2021.
10

---

Published as a conference paper at ICLR 2022
Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabilistic models. arXiv preprint arXiv:2106.00133. 2021.
Max WY Lam, Jun Wang, Rongjie Huang, Dan Su, and Dong Yu. Bilateral denoising diffusion models, arXiv preprint arXiv:2108. /154,, 2021.
Haoying Li, Yifan Yang, Meng Chang, Huajun Feng, Zhihai Xu, Qi Li, and Yueting Chen. Srdiff: Single image super-resolution with diffusion probabilistic models. arXiv preprint 1rXh::2104.14951, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 2017.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint Stti::2001.023880 2021.
Eliya Nachmani, Robin San Roman, and Lior Wolf. Non gaussian denoising diffusion models. arXiv preprint arXiv:2106.07588 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. ICML, 2021.
Chenhao NIu, Yang Song, Jiaming Song, Shengjia Zhao, Aditya Grover, and Stefano Ermon. Per- mutation invariant graph generation via score-based generative modeling. In International Con- ference on Artificial Intelligence and Statistics, pp. 4474 4484. PMLR, 2020.
Prajit Ramachandran, Tom Le Paine, Pooya Khorrami, Mohammad Babaeizadeh, Shiyu Chang, Yang Zhang, Mark A Hasegawa-Johnson, Roy H Campbell, and Thomas S Huang. Fast genera- tion for convolutional autoregressive models. arXiv preprint FrXivv1704.06000 2017.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. arXiv preprint arXiv:2 2021.
Robin San-Roman, Eliya Nachmani, and Lior Wolf. Noise estimation for generative diffusion mod- cls. arXiv preprint arXiv:2104. 026000 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2256-2263 2015.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. Interna- tional Conference on Learning Representations, 202 la.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, pp. 1 11895-11907, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative. Advances in Neural Information Processing Systems, 2020.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. arXiv e-prints, pp. arXiv-2101, 2021b.
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, 2021c
Yuxuan Song, Qiwei Ye, Minkai Xu, and Tie-Yan Liu. Discriminator contrastive divergence: Semi-amortized generative modeling by exploring energy of the discriminator. arXiv preprint arXiv:2004.00 01704, 2020.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arXiv preprint 1905v 09888, 2019a.

---

Published as a conference paper at ICLR 2022
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory. pp. 3084-3114. PMLR, 2019b.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. arXiv preprint arXii:2106.05940 2021.
Pascal Vincent. connection between score matching and denoising autoencoders. Neural Compu- lation, 23(7):1661-1674, 2011.
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning efficiently sam- ple from diffusion probabilistic models. arXiv preprint 2XXi:22106.088222 2021.
12

---

Published as a conference paper at ICLR 2022
A PROBABILITY FLOW ODE IN TERMS OF LOG-SNR
Song et al. (2021c) formulate the forward diffusion process in terms of an SDE of the form
dz = f(z, t)dt + g(t)dW,
(10)
and show that samples from this diffusion process can be generated by solving the associated prob- ability flow ODE:
1 (t) log pt
dz= [f - where in practice V, (z) is approximated by learned denoising model using
(11)
V. (z) ~ (Z)) Ze
(12)
of
Following Kingma al. (2021) we have f(, d'oga di
2, and (t) da? dt
2dllog oi. Assum-
dt
ing a variance preserving diffusion process with a (without loss of generality, see Kingma et al. (2021)), we get
of sigmoid (At) for X logla?/of
[(2,1) = d log ai: 2t 2
1 doog az dA
1
dd
dd
zt = 2
- 0z
2 = q
2t-
(13)
dt
dl
dt
dt
2
dt
Similarly, we get
gt) = ? d lo at 0g = do dd
dd
dx
dp
yP
dt dt
dd dt
ot dt
= (o 0,
dt
t dt
dt
(14)
Plugging these into the probability flow ODE then gives
dz = [f(z,t) 29 2 (t) & log (z()ddt
(15)
1
2 0 22 + V log PA (2)]dX
(16)
Plugging in our function approximation from Equation 12 gives
1
a2xo (2x) - 2x
dz
2x +
dd
(17)
2
0
1
(ax)x(Z2)) (zx) + (0 1)zx(dA
(18)
2
2 (os8o(Z2)) azzzldX
(19)
B DDIM IS AN INTEGRATOR OF THE PROBABILITY FLOW ODE
The DDIM update rule (Song& & Ermon, 2020) is given by
0, | xx + (0))
2,
(20)
ot
for s < t. Taking the derivative of this expression with respect to As assuming again a variance preserving diffusion process, and using day
d = 2arot and da, = - 2o gives dA
di. dd, [zz acx)(zz)] + day, xo (zi)
zx. d
(21)
dd,
a
o.
2 ac)00z1)] + / ,0,0,,0000
(22)
2 ax
Evaluating this derivative at s = then gives
2
1 a
2
lz axxx(z>)) + _010)
1
(23)
d
==1
2
1
1
al - axxx((z)]
+ ax (2)) 2
(24)
2
(x)))(z>) az
(25)
Comparison with Equation 19 now shows that DDIM follows the probability flow ODE up to first order, and can thus be considered an integration rule for this ODE.
13

---

Published as a conference paper at ICLR 2022
C EVALUATION OF INTEGRATORS OF THE PROBABILITY FLOW ODE
In preliminary investigation tried several numerical integrators for the probability flow ODE. As our model we used pre-trained class-conditional 128x128 ImageNet model following the de- scription in Ho et al. (2020). We tried simple Euler integrator, RK4 (the "classic" 4th order Runge- Kutta integrator), and DDIM (Song et al., 2021a). addition we compared to a Gaussian sampler with variance equal to the lower bound given by Ho al. (2020). We calculated FID scores on just 5000 samples, hence our results in this experiment are not comparable to results reported in the literature. This preliminary investigation gave the results listed in Table and identified DDIM as the best integrator in terms of resulting sample quality.
Sampler
Number of steps
FID
Stochastic
1000
13.35
Euler
1000
16.5
RK4
1000
16.33
DDIM
1000
15.98
Stochastic
100
18.44
Euler
100
23.67
RK4
100
18.94
DDIM
100
16.35
Table 3: Preliminary FID scores on 128 x 128 ImageNet for various integrators of the probability flow ODE, and compared against a stochastic sampler. Model specification and noise schedule follow Ho et al. (2020).
D EXPRESSION OF DDIM IN ANGULAR PARAMETERIZATION
We can simplify the DDIM update rule by expressing it in terms of o; = (rrtan(or//a)) rather than in terms of time or log-SNR λ₁ show here.
Given our definition of d. and assuming a variance preserving diffusion process, we have as = cos(e), 0g sin( and hence 2Φ cos(o)x sin(o)c. We now define the velocity of z
dz
dcos( dó dø
dsin(Ø)
V
x
dø = cos(Ø) - sin(o)xx
(26)
Rearranging e, X, v, we then get
sin())x= = cos(()c -v
(27)
cos(
(in(o) (z - cos(o)x) - -V
(28)
(in cos(o)z - cos (0)x - sin(o)v,
(29)
(sin + cos (0) x cos sin
(30)
and similarly we get sin( + cos(@)va Furthermore, we define the predicted velocity as
vo(zØ) = cos(o)εe(zz) sin(2) (2)) X (z.). where (0())) = (Z.g (z,)))))))))))))))))) sin()).
(31)
Rewriting the DDIM update rule the introduced terms then gives
ZØ, = cos( (x) + sin (zo.)
(32)
cos(o,)(con) (b) )zg, sin dh Vo (zo.)) ))
cos(or)vo(zs.)
(33)
= [cos(oo) c0s( son( (d.) sin (0:) Izo. + [sin(d.) cos( cos(o.) sin( Vi vo
(34)
14

---

Published as a conference paper at ICLR 2022
Finally, use the trigonometric identities
cos(Ø,) sin( sin (Q.) cos( = cos(o, - ②1)
(35)
sin(d.) (o. cos( sin((:) sin(d, 01),
(36)
to find that
cos(Ø> - Φ2) + sin( dh dr)
(37)
or equivalently
2 8 cos(6)zz. sin (in())))(Zz.)
(38)
Viewed from this perspective, DDIM thus evolves Zg, by moving circle in the (Zo.. Vo. ba- sis, along the -Vo direction. The relationship between ZΦ Vt, Oigy ot X, is visualized in Figure 5.
X
OK
Z
E
6
Figure 5: Visualization of reparameterizing the diffusion process in terms of ø and v
E SETTINGS USED IN EXPERIMENTS
Our model architectures closely follow those described by Dhariwal Nichol (2021). For 64 x 64 ImageNet their model exactly, with 192 channels the highest resolution. All other models are slight variations with different hyperparameters.
architecture with fixed number of channels all resolutions of 256. The model consists of a UNet that internally downsamples the data twice, to 16 x 16 and to x 8. At each resolution apply 3 residual blocks, like described by Dhariwal & Nichol (2021). We single-headed attention, and only apply this at the 16 x 16 and 8 x 8 resolutions. We use dropout of 0.2 when training the original model. No dropout is used during distillation.
For CIFAR-10 10
For LSUN we use a model similar to that for ImageNet, but with a reduced number of 128 channels the 64 x 64 resolution. Compared to ImageNet we have an additional level in the UNet, corre- sponding to the input resolution of 128 x 128, which we process using 3 residual blocks with 64 channels. We only use attention layers for the resolutions of 32 32 and lower.
For CIFAR-I0 we take the output of the model to represent a prediction of x directly, as discussed Section 4. For the other data sets we used the combined prediction of (x, c) like described in that section also. All original models are trained with Adam with standard settings (learning rate of 3 10-4), using a parameter moving average with constant 0.9999 and very slight decoupled weight decay (Loshchilov & Hutter, 2017) with constant of 0.001. We clip the norm of gradients to global norm of before calculating parameter updates. For CIFAR- 10 we train for 800k parameter updates, for ImageNet we train for 50k updates per iteration, except for the distillation to 2 and sampling steps, for which use 100k updates. We linearly anneal the leaming rate from 10- zero during each iteration.
550k updates, and for LSUN use 400k updates. During distillation
15

---

Published as a conference paper at ICLR 2022
We use batch size of 128 for CIFAR- 10 and 2048 for the other data sets. We experiments on TPUv4, using TPU chips for CIFAR-I0, and 64 chips for the other data sets. The total time required to first train and then distill model varies from about day for CIFAR-10, to about 5 days for ImageNet.
STOCHASTIC SAMPLING WITH DISTILLED MODELS
Our progressive distillation procedure was designed to be used with the DDIM sampler, but the resulting distilled model could in principle also be used with a stochastic sampler. Here we evaluate distilled model for 64x64 ImageNet using the optimized stochastic sampler also used in Section 5.2. The results are presented Figure 6.
64x64 ImageNet
Distilled DDIM
Distilled Stochastic
Undistilled Stochastic
10
8
FID
5
3
2
2
8
32
128 256
sampling steps
Figure 6: FID of generated samples from distilled and undistilled models, using DDIM or stochastic sampling. For the stochastic sampling results we present the best FID obtained by grid-search over 11 possible noise levels, spaced log-uniformly between the upper and lower bound on the variance derived by Ho et al. (2020). The performance of the distilled model with stochastic sampling found to lie in between the undistilled original model with stochastic sampling and the distilled DDIM sampler: For small numbers of sampling steps the DDIM sampler performs better with the distilled model, for large numbers of steps the stochastic sampler performs better.
G DERIVATION OF THE DISTILLATION TARGET
The key difference between our progressive distillation algorithm proposed Section and the standard diffusion training procedure is in how determine the target value for denoising model. In standard diffusson training, the target for denoising is the clean data x. In progressive distillation it is the value X the student denoising model would need to predict in order to match the teacher model when sampling. Here we derive what this target needs to be.
Using notation E' we have that the teacher model samples the next set of noisy data Ze" given the current noisy data z, by taking steps of DDIM. The student tries sample the same value in only step of DDIM. Denoting the student denoising prediction by &, and its one-step sample by Z application of the DDIM sampler (see equation 8), gives:
0.5/N and t" = t 1/N, when training student with N sampling steps,
Te'
2,"0 = ar'x* +
(ze - axx).
(39)
σt
16

---

Published as a conference paper at ICLR 2022
In order for the student sampler to match the teacher sampler, must set Znr equal to z,". This gives
0r"
2p = 0008+
(z: -x²) = Zp"
(40)
0:
=(am"
0e" C11
0t"
Zt Z,!'
(41)
0t
01
Oe"
Te" Zt
Oty
at
x Zz!
(42)
Tt
Ot
21"
x
(43)
44"
0y"
In other words, if our student denoising model exactly predicts X defined in equation 43 above, then the one-step student sample Z is identical to the two-step teacher sample Zzn. In order to have student model approximate this ideal outcome, thus train it to predict x from 2, as well as possible, using the standard squared error denoising loss (see Equation 9).
Note that this possibility of matching the two-step teacher model with one-step student model is unique to deterministic samplers like DDIM: the composition of two standard stochastic DDPM sampling steps (Equation 5) forms non-Gaussian distribution that falls outside the family of Gaus- sian distributions that can be modelled by a single DDPM student step: A multi-step stochastic DDPM sampler can thus not be distilled into few-step sampler without some loss in fidelity. This is in contrast with the deterministic DDIM sampler: here both the two-step DDIM teacher update and the one-step DDIM student update represent deterministic mappings implemented by neural net, which is why the student is able to accurately match the teacher.
Finally, note that we lose something during the progressive distillation process: while the original model was trained to denoise 2; for any given continuous time f. the distilled student models are only ever evaluated on a small discrete set of times t. The student models thus lose generality as distillation progresses. At the same time, it's this loss of generality that allows the student models to free up enough modeling capacity to accurately match the teacher model without increasing their model size.
H ADDITIONAL RANDOM SAMPLES
In this section we present additional random samples from our diffusion models obtained through progressive distillation. We show samples for distilled models taking 256. 4, and sampling steps. All samples are uncurated.
As explained in Section 3, our distilled samplers implement a deterministic mapping from input noise to output samples (also see Appendix G). To facilitate comparison of this mapping for varying numbers of sampling steps, we generate all samples using the same random input noise, and we present the samples side-by-side. As these samples show, the mapping is mostly preserved when moving from many steps to a single step: The same input noise is mapped to the same output image, with slight loss in image quality, as the number of steps is reduced. Since the mapping is preserved while reducing the number of steps, our distilled models also preserve the excellent sample diversity of diffusion models (see e.g. Kingma et al. (2021)).
17

---

Published as a conference paper at ICLR 2022
(a) 256 sampling steps
(b) sampling steps
(c) sampling step
Figure 7: Random samples from our distilled CIFAR-10 models, for fixed random seed and for varying number of sampling steps.
(a) 256 sampling steps
(b) sampling steps
(c) sampling step
Figure 8: Random samples from our distilled 64 x 64 ImageNet models, conditioned on the 'coral reef' class, for fixed random seed and for varying number of sampling steps.
(a) 256 sampling steps
(b) sampling steps
(c) sampling step
Figure 9: Random samples from our distilled 64 64 ImageNet models, conditioned on the "sports car' class, for fixed random seed and for varying number of sampling steps.
18

---

Published as a conference paper at ICLR 2022
(a) 256 sampling steps
(b) sampling steps
(c) 1 sampling step
Figure 10: Random samples from our distilled LSUN bedrooms models, for fixed random seed and for varying number of sampling steps.
(a) 256 sampling steps
(b) sampling steps
(c) sampling step
Figure 11: Random samples from our distilled LSUN church-outdoor models, for fixed random seed and for varying number of sampling steps.
19

---

Published as a conference paper at ICLR 2022
I ABLATION WITH FASTER DISTILLATION SCHEDULES
In order to further reduce the computational requirements for our progressive distillation approach, perform an ablation study CIFAR-IO,, where we decrease the number of parameter updates we use to train each new student model. In Figure 12 we present results for taking 25 thousand, 10 thousand, or thousand optimization steps, instead of the 50 thousand we suggested in Section 3. As the results show, we can drastically decrease the number of optimization steps taken, and still get very good performance when using 2 sampling steps. When taking very few sampling steps, the loss in performance becomes more pronounced when training the student for only short time.
In addition to just decreasing the number of parameter updates, we also experiment with schedule where we train each student on 4 times fewer sampling steps than its teacher, rather than the 2 times we propose Section Here the denoising target is still derived from taking 2 DDIM steps with the teacher model usual, since taking 4 teacher steps would negate most of the computational savings. As Figure 12 shows, this does not work as well: if the computational budget is limited, it's better to take fewer parameter updates per halving of the number of sampling steps then to skip distillation iterations altogether.
In Figure 13 show the results achieved with faster schedule for the ImageNet and LSUN datasets. Here also, we achieve excellent results with a faster distillation schedule.
CIFAR-10 10
20
Distilled 50k
Distilled 25k
Distilled 10k
10
Distilled 5k
Distilled 4x
8
7
FID
6
5
4
3
2
2
8
16 32
64 128 256 512
sampling steps
Figure 12:
Comparing our proposed schedule for progressive distillation taking 50k parameter up- dates to train student every time the number of steps is halved, versus fast sampling schedules taking fewer parameter updates (25k, 10k, 5k), and a fast schedule dividing the number of steps by 4 for every new student instead of by 2. All reported numbers For each schedule we selected the optimal learning rate from [5e -5 5, le 2e 3e 3e
averages over 4 random seeds.
20

---

Published as a conference paper at ICLR 2022
64x66 ImageNet
Bedrosms
Church-Oundorr
20
20
20
50% uplinis updanes
50k updates
SOk upsiness
10k
FID
sampling
sampling ling steps
Figure 13: Comparing our proposed schedule for progressive distillation taking 50k parameter up- dates to train new student every time the number of steps is halved, versus a fast sampling schedule taking 10k parameter updates. For each reported number of steps we selected the optimal learning rate from [5e 5, 14_4, 2e 3e *1. Results are for single random seed.
21


Reviews:
[{"review_id": "lchbhViJe-t", "reviewer": "Reviewer_HZnL", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "**Strengths**:\n- The paper successfully addresses a crucial downside of generative diffusion models, i.e., their slow sampling speed. This is a problem of very high relevance.\n- The idea to use the deterministic DDIM [1] sampling scheme and distill 2 teacher steps into a single student step in each round of distillation is interesting and novel. I think it makes the distillation process easier compared to the situation where we would have to deal with a fully stochastic sampler during distillations. The overall algorithm is relatively straight-forward, which is good.\n- The paper demonstrates strong empirical results, outperforming all previous works that aimed specifically to accelerate sampling from diffusion models. In particular, reaching FID scores $\\leq3.0$ in as few as 4 steps on CIFAR-10 is a very strong result.\n\n**Weaknesses**:\n- The work tackles image generation and aims purely at reaching high FID in as few as possible sampling steps. If our only goal is to reach high perceptual quality as measured by FID with quick sampling, then one could also simply use state-of-the-art GANs [2], which reach FID scores $\\leq3.0$ in a single shot (on CIFAR-10). I think that one of the crucial advantages of generative diffusion models is that they reach particularly strong diversity, do not easily drop modes, and better cover the full data distribution in their standard formulation (i.e. no distillation), unlike GANs, which are known to suffer from mode dropping and related challenges. It is unclear whether these advantages of diffusion models are preserved or lost in the proposed distillation process. This aspect is not discussed in the paper and no experiments to investigate this aspect are presented. This suggests the following: (a) Is it possible to somehow measure the likelihood of held-out validation data under the final distilled model? (b) As I assume that this may be difficult, could we calculate standard recall metrics [3] and show that distribution coverage is not significantly decreased in the distillation (and maybe also compare to GANs)? (c) To provide some further evidence towards diversity preservation and mode coverage, we could also run 2D toy experiments on a mixture of Gaussians and show that the distilled model still faithfully covers all modes (standard GANs often tend to miss modes in such settings).\n- The explanation why $\\epsilon$-prediction is not well-suited for the distillation approach makes sense. However, the proposed solutions, this is, the different proposed model parametrizations and loss weightings, seem to come primarily from trial-and-error and aren't overly well motivated. Is there anything we can say about which parametrizations and loss weightings should be optimal with respect to certain criteria? Can we provide more insights here?\n- I sort of agree with the intuition that using the DDIM framework for the distillation protocol makes sense. However, what would happen if we did a similar distillation while using the standard stochastic DDPM sampling? Would this still be possible? A more thorough explanation or an experiment would be very interesting to clearly show that it would not work and provide a more solid motivation for choosing DDIM-based distillation.\n- The method is evaluated on image generation only. Will the distillation approach also work for diffusion models in other domains? What could be additional challenges that may potentially occur? In fact, I think it would be very promising to also apply the method on other data types, because in non-image domains there is less \"competition\" from the GAN literature.\n\n**Additional Questions and Suggestions**:\n- In the interest of a clean presentation, I would suggest to properly introduce the $\\epsilon$ in equation (2).\n- In cection 3: \"By inverting a single step of DDIM, we then calculate the value the student model would need to predict in order to step from $z_t$ to $z_{t-1/N}$ in a single step.\" seems to correspond to the operation at \"Teacher $\\hat{x}$ target\" in algorithm 2. I would consider explaining this in a bit more detail with more background and explicitly derive the expression for $\\tilde{x}$ in algorithm 2. I think it could be made clearer and easier to follow and this is one of the most central operations in the proposed algorithm.\n- If we distill the slow iterative sampling from diffusion models into these few-step samplers, one may intuitively expect that we pay a price in terms of requiring bigger network capacity to make up for the fewer synthesis steps. This does not seem to be the case, however. Why not? Note that this question is also related to the first point under *Weaknesses* above (i.e., maybe we don't use a bigger model, but have to pay a price with regards to diversity and mode coverage?).\n\n\n[1] Song et al. \"Denoising Diffusion Implicit Models\", 2020.\n\n[2] Karras et al. \"Training Generative Adversarial Networks with Limited Data\", 2020.\n\n[3] Kynkäänniemi et al. \"Improved Precision and Recall Metric for Assessing Generative Models\", 2019.", "summary_of_the_paper": "Diffusion models have recently emerged as a promising class of generative models. Their main drawback is the slow, iterative sampling, which often requires hundreds or thousands of neural network calls. The paper proposes a method to progressively distill this slow iterative generation process into much faster synthesis relying on fewer denoising steps, while almost preserving generation quality as measured by metrics quantifying perceptual image quality. The method progressively reduces the number of sampling steps required by a factor of 2 in multiple rounds of distillation and relies on a teacher-student framework. The student model learns to model 2 steps of the teacher in a single step. The method leverages the deterministic synthesis scheme from Denoising Diffusion Implicit Models for that purpose. To get the method to work successfully, new score model parametrizations and score matching loss weightings are proposed. The quantitative results are promising, for example achieving FID scores $\\leq3.0$ on CIFAR-10 in as few as 4 steps.", "summary_of_the_review": "The distillation idea in itself isn't entirely novel, as it was proposed by [4], too. However, I certainly acknowledge that this work shows how to do it in an elegant, scalable and high-performance manner using a progressive protocol. Overall, the paper addresses an important problem and shows promising quantitative results. Hence, I am leaning towards recommending acceptance. That said, some weaknesses remain. I would be willing to raise my score, if the weak points were addressed in a satisfactory manner.\n\n[4] Luhmann and Luhmann, \"Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed\", 2021.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RkXwFOE29eM", "reviewer": "Reviewer_3Sjb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "## Strengths\n\n1. Speeding up diffusion models is an important research direction, and achieving competitive performance with as few as 4 iterations is a remarkable achievement, which I expect to engender major impacts in the community. \n2. Design choices like different parameterization of the score model, and different weighting functions in the score matching loss are clearly discussed. Their impact on model performance is also investigated rigorously in Table 1.\n3. I especially like the discussion of DDIM as a numerical integrator for the probability flow ODE, and the experimental study on the performance of different numerical ODE solvers in the appendix.\n4. Experiments demonstrate clear advantage over competing methods. It is surprising that even one iteration of the method is already sufficient for generating realistic images.\n\n## Weaknesses\n\nThis paper is mostly focused on the variance preserving diffusion process. It is unclear how the distillation approached introduced in this paper can be applied to other SDEs such as VE SDEs and subVP SDEs, as well as how to choose the weighting function and score function parameterization in those cases.", "summary_of_the_paper": "This work proposes to speed up diffusion models by progressively distilling a score network with deterministic dynamics. Authors report significant improvement of sample quality using very few iterations. Authors also provide a proof that DDIM is a special numerical integrator for the probability flow ODE, connecting two existing approaches. In addition, empirical studies on the weighting function and parameterization of score functions are also discussed.", "summary_of_the_review": "Well written paper with impressive experimental results and clear description of design choices.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QF3uoyKF5hj", "reviewer": "Reviewer_orkh", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths:\n\n(1) Overall, the presentation is clear and easy to understand, though there exist some minor issues with typos and confusing sentences (see comments below).\n\n(2) The proposed progressive distillation is novel and interesting. Though the distillation idea to reduce the sampling time of diffusion models is not new, doing so in a progressive way seems to be more effective and efficient (which has been clearly discussed in the paper).\n\n(3) The new parameterizations and training losses of diffusion models have been well motivated and explained to accommodate the proposed progressive distillation.  \n\n(4) Experiments have well supported the main claims and the effectiveness of the proposed method in reducing sampling time. In particular, the comparison with previous fast sampling methods on CIFAR-10 demonstrates that it largely outperforms these strong methods regarding FID scores with a few model evaluations. \n\nWeaknesses:\n\n(1) The progressive distillation process seems to need a much larger computational cost than many previous fast sampling methods, such as DDIM and DDPM respacing. As stated in the paper, its training budget is almost the same as training a diffusion model from scratch. I wonder how this concern can be addressed in practice?\n\n(2) The main claim is that the method can reduce the number of model evaluations in sampling to as small as 4 or 8 steps while retaining high image quality. Currently, the highest resolution of considered datasets is 128x128. I wonder if the claim still holds for datasets with higher resolution? Does the resolution or complexity of the dataset impact the final steps of model evaluations? \n\n(3) Minor issues in the writing. 1) Typos. For example, in the abstract, “as little as 4 steps” => “as few as 4 steps”. 2) Repeated references. For example, \t“Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis.”, and “Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. ”. 3) Confusing sentences. For example, what do you mean by saying “... unlike the original data point x, since multiple different data points x could conceivably have led to observing noisy data $z_t$”? Also, when saying “we found this to work slightly better than starting from a non-zero signal-to-noise ratio as used by e.g. Ho et al. (2020)”, does it refer to the undistilled sampler or the distilled sampler?", "summary_of_the_paper": "This work studies the problem of improving the sampling speed of diffusion models, which is currently a timely and challenging topic. To tackle this problem, this work first proposes several new parameterizations of diffusion models for better stability along with a fast sampling, and then proposes a progressive distillation method that progressively distills two DDIM steps of teacher diffusion models to one DDIM step of student diffusion models. In the end, the method can reduce the number of model evaluations in sampling to as small as 4 or 8 steps while retaining high image quality. Experiments on CIFAR10, ImageNet (64x64), LSUN-bedrooms (128x128) and LSUN-Church (128x128) were conducted to show the effectiveness of the progrssive distillation in reducing the sampling time.\n", "summary_of_the_review": "The proposed method is novel and effective in reducing the sample time of diffusion models, and experiments on several datasets well support its effectiveness and better performance than prior fast sampling methods. Still, I have some concerns about the method regarding computational cost and generalization to larger datasets. Thus, my initial rating is a weak accept. I’m willing to increase the score if my concerns can be addressed. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1Sqf-VFXD4Y", "reviewer": "Reviewer_2wPQ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths\n- An important topic given that diffusion models achieve very good results but that their sampling time is a barrier to wider adoption in practice.\n- The idea of progressively distilling the model seems to be original and useful to avoid the need to evaluate the teacher model for many steps.\n- Distillation to very few number of steps requires adaptions on the parameterization and loss weighting. Different plausible choices are presented and evaluated against each other. I particularly liked that this was evaluated on the original diffusion models without distillation to avoid confounding factors and to also provide insight on what works for these models themselves. Knowing that x-prediction works just as well is good to know since it can simplify training quite a bit.\n- Overall, there are quite a few design choices for diffusion models which haven't been evaluated exhaustively and many of the results in this paper will help towards a better understanding on these choices (e.g. parameterization, weighting function, using signal-noise-ratio of 0 at t=1, different integrators evaluated in the supplementary).\n- It also shows a nice interpretation of DDIM as an integrator of the probability flow ODE.\n- Impressive results on retaining sample quality with very few steps, especially compared to previous approaches.\n- The overall presentation is very clear, including the background on diffusion models, the proposed approach and difficulties, the algorithm, the results and the discussion of related works which also include very recent works.\n\nWeaknesses\n- The effect of distilling iteratively/progressively is not fully explored. Results in Tab. 2 suggest that when distilling to a single step model, the \"one-step\" distillation approach of Luhman&Luhman, 2021, performs similar to the presented approach. While potential advantages of not having to create a training datasets using the original model are mentioned in the Related Work Section, it would be nice to see a more detailed analysis, which, e.g. compares training costs between the two approaches. It would also be interesting to see the effect of choosing a different schedule for the progressive distillation, for example reducing the number of integration steps by a factor of four instead of two for each new student. This could have effects on both the time required for the distillation process as well as the quality achieved by the distilled models. In general, a bit more details on the time required for the distillation process and comparing that to other approaches might be useful.\n\nMinor\n- A bit contrary to what is claimed in the second paragraph of the introduction, [Dhariwal and Nichol, 2021b] also report good results with 25 steps in a class conditional setting (even though I agree with the general sentiment that tasks with less conditioning are more challenging).", "summary_of_the_paper": "The paper presents an approach to improve sampling speed of diffusion models. The idea is to iteratively learn integrators for the probability flow ODE which, given a diffusion model, maps noise to data samples. In each iteration, the integrator is trained to integrate steps of the ODE over twice the length than the previous model. Thus, the approach resembles an iterative distillation process, where a faster student integrator is learned from a slower teacher integrator. After the student has finished learning, it becomes the teacher for the next distillation iteration.\n\nEnabling this distillation requires changes to the parameterization and loss weighting, for which different, intuitive options are presented. Experiments evaluate the different design choices and compare the resulting models to other recent works on the topic, which demonstrates significant gains of the proposed approach in the regime where a small number of evaluations are used.", "summary_of_the_review": "Given the impressive results obtained with diffusion models, the paper addresses an important and timely topic. The proposed approach significantly outperforms other methods in terms of quality achieved with distilled diffusion models at small number of evaluation steps, which, ultimately, is one of the main goals on this topic (even though the distillation process might be relatively expensive). The presentation is very clear and easy to follow and different possible design choices are well motivated and intuitive to understand. The different choices are evaluated in well designed experiments and comparisons are very convincing of the benefits of the presented approach and also include very recent works on the topic. The supplementary material also contains interesting material on different formulations and connections between different approaches. Thus, I recommend accepting this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a faster sampling method for diffusion based generative models which are usually slow in practice. The key idea is based a progressive distillation approach (e.g., how to distill a 4 step sampler into a 1 step sampler). The paper studies the various design choices for diffusion models which existing work hasn't looked at that deeply and sheds light on the effects of these choices. The paper also shows that DDIM can be seen as a numerical integrator for probability flow ODE. The experimental results are impressive. 

There were some concerns such as the effect of progressive distillation and the overhead of distilling the diffusion model but the authors provided a satisfactory response and backed it up with additional results.

Overall, this is a nice paper on making diffusion based generative models generate faster samples and also provides novel insights into the behavior of these models under various design choices. Given the significant recent interest in these models which are pretty impressive in terms of generation quality but slow, the paper indeed makes a timely contribution which will fuel further interest in these models.  All the reviewers have voted for acceptance. Based on my own reading, the reviewers' assessments, the discussions, and the authors' response, I would vote for acceptance.
Decision: Accept (Spotlight)