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 introduces a family of diffusion-based generative models for density estimations on image datasets. Here, the forward diffusion is directly defined by a set of conditional densities of diffused random variables for a given data. In particular, the paper proposes that $t$-th density for a given data follows a normal distribution; its mean and variance will be determined by data and a monotonic function of $t$, called $SNR(t)$. The generative models aim at learning diffusion processes reversed in time by maximizing its ELBO.  The paper first draws connections between the proposed models and previous diffusion-based models. For example, the authors show that the proposed method's discrete-time versions include previous finite-length Markov chain models as special cases. The paper then demonstrates that the ELBOs of the proposed continuous-time models exist. In here, the resulting ELBOs include an integration (wrt time) of a reconstruction loss weighted by the time-derivative of $SNR(t)$. Based on this, the paper shows that the losses of the previous continuous-time diffusion-based models follow the same integration form but with different weighting functions.  Second, the paper presents interesting properties of the proposed method due to $SNR(t)$. In particular, remind that $SNR(t)$ is invertible because of its monotonicity, and the continuous-time loss's integration contains the time-derivative of $SNR(t)$. The paper shows that the integral can be re-written wrt $SNR(t)$'s output by change-of-variable. In this form, the integral only depends on the range of the integration but doesn't depend on the shape of $SNR(t)$ within the range. As a result, acknowledging the invariant of the loss on the $SNR(t)$'s output space, the paper proposes to parameterize the score models conditioning on log-scaled $SNR(t)$ instead of $t$.  Third, the paper proposes to train $SNR(t)$ for the discrete and continuous-time models. For the monotonicity, $SNR(t)$ is defined as a monotonic network wrt $t$. For discrete-time models, the paper proposes to train $SNR(t)$ networks by maximizing the ELBO together with the diffusion models. For continuous-time models, $SNR(t)$ is trained by minimizing the variance of the Monte-Carlo estimation of the continuous-time ELBOs.  Furthermore, the paper introduces various practical techniques to improve training the models. For example, instead of iid sampling of $t$, the authors propose to use a low-discrepancy sampler of $t$ within a mini-batch. The paper also proposes to use the Fourier feature of inputs: feeding the concatenation of inputs and its Fourier feature to diffusion models.  In the experiments, the paper demonstrates that the proposed models achieve SOTA likelihoods on image datasets, such as CIFAR-10 and ImageNet.