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 generalizes a family of score-based generative models that rely on sequences of noise scalings of the data and extends them to the continuous domain, which leads to an SDE-based framework. By using score-matching, the forward SDE, which transforms data into a tractable noise distribution, can be reversed and thus used as a generative model. This is then improved by employing a two-phase algorithm with a prediction step, followed by a tunable number of correction steps. Further, reformulating the problem as a neural ODE allows for exact likelihood computations and reduces the number of required function evaluations. The framework enables unconditional, as well as conditional samples.