Variational Diffusion Models

1 Introduction . Likelihood-based generative modeling is a central task in machine learning that is the basis for a wide range of applications ranging from speech synthesis [ Oord et al. , 2016 ] , to translation [ Sutskever et al. , 2014 ] , to compression [ MacKay , 2003 ] , to many others . Autoregressive models have long been the dominant model class on this task due to their tractable likelihood and expressivity , as shown in Figure 1 . Diffusion models have recently shown impressive results in image [ Ho et al. , 2020 , Song et al. , 2021b , Nichol and Dhariwal , 2021 ] and audio generation [ Kong et al. , 2020 , Chen et al. , 2020 ] in terms of perceptual quality , but have yet to match autoregressive models on density estimation benchmarks . In this paper we make several technical contributions that allow diffusion models to challenge the dominance of autoregressive models in this domain . Our main contributions are as follows : • We introduce a flexible family of diffusion-based generative models that achieve new stateof-the-art log-likelihoods on standard image density estimation benchmarks ( CIFAR-10 and ImageNet ) . This is enabled by incorporating Fourier features into the diffusion model and using a learnable specification of the diffusion process , among other modeling innovations . • We improve our theoretical understanding of density modeling using diffusion models by analyzing their variational lower bound ( VLB ) , deriving a remarkably simple expression in terms of the signal-to-noise ratio of the diffusion process . This result delivers new insight * Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . into the model class : for the continuous-time ( infinite-depth ) setting we prove a novel invariance of the generative model and its VLB to the specification of the diffusion process , and we show that various diffusion models from the literature are equivalent up to a trivial time-dependent rescaling of the data . 2 Related work . Our work builds on diffusion probabilistic models ( DPMs ) [ Sohl-Dickstein et al. , 2015 ] , or diffusion models in short . DPMs can be viewed as a type of variational autoencoder ( VAE ) [ Kingma and Welling , 2013 , Rezende et al. , 2014 ] , whose structure and loss function allows for efficient training of arbitrarily deep models . Interest in diffusion models has recently reignited due to their impressive image generation results [ Ho et al. , 2020 , Song and Ermon , 2020 ] . Ho et al . [ 2020 ] introduced a number of model innovations to the original DPM , with impressive results on image generation quality benchmarks . They showed that the VLB objective , for a diffusion model with discrete time and diffusion variances shared across input dimensions , is equivalent to multiscale denoising score matching , up to particular weightings per noise scale . Further improvements were proposed by Nichol and Dhariwal [ 2021 ] , resulting in better log-likelihood scores . Gao et al . [ 2020 ] show how diffusion can also be used to efficiently optimize energy-based models ( EBMs ) towards a close approximation of the log-likelihood objective , resulting in high-fidelity samples even after long MCMC chains . Song and Ermon [ 2019 ] first proposed learning generative models through a multi-scale denoising score matching objective , with improved methods in Song and Ermon [ 2020 ] . This was later extended to continuous-time diffusion with novel sampling algorithms based on reversing the diffusion process [ Song et al. , 2021b ] . Concurrent to our work , Song et al . [ 2021a ] , Huang et al . [ 2021 ] , and Vahdat et al . [ 2021 ] also derived variational lower bounds to the data likelihood under a continuous-time diffusion model . Where we consider the infinitely deep limit of a standard VAE , Song et al . [ 2021a ] and Vahdat et al . [ 2021 ] present different derivations based on stochastic differential equations . Huang et al . [ 2021 ] considers both perspectives and discusses the similarities between the two approaches . An advantage of our analysis compared to these other works is that we present an intuitive expression of the VLB in terms of the signal-to-noise ratio of the diffused data , leading to much simplified expressions of the discrete-time and continuous-time loss , allowing for simple and numerically stable implementation . This also leads to new results on the invariance of the generative model and its VLB to the specification of the diffusion process . We empirically compare to these works , as well as others , in Table 1 . Previous approaches to diffusion probabilistic models fixed the diffusion process ; in contrast optimize the diffusion process parameters jointly with the rest of the model . This turns the model into a type of VAE [ Kingma and Welling , 2013 , Rezende et al. , 2014 ] . This is enabled by directly parameterizing the mean and variance of the marginal q ( zt|z0 ) , where previous approaches instead parameterized the individual diffusion steps q ( zt+✏|zt ) . In addition , our denoising models include several architecture changes , the most important of which is the use of Fourier features , which enable us to reach much better likelihoods than previous diffusion probabilistic models . 3 Model . We will focus on the most basic case of generative modeling , where we have a dataset of observations of x , and the task is to estimate the marginal distribution p ( x ) . As with most generative models , the described methods can be extended to the case of multiple observed variables , and/or the task of estimating conditional densities p ( x|y ) . The proposed latent-variable model consists of a diffusion process ( Section 3.1 ) that we invert to obtain a hierarchical generative model ( Section 3.3 ) . As we will show , the model choices below result in a surprisingly simple variational lower bound ( VLB ) of the marginal likelihood , which we use for optimization of the parameters . 3.1 Forward time diffusion process . Our starting point is a Gaussian diffusion process that begins with the data x , and defines a sequence of increasingly noisy versions of x which we call the latent variables zt , where t runs from t = 0 ( least noisy ) to t = 1 ( most noisy ) . The distribution of latent variable zt conditioned on x , for any t 2 [ 0 , 1 ] is given by : q ( zt|x ) = N ↵tx , 2 t I , ( 1 ) where ↵t and 2t are strictly positive scalar-valued functions of t. Furthermore , let us define the signal-to-noise ratio ( SNR ) : SNR ( t ) = ↵2t / 2 t . ( 2 ) We assume that the SNR ( t ) is strictly monotonically decreasing in time , i.e . that SNR ( t ) < SNR ( s ) for any t > s. This formalizes the notion that the zt is increasingly noisy as we go forward in time . We also assume that both ↵t and 2t are smooth , such that their derivatives with respect to time t are finite . This diffusion process specification includes the variance-preserving diffusion process as used by [ Sohl-Dickstein et al. , 2015 , Ho et al. , 2020 ] as a special case , where ↵t = p 1 2t . Another special case is the variance-exploding diffusion process as used by [ Song and Ermon , 2019 , Song et al. , 2021b ] , where ↵2t = 1 . In experiments , we use the variance-preserving version . The distributions q ( zt|zs ) for any t > s are also Gaussian , and given in Appendix A . The joint distribution of latent variables ( zs , zt , zu ) at any subsequent timesteps 0 s < t < u 1 is Markov : q ( zu|zt , zs ) = q ( zu|zt ) . Given the distributions above , it is relatively straightforward to verify through Bayes rule that q ( zs|zt , x ) , for any 0 s < t 1 , is also Gaussian . This distribution is also given in Appendix A . 3.2 Noise schedule . In previous work , the noise schedule has a fixed form ( see Appendix H , Fig . 4a ) . In contrast , we learn this schedule through the parameterization 2t = sigmoid ( ⌘ ( t ) ) ( 3 ) where ⌘ ( t ) is a monotonic neural network with parameters ⌘ , as detailed in Appendix H. Motivated by the equivalence discussed in Section 5.1 , we use ↵t = p 1 2t in our experiments for both the discrete-time and continuous-time models , i.e . variance-preserving diffusion processes . It is straightforward to verify that ↵2t and SNR ( t ) , as a function of ⌘ ( t ) , then simplify to : ↵2t = sigmoid ( ⌘ ( t ) ) ( 4 ) SNR ( t ) = exp ( ⌘ ( t ) ) ( 5 ) 3.3 Reverse time generative model . We define our generative model by inverting the diffusion process of Section 3.1 , yielding a hierarchical generative model that samples a sequence of latents zt , with time running backward from t = 1 to t = 0 . We consider both the case where this sequence consists of a finite number of steps T , as well as a continuous time model corresponding to T ! 1 . We start by presenting the discrete-time case . Given finite T , we discretize time uniformly into T timesteps ( segments ) of width ⌧ = 1/T . Defining s ( i ) = ( i 1 ) /T and t ( i ) = i/T , our hierarchical generative model for data x is then given by : p ( x ) = Z z p ( z1 ) p ( x|z0 ) TY i=1 p ( zs ( i ) |zt ( i ) ) . ( 6 ) With the variance preserving diffusion specification and sufficiently small SNR ( 1 ) , we have that q ( z1|x ) ⇡ N ( z1 ; 0 , I ) . We therefore model the marginal distribution of z1 as a spherical Gaussian : p ( z1 ) = N ( z1 ; 0 , I ) . ( 7 ) We wish to choose a model p ( x|z0 ) that is close to the unknown q ( x|z0 ) . Let xi and z0 , i be the i-th elements of x , z0 , respectively . We then use a factorized distribution of the form : p ( x|z0 ) = Y i p ( xi|z0 , i ) , ( 8 ) where we choose p ( xi|z0 , i ) / q ( z0 , i|xi ) , which is normalized by summing over all possible discrete values of xi ( 256 in the case of 8-bit image data ) . With sufficiently large SNR ( 0 ) , this becomes a very close approximation to the true q ( x|z0 ) , as the influence of the unknown data distribution q ( x ) is overwhelmed by the likelihood q ( z0|x ) . Finally , we choose the conditional model distributions as p ( zs|zt ) = q ( zs|zt , x = x̂✓ ( zt ; t ) ) , ( 9 ) i.e . the same as q ( zs|zt , x ) , but with the original data x replaced by the output of a denoising model x̂✓ ( zt ; t ) that predicts x from its noisy version zt . Note that in practice we parameterize the denoising model as a function of a noise prediction model ( Section 3.4 ) , bridging the gap with previous work on diffusion models [ Ho et al. , 2020 ] . The means and variances of p ( zs|zt ) simplify to a remarkable degree ; see Appendix A .

Summary: The paper proposed an improvement to DPM. The main contribution is summarized below.  - Sec 3: The diffusion and reverse processes in DPM are re-written with the proposed SNR function. It can be seen as a high-level framework of the noise schedule.  - Sec 4: The ELBO (for both discrete and continuous DPMs) is then re-written in terms of SNR. The authors had a few other findings: e.g., a larger T is better, and equivalence between continuous DPMs.  - Sec 5: The SNR function is parameterized by a neural network and targets at minimizing the estimation variance of ELBO.  - Sec 6: Additional features computed from Fourier transformation are added to data. - Sec 7: Experiments on CIFAR and ImageNet show improvements on BPD.