Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent
Authors: Giannis Daras
Yuval Dagan
Alex Dimakis
Constantinos Costis Daskalakis
Venue: NeurIPS 2023 poster (2023)
Keywords: diffusion models
sampling drift
Fokker-Planck
invariances
Stochastic Differential Equations
Martingales
TL;DR: We propose a novel training objective that enforces the network to have consistent predictions.

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
Giannis Daras*
Yuval Dagan"
Department of Computer Science University of Texas at Austin
Electrical Engineering and Computer Science University of California, Berkeley
Alexandros G. Dimakis
Constantinos Daskalakis
Department of ECE University of Texas at Austin
Electrical Engineering and Computer Science Massachusetts Institute of Technology
Abstract
Imperfect score-matching leads to shift between the training and the sampling distribution of diffusion models. Due 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 Consistency property (CP) which states that predictions of the model its generated data consistent across time. Theoretically, we show that the differential equation that describes CP together with the that describes a conservative vector field, have 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 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:/githubbcom//aamisaasaa//
1 Introduction
The liffusion-based
47
18 approach to generative models has been successful across various
modalities, including images 41
44 14
39
29
51 16 10 IT] videos [19
20 21, audio [32]
3D structures [40], proteins I 53
45
and medical applications
23 3].
Diffusion models generate data by first drawing a sample from 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 f. = down = 0, stochastic process {xt) (=e)rejo,1) wherein 20 is distributed according to the target distribution and, for all f,
It ~ Ps where Pt = po e N 12))
(I)
where @ denotes the convolution operator. That is, Pe is the distribution resulting from corrupting sample from Po with noise sampled from (000 of Ia), where is given by an increasing function of such that 0p = 0 and a is sufficiently large so that P1 nearly indistinguishable from pure noise. We note that diffusion models have been generalized to other types of corruptions by the recent works
'These authors contributed equally to this work.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

of Daras et al. [12), Bansal et al. A Hoogeboom and Salimans (22] Deasy et al. [B] Nachmani [38.
al.
In order to sample from a diffusion model, i.e. sample the afore-described process in reverse time, it suffices know the score function s(x, t) V. log p(r, where p(x, t) is the density of 2t ~ Pt- Indeed, given a sample Iς Pt. one can use the score function at 2t, i.c. s(x;,) t), to generate a sample from dt by taking an infinitesimal step of stochastic or an ordinary differential equation [5I 48 or by using Langevin dynamics [7 [50 Hence, in order to train a diffusion model to sample from target distribution of interest po it suffices to learn the score function s" "xx t) using samples from the corrupted distributions p; resulting from p and particular noise schedule dt. Notice that those samples can be easily drawn given samples from pū-
The Sampling Drift Challenge: Unfortunately, the true score function s" (z, t) is not perfectly learned during training. Thus, at generation time, the samples It drawn using the learned score function, so (z. t), in the ways discussed above, drift astray in distribution from the true corrupted distributions P.. This drift becomes larger for smaller due compounding of errors and is accentuated by the fact that the further away a sample x; is from the likely support of the true pi the larger the error [ () - 8 (are, I)l between the learned and the true score function at Ir, which feeds into an even larger drift between the distribution of dy from p. for L' < I; see e.g.
[46] 18 39 (). These challenges motivate the question:
Question. How can one train diffusion models to improve the error [1 t) - s" (a, t)ll between the
learned and true score function on inputs (x, where x is unlikely under the target noisy distribution Pi?
direct approach to this challenge is to train our model to minimize the afore-described error on pairs (ar, t) where is sampled from distributions other than p However, there is straightforward way to do so, because we do not have direct access to the values of the true score function s (xx, t).
This motivates us to propose 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 Ps. 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 CP improves sample quality. Further, we provide an ablation study that further provides evidence to this phenomenon of score propagation.
The true score function s² (zr, 1) is closely related to another function, called the optimal denoiser, which predicts a clean sample 20 Pă from noisy observation Ię Tin + am where the noise is r ~ (O, Ta). The optimal denoiser (under the &g loss) is the conditional expectation:
Our Approach:
A (z, t)
E Po
[TO toLn (0,a
Tt = x].
and the true score function can be obtained from the optimal denoiser as follows: s s'(x, t) (h" (z, =)) This result is known as Tweedie's 's Formula [3]. Indeed, the standard training
technique, via mathhng, explicitly trains for the score through the denoiser h* [54] 15 36 30 B3).
We are now ready to state our Consistency Property (CP). We will say that a (denoising) function ho(x, 1) satisfies CP iff
Eg Ir; = x] (h( VI E 1).1).
where the expectation is with respect to a sample from the learned reverse process, defined in terms of the implied score function sa(z, t) = (ha (z. t) - / when this is initialized at .r, I and run backwards in time to sample zg. See Eq. (3) for the precise stochastic differential equation and its justification. In particular, ho satisfies CP if the prediction he (a., f) of the conditional expectation of the clean image To given It = x equals the expected value of an image that is generated by the
"Some of these methods, such as Langevin dynamics, require also to know the score function in the neighborhood of 2z-

---

learned reversed process, starting from 24 = z. Equivalently, one can formulate this property as requiring z, to follow a reverse Martingale (see Lemma 3.1)
While there are several other properties that the score function of a diffusion process must sat- isfy, e.g. the Fokker-Planck equation [3] our first theoretical result is that the ho (x satisfying the Consistency Property suffices (in conjunction with the conservativeness of its score function so sa (hg (z, t) z)/ oo² to guarantee that 59 must be the score function of a diffusion process (and must thus satisty any other property that diffusion process must satisfy). If additionally se t) equals the score function s (z, of target diffusion process
single time f ty and an open subset of x: Rdd 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 ha(x, t) satisfies CP and its corresponding score function so(a,t) (ho (ar,t) o) is conservative field, then so t) is the score function of a diffusion process, i.e. the generation process using score function 8g. is the inverse of a diffusion process. If additionally so(x, (2, t) s" (x. t) for single = and for all in open subset of Rd where s* is the score function of a target diffusion process, then se (x, 1) = =s" (x, 1) everywhere.
Simply put, the above statement states that: i) satisfying CP and being 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 learn for a single to 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:
1. We identify an invariant property, consistency of the denoiser he, that any perfectly trained model should satisfy.
2. We prove that if the denoiser (a f) satisfies CP and its implied score function sp(x; 1) = ())) is a conservative field, then se (2r, f) 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.
3. We prove that optimizing for the score in 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.
4. We propose 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.
5. We show experimentally that, paired with the original Denoising Score Matching (DSM) loss, our objective improves generation quality on conditional and unconditional generation CIFAR-10, and in AFHQ and FFHQ
6. We conduct an ablation study which showcases that even if we do not optimize for DSM for some values of f, 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 [8] The function that one aims to learn, called the score function (defined below), is equivalent (up to a linear transformation) a denoising function [15] 54, whose
goal is to denoise an image that was injected with noise. In particular, for some target distribution Po. one's goal to learn the following function h: Rdd [0, RRd:
h t) Eso-po 2 (2o. [Fo I x].
(2)
In other words, the goal is to predict the expected "clean" image xo given a corrupted version of it, assuming that the image was sampled from po and its corruption was done by adding noise to
3

---

from [O, o La), where of is non-negative and increasing function of t. Given such a function h, we can generate samples from Po by solving a Stochastic Differential Equation (SDE) that depends on h [5I]. Specifically, one starts by sampling x1 from some fixed distribution and then runs the following SDE backwards in time:
dr g(())² h( 1) It dt + g(t)dBI.
of
where Be is reverse-time Brownian motion and g(1)² do? To explain how Eq. (3 derived, consider the forward SDE that starts with Clean image zo and slowly injects noise:
dx, g(L)dB,, xo po:
where Bt is standard Brownian motion. We notice here that the 2, under Eg. 4 is (ao, of where 20 ~ Po. 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 reversed -imee diffusion:
dx, ( logp(zr + g(O)BB,,
(5)
where p(r,, t) is the density of z3 at time We note that s(x, t) == Vz logp(z,t t) is called the score function of Iy at time t. Using Tweedie's lemma [15], one obtains the following relationship between the denoising function h and the score function:
Mxx t)
V logp(z,
(6)
oi
Substituting Eq.
66 in Eq. (5) one obtains Eq. 30.
Training via denoising score matching. The standard way to train he is via denoising score
matching. This is performed by obtaining samples of zo ~ Po and ie ~ NN of and training to
minimize
SM 1, (0) = Po ( 1) - roll
where the optimization is over some family of functions, {holece. It was shown by Vincent [54] that optimizing Eq. (2) is equivalent optimizing he in mnea--squared on random point 2; that is a noisy image, It e N N(ro, Ta) where Po:
Ex,, |1 t) h- (xx,
where h1 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 zo conditioned on 2, = æ) is consistent with the average point za generated using h and conditioning on x; r. Recall from the previous section that generation according to h conditioning on x is done by running the following SDE backwards in time conditioning on 24 z:
dr; -t)) h(2,, t) T3 dt + g(t)
of
CP is therefore defined as follows:
Property 1 (Consistency Property.). A function h: Rd x (0, 1] - R is said to satisfy CP iff for all t e (0, 1] and all x E Rd,
h(r,t) = E, [2o 2, z].
(8)
where Ex [To I, x) corresponds to the conditional expectation of xo in the process that starts with Ie and samples za by running the SDE of Eq. backwards in time (where note that the SDE uses h).
'A Brownian motion B, is stochatic process that injects noise while one goes forward in time, and a reverse-time Brownian motion B, injects noise while one goes backwards in time.

---

The following lemma states that Property| holds if and only if the model prediction, h(x, t), h(re, t) a ceverse-aatinngle under the same process of Eq. 7
Lemma 3.1. Property holds and only if the following two properties hold:
The function h is a reverse-Marte Phrerse-Martingglee namely: for all L >> C' and for any x:
h(x, () Fe,, ((xx)t))))))) =x]; x).
where the expectation is over ze that is sampled according Eq. with the same function h, given the initial condition It T.
For all x E RR4, (x,00)
The proof of this lemma is included in Section B.2 Further, introduce be required for our theoretical results: the learned vector-field should be conservative.
property that will
Property (Conservative vector field Score Property.). Let h: R x [o. 1] + R We say that induces conservative vector field (or that satisfies the score property) if for any (0, 1] there exists some probability density p(:;() such that
(1x, t) - of
Vlog p(x,
We note that the optimal denoiser, i.e., defined in Eq. ( satisfies both the properties 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 single (0. 1] using the results of Lai et al. [33].
First, we show that this enforces h to indeed correspond to a denoising function, namely, satisfies Eq. (2) for some distribution pó over Tig- However, this does imply that Po the correct underlying distribution that we are trying to learn. Indeed, these properties can apply to any distribution Yet, we can show that if we learn h correctly for some inputs and if these properties apply everywhere then is learned correctly everywhere.
Next, show the theoretical consequences of enforcing Properties D and 2
+ Rdd be bounded continuous function. Then: Theorem 3.2. Let h: Rd x [0, 1] -
1. The function h satisfies both Properties 7 andd) if and only h defined by Eq. for some distribution po-
2. Assume that h satisfies Properties ! and 2 Further, let he be another function that cor- responds to Eq. ( with some initial distribution põ- Assume that h h² some open set U C IRd and some fixed to (0, 1), namely, h(x, to) hh" hh"lz,, to) for all a E U. Then, 1r* (x, = h(x,t) t) for all x and all
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 to appropriate scaling (see e.g. 31 28
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.I which can be observed locally with only partial trajectories" We suggest the following loss function, for some fixed t, and a::
1 ||5 (ho(he x. ho t)]
CP L (0)
2
where the expectation Eg[. x:] is taken according process Eq. parameterized by he with the initial condition Z2 x. Differentiating this expectation, one gets the following (see Section B.]
Accordiig to Lemma 3.1 in order to completely train for Property one has to also enforce (a 0) however, this taken from the denoising matching objective Eq. 2

---

for full derivation):
Lo (0) = Eg [ho(ze, ha(22,
21 = x] Es ho(oe, V' log (polxe | Is = =x)) +
Voho(av, (') Voho(oz t)
It I
where Py corresponds to the same probability measure where the expectation Egg is taken from and Vahe corresponds to the Jacobian matrix of he where the derivatives are taken with respect to O. Notice, however, that computing the expectation accurately might require large number of samples. Instead, it is possible to obtain a stochastic gradient of this target by taking two samples, zy and x'e,, independently, from the conditional distribution of xy conditioned on 24 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 be written
1 Ve E (zo t')] ho t)]1
VOLI. =
2
Eg [1o (') ho (2x, 1)) Eg (V log (p(are) )) hg (xp,,())
In order to save computation time, trained by taking gradient steps with respect 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 update is made, as desired.
It remains to determine how to select L, and Ty.. Notice that has to vary throughout the whole range of [0, 1] whereas can either vary over [0, I). however, it is sufficient to take 1 E [ f]. However, the further away & and t' are, need to 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 z; can be chosen arbitrarily and two possible choices are to sample from the target noisy distribution pt or from the model. See Section D Algorithm [ for a pseudocode with more implementation details.
Remark 4.1. It is important sample xe conditioned æ according the specific SDE Eq. (7) While variety of alternative SDEs exist which preserve the same marginal distribution at any f, they might not preseve the conditionals.
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 of- ootte--f-the-arr on conditional generation on CIFAR-I0 and unconditional generation on CIFAR-10, AFHQ (64 x 64 resolution). FFHQ (64 x 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:
L (0) E,
E 20 20
SM 20 (0) + A Ezzeop Bes-a((--ed 11, (8)
CP
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) [8]]8]] the Variance Exploding [5]] 50 49], e.g. we write EDM-VE 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 CIFAR- 10 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 computational

---

resources, which increased the FID from 1.96 (reported value) to 2.29. All models were trained for 200k iterations, in Karras al. 28 Finally, we retrain a baseline model on FFHQ for 150k iterations and we finetune it for 5k steps using our proposed objective.
Implementation Choices and Computational Requirements. As mentioned, when enforcing CP, we are free to choose L' anywhere in the interval [0, 4]. When I, l' far away, sampling z from the distribution p. (z) (r)) requires many sampling steps (to reduce discretization errors). Since this needs to done for every Gradient Descent update, the training time increases significantly. Instead, notice that local consistency implies global consistency. Hence, first fix the number of sampling steps to run in every training iteration and then we sample uniformly in the interval [t c, f for some specified For all our experiments, we fix the number of sampling steps 6 which roughly increases the training time needed by 1.5x. We train all our models on a DGX server with A100 GPUs with 80GBs of memory each.
5.1 Consistency Property Testing
We 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 achieves lower Consistency emror, i.e. lower CP "t,",x, To check this, we do the following two tests: i) we fix 1 and we show how T. ,t, changes as t' changes in J0, IJ. ii) fix L' = and show how the loss is changing you change t [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 La 1b respectively, for the models trained on AFHQ. We include additional results for CIFAR- 10, FFHQ in Figures 6 of the Appendix. As shown, indeed regularizing for the CP Loss drops the CPC "t, as expected. See Section C.1 for additional details and discussion.
Comuireecy Tropeit Teating (4TT1))
Conplstancy Property Thatin (ATHIC)
1EG=-yyttnilll Dan 61 01-Aa Aobiy Li00
2itapl elj?
0600
6000
(a) Consistency Property Testing AFHQ. The plot illustrates how the Consistency Loss, L j21 bchaves for t' 0, t changes.
(b) Consistency Property Testing AFHQ, The plot illustrates how the Consistency Loss, LCP behaves for t 0, as t' changes.
L
Figure Consistency Property Testing on AFHQ.
Performance. We evaluate the performance of our models. Following the methodology of Karras et al. (28), we generate 150k images from each model and we report the minimum FID computed on three sets of 50k images each. We keep checkpoints during training and we report FID for 30k, 70k, 100k, 150k, 180k and 200k iterations in Table [] We also report the best FID found for each model, after evaluating checkpoints every iterations (i.e. we evaluate 40 models spanning 200k steps of training). As shown in the Table, the proposed CP regularization yields improvements throughout the training. the case of CIFAR-10 (conditional and unconditional) where the re-trained baseline trained with exactly the same hyperparameters as the models in the EDM (28) paper,
7

---

(a) Uncurated images by our model trained on AFIQ. FID: 2.21, NFEs: 79.
(b) Uncurated images by our conditional CIFAR-
10 model. FID: 1.77, NFEs: 35.
Figure Comparison of uncurated images generated by two different models.
Model
30k
70k
100k
150k
180k
200k
Best
CDM-VP (Ours)
3.00
2.44
2.30
2.31
2.25
2.44
2.21
EDM- VP (retrained)
3.27
2.41
2.61
2.43
2.29
2.61
2.26
EDM-VP (reported)
1.96
AFHQ
EDM-VE (reported)
2.16
NCSNv3-VP (reported)"
2.58
NCSNv3-VE (reported)*
18.52
CDM-VP (Ours)
2.44
1.94
1.88
1.88
1.80
1.82
1.77
EDM-VP (retrained)
2.50
1.99
1.94
1.85
1.86
1.90
1.82
EDM-VP (reported)
CIFAR-10 (cond.)
1.79
EDM-VE (reported)
1.79
NCSNv3-VP (reported)
2.48
NCSNv3-VE (reported)
3.11
CDM-VP (Ours)
2.83
2.21
2.14
2.08
1.99
2.03
1.95
EDM-VP (retrained)
2.90
2.32
2.15
2.09
2.01
2.13
2.01
EDM-VP (reported)
CIFAR-10 (uncond.)
1.97
EDM-VE (reported)
1.98
NCSNv3-VP (reported)
3.01
NCSNv3-VE (reported)
3.77
CDM-VP (finetuned)
2.61
FFHQ
EDM-VP (retrained)
2.68
Table FID results for deterministic sampling, using the Karras al. 28 second-orded samplers. For the CIFAR- 10 models, we do 35 function evaluations and for AFHQ 79.
our CDM models achieve a new state-of-the-arr. We further show that our CP regularization can be applied on top a pre-trained model. Specifically, we train a baseline EDM-VP model on FFHQ 64 64 for 150k using vanilla Denoising Score Matching. We then do 5k 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-traine trained models, potentially even at large scale, e.g. could apply this idea with text-to-image models such as Stable Diffusion leave this direction for future work.
42 We
Uncurated samples from our best models AFHQ, CIFAR-10 and FFHQ are given in Figures 2a 2b
8 One benefit of the deterministic samplers is the unique identifiability property [5T). 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 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-man mote-matching signal.
Model
FID
EDM (baseline)
5.81
CDM, all times
5.45
Figure 3: Visual comparison of EDM model (top) and CDM model (Ours, bottom) us- ing deterministic sampling initiated with the same noise. As seen, the CP regularization fixes several geometric inconsistencies and artifacts in the generated images. In order to obtain this comparison, we generated images both with EDM and CDM. We found that in 14 of these images, CDM provided considerable improvement and for the remain- der there was no significant difference.
CDM, for some
6.59
CDM, for some early stopped sampling
14.52
Table 2:
Ablation study on removing the DSM loss for some Table reports FID re- sults after 10k steps of training on CIFAR-10.
Ablation Study for Theoretical Predictions. One interesting implication of Theorem 3.2 that it suggests that we only need to learn the score perfectly on some fixed to and then the CP implies that the score learned everywhere (for all and in the whole space). This motivates the following experiment: instead of using our loss the weighted sum of DSM and our CP regularization for all I, we will not use DSM for S lehreshold, for some 4hreshold that we test our theory for.
We pick fthreshold 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 10k steps on (conditional) CIFAR-10. -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 f 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 < /threshold and we just output hip (Ztauestooss 4threshoia). The numbers are summarized in Table 2 As shown, the theory is predictive since early stopping the generation at time gives significantly worse results than continuing the sampling through the times that were never explicitly trained for approximating the score (i.c. 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 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. [ analyze how the l2 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). 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 conseevservtttvttyy of the induced vector field, due to imperfect score matching. However, they do not formally test or enforce this property. Chao et al. [ shows that failure to satisfy the conservativity property can harm the performance, and they propose a modification to relieve this degradation. Lai et al. [3]] 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, 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), 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 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 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 objective that enforces the trained network to follow 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 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 bascline improvements on AFHQ and FFHQ
There are limitations of our method and several directions for future work. The proposed regulariza- tion 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 201 1844, 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-1901 292, DMS-2022443 and DMS2134108, a Simons Investigator Award, the Simons Collaboration on the Theory of Algorithmic Faimness and a DSTA grant.
References
[1] Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15013 2022.
[2] Brian D.O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):3133326, 1982.
10

---

[3] Marius Arvinte, Ajil Jalal, Giannis Daras, Eric Price, Alex Dimakis, and Jonathan I Tamir. Single-shot adaptation using score-bac baree- models for mri reconstruction. In International Society for Magnetic Resonance in Medicine, Annual Meeting, 2022.
[4] Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie $ Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold Diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.092: 09392, 2022.
[5] Chen-Hao Chao, Wei-Fang Sun, Bo- Wun Cheng, and Chun- Yi Lee. On investigating the conservative property of score-based generative models. 2023.
[6] 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. 1215, 2022.
[7] Sitan Chen, Giannis Daras, and Alexandros G Dimakis. Restoration-degraaatiion beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. arXiv preprint arXiv:2303.03388 2023.
[8] Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning arXiv preprint arXiv:2208. 04200,, 2022.
[9] 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.
[10] Giannis Daras and Alexandros G Dimakis. Multiresolution textual inversion. arXiv preprint arXiv:2211.17115 2022.
[11] Giannis Daras, Yuval Dagan, Alexandros G Dimakis, and Constantinos Daskalakis. Score- guided intermediate layer optimization: Fast langevin mixing for inverse problem. arXiv preprint arXiv:2206.09104 2022.
[12] Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G Dimakis, and Peyman Milan- far. Soft diffusion: Score matching for general corruptions. arXiv preprint arXiv:2209.05444 2022.
[13] Jacob Deasy, Nikola Simidjievski, and Pietro Liò. Heavy-tailed denoising score matching. arXiv preprint arX.v:2112.009788 2021.
[14] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in Neural Information Processing Systems, 34:8780-8799 8794, 2021.
[15] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 1604496):602-1611 2011.
[16] 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
[17] Ulf Grenander and Michael I Miller. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series (Methodologicall, 56(4)559988) 1994.
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 6851, 2020.
[19] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao. Alexey Gritsenko, Diederik Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:22 10.02303, 2022.
[20] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. 2rXiv:2204.03458, 2022.

---

[21] 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.1586. 2022.
[22] Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. WMXiv2200.005577 2022.
arXiv preprint
[23] 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: 1014448444999 2021.
[24] Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taille Taillee,,,,,er, Tal Kachman, and loannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint 1050:0005.14088, 2021.
[25] 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.
[26] Tero Karras, Samuli Laine, Mika 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 81 $110-8119 19,9, 2020.
[27] 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-866 -63,, 2021.
[28] Tero Karras, Mika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- generative models. arXiv preprint 2rXiv:2206.00344, 2022.
[29] Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and II-Chul Moon. Soft truncation: A universal training technique of ccore-based diffusion model for high precision score estimation. In International Conference on Machine Learning, pages 11201-11228.
PMLR, 2022.
[30] Kwanyoung Kim and Jong Chul Ye. Noise2score: tweedie's approach self-supervised image denoising without clean images. Advances in Neural Information Processing Systems, 34: 864-874, 2021.
[31] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21707, 2021.
[32] 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.
[33] Chieh-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Regularizing score- score-based models with score fokker lokker-planck equations. In NeurlPS 2022 Workshop Score-Based Methods, 2022.
[34] 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.0000 00367, 2023.
[35] Calvin Luo. Understanding diffusion models: A unified perspective.
arXiv preprint
arXiv:2208.1 172208.11977 1970, 2022.
[36] Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon. Estimating high order gradients of the data distribution by denoising. Advances Neural Information Processing Systems, 34: 2021.
[37] Sachit Menon, Alexandru Damian, Shijia Hu, Nikhil Ravi, and Cynthia Rudin. Pulse: Self- supervised photo upsampling via latent space exploration of generative models. Proceedings of the ieee/cvf conference computer vision and pattern recognition, pages 2437-2445, 2020.
12

---

[38] Eliya Nachmani, Robin San Roman, and Lior Wolf. Denoising diffusion gamma models. arXiv preprint arXiv:2110.05948, 2021.
[39] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabílistic
models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[40] Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv, 2022.
[41] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint WwXiv:2204.06122 2022.
[42] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[43] 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.
[44] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205. /448, 2022.
[45] Arne Schneuing, Yuanqi Du, Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blundell, Pietro Lió, Carla Gomes, Max Welling, al. Structure- drug design with equivariant diffusion models. arXiv preprint 1. 2022.
[46] 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 Computer Vision and Pattern Recognition, pages 11492-11501, 2022.
[47] Jascha Sohl-Diccstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2226 2265. PMLR, 2015.
[48] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[49] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[50] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 13:12438-12444 2020.
[51] 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.
[52] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint 2333::2333.0166 2023.
[53] Brian 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.04/ 00/ 19,41199910 2022.
[54] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-167 2011.
13

---

A Proof of Theorem 3.2
In Section A. we present 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 f satisfies Properties | and and will show that h is defined by Eq. (2) for some distribution po (while the other direction in the equivalence follows trivially from the definitions of these properties). Motivated by Eq. (6). define the function
8: IR x (0, 1] according to
h(x, x
8 (a, 1)
(10)
of
We will first show that s satisfies the partial differential equation
as
Ōt = g(1) J&8 2 As
(11)
where Sa E Rddd is the Jacobian of (Ja)ij
Ds. and each coordinate of As E is the 25
Laplacian of coordinate s, (As); = 2j=1
8²
order to obtain Eq. (I]. first, we use
az
generalization of Ito's lemma, which states that for an SDE
dx, = p( l)dt + g()d
(12)
and for f: Rd [0 + RR. [xx,, 1) satisfies the SDE
0f Jrn g(t)² -5) di + g(B)
df (xx, ()
at
If f is reverse-Martingale then the term that multiplies dt has to equal zero, namely,
0s
Ot JIp s(1) = 0.
By Lemma 3.1 M(x, 1) is reverse Martingale, therefore we can substitute f and substitute -1)) 2s according to Eq. (7). to deduce that
ah at 9() Jns g(1)²
Ah 0.
Substituting h(xx,) t) = sz, t) + according to Eq. (6) yields Eq. as required.
Next, we show that any 5' that is the score-function (i.e. gradient of log probability) of some diffusion
process that follows the SDE Eq. 44) also satisfies Eq. (1I). To obtain this, one can use the Fokker- Planck equation, whose special case states that the density function p(z, of any stochastic process that satisfies the SDE Eq. (4) satisfies the PDE
Öp
9()) 2 Ap
Öt
where A corresponds to the Laplacian operator. Using this one can obtain a PDE for V, log p which happens to exactly Eq. (TT) the process defined by Eq. 4
Next, we use Property to deduce that there exists some densities p(,, t) for t E [0, 1] such that
h(x, t) = Vo t). of
s(x, t)
Denote by p/ (xr !) the score function of the diffusion process that is defined by the SDE of Eq. with the initial condition that p(x, 0) = ((, 0) for all æ. Denote by s' s(x V, log (x, 1) the score function of p'. As we proved above, both and s' satisfy the PDE Eq. II and the same initial
14

---

condition at t = 0. By the uniqueness of the PDE, it holds that s(z,t) = s' t) for all t 2 to.
Denote by h" the function that satisfies Eq. 2 with the initial condition 2g Po- By Eq. (6).
h" h" t) s (x,,
of
By Eq. (10 and since s s', it follows that h and this is what we wanted to prove.
We proceed with proving part of the theorem. We use the notion of an analytic function on R that a function f: Rd - I such that at any zo Rd the Taylor series of centered zn converges for all x E I2 to / [xr). We use the property that an analytic function is uniquely determined by its value on any open subset: If f and are functtio functions that identify in some open subset U C R2 then everywhere. We prove this statement in the remainder of this paragraph, and as Taylor series around some Ti E U. The Taylor series of f and identify: indeed, these series functions of the derivatives of and g which functions of only the values in U. Since and equal their Taylor series, they are equal.
follows: Represent
Next, we will show that for any diffusion process that is defined by Eq. (4). the probability density of p(x, (o) at any time lo > 0 analytic as function of x. Recall that the distribution of xo is defined in Eq. (4) as po and it holds that the distribution of Tto is obtained from po by adding Gaussian
noise N (0, If and its density at any x equals
p(x, 10) = S. VTTT
1
(x - a)) 20?
exp
dpo(a).
Since the function exp(-(x- a)² (2o analytic, one could deduce that p(ar, (o) is also analytic. Further, p(z,, fo) > 0 for all æ which implies that there is no singularity for log ((r, ta) which can be used deduce that log p(2, to) also analytic and further that V log p(r, to) analytic as well.
We use the first part of the theorem to deduce that is the score function of some diffusion process hence it is analytie. By assumption, identifies with some target score function s* in some open subset U C Rd at some ln. which, by the fact that s(x, lo) and 8 (xx, 4n) are analytic, implies that s(r, to) s" (z,to) to) for all x. Finally, since s and both satisfy the PDE Eq. (IT) and they satsify the same initial condition at to, it holds that by uniqueness of the PDE s(x, t) s (æ, t) for all and t.
A.2 Preliminaries
Preliminaries on diffuusin processes In the next definition we define for a function F: Rd - Rd its Jacobian Jr. its divergence F and its Laplacian FF that is computed separately on each coordinate of F:
Definition A.I. Given a function = Jacobian:
Rd + Rd, denote by Jr: Rkd - Rdxd its (Jr) 0f((z) Or.-
The divergence of F is defined
V F(x) = Σ af(())
1=1
Denote by AF; IRd + Rd the function whose ith entry is the Laplacian of fs:
12
f2)x
(AF(x));
Ox7
If F is function of both Rd and / E R, then Jr. Af and V F correspond to F as a function of x, whereas is kept fixed. particular,
(f) (Jr(z, (o(xx))))= Oz,
f2,
0f(lr,)
(AF(I,t)) Σ
F
dn²
dr.
15

---

We use the celebrated Ito's lemma and some of its immediate generalizations:
Lemma A.2 (Ito's Lemma), Let de be a stochastic process at. Rd, that is defined by the following SDE:
dx, = u(xt; 1)dt + g())dBr, where B, a standard Brownian motion. Let f: Rd x R R. Then,
dt (f t) = + v.f u(xe,t) + g(t)²
df
2 Af) dt + g(t)Vof dB.
Further, F: Rd x R R is multi-valued function, then
ddF
g(1)
dF
dt Jrp
AF dt + g(t)JrdB.
2
Lastly, if x4 is instead defined with a reverse noise,
dx, = u(z,, t)dt + g
then the multi-valued Ito's 's's lemma is modified as follows:
c.F dt JFH
g()) 2 AF) de g(t)JrdB,.
dF (F 8)
(13)
Lastly, we present the Pokker-Plande equation which states that the probability distribution that corresponds to diffusion processes satisfy certain partial differential equation: Lemma A.3 (Fokker-Plann) equation). Let xt be defined by
dæ; - u( d) + g())dB..
where ze. p(x,t)= Rd and B, is Brownian motion in F. Denote by p(x,t) the density at point on time t. Then,
① at (x,tt) -V (1(z, t)p(x,t)) s(t)²
Apla,t) pv p- pv pt g(t)2 2 Ap.
Preliminaries on analytic functions
Definition A.4. A function f: IRd - R is analytic on Rd if for any No. x E Rdd the Taylor series of f around Tg. evaluated at æ. converges to fx)) We say that F (fi...., analytic function if fi is analytic for all i nI.
fi....,fa) fn): Rd Rd is an
The following holds:
Lemma A.5. If F, G: Rd - Red are two analytic functions and if F G for all U where C TR.. U + 0, is an open set, then F = on all Rdd
This is well known result and a proof sketch was given in Section
The following is Folklore lemma the uniqueness of the solutions the The heat equation. heat equation:
Lemma A.6. Let p and pf be two continuous functions on I [to, that satisfy the heat equation
Op
g(t)²
Ap.
(14)
Øt 2
Further, assume that p(.,to) (..ttt)) Then, = p' for all t E [to. 1].
A.3
Main proof
In what appears below we denote
h(x, t) of
s(x,
(15)
start by claiming that if h satisfies Property then s satisfies the PDE Eq.
(proof in
Section A.4
16

---

Lemma A.7. Let h satisfy Property and define s according to Eq. 15) Then, s satisfies Eq.
Next, we claim that the score function of any diffusion process satisfies the PDE Eq. (proof in Section A.55)
Lemma A.8. Let be the score function of some diffusion process that is defined by Eq. 4 Then, satisfies the PDE Eq. (11.
To complete the first part of the proof, denote by p(:: t) the probability distribution such that s(x, () = V log p(a, t), whose existence follows from Property We would like to argue that {p(,, he(0..1 corresponds the probability density of the diffusion
dæ, = g(t)dBB.
(16)
It suffices to show that for any to > 0, {p(;,t))E((to.. corresponds to the same diffusion. To show the latter, let to E (0, 1) and consider the diffusion process according to Eq. (I6) with the initial condition that Tto ~ pp-;;440 fo). Denote its score function by s' and notice that satisfies the PDE Eq. (TT) and the initial condition s' s'(2r, (n) = Va æ log p(ar, to) =(xx, fo), where the first equality follows from the definition of a score function and the second from the construction of p(a, fo). Further, recall
that s(x, t) satisfies the same PDE Eq. (II by Lemma Next we will show that = s' for all
t E [to, 1]. and this will follow from the following lemma: (proof in Section (A66)
Lemma A.9. Let s and s' be two solutions for the PDE (TT) on the domain TR. |tg: 1] that
satisfy the same initial condition at to: s(xr, to) = 8' (xx to) all x. Further, assume that for all t [to. 1] there exist probability densities p(-,t) and p' ('(,t) such that s(x,,) V. log p(x, t) and
s' 3'(x, () = v log r' (ax for all x. Then, 8 = s' on all of Rd [o.11.
Then, by uniqueness of the PDE one obtains that s = s' for all E [to, 1]. Hence, s is the score of a diffusion for all t to and this holds for any to O, 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 lo > 0 and some open subset C TRd = 8", namely s(x, 4o) = 8 (2, 4)) for all in > O and all x E U. First, we would like to argue that if s(x, t) is the score function of some diffusion process that satisfies Eq. then for any to > holds that s(x, to) is an analytic function (proof in Section A.7))
Lemma A.10. Let xe obey the SDE Eq. 4 with the initial condition To ~ -o- Let > and let s(x, E) denote the score function of xv, namely, s(x,t) Va log p(x, t) where p(x, t) is the density of It. Assume that po is a bounded- counded-support distribution. Then, s(x, t) is an analytic function.
Since both and s* are scores of diffusion processes, then s(x, ty) and s* (ar, to) are analytic functions. Using the fact that for all x. Let and p" denote the densities that correspond to the score functions and s* and by definition of a score function, we obtain that for all x,
s* on U }to} and using Lemma A.5 we derive that s(x, = s* (z, to)
log p(x, to) = s(x, to) s" (x, V fx,,t,
which implies, by integration, that
logp(r, to) = log (s,, to) +
for some constant e E R. However, Indeed,
1 = f p(z, to) ʃ / = Jo +e (1tt)( J (x, = e
which implies that = 0 as required. As consequence, the following lemma implies that p(z, 0) p" (xx 0) for all x (proof in Section A.8h.
Lemma A.11. Let Ie and yt be stochastic processes that follow Eq. 4 with initial conditions Do ~ po and yo ~ pbu and assume that pn and pgg are bounded-support. Assume that for some to > O, Itg and ytu have the same distribution. Then, 1i0 1/--
Without loss of generality, one can replace 0 with any E (0. to). to obtain that p(z, i) pxx for any E [0, to] Now, p(x, to) is analytic, from Lemma A.5 hence it is continuous. Consequently, Lemma A.6 implies that p = p" in IRd x (Lo, 1]. This concludes that p = p² in all the domain, which implics that V logp= V loggp" s", as required.
17

---

A.4 Proof of Lemma A.7
We use Ito's lemma, and in particular Eq. 13p. to get PDE for the function h(2 t) where 2t satisfies the stochastic process
dr, = - t) + g())BB,
Ito's formula yields that
[Oh
s(1)2 2 Sh dt + J.I.BB..
dh(
01 s(t)2
Since (h, 8) satisfies Property | and using Lemma 3.1 h is reverse martingale which implies that the term that multiplies dt has to equal zero. In particular, we have that
Oh
Jhx g(1)²
Ah = 0.
(17)
Øt
2
By Eq. (15).
h - x
oi
Therefore,
h= ofS.
Substituting this in Eq. 117 and using the relation do /t g()) that follows from Eq. (5 one obtains that
g(t)2
0 Ot
(x + os g(t)²
J. 2-0708
A(x o7
o g(t)2 g(t)²s s + oi 0s
2 As
Ot
Ōs
g(t)²oz
o o
As.
0t
2
Dividing by of. we get that
Os g(t)) J,s
g(t)2 As O,
which is what 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 logp:
Lemma A.12. Let xı be defined by
dx. = ulxt, t)dt + g())dB,.
Then,
i log Öt
p pV logp g(t)2 1l
g(t)2 log
+
2
2
Proof. We would like to replace the partial derivatives of p that appear Lemma A.3 with the partial derivatives of log Using the formula
o olog 1 0p
p Ot
one obtains that
Op logp
Ot Öt
Similarly,
Op logg
(18)
Or: Oz:
18

---

which also implies that
Vp pVlog,, v-p=pt logp.
Differentiating Eq.
18 again with respect to 24 and applying Eq. (88) once more, one obtains that
8.p
p
i Flogp
Op θ Glogp
8 Glogp
(
( logp
8a log
0x
Oz.
p
Ori
Or:
+p dx.
+
Ox7
Summing over i, one obtains that
(Qlos))
[OlogP] 2
8 loo
Ap=p2 = >
pl V log pl| pp logp.
(19)
Ox:
dx?
i=1
Substituting the partials derivatives inside the Fokker-Planck equation Lemma A.3 that
one obtains
i logg -V- p u(pV 1)g
r(t)²
(plVVoggPl' + p4 log
Ot
Dividing by p. one gets that
logg
V p- vv logp g(t) V log g(t)² logp
0t
2
as required.
We are ready to prove Lemma A.8
Substituting μr = 0 in Lemma A.12
one obtains that
i logg
((t log
g(t)² logp
+
Ōt
2
2
Taking the gradient with respect to æ, one obtains that
i logp
Voggll"
g A log
(20)
a
Since 0/Oz, commutes with 0/Ot, it holds that
Ologp
① V logp-
Os
(21)
Ot
Öt
0t
recalling that by definition 8 logp. Further,
0 Olog Øx:
82 log 2 ologp 0r,0z, 0z;
V
ll
=
2( 2(H V logp);
Or.
Ox.
1=1
where for any function f: 12d + R. Hj is the Hessian function of f that is defined by 02
(Hy)ss dx,0x,
This implies that
/ logpll = 2HH V logp.
Further, notice that
H, = Jas
which implies that Lastly, we get that by the commutative property of partial derivatives.
2 log V logp = 2J,s.
(22)
V log 4V log As.
(23)
Substituting Eq.
21 Eq.
22
and Eq.
23 in Eq. 20).. one obtains that
as
(t) J.s
Ōt
2
required.
19

---

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 e and s satisfy
0s
1
2(1²
s(1)
Ja8 + (x)
(VIl))2 As)
2
By substituting
V log
0V log
g(t)²
AV logp).
Ōt
2
By exchanging the order of derivatives, obtain that
Ölogp
s(t)
(IIV logpl + A logg).
ôt
By integrating, this implies that
Ölogp
g(t) (I log A log c(t),
0t
where c(t) depends only t. Eq.
19 shows that
Ap
logg
| logpl?
By substituting this in the equation above, we obtain that
Olog
g(t)² Ap
+ c(t).
ØL 2 p
By multiplying both sides with we get that
Op
2 loggp
g(t)²
Ap c(t).
(24)
at
- ØL
2
Since p is probability distribution,
Jpc 1(2x1)dx== =1,
therefore,
f
Öp(x, t)
② / p(z, p(z, t)dr 01
Ot dr Øt
0.
Ot
Integrating over Eq.
24
we obtain that
g(t)² 2
Ap c()da - 0+ c(()de,
-
where the last equation holds since the integral of a Laplacian of probability density integrates to 0. follows that c(l) which implies that
0p
g(t) 4p.
(25)
0t 2
and the same PDE holds where p' replaces P. and this follows without loss of generality. Further, since log p and logy' are differentiable, it holds that p(..() and p' (- () are continuous for all fixed This implies that p and p' are continuous as functions of x and since they both satisfy the heat equation Eq. [T4. Consequently, Lemma A.6 implies that p p' on I x (lo, 1]. Finally,
v logp
V log p' s', as required.
20

---

A.7 Proof of Lemma A.10
First, recall that since 2; satisfies Eq. (4) with the initial condition 2o ~ po. then ha ~ Ho+ of namely, is the addition of a random variable drawn from fl and an independent Gaussian NO0,0 I). Therefore, the density of 24, which we denote by p(z, t), equals
p a) Eurppo [v
1
(- [111-- 11---al aal2 )
exp
V2
20
Using the equation
V. log f(x) = Vaf(x)
f(x)
we get that
V.p(r,a)
Eawuo
E
v2n0, of
exp
|a aa-all 20;
s(x, a) = Va logp(a,
(26)
p(x,a)
Bawuo
exp
1r-oo al 2
/250,
By using the fact that the Taylor formula for e" equals
i!
1=0
we obtain that the right hand side of Eq. 26
equals
Eawpa
E-a
Eiso
(-1)'
|
Eampo
Z-a
(-1)
V2n0,
of
20
-1=0
2!
20
(27)
Earuo
V2Ter
Zi=0
(-1) 1 Ir-2
2a
Baruo
Ei
(-1) x--al
2
We will use the following property of analytic functions: f and g are analytic functions over Rd and g(x) * 0 for all x then f/g is analytic over R& Since the denominator at the right hand side of Eq. (27) is nonzero, 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 No bounded, hence there is some M > such that || S M for any x in the support. Then,
(-1) i! ( || all2) 20,
-
)
2
a
M
o
o.!!!
This bound is independent on a, and summing these abvolute values of coefficients for i E N, one obtains 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 equals
0C
(-1) i!
all 207
Σ
Eawwo
(28)
This is the Taylor series around 0 of the above-described denominator it converges to the value of the denominator at any æ. While this Taylor series is taken around 0, we note the Taylor series around any other point uo E IR" converges as well. This can be shown by shifting the coordinate system by constant vector such that zi shifts to 0 and applying the same proof. One deduces that the Taylor series for the denominator around any point To converges all Rd. 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 well required.
A.8 Proof of Lemma A.11
Let / >> 0, denote by ft and 14 the distributions of xe and a/, respectively, and by p(a, 1) and p( I) the densities of these variables. Then, /4, 14n + --(0,, of I), namely, μt is obtained by adding an independent sample from di with an independent, A/(0, of variable, and similarly for 14. Hence, the density p(æ, t) is the convolution of the densities p(z, 0) with the density of a Gaussian (0, of
21

---

Denote by p(y. t) the Fourier transform of the density p(z, t) with respect to x' (while keeping and similarly define as the Fourier transform of p'. Denote by and by و the density of N (0, of and its Fourier transform, respectively. Denote the convolution of two functions by the operator Then,
t fixed)
p(z, t) p(z, 0) g(x),
pl(z,t) = p'(,, g(z).
Since the Fourier transform turns convolutions into multiplications, one obtains that
ply 1) - p(v. 0)ă(μ), fly 1) - -i( (v, 0)ğ(y).
Since p(x, t) = pz, t) we obtain that p(y. t) = ('y,, t). Consequently, ply (v)(v) = ÿ (y. 0)g(y)
Since the Fourier transform of Gaussian is nonzero, we divide by O(v) get that fyy 0) p (v, 0).
This implies that the Fourier transform of p(z, 0) equals that of p' (x. 0) hence p(v. 0) = p2(x, 0) for all x, as required.
B
Other proofs
B.1
Differentiating the loss function
Denote our parameter space as θ C IR.. order to differentiate LEM (0) with respect to O, we make the following calculations below, and we notice that IE, used to denote an expectation with respect the distribution of æ4ea according to Eq. 7 with
and the initial condition Jt = x. In other words, the expectation is over æ[v.e] that is taken with respect to the sampler that parameterized by with the initial condition 2, z. We denote by ( Ie x) the corresponding density of x]e For any function (fi.... fn): 12 denote by Ve the Jacobian matrix of f, where
0f.
(Vof)s
00,"
For notational consistency, if is a single-valued function, namely, if 1, then /o column vector. We begin with the following:
Vo Eo [ho(xe,
= Vo
aad ho(xe, (')ooe xz = x) x)dæę
Voho (x [ x)dxr + / holze, ('oooet"VVoooopeeee z)dxe
Rd
V Vopo(TTe .9/tse = x))
E. (aho t')] E. ho(ww, 1')
po (2 .e| x) Es (Te, t')] Ey hhe (Nv, (') log (po (Tteee)))) I aee = xx)]
Differentiating the whole loss. we get the following:
V VoLSM SM = (0)
2 Vo (Eg ( (tlaareett")) t')] ho +))²
(Eolha(Zre (Ft,, 1')] ho(2, 1)) (Vo Elhoaar, (')] Voho(a, 1))
Eg [ho(se, (') ho(a, 1)] Te (Voho(vr, Voho 1)
Eo [ho
[)] Eg [ho V' log (po( x)) 2 =:)]
Let us compute the gradient of the log density. We use the discrete process, and let us assume that t to t to t' are the sampling times. Then,
po Iz = x)
(at, Dt, ).
22

---

We assume that
Po( 14-1) ((
Then,
| - (Iee. -1n-1)) 111
Pa(z 22 = 2)) x II
exp 2m1
29
Therefore
1
(1a, - Dzz 111
log (Z( =0==00 z) C+
Σ
2g
i=1
where C' corresponds to the normalizing factor that is independent of 8. Differentiating, get that
(No (Fe. =),_))) Vela,
Vy log Ie x)
97
B.2 Proof of Lemma 3.1
In what appears below, the expectation E[- x; = x]] is taken with respect to the distribution obtained by Eq. 7 namely, the backward SDD that corresponds to the function with the initial condition Te = r. Similarly, EE|- Tee is taken with the initial condition at Te. To prove the first direction in the equivalence, assume that Property holds and our goal is to prove the consequences as described in the lemma. To prove the first consequence, by the law of total expectation and by the fact that 2a Ze Zg a Markov chain, namely, 2n and my are independent conditioned on Iy', we obtain that
x2 = x] = [E((ae, 1 = x²) hr, 1) = E/ x = x²] [E[ [e)]
To prove the second consequence, by Property
(2r, T. (2n0
zi 2]] 30.
This concludes the first direction in the equivalence.
2t = x] and that h(,, 0) æ and To prove the second direction, assume that h(x, L) E1h(zr, (')
notice that by substituting /' = O we derive the following:
h( t) B[h(zo, 0) x, = z] = E|x
Øx æ].
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 measure the violation of MP. Standard deviation was small that it was not visible in the plot.
denotes the standard deviation of the noise at time 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.
23

---

Cansistency Property Testing (CIFARIO)
Sampler (300)
0.0025
Ours
0.0020
E Paloej [Xo] - hhg (x t)l
0.0010
0.0000
Figure 4: Consistency Property Testing CIFAR-10. The plot illustrates how the Consistency Loss, m. behaves for = 0, as changes.
LCP
Consistency Property Testing (CIFAR10) Sampler
Ours
( (Xe, t)] ha 1g 1g 1.
0.0200
0.0015
0.0000
0.0000
Figure 5: Consistency Property Testing CIFAR-10. The plot illustrates how the Consistency Loss,
CP
ay behaves for as t' changes.

---

Property (FFHO)
Backward Sampler (1000) steps
0.005
Ours
Exx Palooje [Xa] hh t)ll
0.002
0.001
Figure 6: Consistency Property Testing on FFHQ The plot illustrates how the Consistency Loss, LCP m. behaves for = 0, as changes.
Consistency Property Testing (FFHQ)
Backward Sampler (1000) steps
0.005
Ours
[E Peteojry whahaaa (Xe, t)] hh (x 1.2.1.011 I.o
0.004
0.002
0.001
Figure 7: Consistency Property Testing on FFHQ The plot illustrates how the Consistency Loss,
CP
ay a behaves for as t' changes.
25

---

C.2
Uncurated Samples
Figure 8: Uncurated generated images by our fine-tur senettnndd model FFHQ. FID: 2.61, NFEs: 79.
D Implementation details
Algorithm 1 Consistent Diffusion Models (CDM) Training
Input: dataset D, noise schedule (o:]0' initial model parameters 0, learning rate v,, Consistency Property (CP) regularization strength \, maximum distance between / and l', and step size Al for discretizing the Reverse SDE.
repeat
Sample xo D
(Clean image) b (Corruption level) (Corrupted image)
Sample uniformly from {o.)²
Sample 3i ~Noo, Ta)
1SM + (T () - xoll
D (Compute Score Matching loss) >> (Level with less corruption)
Sample t' Ult t]
Stopgrad
Sample x Iy by running Se. steps of the Reverse SDE (Sample less corrupted images)
(1) (2)
Resumegrad
C hos (2 ,t' - ho(zr)t)) (ho(=)?), (') ho( Enforce CP with stochastic gradient.
0 0 V (1SS AL.P)
(Update parameters)
until convergence
E Limitations
The capacity for generative models to exert consequential societal influence in myriad ways is undeniable, and it also brings along 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 exert direct influence on these particular biases. It is imperative to acknowledge that addressing such biases presents a substantial challenge.
26


Reviews:
[{"review_id": "O0c7TIE8yn", "reviewer": "Reviewer_UfbR", "summary": "This 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.", "questions": "1. 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?", "limitations": "Limitations were clearly stated by the authors and the discussion of future work is sensible.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "It 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank a lot the Reviewer for appreciating the novelty, the theoretical and the empirical results and the presentation of our work!\n\nWe will compress the background section and combine Figures 1a, 1b into a single Figure as recommended by the Reviewer.\n\n> The 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\nThat’s a fair comment, we thank the Reviewer for raising it. In the one page PDF response that accompanies this rebuttal, we include the Algorithm we use in practice to train models with our MP regularization. \n\nWe will make this Section more user-friendly by starting right away with this algorithm and by describing step by step how the method works. We will briefly give some intuition of what this loss is doing, which is to make sure that the network's prediction remains on average the same as we progress along the trajectories of the reverse SDE. We will abstract away the calculations needed to derive this (and point to the Appendix for the detailed calculations).\n\n> Could the author provides a bit more detailed explanations to link the results in Table 1 (CIFAR10 parts) and Table 2?\n\nAbsolutely! The results in Table 1 are obtained by training with the loss function in Section 5, i.e. with $L_{\\lambda}^{\\mathrm{ours}}(\\theta)$. This loss has two terms, the standard denoising score matching term and the MP regularization. This is the recommended way to train diffusion models with MP regularization and it leads to improved performance in image generation compared to the baseline that only uses the denoising score matching loss.\n\nTable 2 summarizes the results of an ablation study in support of our Theorem 3.2. According to this theorem, if we can learn the score perfectly only for some noise level (indexed by $t$), we satisfy the MP property everywhere and we have learned a conservative vector field, then we know the score perfectly everywhere. \n\nIn order to test whether this theoretical result captures (at least partially) practical behaviour of diffusion training, we train using the DSM objective only for $80$% of the diffusion time and we do not use it at all for last $20$% of the diffusion trajectory (on the side of clean images). We show that without the MP regularization for the last $20$% of the diffusion trajectory (MDM with early stopped sampling) the FID increases a lot (i.e. performance drops significantly). However, using the MP regularization for the last $20$%, the performance doesn’t drop significantly compared to using both the MP regularization and the DSM loss for the whole trajectory (as we did in Table 1). We hope this clarifies things and we will update the text accordingly to make this clearer.\n\n> 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\nWe followed the evaluation protocol (and the code) of the EDM paper that does not report such metrics. We will include those numbers for both the baseline and our method in the camera-ready version of our work. We thank the Reviewer for the suggestion! \n\nWe will be happy to answer additional questions if any throughout the Paper Discussion period. We thank again the Reviewer for their constructive feedback!"}, {"review_id": "mVz8j0N2hl", "reviewer": "Reviewer_G7yx", "summary": "Diffusion 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", "questions": "* 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\n", "limitations": "* 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).", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* 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.", "weaknesses": "* 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the extremely detailed and constructive feedback. We will incorporate all the thoughtful comments regarding the presentation of our work.\n\nWe agree that there are some similarities with the concurrent CM work and we will gladly discuss it. The motivation is different: CM attempts to solve the reverse ODE in one step. We attempt to improve generation quality. The two methods are similar in that they enforce some 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 MP condition but on the deterministic sampler (for which the expectation becomes the point itself). \n\nWe think that the similarities between these two works should be celebrated: by applying the MP regularization to the deterministic and the stochastic sampler you get two complimentary, nice properties: faster sampling and improved quality. \n\n> Also, MP seems to apply the same parameters $\\theta$ to both parts of the consistency function, [...]\n\nThis is true, we do not use EMA.\n\n> 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\nWe can get an unbiased estimator of the gradient, by only using two samples to estimate the expectation. This is what we do in all the experiments to avoid computational overhead (lines 195-197).\n\n> The paper should specify clearly what the difference is between $B_t$ and $\\bar{B}_t$ [...]\n\n$B_t$ is a Brownian motion whereas $\\bar{B}_t$ is a Brownian motion that runs backwards in time. We will clarify it in the paper.\n\n> In L33, why has the notation switched to using superscripts $p_0^*$? \n\nWe use the notation $p_0^*$ when we discuss a distribution that we would like to learn whereas we use $p_0$ when the discussion is only about the mathematical properties of diffusion processes. \n\n> In L38, I think it would be clearer if the learned score function were denoted by writing its parameters [...]\n\nWe will update the text accordingly to make this distinction clear. \n\n> Why are the \"baseline\" and \"Ours\" curves following the same pattern (e.g., the same fluctuations)? [...]\n\nThis figure shows how much the Martingale Property is violated as we keep $t’=0$ and we change $t$ from $0$ to $T$. $\\sigma_t$ denotes the standard deviation of the noise at time $t$ and it is an increasing function of $t$. For the perfect denoiser, the error would have been $0$ everywhere. However, even after the MP regularization, the learning is not perfect and as the distance of $t$ and $t’$ increases the property gets violated. We should not necessarily expect this function to be monotonic. The non-monotonic increase might be attributed to the model being closer to the ideal denoiser at certain noise levels and perhaps having the capacity to correct previous mistakes. The message of the plot is that the MDM models violate MP less. The fluctuations are similar. This could be related to the scheduling of the noise; both models might struggle to maintain MP at the same levels. Another explanation could be that we are adding the same noise realization. For this plot, we first 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 so small that it was not visible in the plot. We will include these details in our next revision.\n\n> Why is the FFHQ dataset not included in Table 1?\n\nFollowing the Reviewer’s recommendation, we will include it in Table 1 with the rest of the results and we will name it as FFHQ (finetuning) to avoid confusion.\n\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\nWe recently came across [this](https://arxiv.org/abs/2209.12753) paper that investigates how to make diffusion models satisfy the conservative property. We will make sure to add this reference.\n\n> The authors do not provide an algorithm box to state the approach formally (not in the main paper nor in the appendix).\n\nWe thank the Reviewer for the suggestion. We include an algorithm box in the one page PDF that accompanies this rebuttal, we will also include it in the main paper in the camera-ready version of our submission. \n\n> An important point regarding notation is that throughout the paper, it is unclear what certain expectations are taken with respect to [...].\n\nIn the first Eq. on page 2, the conditional expectation is with respect to the joint distribution between $x_0$ that is sampled from $p_0^*$ and $x_t$ that is obtained by adding noise: $x_t = x_0 + \\sigma_t \\eta$. In the second Eq., one again computes the conditional expectation of $x_0$ given $x_t = x$, but here, the conditional expectation is with respect to the joint distribution where $x_t$ is fixed and $x_0$ is sampled by running the differential equation described by Eq. (3) backward in time.\nWe will improve the readability of the paper according to the suggestions.\n\n> It is not clear how representative the examples in Figure 3 are [...]\n\nWe generated 64 images from the baseline and our method and by inspection we found that we considerably improved 14 of them. For the rest of the samples we did not see considerable differences. \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\nWe drop it because backpropagating through the sampler is computationally intensive. We note that CM does a similar stopgrad operation.\nIntuitively, dropping it might not matter much because of the score-matching term already giving sufficient signal to improve the Martingale loss. We will add this discussion to the paper.\n\n\n\n\n\n"}, {"review_id": "hBwrbvjQYN", "reviewer": "Reviewer_Qt1z", "summary": "Using 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.", "questions": "- 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?", "limitations": "It 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?", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their feedback. We are happy that the Reviewer appreciated the theoretical and experimental aspects of our work. \n\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\nWe agree that Theorem 3.2 is not actionable by itself. However, the theorem leads to our method (Section 4) which is very actionable since it gives a new loss function and then is used to train better diffusion models. \n   \nWhile we agree that the MP assumption will not hold exactly, satisfying this property is easier compared to learning the score of the true distribution. Namely, satisfying the MP property implies that we have learned the score of some distribution, not necessarily the score of the data distribution. The whole training and sampling of diffusion models has been based on the assumption that the latter is perfectly learned. Even though this assumption does not hold exactly in practice, it has led to a useful framework for generative modeling. Theoretical works have modeled how errors caused by violation of this assumption propagate, e.g. see [1, 2]. We think that a similar analysis could be carried out for the case where MP is approximately satisfied. Yet, studying error propagation in PDEs requires a new analysis and we will leave it for future work.\n\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\nThe goal of the theory section of this paper is to provide mathematical insight into the sampling drift problem. We develop it to suggest an algorithm to facilitate handling distribution shift in a principled fashion. \n\n\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\nWe have experiments on CIFAR-10 (conditional and unconditional), AFHQ and FFHQ (finetuning). While we agree that the FID benefit is relatively small, we also want to point out that improving the state-of-the-art in such standard datasets is extremely hard and the margins of improvement are already very thin. \n\n> In practice, how well does MP violation correlate to poor image quality?\n\nThat's a great question. In our loss function, we have two terms: the MP regularization and the score-matching loss. By increasing too much the weight of the MP term, we can even have detrimental effects in the generation quality -- MP (+ conservativeness) only ensures that we are learning some score function. However, for a properly chosen, fixed weight for the MP term, we get consistent performance benefits in generation quality when the MP loss drops. We will gladly discuss this in the paper and provide the training plots that demonstrate this.\n\n\n> How does enforcing the MP loss affect the image quality as NFE changes?\n\nGreat question! The MP regularized models outperform the baselines in all NFE regimes. We observed that the performance gap is even larger for low NFEs (MP models seem more robust to decreasing the number of steps), but the optimal performance is obtained at the same number of NFEs for both the baseline and the regularized models.\n\n> It 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?\n\nThere are some similarities with the concurrent CM work and we will gladly cite it and discuss it in our next revision. The motivation is different: CM attempts to solve the reverse ODE in one step. We attempt to improve generation quality. The two methods are similar in that they enforce some 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 MP condition but on the deterministic sampler (for which the expectation becomes the point itself).\n\nWe think that the similarities between these two works should be celebrated: by applying the MP regularization to the deterministic and the stochastic sampler you get two complimentary, nice properties: faster sampling and improved quality.\n\n\n[1] Sampling is as easy as learning the score\n\n[2] Restoration-Degradation Beyond Linear Diffusions: A Non-Asymptotic Analysis For DDIM-Type Samplers\n\n\nWe hope that our response addresses the concerns of the Reviewer and that it can lead to an increase in the rating of our work. If there are any additional questions, we will remain available throughout the Paper Discussion period!"}, {"review_id": "ON62BEkl8L", "reviewer": "Reviewer_w9Z2", "summary": "This 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.", "questions": "* 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$", "limitations": "The authors adequately addressed the limitations of this paper.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "* 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.", "weaknesses": "* 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank a lot the Reviewer for their feedback! We are glad that the Reviewer appreciated certain aspects of our work. In what follows, we attempt to address some remaining concerns.\n\n\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\nThere are some similarities with the concurrent CM work and we will gladly cite it and discuss it in our next revision. The motivation is different: CM attempts to solve the reverse ODE in one step. We attempt to improve generation quality. The two methods are similar in that they enforce some 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 MP condition but on the deterministic sampler (for which the expectation becomes the point itself).\n\nWe think that the similarities between these two works should be celebrated: by applying the MP regularization to the deterministic and the stochastic sampler you get two complimentary, nice properties: faster sampling and improved quality.\n\n> The FID performance gain is small, considering the additional training budget consumed.\n\nWhile we agree that the FID benefit is relatively small, we also want to point out that improving the state-of-the-art in such standard datasets is extremely hard and the margins of improvement are already very thin. \n\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\nThis is not exactly true. Only training with MDM loss guarantees that we learn the score function of some distribution, but not necessarily the score function of the true distribution. Our theoretical results suggest that we can use the score-matching loss only for certain noise levels and the MDM loss everywhere else and then learn (in theory) the score of the data distribution. The mere fact that we can get a reasonably good performance by not using the score-matching term for the last $20$% of the diffusion trajectory (in the side of clean images) was actually surprising to us and we consider it a validation test for our theoretical result.\n\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\nThat is true. However, backpropagating through the sampler is computationally intensive. The concurrent Consistency Models paper that the Reviewer brought to our attention also uses a stopgrad operation that has a similar effect.\n\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\nIn all the results in the paper, the MDM models and the baseline models were trained for the same number of steps (200K steps, following the setup of the EDM paper). For all our experiments, the optimal performance was obtained before the 200K steps -- the performance flattens or even slightly deteriorates after a while. Hence, we do not have any reason to believe that the performance would improve further for either the baseline or our models. \n\n>  Are there any additional results produced for higher-resolution datasets?\n\nUnfortunately, the computing required to train diffusion models from scratch for higher resolutions becomes enormous. For reference, the official EDM repository states that training on ImageNet requires 32 A100 GPUs for 13 days. This is way beyond our computational budget. However, as the Reviewer suggests, it is quite possible that our method would give even more significant benefits for higher resolution datasets.\n\nFinally, we want to thank the Reviewer for carefully reading our paper and for finding our typos. We commit to fixing them in the camera-ready version of our work. \n\nIf there are additional questions, we would be happy to answer them!\n"}]

Author rebuttal:
We thank the Reviewers for their constructive feedback. We reply separately to each of the Reviewers. Following the recommendation of Reviewer G7yx, we also include a formal statement of our training algorithm, given in the one-page PDF accompanying this rebuttal.

We thank again the Reviewers and we will gladly answer additional questions, if any!

Meta-review: The authors develop a technique to address errors in generation in a diffusion model translating to evaluations of the score model outside of where it was trained. The technique develops a new loss term to enforce a martingale property of the reverse diffusion. When combined with existing loss functions, the approach leads to improved results. After all of the discussion, all reviewers are positive.

I'd also suggest the authors change the acronym from MDM to something else as MDM was used recently in a diffusion paper for the term multivariate diffusion models at ICLR 2023 (https://openreview.net/pdf?id=osei3IzUia).
Decision: Accept (poster)