Improved Denoising Diffusion Probabilistic Models

1 INTRODUCTION . Sohl-Dickstein et al . ( 2015 ) introduced diffusion probabilistic models ( `` diffusion models '' for brevity ) , a class of generative models which match a data distribution by learning to reverse a gradual , multi-step noising process . More recently , Ho et al . ( 2020 ) showed an equivalence between these models and score based generative models ( Song & Ermon , 2019 ; 2020 ) , which learn a gradient of the log-density of the data distribution using denoising score matching ( Hyvärinen , 2005 ) . It has recently been shown that this class of models can produce high-quality images ( Ho et al. , 2020 ; Song & Ermon , 2020 ; Jolicoeur-Martineau et al. , 2020 ) and audio ( Chen et al. , 2020b ; Kong et al. , 2020 ) , but it has yet to be shown that diffusion models can achieve competitive log-likelihoods . Furthermore , while Ho et al . ( 2020 ) showed extremely good results on the CIFAR-10 ( Krizhevsky , 2009 ) and LSUN ( Yu et al. , 2015 ) datasets , it is unclear how well diffusion models scale to datasets with higher diversity such as ImageNet . Finally , while Chen et al . ( 2020b ) found that diffusion models can efficiently generate audio using a small number of sampling steps , it has yet to be shown that the same is true for images . In this paper , we show that diffusion models can achieve competitive log-likelihoods while maintaining good sample quality , even on high-diversity datasets like ImageNet . Additionally , we show that our improved models can produce competitive samples an order of magnitude faster than those from Ho et al . ( 2020 ) . We achieve these results by combining a simple reparameterization of the reverse process variance , a hybrid learning objective that combines the variational lower-bound with the simplified objective from Ho et al . ( 2020 ) , and a novel noise schedule which allows the model to better leverage the entire diffusion process . We find surprisingly that , with our hybrid objective , our models obtain better log-likelihoods than those obtained by optimizing the log-likelihood directly , and discover that the latter objective has much more gradient noise during training . We show that a simple importance sampling technique reduces this noise and allows us to achieve better log-likelihoods than with the hybrid objective . Using our trained models , we study how sample quality and log-likelihood change as we adjust the number of diffusion steps used at sampling time . We demonstrate that our improved models allow us to use an order of magnitude fewer steps at test time with only a modest change in sample quality and log-likelihood , thus speeding up sampling for use in practical applications . Finally , we evaluate the performance of these models as we increase model size , and observe trends that suggest predictable improvements in performance as we increase training compute . 2 DENOISING DIFFUSION PROBABILISTIC MODELS . We briefly review the formulation of diffusion models from Ho et al . ( 2020 ) . This formulation makes various simplifying assumptions , such as a fixed noising process q which adds diagonal Gaussian noise at each timestep . For a more general derivation , see Sohl-Dickstein et al . ( 2015 ) . 2.1 DEFINITIONS . Given a data distribution x0 ∼ q ( x0 ) , we define a forward noising process q which produces latents x1 through xT by adding Gaussian noise at time t with variance βt ∈ ( 0 , 1 ) as follows : q ( x1 , ... , xT |x0 ) : = T∏ t=1 q ( xt|xt−1 ) ( 1 ) q ( xt|xt−1 ) : = N ( xt ; √ 1− βtxt−1 , βtI ) ( 2 ) Given sufficiently large T and a well behaved schedule of βt , the latent xT is nearly an isotropic Gaussian distribution . Thus , if we know the exact reverse distribution q ( xt−1|xt ) , we can sample xT ∼ N ( 0 , I ) and run the process in reverse to get a sample from q ( x0 ) . However , since q ( xt−1|xt ) depends on the entire data distribution , we approximate it using a neural network : pθ ( xt−1|xt ) : = N ( xt−1 ; µθ ( xt , t ) , Σθ ( xt , t ) ) ( 3 ) The combination of q and p is a variational auto-encoder ( Kingma & Welling , 2013 ) , and we can write the variational lower bound ( VLB ) as follows : Lvlb : = L0 + L1 + ... + LT−1 + LT ( 4 ) L0 : = − log pθ ( x0|x1 ) ( 5 ) Lt−1 : = DKL ( q ( xt−1|xt , x0 ) || pθ ( xt−1|xt ) ) ( 6 ) LT : = DKL ( q ( xT |x0 ) || p ( xT ) ) ( 7 ) Aside from L0 , each term of Equation 4 is a KL divergence between two Gaussian distributions , and can thus be evaluated in closed form . To evaluate L0 for images , we assume that each color component is divided into 256 bins , and we compute the probability of pθ ( x0|x1 ) landing in the correct bin ( which is tractable using the CDF of the Gaussian distribution ) . Also note that while LT does not depend on θ , it will be close to zero if the forward noising process adequately destroys the data distribution so that q ( xT |x0 ) ≈ N ( 0 , I ) . It is useful to define and derive several other quantities which are relevant to the forward noising process , so we repeat them here from Ho et al . ( 2020 ) : αt : = 1− βt ( 8 ) ᾱt : = t∏ s=0 αs ( 9 ) β̃t : = 1− ᾱt−1 1− ᾱt βt ( 10 ) µ̃t ( xt , x0 ) : = √ ᾱt−1βt 1− ᾱt x0 + √ αt ( 1− ᾱt−1 ) 1− ᾱt xt ( 11 ) q ( xt|x0 ) = N ( xt ; √ ᾱtx0 , ( 1− ᾱt ) I ) ( 12 ) q ( xt−1|xt , x0 ) = N ( xt−1 ; µ̃ ( xt , x0 ) , β̃tI ) ( 13 ) 2.2 TRAINING IN PRACTICE . Equation 12 provides an efficient way to jump directly to an arbitrary step of the forward noising process . This makes it possible to randomly sample t during training . Ho et al . ( 2020 ) uniformly sample t for each image in each mini-batch . There are many different ways to parameterize µθ ( xt , t ) . The most obvious option is to predict µθ ( xt , t ) directly with a neural network ; alternatively , the network could predict x0 , and this output could then be fed through µ̃ ( xt , x0 ) ; finally , the network could predict the noise added to x0 , and this noise could be used to predict x0 via x0 = 1 √ αt ( xt − βt√ 1− ᾱt ) ( 14 ) Ho et al . ( 2020 ) found that predicting worked best , especially when combined with a reweighted loss function : Lsimple = Et , x0 , [ || − θ ( xt , t ) ||2 ] ( 15 ) This objective can be seen as a reweighted form of Lvlb ( without the terms affecting Σθ ) . The authors found that optimizing this reweighted objective resulted in much better sample quality than optimizing Lvlb directly , and explain this by drawing a connection to generative score matching ( Song & Ermon , 2019 ; 2020 ) . One subtlety is that Lsimple provides no learning signal for Σθ ( xt , t ) . This is irrelevant , however , since Ho et al . ( 2020 ) achieved their best results by fixing the variance to σ2t I rather than learning it . They found that they achieve similar sample quality using either σ2t = βt or σ 2 t = β̃t , which are two extremes given by q ( x0 ) being either isotropic Gaussian noise or a delta function , respectively . 3 IMPROVING THE LOG-LIKELIHOOD . While Ho et al . ( 2020 ) found that diffusion models can generate high-fidelity samples according to FID ( Heusel et al. , 2017 ) and Inception Score ( Salimans et al. , 2016 ) , they were unable to achieve competitive log-likelihoods with these models . Log-likelihood is a widely used metric in generative modeling , and it is generally believed that optimizing log-likelihood forces generative models to capture all of the modes of the data distribution ( Razavi et al. , 2019 ) . Additionally , recent work ( Henighan et al. , 2020 ) has shown that small improvements in log-likelihood can have a dramatic impact on sample quality and learnt feature representations . Thus , it is important to explore why diffusion models seem to perform poorly on this metric , since this may suggest a fundamental shortcoming such as bad mode coverage . This section explores several modifications to the algorithm described in Section 2 that , when combined , allow diffusion models to achieve much better log-likelihoods on image datasets , suggesting that these models enjoy the same benefits as other likelihood-based generative models . To study the effects of different modifications , we train fixed model architectures with fixed hyperparameters ( Appendix A ) on the ImageNet 64 × 64 ( van den Oord et al. , 2016a ) and CIFAR-10 ( Krizhevsky , 2009 ) datasets . While CIFAR-10 has seen more usage for this class of models , we chose to study ImageNet 64× 64 as well because it provides a good trade-off between diversity and resolution , allowing us to train models quickly without worrying about overfitting . Additionally , ImageNet 64×64 has been studied extensively in the context of generative modeling ( van den Oord et al. , 2016b ; Menick & Kalchbrenner , 2018 ; Child et al. , 2019 ; Roy et al. , 2020 ) , allowing us to compare diffusion models directly to many other generative models . The setup from Ho et al . ( 2020 ) ( optimizing Lsimple while setting σ2t = βt and T = 1000 ) achieves a log-likelihood of 3.99 bits/dim on ImageNet 64 × 64 after 200K training iterations . We found in early experiments that we could get a boost in log-likelihood by increasing T from 1000 to 4000 ; with this change , the log-likelihood improves to 3.77 bits/dim . For the remainder of this section , we use T = 4000 , but we explore this choice in Section 4 . 3.1 LEARNING Σθ ( xt , t ) In Ho et al . ( 2020 ) , the authors set Σθ ( xt , t ) = σ2t I , where σt is not learned . Oddly , they found that fixing σ2t to βt yielded roughly the same sample quality as fixing it to β̃t . Considering that βt and β̃t represent two opposite extremes , it is reasonable to ask why this choice doesn ’ t affect samples . One clue is given by Figure 1a , which shows that βt and β̃t are almost equal except near t = 0 , i.e . where the model is dealing with imperceptible details . Furthermore , as we increase the number of diffusion steps , βt and β̃t seem to remain close to one another for more of the diffusion process . Figure 1a : The ratio β̃t/βt for every diffusion step for diffusion processes of different lengths . 0 500 1000 1500 2000 2500 3000 3500 4000 diffusion step Figure 2 : Latent samples from linear ( top ) and cosine ( bottom ) schedules respectively at linearly spaced values of t from 0 to T . The latents in the last quarter of the linear schedule are almost purely noise , whereas the cosine schedule adds noise more slowly This suggests that , in the limit of infinite diffusion steps , the choice of σt might not matter at all for sample quality . In other words , as we add more diffusion steps , the model mean µθ ( xt , t ) determines the distribution much more than Σθ ( xt , t ) . While the above argument suggests that fixing σt is a reasonable choice for the sake of sample quality , it says nothing about log-likelihood . In fact , Figure 1b shows that the first few steps of the diffusion process contribute the most to the variational lower bound . Thus , it seems likely that we could improve log-likelihood by using a better choice of Σθ ( xt , t ) . To achieve this , we must learn Σθ ( xt , t ) without the instabilities encountered by Ho et al . ( 2020 ) . Since Figure 1a shows that the reasonable range for Σθ ( xt , t ) is very small , it would be hard for a neural network to predict Σθ ( xt , t ) directly , even in the log domain , as observed by Ho et al . ( 2020 ) . Instead , we found it better to parameterize the variance as an interpolation between βt and β̃t in the log domain . In particular , our model outputs a vector v containing one component per dimension , and we turn this output into variances as follows : Σθ ( xt , t ) = exp ( v log βt + ( 1− v ) log β̃t ) ( 16 ) We did not apply any constraints on v , theoretically allowing the model to predict variances outside of the interpolated range . However , we did not observe the network doing this in practice , suggesting that the bounds for Σθ ( xt , t ) are indeed expressive enough . Since Lsimple doesn ’ t depend on Σθ ( xt , t ) , we define a new hybrid objective : Lhybrid = Lsimple + λLvlb ( 17 ) For our experiments , we set λ = 0.001 to prevent Lvlb from overwhelming Lsimple . Along this same line of reasoning , we also apply a stop-gradient to the µθ ( xt , t ) output for the Lvlb term . This way , Lvlb can guide Σθ ( xt , t ) while Lsimple is still the main source of influence over µθ ( xt , t ) .

Summary: The paper found several methods to improve log likelihood of diffusion models while maintain their sample quality, including cosine instead of linear noise schedule, using a hybrid objective to learn parameters of the covariance function, and using importance sampling to improve the gradient noise. The authors also explore how sample quality and log likelihood scale with the number of diffusion steps and model capacity. Experiments on 64x64 ImageNet dataset show competitive llh while keeping the sample quality.