Score-based generative modeling secretly minimizes the Wasserstein distance
Anonymous Author(s)
AffiliationAddress email
Abstract
Score-based generative models are shown to achieve remarkable empirical performances in various applications such as image generation and audio synthesis.
However, a theoretical understanding of score-based diffusion models is still incomplete.
Recently, Song et al.
showed that the training objective of score-basedgenerative models is equivalent to minimizing the Kullback-Leibler divergence ofthe generated distribution from the data distribution.
In this work, we show thatscore-based models also minimize the Wasserstein distance between them undersuitable assumptions on the model.
Specifically, we prove that the Wassersteindistance is upper bounded by the square root of the objective function up to multiplicative constants and a fixed constant offset.
Our proof is based on a novelapplication of the theory of optimal transport, which can be of independent interestto the society.
Our numerical experiments support our findings.
By analyzing ourupper bounds, we provide a few techniques to obtain tighter upper bounds
Introduction
The idea of gradually perturbing the data distribution to define the forward process and learningthe reverse process has recently demonstrated remarkable performances in generative modeling.
Score-based generative models [33, 34, 35] and diffusion probabilistic models [30, 17] have beenused in numerous applications, including image generation [26, 12], and audio synthesis [9, 21].
Both model families have the following two processes.
The given data distribution is incrementallyperturbed by adding small noise in the forward process.
This forward process can be understood asthe standard diffusion process [35].
The reverse process requires the so-called score function [2], thegradient of the log probability density function.
To approximate the score function, it is common tominimize a weighted combination of score matching losses.
Once the score function is approximated,one can generate samples from randomly generated noise by applying the learned reverse process.
Note that score-based models do not explicitly minimize any sort of similarity between the modeldistribution (the distribution of the generated samples) and the data distribution.
Therefore, itis unclear whether or not scored-based generative models can minimize divergence (or distance)between the model distribution and the data distribution.
In very recent work [32], the authors showthat optimizing the score matching loss is indeed equivalent to minimizing an upper bound on the
Kullback–Leibler divergence (KL divergence).
Therefore, it is expected that the model distribution issimilar (in the sense of KL divergence) to the data distribution if the score matching loss is smallenough.
A natural question follows – do score-based models minimize any other divergence or distancethan the KL divergence?
At first glance, this seems unlikely.
For instance, Generative Adversarial
Networks (GANs) [16] were shown to minimize different divergences depending on the choiceof the loss function.
The original GAN loss function turns out to capture the Jensen–Shannon divergence [16], while the W-GAN loss function captures the Wasserstein distance [3].
We refer thereaders to [10] for more details.
On the other hand, a unique aspect of score-based models is that the score function is defined in thesame space where data lies in.
In stark contrast, all the other popular generative models such as GANs[16], VAEs [19], and Normalizing Flows [20] use loss functions that only care about likelihoods anddo not directly operate in the knowledge of the underlying data space.
In this work, we answer this question by revealing the following property of score-based modeling:
Score-based modeling minimizes the Wasserstein distance.
We note that this is quite surprising as there is no direct relationship between the KL divergence andthe Wasserstein distance, in general.
Therefore, this finding is not implied by the aforementionedfinding.
[Figure]
Using the theory of optimal transport, we investigate the evolution of data distribution based on
Partial differential equations (PDEs).
This allows us to estimate the Wasserstein distance betweenthe given data distribution and the distribution of the generated samples.
Our work theoreticallyguarantees that the score-based generative models approximate data distributions in the space ofprobability measures equipped with the Wasserstein distance.
Our main contributions and findingsare summarized as follows:
• (Theorem 1 & Corollary 1) Using the theory of optimal transport, we prove that the Wasser-
stein distance between the data distribution and the model distribution is bounded from
above.
The upper bound is given as the square root of the score-matching loss J SM up
to multiplicative constants and a fixed constant offset.
The main idea is to estimate the
time-derivative of the Wasserstein metric by using the form of the continuity equations.
• (Theorem 2)
We prove J SM ≤ J DSM (as defined in eq. (5) and eq.
(14), respectively)under suitable assumptions, which gives us a more practical upper bound since J DSM ismuch easier to estimate in practice and is widely used as a loss function to train score-based models. • (Theorem 3)
We analyze how the model distribution changes if the data distribution is perturbed initially.
This guarantees that the score based-model is robust under noise. • (Section 4)
Our numerical experiments support the upper bound of our analysis.
We also
explore the effect on the weight decay coefficient in training and different choices of the number of total timesteps T .
Notations: x = ( x 1 , x 2 , · · · , x d ) ⊤ ∈ R d and t ∈ [0 , T ] denote the spatial variable and the timevariable, respectively.
Here, R d is the d -dimensional Euclidean space and T > 0 is the time horizon.
The gradient of a real-valued function ρ with respect to the spatial variable and the time-derivativeof ρ are denoted by ∇ ρ = (cid:16) ∂ρ∂x 1 , ∂ρ∂x 2 , · · · , ∂ρ∂x d (cid:17) ⊤ , and ∂ t ρ , respectively.
The Laplacian of ρ isgiven as ∆ ρ = ∇ · ( ∇ ρ ) .
Here, ∇ · F = ∂F 1 ∂x 1 + ∂F 2 ∂x 2 + · · · + ∂F d ∂x d denotes the divergence of
F = ( F 1 , F 2 , · · · , F d ) with respect to the spatial variable x . The Hessian of ρ is a d × d matrix given by ( D 2 ρ ) ij = ∂
2 ρ ∂x i ∂x j for 1 ≤ i, j ≤ d . We indicate Gaussian distribution with mean µ andcovariance Σ by N ( µ, Σ) , and L 2 norm by ∥ · ∥ .
We denote the data distribution and the distributionof the generated samples by p 0 and q 0 .
Preliminaries and related works
In this section, we analyze the forward process and the reverse process in the score-based modelthrough Partial differential equations (PDEs).
In particular, we investigate the time evolution ofmarginal distributions through the two processes.
2.1 The forward process
As in [35], the forward process can be modeled by the standard diffusion process { x ( t ) }
Tt =0 fort ∈ [0 , T ] such that x (0) ∼ p 0 and
[Equation]
Here, w is the standard Wiener process, f : R d × [0 , T ] → R d is the drift coefficient of x ( t ) , andg : [0 , T ] → R is the diffusion coefficient of x ( t ) .
Let p t = p ( · , t ) be a probability density function of the random variable x ( t ) on R d at time t ∈ [0 , T ] .
Then, the function p : R d × [0 , T ] → [0 , + ∞ ) satisfies the following partial differential equations(PDEs), namely the Fokker–Planck equation or the forward Kolmogorov equation as in [5, Theorem2.2]:
[Equation] [Equation]
) where the pdf p 0 represents the data distribution.
Under suitable assumptions on f and g , a solutionto the above initial value problem uniquely exists.
We refer the readers to [6] for more details.
2.2 Understanding the reverse process through PDEs
To generate samples from t
= T to 0 , the score-based model uses the reverse-time SDE [2]:
[Equation]
Here, ¯ w is a standard Wiener process when time flows backward from t =
T to t = 0 .
As in the previous section, for given x ( T ) ∼ p T , a probability density function of the random variablex ( t ) from the above process (3) satisfies the final value problem:
[Equation]
Using this simple identity, ∆ p = ∇ · ( ∇ p ) = ∇ · (cid:16) p ∇ p p (cid:17) =
∇ · ( p ∇ (log p )) , it can be seen that thedifferential equations in (2) and (4) are indeed identical.
2.3 The objective function
The score of the probability distribution, ∇ log p t , is approximated by a function s θ : R d × [0 , T ] → 100 R d .
For instance, in [32], the authors consider the weighted mean squared error (MSE):
[Equation] [Equation]
where λ : [0 , T ] → (0 , ∞ ) is a positive weighting function.
Replacing ∇ log p in (4) with s θ , we deduce a new reverse process.
For given noise distribution q T , aprobability density function q t = q ( · , t ) at time t ∈ [0 , T ] solves
[Equation]
Here, we introduced a new function q : R d × [0 , T ] → [0 , + ∞ )
2.4 Related works
Optimal transport The Wasserstein distance has a solid connection to the Fokker–Planck equation.
In particular, the Fokker–Planck equation can be understood as the gradient flow in the space ofprobability measures equipped with the 2 -Wasserstein distance [1, 18, 27].
For two probabilitymeasures µ and ν on R d , we define the 2 -Wasserstein or Monge-Kantorovich distance as
[Equation]
Here, Π( µ, ν ) is the set of all couplings π of µ and ν where a Borel probability measure π on R d × R dsatisfies π ( A × R d ) = µ ( A ) and π ( R d × B ) =
ν ( B ) for all measurable subsets A, B ⊂ R d .
Under suitable assumptions on the drift coefficient and the diffusion coefficient, a solution to the
Fokker–Planck equation satisfies the so-called Wasserstein contraction property [7, 8, 29], whichplays an important role in our analysis.
More discussions can be found in Section 3.3 and Section 5.
Score matching and diffusion probabilistic models Score matching models [33, 35] use the scorefunction to reverse the noise corruption in a continuous time range by solving stochastic differentialequations (SDEs).
Diffusion probabilistic models [30, 17] adopt a similar idea of modeling tractabledenoising process to reverse each discrete step of noise corruption with probabilistic models.
In fact,[35] show that diffusion models are score matching models with certain forms of SDE, which unifiesthe whole framework.
Diffusion models are gaining increasing attention: recent advances [31, 26, 12]show that diffusion models can generate even better images than GANs while enjoying a more stabletraining process, which encourages it to be widely applied in various tasks, such as text-to-image andimage editing [25, 24], audio generation [9, 21], etc.
Approximation in the Wasserstein space
In this section, we estimate the Wasserstein distance between a given data distribution p 0 and themarginal of the generated samples q 0 obtained by the score-based model. The proofs of the mainresults are deferred to Appendix A and Appendix D.
The key ideas are illustrated in Section 5.
3.1 The main results
We present our main results under the following assumptions. (A1)
The drift coefficient f : R d ×[0 , T ] → R d is Lipschitz continuous in the space variable x : there exists a positive constant L f ( t ) ∈(0 , ∞ ) , depending on t ∈ [0 , T ] , such that for all x, y ∈ R d
[Equation]
(A2) s θ : R d × [0 , T ] → R d satisfies the one-sided Lipschitz condition [13, Definition 2.1]: thereexists a constant L s ( t ) ∈ R , depending on t ∈ [0 , T ] , satisfying
[Equation]
for all x, y ∈ R d . The usual Lipschitz constant of s θ satisfies the above inequality (9).
Unlike theusual Lipschitz constant, the one-sided Lipschitz constant L s ( t ) is not necessarily nonnegative (see Appendix B).
More technical assumptions can be found in Appendix A.
We now present a generalresult.
Theorem 1.
Let p 0 and q 0 be the data distribution and a pdf obtained by the score-based model asin Section 2.
Then the Wasserstein distance between p 0 and q 0 is bounded as follows: [Equation] [Equation]
Here, a function I : [0 , T ] → (0 , + ∞ ) defined by
[Equation]
is an integrating factor of the differential inequality (22) given later in Proposition 1. The main ideaof the proofs is to estimate the derivatives of the Wasserstein distance along with the time evolution,which is illustrated in Section 5.
The first term on the right-hand side of (10) naturally arises from the estimation of W 2 ( p 0 , q 0 ) .
It is similar to a loss function J SM commonly used in the score-based models. Applying the Cauchy-Schwarz inequality, we obtain the upper bound in terms of the weighted MSE J SM .
Corollary 1.
Let p 0 and q 0 be given in Theorem 1.
If λ ( t ) =
g ( t ) 2 for all t ∈ [0 , T ] , then thefollowing inequality holds:
[Equation]
We defer the proofs of Theorem 1, and Corollary 1 to Appendix A.
Remark 1.
From the above result, the Wasserstein distance converges to zero as the loss function
J SM goes to zero if p T = q T .
However, due to the tension between I ( T ) and W 2 ( p T , q T ) , a generalconvergence result remains open (see Section 6 and Appendix H).
Remark 2.
If g 4 I 2 λ − 1 is integrable in [0 , T ] , the similar inequality as in (32) holds for other choicesof λ . See Corollary 4 in Appendix A.
Remark 3.
Note that the above upper bound depends on the time horizon T .
Numerical experimentsfor different choices of T are provided in Section 4.3.2. We further investigate the offset I ( T ) W 2 ( p T , q T ) for Denoising Diffusion Probabilistic Models [17]. Corollary 2. Under the same setting as in Corollary 1, let f ( t ) =
− β ( t ) 2 x , g ( t ) = σ (cid:112) β ( t ) , andq T = ϕ ∼ N (0 , σ 2 I ) . Then, we have [Equation]
Here, σ > 0 , and β : [0 , T ] → (0 , + ∞ ) is a positive function.
The proof of Corollary 2 and the analysis of L s can be found in Appendix H. As the arguments aretechnical, we provide a formal analysis below.
Remark 4. If s θ is close enough to ∇ log p t ( x ) , then we have L s ( t ) ∼ ( − σ − 2 + ∥ D 2 (log h t ) ∥ ∞ ) . Here, h t := p t /ϕ exponentially converges to 1 as t → ∞ (see Lemma 4).
Thus, L s ( t ) converges to− σ − 2 , as observed in our experiments with σ = 1 in Fig. 4a. Consequently, the second term in (13)decays as T increases (see Fig. 4b). 3.2 The relation between J SM and J DSM The above results give the guarantee that J SM can upper-bound Wasserstein distance.
However, J SMis generally intractable due to the unknown quantity ∇ x log p t ( x ) [32], and is rarely used in practice.
Alternatively, there is a conditional version of the score-matching loss J DSM , which is equivalentto J SM up to a constant, as shown in [36].
J DSM is also widely used as a loss function to trainscore-based models:
[Equation]
where p 0 t indicates the conditional probability of x at step t given x (0) .
Now we present our newresult, which shows that J DSM is also an upper bound of J SM under the certain assumption.
Theorem 2.
If p 0 t satisfies
[Equation]
then we have
[Equation]
It is worth noting that Denoising Diffusion Probabilistic Models (DDPM)
[17] satisfy eq.
Corollary 3.
Let p ( x ( t ) | x ( t − 1))
=
N ( √ 1 − β t x
( t − 1) , β t I ) as in Denoising Diffusion Probabilistic Models (DDPM)
[17], which corresponds to f ( x, t ) = − 12 β t x , and g ( t ) = √ β t .
Here,β t = β ( t ) > 0 is the noise schedule for all t , and let λ ( t ) =
g 2 ( t ) .
Then eq.
(16) holds.
The proofs of Theorem 2 and Corollary 3 are in Appendix D. We also provide numerical resultswhich show J SM ≤ J DSM for DDPM in Appendix E.
Remark 5.
Using Theorem 2 and Corollary 3, we can present an upper bound of the Wassersteindistance in terms of J DSM , where both W 2 ( p 0 , q 0 ) and J DSM are tractable.
This enables us toempirically verify this upper bound via DDPM in Section 4.
3.3 Perturbation and sensitivity analysis
Let ˜ p 0 be a small perturbation of the original data distribution p 0 .
For a given corrupted distribution˜ p 0 , denote the probability density functions in the forward process and the reverse process by ˜ p t and˜ q t , respectively.
We investigate how this perturbation changes the score-based model.
It is well known from the Wasserstein contraction property [8] that the Wasserstein distance betweentwo probability distributions in the forward process can be estimated as follows: for all t ∈ [0 , T ]
[Equation]
[Equation]
In the reverse process, the score function in the drift coefficients may vary, and thus the aboveinequality does not directly give us the upper bound on W 2 ( q 0 , ˜ q 0 ) .
Instead, our analysis in theprevious sections allows us to estimate W 2 ( q 0 , ˜ q 0 ) .
It is worth noting that the Wasserstein distance satisfies the triangle inequality, unlike other statisticalmeasures, including the Kullback–Leibler (KL) divergence.
Theorem 3.
Let ˜ p 0 be the corrupted data distribution and ˜ q 0 be a probability density functionobtained by the score-based model starting from ˜ p 0 .
It holds that
[Equation]
Here, ˜ J SM and ˜ I denote the weighted MSE and the integrating factor from the perturbed model.
Remark 6.
The Wasserstein distance between the original data distribution p 0 and the modeldistribution ˜ q 0 starting from the corrupted data ˜ p 0 can also be estimated similarly.
Experiments
Here we present experiments on simulation datasets to verify the upper bound in eq.
4.1 Experiment settings
Datasets Here we adopt three 2D datasets for simulation: One cluster Gaussian N ( 0 , 0 .
1 I ) , twomoons in [28], and four clusters Gaussian mixture N (( ± 0 .
5 , ± 0 .
5) ⊤ , 0 .
01 I ) with equal weights foreach cluster.
Visualization for these datasets is in Fig 2.
For forward dynamics, we use discrete timesteps with T = 10 .
For t ∈ [1 , 10] , p ( x ( t ) | x ( t − 1)) =
N ( √ 1 − β t x ( t − 1) , β t I ) as in DDPM[17], i.e., f ( x, t ) = − 1 2 β t x , and g ( t ) = √ β t .
For β t , weadopt a sigmoid schedule from β 1 = 10 − 5 to β T = 10 − 2 .
Training and evaluation We use a 3-layer neural network as the score matching model, with ReLUnonlinearity and skip-connection at the final output.
Each layer is composed of a linear layer with 128hidden neurons, and an embedding layer for 10 timesteps.
For optimizer, we use AdamW [22] with learning rate = 0 .
001 and weight decay coefficient 0 .
To evaluate Wasserstein distance, we use POT [15] to compute W 2 ( p 0 , q 0 ) at the end of each training epoch.
To avoid the term of W 2 ( p T , q T )in eq. (16) for verification, we first perform forward pass on sampled x (0) to get x ( T ) , and performdata generation on the same x ( T ) to fix q T = p T and W 2 ( p T , q T ) = 0 .
For loss function, we use J DSM with λ ( t ) =
g 2 ( t ) and batch size = 512.
To get both lower and higher losses to verify ourbound, we first maximize the loss for 15 epochs, and then minimize the loss till convergence.
4.2 Results
To verify the upper bound in eq.
(16), we take log on each side of the equation with W 2 ( p T , q T ) =
0and get:
[Equation]
We present a log-log plot to verify the slope and the intercept in the upper bound, respectively: see
Fig 2.
The log-log plots are obtained by evaluating Wasserstein distance and the loss function duringtraining.
The red dots are the sampled data from each dataset, the green dots are generated data afterloss maximization, and the purple dots are the generated data sampled at convergence.
From thelog-log plots, we can observe that the theoretical line (the black line) is a valid upper bound of theempirical one (the blue dots).
[Figure]
4.3 Exploring the upper bound
Here we explore different factors that can affect the tightness of the theoretical upper bound.
Asshown in Fig 2, the gap between the theoretical bound and the empirical one mainly results from theintercept in eq.
Recall our intercept is 12 log (cid:16)(cid:82) T
g 2 ( t ) I ( t ) 2 dt (cid:17) for I ( t ) given in (11).
4.3.1 Effect of weight regularization
Since g ( t ) and L f are fixed, only the one-sided Lipchitz constant L s ( r ) of the model s θ is subject tothe training process, which is estimated via grid search in our experiments.
Notice that the loss function J DSM aims to minimize the distance between s θ ( x, t ) and∇
x log p 0 t ( x | x (0)) .
Thus, L s at convergence is determined by dynamics ∇ x log p 0 t ( x | x (0)) .
Ifthe one-sided Lipschitz constant for ∇ x log p 0 t ( x | x (0)) is small, then the intercept would be expected to be small (see example in Fig 2a of Gaussian cluster with moderately high variance). Ifthe one-sided Lipschitz constant for ∇ x log p 0 t ( x | x (0)) is large, to minimize the loss, the interceptwould also be large (see example in Fig 2c of Gaussian clusters with lower variance).
However, it canbe possible to make the upper bound tighter while sacrificing the training loss by constraining L s .
To constrain L s , we adopt spectral normalization (as used in [4] to regularize Lipschitz constant),weight clipping (used in [3] to regularize Lipschitz constant), and simple weight decay.
As we canobserve in Fig. 3, loss at convergence will increase when we impose very strong regularization onthe weights, while the intercept decreases.
Importantly, the intercept approaches the empirical onewithout not exceeding it.
See more examples of weight decay in Appendix G, where we also presentthe log-log training plots for weight decay = 0 .
5 , 1 , 5 respectively.
[Figure]
4.3.2 Effect of the number of total timesteps
[Figure]
In our Remark 4, we argue that the offset in our upper bound would decay to 0 when choosingappropriate SDE and sufficiently large T .
Here we provide empirical evidence for our claim withoutfixing p T = q T .
We use the four clusters Gaussian dataset and fix β to be sigmoid schedule between[10 − 5 , 10 − 2 ] for any T .
In Fig.
4, we plot L s ( t ) for T = 100 to show that L s ( t ) would convergeto − 1 when t → T .
As a result, since W 2 ( p T , q T ) decays exponentially (see eq. 48), we also showthat as T increases, our offset I ( T ) W 2 ( p T , q T ) would also decay in a rate that is approximatelyexponential when T is sufficiently large, as the model converges to the real data distribution since
W 2 ( p 0 , q 0 ) ≈ 0 .
The main idea of the proofs
We investigate the time evolution of probability measures arising in the score-based models.
Usingthe optimal transport theory, we differentiate the Wasserstein distances, W 2 ( p t , q t ) with respect totime.
The main idea is to express these equations (2) and (6) in the form of the continuity equation:
[Equation]
(20) where v is a so-called velocity field.
Let us denote the velocity fields of p , and q by
[Equation]
Then, the equations (2), and (6) can be written in the form of (20) with v [ p ] and v [ q ] , respectively.
Using these vector fields, we compute the time derivative of the Wasserstein distance between twoevolving probability measures.
For velocity fields v [ ρ ( i ) ] depending on ρ ( i ) , let ρ ( i ) for i = 1 , 2 betwo solutions of the continuity equations ∂ t ρ ( i ) + ∇ · ( ρ ( i ) v [ ρ ( i ) ])
= 0 .
Then, the relation betweenthe Wasserstein distance and velocity fields arises from the theory of optimal transport:
[Equation]
T I(T)W(p_T,q_T) W(p_0,q_0) where π t is an optimal transport plan between ρ (1) t and ρ (2) t .
See [1, Theorem 8.4.7] and [29, Corollary5.25].
This equality plays an important role when we estimate the Wasserstein distance.
The followingproposition is a consequence of (21) and Lemma 1 in Appendix C.
Proposition 1.
Let p t and q t be solutions to (2) and (6) , respectively, in [0 , T ] .
Then,
[Equation]
p t , q t ) ≤
[Equation]
L f + L s g 2 ) W
[Equation] p t , q t ) + g 2 b 12 where b ( t ) : [Equation] E p
[Equation]
Discussions
Convergence of the Wasserstein distance If p T =
q T , then our result yields that the Wassersteindistance W 2 ( p 0 , q 0 ) converges to zero as J SM goes to zero.
However, if p T is not equal to q T andf and g are not explicitly given, then we find that it is not easy to obtain such a convergence result.
This is due to the nature of the Langevin dynamics in the reverse process.
Typically, additionalassumptions such as the log-concavity are required to guarantee the Langevin dynamics convergence,for instance, [11].
In addition, the dependence of the neural network s θ makes it hard to evaluate the precise upperbound.
We expect that appropriate assumptions on s θ and the initial distribution p 0 may resolve thisissue, but we do not pursue them in this paper.
Estimating L s
Estimating our theoretical bound requires estimating the one-sided Lipschitz constant L s , which is NP-hard [37], making it difficult to estimate in high dimensional data.
Note that thetwo-sided Lipschitz constant is an upper bound of the one-sided one, which can be more efficientlyestimated as shown in [14].
However, it would result in a looser theoretical upper bound if we adoptthe two-sided Lipschitz constant, and estimating the one-sided Lipschitz constant does not enjoy thesemi-definiteness of estimating the two-sided one in [14].
Conclusion
In this paper, we mainly discuss the relationship between the loss function of score-based modelsand the Wasserstein distance.
With the help of the theory in optimal transport, we present a novelupper bound on Wasserstein distance in terms of the loss function used in practice.
We also providenumerical experiments to verify our theoretical bound with various simulation datasets.
Moreover,we explore factors that affect the estimation of the upper bound and show the corresponding influenceon the training process.
A Assumptions and Proofs
In addition to (A1) and (A2) in Section 3.1, we assume the following:(A3) For the initial distribution p 0 of (2) and the final distribution q T of (6), the following quantitiesare finite:
[Equation]
where Λ( x ) := log max( ∥ x ∥ , 1) . (A4) There exist constants C and C 2 such that
[Equation]
(A5) There exists a positive constant C 3 such that
[Equation]
(A6) (cid:82) T
[Equation]
{+E p t [ f 2 ] dt and (cid:82) T 
+}
E q t [( f − g 2 s θ ) 2 ] dt are finite. The above assumptions guarantee the integrability of the score function ∇ log p . Proposition 2. [6, Theorem 4.2.1]
A solution p t to (2) satisfies ∇ p t ∈ L 1 ( R d ) and [Equation] T [Equation]
[Equation]
We give proofs in smooth settings.(A7) p 0 , q T ∈ C 2 ( R d ) , and g ∈ C 2 ([0 , T ]) .
f and s θ are C 2 in both x ∈ R d and t ∈ [0 , T ] .(A8) A solution p t to (2) and a solution q t to (6) are positive C 2 probability densities decaying rapidlyin R d : there exists k > 0 such that p t ( x ) = O (exp( −∥ x ∥ k )) and q t ( x ) = O (exp( −∥ x ∥ k )) for all t ∈[0 , T ] .
{+Furthermore, the vector fields v [ p t ] = f − 12 g 2 ∇ log p t , and v [ q t ] = (cid:0) f − g 
 s θ (cid:1) + 12 g 2 ∇ log q tgrow at most linearly as ∥ x ∥ → ∞ . Proof of Theorem 1. Recall the differential inequality (22) below from Proposition 1:+}
[Equation] It can be solved by introducing the integrating factor [Equation] where b ( t ) :
= E p t (cid:2) ∥∇ log p t ( x ) − s θ ( x, t ) ∥ 2 (cid:3) . As ddt I ( t ) = ( L f ( t ) + L s ( t ) g 2 ( t )) I ( t ) , the above inequality (22) can be written as
[Equation]
Integrating both sides from 0 to T , we obtain that [Equation] As I (0) = 1 , we conclude that
[Equation]
Proof of Corollary 1.
Let [Equation] Here, I ( t ) and b ( t ) = E p t (cid:2) ∥∇ log p t ( x ) − s θ ( x, t ) ∥ 2 (cid:3) are given in (27).
The Cauchy-Schwarzinequality yields
[Equation]
Since λ = g 2 and J SM given in (5) satisfies
[Equation]
we conclude (32):
[Equation]
Remark 7.
It is expected that the parallel results hold for weak solutions to (2) and (4) by usingsuitable approximations.
This approach has been studied for the Wasserstein contraction property in[8, Section 6.2].
From (30) and (31) in the proof of Corollary 1, we obtain the following result for other choices of λ .
Corollary 4.
Let p 0 and q 0 be given in Theorem 1.
Suppose that g 4 I 2 λ − 1 is integrable in [0 , T ] .
Then the following inequality holds:
[Equation]
B One-sided Lipschitzness
For an arbitrary Lipschitz function F : R d → R d , the Cauchy-Schwarz inequality yields that
[Equation]
where L is the Lipschitz constant of F .
Therefore, all Lipschitz function satisfies the one-sided
Lipschitz condition:
[Equation]
As pointed out earlier in Section 3.1, the one-sided Lipschitz constant is not necessarily to be positive.
For instance, if F ( x ) = − ax + b for a > 0 and b ∈ R d , then − a < 0 can be the one-sided Lipschitzconstant of F while its Lipschitz constant is a > 0 .
Figure 5 visualizes this.
C Technical lemmas
Lemma 1.
Let π t be an optimal transport plan between p t and q t .
Then, we have
E π t [
[Equation]
y
[Equation]
(35) where b ( t ) :=
E p t (cid:2) ∥∇ log p t ( x ) − s θ ( x, t ) ∥ 2 (cid:3) .
[Figure]
Proof.
The left-hand side of (35) is given by
[Equation]
In Lemma 2 below, we prove that the last term is less than or equal to zero.
In what follows, we estimate the first two terms.
First, using the Lipschitzness of f in space, we get
[Equation]
The last equality follows from the fact that π t is an optimal plan between p t and q t .
Next, the second term g 2 E π t [( x − y ) · ( ∇ log p t ( x ) − s θ ( y, t ))] is the sum of the following two terms:
[Equation]
and [Equation] As shown above, the former one I 1 is bounded from above by g 2 L s W 22 ( p t , q t ) . It suffices to find an upper bound on the latter one I 2 .
By the Cauchy-Schwarz inequality, we have
[Equation]
As the marginals of π t are p t and q t , it holds that
[Equation]
As a consequence, we conclude that
[Equation]
where b ( t ) =
E p t (cid:2) ∥∇ log p t ( x ) − s θ ( x, t ) ∥ 2 (cid:3) .
Before proving the lemma below, let us recall some basic definitions from the theory of optimaltransport.
The Wasserstein distance defined in (7) has an equivalent formulation:
[Equation]
The optimizer of the above problem is called the optimal map from µ to ν .
It is well known that thereexists a convex function ϕ such that T = ∇ ϕ .
Lemma 2.
E π t [( x − y ) · ( ∇ log q t ( y ) − ∇ log p t ( x ))] is nonpositive.
Proof.
Let T t be an optimal transport map from p t to q t and a convex function ϕ t satisfy ∇ ϕ t =
T tfor all t ∈ [0 , T ] .
As in the proof of [7, Theorem 1], we have
[Equation]
[Equation]
where ϕ ∗ t is a convex conjugate of ϕ t .
The convexity of ϕ t yields that ∆ ϕ t + ∆ ϕ ∗ t ( ∇ ϕ t ) − 2 d andwe conclude.
D Proofs of Theorem 2 and Corollary
Definitions For a given t , define
[Equation]
38) and [Equation] Lemma 3. (Appendix in [36])
[Equation] [Equation] Proof of Theorem 2. From [36], we have
[Equation]
Note that
[Equation]
] = E x [Equation] x log p t (
[Equation]
⊤ ( [Equation] x log [Equation]
( x ))] − E x (0) [ E x | x (0) [( ∇ x log p 0 t ( x | x (0))) ⊤ ( ∇ x log p 0 t ( x | x (0))) | x (0)]] = Var[( ∇ x log p t ( x ))
⊤ ] − E
[Var[( ∇ x log p 0 t ( x | x (0)))
⊤ ] | x (0)]
+ ( E x
[Equation]
where the inequality comes from the law of total variance and our condition: [Equation] Then we have [Equation]
= ( E x [ ∇ x log p t [Equation] )] − E x [Equation]
E x | x (0) [ ∇ x log p 0 t ( x | x (0))
| x (0)]) ⊤ (
[Equation]
+
[Equation]
= 0 , (44) where first inequality comes from Jensen’s inequality, and the last equality comes from eq.
Recall that J SM ( θ, λ ) =
(cid:82) T 0 λ ( t )
J SM ( θ, t ) dt and J DSM ( θ, λ ) =
(cid:82) T 0 λ ( t ) J DSM ( θ, t )
dt , λ ( t ) > 0 ,we have that J DSM ≥ J SM .
Plugging it in eq.
(32), we can get eq.
Proof of Corollary 3.
We have p 0 t ( x | x (0)) = N ( √ ¯ α t x (0) , (1 − ¯ α t ) I ) where ¯ α t =
(cid:81)
tr =1 (1 − β t ) ,which can be inferred from p ( x ( t ) | x ( t − 1))
= N ( √ 1 − β t x ( t − 1) , β t I ) . Thus we can show that E [ ∇ x log p 0 t ( x | x (0))
⊤ | x (0)] is constant with respect to x (0) :
recallp 0 t ( x | x (0))
=
N ( √ ¯ α t x (0) , (1 − ¯ α t ) I ) , we have ∇ x log p 0 t ( x | x (0))
=
− ((1 − ¯ α t ) I ) − 1 ( x ( t ) −
√ ¯ α t x (0)) , which is a linear function of x ( t ) − √ ¯ α t x (0) .
Using the Gaussian density function, wehave:
[Equation]
As a result, Var[ E [( ∇ x log p 0 t ( x | x (0)))
⊤ | x (0)]] = 0 , which satisfies the condition of eq.
(15) in
Theorem 2.
Note that here the assumption of f and g is only a sufficient condition for
Var[ E [( ∇ x log p 0 t ( x | x (0)))
⊤ | x (0)]] = 0 .
In fact, any conditional distribution p 0 t that satisfies Var[ E [( ∇ x log p 0 t ( x | x (0)))
⊤ | x (0)]]
= 0 can lead to the same conclusion.
E Numerical results on J DSM upper-bounding J SM in DDPM
To verify J SM ≤ J DSM in (16) for DDPM, we adopt the same datasets as in Fig 1, and the sametraining and evaluation settings in Section 4.1.
Moreover, to estimate J SM numerically, we estimate p t ( x ) by performing Gaussian kernel densityestimation with bandwidth = 0 .
05 on sampled data.
∇ x p t ( x ) is estimated by central differenceapproximation with interval = 0 .
The resulting plots of J DSM and J SM are shown in Fig 6, whichshows that J DSM is an upper bound of J SM in DDPM during training, where p 0 is the dataset andq 0 is the generated data distribution at the convergence of training.
[Figure]
modification of J DSM loss.
F
{+Full plot of L s ( t ) when T = 
+}
See Fig.
G
Log-log plots with weight decay
See Fig. 8.
H Further analysis of the upper bound
H.1 Proof of Corollary 2
We further investigate the upper bound that appeared in the previous sections.
To illustrate this, let usconsider the following case:
[Equation]
(46) for a positive function β : [0 , T ] → (0 , + ∞ ) .
[Figure]
Proof of Corollary 2.
In the case of (46), the stationary solution to (2) is given as the probabilitydistribution of N (0 , σ 2 I ) :
[Equation]
Based on the contraction property [8], we quantify the Wasserstein distance between p T and ϕ .
[Equation]
Using the above, the definition of I ( t ) , and q T
= ϕ , we conclude (13).
It worth noting that as a consequence of (48), W 2 ( p T , ϕ ) is small for an appropriate choice of T andβ ( t ) .
H.2 Exponential convergence of h t
For simplicity of notations, we define the norm in L 2 ( ϕ ) as follows,
[Equation]
where ϕ is given in (47).
In addition, assume thatβ ( t ) > c > 0 for all t ≥ 0 .
[Equation]
[Equation]
Lemma 4.
Under the same setting as in Corollary 2, we have
[Equation]
where h t =
p t /ϕ for some constant λ > 0 .
In particular, h t exponentially converges to 1 in L 2 ( ϕ )for β satisfying (50) .
Remark 8.
The convergence of h t to 1 can be shown under more general assumption: β >satisfying
[Equation]
For a constant function β , Lemma 4 is proven in [23].
For the sake of completeness, we provide theproof, which is a small modification of [23, Section 2].
Proof of Lemma 4.
In our case, the equation (2) of p t is given by
[Equation]
By direct computations, we obtain the equation of h ( x, t ) =
h t ( x ) as follows:
[Equation]
We estimate ∥ h − 1 ∥ 2 L 2 ( ϕ ) by differentiating it with respect to time.
Using the integration by partsand Poincaré inequality, we obtain that
[Equation]
This yields (51).
Lastly, for β satisfying (50), ∥ h t − 1 ∥ 2 L 2 ( ϕ ) ≤ exp( − σ 2 λct )
∥
h 0 − 1 ∥ 2 L 2 ( ϕ ) .
Thus,we conclude the exponential convergence of h to 1 .
Further analysis is plausible based on the techniques in the study of partial differential equations.
Remark 9.
As p t is given as the convolution between p 0 and the Gaussian distribution, it is smoothfor t > 0 .
Therefore, h t is also smooth, and the higher order derivatives of h t are all bounded.
As aconsequence, the above result combining with Gagliardo–Nirenberg interpolation inequality yieldthat the gradient of h , Dh , and the Hessian of h , D 2 h , also converge to zero in L 2 ( ϕ ) .
Remark 10.
Proving the uniform convergence of h , Dh or D 2 h requires an additional technicalassumption that the support of h is bounded.
Under the assumption, another interpolation inequality,
Agmon’s inequality yields the desired uniform convergence result.
H.3 Estimation of L s
In this subsection, we investigate the estimation of L s .
If J SM is sufficiently small, then s θ is closeto ∇ log p t .
We first investigate the one-sided Lipschitz constant of ∇ log p t .
Lemma 5.
Under the same setting as in Corollary 2, ∇ log p t satisfies the one-sided Lipschitzcondition with a constant ( − σ − 2 + ∥ D 2 (log h ) ∥ ∞ ) i.e.,
[Equation]
55) where ∥ · ∥ ∞ denotes the supremum of the matrix norm,
[Equation]
Proof.
From the definition of h t , we have
[Equation]
(57) for some constant c .
As a consequence,
[Equation]
To prove (55), it suffices to estimate ( ∇ log h t ( x ) − ∇ log h t ( y )) · ( x − y ) .
From the fundamentaltheorem of calculus, we have
[Equation]
Using | z ⊤ D 2 (log h t ( w ))
z | ≤ ∥ D 2 (log h ) ∥ ∞
∥ z ∥ 2 for all w, z ∈ R d , we have
[Equation]
60) and conclude (55).
Remark 11.
Based on the similar relation as in (59) , the difference between L s and − σ − 2 +∥ D 2 (log h ) ∥ ∞ can be estimated.
More precisely, we have L s ( t ) =
( − σ − 2 + ∥ D 2 (log h ) ∥ ∞ ) + ϵ ( t ) .
Here, ϵ ( t ) depends on the difference between s θ and ∇ log p t .
Therefore, it is expected that the upperbound of (cid:82) T 0 ϵ ( t ) dt is given by J SM under suitable regularity assumptions. I Verifying the upper bound in high-dimensional data Here we present numerical results on some moderately high dimensional data, where estimating both W 2 ( p 0 , q 0 ) and L s would become more computationally expensive
Here we present the log-log plots for Gaussian cluster N ( 0 , 0 .
1 I ) with d = 10 , 20 and 100 respectively in Fig 9, where the training configuration is the same as Section 4. As we can observe, ourupper bound is still a valid one of the empirical Wasserstein distance, though the gap between theempirical one and the theoretical one becomes larger when the dimensionality is higher.
[Figure]
We also trained the diffusion model on MNIST dataset for 100 epochs with T = 100 , while wemodify the neural network architecture to UNet, omit the loss maximization training, and followingother experimental configurations as in [17]. As in Section 4, we still sample from p T =
q T whencalculating W 2 ( p 0 , q 0 ) , in order to satisfy W 2 ( p T , q T )
{+≈ 
+}
for the log-log plot. The generated imagesand the corresponding log-log plot obtained during training are shown in Fig 10, where we also showthe validity of our upper bound on image data.
[Figure]
J Computing resources
We use Google Colab ( https://research.google.com/colaboratory ) for our training.