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 paper studies score-based denoising-diffusion generative models and makes one major contribution. It shows how the training objective, a mixture or score-matching losses defined along different times of the diffusion, can be adapted to enable maximum likelihood training (using a lower bound objective). This is rigorously derived and proven analytically. The result is a specific weighting of the score-matching losses for different times t. The paper validates the derived methodology and empirically demonstrates very strong likelihood performance, state-of-the-art among models without data augmentations.