SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Giannis Daras*  
Department of Computer Science  
University of Texas at Austin
Alexandros G. Dimakis  
Department of ECE  
University of Texas at Austin
Yuval Dagan* Electrical Engineering and Computer Science University of California, Berkeley
Constantinos Daskalakis  
Electrical Engineering and Computer Science  
Massachusetts Institute of Technology
Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. However, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a Consistency property (CP) which states that predictions of the model on its own generated data are consistent across time. Theoretically, we show that the differential equation that describes CP together with the one that describes a conservative vector field, have a unique solution given some initial condition. Consequently, if the score is learned well on 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.
The diffusion-based [47, 49, 18] approach to generative models has been successful across various modalities, including images [41, 44, 14, 39, 29, 51, 43, 16, 10, 11], videos [19, 20, 21], audio [32], 3D structures [40], proteins [1, 53, 45, 9], and medical applications [23, 3].
Diffusion models generate data by first drawing a sample from a noisy distribution and slowly denoising this sample to ultimately obtain a sample from the target distribution. This is achieved by sampling, in reverse from time  $t = 1$  down to  $t = 0$ , a stochastic process  $\{x_{t}\}_{t\in [0,1]}$  wherein  $x_0$  is distributed according to the target distribution  $p_0$  and, for all  $t$ ,
$$
x _ {t} \sim p _ {t} \text {w h e r e} p _ {t} := p _ {0} \oplus \mathcal {N} \left(0, \sigma_ {t} ^ {2} I _ {d}\right), \tag {1}
$$
where  $\oplus$  denotes the convolution operator. That is,  $p_t$  is the distribution resulting from corrupting a sample from  $p_0$  with noise sampled from  $\mathcal{N}(0,\sigma_t^2 I_d)$ , where  $\sigma_t$  is given by an increasing function of  $t$ , such that  $\sigma_0 = 0$  and  $\sigma_1$  is sufficiently large so that  $p_1$  is nearly indistinguishable from pure noise. We note that diffusion models have been generalized to other types of corruptions by the recent works
of Daras et al. [12], Bansal et al. [4], Hoogeboom and Salimans [22], Deasy et al. [13], Nachmani et al. [38].
In order to sample from a diffusion model, i.e. sample the afore-described process in reverse time, it suffices to know the score function  $s(x,t) = \nabla_{x}\log p(x,t)$ , where  $p(x,t)$  is the density of  $x_{t} \sim p_{t}$ . Indeed, given a sample  $x_{t} \sim p_{t}$ , one can use the score function at  $x_{t}$ , i.e.  $s(x_{t},t)$ , to generate a sample from  $p_{t - dt}$  by taking an infinitesimal step of a stochastic or an ordinary differential equation [51, 48], or by using Langevin dynamics [17, 50]. Hence, in order to train a diffusion model to sample from a target distribution of interest  $p_0^*$  it suffices to learn the score function  $s^{*}(x,t)$  using samples from the corrupted distributions  $p_t^*$  resulting from  $p_0^*$  and a particular noise schedule  $\sigma_t$ . Notice that those samples can be easily drawn given samples from  $p_0^*$ .
The Sampling Drift Challenge: Unfortunately, the true score function  $s^*(x,t)$  is not perfectly learned during training. Thus, at generation time, the samples  $x_t$  drawn using the learned score function,  $s_\theta(x,t)$ , in the ways discussed above, drift astray in distribution from the true corrupted distributions  $p_t^*$ . This drift becomes larger for smaller  $t$  due to compounding of errors and is accentuated by the fact that the further away a sample  $x_t$  is from the likely support of the true  $p_t^*$  the larger the error  $\| s_\theta(x_t,t) - s^*(x_t,t) \|$  between the learned and the true score function at  $x_t$ , which feeds into an even larger drift between the distribution of  $x_{t'}$  from  $p_{t'}^*$  for  $t' < t$ ; see e.g. [46, 18, 39, 6]. These challenges motivate the question:
Question. How can one train diffusion models to improve the error  $\| s_{\theta}(x,t) - s^{*}(x,t)\|$  between the learned and true score function on inputs  $(x,t)$  where  $x$  is unlikely under the target noisy distribution  $p_t^*$ ?
A direct approach to this challenge is to train our model to minimize the afore-described error on pairs  $(x,t)$  where  $x$  is sampled from distributions other than  $p_t^*$ . However, there is no straightforward way to do so, because we do not have direct access to the values of the true score function  $s^*(x,t)$ .
This motivates us to propose a training method to mitigate sampling drift by enforcing that the learned score function satisfies an invariant, that we call the Consistency Property (CP). This property can be optimized without using any samples from the target distribution  $p_0^*$ . We will show that theoretically, enforcing CP on drifted points, in conjunction with minimizing the standard score matching objective on non drifted points, suffices to learn the correct score everywhere - at least in the theoretical limit where the error approaches zero and when one also enforces conservative vector field. We also provide experiments illustrating that regularizing the standard score matching objective using our CP improves sample quality. Further, we provide an ablation study that further provides evidence to this phenomenon of score propagation.
Our Approach: The true score function  $s^*(x, t)$  is closely related to another function, called the optimal denoiser, which predicts a clean sample  $x_0 \sim p_0^*$  from a noisy observation  $x_t = x_0 + \sigma_t \eta$  where the noise is  $\eta \sim \mathcal{N}(0, I_d)$ . The optimal denoiser (under the  $\ell_2$  loss) is the conditional expectation:
$$
h^{*}(x,t):= \mathbb{E}_{\substack{x_{0}\sim p_{0}^{*}\\ x_{t} = x_{0} + \sigma_{t}\eta \\ \eta \sim \mathcal{N}(0,I_{d})}}[x_{0}\mid x_{t} = x],
$$
and the true score function can be obtained from the optimal denoiser as follows:  $s^*(x,t) = (h^*(x,t) - x) / \sigma_t^2$ . This result is known as Tweedie's Formula [15]. Indeed, the standard training technique, via score-matching, explicitly trains for the score through the denoiser  $h^*$  [54, 15, 36, 30, 35].
We are now ready to state our Consistency Property (CP). We will say that a (denoising) function  $h_\theta(x,t)$  satisfies CP iff
$$
\mathbb {E} _ {\theta} \left[ x _ {0} \mid x _ {t} = x \right] = h _ {\theta} (x, t), \forall t \in [ 0, 1 ], \forall x,
$$
where the expectation is with respect to a sample from the learned reverse process, defined in terms of the implied score function  $s_{\theta}(x,t) = (h_{\theta}(x,t) - x) / \sigma_t^2$ , when this is initialized at  $x_{t} = x$  and run backwards in time to sample  $x_0$ . See Eq. (3) for the precise stochastic differential equation and its justification. In particular,  $h_\theta$  satisfies CP if the prediction  $h_\theta (x,t)$  of the conditional expectation of the clean image  $x_0$  given  $x_{t} = x$  equals the expected value of an image that is generated by the
<sup>2</sup>Some of these methods, such as Langevin dynamics, require also to know the score function in the neighborhood of  $x_{t}$ .
learned reversed process, starting from  $x_{t} = x$ . Equivalently, one can formulate this property as requiring  $x_{t}$  to follow a reverse Martingale (see Lemma 3.1).
While there are several other properties that the score function of a diffusion process must satisfy, e.g. the Fokker-Planck equation [33], our first theoretical result is that the  $h_{\theta}(x,t)$  satisfying the Consistency Property suffices (in conjunction with the conservativeness of its score function  $s_{\theta}(x,t) = (h_{\theta}(x,t) - x) / \sigma_t^2$ ) to guarantee that  $s_{\theta}$  must be the score function of a diffusion process (and must thus satisfy any other property that a diffusion process must satisfy). If additionally  $s_{\theta}(x,t)$  equals the score function  $s^{*}(x,t)$  of a target diffusion process at a single time  $t = t_0$  and an open subset of  $x \in \mathbb{R}^d$ , then it equals  $s^{*}$  everywhere. We comment that the formal theorem is proved for an idealistic setting when the error is (or approaches) zero. Still, it is likely to believe that even in the finite-error regime, training with DSM in-sample and enforcing CP off-sample is expected to improve the score function values off-sample. The statement is summarized as follows below:
Theorem 1.1 (informal). If some denoiser  $h_\theta(x,t)$  satisfies CP and its corresponding score function  $s_\theta(x,t) = (h_\theta(x,t) - x) / \sigma_t^2$  is a conservative field, then  $s_\theta(x,t)$  is the score function of a diffusion process, i.e. the generation process using score function  $s_\theta$ , is the inverse of a diffusion process. If additionally  $s_\theta(x,t) = s^*(x,t)$  for a single  $t = t_0$  and for all  $x$  in an open subset of  $\mathbb{R}^d$ , where  $s^*$  is the score function of a target diffusion process, then  $s_\theta(x,t) = s^*(x,t)$  everywhere.
Simply put, the above statement states that: i) satisfying CP and being a conservative vector field is enough to guarantee that the sampling process is the inverse of some diffusion process and ii) to learn the score function everywhere it suffices to learn it for a single  $t_0$  and an open subset of  $x$ 's.
We propose a loss function to train for the Consistency Property and we show experimentally that regularizing the standard score matching objective using our property leads to better models.
Summary of Contributions:
1. We identify an invariant property, consistency of the denoiser  $h_\theta$ , that any perfectly trained model should satisfy.
2. We prove that if the denoiser  $h_{\theta}(x,t)$  satisfies CP and its implied score function  $s_{\theta}(x,t) = (h_{\theta}(x,t) - x) / \sigma_t^2$  is a conservative field, then  $s_{\theta}(x,t)$  is the score function of some diffusion process, even if there are learning errors with respect to the score of the target process, which generates the training data.
3. We prove that optimizing for the score in a subset of the domain and enforcing these two properties, guarantees that the score is learned correctly in all the domains, in the limit where the error approaches zero.
4. We propose a novel training objective that enforces the Consistency Property. Our new objective optimizes the network to have consistent predictions on data points from the learned distribution.
5. We show experimentally that, paired with the original Denoising Score Matching (DSM) loss, our objective improves generation quality on conditional and unconditional generation on CIFAR-10, and in AFHQ and FFHQ.
6. We conduct an ablation study which showcases that even if we do not optimize for DSM for some values of  $t$ , satisfying CP enforces good score approximation there.
Diffusion processes, score functions and denoising. Diffusion models are trained by solving a supervised regression problem [49, 18]. The function that one aims to learn, called the score function (defined below), is equivalent (up to a linear transformation) to a denoising function [15, 54], whose goal is to denoise an image that was injected with noise. In particular, for some target distribution  $p_0$ , one's goal is to learn the following function  $h\colon \mathbb{R}^d\times [0,1]\to \mathbb{R}^d$ :
$$
h (x, t) = \mathbb {E} _ {x _ {0} \sim p _ {0}, x _ {t} \sim \mathcal {N} \left(x _ {0}, \sigma_ {t} ^ {2} I _ {d}\right)} [ x _ {0} \mid x _ {t} = x ]. \tag {2}
$$
In other words, the goal is to predict the expected "clean" image  $x_0$  given a corrupted version of it, assuming that the image was sampled from  $p_0$  and its corruption was done by adding noise to it
from  $\mathcal{N}(0, \sigma_t^2 I_d)$ , where  $\sigma_t^2$  is a non-negative and increasing function of  $t$ . Given such a function  $h$ , we can generate samples from  $p_0$  by solving a Stochastic Differential Equation (SDE) that depends on  $h$  [51]. Specifically, one starts by sampling  $x_1$  from some fixed distribution and then runs the following SDE backwards in time:
$$
d x _ {t} = - g (t) ^ {2} \frac {h \left(x _ {t} , t\right) - x _ {t}}{\sigma_ {t} ^ {2}} d t + g (t) d \bar {B} _ {t}, \tag {3}
$$
where  $\overline{B}_t$  is a reverse-time Brownian motion $^3$  and  $g(t)^2 = \frac{d\sigma_t^2}{dt}$ . To explain how Eq. (3) was derived, consider the forward SDE that starts with a clean image  $x_0$  and slowly injects noise:
$$
d x _ {t} = g (t) d B _ {t}, x _ {0} \sim p _ {0}, \tag {4}
$$
where  $B_{t}$  is a standard Brownian motion. We notice here that the  $x_{t}$  under Eq. (4) is  $\mathcal{N}(x_0,\sigma_t^2 I_d)$  where  $x_0\sim p_0$ , so it has the same distribution as in Eq. (2). Remarkably, such SDEs are reversible in time [2]. Hence, the diffusion process of Eq. (4) can be viewed as a reversed-time diffusion:
$$
d x _ {t} = - g (t) ^ {2} \nabla_ {x} \log p \left(x _ {t}, t\right) d t + g (t) d \bar {B} _ {t}, \tag {5}
$$
where  $p(x_{t}, t)$  is the density of  $x_{t}$  at time  $t$ . We note that  $s(x, t) \coloneqq \nabla_{x} \log p(x, t)$  is called the score function of  $x_{t}$  at time  $t$ . Using Tweedie's lemma [15], one obtains the following relationship between the denoising function  $h$  and the score function:
$$
\nabla_ {x} \log p (x, t) = \frac {h (x , t) - x}{\sigma_ {t} ^ {2}}. \tag {6}
$$
Substituting Eq. (6) in Eq. (5), one obtains Eq. (3).
Training via denoising score matching. The standard way to train  $h_\theta$  is via denoising score matching. This is performed by obtaining samples of  $x_0 \sim p_0$  and  $x_t \sim \mathcal{N}(x_0, \sigma_t^2 I_d)$  and training to minimize
$$
\mathbb {E} _ {x _ {0} \sim p _ {0}, x _ {t} \sim \mathcal {N} (x _ {0}, \sigma_ {t} ^ {2} I _ {d})} L _ {t, x _ {t}, x _ {0}} ^ {\mathrm {S M}} (\theta) = \mathbb {E} _ {x _ {0} \sim p _ {0}, x _ {t} \sim \mathcal {N} (x _ {0}, \sigma_ {t} ^ {2} I _ {d})} \| h _ {\theta} (x _ {t}, t) - x _ {0} \| ^ {2},
$$
where the optimization is over some family of functions,  $\{h_{\theta}\}_{\theta \in \Theta}$ . It was shown by Vincent [54] that optimizing Eq. (2) is equivalent to optimizing  $h_\theta$  in mean-squared-error on a random point  $x_t$  that is a noisy image,  $x_t \sim \mathcal{N}(x_0, \sigma_t^2 I_d)$  where  $x_0 \sim p_0$ :
$$
\mathbb {E} _ {x _ {t}} \left\| h _ {\theta} \left(x _ {t}, t\right) - h ^ {*} \left(x _ {t}, t\right) \right\| ^ {2},
$$
where  $h^*$  is the true denoising function from Eq. (2).
We define below the Consistency Property that a function  $h$  should satisfy. Simply put, it states that the output of  $h(x,t)$  (which is meant to approximate the conditional expectation of  $x_0$  conditioned on  $x_t = x$ ) is consistent with the average point  $x_0$  generated using  $h$  and conditioning on  $x_t = x$ . Recall from the previous section that generation according to  $h$  conditioning on  $x_t = x$  is done by running the following SDE backwards in time conditioning on  $x_t = x$ :
$$
d x _ {t} = - g (t) ^ {2} \frac {h \left(x _ {t} , t\right) - x _ {t}}{\sigma_ {t} ^ {2}} d t + g (t) ^ {2} d \bar {B} _ {t}, \tag {7}
$$
CP is therefore defined as follows:
Property 1 (Consistency Property.). A function  $h \colon \mathbb{R}^d \times [0,1] \to \mathbb{R}^d$  is said to satisfy  $CP$  iff for all  $t \in (0,1]$  and all  $x \in \mathbb{R}^d$ ,
$$
h (x, t) = \mathbb {E} _ {h} [ x _ {0} \mid x _ {t} = x ], \tag {8}
$$
where  $\mathbb{E}_h[x_0\mid x_t = x]$  corresponds to the conditional expectation of  $x_0$  in the process that starts with  $x_{t} = x$  and samples  $x_0$  by running the SDE of Eq. (7) backwards in time (where note that the SDE uses  $h$ ).
$^3$ A Brownian motion  $B_{t}$  is a stochastic process that injects noise while one goes forward in time, and a reverse-time Brownian motion  $\bar{B}_t$  injects noise while one goes backwards in time.
The following lemma states that Property 1 holds if and only if the model prediction,  $h(x,t)$ ,  $h(x_{t},t)$  is a reverse-Martingale under the same process of Eq. (7).
Lemma 3.1. Property 1 holds if and only if the following two properties hold:
- The function  $h$  is a reverse-Martingale, namely: for all  $t > t'$  and for any  $x$ :
$$
h (x, t) = \mathbb {E} _ {h} \left[ h \left(x _ {t ^ {\prime}}, t ^ {\prime}\right) \mid x _ {t} = x \right],
$$
where the expectation is over  $x_{t}$ , that is sampled according to Eq. (7) with the same function  $h$ , given the initial condition  $x_{t} = x$ .
- For all  $x \in \mathbb{R}^d$ ,  $h(x,0) = x$ .
The proof of this lemma is included in Section B.2. Further, we introduce one more property that will be required for our theoretical results: the learned vector-field should be conservative.
Property 2 (Conservative vector field / Score Property.). Let  $h \colon \mathbb{R}^d \times [0,1] \to \mathbb{R}^d$ . We say that  $h$  induces a conservative vector field (or that satisfies the score property) if for any  $t \in (0,1]$  there exists some probability density  $p(\cdot ,t)$  such that
$$
\frac {h (x , t) - x}{\sigma_ {t} ^ {2}} = \nabla \log p (x, t).
$$
We note that the optimal denoiser, i.e.,  $h$ , defined as in Eq. (2), satisfies both of the properties we introduced. In the paper, we will focus on enforcing CP and we are going to assume conservativeness for our theoretical results. This assumption can be relaxed to hold only at a single  $t \in (0,1]$  using the results of Lai et al. [33].
Next, we show the theoretical consequences of enforcing Properties 1 and 2. First, we show that this enforces  $h$  to indeed correspond to a denoising function, namely,  $h$  satisfies Eq. (2) for some distribution  $p_0'$  over  $x_0$ . However, this does not imply that  $p_0$  is the correct underlying distribution that we are trying to learn. Indeed, these properties can apply to any distribution  $p_0$ . Yet, we can show that if we learn  $h$  correctly for some inputs and if these properties apply everywhere then  $h$  is learned correctly everywhere.
Theorem 3.2. Let  $h \colon \mathbb{R}^d \times [0,1] \to \mathbb{R}^d$  be a bounded continuous function. Then:
1. The function  $h$  satisfies both Properties 1 and 2 if and only if  $h$  is defined by Eq. (2) for some distribution  $p_0$ .
2. Assume that  $h$  satisfies Properties 1 and 2. Further, let  $h^*$  be another function that corresponds to Eq. (2) with some initial distribution  $p_0^*$ . Assume that  $h = h^*$  on some open set  $U \subseteq \mathbb{R}^d$  and some fixed  $t_0 \in (0,1]$ , namely,  $h(x,t_0) = h^*(x,t_0)$  for all  $x \in U$ . Then,  $h^*(x,t) = h(x,t)$  for all  $x$  and all  $t$ .
Remark 3.3. While our theorem uses Eq. (2) which describes the VE-SDE (Variance Exploding SDE), it is also valid for VP-SDE (Variance Preserving SDE), as these two SDEs are equivalent up to appropriate scaling (see e.g. [31, 28]).
Theorem 3.2 motivates enforcing CP on the learned model. We notice that the CP Equation Eq. (8) may be expensive to train for, because it requires one to generate whole trajectories. Rather, we use the equivalent Martingale assumption of Lemma 3.1, which can be observed locally with only partial trajectories<sup>4</sup>. We suggest the following loss function, for some fixed  $t$ ,  $t'$  and  $x$ :
$$
L _ {t, t ^ {\prime}, x} ^ {\mathrm {C P}} (\theta) = \frac {1}{2} | | \mathbb {E} _ {\theta} [ h _ {\theta} (x _ {t ^ {\prime}}, t ^ {\prime}) \mid x _ {t} = x ] - h _ {\theta} (x, t) | | ^ {2},
$$
where the expectation  $\mathbb{E}_{\theta}[\cdot |x_t = x]$  is taken according to process Eq. (7) parameterized by  $h_\theta$  with the initial condition  $x_{t} = x$ . Differentiating this expectation, one gets the following (see Section B.1
4According to Lemma 3.1, in order to completely train for Property 1, one has to also enforce  $h_\theta (x,0) = x$  however, this is taken care from the denoising score matching objective Eq. (2).
for full derivation):
$$
\begin{array}{l} \nabla L _ {t, t ^ {\prime}, x} ^ {\mathrm {C P}} (\theta) = \mathbb {E} _ {\theta} \left[ h _ {\theta} (x _ {t ^ {\prime}}, t ^ {\prime}) - h _ {\theta} (x _ {t}, t) \mid x _ {t} = x \right] ^ {\top} \mathbb {E} _ {\theta} \left[ h _ {\theta} (x _ {t ^ {\prime}}, t ^ {\prime}) \nabla_ {\theta} \log \left(p _ {\theta} (x _ {t ^ {\prime}} \mid x _ {t} = x)\right) + \right. \\ \nabla_ {\theta} h _ {\theta} (x _ {t ^ {\prime}}, t ^ {\prime}) - \nabla_ {\theta} h _ {\theta} (x _ {t}, t) \left| x _ {t} = x \right], \\ \end{array}
$$
where  $p_{\theta}$  corresponds to the same probability measure where the expectation  $\mathbb{E}_{\theta}$  is taken from and  $\nabla_{\theta}h_{\theta}$  corresponds to the Jacobian matrix of  $h_\theta$  where the derivatives are taken with respect to  $\theta$ . Notice, however, that computing the expectation accurately might require a large number of samples. Instead, it is possible to obtain a stochastic gradient of this target by taking two samples,  $x_{t'}$  and  $x_{t'}'$ , independently, from the conditional distribution of  $x_{t'}$  conditioned on  $x_{t} = x$  and replace each of the two expectations in the formula above with one of these two samples.
We further notice the gradient of the CP loss can be written as
$$
\nabla_ {\theta} L _ {t, t ^ {\prime}, x} ^ {\mathrm {C P}} (\theta) = \frac {1}{2} \underbrace {\nabla_ {\theta} \| \mathbb {E} _ {\theta} [ h _ {\theta} (x _ {t ^ {\prime}} , t ^ {\prime}) ] - h _ {\theta} (x , t) \| ^ {2}} _ {A} +
$$
$$
\mathbb {E} _ {\theta} \left[ h _ {\theta} \left(x _ {t ^ {\prime}}, t ^ {\prime}\right) - h _ {\theta} (x, t) \right] ^ {\top} \mathbb {E} _ {\theta} \left[ \nabla_ {\theta} \log \left(p \left(x _ {t ^ {\prime}}\right)\right) h _ {\theta} \left(x _ {t ^ {\prime}}, t ^ {\prime}\right) \right] \tag {9}
$$
In order to save on computation time, we trained by taking gradient steps with respect to only the first summand (term A) in this decomposition. This term appears in line (4), while we ignored the term in line (9). Notice that if CP is preserved then this term becomes zero, which implies that no update is made, as desired.
It remains to determine how to select  $t, t'$  and  $x_{t'}$ . Notice that  $t$  has to vary throughout the whole range of [0, 1] whereas  $t'$  can either vary over  $[0, t]$ , however, it is sufficient to take  $t' \in [t - \epsilon, t]$ . However, the further away  $t$  and  $t'$  are, we need to run more steps of the reverse SDE to avoid large discretization errors. Instead, we enforce the property only on small time windows using that consistency over small intervals implies global consistency. We notice that  $x_t$  can be chosen arbitrarily and two possible choices are to sample it from the target noisy distribution  $p_t$  or from the model. See Section D, Algorithm 1 for a pseudocode with more implementation details.
Remark 4.1. It is important to sample  $x_{t'}$  conditioned on  $x_{t}$  according to the specific SDE Eq. (7). While a variety of alternative SDEs exist which preserve the same marginal distribution at any  $t$ , they might not preserve the conditionals.