Title: Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion models

sampling drift

Fokker-Planck

invariances

Stochastic Differential Equations

Martingales
Abstract: Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. However, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a \emph{Consistency} property (CP) which states that predictions of the model on its own
generated data are consistent across time. Theoretically, we show that the differential equation that describes CP together with the one that describes a conservative vector field, have a unique solution given some initial condition. Consequently, if the score is learned well on non-drifted points via DSM (enforcing the true initial condition) then enforcing CP on drifted points propagates true score values. Empirically, we show that enforcing CP improves the generation quality for conditional and unconditional generation on CIFAR-10, and in AFHQ and FFHQ. 
We open-source our code and models: https://github.com/giannisdaras/cdm.

Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent

Abstract

Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. However, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a Consistency property (CP) which states that predictions of the model on its own generated data are consistent across time. Theoretically, we show that the differential equation that describes CP together with the one that describes a conservative vector field, have a unique solution given some initial condition. Consequently, if the score is learned well on nondrifted points via DSM (enforcing the true initial condition) then enforcing CP on drifted points propagates true score values. Empirically, we show that enforcing CP improves the generation quality for conditional and unconditional generation on CIFAR-10, and in AFHQ and FFHQ. We open-source our code and models: https://github.com/giannisdaras/cdm.

1 Introduction

The diffusion-based [47, 49, 18] approach to generative models has been successful across various modalities, including images [41, 44, 14, 39, 29, 51, 43, 16, 10, 11], videos [19, 20, 21], audio [32], 3D structures [40], proteins [1, 53, 45, 9], and medical applications [23, 3].

Diffusion models generate data by first drawing a sample from a noisy distribution and slowly denoising this sample to ultimately obtain a sample from the target distribution. This is achieved by sampling, in reverse from time t = 1 down to t = 0 , a stochastic process { x t } t ∈ [0 , 1] wherein x 0 is distributed according to the target distribution p 0 and, for all t ,

where ⊕ denotes the convolution operator. That is, p t is the distribution resulting from corrupting a sample from p 0 with noise sampled from N (0 , σ 2 t I d ) , where σ t is given by an increasing function of t , such that σ 0 = 0 and σ 1 is sufficiently large so that p 1 is nearly indistinguishable from pure noise. We note that diffusion models have been generalized to other types of corruptions by the recent works

of Daras et al. [12], Bansal et al. [4], Hoogeboom and Salimans [22], Deasy et al. [13], Nachmani et al. [38].

In order to sample from a diffusion model, i.e. sample the afore-described process in reverse time, it suffices to know the score function s ( x, t ) = ∇ x log p ( x, t ) , where p ( x, t ) is the density of x t ∼ p t . Indeed, given a sample x t ∼ p t , one can use the score function at x t , i.e. s ( x t , t ) , to generate a sample from p t -dt by taking an infinitesimal step of a stochastic or an ordinary differential equation [51, 48], or by using Langevin dynamics [17, 50]. 2 Hence, in order to train a diffusion model to sample from a target distribution of interest p ∗ 0 it suffices to learn the score function s ∗ ( x, t ) using samples from the corrupted distributions p ∗ t resulting from p ∗ 0 and a particular noise schedule σ t . Notice that those samples can be easily drawn given samples from p ∗ 0 .

The Sampling Drift Challenge: Unfortunately, the true score function s ∗ ( x, t ) is not perfectly learned during training. Thus, at generation time, the samples x t drawn using the learned score function, s θ ( x, t ) , in the ways discussed above, drift astray in distribution from the true corrupted distributions p ∗ t . This drift becomes larger for smaller t due to compounding of errors and is accentuated by the fact that the further away a sample x t is from the likely support of the true p ∗ t the larger the error ∥ s θ ( x t , t ) -s ∗ ( x t , t ) ∥ between the learned and the true score function at x t , which feeds into an even larger drift between the distribution of x t ′ from p ∗ t ′ for t ′ < t ; see e.g. [46, 18, 39, 6]. These challenges motivate the question:

Question. How can one train diffusion models to improve the error ∥ s θ ( x, t ) -s ∗ ( x, t ) ∥ between the learned and true score function on inputs ( x, t ) where x is unlikely under the target noisy distribution p ∗ t ?

A direct approach to this challenge is to train our model to minimize the afore-described error on pairs ( x, t ) where x is sampled from distributions other than p ∗ t . However, there is no straightforward way to do so, because we do not have direct access to the values of the true score function s ∗ ( x, t ) .

This motivates us to propose a training method to mitigate sampling drift by enforcing that the learned score function satisfies an invariant, that we call the Consistency Property (CP). This property can be optimized without using any samples from the target distribution p ∗ 0 . We will show that theoretically, enforcing CP on drifted points, in conjunction with minimizing the standard score matching objective on non drifted points, suffices to learn the correct score everywhere - at least in the theoretical limit where the error approaches zero and when one also enforces conservative vector field. We also provide experiments illustrating that regularizing the standard score matching objective using our CP improves sample quality. Further, we provide an ablation study that further provides evidence to this phenomenon of score propagation.

Our Approach: The true score function s ∗ ( x, t ) is closely related to another function, called the optimal denoiser , which predicts a clean sample x 0 ∼ p ∗ 0 from a noisy observation x t = x 0 + σ t η where the noise is η ∼ N (0 , I d ) . The optimal denoiser (under the ℓ 2 loss) is the conditional expectation:

and the true score function can be obtained from the optimal denoiser as follows: s ∗ ( x, t ) = ( h ∗ ( x, t ) -x ) /σ 2 t . This result is known as Tweedie's Formula [15]. Indeed, the standard training technique, via score-matching , explicitly trains for the score through the denoiser h ∗ [54, 15, 36, 30, 35].

We are now ready to state our Consistency Property (CP). We will say that a (denoising) function h θ ( x, t ) satisfies CP iff

where the expectation is with respect to a sample from the learned reverse process , defined in terms of the implied score function s θ ( x, t ) = ( h θ ( x, t ) -x ) /σ 2 t , when this is initialized at x t = x and run backwards in time to sample x 0 . See Eq. (3) for the precise stochastic differential equation and its justification. In particular, h θ satisfies CP if the prediction h θ ( x, t ) of the conditional expectation of the clean image x 0 given x t = x equals the expected value of an image that is generated by the

learned reversed process, starting from x t = x . Equivalently, one can formulate this property as requiring x t to follow a reverse Martingale (see Lemma 3.1).

While there are several other properties that the score function of a diffusion process must satisfy, e.g. the Fokker-Planck equation [33], our first theoretical result is that the h θ ( x, t ) satisfying the Consistency Property suffices (in conjunction with the conservativeness of its score function s θ ( x, t ) = ( h θ ( x, t ) -x ) /σ 2 t ) to guarantee that s θ must be the score function of a diffusion process (and must thus satisfy any other property that a diffusion process must satisfy). If additionally s θ ( x, t ) equals the score function s ∗ ( x, t ) of a target diffusion process at a single time t = t 0 and an open subset of x ∈ R d , then it equals s ∗ everywhere. We comment that the formal theorem is proved for an idealistic setting when the error is (or approaches) zero. Still, it is likely to believe that even in the finite-error regime, training with DSM in-sample and enforcing CP off-sample is expected to improve the score function values off-sample. The statement is summarized as follows below:

Theorem 1.1 (informal) . If some denoiser h θ ( x, t ) satisfies CP and its corresponding score function s θ ( x, t ) = ( h θ ( x, t ) -x ) /σ 2 t is a conservative field, then s θ ( x, t ) is the score function of a diffusion process, i.e. the generation process using score function s θ , is the inverse of a diffusion process. If additionally s θ ( x, t ) = s ∗ ( x, t ) for a single t = t 0 and for all x in an open subset of R d , where s ∗ is the score function of a target diffusion process, then s θ ( x, t ) = s ∗ ( x, t ) everywhere.

Simply put, the above statement states that: i) satisfying CP and being a conservative vector field is enough to guarantee that the sampling process is the inverse of some diffusion process and ii) to learn the score function everywhere it suffices to learn it for a single t 0 and an open subset of x 's.

We propose a loss function to train for the Consistency Property and we show experimentally that regularizing the standard score matching objective using our property leads to better models.

Summary of Contributions:

We identify an invariant property, consistency of the denoiser h θ , that any perfectly trained model should satisfy.

We prove that if the denoiser h θ ( x, t ) satisfies CP and its implied score function s θ ( x, t ) = ( h θ ( x, t ) -x ) /σ 2 t is a conservative field, then s θ ( x, t ) is the score function of some diffusion process, even if there are learning errors with respect to the score of the target process, which generates the training data.

We prove that optimizing for the score in a subset of the domain and enforcing these two properties, guarantees that the score is learned correctly in all the domain, in the limit where the error approaches zero.

We propose a novel training objective that enforces the Consistency Property. Our new objective optimizes the network to have consistent predictions on data points from the learned distribution.

We show experimentally that, paired with the original Denoising Score Matching (DSM) loss, our objective improves generation quality on conditional and unconditional generation on CIFAR-10, and in AFHQ and FFHQ.

We conduct an ablation study which showcases that even if we do not optimize for DSM for some values of t , satisfying CP enforces good score approximation there.

2 Background

Diffusion processes, score functions and denoising. Diffusion models are trained by solving a supervised regression problem [49, 18]. The function that one aims to learn, called the score function (defined below), is equivalent (up to a linear transformation) to a denoising function [15, 54], whose goal is to denoise an image that was injected with noise. In particular, for some target distribution p 0 , one's goal is to learn the following function h : R d × [0 , 1] → R d :

In other words, the goal is to predict the expected 'clean' image x 0 given a corrupted version of it, assuming that the image was sampled from p 0 and its corruption was done by adding noise to it

from N (0 , σ 2 t I d ) , where σ 2 t is a non-negative and increasing function of t . Given such a function h , we can generate samples from p 0 by solving a Stochastic Differential Equation (SDE) that depends on h [51]. Specifically, one starts by sampling x 1 from some fixed distribution and then runs the following SDE backwards in time:

where B t is a reverse-time Brownian motion 3 and g ( t ) 2 = dσ 2 t dt . To explain how Eq. (3) was derived, consider the forward SDE that starts with a clean image x 0 and slowly injects noise:

where B t is a standard Brownian motion. We notice here that the x t under Eq. (4) is N ( x 0 , σ 2 t I d ) , where x 0 ∼ p 0 , so it has the same distribution as in Eq. (2). Remarkably, such SDEs are reversible in time [2]. Hence, the diffusion process of Eq. (4) can be viewed as a reversed-time diffusion:

where p ( x t , t ) is the density of x t at time t . We note that s ( x, t ) := ∇ x log p ( x, t ) is called the score function of x t at time t . Using Tweedie's lemma [15], one obtains the following relationship between the denoising function h and the score function:

Substituting Eq. (6) in Eq. (5), one obtains Eq. (3).

Training via denoising score matching. The standard way to train h θ is via denoising score matching . This is performed by obtaining samples of x 0 ∼ p 0 and x t ∼ N ( x 0 , σ 2 t I d ) and training to minimize

where the optimization is over some family of functions, { h θ } θ ∈ Θ . It was shown by Vincent [54] that optimizing Eq. (2) is equivalent to optimizing h θ in mean-squared-error on a random point x t that is a noisy image, x t ∼ N ( x 0 , σ 2 t I d ) where x 0 ∼ p 0 :

where h ∗ is the true denoising function from Eq. (2).

3 Theory

We define below the Consistency Property that a function h should satisfy. Simply put, it states that the output of h ( x, t ) (which is meant to approximate the conditional expectation of x 0 conditioned on x t = x ) is consistent with the average point x 0 generated using h and conditioning on x t = x . Recall from the previous section that generation according to h conditioning on x t = x is done by running the following SDE backwards in time conditioning on x t = x :

CP is therefore defined as follows:

Property 1 ( Consistency Property .) . A function h : R d × [0 , 1] → R d is said to satisfy CP iff for all t ∈ (0 , 1] and all x ∈ R d ,

where E h [ x 0 | x t = x ] corresponds to the conditional expectation of x 0 in the process that starts with x t = x and samples x 0 by running the SDE of Eq. (7) backwards in time (where note that the SDE uses h ).

The following lemma states that Property 1 holds if and only if the model prediction, h ( x, t ) , h ( x t , t ) is a reverse-Martingale under the same process of Eq. (7).

Lemma 3.1. Property 1 holds if and only if the following two properties hold:

The function h is a reverse-Martingale, namely: for all t > t ′ and for any x :

where the expectation is over x t ′ that is sampled according to Eq. (7) with the same function h , given the initial condition x t = x .

For all x ∈ R d , h ( x, 0) = x .

The proof of this lemma is included in Section B.2. Further, we introduce one more property that will be required for our theoretical results: the learned vector-field should be conservative.

Property 2 ( Conservative vector field / Score Property. ) . Let h : R d × [0 , 1] → R d . We say that h induces a conservative vector field (or that satisfies the score property) if for any t ∈ (0 , 1] there exists some probability density p ( · , t ) such that

We note that the optimal denoiser, i.e., h , defined as in Eq. (2), satisfies both of the properties we introduced. In the paper, we will focus on enforcing CP and we are going to assume conservativeness for our theoretical results. This assumption can be relaxed to hold only at a single t ∈ (0 , 1] using the results of Lai et al. [33].

Next, we show the theoretical consequences of enforcing Properties 1 and 2. First, we show that this enforces h to indeed correspond to a denoising function, namely, h satisfies Eq. (2) for some distribution p ′ 0 over x 0 . However, this does not imply that p 0 is the correct underlying distribution that we are trying to learn. Indeed, these properties can apply to any distribution p 0 . Yet, we can show that if we learn h correctly for some inputs and if these properties apply everywhere then h is learned correctly everywhere.

Theorem 3.2. Let h : R d × [0 , 1] → R d be a bounded continuous function. Then:

The function h satisfies both Properties 1 and 2 if and only if h is defined by Eq. (2) for some distribution p 0 .

Assume that h satisfies Properties 1 and 2. Further, let h ∗ be another function that corresponds to Eq. (2) with some initial distribution p ∗ 0 . Assume that h = h ∗ on some open set U ⊆ R d and some fixed t 0 ∈ (0 , 1] , namely, h ( x, t 0 ) = h ∗ ( x, t 0 ) for all x ∈ U . Then, h ∗ ( x, t ) = h ( x, t ) for all x and all t .

Remark 3.3 . While our theorem uses Eq. (2) which describes the VE-SDE (Variance Exploding SDE), it is also valid for VP-SDE (Variance Preserving SDE), as these two SDEs are equivalent up to appropriate scaling (see e.g. [31, 28]).

4 Method

Theorem 3.2 motivates enforcing CP on the learned model. We notice that the CP Equation Eq. (8) may be expensive to train for, because it requires one to generate whole trajectories. Rather, we use the equivalent Martingale assumption of Lemma 3.1, which can be observed locally with only partial trajectories 4 . We suggest the following loss function, for some fixed t, t ′ and x :

where the expectation E θ [ · | x t = x ] is taken according to process Eq. (7) parameterized by h θ with the initial condition x t = x . Differentiating this expectation, one gets the following (see Section B.1

for full derivation):

where p θ corresponds to the same probability measure where the expectation E θ is taken from and ∇ θ h θ corresponds to the Jacobian matrix of h θ where the derivatives are taken with respect to θ . Notice, however, that computing the expectation accurately might require a large number of samples. Instead, it is possible to obtain a stochastic gradient of this target by taking two samples, x t ′ and x ′ t ′ , independently, from the conditional distribution of x t ′ conditioned on x t = x and replace each of the two expectations in the formula above with one of these two samples.

We further notice the gradient of the CP loss can be written as

In order to save on computation time, we trained by taking gradient steps with respect to only the first summand (term A) in this decomposition. This term appears in line (4), while we ignored the term in line (9). Notice that if CP is preserved then this term becomes zero, which implies that no update is made, as desired.

It remains to determine how to select t, t ′ and x t ′ . Notice that t has to vary throughout the whole range of [0 , 1] whereas t ′ can either vary over [0 , t ] , however, it is sufficient to take t ′ ∈ [ t -ϵ, t ] . However, the further away t and t ′ are, we need to run more steps of the reverse SDE to avoid large discretization errors. Instead, we enforce the property only on small time windows using that consistency over small intervals implies global consistency. We notice that x t can be chosen arbitrarily and two possible choices are to sample it from the target noisy distribution p t or from the model. See Section D, Algorithm 1 for a pseudocode with more implementation details.

Remark 4.1 . It is important to sample x t ′ conditioned on x t according to the specific SDE Eq. (7). While a variety of alternative SDEs exist which preserve the same marginal distribution at any t , they might not preseve the conditionals.

5 Experiments

For all our experiments, we rely on the official open-sourced code and the training and evaluation hyper-parameters from the paper ' Elucidating the Design Space of Diffusion-Based Generative Models ' [28] that, to the best of our knowledge, holds the current state-of-the-art on conditional generation on CIFAR-10 and unconditional generation on CIFAR-10, AFHQ ( 64 × 64 resolution), FFHQ ( 64 × 64 resolution). We refer to the models trained with our regularization as 'CDM (Ours)' and to models trained with vanilla Denoising Score Matching (DSM) as 'EDM' models. 'CDM' models are trained with the weighted objective:

while the 'EDM' models are trained only with the first term of the outer expectation. We also denote in the name whether the models have been trained with the Variance Preserving (VP) [51, 18] or the Variance Exploding [51, 50, 49], e.g. we write EDM-VP. Finally, for completeness, we also report scores from the models of Song et al. [51], following the practice of the EDM paper. We refer to the latter baselines as 'NCSNv3' baselines.

We train diffusion models, with and without our regularization, for conditional generation on CIFAR10 and unconditional generation on CIFAR-10 and AFHQ ( 64 × 64 resolution). For the re-trained models on CIFAR-10, we use exactly the same training hyperparameters as in Karras et al. [28] and we verify that our re-trained models match (within 1% ) the FID numbers mentioned in the paper. For AFHQ, we dropped the batch size from the suggested value of 512 to 256 to save on computational

resources, which increased the FID from 1 . 96 (reported value) to 2 . 29 . All models were trained for 200 k iterations, as in Karras et al. [28]. Finally, we retrain a baseline model on FFHQ for 150 k iterations and we finetune it for 5 k steps using our proposed objective.

Implementation Choices and Computational Requirements. As mentioned, when enforcing CP, we are free to choose t ′ anywhere in the interval [0 , t ] . When t, t ′ are far away, sampling x ′ t from the distribution p θ t ′ ( x ′ t | x t ) requires many sampling steps (to reduce discretization errors). Since this needs to be done for every Gradient Descent update, the training time increases significantly. Instead, we notice that local consistency implies global consistency. Hence, we first fix the number of sampling steps to run in every training iteration and then we sample t ′ uniformly in the interval [ t -ϵ, t ] for some specified ϵ . For all our experiments, we fix the number of sampling steps to 6 which roughly increases the training time needed by 1 . 5 x. We train all our models on a DGX server with 8 A 100 GPUs with 80 GBs of memory each.

5.1 Consistency Property Testing

We are now ready to present our results. The first thing that we check is whether regularizing for CP actually leads to models that are more consistent with their predictions, as the property implies. Specifically, we want to check that the model trained with L ours λ achieves lower Consistency error, i.e. lower L CP t,t ′ ,x t . To check this, we do the following two tests: i) we fix t = 1 and we show how L CP t,t ′ ,x t changes as t ′ changes in [0 , 1] , ii) we fix t ′ = 0 and we show how the loss is changing as you change t in [0 , 1] . Intuitively, the first test shows how the violation of CP splits across the sampling process and the second test shows how much you finally ( t ′ = 0 ) violate the property if the violation started at time t . The results are shown in Figures 1a, 1b, respectively, for the models trained on AFHQ. We include additional results for CIFAR-10, FFHQ in Figures 4, 5, 6, 7 of the Appendix. As shown, indeed regularizing for the CP Loss drops the L CP t,t ′ ,x t as expected. See Section C.1 for additional details and discussion.

0

10

20

30

40

50

60

70

80

t

0.000

0.001

0.002

0.003

0.004

0.005

||

x

0

p

(

x

0|

xt

)

[

x

0]

h

(

xt

,

t

)||

2

Baseline

Ours

Consistency Property Testing (AFHQ)

(a) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, L CP t,t ′ ,x t , behaves for t ′ = 0 , as t changes.

0

10

20

30

40

50

60

70

80

t

0.000

0.001

0.002

0.003

0.004

0.005

||

xt

p

(

xt

|

x

1)

[

h

(

xt

,

t

)]

h

(

x

1, 1.0)||

2

Baseline

Ours

Consistency Property Testing (AFHQ)

(b) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, L CP t,t ′ ,x t , behaves for t = 0 , as t ′ changes.

Performance. We evaluate the performance of our models. Following the methodology of Karras et al. [28], we generate 150 k images from each model and we report the minimum FID computed on three sets of 50 k images each. We keep checkpoints during training and we report FID for 30k , 70k , 100k , 150k , 180k and 200k iterations in Table 1. We also report the best FID found for each model, after evaluating checkpoints every 5k iterations (i.e. we evaluate 40 models spanning 200 k steps of training). As shown in the Table, the proposed CP regularization yields improvements throughout the training. In the case of CIFAR-10 (conditional and unconditional) where the re-trained baseline was trained with exactly the same hyperparameters as the models in the EDM [28] paper,

our CDM models achieve a new state-of-the-art. We further show that our CP regularization can be applied on top of a pre-trained model. Specifically, we train a baseline EDM-VP model on FFHQ 64 × 64 for 150 k using vanilla Denoising Score Matching. We then do 5 k steps of finetuning, with and without our CP regularization and we measure the FID score of both models. The baseline model achieves FID 2 . 68 while the model finetuned with CP regularization achieves 2 . 61 . This experiment shows the potential of applying our CP regularization to pre-trained models, potentially even at large scale, e.g. we could apply this idea with text-to-image models such as Stable Diffusion [42]. We leave this direction for future work.

Uncurated samples from our best models on AFHQ, CIFAR-10 and FFHQ are given in Figures 2a, 2b and 8. One benefit of the deterministic samplers is the unique identifiability property [51]. Intuitively, this means that by using the same noise and the same deterministic sampler, we can directly compare

visually models that might have been trained in completely different ways. We select a couple of images from Figure 2a (AFHQ generations) and we compare the generated images from our model with the ones from the EDM baseline for the same noises. The results are shown in Figure 3. As shown, the CP regularization fixes several geometric inconsistencies for the picked images. We underline that the shown images are examples for which CP regularization helped and that potentially there are images for which the baseline models give more realistic results.

We note that an appropriate regularization parameter λ for our consistency loss has to be chosen. We found that an excessively large λ harms the performance: while the CP regularization enforces the model to obey some diffusion process, it does not enforce it to obey the true diffusion process, and a large λ might disturb with the score-matching signal.

Ablation Study for Theoretical Predictions. One interesting implication of Theorem 3.2 is that it suggests that we only need to learn the score perfectly on some fixed t 0 and then the CP implies that the score is learned everywhere (for all t and in the whole space). This motivates the following experiment: instead of using as our loss the weighted sum of DSM and our CP regularization for all t , we will not use DSM for t ≤ t threshold , for some t threshold that we test our theory for.

We pick t threshold such that for 20% of the diffusion (on the side of clean images), we do not train with DSM. For the rest 80% we train with both DSM and our CP regularization. Since this is only an ablation study, we train for only 10 k steps on (conditional) CIFAR-10. We report FID numbers for three models: i) training with only DSM, ii) training with DSM and CP regularization everywhere, iii) training with DSM for 80% of times t and CP regularization everywhere. In our reported models, we also include FID of an early stopped sampling of the latter model, i.e. we do not run the sampling for t < t threshold and we just output h θ ( x t threshold , t threshold ) . The numbers are summarized in Table 2. As shown, the theory is predictive since early stopping the generation at time t gives significantly worse results than continuing the sampling through the times that were never explicitly trained for approximating the score (i.e. we did not use DSM for those times). That said, the best results are obtained by combining DSM and our CP regularization everywhere, which is what we did for all the other experiments in the paper.

6 Related Work

The fact that imperfect learning of the score function introduces a shift between the training and the sampling distribution has been well known. Chen et al. [6, 7] analyze how the l 2 error in the approximation of the score function propagates to Total Variation distance error bounds between the true and the learned distribution. Several methods for mitigating this issue have been proposed, but the majority of the attempts focus on changing the sampling process [51, 28, 24, 46]. A related work is the Analog-Bits paper [8] that conditions the model during training with past model predictions.

Karras et al. [28] discusses potential violations of invariances, such as the non-conservativity of the induced vector field, due to imperfect score matching. However, they do not formally test or enforce this property. Chao et al. [5] shows that failure to satisfy the conservativity property can harm the performance, and they propose a modification to relieve this degradation. Lai et al. [33] study the problem of regularizing diffusion models to satisfy the Fokker-Planck equation. While we show in Theorem 3.2 that perfect conservative training enforces the Fokker-Planck equation, we notice that their training method is different: they suggest enforcing the equation locally by using the finite differences method to approximate the derivatives. Further, they do not train on drifted data. Instead, we notice that our CP loss is well suited to handle drifted data since it operates across trajectories generated by the model.

A concurrent work by Song et al. [52] proposes Consistency Models (CM), a new class of generative models that output directly the solution of the Probability Flow ODE. This idea resembles the consistency in the model outputs that we enforce through CP, but the motivation is different: CM attempts to accelerate sampling and we attempt to improve generation quality. The two methods are similar in that they enforce a property that the model should satisfy. In the case of the CM, the ODE solver should produce the same solution when evaluated at points that belong to the same trajectory. Our property is that on expectation the predictions should not be changing for points that have the same origin. One way to view it is that CM applies our consistency condition but on the deterministic sampler (for which the expectation becomes the point itself). Subsequent work by Lai et al. [34] compares our CDM to Consistency Models [52] and to Fokker Planck regularization [33].

7 Conclusions and Future Work

We proposed an objective that enforces the trained network to follow a reverse Martingale, thereby having self-consistent predictions over time. We optimize this objective with points from the sampling distribution, effectively reducing the sampling drift observed in prior empirical works. Theoretically, we show that CP implies that we are sampling from the reverse of some diffusion process. Together with the assumption that the network has learned the score correctly in a subset of the domain, we can prove that CP (together with conservativity of the vector field) implies that the score is learned correctly everywhere - in the limit where the error approaches zero. Empirically, we use our objective to obtain state-of-the-art for CIFAR-10 and baseline improvements on AFHQ and FFHQ.

There are limitations of our method and several directions for future work. The proposed regularization increases the training time. It would be interesting to explore how to enforce CP in more effective ways in future work. Further, our method does not test nor enforce that the induced vector-field is conservative, which is a key theoretical assumption. Our method guarantees only indirectly improve the performance in the samples from the learned distribution by enforcing some invariant. Finally, our theoretical result holds in the limit where the error of our regularized objective approaches zero and it would be meaningful to theoretically study also the constant-error regime.

8 Acknowledgments

This research has been supported by NSF Grants CCF 1763702, AF 1901292, CNS 2148141, Tripods CCF 1934932, IFML CCF 2019844, the Texas Advanced Computing Center (TACC) and research gifts by Western Digital, WNCG IAP, UT Austin Machine Learning Lab (MLL), Cisco and the Archie Straiton Endowed Faculty Fellowship. Giannis Daras has been supported by the Onassis Fellowship, the Bodossaki Fellowship and the Leventis Fellowship. Constantinos Daskalakis has been supported by NSF Awards CCF-1901292, DMS-2022448 and DMS2134108, a Simons Investigator Award, the Simons Collaboration on the Theory of Algorithmic Fairness and a DSTA grant.

References

Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019 , 2022.

Brian D.O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Marius Arvinte, Ajil Jalal, Giannis Daras, Eric Price, Alex Dimakis, and Jonathan I Tamir. Single-shot adaptation using score-based models for mri reconstruction. In International Society for Magnetic Resonance in Medicine, Annual Meeting , 2022.

Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold Diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.09392 , 2022.

Chen-Hao Chao, Wei-Fang Sun, Bo-Wun Cheng, and Chun-Yi Lee. On investigating the conservative property of score-based generative models. 2023.

Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209.11215 , 2022.

Sitan Chen, Giannis Daras, and Alexandros G Dimakis. Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. arXiv preprint arXiv:2303.03384 , 2023.

Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint arXiv:2208.04202 , 2022.

Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210.01776 , 2022.

Giannis Daras and Alexandros G Dimakis. Multiresolution textual inversion. arXiv preprint arXiv:2211.17115 , 2022.

Giannis Daras, Yuval Dagan, Alexandros G Dimakis, and Constantinos Daskalakis. Scoreguided intermediate layer optimization: Fast langevin mixing for inverse problem. arXiv preprint arXiv:2206.09104 , 2022.

Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G Dimakis, and Peyman Milanfar. Soft diffusion: Score matching for general corruptions. arXiv preprint arXiv:2209.05442 , 2022.

Jacob Deasy, Nikola Simidjievski, and Pietro Liò. Heavy-tailed denoising score matching. arXiv preprint arXiv:2112.09788 , 2021.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association , 106(496):1602-1614, 2011.

Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H. Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion, 2022. URL https://arxiv.org/abs/2208.01618 .

Ulf Grenander and Michael I Miller. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series B (Methodological) , 56(4):549-581, 1994.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303 , 2022.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv:2204.03458 , 2022.

Wenyi Hong, Ming Ding, Wendi Zheng, Xinghan Liu, and Jie Tang. Cogvideo: Large-scale pretraining for text-to-video generation via transformers. arXiv preprint arXiv:2205.15868 , 2022.

Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. arXiv preprint arXiv:2209.05557 , 2022.

Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir. Robust compressed sensing mri with deep generative priors. Advances in Neural Information Processing Systems , 34:14938-14954, 2021.

Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2105.14080 , 2021.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4401-4410, 2019.

Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 8110-8119, 2020.

Tero Karras, Miika Aittala, Samuli Laine, Erik Härkönen, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Alias-free generative adversarial networks. Advances in Neural Information Processing Systems , 34:852-863, 2021.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. arXiv preprint arXiv:2206.00364 , 2022.

Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and Il-Chul Moon. Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation. In International Conference on Machine Learning , pages 11201-11228. PMLR, 2022.

Kwanyoung Kim and Jong Chul Ye. Noise2score: tweedie's approach to self-supervised image denoising without clean images. Advances in Neural Information Processing Systems , 34: 864-874, 2021.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations , 2021.

Chieh-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Regularizing score-based models with score fokker-planck equations. In NeurIPS 2022 Workshop on Score-Based Methods , 2022.

Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Naoki Murata, Yuki Mitsufuji, and Stefano Ermon. On the equivalence of consistency-type models: Consistency models, consistent diffusion models, and fokker-planck regularization. arXiv preprint arXiv:2306.00367 , 2023.

Calvin Luo. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970 , 2022.

Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems , 34: 25359-25369, 2021.

Sachit Menon, Alexandru Damian, Shijia Hu, Nikhil Ravi, and Cynthia Rudin. Pulse: Selfsupervised photo upsampling via latent space exploration of generative models. In Proceedings of the ieee/cvf conference on computer vision and pattern recognition , pages 2437-2445, 2020.

Eliya Nachmani, Robin San Roman, and Lior Wolf. Denoising diffusion gamma models. arXiv preprint arXiv:2110.05948 , 2021.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv , 2022.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. 2022.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487 , 2022.

Arne Schneuing, Yuanqi Du, Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blundell, Pietro Lió, Carla Gomes, Max Welling, et al. Structure-based drug design with equivariant diffusion models. arXiv preprint arXiv:2210.13695 , 2022.

Vikash Sehwag, Caner Hazirbas, Albert Gordo, Firat Ozgenel, and Cristian Canton. Generating high fidelity data from low-density regions using diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11492-11501, 2022.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems , 33:12438-12448, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021.

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 , 2023.

Brian L Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. arXiv preprint arXiv:2206.04119 , 2022.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

A Proof of Theorem 3.2

In Section A.1 we present a proof overview, in Section A.2 we present some preliminaries to the proof, in Section A.3 we include the proof, with proofs of some lemmas omitted and in the remaining sections we prove these lemmas.

A.1 Proof overview

We start with the first part of the theorem. We assume that h satisfies Properties 1 and 2 and we will show that h is defined by Eq. (2) for some distribution p 0 (while the other direction in the equivalence follows trivially from the definitions of these properties). Motivated by Eq. (6), define the function s : R d × (0 , 1] according to

We will first show that s satisfies the partial differential equation

where J s ∈ R d × d is the Jacobian of s , ( J s ) ij = ∂s i x j and each coordinate i of △ s ∈ R d is the Laplacian of coordinate i of s , ( △ s ) i = ∑ n j =1 ∂ 2 s i ∂x 2 j . In order to obtain Eq. (11), first, we use a generalization of Ito's lemma, which states that for an SDE

and for f : R d × [0 , 1] → R d , f ( x t , t ) satisfies the SDE

If f is a reverse-Martingale then the term that multiplies dt has to equal zero, namely,

By Lemma 3.1, h ( x t , t ) is a reverse Martingale, therefore we can substitute f = h and substitute µ = -g ( t ) 2 s according to Eq. (7), to deduce that

Substituting h ( x, t ) = σ 2 t s ( x, t ) + x according to Eq. (6) yields Eq. (11) as required.

Next, we show that any s ′ that is the score-function (i.e. gradient of log probability) of some diffusion process that follows the SDE Eq. (4), also satisfies Eq. (11). To obtain this, one can use the FokkerPlanck equation, whose special case states that the density function p ( x, t ) of any stochastic process that satisfies the SDE Eq. (4) satisfies the PDE

where △ corresponds to the Laplacian operator. Using this one can obtain a PDE for ∇ x log p which happens to be exactly Eq. (11) if the process is defined by Eq. (4).

Next, we use Property 2 to deduce that there exists some densities p ( · , t ) for t ∈ [0 , 1] such that

Denote by p ′ ( x, t ) the score function of the diffusion process that is defined by the SDE of Eq. (4) with the initial condition that p ( x, 0) = p ′ ( x, 0) for all x . Denote by s ′ ( x, t ) = ∇ x log p ′ ( x, t ) the score function of p ′ . As we proved above, both s and s ′ satisfy the PDE Eq. (11) and the same initial

condition at t = 0 . By the uniqueness of the PDE, it holds that s ( x, t ) = s ′ ( x, t ) for all t ≥ t 0 . Denote by h ∗ the function that satisfies Eq. (2) with the initial condition x 0 ∼ p 0 . By Eq. (6),

By Eq. (10) and since s = s ′ , it follows that h = h ∗ and this is what we wanted to prove.

We proceed with proving part 2 of the theorem. We use the notion of an analytic function on R d : that is a function f : R d → R such that at any x 0 ∈ R d , the Taylor series of f centered at x 0 converges for all x ∈ R d to f ( x ) . We use the property that an analytic function is uniquely determined by its value on any open subset: If f and g are analytic functions that identify in some open subset U ⊂ R d then f = g everywhere. We prove this statement in the remainder of this paragraph, as follows: Represent f and g as Taylor series around some x 0 ∈ U . The Taylor series of f and g identify: indeed, these series are functions of the derivatives of f and g which are functions of only the values in U . Since f and g equal their Taylor series, they are equal.

Next, we will show that for any diffusion process that is defined by Eq. (4), the probability density of p ( x, t 0 ) at any time t 0 > 0 is analytic as a function of x . Recall that the distribution of x 0 is defined in Eq. (4) as p 0 and it holds that the distribution of x t 0 is obtained from p 0 by adding a Gaussian noise N (0 , σ 2 t I ) and its density at any x equals

Since the function exp( -( x -a ) 2 / (2 σ 2 t )) is analytic, one could deduce that p ( x, t 0 ) is also analytic. Further, p ( x, t 0 ) > 0 for all x which implies that there is no singularity for log p ( x, t 0 ) which can be used to deduce that log p ( x, t 0 ) is also analytic and further that ∇ x log p ( x, t 0 ) is analytic as well.

We use the first part of the theorem to deduce that s is the score function of some diffusion process hence it is analytic. By assumption, s identifies with some target score function s ∗ in some open subset U ⊆ R d at some t 0 , which, by the fact that s ( x, t 0 ) and s ∗ ( x, t 0 ) are analytic, implies that s ( x, t 0 ) = s ∗ ( x, t 0 ) for all x . Finally, since s and s ∗ both satisfy the PDE Eq. (11) and they satsify the same initial condition at t 0 , it holds that by uniqueness of the PDE s ( x, t ) = s ∗ ( x, t ) for all x and t .

A.2 Preliminaries

Preliminaries on diffusion processes In the next definition we define for a function F : R d → R d its Jacobian J F , its divergence ∇· F and its Laplacian △ F that is computed separately on each coordinate of F :

Definition A.1. Given a function F = ( f 1 , . . . , f n ): R d → R d , denote by J F : R d → R d × d its Jacobian:

The divergence of F is defined as

Denote by △ F : R d → R d the function whose i th entry is the Laplacian of f i :

If F is a function of both x ∈ R d and t ∈ R , then J F , △ f and ∇· F correspond to F as a function of x , whereas t is kept fixed. In particular,

We use the celebrated Ito's lemma and some of its immediate generalizations:

Lemma A.2 (Ito's Lemma) . Let x t be a stochastic process x t ∈ R d , that is defined by the following SDE:

where B t is a standard Brownian motion. Let f : R d × R → R . Then,

Further, if F : R d × R → R d is a multi-valued function, then

Lastly, if x t is instead defined with a reverse noise,

then the multi-valued Ito's lemma is modified as follows:

Lastly, we present the Fokker-Planck equation which states that the probability distribution that corresponds to diffusion processes satisfy a certain partial differential equation:

Lemma A.3 (Fokker-Planck equation) . Let x t be defined by

where x t , µ ( x, t ) ∈ R d and B t is a Brownian motion in R d . Denote by p ( x, t ) the density at point x on time t . Then,

Preliminaries on analytic functions

Definition A.4. A function f : R d → R is analytic on R d if for any x 0 , x ∈ R d , the Taylor series of f around x 0 , evaluated at x , converges to f ( x ) . We say that F = ( f 1 , . . . , f n ): R d → R d is an analytic function if f i is analytic for all i ∈ { 1 , . . . , n } .

The following holds:

̸

Lemma A.5. If F, G : R d → R d are two analytic functions and if F = G for all x ∈ U where U ⊆ R d , U = 0 , is an open set, then F = G on all R d .

This is a well known result and a proof sketch was given in Section 3.

The heat equation. The following is a Folklore lemma on the uniqueness of the solutions to the heat equation:

Lemma A.6. Let p and p ′ be two continuous functions on R d × [ t 0 , 1] that satisfy the heat equation

Further, assume that p ( · , t 0 ) = p ′ ( · , t 0 ) . Then, p = p ′ for all t ∈ [ t 0 , 1] .

A.3 Main proof

In what appears below we denote

We start by claiming that if h satisfies Property 1, then s satisfies the PDE Eq. (11): (proof in Section A.4)

Lemma A.7. Let h satisfy Property 1 and define s according to Eq. (15) . Then, s satisfies Eq. (11) .

Next, we claim that the score function of any diffusion process satisfies the PDE Eq. (11): (proof in Section A.5)

Lemma A.8. Let s be the score function of some diffusion process that is defined by Eq. (4) . Then, s satisfies the PDE Eq. (11) .

To complete the first part of the proof, denote by p ( · , t ) the probability distribution such that s ( x, t ) = ∇ log p ( x, t ) , whose existence follows from Property 2. We would like to argue that { p ( · , t ) } t ∈ (0 , 1] corresponds the probability density of the diffusion

It suffices to show that for any t 0 > 0 , { p ( · , t ) } t ∈ ( t 0 , 1] corresponds to the same diffusion. To show the latter, let t 0 ∈ (0 , 1) and consider the diffusion process according to Eq. (16) with the initial condition that x t 0 ∼ p ( · , t 0 ) . Denote its score function by s ′ and notice that it satisfies the PDE Eq. (11) and the initial condition s ′ ( x, t 0 ) = ∇ x log p ( x, t 0 ) = s ( x, t 0 ) , where the first equality follows from the definition of a score function and the second from the construction of p ( x, t 0 ) . Further, recall that s ( x, t ) satisfies the same PDE Eq. (11) by Lemma A.4. Next we will show that s = s ′ for all t ∈ [ t 0 , 1] , and this will follow from the following lemma: (proof in Section A.6)

Lemma A.9. Let s and s ′ be two solutions for the PDE (11) on the domain R d × [ t 0 , 1] that satisfy the same initial condition at t 0 : s ( x, t 0 ) = s ′ ( x, t 0 ) for all x . Further, assume that for all t ∈ [ t 0 , 1] there exist probability densities p ( · , t ) and p ′ ( · , t ) such that s ( x, t ) = ∇ x log p ( x, t ) and s ′ ( x, t ) = ∇ x log p ′ ( x, t ) for all x . Then, s = s ′ on all of R d × [ t 0 , 1] .

Then, by uniqueness of the PDE one obtains that s = s ′ for all t ∈ [ t 0 , 1] . Hence, s is the score of a diffusion for all t ≥ t 0 and this holds for any t 0 > 0 , hence this holds for any t > 0 . This concludes the proof of the first part of the theorem.

For the second part, let s ∗ denote some score function of a diffusion process that satisfies Eq. (4). Assume that for some t 0 > 0 and some open subset U ⊆ R d , s = s ∗ , namely s ( x, t 0 ) = s ∗ ( x, t 0 ) for all t 0 > 0 and all x ∈ U . First, we would like to argue that if s ( x, t ) is the score function of some diffusion process that satisfies Eq. (4), then for any t 0 > 0 it holds that s ( x, t 0 ) is an analytic function (proof in Section A.7)

Lemma A.10. Let x t obey the SDE Eq. (4) with the initial condition x 0 ∼ µ 0 . Let t > 0 and let s ( x, t ) denote the score function of x t , namely, s ( x, t ) = ∇ x log p ( x, t ) where p ( x, t ) is the density of x t . Assume that µ 0 is a bounded-support distribution. Then, s ( x, t ) is an analytic function.

Since both s and s ∗ are scores of diffusion processes, then s ( x, t 0 ) and s ∗ ( x, t 0 ) are analytic functions. Using the fact that s = s ∗ on U ×{ t 0 } and using Lemma A.5 we derive that s ( x, t 0 ) = s ∗ ( x, t 0 ) for all x . Let p and p ∗ denote the densities that correspond to the score functions s and s ∗ and by definition of a score function, we obtain that for all x ,

which implies, by integration, that

for some constant c ∈ R . However, c = 0 . Indeed,

which implies that c = 0 as required. As a consequence, the following lemma implies that p ( x, 0) = p ∗ ( x, 0) for all x (proof in Section A.8):

Lemma A.11. Let x t and y t be stochastic processes that follow Eq. (4) with initial conditions x 0 ∼ µ 0 and y 0 ∼ µ ′ 0 and assume that µ 0 and µ ′ 0 are bounded-support. Assume that for some t 0 > 0 , x t 0 and y t 0 have the same distribution. Then, µ 0 = µ ′ 0 .

Without loss of generality, one can replace 0 with any ˜ t ∈ (0 , t 0 ) , to obtain that p ( x, ˜ t ) = p ∗ ( x, ˜ t ) for any ˜ t ∈ [0 , t 0 ] . Now, p ( x, t 0 ) is analytic, from Lemma A.5, hence it is continuous. Consequently, Lemma A.6 implies that p = p ∗ in R d × [ t 0 , 1] . This concludes that p = p ∗ in all the domain, which implies that s = ∇ log p = ∇ log p ∗ = s ∗ , as required.

A.4 Proof of Lemma A.7

We use Ito's lemma, and in particular Eq. (13), to get a PDE for the function h ( x t , t ) where x t satisfies the stochastic process

Ito's formula yields that

Since ( h, s ) satisfies Property 1 and using Lemma 3.1, h is a reverse martingale which implies that the term that multiplies dt has to equal zero. In particular, we have that

By Eq. (15),

Therefore,

Substituting this in Eq. (17) and using the relation dσ 2 t /dt = g ( t ) 2 that follows from Eq. (15), one obtains that

Dividing by σ 2 t , we get that

which is what we wanted to prove.

A.5 Proof of Lemma A.8

We present as a consequence of the Fokker-Plank equation (Lemma A.3) a PDE for the log density log p :

Then,

Proof. We would like to replace the partial derivatives of p that appear in Lemma A.3 with the partial derivatives of log p . Using the formula

one obtains that

Similarly,

which also implies that

This implies that

Further, notice that which implies that

Differentiating Eq. (18) again with respect to x i and applying Eq. (18) once more, one obtains that

Summing over i , one obtains that

Substituting the partials derivatives of p inside the Fokker-Planck equation in Lemma A.3, one obtains that

Dividing by p , one gets that

as required.

We are ready to prove Lemma A.8: Substituting µ = 0 in Lemma A.12, one obtains that

Taking the gradient with respect to x , one obtains that

Since ∂/∂x i commutes with ∂/∂t , it holds that

recalling that by definition s = ∇ log p . Further,

where for any function f : R d → R , H f is the Hessian function of f that is defined by

Lastly, we get that by the commutative property of partial derivatives,

Substituting Eq. (21), Eq. (22) and Eq. (23) in Eq. (20), one obtains that

as required.

A.6 Proof of Lemma A.9

We will prove that p and p ′ satisfy the same PDE (which is the heat equation). Recall that s and s ′ satisfy

By substituting s = ∇ log p ,

By exchanging the order of derivatives, we obtain that

By integrating, this implies that

where c ( t ) depends only on t . Eq. (19) shows that

By substituting this in the equation above, we obtain that

By multiplying both sides with p , we get that

Since p is a probability distribution,

Integrating over Eq. (24) we obtain that

where the last equation holds since the integral of a Laplacian of probability density integrates to 0 . It follows that c ( t ) = 0 which implies that

and the same PDE holds where p ′ replaces p , and this follows without loss of generality. Further, since log p and log p ′ are differentiable, it holds that p ( · , t ) and p ′ ( · , t ) are continuous for all fixed t . This implies that p and p ′ are continuous as functions of x and t since they both satisfy the heat equation Eq. (14). Consequently, Lemma A.6 implies that p = p ′ on R d × [ t 0 , 1] . Finally, s = ∇ log p = ∇ log p ′ = s ′ , as required.

therefore,

A.7 Proof of Lemma A.10

First, recall that since x t satisfies Eq. (4) with the initial condition x 0 ∼ µ 0 , then x t ∼ µ 0 + N (0 , σ 2 t I ) , namely, x t is the addition of a random variable drawn from µ 0 and an independent Gaussian N (0 , σ 2 t I ) . Therefore, the density of x t , which we denote by p ( x, t ) , equals

Using the equation we get that

By using the fact that the Taylor formula for e x equals

we obtain that the right hand side of Eq. (26) equals

̸

We will use the following property of analytic functions: if f and g are analytic functions over R d and g ( x ) = 0 for all x then f/g is analytic over R d . Since the denominator at the right hand side of Eq. (27) is nonzero, it suffices to prove that the numerator and the denominator are analytic. We will prove for the denominator and the proof for the numerator is nearly identical. By assumption of this lemma, the support of µ 0 is bounded, hence there is some M > 0 such that ∥ x ∥ ≤ M for any x in the support. Then,

This bound is independent on a , and summing these abvolute values of coefficients for i ∈ N , one obtains a convergent series. Hence we can replace the summation and the expectation in the denominator at the right hand side of Eq. (27) to get that it equals

This is the Taylor series around 0 of the above-described denominator it converges to the value of the denominator at any x . While this Taylor series is taken around 0 , we note the Taylor series around any other point x 0 ∈ R n converges as well. This can be shown by shifting the coordinate system by a constant vector such that x 0 shifts to 0 and applying the same proof. One deduces that the Taylor series for the denominator around any point x 0 converges on all R d , which implies that the denominator in the right hand side of Eq. (27) is analytic. The numerator is analytic as well by the same argument. Therefore the ratio, which equals s ( x, t ) , is analytic as well as required.

A.8 Proof of Lemma A.11

Let t > 0 , denote by µ t and µ ′ t the distributions of x t and x ′ t , respectively, and by p ( x, t ) and p ′ ( x, t ) the densities of these variables. Then, µ t = µ 0 + N (0 , σ 2 t I ) , namely, µ t is obtained by adding an independent sample from µ 0 with an independent N (0 , σ 2 t I ) variable, and similarly for µ ′ t . Hence, the density p ( x, t ) is the convolution of the densities p ( x, 0) with the density of a Gaussian N (0 , σ 2 t I ) .

Denote by ˆ p ( y, t ) the Fourier transform of the density p ( x, t ) with respect to x (while keeping t fixed) and similarly define ˆ p ′ as the Fourier transform of p ′ . Denote by g and by ˆ g the density of N (0 , σ 2 t I ) and its Fourier transform, respectively. Denote the convolution of two functions by the operator ∗ . Then,

Since the Fourier transform turns convolutions into multiplications, one obtains that

Since p ( x, t ) = p ′ ( x, t ) we obtain that ˆ p ( y, t ) = ˆ p ′ ( y, t ) . Consequently,

Since the Fourier transform of a Gaussian is nonzero, we can divide by ˆ g ( y ) to get that

This implies that the Fourier transform of p ( x, 0) equals that of p ′ ( x, 0) hence p ( x, 0) = p ′ ( x, 0) for all x , as required.

B Other proofs

B.1 Differentiating the loss function

Denote our parameter space as Θ ⊆ R m . In order to differentiate L SM t,t ′ ,x ( θ ) with respect to θ ∈ Θ , we make the following calculations below, and we notice that E θ is used to denote an expectation with respect to the distribution of x [ t ′ ,t ] according to Eq. (7) with s = s θ and the initial condition x t = x . In other words, the expectation is over x [ t ′ ,t ] that is taken with respect to the sampler that is parameterized by θ with the initial condition x t = x . We denote by p θ ( x [ t ′ ,t ] | x t = x ) the corresponding density of x [ t ′ ,t ] . For any function f = ( f 1 , . . . , f n ): Θ → R n , denote by ∇ θ f the Jacobian matrix of f , where

For notational consistency, if f is a single-valued function, namely, if n = 1 , then ∇ θ f is a column vector. We begin with the following:

Differentiating the whole loss, we get the following:

Let us compute the gradient of the log density. We use the discrete process, and let us assume that t = t 0 > t 1 > · · · > t k = t ′ are the sampling times. Then,

We assume that

Then,

Therefore

where C corresponds to the normalizing factor that is independent of θ . Differentiating, we get that

B.2 Proof of Lemma 3.1

In what appears below, the expectation E [ · | x t = x ] is taken with respect to the distribution obtained by Eq. (7), namely, the backward SDE that corresponds to the function s , with the initial condition x t = x . Similarly, E [ · | x t ′ ] is taken with the initial condition at x t ′ . To prove the first direction in the equivalence, assume that Property 1 holds and our goal is to prove the two consequences as described in the lemma. To prove the first consequence, by the law of total expectation and by the fact that x t -x t ′ -x 0 is a Markov chain, namely, x 0 and x t are independent conditioned on x t ′ , we obtain that

To prove the second consequence, by Property 1

This concludes the first direction in the equivalence.

To prove the second direction, assume that h ( x, t ) = E [ h ( x t ′ , t ′ ) | x t = x ] and that h ( x, 0) = x and notice that by substituting t ′ = 0 we derive the following:

as required.

C Additional Results

C.1 Property Testing

Discussion. We add some details about the plots:

We note that the plots are not monotonic. This might be attributed to the model being closer to the ideal denoiser at certain noise levels and perhaps for having the capacity to correct previous mistakes.

For each plot, we sample 32 images, and for each t we get 32 stochastic samples conditioned on each one of the images and we measure the violation of MP. Standard deviation was small that it was not visible in the plot.

σ t denotes the standard deviation of the noise at time t and it is a monotonic function of t .

We underline that for the perfect denoiser, the error would have been 0 everywhere. However, even after the MP regularization, the learning is not perfect hence there are errors.

Consistency Property Testing (CIFAR10) Backward Sampler (300) steps

0

10

20

30

40

50

60

70

80

t

0.0000

0.0005

0.0010

0.0015

0.0020

0.0025

||

x

0

p

(

x

0|

xt

)

[

x

0]

h

(

xt

,

t

)||

2

Baseline

Ours

Consistency Property Testing (CIFAR10) Backward Sampler (300) steps

0

10

20

30

40

50

60

70

80

t

0.0000

0.0005

0.0010

0.0015

0.0020

0.0025

||

xt

p

(

xt

|

x

1)

[

h

(

xt

,

t

)]

h

(

x

1, 1.0)||

2

Baseline

Ours

Consistency Property Testing (FFHQ) Backward Sampler (1000) steps

0

10

20

30

40

50

60

70

80

t

0.000

0.001

0.002

0.003

0.004

0.005

||

x

0

p

(

x

0|

xt

)

[

x

0]

h

(

xt

,

t

)||

2

Baseline

Ours

Consistency Property Testing (FFHQ) Backward Sampler (1000) steps

0

10

20

30

40

50

60

70

80

t

0.000

0.001

0.002

0.003

0.004

0.005

||

xt

p

(

xt

|

x

1)

[

h

(

xt

,

t

)]

h

(

x

1, 1.0)||

2

Baseline

Ours

C.2 Uncurated Samples

D Implementation details

Algorithm 1 Consistent Diffusion Models (CDM) Training

Input: dataset D , noise schedule { σ t } T 0 , initial model parameters θ , learning rate η , Consistency Property (CP) regularization strength λ , maximum distance ϵ between t and t ′ , and step size ∆ t for discretizing the Reverse SDE.

repeat

E Limitations

The capacity for generative models to exert consequential societal influence in myriad ways is undeniable, and it also brings along a multiplicity of inherent risks [37, 25, 26, 27]. These models, for instance, may be exploited to fabricate counterfeit images, and furthermore, they have the potential to intensify societal prejudices. This work does not seem to exert a direct influence on these particular biases. It is imperative to acknowledge that addressing such biases presents a substantial challenge.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduced Martingale Property (MP) in diffusion model and the training strategy to apply this property to modify the score network outputs. In short, the martingale property states the condition that the SDE trajectory conditioned in a sample at time t is equal to the denoising function. After showing that MP transfers the score functions of small t to that of large t, this paper adds additional loss term that considers MP that the denoising function at time t' to be equal to that at time t, where t' < t. With this loss term added with the EDM model, the image generation quality in terms of FID is improved.\n\nSTRENGTHS:\n* This paper introduces a unified technique to improve the diffusion model with simple fine-tuning with short SDE trajectory starting from noisy data. Training with this new loss improves the sampling performance with some margin for CIFAR and AFHQ datasets.\n* Compared to the concurrent work [1] that focused on the small-NFE performances, this directly learns \"better\" score function by reducing the learning variance, where the variance between x_t' and x_t is much smaller than that of x_0 and x_t.\n\nWEAKNESSES:\n* The FID performance gain is small, considering the additional training budget consumed. \n* Theoretically, training solely with MDM loss term should show fair generation quality, since when the MP is satisfied, the model should be learned properly. However, the performance became worse, as Table 2 have shown.\n* In the date of submission, there is a concurrent work [1] that also modifies the score network output by distilling knowledge from less noisy data.\n* The second term in the equation after line 198 is not yet considered. Adding (even after some approximation) this term may further improve the method.\n\n[1] Y. Song et al., \"Consistency models\", https://arxiv.org/abs/2303.01469\n\nQUESTIONS:\n* Is there any benchmarks when the MDM is trained longer without the DSM loss (extending the result in Table 2) than the result presented in the paper?\n* Because of the second strength (variance reduction by reducing the noise from $x_0\\to x_t$ to $x_{t'}\\to x_t$) this work may perform better in larger datasets. Are there any additional results produced for higher-resolution datasets?\n\n=====\n\nCorrections.\n\nLine 145: conditionning --> conditioning\n\nLine 163: is satisfies --> satisfies\n\nLine 203: sufficient --> is sufficient\n\nLine 262: 64x64 --> $64\\times64$\n\nLIMITATIONS:\nThe authors adequately addressed the limitations of this paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUsing approximate score functions and discretization can lead to compounding distribution shift as samples drift toward less likely regions of the training distribution. The authors propose to address this problem by enforcing that the learned denoiser satisfies an invariant they call the Martingale Property: $E_{gen} [x_0 | x_t = x] = h(x, t)$, where $h$ is the optimal (true) denoiser. The authors prove that if the learned denoiser satisfies MP and the score function is conservative, then there exists some underlying corresponding diffusion process; moreover, uniqueness holds. This implies that if in addition the learned denoiser matches the optimal denoiser at any time in a small region of the spatial domain, then the learned and optimal denoisers must match everywhere. Empirically, an extra loss term encouraging MP is added to the standard score matching loss, and results suggest this improves generated samples.\n\nSTRENGTHS:\n- The presentation of the theoretical results is clear.\n- Experimental results are promising. Although MP is expensive to enforce exactly, the experiments demonstrate that the proposed MP loss is effective in practice.\n\nWEAKNESSES:\n- The main theorem (Thm. 3.2) does not seem to be very actionable. The results assume that MP holds exactly, which is an unreasonable assumption when considering diffusion model training in practice.\n- Some of the theoretical results seem to have limited novelty. For example, uniqueness (Lemma A.9), which underpins the second claim of the main theorem, appears to be a standard argument.\n- The experiments are promising but somewhat limited. Results are reported for only two datasets (including CIFAR10 cond vs. uncond) and the benefits of the MP loss is not very clear.\n\nQUESTIONS:\n- In practice, how well does MP violation correlate to poor image quality?\n- How does enforcing the MP loss affect the image quality as NFE changes?\n\nLIMITATIONS:\nIt could be helpful to clarify the relationship between MP and the consistency property from [1]. In practice, does enforcing the martingale property in expectation along generated trajectories behave similarly to the consistency loss, which is enforced pointwise along ODE trajectories?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDiffusion models are trained to denoise images that have been generated using a particular schedule of additive Gaussian noise; however, when sampling from a trained diffusion model, one applies a chain of denoising steps, each one operating on the output of the previous step. This may lead to sampling drift, where the distribution of the images at intermediate denoising steps differs from the distribution encountered during training. This sampling drift is hypothesized to be detrimental to the sampling process.\n\nThis paper proposes an approach to mitigate this drift, by enforcing a Martingale property (MP), which states that the denoising function $h_{\\theta}(x, t)$ for all $t \\in (0, 1]$ and $x \\in \\mathbb{R}^d$ outputs $h_{\\theta}(x, t) = \\mathbb{E}[x_0 \\mid x_t = x]$, where this expectation is over the reverse diffusion process that starts from noisy example $x_t$ and runs the diffusion SDE, which applies the learned denoiser $h_{\\theta}(x, t)$. They show that to enforce the MP, $h_{\\theta}$ must satisfy a consistency property $h_{\\theta}(x, t) = \\mathbb{E}\\_{h} \\left[ h_{\\theta}(x_{t'}, t') \\mid x_t = x \\right]$ for all $t > t'$, where the expectation is over all $x_{t'}$ sampled using the reverse diffusion process with denoiser $h_{\\theta}$. In addition to the consistency property, they require the boundary condition $h(x, 0) = x$, such that at time $t=0$, the denoiser simply outputs the original input.\n\nIn order to encourage a denoiser to satisfy the MP, the authors propose a loss function that takes two successive points along the reverse diffusion trajectory, $x_t$ and $x_{t'}$, where $t'$ is slightly smaller than $t$, and minimizes the squared error between the denoiser output on those two points, $\\mathcal{L}\\_{t, t', x}(\\theta) = \\frac{1}{2} \\left( \\mathbb{E}\\_{\\theta}\\left[ h_{\\theta}(x_{t'}, t') \\mid x_t = x \\right] - h_{\\theta}(x_t, t) \\right)^2$. This term acts as a regularizer (with a weighting coefficient $\\lambda$) on top of the standard denoising objective.\n\nThe paper presents results on several real-world datasets, including CIFAR-10, AFHQ, and FFHQ. Overall, it is not completely convincing that enforcing the Martingale property is worth the effort (of computing the datapoint pairs $(x_t, x_{t'})$ and potentially tuning the contribution of the new loss term). The results do not seem to improve dramatically when using the MP term.\n\nSTRENGTHS:\n* The paper addresses an important problem with good motivation, to mitigate drift when sampling from a trained diffusion model. The idea to enforce the Martingale property is very nice, and the authors develop this idea into a practical method that adds an additional loss term to the standard denoising objective.\n\n* The proposed approach is theoretically justified, and fairly intuitive. The authors give a thorough discussion of the necessary assumptions and theoretical properties of their approach.\n\n* The authors apply the Martingale property regularizer to train diffusion models from scratch as well as to fine-tune pre-trained diffusion models. They evaluate the approach on three datasets---CIFAR-10, AFHQ, and FFHQ---and they show slight improvements over two baselines in terms of FID scores.\n\n* The authors also verify that their approach decreases the Martingale loss, which measures consistency in the denoised outputs, comparing this loss between a baseline model and one trained with the Martingale regularizer.\n\nWEAKNESSES:\n* The proposed framework has some similarities to Consistency Models (Song et al., 2023) and it would be helpful if the authors could add more discussion of this in the paper. I don't think the two approaches are identical. The core similarities are in the reverse-Martingale property, which looks basically the same as the consistency property for CM, that states that the output of the function should be equivalent whether it is computed on $x_t$ or $x_{t'}$ for any $t' < t$. Also, both CM and this work have a similar boundary condition for $t=0$. The motivations for the approaches are different, as CM aims to produce a 1-step sampler that maps directly from noise to a clean image, while the Martingale property is used to mitigate sampling drift to improve sample quality when using standard sequential sampling chains. Both approaches sample two points along the denoising trajectory, $x_t$ and $x_{t'}$, which are used to compute a consistency-like loss.\n\n* One apparent difference between MP and CM is that the MP uses an expectation over samples from the reverse diffusion process for a single input $x$. Also, MP seems to apply the same parameters $\\theta$ to both parts of the consistency function, while CM uses an exponential moving average of the parameters for one part? In MP, what is the number of samples used to approximate the expectation in the first part, $\\mathbb{E}\\_{\\theta}[h_{\\theta}(x_{t'}, t') \\mid x_t = x]$?\n\n* The paper should specify clearly what the difference is between $\\bar{B}_t$ and $B_t$ in equations 3 and 4.\n\n* In L33, why has the notation switched to using superscripts $p^*_0$? Does this denote anything different from the non-superscripted version $p_0$ used just before this sentence?\n\n* In L38, I think it would be clearer if the learned score function were denoted by writing its parameters $s_{\\theta}(x, t)$ rather than just $s(x, t)$, to clarify how it differs from the true score function $s^*(x, t)$. Otherwise, this distinction is subtle and unclear.\n\n* Equation numbers need to be added to the two equations in the introduction (page 2). Also, it is confusing that the order of the left- and right-hand sides are flipped between the first and second equation on page 2. I think it would be clearer for the second equation to be $h(x, t) = \\mathbb{E}[x_0 \\mid x_t = x], \\forall t \\in [0, 1] \\forall x$.\n\n* The superscripts used to denote the loss functions $L^1_{t, x_t, x_0}$ and $L_{t, t', x}^2$ are confusing because they look like they would represent $L_1$ and $L_2$ losses, respectively, while they are actually just indexes for the first and second loss terms. It would be better to rename them $L^{\\text{SM}}$ for the score matching loss and $L^{\\text{MP}}$ for the Martingale property loss.\n\n* The title, caption, and labels are too small in Figures 1(a,b). Why are the \"baseline\" and \"Ours\" curves following the same pattern (e.g., the same fluctuations)? Why are the curves non-monotonic? Are the plots shown for one example or are they an average over many examples? Could the authors report min/max values or the standard deviation over different examples for these curves? What do the values of $\\sigma_t$ from 0 to 80 represent?\n\n* Why is the FFHQ dataset not included in Table 1?\n\n**Minor**\n* L9 typo: \"describes conservative\" --> \"describes a conservative\"\n\n* L14 typo: \"in CIFAR-10\" --> \"on CIFAR-10\"\n\n* In Eq. 1, explain the notation $\\oplus$ as it is fairly non-standard.\n\n* The conventional notation for normal distributions is $\\mathcal{N}$, not $N$.\n\n* L24: \"$\\sigma_t$ is an increasing function\" --> $\\sigma_t$ is a value, not a function, so instead one can say that $\\sigma_t$ is given by an increasing function?\n\n* L44: It is strange to say \"Question 1\" when there are no other questions.\n\n* L84: \"The formal statement is summarized as follows below: Theorem 1.1 (informal)\" --> The \"formal\" statement immediately says \"informal.\"\n\n* In Eq. 2, why use this strange notation, rather than following the standard convention of using subscripts of $\\mathbb{E}$, that is $\\mathbb{E}_{x_0 \\sim p_0, x_t \\sim \\mathcal{N}(x_0, \\sigma_t^2 I_d)}$?\n\n* L41 typo: \"the larger is also the error\" --> \"the larger the error\"\n\n* In L52, \"relates multiple inputs to $s(\\cdot, \\cdot)$\" --> Before seeing the final consistency objective, it is unclear what is meant by this.\n\n* L58 typo: \"this phenomena\" --> \"this phenomenon\"\n\n* L59 typo: \"called optimal denoiser\" --> \"called the optimal denoiser\"\n\n* L64 typo: \"technique via score-matching\" --> \"technique of score-matching\"\n\n* L167 typo: \"can be relieved\" --> \"can be relaxed\"\n\n* L182: I do not think that the term VE-SDE or VP-SDE have been defined yet, before they are used.\n\n* L188 typo: \"trajectories:\" --> \"trajectories.\"\n\n* L195 typo: \"$x_{t'}$ and $x_{t'}$\" --> these are the same.\n\n* In the last equation on page 5, it is hard for the reader to see which is the first and second term of the summand. This equation should be numbered, and the summand terms could be labeled explicitly, so that it's clear which one is used in practice.\n\n* Typo in equation at the bottom of page 5: there is an extra bracket in $\\mathbb{E}_{\\theta}[]$.\n\n* L207 typo: \"choises\" --> \"choices\"\n\n* L228 typo: \"droped\" --> \"dropped\"\n\n* L262 typo: \"$64x64$\" --> \"$64 \\times 64$\"\n\n* Table 1 has a significant amount of empty space, as the last 4 rows for each dataset block are empty in 6 out of 7 columns. Also, most of the numbers in the final column come from other sources.\n\nQUESTIONS:\n* An important point regarding notation is that throughout the paper, it is unclear what certain expectations are taken with respect to. In the two equations on page 2, neither expectation has any subscript to indicate this, yet it is very important (it is written in the text below, but it is also crucial to be precise in the equations). Also, it is slightly confusing that the expectations are not written consistently: in some places, there is no subscript, in others the subscript is $\\mathbb{E}_h$, and in others it is $\\mathbb{E}\\_{\\theta}$. I am not sure whether any of these makes it clear that the sample $x_0$ is obtained from the learned reverse process defined in terms of $h$? Because $\\mathbb{E}_h$ implies that the expectation is over all possible functions $h$. Would it make sense to use alternative notation like $\\mathbb{E}\\_{x \\sim \\text{Reverse}(x_t, h\\_{\\theta})}$?\n\n* It is not clear how representative the examples in Figure 3 are, because it may be possible to find samples that are improved by MP and others that are not; did MP regularization help with the majority of images?\n\n* It would be good to add a discussion on why the second term of the last equation on page 5 can be dropped?\n\nLIMITATIONS:\n* As the authors acknowledge, they do not verify the assumption that the learned vector field is conservative, which is necessary for the theoretical results.\n\n* The authors do not provide an algorithm box to state the approach formally (not in the main paper nor in the appendix).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a new method to address the issue of imperfect score-matching in diffusion models by enforcing a Martingale Property (MP) during training. This MP property ensures that the model's predictions on its self-generated data are consistent with the generated outputs. The key contributions include the identification of an invariant property - the denoiser satisfies MP property - that any perfectly trained model should satisfy, proof of this principle is provided. Further, the paper introduces a novel training objective that enforces MP, optimizing the network to consistently predict data points from the learned distribution. Empirically, the paper showcases that using the proposed objective, paired with the original DSM loss, improves the generation quality in both conditional and unconditional generation contexts, validated on CIFAR-10, AFHQ, and FFHQ datasets.\n\nSTRENGTHS:\nThe paper is well organised and presented, easy to follow.\n\nThe work is novel, enforcing the MP property cast the drift issue into a relatively straightforward question. The methodology is appropriately described and presented and should be of interests to researchers in diffusion model field.\n\nIn my view, general clarity and quality of this paper is good, mathematical formulas are clear and the claims are generally supported by the results.\n\nWEAKNESSES:\nIt seems too much space was occupied by background introduction in the first 3.5 pages, might be good to further compress them.\n\nThe objective function presented in Section 4 makes sense if one follows the paper from the beginning, but might be a bit difficult for quick interpretation and implementation, would be nice if a further simplified version (or even a special case) under some certain conditions can be presented\n\nCan Figure 1 be combined as a single figure? - it doesn't add too much information to show them in two separate figures.\n\nQUESTIONS:\n1. Could the author provides a bit more detailed explanations to link the results in Table 1 (CIFAR10 parts) and Table 2?\n\n2. Is it possible to provide results using metrics other than FID, such as MMD or JS etc to evaluate the data distribution from a different perspective?\n\nLIMITATIONS:\nLimitations were clearly stated by the authors and the discussion of future work is sensible.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}