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

Abstract:
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.

text: 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.

heading: 1 Introduction
text: 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.

heading: 2 Background
text: 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.

heading: 3 Marginal density and stochastic instantaneous change of variable
text: 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.

heading: 4 Inferring latent Brownian motion
text: As our goal is to estimate likelihood, we would like to compute the log density value using (11). However, this involves integrating out all possible Brownian paths, which is intractable. To resolve this, we view the Brownian motion as a latent variable, and perform inference by assigning higher probability to sample paths that are more likely to generate the observation. One can view this as a VAE, except we have an infinite dimensional latent variable.
Formally, let (⌦,F ,P) be the underlying probability space for which B0s is a Brownian motion. Suppose Q is another probability measure on (⌦,F) equivalent to P; that is, P and Q are similar in
the sense that they have the same measure zero sets. This allows us to apply the change-of-measure trick and lower bound the log-likelihood with a finite quantity using Jensen’s inequality:
log p(x, T ) EQ " log
dP dQ + log p0(YT )
Z T
0 r · µ ds
Y0 = x # . (13)
Note that dPdQ is the Radon-Nikodym derivative of P wrt Q. When both measures are absolutely continuous wrt a third measure, say Lebesgue, then the derivative can be expressed as the ratio of the two densities. However, since we are dealing with an infinite dimensional space, we are immediately faced with the following problems:
1. Is there a measure Q (equiv. to P) for which dPdQ can be easily computed, or at least numerically approximated?
2. Can we find a reparameterization (similar to the Gaussian reparameterization) of B0s under the new law Q to estimate the gradient needed for training?
We resort to the Girsanov theorem, which describes a general framework for dealing with the change of measure of Gaussian random variables under additive perturbation. It allows us to consider the law of a diffusion process as Q. See Appendix E for an explanation using the more familiar notion of probability densities.
Theorem 2 (Girsanov theorem, Theorem 8.6.3 of Øksendal (2003)). Let B̂s be an Itô process solving
dB̂s = a(!, s) ds+ dB 0 s, (14)
for ! 2 ⌦, 0  s  T and B̂0 = 0, where a(!, s) satisfies the Novikov’s condition E h exp ⇣ 1 2 R T 0 a 2 ds ⌘i < 1. Then B̂s is a Brownian motion wrt Q where
dQ dP (!) := exp
Z T
0 a(!, s) · dB0s
1
2
Z T
0 ||a(!, s)||22 ds
! . (15)
Equation (14) provides a standarization formula of B0s under Q, which means we can “invert” it to reparameterize B0s. This leads to the following lower bound. Theorem 3 (Continuous-time ELBO). Let Q be defined via the density (15). Then the RHS of (13) can be rewritten as
E " 1
2
Z T
0 ||a(!, s)||22 ds+ log p0(YT )
Z T
0 r · µ ds
Y0 = x # =: E1, (16)
where the expectation is taken wrt the Brownian motion B̂s, and Ys solves4
dY = ( µ+ a) ds+ dB̂s. (17)
We call Ys solving (17) the inference SDE, and E1 the continuous-time ELBO (CT-ELBO). Remark 2 (Computation). This lower bound can be numerically estimated by using any black box SDE solver, by augmenting the dynamic of y with the accumulation of ||a||2 and r · µ. Computing the divergence term r · µ directly can be expensive, but it can be efficiently estimated using the Hutchinson trace estimator (Hutchinson, 1989) along with reverse-mode automatic differentiation, similar to Grathwohl et al. (2018). As the parameters of both the generative and inference models are decoupled from the random variable B̂s, their gradients can be estimated via the reparameterization trick (Kingma & Welling, 2014; Rezende et al., 2014). Furthermore, backpropagation can be computed using an adjoint method with a constant memory cost (Li et al., 2020).
Remark 3 (Drift a). (i) In general, the drift term of the approximate posterior to the latent Brownian motion can be amortized, so that it will encode the information of individual datum x. (ii) The regularization ||a||2 ensures that a is kept close to 0, since it represents the deviation of the measure
4Note that µ and run backward in time from T , whereas a runs forward.
it induces (i.e. Q) from the classical Wiener measure (which is a centered Gaussian measure). (iii) When the diffusion coefficient is 0, the inference SDE reduces to the reverse dynamic of the generative ODE, and if a ⌘ 0 in this case, the lower bound is tight. (iv) There is generally no constraint on the form of a(!, s), so one can potentially augment it with additional dimensions to have a non-Markovian inference SDE. For simplicity, we let the inference SDE be a Markovian model, i.e. a = a(y, s). This is justified by the following theorem.
Theorem 4 (Variational gap and optimal inference SDE). The variational gap can be written as
log p(x, T ) E1 =
Z T
0 E ⇥ ||a(!, s) >r log p(Ys, T s)|| 2 ⇤ ds. (18)
In particular, E1 = log p(x, T ) if and only if a(!, s) can be written as a(!, s) = a(Ys(!), s) for almost every s 2 [0, T ] and ! 2 ⌦, and a(y, s) = >r log p(y, T s) almost everywhere.
Remark 4 (Variational gap). Even though the inference SDE seemingly takes a simple form, it is sufficiently flexible in that this type of variational problem can be generally solved by taking the supremem over all progressively measurable processes a(!, s) (Boué et al., 1998). In fact, the above theorem shows that E1 = log p(x, T ) if and only if a(y, s) = >r log p(y, T s). This means a non-amortized Markovian inference process is powerful enough.

heading: 5 Infinitely deep hierarchical VAE
text: Before we make the connection to score matching, we formally address the common belief that “diffusion models can be viewed as the continuous limit of hierarchical VAEs” (Tzen & Raginsky, 2019), and show that the CT-ELBO consistently extends their discrete-time counterpart. We do so by inspecting the ELBO of a hierarchical VAE defined as discretized5 generative and inference SDEs. We assume the generative model (i.e. the decoder) follows the transitional distributions
p(xi+1|xi) = N (xi+1; µ̃i(xi), ̃ 2 i ) (19)
̃ 2 i = t 2(i t), (20)
where t = T/L is the step size and L is the number of layers. For the inference model (i.e. the encoder), we assume
q(xi|xi+1) = N (xi; µ̂i+1(xi+1), ̂ 2 i+1) (21)
µ̂i(x) = x+ t( µ(x, i t) + (i t)a(x, T i t)) ̂ 2 i = t 2(i t). (22)
These transition kernels constitute a hierarchical variational autoencoder of L stochastic layers, whose marginal likelihood can be lower bounded by
log p(xL) Eq " log p(x0) + L 1X
i=0
log p(xi+1|xi)
q(xi|xi+1)
# =: EL, (23)
which we refer to as the discrete-time ELBO (DT-ELBO). The reconstruction error of the stochastic layer can be seen as some form of finite difference approximation to differentiation, which gives rise to r · µ in the CT-ELBO in the infinitesimal limit (as t approaches 0). The regularization of ||a||2 pops up when we compare the difference between µ̃i and µ̂i using the Gaussian reparameterization to compute the reconstruction error. We formalize this idea in the following theorem.
Theorem 5 (Consistency). Assume µ, , 2, a, ||a||2 and their derivatives up to the fourth order are all bounded and continuous, and that is non-singular. Then EL ! E1 as L ! 1.
This theorem tells us that the CT-ELBO we derive for continuous-time diffusion models is not that different from the traditional ELBO, and that maximizing the CT-ELBO can be seen as training an infinitely deep hierarchical VAE. We present the proof in Appendix F, which formalizes the above intuition, using Taylor’s theorem to control the polynomial approximation error, which will go to 0 as the step size t vanishes when the number of layers L increases to infinity.

heading: 6 Score-based generative modeling
text: Recall that our goal is to analyze the plug-in reverse SDE and draw connection to score matching. To this end, we reparameterize the generative (4) and inference (17) SDEs as
dX = (gg>s✓ f) dt+ g dBt and dY = f ds+ g dB̂s, (24)
by letting a = g>s✓, where the time variable is reversed (T t) for the generative process, and forward in time (s) for inference. The ELBO (16) can be rewritten as
E 1 = EYT [log p0(YT ) |Y0 = x]
Z T
0 EYs
 1
2 ||s✓||2gg> +r · (gg >s✓ f) Y0 = x ds. (25)
Comparing the integrand to the implicit score matching loss in Table 1, we immediately see that the network s✓ approximates r log q(y, s), the score function of the marginal density of Ys. That is, matching the score of q(y, t) amounts to maximizing the lower bound on the marginal likelihood of the plug-in reverse SDE.
Recently, Durkan & Song (2021)6 also attempt to establish the equivalency between maximum likelihood and score matching, by showing the following relationship between the forward KL divergence and a weighted sum of score matching loss (aka the Fisher divergence):
DKL(q(y, 0)||r(y, 0)) = 1
2
Z T
0 Eq(·,s)
h ||r log r(Ys, s) r log q(Ys, s)|| 2 gg> i ds, (26)
where r(y, s) is the density of Ys solving the same inference SDE with the initial condition y0 ⇠ r(·, 0), assuming q(y, T ) = r(y, T ). However, it is inaccurate to claim that score matching is equivalent to maximum likelihood. This is because if we simply let r(y, 0) = p(y, T ), i.e. the density of the generative SDE evaluated at y, r(y, s) will not necessarily be the same as either p(y, T s) or s✓(y, s). This means the KL divergence is not equal to the integral of the weighted score matching loss E[ 12 ||s✓ r log q|| 2 gg> ]. In fact, the latter corresponds to a lower bound on the likelihood (the cross-entropy term of the KL) up to some constant, as equation (25) suggests.
More generally, we can apply our analysis to a family of plug-in reverse SDEs indexed by some parameter  1:
dX = 1 2
g 2s✓ f dt+ p 1 g dBt and dY = f
2 g
2 r log q ds+ p 1 g dB̂s,
(27)
5We follow the Euler-Maruyama (EM) scheme. Other discretization scheme may also work; we leave that for future work.
6This refers to the v1 of the paper on arXiv. This version was later on replaced with a new version where they derived a similar bound as ours.
where we assume g is diagonal for simplicity. We defer the formal discussion to Appendix G, but the essence is that this inference SDE induces the same marginal distribution as (1), and the generative SDE is its corresponding plug-in reverse. Equation (27) includes the original plug-in reverse SDE (24) and an equivalent ODE as special cases with = 0 and = 1. Denote its corresponding CT-ELBO by E1 . Specifically, (25) becomes E
1 0 . Then we have the following relationship.
Theorem 6 (Plug-in reverse SDE ELBO, abridged). For < 1,
EY0 [E1 ] = EY0 [E10 ] ✓
2
4(1 )
◆Z T
0 EYs
 1
2 ||s✓(Ys, s) r log q(Ys, s)||2g2 ds (28)
We state the full theorem in Appendix H, where we rearrange the terms to show that the average CT-ELBO of the -plug-in reverse SDE is also equivalent to the ISM loss, similarly to (25) but up to some multiplying and additive constants. The implication is that while minimizing the score matching loss, we implicitly maximize the likelihood of a continuum of plug-in reverse SDEs which include the ODE as a limiting case ( ! 1). See Figure 2 (right) for illustration. This suggests the likelihood of the equivalent ODE can be improved by minimizing the score matching loss, as the ODE’s likelihood will be close to plug-in reverse SDEs with ⇡ 1, which explains the good likelihood of the equivalent ODE reported in Song et al. (2021). In practice, we can only estimate the ELBO of the case = 0 since otherwise there will be some constant we do not have access to, but their gradients can all be estimated via score matching.

heading: 6.1 Computational trade-off
text: Having a general framework for estimating the likelihood of diffusion processes allows us to compare a wide family of models, including continuous-time flows and plug-in reverse SDEs trained by score matching. We compare the two by measuring the negative ELBO throughout training to highlight their computation-estimation trade-off. We train the models on the Swiss roll toy data. For continuous-time flow, we set = 0, using the Hutchinson trace estimator following Grathwohl et al. (2018). The ELBO in this case is tight since a will be penalized to be 0. We use the torchdiffeq library (Chen et al., 2018) for numerical integration for fairer comparison7. For plug-in reverse SDE, we train the drift network a using SSM and DSM (for DSM the loss is weighted to reduce variance, which introduces some bias; see the next subsection). We use the variance-preserving inference SDE from Song et al. (2021), which allows us to sample Ys using a closed form formula, for s sampled uniformly between [0, T ]. The trained models are visualized in Figure 1, the learning curves presented in Figure 3.
From the learning curve figures, we see that neg-likelihood decreases rapidly for the continuous-time flow in terms of the number of parameter updates. But once the x-axis is normalized by runtime, the convergence speed becomes almost indistinguishable. This is because for continuous-time flows, numerical integration takes time, whereas for plug-in reverse SDEs, we train on a random time step s; that is, within a fixed amount of time the latter can make more parameter updates at the cost of noisier gradients. Note that both models have constant memory cost (wrt T or L, the number of integration steps), so a large batch size can be used to reduce variance for training.

heading: 6.2 Bias and variance trade-off
text: The integral in equation (25) can be estimated by sampling (Ys, s), and using the Hutchinson trace estimator to estimate the divergence, which corresponds to implicit score matching. However, in practice the variance of this estimator is very high when the norm of the Jacobian rs✓ is large. Another popular approach is to use the denoising estimator (recall the identity from (3)),
EYs  1
2 ||s✓(Ys, s) r log q(Ys|Y0)||2gg> Y0 = x . (29)
The inference SDE is typically chosen so that Ys can be easily sampled, e.g. following N (µs, 2s), where µs and s are functions of Y0 and s. In this case, if we reparameterize Ys = µs + s✏ where
7Black-box SDE solvers such as torchsde (Li et al., 2020) might not be optimized for the deterministic case, since their stochastic adjoint method scales O(L logL) in time whereas deterministic numerical solvers are usually faster. This matters for our runtime comparison.
✏ ⇠ N (0, I), then the score becomes r log q = ✏ s . Since s ! 0 as s ! 0, this estimator normally has unbounded variance. Song & Ermon (2019); Song et al. (2021) propose to remedy this by multiplying the DSM loss by 2s/g2 (assuming g is a scalar for simplicity), so that the target has constant magnitude on average E[ 12 || ss✓ + ✏||
2], which would result in a biased gradient estimate with much smaller variance. We can debias this estimator by sampling s ⇠ q(s) / g2/ 2s . This ratio, however, is usually not normalizable in practice (as it integrates to 1). As an alternative, we consider the following unnormalized density q̃✏(s) = g2(s✏)/ 2s✏ for s 2 [0, s✏], and q̃✏(s) = g 2(s)/ 2s for s 2 [s✏, T ]. We experiment with this debiased procedure by sampling s ⇠ q✏ / q̃✏, for f and g chosen to be the variance-preserving SDE. s✏ is small so that the bias is negligible.
We train the model on MNIST (LeCun et al., 1998) and CIFAR10 (Krizhevsky et al., 2009). We present the learning curves and the standard error of the estimate of the ELBO in Figure 4. The lower bound is estimated using the Hutchinson trace estimator with s sampled uniformly from [0, T ], with the same batch size, so the only thing that will affect the dispersion is the magnitude of rs✓. Since smaller values of s are more likely to be sampled under q✏, the debiased model will see samples with less perturbation more often. On the contrary, sampling s uniformly will bias the model to learn from noisier data, causing the learned score to be smoother. We also experiment with parameterizing s✓ vs parameterizing a. We find the latter parameterization to be helpful since the relationship s✓ = g 1a has the effect of negating the multiplier s in the reweighted loss, i.e. E[ 12 || s g a + ✏||
2]. This is similar to the noise conditioning technique introduced in Song & Ermon (2020).

heading: 7 Related work
text: Diffusion-based generative models Our work lays a theoretical foundation for Song et al. (2021), which recognizes that conditional denoising score matching (Song & Ermon, 2019, 2020) and
discrete-time diffusion-based generative models (Sohl-Dickstein et al., 2015; Goyal et al., 2017; Ho et al., 2020) can be viewed as learning to revert an inference process (using the plug-in reverse SDE). Different from Ho et al. (2020), which shows the ELBO of discrete time diffusion process can be likened to DSM (Section 3.2 of the paper), we show that ISM loss naturally arises from the Fokker-Planck equation of the marginal density, via the Fenman-Kac representation and the Girsanov change of measure. This line of work has been successfully applied to modeling high dimensional natural images (Dhariwal & Nichol, 2021; Saharia et al., 2021), audio (Kong et al., 2020), 3D point cloud (Cai et al., 2020; Zhou et al., 2021), and discrete data (Hoogeboom et al., 2021).
Time-reversal of diffusion processes Plenty of works have studied the reverse-time diffusion processes (2), including Anderson (1982); Föllmer (1985); Elliott & Anderson (1985); Haussmann & Pardoux (1986). These are different from our marginal-equivalent (reverse) processes (27) when s✓ = r log q, since the latter is related by the marginals only.
Score matching for energy-based models Besides the connection to diffusion models, score matching is also often used as a method for learning energy based models (EBM)— see Song & Kingma (2021) for a comprehensive review on useful techniques—. When used as an EBM, sampling from the conditional score model can be achieved by running the annealed Langevin diffusion (Neal, 2001), which is connected to free-energy estimaton in physics (Jarzynski, 1997), wherein the path integral is essentially a Feynman Kac representation.
De Bruijn’s identity To connect maximum likelihood and score matching, Durkan & Song (2021) shows that KL divergence can be represented as an integral of weighted Fisher divergence, generalizing the case of Lyu (2009) where the inference perturbation is a simple Brownian motion. This type of formulas fall into the category of de Bruijn’s identity (Cover, 1999) for relative entropy. A similar differential form result can be found in Wibisono et al. (2017).
Learning SDEs Tzen & Raginsky (2019); Li et al. (2020) also propose to learn a neural SDE by applying Girsanov’s theorem. The key difference is that they treat the SDE entirely as a latent variable, with an additional emission probability, whereas we use the Feynman-Kac formula to directly express the marginal density as an expectation, side-stepping the need to smooth out the density using the emission probability (which will be a Dirac point mass in our case). In their case, the inference direction is the same as the generative direction, since they infer the latent SDE directly, whereas we apply Girsanov to the Feynman-Kac diffusion (opposite the generative direction). Xu et al. (2021) further apply neural SDE as an infinitely deep Bayesian neural network.

heading: 8 Conclusion and Discussion
text: In this work, we derive a general variational framework for estimating the marginal likelihood of continuous-time diffusion models. This framework allows us to study a wide spectrum of models, including continuous-time normalizing flows and score-based generative models. Using our framework, we show that performing score matching with a particular choice of mixture weighting is equivalent to maximizing a lower bound on the marginal likelihood of a family of plug-in reverse SDEs, of which the one used in Song et al. (2021) and the equivalent ODE are special cases. Empirically, we validate our theory by monitoring the ELBO while performing score matching, and discuss the implication of the choice of mixture weighting and the potential of debiasing via nonuniform sampling. We emphasize that our theory does not explain the impressive sample quality of this family of models, which is still an open research problem and we leave it for future work.
This work introduces a general framework to estimate the likelihood of diffusion-based models, which allows the parameters of both the generative and inference SDEs to be learned, using a numerical solver with constant memory cost (as per Remark 2). The training time can be reduced via the connection to score matching and the reverse-time parameterization (24) as long as f and g take a simple form (so that Yt can be sampled without numerical integration). For example, one can generalize the Ornstein-Uhlenbeck process to have non-linear (in time) f and g, similar to the variance-presering SDE, by parameterizing the integral of f using a monotone network (Sill, 1998; Kay & Ungar, 2000; Daniels & Velikova, 2010; Huang et al., 2018). This has been explored in a concurrent work by Kingma et al. (2021) in a different framework.

heading: Acknowledgements
text: We would like to thank David Kanaa, Ricky Chen, Simon Verret, Rémi Piché-Taillefer, Alexia Jolicoeur-Martineau, and Faruk Ahmed for giving their feedback on this manuscript. We would also like to thank the INNF+ 2021 reviewers and NeurIPS 2021 reviewers for their constructive suggestions, which help us improve the clarity of the paper. Chin-Wei is supported by the Google PhD fellowship.
We also acknowledge the Python community (Van Rossum & Drake Jr, 1995; Oliphant, 2007) for developing the tools that enabled this work, including numpy (Oliphant, 2006; Van Der Walt et al., 2011; Walt et al., 2011; Harris et al., 2020), PyTorch (Paszke et al., 2019), Matplotlib (Hunter, 2007), seaborn (Waskom et al., 2018), and SciPy (Jones et al., 2014).

Review summaries:
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.

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.