Maximum Likelihood Training of Score-Based Diffusion Models

1 Introduction . Score-based generative models [ 44 , 45 , 48 ] and diffusion probabilistic models [ 43 , 19 ] have recently achieved state-of-the-art sample quality in a number of tasks , including image generation [ 48 , 11 ] , audio synthesis [ 5 , 27 , 37 ] , and shape generation [ 3 ] . Both families of models perturb data with a sequence of noise distributions , and generate samples by learning to reverse this path from noise to data . Through stochastic calculus , these approaches can be unified into a single framework [ 48 ] which we refer to as score-based diffusion models in this paper . The framework of score-based diffusion models [ 48 ] involves gradually diffusing the data distribution towards a given noise distribution using a stochastic differential equation ( SDE ) , and learning the time reversal of this SDE for sample generation . Crucially , the reverse-time SDE has a closed-form expression which depends solely on the time-dependent gradient field ( a.k.a. , score ) of the perturbed data distribution . This gradient field can be efficiently estimated by training a neural network ( called a score-based model [ 44 , 45 ] ) with a weighted combination of score matching losses [ 23 , 56 , 46 ] as the objective . A key advantage of score-based diffusion models is that they can be transformed into continuous normalizing flows ( CNFs ) [ 6 , 15 ] , thus allowing tractable likelihood computation with numerical ODE solvers . ˚Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Compared to vanilla CNFs , score-based diffusion models are much more efficient to train . This is because the maximum likelihood objective for training CNFs requires running an expensive ODE solver for every optimization step , while the weighted combination of score matching losses for training score-based models does not . However , unlike maximum likelihood training , minimizing a combination of score matching losses does not necessarily lead to better likelihood values . Since better likelihoods are useful for applications including compression [ 21 , 20 , 51 ] , semi-supervised learning [ 10 ] , adversarial purification [ 47 ] , and comparing against likelihood-based generative models , we seek a training objective for score-based diffusion models that is as efficient as score matching but also promotes higher likelihoods . We show that such an objective can be readily obtained through slight modification of the weighted combination of score matching losses . Our theory reveals that with a specific choice of weighting , which we term the likelihood weighting , the combination of score matching losses actually upper bounds the negative log-likelihood . We further prove that this upper bound becomes tight when our score-based model corresponds to the true time-dependent gradient field of a certain reverse-time SDE . Using likelihood weighting increases the variance of our objective , which we counteract by introducing a variance reduction technique based on importance sampling . Our bound is analogous to the classic evidence lower bound used for training latent-variable models in the variational autoencoding framework [ 26 , 39 ] , and can be viewed as a continuous-time generalization of [ 43 ] . With our likelihood weighting , we can minimize the weighted combination of score matching losses for approximate maximum likelihood training of score-based diffusion models . Compared to weightings in previous work [ 48 ] , we consistently improve likelihood values across multiple datasets , model architectures , and SDEs , with only slight degradation of Fréchet Inception distances [ 17 ] . Moreover , our upper bound on negative log-likelihood allows training with variational dequantization [ 18 ] , with which we reach negative log-likelihood of 2.83 bits/dim on CIFAR-10 [ 28 ] and 3.76 bits/dim on ImageNet 32ˆ32 [ 55 ] with no data augmentation . Our models present the first instances of normalizing flows which achieve comparable likelihood to cutting-edge autoregressive models . 2 Score-based diffusion models . Score-based diffusion models are deep generative models that smoothly transform data to noise using a diffusion process , and synthesize samples by learning and simulating the time-reversal of this diffusion . The overall approach is illustrated in Fig . 1 . 2.1 Diffusing data to noise with an SDE . Let ppxq denote the unknown distribution of a dataset consisting of D-dimensional i.i.d . samples . Score-based diffusion models [ 48 ] employ a stochastic differential equation ( SDE ) to diffuse ppxq towards a noise distribution . The SDEs are of the form dx “ fpx , tqdt ` gptqdw , ( 1 ) where fp¨ , tq : RD Ñ RD is the drift coefficient , gptq P R is the diffusion coefficient , and w P RD denotes a standard Wiener process ( a.k.a. , Brownian motion ) . Intuitively , we can interpret dw as infinitesimal Gaussian noise . The solution of an SDE is a diffusion process txptqutPr0 , T s , where r0 , T s is a fixed time horizon . We let ptpxq denote the marginal distribution of xptq , and p0tpx1 | xq denote the transition distribution from xp0q to xptq . Note that by definition we always have p0 “ p when using an SDE to perturb the data distribution . The role of the SDE is to smooth the data distribution by adding noise , gradually removing structure until little of the original signal remains . In the framework of score-based diffusion models , we choose fpx , tq , gptq , and T such that the diffusion process txptqutPr0 , T s approaches some analytically tractable prior distribution πpxq at t “ T , meaning pT pxq « πpxq . Three families of SDEs suitable for this task are outlined in [ 48 ] , namely Variance Exploding ( VE ) SDEs , Variance Preserving ( VP ) SDEs , and subVP SDEs . 2.2 Generating samples with the reverse SDE . Sample generation in score-based diffusion models relies on time-reversal of the diffusion process . For well-behaved drift and diffusion coefficients , the forward diffusion described in Eq . ( 1 ) has an associated reverse-time diffusion process [ 1 , 16 ] given by the following SDE dx “ “ fpx , tq ´ gptq2∇x log ptpxq ‰ dt ` gptqdw̄ , ( 2 ) where w̄ is now a standard Wiener process in the reverse-time direction . Here dt represents an infinitesimal negative time step , meaning that the above SDE must be solved from t “ T to t “ 0 . This reverse-time SDE results in exactly the same diffusion process txptqutPr0 , T s as Eq . ( 1 ) , assuming it is initialized with xpT q „ pT pxq . This result allows for the construction of diffusion-based generative models , and its functional form reveals the key target for learning : the time-dependent score function ∇x log ptpxq . Again , see Fig . 1 for a helpful visualization of this two-part formulation . In order to estimate ∇x log ptpxq from a given dataset , we fit the parameters of a neural network sθpx , tq , termed a score-based model , such that sθpx , tq « ∇x log ptpxq for almost all x P RD and t P r0 , T s. Unlike many likelihood-based generative models , a score-based model does not need to satisfy the integral constraints of a density function , and is therefore much easier to parameterize . Good score-based models should keep the following least squares loss small JSMpθ ; λp¨qq : “ 1 2 ż T 0 Eptpxqrλptq ‖∇x log ptpxq ´ sθpx , tq‖ 2 2sdt , ( 3 ) where λ : r0 , T s Ñ Rą0 is a positive weighting function . The integrand features the well-known score matching [ 23 ] objective Eptpxqr‖∇x log ptpxq ´ sθpx , tq‖ 2 2s . We therefore refer to Eq . ( 3 ) as a weighted combination of score matching losses . With score matching techniques [ 56 , 46 ] , we can compute Eq . ( 3 ) up to an additive constant and minimize it for training score-based models . For example , we can use denoising score matching [ 56 ] to transform JSMpθ ; λp¨qq into the following , which is equivalent up to a constant independent of θ : JDSMpθ ; λp¨qq : “ 1 2 ż T 0 Eppxqp0tpx1|xqrλptq ∥∥∇x1 log p0tpx1 | xq ´ sθpx1 , tq∥∥22sdt . ( 4 ) Whenever the drift coefficient fθpx , tq is linear in x ( which is true for all SDEs in [ 48 ] ) , the transition density p0tpx1 | xq is a tractable Gaussian distribution . We can form a Monte Carlo estimate of both the time integral and expectation in JDSMpθ ; λp¨qq with a sample pt , x , x1q , where t is uniformly drawn from r0 , T s , x „ ppxq is a sample from the dataset , and x1 „ p0tpx1 | xq . The gradient ∇x1 log p0tpx1 | xq can also be computed in closed form since p0tpx1 | xq is Gaussian . After training a score-based model sθpx , tq with JDSMpθ ; λp¨qq , we can plug it into the reverse-time SDE in Eq . ( 2 ) . Samples are then generated by solving this reverse-time SDE with numerical SDE solvers , given an initial sample from πpxq at t “ T . Since the forward SDE Eq . ( 1 ) is designed such that pT pxq « πpxq , the reverse-time SDE will closely trace the diffusion process given by Eq . ( 1 ) in the reverse time direction , and yield an approximate data sample at t “ 0 ( as visualized in Fig . 1 ) . 3 Likelihood of score-based diffusion models . The forward and backward diffusion processes in score-based diffusion models induce two probabilistic models for which we can define a likelihood . The first probabilistic model , denoted as pSDEθ pxq , is given by the approximate reverse-time SDE constructed from our score-based model sθpx , tq . In particular , suppose tx̂θptqutPr0 , T s is a stochastic process given by dx̂ “ “ fpx̂ , tq ´ gptq2sθpx̂ , tq ‰ dt ` gptqdw̄ , x̂θpT q „ π . ( 5 ) We define pSDEθ as the marginal distribution of x̂θp0q . The probabilistic model pSDEθ is jointly defined by the score-based model sθpx , tq , the prior π , plus the drift and diffusion coefficients of the forward SDE in Eq . ( 1 ) . We can obtain a sample x̂θp0q „ pSDEθ by numerically solving the reverse-time SDE in Eq . ( 5 ) with an initial noise vector x̂θpT q „ π . The other probabilistic model , denoted pODEθ pxq , is derived from the SDE ’ s associated probability flow ODE [ 32 , 48 ] . Every SDE has a corresponding probability flow ODE whose marginal distribution at each time t matches that of the SDE , so that they share the same ptpxq for all time . In particular , the ODE corresponding to the SDE in Eq . ( 1 ) is given by dx dt “ fpx , tq ´ 1 2 gptq2∇x log ptpxq . ( 6 ) Unlike the SDEs in Eq . ( 1 ) and Eq . ( 2 ) , this ODE describes fully deterministic dynamics for the process . Notably , it still features the same time-dependent score function ∇x log ptpxq . By approximating this score function with our model sθpx , tq , the probability flow ODE becomes dx̃ dt “ fpx̃ , tq ´ 1 2 gptq2sθpx̃ , tq . ( 7 ) In fact , this ODE is an instance of a continuous normalizing flow ( CNF ) [ 15 ] , and we can quantify how the ODE dynamics transform volumes across time in exactly the same way as these traditional flow-based models [ 6 ] . Given a prior distribution πpxq , and a trajectory function x̃θ : r0 , T s Ñ RD satisfying the ODE in Eq . ( 7 ) , we define pODEθ as the marginal distribution of x̃θp0q when x̃θpT q „ π . Similarly to pSDEθ , the model p ODE θ is jointly defined by the score-based model sθpx , tq , the prior π , and the forward SDE in Eq . ( 1 ) . Leveraging the instantaneous change-of-variables formula [ 6 ] , we can evaluate log pODEθ pxq exactly with numerical ODE solvers . Since pODEθ is a CNF , we can generate a sample x̃θp0q „ pODEθ by numerically solving the ODE in Eq . ( 7 ) with an initial value x̃θpT q „ π . Although computing log pODEθ pxq is tractable , training pODEθ with maximum likelihood will require calling an ODE solver for every optimization step [ 6 , 15 ] , which can be prohibitively expensive for large-scale score-based models . Unlike pODEθ , we can not evaluate log p SDE θ pxq exactly for an arbitrary data point x . However , we have a lower bound on log pSDEθ pxq which allows both efficient evaluation and optimization , as will be shown in Section 4.2 .

Summary: The authors show that for a specific weighting of score matching losses at various noise levels, the objective upper bounds the negative log-likelihood. They further bound the log-likelihood of individual datapoints to develop a continuous-time generalization of the ELBO.  Using the proposed objective the authors optimize score-based diffusion models towards likelihood. The authors further make use of 1) importance sampling of time steps for reduced variance, 2) variational dequantization for a tighter bound on the likelihood and 3) the use of the equivalent ODE during evaluation, which allows exact likelihood computation (apart from the dequantization) for an approximation to the SDE model. As a result of this, state-of-the-art likelihood performance is presented on CIFAR-10 and ImageNet 32x32, beating even the best autoregressive models.