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