SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sitan Chen* Sinho Chewi† Holden Lee‡ Yuanzhi Li§ Jianfeng Lu¶ Adil Salim
We provide the first polynomial-time convergence guarantees for the probability flow ODE implementation (together with a corrector step) of score-based generative modeling with an OU forward process. Our analysis is carried out in the wake of recent results obtaining such guarantees for the SDE-based implementation (i.e., denoising diffusion probabilistic modeling or DDPM), but requires the development of novel techniques for studying deterministic dynamics without contractivity. Through the use of a specially chosen corrector step based on the underdamped Langevin diffusion, we obtain better dimension dependence than prior works on DDPM  $(O(\sqrt{d})$  vs.  $O(d)$ , assuming smoothness of the data distribution), highlighting potential advantages of the ODE framework.
Score-based generative models (SGMs) [Soh+15; SE19; HJA20; DN21; Son+21a; Son+21b; VKK21] are a class of generative models which includes prominent image generation systems such as DALL-E 2 [Ram+22]. Their startling empirical success at data generation across a range of application domains has made them a central focus of study in the literature on deep learning [Aus+21; DN21; Kin+21; Shi+21; CSY22; Gna+22; Rom+22; Son+22; BV23; WHZ23]. In this paper, we aim to provide theoretical grounding for such models and thereby elucidate the mechanisms driving their remarkable performance.
Our work follows in the wake of numerous recent works which have provided convergence guarantees for denoising diffusion probabilistic models (DDPMs) [De +21; BMR22; De 22; LLT22; Liu+22; Pid22; WY22; Che+23a; CLL23; LLT23] and denoising diffusion implicit models (DDIMs) [CDD23]. We briefly recall that the generating process for SGMs is the time reversal of a certain diffusion process, and that DDPMs hinge upon implementing the reverse diffusion process as a stochastic differential equation (SDE) whose coefficients are learned via neural network training and the statistical technique of score matching [Hyv05; Vin11] (more detailed background is provided in §2). Among these prior works, the concurrent results of [Che+23a; LLT23] are remarkable because they require minimal assumptions on the data distribution (in particular, they do not assume log-concavity or similarly restrictive conditions) and they hold when the errors incurred during score matching are only bounded in an  $L^2$  sense, which is both natural in view of the derivation of score matching (see [Hyv05; Vin11]) and far more realistic. Subsequently, the work of [CLL23] significantly sharpened the analysis in the case when no smoothness assumptions are imposed on the data distribution.
*Harvard University, sitan@seas.harvard.edu
†Institute for Advanced Study, schewi@ias.edu
$^{\ddagger}$ Johns Hopkins University, hlee283@jhu.edu
Microsoft Research, yuanzhili@microsoft.com
$^{\dagger}$ Duke University, jianfeng@math.duke.edu
Microsoft Research, adilsalim@microsoft.com
<sup>7</sup>It is unreasonable, for instance, to assume that the score errors are bounded in an  $L^{\infty}$  sense, since we cannot hope to learn the score function in regions of the state space which are not well-covered by the training data.
Taken together, these works paint an encouraging picture of our understanding of DDPMs which takes into account both the diversity of data in applications (including data distributions which are highly multimodal or supported on lower-dimensional manifolds), as well as the non-convex training process which is not guaranteed to accurately learn the score function uniformly in space.
Besides DDPMs, instead of implementing the time reversed diffusion as an SDE, it is also possible to implement it as an ordinary differential equation (ODE), called the probability flow ODE [Son+21b]; see §2. The ODE implementation is often claimed to be faster than the SDE implementation [Lu+22; ZC23], with the rationale being that ODE discretization is typically more accurate than SDE discretization, so that one could use a larger step size. Indeed, the discretization error usually depends on the regularity of the trajectories, which is  $\mathcal{C}^1$  for ODEs but only  $\mathcal{C}^{\frac{1}{2}-}$  for SDEs (i.e., Hölder continuous with any exponent less than  $\frac{1}{2}$ ) due to the roughness of the Brownian motion driving the evolution.
Far from being able to capture this intuition, current analyses of SGMs cannot even provide a polynomial-time analysis of the probability flow ODE. The key issue is that under our minimal assumptions (i.e., without log-concavity of the data distribution), the underlying dynamics of either the ODE or SDE implementation are not contractive, and hence small errors quickly accumulate and are magnified. The aforementioned analyses of DDPMs managed to overcome this challenge by leveraging techniques specific to the analysis of SDEs, through which we now understand that stochasticity plays an important role in alleviating error accumulation. It is unknown, however, how to carry out the analysis for the purely deterministic dynamics inherent to the probability flow ODE.
Our first main contribution is to give the first convergence guarantees for SGMs with OU forward dynamics in which steps of the discretized probability flow ODE—referred to as predictor steps—are interleaved with corrector steps which runs the overdamped Langevin diffusion with estimated score, as pioneered in [Son+21b]. Our results are akin to prior works on DDPMs in that they hold under minimal assumptions on the data distribution and under  $L^2$  bounds on the score estimation error, and our guarantees scale polynomially in all relevant problem parameters. Here, the corrector steps inject stochasticity which is crucial for our proofs; however, we emphasize that the use of corrector steps does not simply reduce the problem to applying existing DDPM analyses. Instead, we must develop an entirely new framework based on Wasserstein-to-TV regularization, which is of independent interest; see §4 for a detailed overview of our techniques. Our results naturally raise the question of whether the corrector steps are necessary in practice, and we discuss this further in §5.
When the data distribution is log-smooth, then the dimension dependence of prior results on DDPMs, as well as our first result for the probability flow ODE with overdamped corrector, both scale as  $O(d)$ . Does this contradict the intuition that ODE discretization is more accurate than SDE discretization? The answer is no; upon inspecting our proof, we see that the discretization error of the probability flow ODE is indeed smaller than what is incurred by DDPMs, and in fact allows for a larger step size of order  $1 / \sqrt{d}$ . The bottleneck in our result stems from the use of the overdamped Langevin diffusion for the corrector steps. Taking inspiration from the literature on log-concave sampling (see, e.g., [Che22] for an exposition), our second main contribution is to propose corrector steps based on the underdamped Langevin diffusion (see §2) which is known to improve the dimension dependence of sampling. In particular, we show that the probability flow ODE with underdamped Langevin corrector attains  $O(\sqrt{d})$  dimension dependence. This dependence is better than what was obtained for DDPMs in [Che+23a; CLL23; LLT23] and therefore highlights the potential benefits of the ODE framework. We note that the benefit to which we refer is at generation time, and not at training time.
Previously, [JP22] have proposed a "noise-denoise" sampler using the underdamped Langevin diffusion, but to our knowledge, our work is the first to use it in conjunction with the probability flow ODE. Although we provide preliminary numerical experiments in the Appendix, we leave it as a question for future work to determine whether the theoretical benefits of the underdamped Langevin corrector are also borne out in practice.
In summary, our contributions are the following.
- We provide the first convergence guarantees for the probability flow ODE with overdamped Langevin corrector (DPOM; Algorithm 1).
- We propose an algorithm based on the probability flow ODE with underdamped Langevin corrector (DPUM; Algorithm 2).
- We provide the first convergence guarantees for DPUM. These convergence guarantees show improvement over (i) the complexity of DPOM  $(O(\sqrt{d})$  vs  $O(d))$  and (ii) the complexity of DDPMs, i.e., SDE implementations of score-based generative models (again,  $O(\sqrt{d})$  vs  $O(d)$ ).
- We provide preliminary numerical experiments in a toy example showing that DPUM can sample from a highly non log-concave distribution (see Appendix). The numerical experiments are not among our main contributions and are provided for illustration only. The Python code can be found in the Supplementary material.
Our main theorem can be summarized informally as follows; see §3 for more detailed statements.
Theorem 1 (Informal). Assume that the score function along the forward process is  $L$ -Lipschitz, and that the data distribution has finite second moment. Assume that we have access to  $\widetilde{O}(\varepsilon/\sqrt{L})$ $L^2$ -accurate score estimates. Then, the probability flow ODE implementation of the reversed Ornstein-Uhlenbeck process, when interspersed with either the overdamped Langevin corrector (DPOM; Algorithm 1) or with the underdamped Langevin corrector (DPUM; Algorithm 2), outputs a sample whose law is  $\varepsilon$ -close in total variation distance to the data distribution, using  $\widetilde{O}(L^3 d/\varepsilon^2)$  or  $\widetilde{O}(L^2\sqrt{d}/\varepsilon)$  iterations respectively.
Our result provides the first polynomial-time guarantees for the probability flow ODE implementation of SGMs, so long as it is combined with the use of corrector steps. Moreover, when the corrector steps are based on the underdamped Langevin diffusion, then the dimension dependence of our result is significantly smaller  $(O(\sqrt{d})$  vs.  $O(d))$  than prior works on the complexity of DDPMs, and thus provides justification for the use of ODE discretization in practice, compared to SDEs.
Our main assumption on the data is that the score functions along the forward process are Lipschitz continuous, which allows for highly non-log-concave distributions, yet does not cover non-smooth distributions such as distributions supported on lower-dimensional manifolds. However, as shown in [Che+23a; CLL23; LLT23], we can also obtain polynomial-time guarantees without this smoothness assumption via early stopping (see Remark 1).
The idea of using a time-reversed diffusion for sampling has been fruitfully exploited in the log-concave sampling literature via the proximal sampler [TP18; LST21; CE22; LC22; FYC23; LC23], as put forth in [Che+22], as well as through algorithmic stochastic localization [EMS22; MW23]. Although we do not aim to be comprehensive in our discussion of the literature, we mention, e.g., [ABV23; Che+23b] for alternative approaches for diffusion models. We also note that the recent work of [CDD23] obtained a discretization analysis for the probability flow ODE (without corrector) in KL divergence, though their bounds have a large dependence on  $d$  and are exponential in the Lipschitz constant of the score integrated over time.
Since the original arXiv submission of this paper, there have been further works studying the probability flow ODE. The work of [BDD23] also studied the probability flow ODE, but without providing discretization guarantees (and with possibly exponential dependencies). The work [Li+23] provides polynomial-time guarantees for the probability flow ODE (without corrector steps), at the cost of larger polynomial dependencies and more stringent score assumptions (namely, bounds on the Jacobian of the score). Also, [PMM23] study another variant of the predictor-corrector framework.
2.1 Score-based generative modeling
Let  $q_{\star}$  denote the data distribution, i.e., the distribution from which we wish to sample. In score-based generative modeling, we define a forward process  $(q_t^{\rightarrow})_{t\geq 0}$  with  $q_0^\rightarrow = q_\star$ , which transforms our data distribution into noise. In this paper, we focus on the canonical choice of the Ornstein-Uhlenbeck
(OU) process,
$$
\mathrm {d} x _ {t} ^ {\rightarrow} = - x _ {t} ^ {\rightarrow} \mathrm {d} t + \sqrt {2} \mathrm {d} B _ {t}, \quad x _ {0} ^ {\rightarrow} \sim q _ {\star}, \quad q _ {t} ^ {\rightarrow} := \operatorname {l a w} \left(x _ {t} ^ {\rightarrow}\right), \tag {1}
$$
where  $(B_{t})_{t\geq 0}$  is a standard Brownian motion in  $\mathbb{R}^d$ . It is well-known that the OU process mixes rapidly (exponentially fast) to its stationary distribution, the standard Gaussian distribution  $\gamma^d$ .
Once we fix a time horizon  $T > 0$ , the time reversal of the SDE defined in (1) over  $[0, T]$  is given by
$$
\mathrm {d} x _ {t} ^ {\leftarrow} = \left(x _ {t} ^ {\leftarrow} + 2 \nabla \ln q _ {t} ^ {\leftarrow} \left(x _ {t} ^ {\leftarrow}\right)\right) \mathrm {d} t + \sqrt {2} \mathrm {d} B _ {t}, \tag {2}
$$
where  $q_{t}^{\leftarrow} \coloneqq q_{T - t}^{\rightarrow}$ , and the reverse SDE is a generative model: when initialized at  $x_0^{\leftarrow} \sim q_0^{\leftarrow}$ , then  $x_{T}^{\leftarrow} \sim q$ . Since  $q_0^{\leftarrow} = q_T^{\rightarrow} \approx \gamma^d$ , the reverse SDE transforms samples from  $\gamma^{d}$  (i.e., pure noise) into approximate samples from  $q_{\star}$ . In order to implement the reverse SDE, however, one needs to estimate the score functions  $\nabla \ln q_t^{\leftarrow}$  for  $t \in [0,T]$  using the technique of score matching [Hyv05; Vin11]. In practice, the score estimates are produced via a deep neural network, and our main assumption is that these score estimates are accurate in an  $L^2$  sense (see Assumption 4). This gives rise to the denoising diffusion probabilistic modeling (DDPM) algorithm.
Notation. Since the reverse process is the primary object of interest, we drop the arrow  $\leftarrow$  from the notation for simplicity; thus,  $q_{t} := q_{t}^{\leftarrow}$ . We will always denote the forward process with the arrow  $\rightarrow$ . For each  $t \in [0,T]$ , let  $s_t$  denote the estimate for the score  $\nabla \ln q_t = \nabla \ln q_t^{\leftarrow}$ .
2.2 Probability flow ODE (predictor steps)
Instead of running the reverse SDE (2), there is in fact an alternative process  $(x_{t})_{t\in [0,T]}$  which evolves according to an ODE (and hence evolves deterministically), and yet has the same marginals as (2). This alternative process, called the probability flow ODE, can also be used for generative modeling.
One particularly illuminating way of deriving the probability flow ODE is to invoke the celebrated theorem, due to [JKO98], that the OU process is the Wasserstein gradient flow of the KL divergence functional (i.e. relative entropy)  $\mathsf{KL}(\cdot \| \gamma^d)$ . From the general theory of Wasserstein gradient flows (see [AGS08; San15]), the Wasserstein gradient flow  $(\mu_t)_{t\geq 0}$  of a functional  $\mathcal{F}$  can be implemented via the dynamics
$$
\dot {z} _ {t} = - \left[ \nabla_ {W _ {2}} \mathcal {F} (\mu_ {t}) \right] (z _ {t}), \qquad z _ {0} \sim \mu_ {0},
$$
in that  $z_{t} \sim \mu_{t}$  for all  $t \geq 0$ . Applying this to  $\mathcal{F} \coloneqq \mathrm{KL}(\cdot \| \gamma^{d})$ , we arrive at the forward process
$$
\dot {x} _ {t} ^ {\rightarrow} = - \nabla \ln \left(\frac {q _ {t} ^ {\rightarrow}}{\gamma^ {d}}\right)\left(x _ {t} ^ {\rightarrow}\right) = - x _ {t} ^ {\rightarrow} - \nabla \ln q _ {t} ^ {\rightarrow} \left(x _ {t} ^ {\rightarrow}\right). \tag {3}
$$
Setting  $x_{t} \coloneqq x_{T - t}^{\rightarrow}$ , it is easily seen that the time reversal of (3) is
$$
\dot {x} _ {t} = x _ {t} + \nabla \ln q _ {t} \left(x _ {t}\right), \quad i. e., \quad \dot {x} _ {t} = x _ {t} + \nabla \ln q _ {T - t} ^ {\rightarrow} \left(x _ {t}\right), \tag {4}
$$
which is called the probability flow ODE. In this paper, the interpretation of the probability flow ODE as a reverse Wasserstein gradient flow is only introduced for interpretability, and the reader who is unfamiliar with Wasserstein calculus can take (4) to be the definition of the probability flow ODE. Crucially, it has the property that if  $x_0 \sim q_0$ , then  $x_t \sim q_t$  for all  $t \in [0, T]$ .
We can discretize the ODE (4). Fixing a step size  $h > 0$ , replacing the score function  $\nabla \ln q_t$  with the estimated score given by  $s_t$ , and applying the exponential integrator to the ODE (i.e., exactly integrating the linear part), we arrive at the discretized process
$$
x _ {t + h} = x _ {t} + \int_ {0} ^ {h} x _ {t + u} \mathrm {d} u + h s _ {t} \left(x _ {t}\right) = \exp (h) x _ {t} + \left(\exp (h) - 1\right) s _ {t} \left(x _ {t}\right). \tag {5}
$$
2.3 Corrector steps
Let  $q$  be a distribution over  $\mathbb{R}^d$ , and write  $U$  as a shorthand for the potential  $-\ln q$ .
Overdamped Langevin. The overdamped Langevin diffusion with potential  $U$  is a stochastic process  $(x_{t})_{t\geq 0}$  over  $\mathbb{R}^d$  given by
$$
\mathrm {d} x _ {t} = - \nabla U (x _ {t}) \mathrm {d} t + \sqrt {2} \mathrm {d} B _ {t}.
$$
The stationary distribution of this diffusion is  $q \propto \exp(-U)$ .
We also consider the following discretized process where  $-\nabla U$  is replaced by a score estimate  $s$ . Fix a step size  $h > 0$  and let  $(\widehat{x}_t)_{t\geq 0}$  over  $\mathbb{R}^d$  be given by
$$
\mathrm {d} \widehat {x} _ {t} = s \left(\widehat {x} _ {\lfloor t / h \rfloor h}\right) \mathrm {d} t + \sqrt {2} \mathrm {d} B _ {t}.
$$
Underdamped Langevin. Given a friction parameter  $\gamma > 0$ , the corresponding underdamped Langevin diffusion is a stochastic process  $(z_{t}, v_{t})_{t \geq 0}$  over  $\mathbb{R}^d \times \mathbb{R}^d$  given by
$$
\mathrm {d} z _ {t} = v _ {t} \mathrm {d} t,
$$
$$
\mathrm {d} v _ {t} = - \left(\nabla U \left(z _ {t}\right) + \gamma v _ {t}\right) \mathrm {d} t + \sqrt {2 \gamma} \mathrm {d} B _ {t}.
$$
The stationary distribution of this diffusion is  $q\otimes \gamma^d$
We also consider the following discretized process, where  $-\nabla U$  is replaced by a score estimate  $s$ . Let  $(\widehat{z}_t,\widehat{v}_t)_{t\geq 0}$  over  $\mathbb{R}^d\times \mathbb{R}^d$  be given by
$$
\mathrm {d} \widehat {z} _ {t} = \widehat {v} _ {t} \mathrm {d} t,
$$
$$
\mathrm {d} \widehat {v} _ {t} = \left(s \left(\widehat {z} _ {\lfloor t / h \rfloor h}\right) - \gamma \widehat {v} _ {t}\right) \mathrm {d} t + \sqrt {2 \gamma} \mathrm {d} B _ {t}. \tag {6}
$$
Diffusions as corrector steps. At time  $t$ , the law of the ideal reverse process (4) initialized at  $q_{0}$  is  $q_{t}$ . However, errors are accumulated through the course of the algorithm: the error from initializing at  $\gamma^{d}$  rather than at  $q_{0}$ ; errors arising from discretization of (4); and errors in estimating the score function. That's why the law of the algorithm's iterate will not be exactly  $q_{t}$ . We propose to use either the overdamped or the underdamped Langevin diffusion with stationary distribution  $q_{t}$  and estimated score as a corrector, in order to bring the law of the algorithm iterate closer to  $q_{t}$ . In the case of the underdamped Langevin diffusion, this is done by drawing an independent Gaussian random variable  $\widehat{v}_{0} \sim \gamma^{d}$ , running the system (6) starting from  $(\widehat{z}_{0},\widehat{v}_{0})$  (where  $\widehat{z}_0$  is the current algorithm iterate) for some time  $t$ , and then keeping  $\widehat{z}_t$ . In our theoretical analysis, the use of corrector steps boosts the accuracy and efficiency of the SGM.
We make the following mild assumptions on the data distribution  $q_{\star}$  and on the score estimate  $s$
Assumption 1 (second moment bound). We assume that  $\mathfrak{m}_2^2 \coloneqq \mathbb{E}_{q_*}[\| \cdot \| ^2 ] < \infty$ .
Assumption 2 (Lipschitz score). For all  $t \in [0, T]$ , the score  $\nabla \ln q_t$  is  $L$ -Lipschitz, for some  $L \geq 1$ .
Assumption 3 (Lipschitz score estimate). For all  $t$  for which we need to estimate the score function in our algorithms, the score estimate  $s_t$  is  $L$ -Lipschitz.
Assumption 4 (score estimation error). For all  $t$  for which we need to estimate the score function in our algorithms,
$$
\mathbb {E} _ {q _ {t}} \left[ \left\| s _ {t} - \nabla \ln q _ {t} \right\| ^ {2} \right] \leq \varepsilon_ {\mathfrak {s c}} ^ {2}.
$$
Assumptions 1, 2, and 4 are standard and were shown in [Che+23a; CLL23; LLT23] to suffice for obtaining polynomial-time convergence guarantees for DDPMs. The new condition that we require in our analysis is Assumption 3, which was used in [LLT22] but ultimately shown to be unnecessary for DDPMs. We leave it as an open question whether this can be lifted in the ODE setting.
Remark 1. As observed in [Che+23a; CLL23; LLT23], Assumption 2 can be removed via early stopping, at the cost of polynomially larger iteration complexity. The idea is that if  $q_{\star}$  has compact support but does not necessarily satisfy Assumption 2 (e.g., if  $q_{\star}$  is supported on a compact and lower-dimensional manifold), then  $q_{\delta}^{\rightarrow}$  will satisfy Assumption 2 if  $\delta > 0$ . By applying our analysis up to time  $T - \delta$  instead of time  $T$ , one can show that a suitable projection of the output distribution is close in Wasserstein distance to  $q_{\star}$  (see [CLL23, Corollary 2.4] or [Che+23a, Corollary 5]). For brevity, we do not consider this extension of our results here.
3.2 Algorithms
We provide the pseudocode for the two algorithms we consider, Diffusion Predictor + Overdamped Modeling (DPOM) and Diffusion Predictor + Underdamped Modeling (DPUM), in Algorithms 1 and 2 respectively. The only difference between the two algorithms is in the corrector step, which we highlight in Algorithm 2. For simplicity, we take the total amount of time  $T$  to be equal to  $N_0 / L + h_{\mathrm{pred}}$  for an integer  $N_0 \geq 1$ , and we assume that  $1 / L$  is a multiple of  $h_{\mathrm{pred}}$  and that  $h_{\mathrm{pred}}$  is a multiple of  $\delta = \Theta \left( \frac{\varepsilon^2}{L^2(d \vee m_2^2)} \right)$ .
We consider two stages: in the first stage, which lasts until time  $N_0 / L = T - h_{\mathrm{pred}}$ , we intersperse predictor epochs (run for time  $1 / L$ , discretized with step size  $h_{\mathrm{pred}}$ ) and corrector epochs (run for time  $\Theta(1 / L)$  for the overdamped corrector or for time  $\Theta(1 / \sqrt{L})$  for the underdamped corrector, and discretized with step size  $h_{\mathrm{corr}}$ ). The second stage lasts from time  $T - h_{\mathrm{pred}}$  to time  $T - \delta$ , and we incorporate geometrically decreasing step sizes for the predictor. Note that this implies that our algorithm uses early stopping.
Algorithm 1: DPOM(T,  $h_{\text{pred}}$ ,  $h_{\text{corr}}$ , s)
Input: Total time  $T$ , predictor step size  $h_{\mathrm{pred}}$ , corrector step size  $h_{\mathrm{corr}}$ , score estimates  $s$   
Output: Approximate sample from the data distribution  $q_{\star}$
Algorithm 2:  $\text{DPUM}(T, h_{\text{pred}}, h_{\text{corr}}, s)$
Input: Total time  $T$ , predictor step size  $h_{\mathrm{pred}}$ , corrector step size  $h_{\mathrm{corr}}$ , score estimates  $s$   
Output: Approximate sample from the data distribution  $q_{\star}$
3 Predictor: Starting from  $\widehat{x}_{n / L}$ , run the discretized probability flow ODE (5) from time  $\frac{n}{L}$  to  $\frac{n + 1}{L}$  with step size  $h_{\mathrm{pred}}$  and estimated scores to obtain  $\widehat{x}^{\prime}_{(n + 1) / L}$ .
4 Corrector: Starting from  $\widehat{x}_{(n+1)/L}^{\prime}$ , run underdamped Langevin Monte Carlo for total time  $\Theta(1/\sqrt{L})$  with step size  $h_{\mathrm{corr}}$  and score estimate  $s_{(n+1)/L}$  to obtain  $\widehat{x}_{(n+1)/L}$ .
5 Predictor: Starting from  $\widehat{x}_{T - h_{\text{pred}}}$ , run the discretized probability flow ODE (5) with step sizes  $h_{\text{pred}} / 2, h_{\text{pred}} / 4, h_{\text{pred}} / 8, \ldots, \delta$  and estimated scores to obtain  $\widehat{x}_{T - \delta}^{\prime}$ .