Original: P ROGRESSIVE D ISTILLATION FOR F AST S AMPLING OF D IFFUSION M ODELS
Revised: P ROGRESSIVE D ISTILLATION FOR F AST S AMPLING OF D IFFUSION M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation.
Revised: Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation.

Original: A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations.
Revised: A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time.
Revised: We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time.

Original: 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 few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps.
Revised: 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 few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps.

Original: 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 efﬁcient solution for generative modeling using diffusion at both train and test time.
Revised: 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 efﬁcient solution for generative modeling using diffusion at both train and test time.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: These models have achieved ImageNet generation results outperforming BigGAN-deep and VQ-VAE-2 in terms of FID score and classiﬁcation 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).
Revised: These models have achieved ImageNet generation results outperforming BigGAN-deep and VQ-VAE-2 in terms of FID score and classiﬁcation 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).

Original: Diffusion models have also succeeded in image super-resolution (Saharia et al., 2021;
Revised: Diffusion models have also succeeded in image super-resolution (Saharia et al., 2021;

Original: 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
Revised: 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

Original: (Niu et al., 2020), and text generation (Hoogeboom et al., 2021;
Revised: (Niu et al., 2020), and text generation (Hoogeboom et al., 2021;

Original: Austin et al., 2021).
Revised: Austin et al., 2021).

Original: A major barrier remains to practical adoption of diffusion models: sampling speed.
Revised: A major barrier remains to practical adoption of diffusion models: sampling speed.

Original: 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 classiﬁer (Dhariwal & Nichol, 2021), the situation is substantially different in settings in which there is less conditioning information.
Revised: 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 classiﬁer (Dhariwal & Nichol, 2021), the situation is substantially different in settings in which there is less conditioning information available.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: We present a procedure to distill the behavior of a N -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.
Revised: We present a procedure to distill the behavior of a N -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.

Original: 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.
Revised: 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.

Original: [Figure]
Revised: [Figure]

Original:  B ACKGROUND ON DIFFUSION MODELS
Revised:  B ACKGROUND ON DIFFUSION MODELS

Original: We consider diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019;
Revised: We consider diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019;

Original: Ho et al., 2020) speciﬁed in continuous time (Tzen & Raginsky, 2019a; Song et al., 2021c; Chen et al., 2021; Kingma et al., 2021).
Revised: Ho et al., 2020) speciﬁed in continuous time (Tzen & Raginsky, 2019a; Song et al., 2021c; Chen et al., 2021; Kingma et al., 2021).

Original: We use x ∼ p ( x ) to denote training data.
Revised: We use x ∼ p ( x ) to denote training data.

Original: A diffusion model has latent variables
Revised: A diffusion model has latent variables

Original: z = { z t | t ∈ [0 , 1] } and is speciﬁed by a noise schedule comprising differentiable functions α t , σ t such that λ t = log
Revised: z = { z t | t ∈ [0 , 1] } and is speciﬁed by a noise schedule comprising differentiable functions α t , σ t such that λ t = log

Original: [ α 2 t /σ 2 t ] , the log signal-to-noise-ratio, decreases monotonically with t .
Revised: [ α 2 t /σ 2 t ] , the log signal-to-noise-ratio, decreases monotonically with t .

Original: These ingredients deﬁne the forward process q ( z | x ) , a Gaussian process satisfying the following
Revised: These ingredients deﬁne the forward process q ( z | x ) , a Gaussian process satisfying the following

Original: Markovian structure:
Revised: Markovian structure:

Original: [Equation]
Revised: [Equation]

Original: where 0 ≤ s < t ≤ 1 and σ 2 t | s = (1 − e λ t − λ s ) σ 2 t .
Revised: where 0 ≤ s < t ≤ 1 and σ 2 t | s = (1 − e λ t − λ s ) σ 2 t .

Original: The role of function approximation in the diffusion model is to denoise z t ∼ q ( z t | x ) into an estimate ˆ x θ ( z t )
Revised: The role of function approximation in the diffusion model is to denoise z t ∼ q ( z t | x ) into an estimate ˆ x θ ( z t )

Original: ≈ x (the function approximator also receives λ t as an input, but we omit this to keep our notation clean).
Revised: ≈ x (the function approximator also receives λ t as an input, but we omit this to keep our notation clean).

Original: We train this denoising model ˆ x θ using a weighted mean squared error loss
Revised: We train this denoising model ˆ x θ using a weighted mean squared error loss

Original: [Equation]
Revised: [Equation]

Original: over uniformly sampled times t ∈ [0 , 1] .
Revised: over uniformly sampled times t ∈ [0 , 1] .

Original: This loss can be justiﬁed 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).
Revised: This loss can be justiﬁed 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).

Original: We will discuss particular choices of weighting function w ( λ t ) later on.
Revised: We will discuss particular choices of weighting function w ( λ t ) later on.

Original: Sampling from a trained model can be performed in several ways.
Revised: Sampling from a trained model can be performed in several ways.

Original: The most straightforward way is discrete time ancestral sampling (Ho et al., 2020).
Revised: The most straightforward way is discrete time ancestral sampling (Ho et al., 2020).

Original: To deﬁne this sampler, ﬁrst note that the forward process can be described in reverse as q ( z s | z t , x ) =
Revised: To deﬁne this sampler, ﬁrst note that the forward process can be described in reverse as q ( z s | z t , x ) =

Original: N ( z s ; ˜ µ s | t ( z t , x ) ,
Revised: N ( z s ; ˜ µ s | t ( z t , x ) ,

Original: ˜ σ 2 s | t I ) (noting s < t ),
Revised: ˜ σ 2 s | t I ) (noting s < t ),

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: We use this reversed description of the forward process to deﬁne the ancestral sampler.
Revised: We use this reversed description of the forward process to deﬁne the ancestral sampler.

Original: Starting at z 1 ∼ N ( 0 , I ) , the ancestral sampler follows the rule
Revised: Starting at z 1 ∼ N ( 0 , I ) , the ancestral sampler follows the rule

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where (cid:15) is standard Gaussian noise, and γ is a hyperparameter that controls how much noise is added during sampling, following Nichol & Dhariwal (2021).
Revised: where (cid:15) is standard Gaussian noise, and γ is a hyperparameter that controls how much noise is added during sampling, following Nichol & Dhariwal (2021).

Original: Alternatively, Song et al.
Revised: Alternatively, Song et al.

Original: (2021c) show that our denoising model ˆ x θ ( z t ) can be used to deterministically map noise z 1 ∼ N ( 0 , I ) to samples x by numerically solving the probability ﬂow ODE :
Revised: (2021c) show that our denoising model ˆ x θ ( z t ) can be used to deterministically map noise z 1 ∼ N ( 0 , I ) to samples x by numerically solving the probability ﬂow ODE :

Original: [Equation]
Revised: [Equation]

Original: where ∇
Revised: where ∇

Original: [Equation]
Revised: [Equation]

Original: Following Kingma et al.
Revised: Following Kingma et al.

Original: (2021), we have f ( z t , t
Revised: (2021), we have f ( z t , t

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: dt z t and g
Revised: dt z t and g

Original: [Equation]
Revised: [Equation]

Original: σ 2 t .
Revised: σ 2 t .

Original: Since ˆ x θ ( z t ) 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 ﬂow (Grathwohl et al., 2018).
Revised: Since ˆ x θ ( z t ) 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 ﬂow (Grathwohl et al., 2018).

Original: Solving the ODE in Equation 6 numerically can be done with standard methods like the Euler rule or the Runge-Kutta method.
Revised: Solving the ODE in Equation 6 numerically can be done with standard methods like the Euler rule or the Runge-Kutta method.

Original: The DDIM sampler proposed by Song et al.
Revised: The DDIM sampler proposed by Song et al.

Original: (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.
Revised: (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.

Original: The update rule speciﬁed by DDIM is
Revised: The update rule speciﬁed by DDIM is

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: and in practice this rule performs better than the aforementioned standard ODE integration rules in our case, as we show in Appendix C.
Revised: and in practice this rule performs better than the aforementioned standard ODE integration rules in our case, as we show in Appendix C.

Original: If ˆ x θ ( z t ) satisﬁes mild smoothness conditions, the error introduced by numerical integration of the probability ﬂow ODE is guaranteed to vanish as the number of integration steps grows inﬁnitely large, i.e. N → ∞ .
Revised: If ˆ x θ ( z t ) satisﬁes mild smoothness conditions, the error introduced by numerical integration of the probability ﬂow ODE is guaranteed to vanish as the number of integration steps grows inﬁnitely large, i.e. N → ∞ .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Here, we therefore propose a method to distill these accurate, but slow, ODE integrators into much faster models that are still very accurate.
Revised: Here, we therefore propose a method to distill these accurate, but slow, ODE integrators into much faster models that are still very accurate.

Original: This idea is visualized in Figure 1, and described in detail in the next section.
Revised: This idea is visualized in Figure 1, and described in detail in the next section.

Original:  P ROGRESSIVE DISTILLATION
Revised:  P ROGRESSIVE DISTILLATION

Original: To make diffusion models more efﬁcient 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.
Revised: To make diffusion models more efﬁcient 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.

Original: Our implementation of progressive distillation stays very close to the implementation for training the original diffusion model, as described by e.g.
Revised: Our implementation of progressive distillation stays very close to the implementation for training the original diffusion model, as described by e.g.

Original: Ho et al.
Revised: Ho et al.

Original: 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.
Revised: 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.

Original: We start the progressive distillation procedure with a teacher diffusion model that is obtained by training in the standard way.
Revised: We start the progressive distillation procedure with a teacher diffusion model that is obtained by training in the standard way.

Original: 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 deﬁnition.
Revised: 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 deﬁnition.

Original: 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 t .
Revised: 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 t .

Original: 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
Revised: 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

Original: that makes a single student DDIM step match 2 teacher DDIM steps.
Revised: that makes a single student DDIM step match 2 teacher DDIM steps.

Original: We calculate this target value by running 2 DDIM sampling steps using the teacher, starting from z t and ending at z t − 1 /N , with N being the number of student sampling steps.
Revised: We calculate this target value by running 2 DDIM sampling steps using the teacher, starting from z t and ending at z t − 1 /N , with N being the number of student sampling steps.

Original: By inverting a single step of DDIM, we then calculate the value the student model would need to predict in order to move from z t to z t − 1 /N in a single step, as we show in detail in Appendix G.
Revised: By inverting a single step of DDIM, we then calculate the value the student model would need to predict in order to move from z t to z t − 1 /N in a single step, as we show in detail in Appendix G.

Original: The resulting target value ˜ x ( z t ) is fully determined given the teacher model and starting point z t , which allows the student model to make a sharp prediction when evaluated at z t .
Revised: The resulting target value ˜ x ( z t ) is fully determined given the teacher model and starting point z t , which allows the student model to make a sharp prediction when evaluated at z t .

Original: In contrast, the original data point x is not fully determined given z t , since multiple different data points x can produce the same noisy data z t : this means that the original denoising model is predicting a weighted average of possible x values, which produces a blurry prediction.
Revised: In contrast, the original data point x is not fully determined given z t , since multiple different data points x can produce the same noisy data z t : this means that the original denoising model is predicting a weighted average of possible x values, which produces a blurry prediction.

Original: By making sharper predictions, the student model can make faster progress during sampling.
Revised: By making sharper predictions, the student model can make faster progress during sampling.

Original: 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.
Revised: 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.

Original: 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 z 1 ∼ N ( 0 , I ) that is used at test time.
Revised: 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 z 1 ∼ N ( 0 , I ) that is used at test time.

Original: We found this to work slightly better than starting from a non-zero signalto-noise ratio as used by e.g. Ho et al.
Revised: We found this to work slightly better than starting from a non-zero signalto-noise ratio as used by e.g. Ho et al.

Original: (2020), both for training the original model as well as when performing progressive distillation.
Revised: (2020), both for training the original model as well as when performing progressive distillation.

Original: Require:
Revised: Require:

Original: Model ˆ x θ ( z t ) to be trained
Revised: Model ˆ x θ ( z t ) to be trained

Original: Require: Data set D
Revised: Require: Data set D

Original: Require: Loss weight function w () while not converged do x ∼ D (cid:46)
Revised: Require: Loss weight function w () while not converged do x ∼ D (cid:46)

Original: Sample data t ∼ U [0 , 1] (cid:46)
Revised: Sample data t ∼ U [0 , 1] (cid:46)

Original: Sample time (cid:15) ∼ N (0 , I ) (cid:46)
Revised: Sample time (cid:15) ∼ N (0 , I ) (cid:46)

Original: Sample noise z t = α t x + σ t (cid:15) (cid:46)
Revised: Sample noise z t = α t x + σ t (cid:15) (cid:46)

Original: Add noise to data ˜ x
Revised: Add noise to data ˜ x

Original: = x (cid:46) Clean data is target for ˆ x λ t
Revised: = x (cid:46) Clean data is target for ˆ x λ t

Original: = log
Revised: = log

Original: [ α 2 t /σ 2 t ] (cid:46) log-SNR
Revised: [ α 2 t /σ 2 t ] (cid:46) log-SNR

Original: L θ = w ( λ t ) (cid:107) ˜ x − ˆ x θ ( z t ) (cid:107)(cid:46) Loss θ ← θ − γ ∇ θ L θ (cid:46) Optimization end while
Revised: L θ = w ( λ t ) (cid:107) ˜ x − ˆ x θ ( z t ) (cid:107)(cid:46) Loss θ ← θ − γ ∇ θ L θ (cid:46) Optimization end while

Original: Require:
Revised: Require:

Original: Trained teacher model ˆ x η ( z t )
Revised: Trained teacher model ˆ x η ( z t )

Original: Require: Data set D
Revised: Require: Data set D

Original: Require: Loss weight function w ()
Revised: Require: Loss weight function w ()

Original: Require:
Revised: Require:

Original: Student sampling steps N for K iterations do θ ← η (cid:46)
Revised: Student sampling steps N for K iterations do θ ← η (cid:46)

Original: Init student from teacher while not converged do x ∼ D
Revised: Init student from teacher while not converged do x ∼ D

Original: [Equation]
Revised: [Equation]

Original: steps of DDIM with teacher
Revised: steps of DDIM with teacher

Original: [Equation]
Revised: [Equation]

Original: Teacher ˆ x target
Revised: Teacher ˆ x target

Original: [Equation]
Revised: [Equation]

Original: end while η ← θ (cid:46) Student becomes next teacher
Revised: end while η ← θ (cid:46) Student becomes next teacher

Original: N ← N/ 2 (cid:46) Halve number of sampling steps end for
Revised: N ← N/ 2 (cid:46) Halve number of sampling steps end for

Original:  D IFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS
Revised:  D IFFUSION MODEL PARAMETERIZATION AND TRAINING LOSS

Original: In this section, we discuss how to parameterize the denoising model ˆ x θ , and how to specify the reconstruction loss weight w ( λ t ) .
Revised: In this section, we discuss how to parameterize the denoising model ˆ x θ , and how to specify the reconstruction loss weight w ( λ t ) .

Original: We assume a standard variance-preserving diffusion process for which σ 2 t = 1 − α 2 t .
Revised: We assume a standard variance-preserving diffusion process for which σ 2 t = 1 − α 2 t .

Original: This is without loss of generalization, as shown by (Kingma et al., 2021, appendix G): different speciﬁcations of the diffusion process, such as the variance-exploding speciﬁcation, can be considered equivalent to this speciﬁcation, up to rescaling of the noisy latents z t .
Revised: This is without loss of generalization, as shown by (Kingma et al., 2021, appendix G): different speciﬁcations of the diffusion process, such as the variance-exploding speciﬁcation, can be considered equivalent to this speciﬁcation, up to rescaling of the noisy latents z t .

Original: We use a cosine schedule α t = cos(0 .
Revised: We use a cosine schedule α t = cos(0 .

Original: 5 πt ) , similar to that introduced by Nichol & Dhariwal (2021).
Revised: 5 πt ) , similar to that introduced by Nichol & Dhariwal (2021).

Original: Ho et al.
Revised: Ho et al.

Original: (2020) and much of the following work
Revised: (2020) and much of the following work

Original: choose to parameterize the denoising model through directly predicting (cid:15) with a neural network ˆ (cid:15) θ ( z t ) , which implicitly sets ˆ x θ ( z t )
Revised: choose to parameterize the denoising model through directly predicting (cid:15) with a neural network ˆ (cid:15) θ ( z t ) , which implicitly sets ˆ x θ ( z t )

Original: =
Revised: =

Original: 1 α t ( z t
Revised: 1 α t ( z t

Original: [Equation]
Revised: [Equation]

Original: σ t ˆ (cid:15) θ ( z t )) .
Revised: σ t ˆ (cid:15) θ ( z t )) .

Original: In this case, the training loss is also usually deﬁned as mean squared error in the (cid:15) -space:
Revised: In this case, the training loss is also usually deﬁned as mean squared error in the (cid:15) -space:

Original: [Equation]
Revised: [Equation]

Original: which can thus equivalently be seen as a weighted reconstruction loss in x -space, where the weighting function is given by w ( λ t )
Revised: which can thus equivalently be seen as a weighted reconstruction loss in x -space, where the weighting function is given by w ( λ t )

Original: = exp( λ t ) , for log signal-to-noise ratio λ t = log[ α 2 t /σ 2 t ] .
Revised: = exp( λ t ) , for log signal-to-noise ratio λ t = log[ α 2 t /σ 2 t ] .

Original: Although this standard speciﬁcation 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
Revised: Although this standard speciﬁcation 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

Original: α 2 t /σ 2 t , but as distillation progresses we increasingly evaluate at lower and lower signal-to-noise ratios.
Revised: α 2 t /σ 2 t , but as distillation progresses we increasingly evaluate at lower and lower signal-to-noise ratios.

Original: As the signal-to-noise ratio goes to zero, the effect of small changes in the neural network output ˆ (cid:15) θ ( z t ) on the implied prediction in x -space is increasingly ampliﬁed, since ˆ x θ ( z t )
Revised: As the signal-to-noise ratio goes to zero, the effect of small changes in the neural network output ˆ (cid:15) θ ( z t ) on the implied prediction in x -space is increasingly ampliﬁed, since ˆ x θ ( z t )

Original: =
Revised: =

Original: 1 α t ( z t − σ t ˆ (cid:15) θ
Revised: 1 α t ( z t − σ t ˆ (cid:15) θ

Original: ( z t )) divides by α t → 0 .
Revised: ( z t )) divides by α t → 0 .

Original: This is not much of a problem when taking many steps, since the effect of early missteps is limited by clipping of the z t iterates, and later updates can correct any mistakes, but it becomes increasingly important as we decrease the number of sampling steps.
Revised: This is not much of a problem when taking many steps, since the effect of early missteps is limited by clipping of the z t iterates, and later updates can correct any mistakes, but it becomes increasingly important as we decrease the number of sampling steps.

Original: Eventually, if we distill all the way down to a single sampling step, the input to the model is only pure noise (cid:15) , which corresponds to a signal-to-noise ratio of zero, i.e. α t = 0 , σ t = 1 .
Revised: Eventually, if we distill all the way down to a single sampling step, the input to the model is only pure noise (cid:15) , which corresponds to a signal-to-noise ratio of zero, i.e. α t = 0 , σ t = 1 .

Original: At this extreme, the link between (cid:15) -prediction and x -prediction breaks down completely: observed data z t = (cid:15) is no longer informative of x and predictions ˆ (cid:15) θ ( z t ) no longer implicitly predict x .
Revised: At this extreme, the link between (cid:15) -prediction and x -prediction breaks down completely: observed data z t = (cid:15) is no longer informative of x and predictions ˆ (cid:15) θ ( z t ) no longer implicitly predict x .

Original: 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.
Revised: 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.

Original: For distillation to work, we thus need to parameterize the diffusion model in a way for which the implied prediction ˆ x θ ( z t ) remains stable as λ t = log
Revised: For distillation to work, we thus need to parameterize the diffusion model in a way for which the implied prediction ˆ x θ ( z t ) remains stable as λ t = log

Original: [ α 2 t /σ 2 t ] varies.
Revised: [ α 2 t /σ 2 t ] varies.

Original: We tried the following options, and found all to work well with progressive distillation:
Revised: We tried the following options, and found all to work well with progressive distillation:

Original: • Predicting x directly.
Revised: • Predicting x directly.

Original: Predicting both x and (cid:15) , via separate output channels { ˜ x θ ( z t ) , ˜ (cid:15) θ ( z t ) } of the neural network, and then merging the predictions via ˆ x = σ 2 t ˜ x θ ( z t ) + α t ( z t − σ t ˜ (cid:15) θ ( z t )) , thus smoothly interpolating between predicting x directly and predicting via (cid:15) .
Revised: Predicting both x and (cid:15) , via separate output channels { ˜ x θ ( z t ) , ˜ (cid:15) θ ( z t ) } of the neural network, and then merging the predictions via ˆ x = σ 2 t ˜ x θ ( z t ) + α t ( z t − σ t ˜ (cid:15) θ ( z t )) , thus smoothly interpolating between predicting x directly and predicting via (cid:15) .

Original: • Predicting v ≡ α t (cid:15) − σ t x , which gives ˆ x = α t z t − σ t ˆ v θ ( z t ) , as we show in Appendix D.
Revised: • Predicting v ≡ α t (cid:15) − σ t x , which gives ˆ x = α t z t − σ t ˆ v θ ( z t ) , as we show in Appendix D.

Original: In Section 5.1 we test all three parameterizations on training an original diffusion model (no distillation), and ﬁnd them to work well there also.
Revised: In Section 5.1 we test all three parameterizations on training an original diffusion model (no distillation), and ﬁnd them to work well there also.

Original: In addition to determining an appropriate parameterization, we also need to decide on a reconstruction loss weighting w ( λ t ) .
Revised: In addition to determining an appropriate parameterization, we also need to decide on a reconstruction loss weighting w ( λ t ) .

Original: The setup of Ho et al.
Revised: The setup of Ho et al.

Original: (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.
Revised: (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.

Original: We consider two alternative training loss weightings:
Revised: We consider two alternative training loss weightings:

Original: • L θ = max ( (cid:107) x − ˆ x t (cid:107) 22 , (cid:107) (cid:15) − ˆ (cid:15) t (cid:107) 22 ) =
Revised: • L θ = max ( (cid:107) x − ˆ x t (cid:107) 22 , (cid:107) (cid:15) − ˆ (cid:15) t (cid:107) 22 ) =

Original: max ( α 2 t σ 2 t , 1) (cid:107) x − ˆ x t (cid:107) 22 ; ‘truncated SNR’ weighting.
Revised: max ( α 2 t σ 2 t , 1) (cid:107) x − ˆ x t (cid:107) 22 ; ‘truncated SNR’ weighting.

Original: • L θ = (cid:107) v t − ˆ v t (cid:107) 22 =
Revised: • L θ = (cid:107) v t − ˆ v t (cid:107) 22 =

Original: (1 + α 2 t σ 2 t ) (cid:107) x − ˆ x t (cid:107) 22 ; ‘SNR+1’ weighting.
Revised: (1 + α 2 t σ 2 t ) (cid:107) x − ˆ x t (cid:107) 22 ; ‘SNR+1’ weighting.

Original: We examine both choices in our ablation study in Section 5.1, and ﬁnd both to be good choices for training diffusion models.
Revised: We examine both choices in our ablation study in Section 5.1, and ﬁnd both to be good choices for training diffusion models.

Original: 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.
Revised: 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.

Original: Our results are for a cosine schedule α t = cos(0 .
Revised: Our results are for a cosine schedule α t = cos(0 .

Original: 5 πt ) , where time is sampled uniformly from [0 , 1] .
Revised: 5 πt ) , where time is sampled uniformly from [0 , 1] .

Original: In Figure 2 we visualize the resulting loss weightings, both including and excluding the effect of the cosine schedule.
Revised: In Figure 2 we visualize the resulting loss weightings, both including and excluding the effect of the cosine schedule.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: 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.
Revised: 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.

Original: We consider various image generation benchmarks, with resolution varying from 32 × 32 to 128 × 128 .
Revised: We consider various image generation benchmarks, with resolution varying from 32 × 32 to 128 × 128 .

Original: All experiments use the cosine schedule α t = cos(0 .
Revised: All experiments use the cosine schedule α t = cos(0 .

Original: 5 πt ) , and all models use a U-Net architecture similar to that introduced by Ho et al.
Revised: 5 πt ) , and all models use a U-Net architecture similar to that introduced by Ho et al.

Original: (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.
Revised: (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.

Original: (2021c).
Revised: (2021c).

Original: Our training setup closely matches the open source code by Ho et al.
Revised: Our training setup closely matches the open source code by Ho et al.

Original: Exact details are given in
Revised: Exact details are given in

Original: Appendix E. We plan to open source the code soon.
Revised: Appendix E. We plan to open source the code soon.

Original: 5. M ODEL PARAMETERIZATION AND TRAINING LOSS
Revised: 5. M ODEL PARAMETERIZATION AND TRAINING LOSS

Original: As explained in Section 4, the standard method of having our model predict (cid:15) , and minimizing mean squared error in the (cid:15) -space
Revised: As explained in Section 4, the standard method of having our model predict (cid:15) , and minimizing mean squared error in the (cid:15) -space

Original: (Ho et al., 2020), is not appropriate for use with progressive distillation.
Revised: (Ho et al., 2020), is not appropriate for use with progressive distillation.

Original: −−−    −  log SNR l og w e i gh t ( e x c l ud
Revised: −−−    −  log SNR l og w e i gh t ( e x c l ud

Original: i ng s c
Revised: i ng s c

Original: h e du l e )
Revised: h e du l e )

Original: SNR weight truncated SNR
Revised: SNR weight truncated SNR

Original: SNR+1 weight −−−     0 .0 .0 .0 .log
Revised: SNR+1 weight −−−     0 .0 .0 .0 .log

Original: SNR w e i gh t ( i n c l ud
Revised: SNR w e i gh t ( i n c l ud

Original: i ng s c
Revised: i ng s c

Original: h e du l e )
Revised: h e du l e )

Original: SNR truncated SNR
Revised: SNR truncated SNR

Original: SNR+
Revised: SNR+

Original: Network Output
Revised: Network Output

Original: Loss Weighting
Revised: Loss Weighting

Original: Stochastic sampler
Revised: Stochastic sampler

Original: DDIM sampler ( x , (cid:15) ) combined
Revised: DDIM sampler ( x , (cid:15) ) combined

Original: SNR 2.54/9.
Revised: SNR 2.54/9.

Original: Truncated SNR 2.47/9.
Revised: Truncated SNR 2.47/9.

Original: SNR+ 2.52/9.
Revised: SNR+ 2.52/9.

Original: x
Revised: x

Original: SNR 2.65/9.
Revised: SNR 2.65/9.

Original: Truncated SNR 2.53/ 9.
Revised: Truncated SNR 2.53/ 9.

Original: SNR+ 2.56/9.
Revised: SNR+ 2.56/9.

Original: (cid:15)
Revised: (cid:15)

Original: SNR 2.59/9.
Revised: SNR 2.59/9.

Original: Truncated SNR
Revised: Truncated SNR

Original: N/A N/A
Revised: N/A N/A

Original: SNR+ 2.56/9.
Revised: SNR+ 2.56/9.

Original: [Table]
Revised: [Table]

Original: 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.
Revised: 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.

Original: Here, we perform a complete ablation experiment of all parameterizations and loss weightings considered in Section 4.
Revised: Here, we perform a complete ablation experiment of all parameterizations and loss weightings considered in Section 4.

Original: For computational efﬁciency, and for comparisons to established methods in the literature, we use unconditional CIFAR-10 as the bench- to
Revised: For computational efﬁciency, and for comparisons to established methods in the literature, we use unconditional CIFAR-10 as the bench- to

Original: Table 1 lists the results of the ablation study.
Revised: Table 1 lists the results of the ablation study.

Original: Overall results are fairly close across different parameterizations and loss weights.
Revised: Overall results are fairly close across different parameterizations and loss weights.

Original: All proposed stable model speciﬁcations achieve excellent performance, with the exception of the combination of outputting (cid:15) with the neural network and weighting the loss with the truncated SNR, which we ﬁnd to be unstable.
Revised: All proposed stable model speciﬁcations achieve excellent performance, with the exception of the combination of outputting (cid:15) with the neural network and weighting the loss with the truncated SNR, which we ﬁnd to be unstable.

Original: Both predicting x directly, as well as predicting v , or the combination ( (cid:15), x ) , could thus be recommended for speciﬁcation of diffusion models.
Revised: Both predicting x directly, as well as predicting v , or the combination ( (cid:15), x ) , could thus be recommended for speciﬁcation of diffusion models.

Original: 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.
Revised: 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.

Original: 5. P ROGRESSIVE DISTILLATION
Revised: 5. P ROGRESSIVE DISTILLATION

Original: We evaluate our proposed progressive distillation algorithm on 4 data sets: CIFAR-10, 64 ×downsampled ImageNet, 128 × 128 LSUN bedrooms, and 128 × 128 LSUN Church-Outdoor.
Revised: We evaluate our proposed progressive distillation algorithm on 4 data sets: CIFAR-10, 64 ×downsampled ImageNet, 128 × 128 LSUN bedrooms, and 128 × 128 LSUN Church-Outdoor.

Original: For each data set we start by training a baseline model, after which we start the progressive distillation procedure.
Revised: For each data set we start by training a baseline model, after which we start the progressive distillation procedure.

Original: For CIFAR-10 we start progressive distillation from a teacher model taking 8192 steps.
Revised: For CIFAR-10 we start progressive distillation from a teacher model taking 8192 steps.

Original: For the bigger data sets we start at 1024 steps.
Revised: For the bigger data sets we start at 1024 steps.

Original: At every iteration of distillation we train forthousand parameter updates, except for the distillation to 2 and 1 sampling steps, for which we use 100 thousand updates.
Revised: At every iteration of distillation we train forthousand parameter updates, except for the distillation to 2 and 1 sampling steps, for which we use 100 thousand updates.

Original: For each model, the computational cost of progressive distillation tosampling steps is comparable or less than for training the original model.We report FID results obtained after each iteration of the algorithm.
Revised: 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.

Revised: In Appendix I we show that this computational cost can be reduce much further still, at a small cost in performance.

Original: In Figure 4 we plot the resulting FID scores (Heusel et al., 2017) obtained for each number of sampling steps.
Revised: In Figure 4 we plot the resulting FID scores (Heusel et al., 2017) obtained for each number of sampling steps.

Original: We compare against the undistilled DDIM sampler, as well as to a highly optimized stochastic baseline sampler.
Revised: We compare against the undistilled DDIM sampler, as well as to a highly optimized stochastic baseline sampler.

Original: For all four data sets, progressive distillation produces near optimal results up to 4 or 8 sampling steps.
Revised: For all four data sets, progressive distillation produces near optimal results up to 4 or 8 sampling steps.

Original: At 2 or 1 sampling steps, the sample quality degrades relatively more quickly.
Revised: At 2 or 1 sampling steps, the sample quality degrades relatively more quickly.

Original: In contrast, the quality of the DDIM and stochastic samplers degrades very sharply after reducing the number of sampling steps below 128.
Revised: In contrast, the quality of the DDIM and stochastic samplers degrades very sharply after reducing the number of sampling steps below 128.

Original: Overall, we conclude that progressive distillation is thus an attractive solution for computational budgets that allow less than or equal to 128 sampling steps.
Revised: Overall, we conclude that progressive distillation is thus an attractive solution for computational budgets that allow less than or equal to 128 sampling steps.

Original: 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 ﬁnd that it achieves performance that falls in between the distilled DDIM sampler and the undistilled stochastic sampler.
Revised: 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 ﬁnd that it achieves performance that falls in between the distilled DDIM sampler and the undistilled stochastic sampler.

Original: 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.
Revised: 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.

Original: Figure 3 shows some random samples from our model obtained at different phases of the distillation process.
Revised: Figure 3 shows some random samples from our model obtained at different phases of the distillation process.

Original: Additional samples are provided in Appendix H
Revised: Additional samples are provided in Appendix H

Original: [Figure]
Revised: [Figure]

Original:  R ELATED WORK ON FAST SAMPLING
Revised:  R ELATED WORK ON FAST SAMPLING

Original: Our proposed method is closest to the work of Luhman & Luhman (2021), who perform distillation of DDIM teacher models into one-step student models.
Revised: Our proposed method is closest to the work of Luhman & Luhman (2021), who perform distillation of DDIM teacher models into one-step student models.

Original: A possible downside of their method is that it requires constructing a large data set by running the original model at its full number ofsampling steps: their cost of distillation thus scales linearly with this number of steps, which can
Revised: 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

Original: [Figure]
Revised: [Figure]

Original:  be prohibitive.
Revised: sampling steps: their cost of distillation thus scales linearly with this number of steps, which can be prohibitive.

Original: 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.
Revised: 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.

Original: DDIM (Song et al., 2021a) was originally shown to be effective for few-step sampling, as was the probability ﬂow sampler (Song et al., 2021c).
Revised: DDIM (Song et al., 2021a) was originally shown to be effective for few-step sampling, as was the probability ﬂow sampler (Song et al., 2021c).

Original: Jolicoeur-Martineau et al.
Revised: Jolicoeur-Martineau et al.

Original: (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.
Revised: (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.

Original: Other work on fast sampling can be viewed as manual or automated methods to adjust samplers or diffusion processes for fast generation.
Revised: Other work on fast sampling can be viewed as manual or automated methods to adjust samplers or diffusion processes for fast generation.

Original: 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.
Revised: 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.

Original: Watson et al.
Revised: Watson et al.

Original: (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.
Revised: (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.

Original: Chen et al.
Revised: Chen et al.

Original: ; Saharia et al.
Revised: ; Saharia et al.

Original: ; Ho et al.
Revised: ; Ho et al.

Original: (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.
Revised: (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.

Original: Their method is effective in highly conditioned settings such as text-to-speech and image super-resolution.
Revised: Their method is effective in highly conditioned settings such as text-to-speech and image super-resolution.

Original: San-Roman et al.
Revised: San-Roman et al.

Original: (2021) train a new network to estimate the noise level of noisy data and show how to use this estimate to speed up sampling.
Revised: (2021) train a new network to estimate the noise level of noisy data and show how to use this estimate to speed up sampling.

Original: Alternative speciﬁcations of the diffusion model can also lend themselves to fast sampling, such as modiﬁed forward and reverse processes (Nachmani et al., 2021; Lam et al., 2021) and training diffusion models in latent space (Vahdat et al., 2021)
Revised: Alternative speciﬁcations of the diffusion model can also lend themselves to fast sampling, such as modiﬁed forward and reverse processes (Nachmani et al., 2021; Lam et al., 2021) and training diffusion models in latent space (Vahdat et al., 2021)

Original: [Table]
Revised: [Table]

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: In addition to the proposed distillation procedure, some of our progress was realized through different parameterizations of the diffusion model and its training loss.
Revised: In addition to the proposed distillation procedure, some of our progress was realized through different parameterizations of the diffusion model and its training loss.

Original: We expect to see more progress in this direction as the community further explores this model class.
Revised: We expect to see more progress in this direction as the community further explores this model class.

Original: R EPRODUCIBILITY STATEMENT
Revised: R EPRODUCIBILITY STATEMENT

Original: We provide full details on model architectures, training procedures and hyperparameters in Appendix E, in addition to our discussion in Section 5.
Revised: We provide full details on model architectures, training procedures and hyperparameters in Appendix E, in addition to our discussion in Section 5.

Original: In Algorithm 2 we provide fairly detailed pseudocode that closely matches our actual implementation.
Revised: In Algorithm 2 we provide fairly detailed pseudocode that closely matches our actual implementation.

Original: We plan to open source our code and trained model checkpoints with the release of the ﬁnal version of this paper.
Revised: We plan to open source our code and trained model checkpoints with the release of the ﬁnal version of this paper.

Original: E THICS STATEMENT
Revised: E THICS STATEMENT

Original: 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.
Revised: 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.

Original: The focus of this paper speciﬁcally is on speeding up generative models at test time in order to reduce their computational demands; we do not have speciﬁc concerns with regards to this contribution.
Revised: The focus of this paper speciﬁcally is on speeding up generative models at test time in order to reduce their computational demands; we do not have speciﬁc concerns with regards to this contribution.

Original: A P ROBABILITY FLOW ODE IN TERMS OF LOG -SNR
Revised: A P ROBABILITY FLOW ODE IN TERMS OF LOG -SNR

Original: Song et al.
Revised: Song et al.

Original: (2021c) formulate the forward diffusion process in terms of an SDE of the form
Revised: (2021c) formulate the forward diffusion process in terms of an SDE of the form

Original: [Equation]
Revised: [Equation]

Original: 10) and show that samples from this diffusion process can be generated by solving the associated probability ﬂow ODE:
Revised: 10) and show that samples from this diffusion process can be generated by solving the associated probability ﬂow ODE:

Original: [Equation]
Revised: [Equation]

Original: where in practice ∇ z log p t ( z ) is approximated by a learned denoising model using
Revised: where in practice ∇ z log p t ( z ) is approximated by a learned denoising model using

Original: [Equation]
Revised: [Equation]

Original: Following Kingma et al.
Revised: Following Kingma et al.

Original: we have f ( z , t ) =
Revised: we have f ( z , t ) =

Original: d log α t dt
Revised: d log α t dt

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: σ
Revised: σ

Original: [Equation]
Revised: [Equation]

Original: Assuming a variance preserving diffusion process with α 2 t = 1 − σ 2 t = sigmoid ( λ t ) for λ t
Revised: Assuming a variance preserving diffusion process with α 2 t = 1 − σ 2 t = sigmoid ( λ t ) for λ t

Original: = log
Revised: = log

Original: [ α 2 t /σ 2 t ] (without loss of generality, see Kingma et al.
Revised: [ α 2 t /σ 2 t ] (without loss of generality, see Kingma et al.

Original: , we get
Revised: , we get

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Similarly, we get
Revised: Similarly, we get

Original: [Equation]
Revised: [Equation]

Original: Plugging these in to the probability ﬂow ODE then gives
Revised: Plugging these in to the probability ﬂow ODE then gives

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Plugging in our function approximation from Equation 12 gives
Revised: Plugging in our function approximation from Equation 12 gives

Original: [Equation]
Revised: [Equation]

Original: B DDIM IS AN INTEGRATOR OF THE PROBABILITY FLOW ODE
Revised: B DDIM IS AN INTEGRATOR OF THE PROBABILITY FLOW ODE

Original: The DDIM update rule (Song & Ermon, 2020) is given by
Revised: The DDIM update rule (Song & Ermon, 2020) is given by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: for s < t .
Revised: for s < t .

Original: Taking the derivative of this expression with respect to λ s , assuming again a variance preserving diffusion process
Revised: Taking the derivative of this expression with respect to λ s , assuming again a variance preserving diffusion process

Original: [Equation]
Revised: [Equation]

Original: and using dα λ dλ = 12 α λ σ 2 λ and dσ λ dλ = − 12 σ λ α 2 λ , gives
Revised: and using dα λ dλ = 12 α λ σ 2 λ and dσ λ dλ = − 12 σ λ α 2 λ , gives

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: Evaluating this derivative at s = t then gives
Revised: Evaluating this derivative at s = t then gives

Original: [Equation]
Revised: [Equation]

Original: Comparison with Equation 19 now shows that DDIM follows the probability ﬂow ODE up to ﬁrst order, and can thus be considered as an integration rule for this ODE.
Revised: Comparison with Equation 19 now shows that DDIM follows the probability ﬂow ODE up to ﬁrst order, and can thus be considered as an integration rule for this ODE.

Original: C E VALUATION OF INTEGRATORS OF THE PROBABILITY FLOW ODE
Revised: C E VALUATION OF INTEGRATORS OF THE PROBABILITY FLOW ODE

Original: In a preliminary investigation we tried several numerical integrators for the probability ﬂow ODE.
Revised: In a preliminary investigation we tried several numerical integrators for the probability ﬂow ODE.

Original: As our model we used a pre-trained class-conditional 128x128 ImageNet model following the description in Ho et al.
Revised: As our model we used a pre-trained class-conditional 128x128 ImageNet model following the description in Ho et al.

Original: We tried a simple Euler integrator, RK4 (the “classic” 4th order
Revised: We tried a simple Euler integrator, RK4 (the “classic” 4th order

Original: Runge–Kutta integrator), and DDIM (Song et al., 2021a).
Revised: Runge–Kutta integrator), and DDIM (Song et al., 2021a).

Original: In addition we compared to a Gaussian sampler with variance equal to the lower bound given by Ho et al.
Revised: In addition we compared to a Gaussian sampler with variance equal to the lower bound given by Ho et al.

Original: We calculated FID scores on just 5000 samples, hence our results in this experiment are not comparable to results reported in the literature.
Revised: We calculated FID scores on just 5000 samples, hence our results in this experiment are not comparable to results reported in the literature.

Original: This preliminary investigation gave the results listed in Table 3 and identiﬁed DDIM as the best integrator in terms of resulting sample quality.
Revised: This preliminary investigation gave the results listed in Table 3 and identiﬁed DDIM as the best integrator in terms of resulting sample quality.

Original: [Table]
Revised: [Table]

Original: D E XPRESSION OF DDIM IN ANGULAR PARAMETERIZATION
Revised: D E XPRESSION OF DDIM IN ANGULAR PARAMETERIZATION

Original: 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.
Revised: 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.

Original: Given our deﬁnition of φ , and assuming a variance preserving diffusion process, we have α φ = cos( φ ) , σ φ = sin( φ ) , and hence z φ = cos( φ ) x + sin( φ ) (cid:15) .
Revised: Given our deﬁnition of φ , and assuming a variance preserving diffusion process, we have α φ = cos( φ ) , σ φ = sin( φ ) , and hence z φ = cos( φ ) x + sin( φ ) (cid:15) .

Original: We can now deﬁne the velocity of z φ as
Revised: We can now deﬁne the velocity of z φ as

Original: [Equation]
Revised: [Equation]

Original: Rearranging (cid:15), x , v , we then get
Revised: Rearranging (cid:15), x , v , we then get

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: and similarly we get (cid:15)
Revised: and similarly we get (cid:15)

Original: =
Revised: =

Original: sin( φ ) z φ + cos( φ ) v φ .
Revised: sin( φ ) z φ + cos( φ ) v φ .

Original: Furthermore, we deﬁne the predicted velocity as
Revised: Furthermore, we deﬁne the predicted velocity as

Original: [Equation]
Revised: [Equation]

Original: where ˆ (cid:15) θ ( z φ )
Revised: where ˆ (cid:15) θ ( z φ )

Original: =
Revised: =

Original: ( z φ − cos( φ )ˆ x θ ( z φ )) / sin( φ ) .
Revised: ( z φ − cos( φ )ˆ x θ ( z φ )) / sin( φ ) .

Original: Rewriting the DDIM update rule in the introduced terms then gives
Revised: Rewriting the DDIM update rule in the introduced terms then gives

Original: [Equation]
Revised: [Equation]

Original: Finally, we use the trigonometric identities
Revised: Finally, we use the trigonometric identities

Original: [Equation]
Revised: [Equation]

Original: 36) to ﬁnd that
Revised: 36) to ﬁnd that

Original: [Equation]
Revised: [Equation]

Original: or equivalently
Revised: or equivalently

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: The relationship between z φ t , v t , α t , σ t , x , (cid:15) is visualized in Figure 5.
Revised: The relationship between z φ t , v t , α t , σ t , x , (cid:15) is visualized in Figure 5.

Original: E S ETTINGS USED IN EXPERIMENTS
Revised: E S ETTINGS USED IN EXPERIMENTS

Original: Our model architectures closely follow those described by Dhariwal & Nichol (2021).
Revised: Our model architectures closely follow those described by Dhariwal & Nichol (2021).

Original: For 64 ×
Revised: For 64 ×

Original: ImageNet we use their model exactly, with 192 channels at the highest resolution.
Revised: ImageNet we use their model exactly, with 192 channels at the highest resolution.

Original: All other models are slight variations with different hyperparameters.
Revised: All other models are slight variations with different hyperparameters.

Original: For CIFAR-10 we use an architecture with a ﬁxed number of channels at all resolutions of 256 .
Revised: For CIFAR-10 we use an architecture with a ﬁxed number of channels at all resolutions of 256 .

Original: The model consists of a UNet that internally downsamples the data twice, to 16 × 16 and to 8 × 8 .
Revised: The model consists of a UNet that internally downsamples the data twice, to 16 × 16 and to 8 × 8 .

Original: At each resolution we apply 3 residual blocks, like described by Dhariwal & Nichol (2021).
Revised: At each resolution we apply 3 residual blocks, like described by Dhariwal & Nichol (2021).

Original: We use single-headed attention, and only apply this at the 16 × 16 and 8 × 8 resolutions.
Revised: We use single-headed attention, and only apply this at the 16 × 16 and 8 × 8 resolutions.

Original: We use dropout of 0 . 2 when training the original model.
Revised: We use dropout of 0 . 2 when training the original model.

Original: No dropout is used during distillation.
Revised: No dropout is used during distillation.

Original: For LSUN we use a model similar to that for ImageNet, but with a reduced number of 128 channels at the 64 × 64 resolution.
Revised: For LSUN we use a model similar to that for ImageNet, but with a reduced number of 128 channels at the 64 × 64 resolution.

Original: 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 withchannels.
Revised: 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 withchannels.

Original: We only use attention layers for the resolutions of 32 × 32 and lower.
Revised: We only use attention layers for the resolutions of 32 × 32 and lower.

Original: For CIFAR-10 we take the output of the model to represent a prediction of x directly, as discussed in Section 4.
Revised: For CIFAR-10 we take the output of the model to represent a prediction of x directly, as discussed in Section 4.

Original: For the other data sets we used the combined prediction of ( x , (cid:15) ) like described in that section also.
Revised: For the other data sets we used the combined prediction of ( x , (cid:15) ) like described in that section also.

Original: All original models are trained with Adam with standard settings (learning rate of 3 ∗ 10 − 4 ), using a parameter moving average with constant 0 .
Revised: All original models are trained with Adam with standard settings (learning rate of 3 ∗ 10 − 4 ), using a parameter moving average with constant 0 .

Original: 9999 and very slight decoupled weight decay (Loshchilov & Hutter, 2017) with a constant of 0 .
Revised: 9999 and very slight decoupled weight decay (Loshchilov & Hutter, 2017) with a constant of 0 .

Original: We clip the norm of gradients to a global norm of 1 before calculating parameter updates.
Revised: We clip the norm of gradients to a global norm of 1 before calculating parameter updates.

Original: For CIFAR-10 we train for 800k parameter updates, for ImageNet we use 550k updates, and for LSUN we use 400k updates.
Revised: For CIFAR-10 we train for 800k parameter updates, for ImageNet we use 550k updates, and for LSUN we use 400k updates.

Original: 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.
Revised: 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.

Original: We linearly anneal the learning rate from 10 − 4 to zero during each iteration.
Revised: We linearly anneal the learning rate from 10 − 4 to zero during each iteration.

Original: We use a batch size of 128 for CIFAR-10 and 2048 for the other data sets.
Revised: We use a batch size of 128 for CIFAR-10 and 2048 for the other data sets.

Original: We run our experiments on TPUv4, using 8 TPU chips for CIFAR-10, and 64 chips for the other data sets.
Revised: We run our experiments on TPUv4, using 8 TPU chips for CIFAR-10, and 64 chips for the other data sets.

Original: The total time required to ﬁrst train and then distill a model varies from about a day for CIFAR-10, to about 5 days for ImageNet.
Revised: The total time required to ﬁrst train and then distill a model varies from about a day for CIFAR-10, to about 5 days for ImageNet.

Original: F S TOCHASTIC SAMPLING WITH DISTILLED MODELS
Revised: F S TOCHASTIC SAMPLING WITH DISTILLED MODELS

Original: 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.
Revised: 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.

Original: Here we evaluate a distilled model for 64x

Revised: Here we evaluate a distilled model for 64x

Original: ImageNet using the optimized stochastic sampler also used in Section 5.2.
Revised: ImageNet using the optimized stochastic sampler also used in Section 5.2.

Original: The results are presented in Figure 6.
Revised: The results are presented in Figure 6.

Original: [Figure]
Revised: [Figure]

Original: G D ERIVATION OF THE DISTILLATION TARGET
Revised: G D ERIVATION OF THE DISTILLATION TARGET

Original: 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.
Revised: 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.

Original: In standard diffusion training, the target for denoising is the clean data x .
Revised: In standard diffusion training, the target for denoising is the clean data x .

Original: In progressive distillation it is the value ˜ x
Revised: In progressive distillation it is the value ˜ x

Original: the student denoising model would need to predict in order to match the teacher model when sampling.
Revised: the student denoising model would need to predict in order to match the teacher model when sampling.

Original: Here we derive what this target needs to be.
Revised: Here we derive what this target needs to be.

Original: Using notation t (cid:48)
Revised: Using notation t (cid:48)

Original: =
Revised: =

Original: t − 0 .
Revised: t − 0 .

Original: 5 /N and t (cid:48)(cid:48)
Revised: 5 /N and t (cid:48)(cid:48)

Original: =
Revised: =

Original: t − 1 /N , when training a student with N sampling steps, we have that the teacher model samples the next set of noisy data z t (cid:48)(cid:48) given the current noisy data z t by taking two steps of DDIM.
Revised: t − 1 /N , when training a student with N sampling steps, we have that the teacher model samples the next set of noisy data z t (cid:48)(cid:48) given the current noisy data z t by taking two steps of DDIM.

Original: The student tries to sample the same value in only one step of DDIM.
Revised: The student tries to sample the same value in only one step of DDIM.

Original: Denoting the student denoising prediction by ˜ x , and its one-step sample by ˜ z t (cid:48)(cid:48) , application of the
Revised: Denoting the student denoising prediction by ˜ x , and its one-step sample by ˜ z t (cid:48)(cid:48) , application of the

Original: DDIM sampler (see equation 8), gives:
Revised: DDIM sampler (see equation 8), gives:

Original: [Equation]
Revised: [Equation]

Original: In order for the student sampler to match the teacher sampler, we must set ˜ z t (cid:48)(cid:48) equal to z t [Equation]
Revised: In order for the student sampler to match the teacher sampler, we must set ˜ z t (cid:48)(cid:48) equal to z t (cid:48)(cid:48) .

Original: This gives
Revised: This gives

Original: [Equation]
Revised: [Equation]

Original: In other words, if our student denoising model exactly predicts ˜ x as deﬁned in equation 43 above, then the one-step student sample ˜
Revised: In other words, if our student denoising model exactly predicts ˜ x as deﬁned in equation 43 above, then the one-step student sample ˜

Original: z t (cid:48)(cid:48) is identical to the two-step teacher sample z t (cid:48)(cid:48) .
Revised: z t (cid:48)(cid:48) is identical to the two-step teacher sample z t (cid:48)(cid:48) .

Original: In order to have our student model approximate this ideal outcome, we thus train it to predict ˜ x from z t as well as possible, using the standard squared error denoising loss (see Equation 9).
Revised: In order to have our student model approximate this ideal outcome, we thus train it to predict ˜ x from z t as well as possible, using the standard squared error denoising loss (see Equation 9).

Revised: 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

Revised: DDPM sampler can thus not be distilled into a few-step sampler without some loss in ﬁdelity.

Revised: 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.

Revised: Finally, note that we do lose something during the progressive distillation process: while the original model was trained to denoise z t for any given continuous time t , the distilled student models are only ever evaluated on a small discrete set of times t .

Revised: The student models thus lose generality as distillation progresses.

Original:  H A DDITIONAL RANDOM SAMPLES 
Revised: 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 A DDITIONAL 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 1 sampling steps.

Original: [Figure]
Revised: 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.

Original:  our distilled CIFAR-10 models, for ﬁxed random seed and for varying number of sampling steps.
Revised: 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. [Figure] [Figure] [Figure]

Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Revised: I A BLATION WITH FASTER DISTILLATION SCHEDULES

Revised: 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.

Revised: In Figure 12 we present results for taking 25 thousand,thousand, or 5 thousand optimization steps, instead of the 50 thousand we suggested in Section 3.

Revised: 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.

Revised: When taking very few sampling steps, the loss in performance becomes more pronounced when training the student for only a short time.

Revised: 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.

Revised: 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.

Revised: 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.

Original: [Figure]
Revised: [Figure]

Revised: sampling steps

Original: [Figure]
Revised: [Figure]

Revised: CIFARDistilled 50k

Revised: Distilled 25k

Revised: Distilled 10k

Revised: Distilled 5k

Revised: Distilled 4 ×