Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance
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-based generative models is equivalent to minimizing the Kullback–Leibler divergence of the generated distribution from the data distribution.
In this work, we show that score-based models also minimize the Wasserstein distance between them under suitable assumptions on the model.
Specifically, we prove that the Wasserstein distance 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 novel application of the theory of optimal transport, which can be of independent interest to the society.
Our numerical experiments support our findings.
By analyzing our upper 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 learning the reverse process has recently demonstrated remarkable performances in generative modeling.
Score-based generative models [33, 34, 35] and diffusion probabilistic models [30, 17] have been used 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 incrementally perturbed by adding small noise in the forward process.
This forward process can be understood as the standard diffusion process [35].
The reverse process requires the so-called score function [2], the gradient of the log probability density function.
We minimize a weighted combination of score matching losses to approximate the score function.
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 model distribution (the distribution of the generated samples) and the data distribution.
Therefore, it is 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 show that 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 is similar (in the sense of the KL divergence) to the data distribution if the score matching loss is small.
A natural question follows – do score-based models minimize any other divergence or distance than 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 choice of the loss function.
The original GAN loss function turns out to capture the Jensen–Shannon
[Figure]
divergence [16], while the W–GAN loss function captures the Wasserstein distance [3].
We refer the readers to [10] for more details
[Figure]
On the other hand
[Figure]
a unique aspect of score-based models is that the score function is defined in the same 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 have nothing to do with the underlying data space.
To gain more insights, we train score-based models on a few synthetic datasets.
Shown in Fig.are the training results.
We plot the (scaled) training objectives in blue.
Shown in green are the KL divergences, and one can note that the KL divergence decreases as the training objective decreases, as predicted in [32].
What we also observe is that as the training objective decreases, the Wasserstein distance also decreases.
This is quite surprising as there is no direct relationship between the KL divergence and the Wasserstein distance, in general.
Therefore, this phenomenon is not immediately implied by the fact that score-based models minimize the KL divergence.
Inspired by this observation, in this work, we prove the following property of score-based modeling:
Score-based modeling minimizes the Wasserstein distance.
More specifically, 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 between the given data distribution and the distribution of the generated samples.
Our work theoretically guarantees that the score-based generative models approximate data distributions in the space of probability measures equipped with the Wasserstein distance.
Our main contributions and findings are summarized as follows:
• (Theorem 1 & Corollary 1) Using the theory of optimal transport, we prove that the Wasserstein 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.
(20), respectively)
under suitable assumptions, which gives us a more practical upper bound since J DSM is 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)
We corroborate our
analysis by conducting
numerical experiments
We
also explore
various
techniques to tighten up the
upper bound
in
practice.
Notations: x = ( x 1 , x 2 , · · · , x d ) ⊤ ∈ R d and t ∈ [0 , T ] denote the spatial variable and the time variable, 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-derivative of ρ are denoted by ∇ ρ = (cid:16) ∂ρ∂x 1 , ∂ρ∂x 2 , · · · , ∂ρ∂x d (cid:17) ⊤ , and ∂ t ρ , respectively.
The Laplacian of ρ is given 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 µ and covariance Σ by N ( µ, Σ) , and L 2 norm by ∥ · ∥ .
We denote the data distribution and the distribution of 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 model through Partial differential equations (PDEs).
In particular, we investigate the time evolution of marginal 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 for t ∈ [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 ) , and g : [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, namely the Fokker–Planck equation or the forward Kolmogorov equation as in [5, Theorem 2.2]:
[Equation] [Equation]
where p 0 represents the given data distribution.
Under suitable assumptions on f and g , a solution to 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 variable x ( t ) from the above process (3) satisfies the final value problem:
[Equation]
Using this simple identity, ∆ p = ∇ · ( ∇ p ) = ∇ · (cid:16)
p ∇ pp (cid:17) =
∇ · ( p ∇ (log p )) , it can be seen that the differential 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 ] → R d .
For instance, in [32], the authors consider the weighted mean squared error (MSE):
[Equation]
(5) 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 , a probability 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 , + ∞ ) to illustrate the reverse process.
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 of probability measures equipped with the 2 -Wasserstein distance [1, 18, 27].
For two probability measures µ 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 d satisfies π ( 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], which plays an important role in our analysis.
More discussions can be found in Section 3.4 and Section 5.
Score matching and diffusion probabilistic models Score matching models [33, 35] use the score function to reverse the noise corruption in a continuous time range by solving stochastic differential equations (SDEs).
Diffusion probabilistic models [30, 17] adopt a similar idea of modeling the tractable denoising 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 unifies the 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 stable training process, which encourages it to be widely applied in various tasks, such as text-to-image and image 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 the marginal of the generated samples q 0 obtained by the score-based model.
The proofs of the main results are deferred to Appendices A, B, and C. The key ideas are illustrated in Section 5.
3.1 Assumptions
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]: there exists a constant L s ( t ) ∈ R , depending on t ∈ [0 , T ] , satisfying for all x, y ∈ R d
[Equation]
The usual Lipschitz constant of s θ satisfies the above inequality (9).
Unlike the usual Lipschitz constant, the one-sided Lipschitz constant L s ( t ) is not necessarily nonnegative (see Appendix D).
(A3) For the
initial distribution p 0 of 2) and the final distribution q T of (6), the following quantities are finite:
[Equation]
where Λ( x ) := log max( ∥ x ∥ , 1) .
(A4)
There exist constants C 1 and C 2 such that
[Equation]
for
all x ∈ R d and t ∈
[Equation]
T ]
[Equation]
(A5)
There exists a positive constant C 3 such that
[Equation]
(A6) (cid:82) T [-
-]0 E p t [ f 2 ] dt and (cid:82) T [-
-]0 E q t
[( f − g 2 s θ ) 2 ] dt are finite.
The above assumptions guarantee the integrability of the score function ∇ log p .
Proposition 1.
[6, Theorem 4.2.1]
A solution p t to (2) satisfies ∇ p t ∈ L 1 ( R d ) and
[Equation]
We give the 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 rapidly in 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 t grow at most linearly as ∥ x ∥ → ∞ .
3.2 The main results
We now present a general result.
Theorem 1.
Let p 0 and q 0 be the data distribution and a pdf obtained by the score-based model as in Section 2.
Then the Wasserstein distance between p 0 and q 0 is bounded as follows:
[Equation]
Here, a function I : [0 , T ] → (0 , + ∞ ) defined by
[Equation]
is an integrating factor of the differential inequality (28) given later in Proposition 2.
The main idea of 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 (14) 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 the following 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 general convergence result remains open (see Section 6 and Appendix B).
Remark 2.
If g 4 I 2 λ − 1 is integrable in [0 , T ] , the similar inequality as in (34) holds for other choices of λ .
See Corollary 4 in Appendix A.
Remark
We further investigate the offset I ( T ) W 2 ( p T , q T ) in Corollary 1.
Consider the following case as in
Denoising Diffusion Probabilistic Models (DDPM)
[17]: for β : [0 , T ] → (0 , + ∞ ) , and σ > 0 , f [Equation] x, t ) [Equation]
β
[Equation]
t )
[Equation]
x and g
[Equation]
t )
[Equation] [Equation] [Equation] In the above case, the stationary solution to (2) is given as the probability distribution of N (0 , σ 2 I ) :
[Equation]
Corollary 2.
Under the same setting as in Corollary 1, let q T
= ϕ .
For f, g given in (17) , we have
[Equation]
The proof of Corollary 2 and the analysis of L s can be found in Appendix B.
Remark 4.
If s θ is close enough to ∇ log p t ( x ) , then we have
[Equation]
Here, h t :=
p t /ϕ exponentially converges to 1 as t → ∞ (see Lemma 3).
As a consequence, for sufficiently large t > 0 , σ 2 L s ( t ) converges to − 1 , as observed in Fig. 4a.
This yields that the second term in (19) decays as T increases (see Fig. 4b).
3.3 The relation between J SM and J DSM
The above results give the guarantee that J SM can upper-bound Wasserstein distance.
However, J SM is 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 equivalent to J SM up to a constant, as shown in [36].
J DSM is also widely used as a loss function to train score-based models:
[Equation]
where p 0 t indicates the conditional probability of x at step t given x (0) .
Now we present our new result, which shows that J DSM is also an upper bound of J SM under a 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 (21).
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 ( t ) 2 .
Then, (22) holds.
The proofs of Theorem 2 and Corollary 3 can be found in Appendix C. We also provide numerical results which show J SM ≤ J DSM for DDPM in Appendix F.
Remark 5.
Using Theorem 2 and Corollary 3, we can present an upper bound of the Wasserstein distance in terms of J DSM , where both W 2 ( p 0 , q 0 ) and J DSM are tractable.
This enables us to empirically verify this upper bound via DDPM in Section 4.
3.4 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.
In particular, we are interested in the upper bound on W 2 ( q 0 , ˜ q 0 ) with respect to W 2 ( p 0 , ˜ p 0 ) .
It is well known from the Wasserstein contraction property [8] that the Wasserstein distance between two probability distributions in the forward process can be estimated as follows: for all t ∈ [0 , T ]
[Equation]
In the reverse process, the score function in the drift coefficients may vary, and thus the above inequality does not directly give us the upper bound on W 2 ( q 0 , ˜ q 0 ) .
Instead, our analysis in the previous 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 statistical measures, including the Kullback–Leibler (KL) divergence.
Combining our main results with
[Equation]
we conclude the following theorem.
Theorem 3.
Under the same setting as in Corollary 1, let ˜ p 0 be the corrupted data distribution and ˜ q 0 be a probability density function obtained 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 model distribution ˜ 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.
Code is available at https://github.com/UW-Madison-Lee-Lab/score-wasserstein .
4.1 Experiment settings
Datasets Here we adopt three 2D datasets for simulation: One cluster Gaussian N ( 0 , 0 .
1 I ) , two moons in [28], and four clusters Gaussian mixture N (( ± 0 .
5 , ± 0 .
5) ⊤ , 0 .
01 I ) with equal weights for each 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 ) = − 12 β t x , and g ( t ) = √ β t .
For β t , we adopt a sigmoid schedule from β 1 = 10 − 5 to β T = 10 − 2 .
Training and evaluation We use a 4-layer neural network as the score matching model, with ReLU nonlinearity and skip-connection at the final output.
Each layer is composed of a linear layer withhidden 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.
(22) for verification, we first perform forward pass on sampled x (0) to get x ( T ) , and perform data 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 ( t ) 2 and batch size = 128.
To get both lower and higher losses to verify our bound, we first maximize the loss for 10 epochs and then minimize the loss till convergence.
4.2 Results 4.2.1 The upper bound
To see the upper bound in eq.
[Figure]
we take the log on each side of the equation with W 2 ( p T , q T ) = 0 :
[Equation]
[Equation]
We present a log-log plot to verify the slope and the intercept in the upper bound respectively while fixing p T =
q T : see Fig 2.
The log-log plots are obtained by evaluating Wasserstein distance and the loss function during training.
The red dots are the sampled data from each dataset, the green dots are generated data after loss maximization, and the purple dots are the generated data sampled at convergence.
From the log-log plots, we can observe that the theoretical line (the black line) is a valid upper bound of the empirical one (the blue dots).
[Figure]
4.2.2 Effect of weight regularization
Here we explore weight constraints that can affect the tightness of the theoretical upper bound while fixing p T = q T .
As shown in Fig 2, the gap between the theoretical bound and the empirical one mainly results from the intercept in eq.
Recall our intercept is 12 log (cid:16)(cid:82) T 0 g ( t )
I ( t ) 2 dt (cid:17) for I ( t ) given in (15).
Since g ( t ) and L f are fixed, only the one-sided Lipchitz constant L s ( r ) of the model s θ is subject to the 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)) .
If the 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). If the one-sided Lipschitz constant for ∇ x log p 0 t ( x | x (0)) is large, to minimize the loss, the intercept would also be large (see example in Fig 2c of Gaussian clusters with lower variance).
However, it can be possible to make the upper bound tighter while sacrificing the training loss by constraining L s .
To constrain L s , we adopt spectral normalization [4], weight clipping with threshold = 0.1[3].
As we can observe in Fig.
3, loss at convergence will increase when we impose very strong regularization on the weights while the intercept decreases.
Importantly, the intercept approaches the empirical one without not exceeding it.
However, if the regularization is too strong, the training loss at convergence would also become larger, resulting in worse quality of the generated data.
See more examples of weight decay in Appendix G, where we also present the log-log training plots with different weight decay coefficients.
[Figure]
4.3 Effect of the number of total time steps: the decay of the offset
In our Remark 4, we argue that the offset in our upper bound would decay to 0 when choosing the appropriate forward SDE and sufficiently large T .
Here we provide empirical evidence for our claim without fixing p T = q T .
We use the four clusters Gaussian dataset and fix β to be a sigmoid schedule in the range of [10 − 5 , 10 − 2 ] for all choices of T .
In Fig.
4, we plot L s ( t ) for T = 100 to show that
L s ( t ) would converge to − 1 when t → T .
Moreover, since W 2 ( p T , q T ) decays exponentially (see eq. 39), we also show that as T increases, our offset I ( T ) W 2 ( p T , q T ) would also decay in a rate that is approximately exponential when T is sufficiently large, as the model converges to the actual data distribution ( W 2 ( p 0 , q 0 ) ≈ 0) .
[Figure]
(a
T = 100 .
We can verify that when T is sufficiently large, L s ( T ) → − 1 as suggested in
Remark 4.
L s ( t ) when t ≤ 50 is cut due to visibility of the threshold 0 and -1.
The full figure is in Appendix E.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
The main idea of the proofs
We investigate the time evolution of probability measures arising in the score-based models.
Using the optimal transport theory, we differentiate the Wasserstein distances, W 2 ( p t , q t ) with respect to time.
The main idea is to express these equations (2) and (6) in the form of the continuity equation:
[Equation]
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 (26) with v [ p ] and v [ q ] , respectively.
Using these vector fields, we compute the time derivative of the Wasserstein distance between two evolving probability measures.
For velocity fields v [ ρ ( i ) ] depending on ρ ( i ) , let ρ ( i ) for i = 1 , 2 be two solutions of the continuity equations
[Equation]
Then, the relation between the Wasserstein distance and velocity fields arises from the theory of optimal transport:
[Equation]
where π t is an optimal transport plan between ρ (1) t and ρ (2) t .
See [1, Theorem 8.4.7] and [29, Corollary 5.25].
This equality plays an important role when we estimate the Wasserstein distance.
The following proposition is a consequence of (27) and Lemma 1 in Appendix A.2.
Proposition 2.
Let p t and q t be solutions to (2) and (6) , respectively, in [0 , T ] .
Then,
[Equation]
(28) where b ( t ) :
= E p t (cid:2) ∥∇ log p t ( x ) − s θ ( x, t ) ∥ 2 (cid:3) .
Discussions
Convergence of the Wasserstein distance If p T =
q T , then our result yields that the Wasserstein distance 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 and f 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, additional assumptions 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 upper bound.
We expect that appropriate assumptions on s θ and the initial distribution pmay resolve this issue, 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.
Also, estimating the one-sided Lipschitz constant does not enjoy the semi-definiteness of estimating the two-sided one in [14].
Note that the two-sided Lipschitz constant is an upper bound of the one-sided one, which can be more efficiently estimated as shown in [14].
However, it would result in a looser theoretical upper bound if we adopt the two-sided Lipschitz constant.
Conclusion
In this paper, we mainly discuss the relationship between the loss function of score-based models and the Wasserstein distance.
With the help of the theory in optimal transport, we present a novel upper bound on Wasserstein distance in terms of the loss function used in practice.
We also provide numerical 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 influence on the training process.
Acknowledgments and Disclosure of Funding
Support for this research was provided by the University of Wisconsin-Madison Office of the Vice
Chancellor for Research and Graduate Education with funding from the Wisconsin Alumni Research
Foundation, and NSF Award DMS-2023239.
A Proofs of Theorem 1 and CorollaryA.1 Proofs
Proof of Theorem 1.
Recall the differential inequality (28) below from Proposition 2:
[Equation]
It can be solved by introducing the integrating factor,
[Equation]
where
[Equation]
t
[Equation]
As ddt I ( t ) =
( L f ( t ) + L s ( t ) g ( t ) 2 ) I ( t )
[Equation]
the above inequality (28) can be written as
[Equation]
Integrating both sides from 0 to T
[Equation]
we obtain that
[Equation]
As I (0) =
[Equation]
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 (29).
The Cauchy-Schwarz inequality yields
[Equation]
Since λ = g 2 and J SM given in (5) satisfies
[Equation]
we conclude (34):
[Equation]
Remark 7.
It is expected that the parallel results hold for weak solutions to (2) and (4) by using suitable approximations.
This approach has been studied for the Wasserstein contraction property in [8, Section 6.2].
From (32) and (33) 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]
A.2 Technical lemmas
Lemma 1.
Let π t be an optimal transport plan between p t and q t .
Then, we have
[Equation] [Equation]
where b ( t ) :
= E p t (cid:2) ∥∇ log p
[Equation]
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 optimal transport.
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 there exists 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 t for all t ∈ [0 , T ] .
As in the proof of [7, Theorem 1], we have
[Equation]
where ϕ ∗ t is a convex conjugate of ϕ t .
The convexity of ϕ t yields that ∆ ϕ t + ∆ ϕ ∗ t ( ∇ ϕ t ) − 2 d
and we conclude.
B Further analysis of the upper bound
Proof of Corollary 2.
Recall from Corollary 1 that
[Equation]
Based on the contraction property [8], we quantify the Wasserstein distance between p T and ϕ .
[Equation]
[Equation]
Using the above, the definition of I ( t ) , and q T
= ϕ , we conclude (19).
It worth noting that as a consequence of (39), W 2 ( p T , ϕ ) is small for an appropriate choice of T and β ( t ) .
B.1 Exponential convergence of h t
For simplicity of notations, we define the norm in L 2 ( ϕ ) as follows,
[Equation]
where ϕ is given in (18).
In addition, assume that
[Equation]
Lemma 3.
Under the same setting as in Corollary 2, we have
[Equation]
42) where h t =
p t /ϕ for some constant λ > 0 .
In particular, h t exponentially converges to 1 in L 2 ( ϕ ) for β satisfying (41) .
For a constant function β , Lemma 3 is proven in [23].
For the sake of completeness, we provide the proof, which is a small modification of [23, Section 2].
Proof of Lemma 3.
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 parts and Poincaré inequality, we obtain that
[Equation]
This yields (42).
Lastly, for β satisfying (41), ∥ 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 .
Remark 8.
The convergence of h t to 1 can be shown under more general assumption: β >satisfying
[Equation]
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 smooth for t > 0 .
Therefore, h t is also smooth, and the higher-order derivatives of h t are all bounded.
As a consequence, the above result combined with Gagliardo–Nirenberg interpolation inequality yield that 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 technical assumption that the support of h is bounded.
Under the assumption, another interpolation inequality,
Agmon’s inequality, yields the desired uniform convergence result.
B.2 Estimation of L s
In this subsection, we investigate the estimation of L s .
If J SM is sufficiently small, then s θ is close to ∇ log p t .
We first investigate the one-sided Lipschitz constant of ∇ log p t .
Lemma 4.
Under the same setting as in Corollary 2, ∇ log p t satisfies the one-sided Lipschitz condition with a constant ( − σ − 2 + ∥ D 2 (log h ) ∥ ∞ ) i.e.,
[Equation] [Equation]
where ∥ · ∥ ∞ denotes the supremum of the matrix norm,
[Equation]
Proof.
From the definition of h t , we have
[Equation]
(48) for some constant c .
As a consequence,
[Equation]
To prove (46), it suffices to estimate ( ∇ log h t ( x ) − ∇ log h t ( y )) · ( x − y ) .
From the fundamental theorem 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]
51) and conclude (46).
Remark 11.
Based on the similar relation as in (50) , 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 upper bound of (cid:82) T 0 ϵ ( t ) dt is given by J SM under suitable regularity assumptions.
C Proofs of Theorem 2 and Corollary
For a given t , let
[Equation]
52) and
[Equation]
Lemma 5.
(Appendix in [36])
[Equation]
( x ) [ ∇ x log p t ( x )]
= E p 0 ( x (0))
p 0 t ( x | x (0))
[ ∇ x log p 0 t ( x | x (0))] .
Proof of Theorem 2.
From Lemma 5 in [36], we have
J SM ( θ, t ) =
J DSM (
[Equation] [Equation]
Note that
[Equation]
where the inequality comes from the law of total variance and our condition:
[Equation]
Then
[Equation]
we have
[Equation]
58) 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.
(34), 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) : recall p 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, we have:
[Equation]
As a result, Var[ E [( ∇ x log p 0 t ( x | x (0)))
⊤ | x (0)]] = 0 , which satisfies the condition of eq.
(21) 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.
[Figure]
D 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.2, 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 Lipschitz constant of F while its Lipschitz constant is a > 0 .
Figure 5 visualizes this.
Note that two-sided Lipschitzness is a subset of one-sided Lipscthizness.
See Figure 6 as an example.
[Figure]
E Full plot of L s ( t ) when T =
See Fig.
7b.
100 t
L s ( t ) -
[Figure]
100 t
L s ( t ) -
[Figure]
F Numerical results on J DSM upper-bounding J SM in DDPM
To verify J SM ≤ J DSM in (22) for DDPM, we adopt the same datasets as in Fig 1, and the same training and evaluation settings in Section 4.1.
Moreover, to estimate J SM numerically, we estimate p t ( x ) by performing Gaussian kernel density estimation with bandwidth = 0 .
05 on sampled data.
∇ x p t ( x ) is estimated by central difference approximation with interval = 0 .
The resulting plots of J DSM and J SM are shown in Fig 8, which shows that J DSM is an upper bound of J SM in DDPM during training, where p 0 is the dataset, and q 0 is the generated data distribution at the convergence of training.
[Figure]
G Log-log plots with weight decay
See Fig.
[Figure]