A Variational Perspective on Diffusion-Based Generative Models and Score Matching

Discrete-time diffusion-based generative models and score matching methods have shown promising results in modeling high-dimensional image data . Recently , Song et al . ( 2021 ) show that diffusion processes that transform data into noise can be reversed via learning the score function , i.e . the gradient of the logdensity of the perturbed data . They propose to plug the learned score function into an inverse formula to define a generative diffusion process . Despite the empirical success , a theoretical underpinning of this procedure is still lacking . In this work , we approach the ( continuous-time ) generative diffusion directly and derive a variational framework for likelihood estimation , which includes continuous-time normalizing flows as a special case , and can be seen as an infinitely deep variational autoencoder . Under this framework , we show that minimizing the score-matching loss is equivalent to maximizing a lower bound of the likelihood of the plug-in reverse SDE proposed by Song et al . ( 2021 ) , bridging the theoretical gap . 1 Introduction . Generative modeling can be thought of as inverting an inference process . If the inference process is invertible , then one can focus on transforming the data into a tractable distribution ( Dinh et al. , 2016 ) . If the inference process is deterministic yet non-invertible , one could learn to invert it stochastically ( Dinh et al. , 2019 ; Nielsen et al. , 2020 ) . Most generally , both inference and generation can be stochastic . This is known as the variational autoencoder ( Kingma & Welling , 2014 ; Rezende et al. , 2014 , VAE ) . Under the variational framework , one has a lot of flexibility in choosing the generative and inference models . Recent work on diffusion-based modeling ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ) can be thought of as removing one degree of freedom , by freezing the inference path . The inference model is a fixed discrete-time Markov chain , that slowly transforms the data into a tractable prior , such as the standard normal distribution . The generative model is another Markov chain that is trained to revert this process iteratively . Diffusion-based models have been shown to perform remarkably well on image synthesis ( Dhariwal & Nichol , 2021 ) , rivaling the performance of stateof-the-art Generative Adversarial Networks ( Brock et al. , 2018 ) . Song et al . ( 2021 ) connect diffusion-based model and score matching ( Hyvärinen & Dayan , 2005 ) , by looking at the stochastic differential equation ( SDE ) associated with the inference process . They realize that the dynamic of the inference process can be inverted if one has access to the score function of the perturbed data , by solving another SDE reversed in time . They then propose to learn the score function of the inference process and substitute the approximate score into the formula of the reverse SDE to obtain a generative model . We call the resulting generative model the plug-in reverse SDE . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Conceptually simple as this learning procedure may seem , little is known about how the score matching loss relates to the plug-in reverse SDE . In this paper , we propose a variational framework suitable for likelihood estimation for general generative diffusion processes , and use this framework to connect score matching with maximum likelihood . We do so by combining two important theorems in stochastic calculus : the Feynman-Kac formula for representing the marginal density of the generative diffusion as an expectation ( Section 3 ) , and the Girsanov theorem for performing inference in function space ( Section 4 ) . We derive a functional evidence lower bound that consistently extends discrete-time diffusion models to have infinite depth , i.e . the number of layers goes to infinity ( Section 5 ) . Finally , by reparameterizing our generative and inference SDEs , we obtain a training objective equivalent to minimizing the ( implicit ) score matching loss ( Section 6 ) . Our theory suggests that by matching the score , one actually maximizes a lower bound on the log marginal density of the plug-in reverse SDE , laying a theoretical foundation for this learning procedure . We further generalize our result to a family of marginal-equivalent plug-in reverse SDEs , including an equivalent ODE as a limiting case . Notation : We use ( Ys , s ) to denote the inference process ( where Y0 is the data ) , and ( Xt , t ) to denote the generative process ( where X0 is a random variable following an unstructured prior ) . We use s and t to distinguish the two directions , and always integrate the differential equations from 0 to T > 0 ( different from the literature , where sometimes one might see integration from T to 0 ) . B̂s and Bt denote the Brownian motions associated with the inference and generative SDEs , respectively . B 0 s is a reparameterization of B̂s ( see Section 4 ) . q ( y , s ) and p ( x , t ) denote the probability density functions of Ys and Xt , respectively . We let s✓ denote a time-indexed parameterized function that will be used to approximate the score r log q ( y , s ) . r is the gradient wrt the spatial variable ( x or y , which we sometimes call position ) , @ t , @ s and @ xi are partial derivatives , and H⇤ denotes Hessian . 2 Background . Assume Y0 follows the data distribution q ( y , 0 ) , and Ys satisfies the Itô SDE ( Øksendal , 2003 ) dY = f ( Y , s ) ds+ g ( Y , s ) dB̂s , ( 1 ) where f and g are chosen such that the density q ( y , s ) will converge to some tractable prior p0 as s ! T . Following Song et al . ( 2021 ) , we assume g is position-independent . It is possible to find a “ reverse ” SDE , whose marginal density evolves according to q ( y , s ) , reversed in time , for example1 dX = ( gg > r log q ( X , T t ) f ) dt+ g dBt . ( 2 ) If X0 ⇠ p0 , then the density p ( x , t ) of Xt is equal to q ( x , T t ) . This means that if we have access to the score function r log q , we can solve the above SDE to obtain XT d = Y0 . Song et al . ( 2021 ) 1See Appendix G for a family of equivalent ( reverse ) SDEs indexed by some parameter , of which equation ( 2 ) is a special case with = 0. propose to approximate the score via a parameterized score function s✓ by minimizing Z T 0 EYs 1 2 ||s✓ ( Ys , s ) r log q ( Ys , s ) ||2⇤ ( s ) ds where the expectation in the integral is known as the explicit score matching ( ESM ) loss LESM , and ⇤ ( s ) is a positive definite matrix2 that serves as a weighting function for the overall loss . LESM is not immediately useful , since we do not have access to the ground truth score r log q . A few alternative losses can be used , which are all equal to one another up to a constant , including implicit score matching ( Hyvärinen & Dayan , 2005 , ISM ) , sliced score matching ( Song et al. , 2020 , SSM ) , and denoising score matching ( Vincent , 2011 , DSM ) . The losses are summarized in Table 1 , and are related through the following identity ( see Appendix A for the derivation ) : LESM 1 2 I ( q ( ys , s ) ) = LISM = LSSM = LDSM 1 2 EY0 [ I ( q ( ys|y0 ) ) ] , ( 3 ) where I ( q ) = E [ ||r log q||2⇤ ] is a constant . After training , Song et al . ( 2021 ) plug s✓ into ( 2 ) to define a generative model . We refer to this SDE as the plug-in reverse SDE . The plug-in reverse SDE has been demonstrated to have impressive empirical results , but a theoretical underpinning of this learning framework is still lacking . For example , it is unclear how the training objective ( minimizing the score matching loss ) relates to the sampling procedure , e.g . whether the probability distribution induced by the plug-in reverse SDE gets closer to the data distribution in the sense of any statistical divergence or metric . We seek to answer the following question in this paper : How will minimizing the score-matching loss impact the plug-in reverse SDE ? We first provide a framework to estimate the likelihood of generative SDEs , and then get back to this question in Section 6 . 3 Marginal density and stochastic instantaneous change of variable . Let Xt be a diffusion process solving the following Itô SDE3 : dX = µ ( X , t ) dt+ ( X , t ) dBt ( 4 ) with the initial condition X0 ⇠ p0 , which induces a family of densities Xt ⇠ p ( · , t ) . We use this SDE as the generative SDE , and we are interested in log p ( x , T ) for maximum likelihood . The density p ( x , t ) follows the Kolmogorov forward ( or the Fokker Planck ) equation : @ tp ( x , t ) = X j @ xj [ µj ( x , t ) p ( x , t ) ] + X i , j @ 2 xi , xj [ Dij ( x , t ) p ( x , t ) ] ( 5 ) with the initial value p ( · , 0 ) = p0 ( · ) , where D = 12 T is the diffusion matrix . We can expand the Fokker Planck and rearrange the terms to obtain @ tp ( x , t ) = 2 4 r · µ ( x , t ) + X i , j @ 2 xi , xjDij ( x , t ) 3 5 p ( x , t ) + X i 2 4 µi ( x , t ) + 2 X j @ xjDij ( x , t ) 3 5 @ xip ( x , t ) + X i , j Dij ( x , t ) @ 2 xi , xjp ( x , t ) ( 6 ) 2We use this matrix to induce a Mahalanobis norm ||x||2⇤ : = x > ⇤x , which will be used in Section 6 . 3For generality , we use the notation µ and to describe a generative SDE , which will be set to g2s✓ f and g when we come back to the discussion of the plug-in reverse SDE in Section 6. so that all coefficients of the same order are grouped together . For simplicity , we assume the diffusion term is independent of x throughout the paper . Then ( 6 ) reduces to @ tp ( x , t ) = ( r · µ ( x , t ) ) p ( x , t ) µ ( x , t ) > rp ( x , t ) +D ( t ) : Hp ( x , t ) ( 7 ) where : denotes the Frobenius inner product between matrices . Even with this simplification , solving ( 7 ) is not trivial . Fortunately , we can estimate this quantity using the Feynman-Kac formula , which tells us that the solution of certain second-order linear partial differential equations have a probabilistic representation . Theorem 1 ( Feynman-Kac representation , Chapter 5.7 of Karatzas & Shreve ( 2014 ) ) . Let T > 0 . Let y and & be the spatial and temporal arguments to the function v 2 C2,1 ( Rd ⇥ [ 0 , T ] ) solving @ & v + cv + b > rv +A : Hv = 0 ( 8 ) with the terminal condition v ( y , T ) = h ( y ) , where A = 12⌘⌘ > for some matrix-valued function ⌘ ( y , & ) . Under the assumption stated in Appendix B , if B0s is a Brownian motion and Ys solves dY = b ( Y , s ) ds+ ⌘ ( Y , s ) dB0s , ( 9 ) with the initial datum Y & = y , then v ( y , & ) = E `` h ( YT ) exp Z T & c ( Ys , s ) ds ! Y & = y # . ( 10 ) To estimate the density p ( · , T ) of ( 7 ) , we can apply the change of variable p ( x , t ) : = v ( x , T t ) by letting the Feynman-Kac ( F-K ) coefficients correspond to their Fokker-Planck ( F-P ) counterparts according to Table 2 . This way , solving ( 8 ) backward is equivalent to solving ( 7 ) forward , and we have the following representation of the marginal density at T : p ( x , T ) = E `` p0 ( YT ) exp Z T 0 r · µ ( Ys , T s ) ds ! Y0 = x # , ( 11 ) where Ys is a diffusion process solving dY = µ ( Y , T s ) ds+ ( T s ) dB0s . ( 12 ) Remark 1 ( Marginalization ) . This representation can be interpreted as a mixture of continuous time flows . Assume a sample path of the Brownian motion is given , and we are interested in how the density evolves following the dynamic ( 4 ) . In the infinitesimal setting , it can been seen as applying the invertible map x 7 ! x+ µ ( x , t ) t+ ( t ) Bi , where Bi : = B ( i+1 ) t Bi t is the Brownian increment . Since the diffusion term is independent of the spatial variable , it can be seen as a constant additive transformation , which is volume preserving , so it will not be taken into account when computing the change of density . The only contribution to the change of density will be from the log-determinant of the Jacobian of id+µ t , which means we can simply apply the instantaneous change of variable formula ( Chen et al. , 2018 ) . This will be the conditional density given the entire { Bt : t 0 } , and marginalizing it out results in the expectation in ( 11 ) . See Appendix C for details . Our framework also works with the general case where depends on x , but the formulae need to be adapted to account for the spatial partial derivatives . See Appendix D for the derivation .

Summary: This paper presents a framework for studying continuous diffusion models. Thanks to this it demonstrates the plug in generative model corresponds to a continuous version of the discrete diffusion models and is equivalent to minimising an ELBO on the likelihood of the model.  It does so by first transforming the change of density associated to the SDE into a second order linear differential equation and then expressing its solution as an expectation over another SDE. Solving this expectation is performed via variation inference on the Brownian paths and yields the continuous-time ELBO. This CT-ELBO is then compared with its discrete time counter part. Finally it is compared to the score matching loss which demonstrates the soundness of such method for learning the generative model and continuous time diffusion models.