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: The paper develops a variational framework for general diffusion processes. A continuous-time ELBO is developed that extends the ELBO of discrete-time diffusion processes. They use this to show that optimizing the score matching loss corresponds to maximizing a lower bound on the log marginal likelihood.  More specifically, the paper makes use of the Fokker-Planck/Kolmogorov forward equation along with the Feynman-Kac formula to express the marginal density of the generative diffusion process. Since the resulting expression is intractable, a change-of-measure is applied, introducing a variational approximation $\mathbb{Q}$ to $\mathbb{P}$. $\mathbb{Q}$ is chosen such that the Radon-Nikodym derivative $\frac{d\mathbb{Q}}{d\mathbb{P}}$ can be expressed using Girsanov theorem. From this, a continuous-time ELBO is obtained. Finally, through a specific choice of the drift coefficient $a$ in the inference SDE, a connection between the ELBO and score matching is established.