Progressive Distillation for Fast Sampling of Diffusion Models

1 INTRODUCTION . 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 . 2 BACKGROUND ON DIFFUSION MODELS . We consider diffusion models ( Sohl-Dickstein et al. , 2015 ; Song & Ermon , 2019 ; Ho et al. , 2020 ) specified in continuous time ( Tzen & Raginsky , 2019a ; Song 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 . 3 PROGRESSIVE DISTILLATION . 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

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