Score-Based Generative Modeling through Stochastic Differential Equations

1 INTRODUCTION . Two successful classes of probabilistic generative models involve sequentially corrupting training data with slowly increasing noise , and then learning to reverse this corruption in order to form a generative model of the data . Score matching with Langevin dynamics ( SMLD ) ( Song & Ermon , 2019 ) estimates the score ( i.e. , the gradient of the log probability density with respect to data ) at each noise scale , and then uses Langevin dynamics to sample from a sequence of decreasing noise scales during generation . Denoising diffusion probabilistic modeling ( DDPM ) ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ) trains a sequence of probabilistic models to reverse each step of the noise corruption , using knowledge of the functional form of the reverse distributions to make training tractable . For continuous state spaces , the DDPM training objective implicitly computes scores at each noise scale . We therefore refer to these two model classes together as score-based generative models . Score-based generative models , and related techniques ( Bordes et al. , 2017 ; Goyal et al. , 2017 ; Du & Mordatch , 2019 ) , have proven effective at generation of images ( Song & Ermon , 2019 ; 2020 ; Ho et al. , 2020 ) , audio ( Chen et al. , 2020 ; Kong et al. , 2020 ) , graphs ( Niu et al. , 2020 ) , and shapes ( Cai ˚Work partially done during an internship at Google Brain . Forward SDE ( data → noise ) Reverse SDE ( noise → data ) score function Figure 1 : Solving a reversetime SDE yields a score-based generative model . Transforming data to a simple noise distribution can be accomplished with a continuous-time SDE . This SDE can be reversed if we know the score of the distribution at each intermediate time step , ∇x log ptpxq . et al. , 2020 ) . To enable new sampling methods and further extend the capabilities of score-based generative models , we propose a unified framework that generalizes previous approaches through the lens of stochastic differential equations ( SDEs ) . Specifically , instead of perturbing data with a finite number of noise distributions , we consider a continuum of distributions that evolve over time according to a diffusion process . This process progressively diffuses a data point into random noise , and is given by a prescribed SDE that does not depend on the data and has no trainable parameters . By reversing this process , we can smoothly mold random noise into data for sample generation . Crucially , this reverse process satisfies a reverse-time SDE ( Anderson , 1982 ) , which can be derived from the forward SDE given the score of the marginal probability densities as a function of time . We can therefore approximate the reverse-time SDE by training a time-dependent neural network to estimate the scores , and then produce samples using numerical SDE solvers . Our key idea is summarized in Fig . 1 . Our proposed framework has several theoretical and practical contributions : Flexible sampling and likelihood computation : We can employ any general-purpose SDE solver to integrate the reverse-time SDE for sampling . In addition , we propose two special methods not viable for general SDEs : ( i ) Predictor-Corrector ( PC ) samplers that combine numerical SDE solvers with score-based MCMC approaches , such as Langevin MCMC ( Parisi , 1981 ) and HMC ( Neal et al. , 2011 ) ; and ( ii ) deterministic samplers based on the probability flow ordinary differential equation ( ODE ) . The former unifies and improves over existing sampling methods for score-based models . The latter allows for fast adaptive sampling via black-box ODE solvers , flexible data manipulation via latent codes , a uniquely identifiable encoding , and notably , exact likelihood computation . Controllable generation : We can modulate the generation process by conditioning on information not available during training , because the conditional reverse-time SDE can be efficiently estimated from unconditional scores . This enables applications such as class-conditional generation , image inpainting , colorization and other inverse problems , all achievable using a single unconditional score-based model without re-training . Unified framework : Our framework provides a unified way to explore and tune various SDEs for improving score-based generative models . The methods of SMLD and DDPM can be amalgamated into our framework as discretizations of two separate SDEs . Although DDPM ( Ho et al. , 2020 ) was recently reported to achieve higher sample quality than SMLD ( Song & Ermon , 2019 ; 2020 ) , we show that with better architectures and new sampling algorithms allowed by our framework , the latter can catch up—it achieves new state-of-the-art Inception score ( 9.89 ) and FID score ( 2.20 ) on CIFAR-10 , as well as high-fidelity generation of 1024ˆ 1024 images for the first time from a score-based model . In addition , we propose a new SDE under our framework that achieves a likelihood value of 2.99 bits/dim on uniformly dequantized CIFAR-10 images , setting a new record on this task . 2 BACKGROUND . 2.1 DENOISING SCORE MATCHING WITH LANGEVIN DYNAMICS ( SMLD ) . Let pσpx̃ | xq : “ N px̃ ; x , σ2Iq be a perturbation kernel , and pσpx̃q : “ ş pdatapxqpσpx̃ | xqdx , where pdatapxq denotes the data distribution . Consider a sequence of positive noise scales σmin “ σ1 ă σ2 ă ¨ ¨ ¨ ă σN “ σmax . Typically , σmin is small enough such that pσminpxq « pdatapxq , and σmax is large enough such that pσmaxpxq « N px ; 0 , σ2maxIq . Song & Ermon ( 2019 ) propose to train a Noise Conditional Score Network ( NCSN ) , denoted by sθpx , σq , with a weighted sum of denoising score matching ( Vincent , 2011 ) objectives : θ˚ “ arg min θ N ÿ i “ 1 σ2i EpdatapxqEpσi px̃|xq “ ‖sθpx̃ , σiq ´∇x̃ log pσipx̃ | xq‖ 2 2 ‰ . ( 1 ) Given sufficient data and model capacity , the optimal score-based model sθ˚px , σq matches ∇x log pσpxq almost everywhere for σ P tσiuNi “ 1 . For sampling , Song & Ermon ( 2019 ) run M steps of Langevin MCMC to get a sample for each pσipxq sequentially : xmi “ xm´1i ` isθ˚px m´1 i , σiq ` ? 2 iz m i , m “ 1 , 2 , ¨ ¨ ¨ , M , ( 2 ) where i ą 0 is the step size , and zmi is standard normal . The above is repeated for i “ N , N ´ 1 , ¨ ¨ ¨ , 1 in turn with x0N „ N px | 0 , σ2maxIq and x0i “ xMi ` 1 when i ă N . As M Ñ 8 and i Ñ 0 for all i , xM1 becomes an exact sample from pσminpxq « pdatapxq under some regularity conditions . 2.2 DENOISING DIFFUSION PROBABILISTIC MODELS ( DDPM ) . Sohl-Dickstein et al . ( 2015 ) ; Ho et al . ( 2020 ) consider a sequence of positive noise scales 0 ă β1 , β2 , ¨ ¨ ¨ , βN ă 1 . For each training data point x0 „ pdatapxq , a discrete Markov chain tx0 , x1 , ¨ ¨ ¨ , xNu is constructed such that ppxi | xi´1q “ N pxi ; ? 1´ βixi´1 , βiIq , and therefore pαipxi | x0q “ N pxi ; ? αix0 , p1´ αiqIq , where αi : “ śi j “ 1p1´ βjq . Similar to SMLD , we can denote the perturbed data distribution as pαipx̃q : “ ş pdatapxqpαipx̃ | xqdx . The noise scales are prescribed such that xN is approximately distributed according to N p0 , Iq . A variational Markov chain in the reverse direction is parameterized with pθpxi´1|xiq “ N pxi´1 ; 1 ? 1´βi pxi ` βisθpxi , iqq , βiIq , and trained with a re-weighted variant of the evidence lower bound ( ELBO ) : θ˚ “ arg min θ N ÿ i “ 1 p1´ αiqEpdatapxqEpαi px̃|xqr‖sθpx̃ , iq ´∇x̃ log pαipx̃ | xq‖ 2 2s . ( 3 ) After solving Eq . ( 3 ) to get the optimal model sθ˚px , iq , samples can be generated by starting from xN „ N p0 , Iq and following the estimated reverse Markov chain as below xi´1 “ 1 ? 1´ βi pxi ` βisθ˚pxi , iqq ` a βizi , i “ N , N ´ 1 , ¨ ¨ ¨ , 1 . ( 4 ) We call this method ancestral sampling , since it amounts to performing ancestral sampling from the graphical model śN i “ 1 pθpxi´1 | xiq . The objective Eq . ( 3 ) described here is Lsimple in Ho et al . ( 2020 ) , written in a form to expose more similarity to Eq . ( 1 ) . Like Eq . ( 1 ) , Eq . ( 3 ) is also a weighted sum of denoising score matching objectives , which implies that the optimal model , sθ˚px̃ , iq , matches the score of the perturbed data distribution , ∇x log pαipxq . Notably , the weights of the i-th summand in Eq . ( 1 ) and Eq . ( 3 ) , namely σ2i and p1´αiq , are related to corresponding perturbation kernels in the same functional form : σ2i91 { Er‖∇x log pσipx̃ | xq‖ 2 2s and p1´ αiq91 { Er‖∇x log pαipx̃ | xq‖ 2 2s . 3 SCORE-BASED GENERATIVE MODELING WITH SDES . Perturbing data with multiple noise scales is key to the success of previous methods . We propose to generalize this idea further to an infinite number of noise scales , such that perturbed data distributions evolve according to an SDE as the noise intensifies . An overview of our framework is given in Fig . 2 . 3.1 PERTURBING DATA WITH SDES . Our goal is to construct a diffusion process txptquTt “ 0 indexed by a continuous time variable t P r0 , T s , such that xp0q „ p0 , for which we have a dataset of i.i.d . samples , and xpT q „ pT , for which we have a tractable form to generate samples efficiently . In other words , p0 is the data distribution and pT is the prior distribution . This diffusion process can be modeled as the solution to an Itô SDE : dx “ fpx , tqdt ` gptqdw , ( 5 ) where w is the standard Wiener process ( a.k.a. , Brownian motion ) , fp¨ , tq : Rd Ñ Rd is a vectorvalued function called the drift coefficient of xptq , and gp¨q : RÑ R is a scalar function known as the diffusion coefficient of xptq . For ease of presentation we assume the diffusion coefficient is a scalar ( instead of a dˆ d matrix ) and does not depend on x , but our theory can be generalized to hold in those cases ( see Appendix A ) . The SDE has a unique strong solution as long as the coefficients are globally Lipschitz in both state and time ( Øksendal , 2003 ) . We hereafter denote by ptpxq the probability density of xptq , and use pstpxptq | xpsqq to denote the transition kernel from xpsq to xptq , where 0 ď s ă t ď T . Typically , pT is an unstructured prior distribution that contains no information of p0 , such as a Gaussian distribution with fixed mean and variance . There are various ways of designing the SDE in Eq . ( 5 ) such that it diffuses the data distribution into a fixed prior distribution . We provide several examples later in Section 3.4 that are derived from continuous generalizations of SMLD and DDPM .

Summary: This paper presents a generative model based on stochastic differential equations (SDEs), which generalizes two other score-based generative models score matching with Langevin dynamics (SMLD) and denoising diffusion probabilistic modeling (DDPM). The recipe to sample from the data distribution is based on (i) the observation that both SMLD and DDPM can be formulated as the discretization of an SDE, (ii) the finding from Anderson (1982) about the reverse of an Ito process, (iii) a score model. A novel aspect of the presented technique is the use of the score model as "predictor", which gives the initial sample from the MCMC sampler that serves as "corrector". Finally, the Ito process induced by the reverse SDE is formulated in a deterministic manner, leading to a neural-ODE based generative model.