Year: 2022

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. 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 state-of-the-art 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 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.

heading: 1 INTRODUCTION
text: Diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) are a class of generative models that have recently delivered impressive sample quality and likelihoods on many standard generative modeling benchmarks. These models have achieved ImageNet generation results outperforming 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 autoregressive image models (Kingma et al., 2021; Song et al., 2021b). Diffusion models have also succeeded in image super-resolution (Saharia et al., 2021; Li et al., 2021) and image inpainting (Song et 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 et al., 2021).
A major barrier remains to practical adoption of diffusion models: sampling speed. While sampling can be accomplished in relatively few steps in strongly conditioned settings, such as text-tospeech (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 is substantially different 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 unconditional and class-conditional image generation, which represent the setting in which diffusion models have been slowest in previous work. We present a procedure to distill the behavior of aN -step DDIM sampler (Song et al., 2021a) for a pretrained diffusion model into a new model with N/2 steps, with little degradation in sample quality. In what we call progressive distillation, we repeat this distillation 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.

heading: 2 BACKGROUND ON DIFFUSION MODELS
text: We consider diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) specified in continuous time (Tzen & Raginsky, 2019a; Song et al., 2021c; Chen et al., 2021; Kingma et al., 2021). We use x ∼ p(x) to denote training data. A diffusion model has latent variables z = {zt | t ∈ [0, 1]} and is specified by a noise schedule comprising differentiable functions αt, σt such that λt = log[α2t /σ 2 t ], the log signal-to-noise-ratio, decreases monotonically with t.
These ingredients define the forward process q(z|x), a Gaussian process satisfying the following Markovian structure:
q(zt|x) = N (zt;αtx, σ2t I), q(zt|zs) = N (zt; (αt/αs)zs, σ2t|sI) (1)
where 0 ≤ s < t ≤ 1 and σ2t|s = (1− e λt−λs)σ2t .
The role of function approximation in the diffusion model is to denoise zt ∼ q(zt|x) into an estimate x̂θ(zt) ≈ x (the function approximator also receives λt as an input, but we omit this to keep our notation clean). We train this denoising model x̂θ using a weighted mean squared error loss
E ,t [ w(λt)‖x̂θ(zt)− x‖22 ] (2)
over uniformly sampled times t ∈ [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 as a form of denoising score matching (Vincent, 2011; Song & Ermon, 2019). We will discuss particular choices of weighting function w(λt) later on.
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 q(zs|zt,x) = N (zs; µ̃s|t(zt,x), σ̃2s|tI) (noting s < t), where
µ̃s|t(zt,x) = e λt−λs(αs/αt)zt + (1− eλt−λs)αsx, σ̃2s|t = (1− e λt−λs)σ2s (3)
We use this reversed description of the forward process to define the ancestral sampler. Starting at z1 ∼ N (0, I), the ancestral sampler follows the rule
zs = µ̃s|t(zt, x̂θ(zt)) + √ (σ̃2s|t) 1−γ(σ2t|s) γ) (4)
= eλt−λs(αs/αt)zt + (1− eλt−λs)αsx̂θ(zt) + √ (σ̃2s|t) 1−γ(σ2t|s) γ) , (5)
where is standard Gaussian noise, and γ is a hyperparameter that controls how much noise is added during sampling, following Nichol & Dhariwal (2021).
Alternatively, Song et al. (2021c) show that our denoising model x̂θ(zt) can be used to deterministically map noise z1 ∼ N (0, I) to samples x by numerically solving the probability flow ODE:
dzt = [f(zt, t)− 1
2 g2(t)∇z log p̂θ(zt)]dt, (6)
where ∇z log p̂θ(zt) = αtx̂θ(zt)−ztσ2t . Following Kingma et al. (2021), we have f(zt, t) = d logαt dt zt and g2(t) = dσ 2 t
dt − 2 d logαt dt σ 2 t . Since x̂θ(zt) is parameterized by a neural network, this equation
is a special case of a neural ODE (Chen et al., 2018), also called a 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 we show in Appendix B, even though it was originally proposed with a different motivation. The update rule specified by DDIM is
zs = αsx̂θ(zt) + σs zt − αtx̂θ(zt)
σt (7)
= e(λt−λs)/2(αs/αt)zt + (1− e(λt−λs)/2)αsx̂θ(zt), (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 x̂θ(zt) 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 → ∞. 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.

heading: 3 PROGRESSIVE DISTILLATION
text: To make diffusion models more efficient at 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 et 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 zt. 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 zt and ending at zt−1/N , withN being the number of student sampling steps. By inverting a single step of DDIM, we then calculate the value the student model would need to predict in order to move from zt to zt−1/N in a single step, as we show in detail in Appendix G. The resulting target value x̃(zt) is fully determined given the teacher model and starting point zt, which allows the student model to make a sharp prediction when evaluated at zt. In contrast, the original data point x is not fully determined given zt, since multiple different data points x can produce the same noisy data zt: this means that the original denoising model is
predicting a weighted average of possible x values, which produces a blurry prediction. By making sharper predictions, the student model can make faster progress during sampling.
After running distillation to learn a student model taking N sampling steps, we can repeat the procedure with N/2 steps: The student model then becomes the new teacher, and a new student model is initialized by making a copy of this model.
Unlike our procedure for training the original model, we always run progressive distillation in discrete time: we sample this discrete time such that the highest time index corresponds to a signal-tonoise ratio of zero, i.e. α1 = 0, which exactly matches the distribution of input noise z1 ∼ N (0, I) that is used at test time. We found this to work slightly better than starting from a non-zero signalto-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 1 Standard diffusion training Require: Model x̂θ(zt) to be trained Require: Data set D Require: Loss weight function w()
while not converged do x ∼ D . Sample data t ∼ U [0, 1] . Sample time ∼ N(0, I) . Sample noise zt = αtx+ σt . Add noise to data
x̃ = x . Clean data is target for x̂ λt = log[α 2 t /σ 2 t ] . log-SNR Lθ = w(λt)‖x̃− x̂θ(zt)‖22 . Loss θ ← θ − γ∇θLθ . Optimization
end while
Algorithm 2 Progressive distillation Require: Trained teacher model x̂η(zt) Require: Data set D Require: Loss weight function w() Require: Student sampling steps N
for K iterations do θ ← η . Init student from teacher while not converged do
x ∼ D t = i/N, i ∼ Cat[1, 2, . . . , N ] ∼ N(0, I) zt = αtx+ σt # 2 steps of DDIM with teacher t′ = t− 0.5/N , t′′ = t− 1/N zt′ = αt′ x̂η(zt) +
σt′ σt (zt − αtx̂η(zt)) zt′′ = αt′′ x̂η(zt′) +
σt′′ σt′ (zt′ − αt′ x̂η(zt′)) x̃ = zt′′−(σt′′/σt)ztαt′′−(σt′′/σt)αt
. Teacher x̂ target λt = log[α 2 t /σ 2 t ] Lθ = w(λt)‖x̃− x̂θ(zt)‖22 θ ← θ − γ∇θLθ end while η ← θ . Student becomes next teacher N ← N/2 . Halve number of sampling steps
end for

heading: 4 DIFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS
text: In this section, we discuss how to parameterize the denoising model x̂θ, and how to specify the reconstruction loss weight w(λt). We assume a standard variance-preserving diffusion process for which σ2t = 1 − α2t . This is without loss of generalization, as shown by (Kingma et al., 2021, appendix G): different specifications of the diffusion process, such as the variance-exploding specification, can be considered equivalent to this specification, up to rescaling of the noisy latents zt. We use a cosine schedule αt = cos(0.5πt), similar to that introduced by Nichol & Dhariwal (2021).
Ho et al. (2020) and much of the following work choose to parameterize the denoising model through directly predicting with a neural network ̂θ(zt), which implicitly sets x̂θ(zt) = 1αt (zt−σt̂θ(zt)). In this case, the training loss is also usually defined as mean squared error in the -space:
Lθ = ‖ − ̂θ(zt)‖22 = ∥∥∥∥ 1σt (zt − αtx)− 1σt (zt − αtx̂θ(zt)) ∥∥∥∥2 2 = α2t σ2t ‖x− x̂θ(zt)‖22, (9)
which can thus equivalently be seen as a weighted reconstruction loss in x-space, where the weighting function is given by w(λt) = exp(λt), for log signal-to-noise ratio λt = log[α2t /σ 2 t ].
Although this standard specification works well for training the original model, it is not well suited for distillation: when training the original diffusion model, and at the start of progressive distillation, the model is evaluated at a wide range of signal-to-noise ratios α2t /σ 2 t , 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 ̂θ(zt) on the implied prediction in x-space is increasingly amplified, since x̂θ(zt) = 1αt (zt − σt̂θ(zt)) divides by αt → 0. This is not much of a problem when taking many steps, since the effect of early missteps is limited by clipping of the zt iterates, and later updates can correct any mistakes, but it becomes increasingly important as we decrease the number of sampling steps. Eventually, if we distill all the way down to a single sampling step, the input to the model is only pure noise , which corresponds to a signal-to-noise ratio of zero, i.e. αt = 0, σt = 1. At this extreme, the link between -prediction and x-prediction breaks down completely: observed data zt = is no longer informative of x and predictions ̂θ(zt) 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.
For distillation to work, we thus need to parameterize the diffusion model in a way for which the implied prediction x̂θ(zt) remains stable as λt = log[α2t /σ 2 t ] varies. We tried the following options, and found all to work well with progressive distillation:
• Predicting x directly. • Predicting both x and , via separate output channels {x̃θ(zt), ̃θ(zt)} of the neural net-
work, and then merging the predictions via x̂ = σ2t x̃θ(zt) + αt(zt − σt̃θ(zt)), thus smoothly interpolating between predicting x directly and predicting via .
• Predicting v ≡ αt − σtx, which gives x̂ = αtzt − σtv̂θ(zt), as we show in Appendix D.
In Section 5.1 we test all three parameterizations on training an original diffusion model (no distillation), and find them to work well there also.
In addition to determining an appropriate parameterization, we also need to decide on a reconstruction loss weighting w(λt). The setup of Ho et al. (2020) weights the reconstruction loss by the signal-to-noise ratio, implicitly gives a weight of zero to data with zero SNR, and is therefore not a suitable choice for distillation. We consider two alternative training loss weightings:
• Lθ = max(‖x− x̂t‖22, ‖ − ̂t‖22) = max( α2t σ2t , 1)‖x− x̂t‖22; ‘truncated SNR’ weighting. • Lθ = ‖vt − v̂t‖22 = (1 + α2t σ2t )‖x− x̂t‖22; ‘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 αt, σt are sampled during training, as this sampling distribution strongly determines the weight the expected loss gives to each signal-to-noise ratio. Our results are for a cosine schedule αt = cos(0.5πt), 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.

heading: 5 EXPERIMENTS
text: 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 to 128× 128. All experiments use the cosine schedule αt = cos(0.5πt), 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 et al. (2021c). Our training setup closely matches the open source code by Ho et al. (2020). Exact details are given in Appendix E. We plan to open source the code soon.

heading: 5.1 MODEL PARAMETERIZATION AND TRAINING LOSS
text: As explained in Section 4, the standard method of having our model predict , and minimizing mean squared error in the -space (Ho et al., 2020), is not appropriate for use with progressive distillation.
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 error in x-space. Here, we perform a complete ablation experiment of all parameterizations 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 benchmark. We measure performance of undistilled models trained from scratch, to avoid introducing too many factors of variation into our analysis.
Table 1 lists the results of the ablation study. Overall results are fairly close across different parameterizations and loss weights. All proposed stable model specifications achieve excellent performance, 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 as predicting v, or the combination ( ,x), could thus be recommended for specification of diffusion
models. Here, predicting v 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.

heading: 5.2 PROGRESSIVE DISTILLATION
text: We evaluate our proposed progressive distillation algorithm on 4 data sets: CIFAR-10, 64 × 64 downsampled ImageNet, 128 × 128 LSUN bedrooms, and 128 × 128 LSUN Church-Outdoor. For each data set we start by training a baseline model, after which we start the progressive distillation procedure. For CIFAR-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 2 and 1 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 of 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. At 2 or 1 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 or 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

heading: 6 RELATED WORK ON FAST SAMPLING
text: 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 that it requires constructing a large data set by running the original model at 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 at the full number of sampling steps: at every iteration of progressive distillation, the number of model evaluations is independent of the number of teacher sampling steps, allowing our method to scale up to large numbers of teacher steps at a 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). Jolicoeur-Martineau 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.
Other work on fast sampling can be viewed as manual or automated methods to adjust samplers or diffusion processes for fast generation. Nichol & Dhariwal (2021); Kong & Ping (2021) describe
methods to adjust a discrete time diffusion model trained on many timesteps into models that can sample in few timesteps. Watson et al. (2021) describe a dynamic programming algorithm to reduce the number of timesteps for a diffusion model in a way that is optimal for log likelihood. Chen et al. (2021); Saharia et al. (2021); Ho et al. (2021) train diffusion models over continuous noise levels and tune samplers post training by adjusting the noise levels of a few-step discrete time reverse diffusion process. Their method is effective in highly conditioned settings such as text-to-speech and image super-resolution. San-Roman et al. (2021) train a new network to estimate the noise level of noisy data and show how to use this estimate to speed up sampling.
Alternative specifications of the diffusion model can also lend themselves to 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).

heading: 7 DISCUSSION
text: We have presented progressive distillation, a method to 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: in 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 of images and also explore progressive distillation of diffusion models for different data modalities such as e.g. audio (Chen et al., 2021).
In addition to the proposed distillation procedure, some of our progress was realized through different parameterizations of the diffusion model and its training loss. We expect to see more progress in this direction as the community further explores this model class.
REPRODUCIBILITY STATEMENT
We provide full details on model architectures, training procedures, and hyperparameters in Appendix E, in addition to our discussion in Section 5. In Algorithm 2 we provide fairly detailed pseudocode that closely matches our actual implementation. We plan to open source our code and trained model checkpoints with the release of the final version of this paper.
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 in order to reduce their computational demands; we do not have specific concerns with regards to this contribution.

heading: A PROBABILITY FLOW ODE IN TERMS OF LOG-SNR
text: 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 probability flow ODE:
dz = [f(z, t)− 1 2 g2(t)∇z log pt(z)]dt, (11)
where in practice∇z log pt(z) is approximated by a learned denoising model using
∇z log pt(z) ≈ αtx̂θ(zt)− zt
σ2t . (12)
Following Kingma et al. (2021) we have f(z, t) = d logαtdt zt and g 2(t) = dσ2t dt −2 d logαt dt σ 2 t . Assuming a variance preserving diffusion process with α2t = 1 − σ2t = sigmoid(λt) for λt = log[α2t /σ2t ] (without loss of generality, see Kingma et al. (2021)), we get
f(z, t) = d logαt dt zt = 1 2 d logα2λ dλ dλ dt zt = 1 2 (1− α2t ) dλ dt zt = 1 2 σ2t dλ dt zt. (13)
Similarly, we get
g2(t) = dσ2t dt − 2d logαt dt σ2t = dσ2λ dλ dλ dt − σ4t dλ dt = (σ4t − σ2t ) dλ dt − σ4t dλ dt = −σ2t dλ dt . (14)
Plugging these into the probability flow ODE then gives
dz = [f(z, t)− 1 2 g2(t)∇z log pt(z)]dt (15)
= 1
2 σ2λ[zλ +∇z log pλ(z)]dλ. (16)
Plugging in our function approximation from Equation 12 gives
dz = 1
2 σ2λ
[ zλ + ( αλx̂θ(zλ)− zλ
σ2λ
)] dλ (17)
= 1
2 [αλx̂θ(zλ) + (σ
2 λ − 1)zλ]dλ (18)
= 1
2 [αλx̂θ(zλ)− α2λzλ]dλ. (19)

heading: B DDIM IS AN INTEGRATOR OF THE PROBABILITY FLOW ODE
text: The DDIM update rule (Song & Ermon, 2020) is given by
zs = σs σt [zt − αtx̂θ(zt)] + αsx̂θ(zt), (20)
for s < t. Taking the derivative of this expression with respect to λs, assuming again a variance preserving diffusion process, and using dαλdλ = 1 2αλσ 2 λ and dσλ dλ = − 1 2σλα 2 λ, gives
zλs dλs = dσλs dλs 1 σt [zt − αtx̂θ(zt)] + dαλs dλs x̂θ(zt) (21)
= −1 2 α2s σs σt [zt − αtx̂θ(zt)] + 1 2 αsσ 2 s x̂θ(zt). (22)
Evaluating this derivative at s = t then gives zλs dλs |s=t = − 1 2 α2λ[zλ − αλx̂θ(zλ)] + 1 2 αλσ 2 λx̂θ(zλ) (23)
= −1 2 α2λ[zλ − αλx̂θ(zλ)] + 1 2 αλ(1− α2λ)x̂θ(zλ) (24) = 1
2 [αλx̂θ(zλ)− α2λzλ]. (25)
Comparison with Equation 19 now shows that DDIM follows the probability flow ODE up to first order, and can thus be considered as an integration rule for this ODE.

heading: C EVALUATION OF INTEGRATORS OF THE PROBABILITY FLOW ODE
text: In a preliminary investigation we tried several numerical integrators for the probability flow ODE. As our model we used a pre-trained class-conditional 128x128 ImageNet model following the description in Ho et al. (2020). We tried a simple Euler integrator, RK4 (the “classic” 4th order Runge–Kutta integrator), and DDIM (Song et al., 2021a). In addition we compared to a Gaussian sampler with variance equal to the lower bound given by Ho et 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 3 and identified DDIM as the best integrator in terms of resulting sample quality.

heading: D EXPRESSION OF DDIM IN ANGULAR PARAMETERIZATION
text: We can simplify the DDIM update rule by expressing it in terms of φt = arctan(σt/αt), rather than in terms of time t or log-SNR λt, as we show here.
Given our definition of φ, and assuming a variance preserving diffusion process, we have αφ = cos(φ), σφ = sin(φ), and hence zφ = cos(φ)x+ sin(φ) . We can now define the velocity of zφ as
vφ ≡ dzφ dφ = d cos(φ) dφ x+ d sin(φ) dφ = cos(φ) − sin(φ)x. (26)
Rearranging ,x,v, we then get
sin(φ)x = cos(φ) − vφ (27)
= cos(φ)
sin(φ) (z− cos(φ)x)− vφ (28)
sin2(φ)x = cos(φ)z− cos2(φ)x− sin(φ)vφ (29) (sin2(φ) + cos2(φ))x = x = cos(φ)z− sin(φ)vφ, (30)
and similarly we get = sin(φ)zφ + cos(φ)vφ.
Furthermore, we define the predicted velocity as
v̂θ(zφ) ≡ cos(φ)̂θ(zφ)− sin(φ)x̂θ(zφ), (31)
where ̂θ(zφ) = (zφ − cos(φ)x̂θ(zφ))/ sin(φ). Rewriting the DDIM update rule in the introduced terms then gives
zφs = cos(φs)x̂θ(zφt) + sin(φs)̂θ(zφt) (32) = cos(φs)(cos(φt)zφt − sin(φt)v̂θ(zφt)) + sin(φs)(sin(φt)zφt + cos(φt)v̂θ(zφt)) (33) = [cos(φs) cos(φt)− sin(φs) sin(φt)]zφt + [sin(φs) cos(φt)− cos(φs) sin(φt)]v̂θ(zφt).
(34)
Finally, we use the trigonometric identities
cos(φs) sin(φt)− sin(φs) cos(φt) = cos(φs − φt) (35) sin(φs) cos(φt)− cos(φs) sin(φt) = sin(φs − φt), (36)
to find that
zφs = cos(φs − φt)zφt + sin(φs − φt)v̂θ(zφt). (37)
or equivalently
zφt−δ = cos(δ)zφt − sin(δ)v̂θ(zφt). (38)
Viewed from this perspective, DDIM thus evolves zφs by moving it on a circle in the (zφt , v̂φt) basis, along the−v̂φt direction. The relationship between zφt ,vt, αt, σt,x, is visualized in Figure 5.

heading: E SETTINGS USED IN EXPERIMENTS
text: Our model architectures closely follow those described by Dhariwal & Nichol (2021). For 64× 64 ImageNet we use their model exactly, with 192 channels at the highest resolution. All other models are slight variations with different hyperparameters.
For CIFAR-10 we use an architecture with a fixed number of channels at all resolutions of 256. The model consists of a UNet that internally downsamples the data twice, to 16 × 16 and to 8 × 8. At each resolution we apply 3 residual blocks, like described by Dhariwal & Nichol (2021). We use single-headed attention, and only apply this at the 16× 16 and 8× 8 resolutions. We use dropout of 0.2 when training the original model. No dropout is used during distillation.
For LSUN we use a model similar to that for ImageNet, but with a reduced number of 128 channels at the 64 × 64 resolution. Compared to ImageNet we have an additional level in the UNet, corresponding to the input resolution of 128 × 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-10 we take the output of the model to represent a prediction of x directly, as discussed in Section 4. For the other data sets we used the combined prediction of (x, ) 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 a constant of 0.001. We clip the norm of gradients to a global norm of 1 before calculating parameter updates. For CIFAR-10 we train for 800k parameter updates, for ImageNet we use 550k updates, and for LSUN we use 400k updates. During distillation we train for 50k updates per iteration, except for the distillation to 2 and 1 sampling steps, for which we use 100k updates. We linearly anneal the learning rate from 10−4 to zero during each iteration.
We use a batch size of 128 for CIFAR-10 and 2048 for the other data sets. We run our experiments on TPUv4, using 8 TPU chips for CIFAR-10, and 64 chips for the other data sets. The total time required to first train and then distill a model varies from about a day for CIFAR-10, to about 5 days for ImageNet.

heading: F STOCHASTIC SAMPLING WITH DISTILLED MODELS
text: 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 a distilled model for 64x64 ImageNet using the optimized stochastic sampler also used in Section 5.2. The results are presented in Figure 6.

heading: G DERIVATION OF THE DISTILLATION TARGET
text: The key difference between our progressive distillation algorithm proposed in Section 3 and the standard diffusion training procedure is in how we determine the target value for our denoising model. In standard diffusion 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 t′ = t− 0.5/N and t′′ = t− 1/N , when training a student with N sampling steps, we have that the teacher model samples the next set of noisy data zt′′ given the current noisy data zt by taking two steps of DDIM. The student tries to sample the same value in only one step of DDIM. Denoting the student denoising prediction by x̃, and its one-step sample by z̃t′′ , application of the DDIM sampler (see equation 8), gives:
z̃t′′ = αt′′ x̃+ σt′′
σt (zt − αtx̃). (39)
In order for the student sampler to match the teacher sampler, we must set z̃t′′ equal to zt′′ . This gives
z̃t′′ = αt′′ x̃+ σt′′
σt (zt − αtx̃) = zt′′ (40)
= ( αt′′ − σt′′
σt αt
) x̃+ σt′′
σt zt = zt′′ (41)(
αt′′ − σt′′
σt αt
) x̃ = zt′′ − σt′′
σt zt (42)
x̃ = zt′′ − σt′′σt zt αt′′ − σt′′σt αt
(43)
In other words, if our student denoising model exactly predicts x̃ as defined in equation 43 above, then the one-step student sample z̃t′′ is identical to the two-step teacher sample zt′′ . In order to have our student model approximate this ideal outcome, we thus train it to predict x̃ from zt 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 a one-step student model is unique to deterministic samplers like DDIM: the composition of two standard stochastic DDPM sampling steps (Equation 5) forms a non-Gaussian distribution that falls outside the family of Gaussian distributions that can be modelled by a single DDPM student step: A multi-step stochastic DDPM sampler can thus not be distilled into a 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 a neural net, which is why the student is able to accurately match the teacher.
Finally, note that we do lose something during the progressive distillation process: while the original model was trained to denoise zt for any given continuous time t, 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.

heading: H ADDITIONAL RANDOM SAMPLES
text: 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 1 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 a 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)).

heading: I ABLATION WITH FASTER DISTILLATION SCHEDULES
text: In order to further reduce the computational requirements for our progressive distillation approach, we perform an ablation study on CIFAR-10, 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 5 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 ≥ 4 sampling steps. When taking very few sampling steps, the loss in performance becomes more pronounced when training the student for only a short time.
In addition to just decreasing the number of parameter updates, we also experiment with a schedule where we train each student on 4 times fewer sampling steps than its teacher, rather than the 2 times we propose in Section 3. Here the denoising target is still derived from taking 2 DDIM steps with the teacher model as 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 we show the results achieved with a faster schedule for the ImageNet and LSUN datasets. Here also, we achieve excellent results with a faster distillation schedule.

Review summaries:
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. 

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.

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.

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.  Enabling 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.