Title: The probability flow ODE is provably fast
Venue: NeurIPS 2023
Accepted: yes
Keywords: DDIM

deterministic samplers

diffusion models

predictor-corrector

probability flow ODE

score-based generative modeling
Abstract: We provide the first polynomial-time convergence guarantees for the probabilistic flow ODE implementation (together with a corrector step) of score-based generative modeling. 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.

The probability flow ODE is provably fast

Abstract

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 ( √ d ) vs. O ( d ) , assuming smoothness of the data distribution), highlighting potential advantages of the ODE framework.

1 Introduction

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. 7 Subsequently, the work of [CLL23] significantly sharpened the analysis in the case when no smoothness assumptions are imposed on the data distribution.

‡

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 C 1 for ODEs but only C 1 2 -for SDEs ( i.e. , Hölder continuous with any exponent less than 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 / √ 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 ( √ 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.

1.1 Our contributions

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 ( √ d ) vs O ( d ) ) and (ii) the complexity of DDPMs, i.e. , SDE implementations of score-based generative models (again, O ( √ 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 ˜ O ( ε/ √ L ) L 2 -accurate score estimates. Then, the probability flow ODE implementation of the reversed OrnsteinUhlenbeck 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 ε -close in total variation distance to the data distribution, using ˜ O ( L 3 d/ε 2 ) or ˜ O ( L 2 √ d/ε ) 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 ( √ 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).

1.2 Related works

The idea of using a time-reversed diffusion for sampling has been fruitfully exploited in the logconcave 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 Preliminaries

2.1 Score-based generative modeling

Let q ⋆ 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 ) t ≥ 0 with q → 0 = q ⋆ , which transforms our data distribution into noise. In this paper, we focus on the canonical choice of the Ornstein-Uhlenbeck

(OU) process,

where ( B t ) t ≥ 0 is a standard Brownian motion in R d . It is well-known that the OU process mixes rapidly (exponentially fast) to its stationary distribution, the standard Gaussian distribution γ d .

Once we fix a time horizon T > 0 , the time reversal of the SDE defined in (1) over [0 , T ] is given by

where q ← t := q → T -t , and the reverse SDE is a generative model: when initialized at x ← 0 ∼ q ← 0 , then x ← T ∼ q . Since q ← 0 = q → T ≈ γ d , the reverse SDE transforms samples from γ d (i.e., pure noise) into approximate samples from q ⋆ . In order to implement the reverse SDE, however, one needs to estimate the score functions ∇ ln q ← t for t ∈ [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 ← from the notation for simplicity; thus, q t := q ← t . We will always denote the forward process with the arrow → .

For each t ∈ [0 , T ] , let s t denote the estimate for the score ∇ ln q t = ∇ ln q ← t .

2.2 Probability flow ODE (predictor steps)

Instead of running the reverse SDE (2), there is in fact an alternative process ( x t ) t ∈ [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) KL ( · ∥ γ d ) . From the general theory of Wasserstein gradient flows (see [AGS08; San15]), the Wasserstein gradient flow ( µ t ) t ≥ 0 of a functional F can be implemented via the dynamics

in that z t ∼ µ t for all t ≥ 0 . Applying this to F := KL ( · ∥ γ d ) , we arrive at the forward process

Setting x t := x → T -t , it is easily seen that the time reversal of (3) is

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 ∼ q 0 , then x t ∼ q t for all t ∈ [0 , T ] .

We can discretize the ODE (4). Fixing a step size h > 0 , replacing the score function ∇ 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

2.3 Corrector steps

Let q be a distribution over 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 ≥ 0 over R d given by

The stationary distribution of this diffusion is q ∝ exp( -U ) .

We also consider the following discretized process where -∇ U is replaced by a score estimate s . Fix a step size h > 0 and let ( ̂ x t ) t ≥ 0 over R d be given by

Underdamped Langevin. Given a friction parameter γ > 0 , the corresponding underdamped Langevin diffusion is a stochastic process ( z t , v t ) t ≥ 0 over R d × R d given by

The stationary distribution of this diffusion is q ⊗ γ d .

We also consider the following discretized process, where -∇ U is replaced by a score estimate s . Let ( ̂ z t , ̂ v t ) t ≥ 0 over R d × R d be given by

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 γ 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 ̂ v 0 ∼ γ d , running the system (6) starting from ( ̂ z 0 , ̂ v 0 ) (where ̂ z 0 is the current algorithm iterate) for some time t , and then keeping ̂ z t . In our theoretical analysis, the use of corrector steps boosts the accuracy and efficiency of the SGM.

3 Results

3.1 Assumptions

We make the following mild assumptions on the data distribution q ⋆ and on the score estimate s .

Assumption 1 (second moment bound) . We assume that m 2 2 := E q ⋆ [ ∥·∥ 2 ] < ∞ .

Assumption 2 (Lipschitz score) . For all t ∈ [0 , T ] , the score ∇ ln q t is L -Lipschitz, for some L ≥ 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,

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 ⋆ has compact support but does not necessarily satisfy Assumption 2 ( e.g. , if q ⋆ is supported on a compact and lower-dimensional manifold), then q → δ will satisfy Assumption 2 if δ > 0 . By applying our analysis up to time T -δ instead of time T , one can show that a suitable projection of the output distribution is close in Wasserstein distance to q ⋆ (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 pred for an integer N 0 ≥ 1 , and we assume that 1 /L is a multiple of h pred and that h pred is a multiple of δ = Θ ( ε 2 L 2 ( d ∨ m 2 2 ) ) .

We consider two stages: in the first stage, which lasts until time N 0 /L = T -h pred , we intersperse predictor epochs (run for time 1 /L , discretized with step size h pred ) and corrector epochs (run for time Θ(1 /L ) for the overdamped corrector or for time Θ(1 / √ L ) for the underdamped corrector, and discretized with step size h corr ). The second stage lasts from time T -h pred to time T -δ , 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 pred , h corr , s ) Input: Total time T , predictor step size h pred , corrector step size h corr , score estimates s Output: Approximate sample from the data distribution q ⋆ 1 Draw ̂ x 0 ∼ γ d . 2 for n = 0 , 1 , . . . , N 0 -1 do 3 Predictor: Starting from ̂ x n/L , run the discretized probability flow ODE (5) from time n L to n +1 L with step size h pred and estimated scores to obtain ̂ x ′ ( n +1) /L . 4 Corrector: Starting from ̂ x ′ ( n +1) /L , run overdamped Langevin Monte Carlo for total time Θ(1 /L ) with step size h corr and score estimate s ( n +1) /L to obtain ̂ x ( n +1) /L . 5 Predictor: Starting from ̂ x T -h pred , run the discretized probability flow ODE (5) with step sizes h pred / 2 , h pred / 4 , h pred / 8 , . . . , δ and estimated scores to obtain ̂ x ′ T -δ . 6 Corrector: Starting from ̂ x ′ T -δ , run overdamped Langevin Monte Carlo for total time Θ(1 /L ) with step size h corr and score estimate s T -δ to obtain ̂ x T -δ . 7 return ̂ x T -δ

Algorithm 2: DPUM ( T, h pred , h corr , s )

Input: Total time T , predictor step size h pred , corrector step size h corr , score estimates s Output: Approximate sample from the data distribution q ⋆ 1 Draw ̂ x 0 ∼ γ d . 2 for n = 0 , 1 , . . . , N 0 -1 do 3 Predictor: Starting from ̂ x n/L , run the discretized probability flow ODE (5) from time n L to n +1 L with step size h pred and estimated scores to obtain ̂ x ′ ( n +1) /L . 4 Corrector: Starting from ̂ x ′ ( n +1) /L , run underdamped Langevin Monte Carlo for total time Θ(1 / √ L ) with step size h corr and score estimate s ( n +1) /L to obtain ̂ x ( n +1) /L . 5 Predictor: Starting from ̂ x T -h pred , run the discretized probability flow ODE (5) with step sizes h pred / 2 , h pred / 4 , h pred / 8 , . . . , δ and estimated scores to obtain ̂ x ′ T -δ . 6 Corrector: Starting from ̂ x ′ T -δ , run underdamped Langevin Monte Carlo for total time Θ(1 / √ L ) with step size h corr and score estimate s T -δ to obtain ̂ x T -δ . 7 return ̂ x T -δ

3.3 Convergence guarantees

Our main results are the following convergence guarantees for the two predictor-corrector schemes described in §3.2:

Theorem 2 ( DPOM ) . Suppose that Assumptions 1-4 hold. If ̂ q denotes the output of DPOM (Algorithm 1) with δ ≍ ε 2 L 2 ( d ∨ m 2 2 ) , then

̂ corr In particular, if we set T = Θ ( ln( d ∨ m 2 2 ε 2 ) ) , h pred = ˜ Θ( ε L 2 d 1 / 2 ) , h corr = ˜ Θ( ε 2 L 3 d ) , and if the score estimation error satisfies ε sc ≤ ˜ O ( ε √ L ) , then we can obtain TV error ε with a total iteration complexity of Θ( L 3 d 2 ) steps.

˜ ε

The five terms in the bound (7) correspond, respectively, to: the convergence of the forward (OU) process; the discretization error from the predictor steps; the discretization error from the corrector steps; the score estimation error; and the early stopping error.

Theorem 2 recovers nearly the same guarantees as the one in [Che+23a; CLL23; LLT23], but for the probability flow ODE with overdamped Langevin corrector instead of the reverse SDE without corrector. Recall also from Remark 1 that our results can easily be extended to compactly supported data distributions without smooth score functions. This covers essentially all distributions encountered in practice. Therefore, our result provides compelling theoretical justification complementing the empirical efficacy of the probability flow ODE, which was hitherto absent from the literature.

However, in Theorem 2, the iteration complexity is dominated by the corrector steps. Next, we show that by replacing the overdamped LMC with underdamped LMC, we can achieve a quadratic improvement in the number of steps, considering the dependence on d . As discussed in the Introduction, this highlights the potential benefits of the ODE framework over the SDE.

Theorem 3 ( DPUM ) . Suppose that Assumptions 1-4 hold. If ̂ q denotes the output of DPUM (Algorithm 2) with δ ≍ ε 2 L 2 ( d ∨ m 2 2 ) , then

In particular, if we set T = Θ ( ln( d ∨ m 2 2 ε 2 ) ) , h pred = ˜ Θ( ε L 2 d 1 / 2 ) , h corr = ˜ Θ( ε L 3 / 2 d 1 / 2 ) , and if the score estimation error satisfies ε sc ≤ ˜ O ( ε √ L ) , then we can obtain TV error ε with a total iteration complexity of ˜ Θ( L 2 d 1 / 2 ε ) steps.

4 Proof overview

Here we give a detailed technical overview for the proof of our main results, Theorems 2 and 3. As in [Che+23a; CLL23; LLT23], the three sources of error that we need to keep track of are (1) estimation of the score function; (2) discretization of time when implementing the probability flow ODE and corrector steps; and (3) initialization of the algorithm at γ d instead of the true law of the end of the forward process, q 0 = q → T . It turns out that (1) is not so difficult to manage as soon as we can control (2) and (3). Furthermore, as in prior work, we can easily control (3) via the data-processing inequality: the total variation distance between the output of the algorithm initialized at q 0 versus at γ d is at most TV ( q → T , γ d ) , which is exponentially small in T by rapid mixing of the OU process. So henceforth in this overview, let us assume that both the algorithm and the true process are initialized at q 0 . It remains to control (2).

Failure of existing approaches. In the SDE implementation of diffusion models, prior works handled (2) by directly bounding a strictly larger quantity, namely the KL divergence between the laws of the trajectories of the algorithm and the true process; by Girsanov's theorem, this has a clean formulation as an integrated difference of drifts. Unfortunately, in the ODE implementation, this KL divergence is infinite: in the absence of stochasticity in the reverse process, these laws over trajectories are not even absolutely continuous with respect to each other.

In search of an alternative approach, one might try a Wasserstein analysis. As a first attempt, we could couple the initialization of both processes and look at how the distance between them changes over time. If ( ̂ x t ) 0 ≤ t ≤ T and ( x t ) 0 ≤ t ≤ T denote the algorithm and true process, then smoothness of the score function allows us to naïvely bound ∂ t E [ ∥ ̂ x t -x t ∥ 2 ] by O ( L ) E [ ∥ ̂ x t -x t ∥ 2 ] . While this ensures that the processes are close if run for time ≪ 1 /L , it does not rule out the possibility that they drift apart exponentially quickly after time 1 /L .

Restarting the coupling-first attempt. What we would like is some way of 'restarting' this coupling before the processes drift too far apart, to avoid this exponential compounding. We now motivate how to achieve this by giving an argument that is incorrect but nevertheless captures the intuition for our approach. Namely, let p t := law( ̂ x t ) denote the law of the algorithm, let P t 0 ,h ODE denote the result of running the ideal probability flow ODE for time h starting from time t 0 , and let ̂ P t 0 ,h ODE denote the same but for the discretized probability flow ODE with estimated score. For h ≲ 1 /L , consider the law of the two processes at time 2 h , i.e.,

The discussion above implies that q 0 P 0 ,h ODE and q 0 ̂ P 0 ,h ODE are close in 2-Wasserstein distance, so by the data-processing inequality, this implies that q 0 P 0 ,h ODE ̂ P h,h ODE and q 0 ̂ P 0 ,h ODE ̂ P h,h ODE are also close. To show that p 2 h and q 2 h in Eq. (8) are close, it thus suffices to show that q 0 P 0 , 2 h ODE and q 0 P 0 ,h ODE ̂ P h,h ODE are close. But these two distributions are given by running the algorithm and the true process for time h , both starting from q 0 P 0 ,h ODE . So if we 'restart' the coupling by coupling the processes based on their locations at time h , rather than time 0 , of the reverse process, we can again apply the naïve Wasserstein analysis.

At this juncture, it would seem that we have miraculously sidestepped the exponential blowup and shown that the expected distance between the processes only increases linearly over time! The issue of course is in the application of the 'data-processing inequality,' which simply does not hold for the Wasserstein distance.

Restarting the coupling with a corrector step. This is where the corrector comes in. The idea is to use short-time regularization : if we apply a small amount of noise to two distributions which are already close in Wasserstein, then they become close in KL divergence, for which a data-processing inequality holds. The upshot is that if the noise doesn't change the distributions too much, then we can legitimately restart the coupling as above and prove that the distance between the processes, now defined by interleaving the probability flow ODE and its discretization with periodic injections of noise, increases only linearly in time.

It turns out that naïve injection of noise, e.g., convolution with a Gaussian of small variance, is somewhat wasteful as it fails to preserve the true process and leads to poor polynomial dependence in the dimension. On the other hand, if we instead run the overdamped Langevin diffusion with potential chosen so that the law of the true process is stationary, then we can recover the linear in d dependence of Theorem 2. Then by replacing overdamped Langevin diffusion with its underdamped counterpart, which has the advantage of much smoother trajectories, we can obtain the desired quadratic speedup in dimension dependence in Theorem 3.

Score perturbation lemma. In addition to the switch from SDE to ODE and the use of the underdamped corrector, a third ingredient is essential to our improved dimension dependence. The former two ensure that the trajectory of our algorithm is smoother than that of DDPMs, so that even over time windows that scale with 1 / √ d , the process does not change too much. By extension, as the score functions are Lipschitz, this means that any fixed score function evaluated over iterates in such a window does not change much. This amounts to controlling discretization error in space .

It is also necessary to control discretization error in time , i.e., proving what some prior works referred to as a score perturbation lemma [LLT22]. That is, for any fixed iterate x , we want to show that the score function ∇ ln q t ( x ) does not change too much as t varies over a small window. Unfortunately, prior works were only able to establish this over windows of length 1 /d . In this work, we improve this to windows of length 1 / √ d (see Lemma 3 and Corollary 1).

In our proof, we bound the squared L 2 norm of the derivative of the score along the trajectory of the ODE. The score function evaluated at y can be expressed as E P 0 | t ( ·| y ) [ ∇ U ] ; here, the posterior distribution P 0 | t ( · | y ) is essentially the prior q ⋆ tilted by a Gaussian of variance O ( t ) . Hence we need to bound the change in the expectation when we change the distribution from P 0 | t to P 0 | t +∆ t ; because ∇ U is L -Lipschitz, we can bound this by the Wasserstein distance between the distributions. For small enough t , P 0 | t is strongly log-concave, and a transport cost inequality bounds this in terms of KL divergence, which is more easily bounded. Indeed, we can bound it with the KL divergence between the joint distributions P 0 ,t and P 0 ,t +∆ t , which reduces to bounding the KL divergence between Gaussians of unequal variance.

However, since our score perturbation lemma degrades near the beginning of the forward process, we require better control of the discretization error during this part of the algorithm, hence leading to our choice of geometrically decreasing step sizes. Alternatively, we could use a two-stage step size schedule, see Remark 4.

5 Conclusion

In this work, we have provided the first polynomial-time guarantees for the probability flow ODE implementation of SGMs with corrector steps and exhibited improved dimension dependence of the ODE framework over prior results for DDPMs ( i.e., the SDE framework). Our analysis raises questions relevant for practice, of which we list a few.

Although we need the corrector steps for our proof, are they in fact necessary for the algorithm to work efficiently in practice?

Is it possible to obtain even better dimension dependence, perhaps using higher-order solvers and stronger smoothness assumptions?

Can we obtain improved dimension dependence even in the non-smooth setting, compared to the result of [CLL23]?

We also list several limitations of our work, namely:

Our analysis only covers the probability flow ODE corresponding to the OU forward process. We leave the study of more general dynamics for future study.

Our guarantees require the score function to be learned to L 2 accuracy ˜ O ( ε/ √ L ) , which is more stringent than the prior works [Che+23a; CLL23; LLT23] and may be an technical artefact of our proof.

We have not validated our theoretical findings with large-scale experiments. In particular, it is still unclear whether flow-based methods can outperform the standard DDPM algorithm in practical, high dimensional settings.

6 Acknowledgments

We want to thank Yin Tat Lee for his valuable comments, shaping the direction of this research project in its early stages. SC was supported by NSF Award 2103300 for part of this work.

References

[CDD23] S. Chen, G. Daras, and A. Dimakis. 'Restoration-degradation beyond linear diffusions: a non-asymptotic analysis for DDIM-type samplers'. In: Proceedings of the 40th International Conference on Machine Learning . Ed. by A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett. Vol. 202. Proceedings of Machine Learning Research. PMLR, July 2023, pp. 4462-4484.

[CE22] Y. Chen and R. Eldan. 'Localization schemes: a framework for proving mixing bounds for Markov chains'. In: arXiv preprint 2203.04163 (2022).

[Che+22] Y. Chen, S. Chewi, A. Salim, and A. Wibisono. 'Improved analysis for a proximal algorithm for sampling'. In: Proceedings of Thirty Fifth Conference on Learning Theory . Ed. by P.-L. Loh and M. Raginsky. Vol. 178. Proceedings of Machine Learning Research. PMLR, July 2022, pp. 2984-3014.

[Che+23a] S. Chen, S. Chewi, J. Li, Y. Li, A. Salim, and A. R. Zhang. 'Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions'. In: The Eleventh International Conference on Learning Representations . 2023.

[Che+23b] Y. Chen, W. Deng, S. Fang, F. Li, N. T. Yang, Y. Zhang, K. Rasul, S. Zhe, A. Schneider, and Y. Nevmyvaka. 'Provably convergent Schrödinger bridge with applications to probabilistic time series imputation'. In: Proceedings of the 40th International Conference on Machine Learning . Ed. by A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett. V ol. 202. Proceedings of Machine Learning Research. PMLR, July 2023, pp. 4485-4513.

[Che22] S. Chewi. Log-concave sampling . Book draft available at https://chewisinho. github.io/ . 2022.

[CLL23] H. Chen, H. Lee, and J. Lu. 'Improved analysis of score-based generative modeling: user-friendly bounds under minimal smoothness assumptions'. In: Proceedings of the 40th International Conference on Machine Learning . Ed. by A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett. Vol. 202. Proceedings of Machine Learning Research. PMLR, July 2023, pp. 4735-4763.

[CSY22] H. Chung, B. Sim, and J. Ye. 'Come-closer-diffuse-faster: accelerating conditional diffusion models for inverse problems through stochastic contraction'. In: 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) . Los Alamitos, CA, USA: IEEE Computer Society, June 2022, pp. 12403-12412.

[De +21] V. De Bortoli, J. Thornton, J. Heng, and A. Doucet. 'Diffusion Schrödinger bridge with applications to score-based generative modeling'. In: Advances in Neural Information Processing Systems . Ed. by M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan. Vol. 34. Curran Associates, Inc., 2021, pp. 17695-17709.

[De 22] V. De Bortoli. 'Convergence of denoising diffusion models under the manifold hypothesis'. In: Transactions on Machine Learning Research (2022).

[DN21] P. Dhariwal and A. Nichol. 'Diffusion models beat GANs on image synthesis'. In: Advances in Neural Information Processing Systems . Ed. by M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan. Vol. 34. Curran Associates, Inc., 2021, pp. 8780-8794.

[EMS22] A. El Alaoui, A. Montanari, and M. Sellke. 'Sampling from the Sherrington-Kirkpatrick Gibbs measure via algorithmic stochastic localization'. In: 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science-FOCS 2022 . IEEE Computer Soc., Los Alamitos, CA, 2022, pp. 323-334.

[FYC23] J. Fan, B. Yuan, and Y. Chen. 'Improved dimension dependence of a proximal algorithm for sampling'. In: Proceedings of Thirty Sixth Conference on Learning Theory . Ed. by G. Neu and L. Rosasco. Vol. 195. Proceedings of Machine Learning Research. PMLR, July 2023, pp. 1473-1521.

[Gna+22] D. Gnaneshwar, B. Ramsundar, D. Gandhi, R. Kurchin, and V. Viswanathan. 'Scorebased generative models for molecule generation'. In: arXiv preprint 2203.04698 (2022).

[GW12] A. Guillin and F.-Y. Wang. 'Degenerate Fokker-Planck equations: Bismut formula, gradient estimate and Harnack inequality'. In: Journal of Differential Equations 253.1 (2012), pp. 20-40.

[HJA20] J. Ho, A. Jain, and P. Abbeel. 'Denoising diffusion probabilistic models'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 6840-6851.

[Hyv05] A. Hyvärinen. 'Estimation of non-normalized statistical models by score matching'. In: J. Mach. Learn. Res. 6 (2005), pp. 695-709.

[JKO98] R. Jordan, D. Kinderlehrer, and F. Otto. 'The variational formulation of the FokkerPlanck equation'. In: SIAM J. Math. Anal. 29.1 (1998), pp. 1-17.

[JP22] A. Jain and B. Poole. 'Journey to the BAOAB-limit: finding effective MCMC samplers for score-based models'. In: NeurIPS 2022 Workshop on Score-Based Methods . 2022.

[Kin+21] D. Kingma, T. Salimans, B. Poole, and J. Ho. 'Variational diffusion models'. In: Advances in Neural Information Processing Systems . Ed. by M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan. Vol. 34. Curran Associates, Inc., 2021, pp. 21696-21707.

[LC22] J. Liang and Y. Chen. 'A proximal algorithm for sampling from non-smooth potentials'. In: arXiv preprint 2110.04597 (2022).

[LC23] J. Liang and Y. Chen. 'A proximal algorithm for sampling'. In: Transactions on Machine Learning Research (2023).

[Li+23] G. Li, Y. Wei, Y. Chen, and Y. Chi. 'Towards faster non-asymptotic convergence for diffusion-based generative models'. In: arXiv preprint 2306.09251 (2023).

[Liu+22] X. Liu, L. Wu, M. Ye, and Q. Liu. 'Let us build bridges: understanding and extending diffusion generative models'. In: arXiv preprint arXiv:2208.14699 (2022).

[LLT22] H. Lee, J. Lu, and Y. Tan. 'Convergence for score-based generative modeling with polynomial complexity'. In: Advances in Neural Information Processing Systems . Ed. by A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho. 2022.

[LLT23] H. Lee, J. Lu, and Y. Tan. 'Convergence of score-based generative modeling for general data distributions'. In: Proceedings of the 34th International Conference on Algorithmic Learning Theory . Ed. by S. Agrawal and F. Orabona. Vol. 201. Proceedings of Machine Learning Research. PMLR, Feb. 2023, pp. 946-985.

[LST21] Y. T. Lee, R. Shen, and K. Tian. 'Structured logconcave sampling with a restricted Gaussian oracle'. In: Proceedings of Thirty Fourth Conference on Learning Theory . Ed. by M. Belkin and S. Kpotufe. Vol. 134. Proceedings of Machine Learning Research. PMLR, Aug. 2021, pp. 2993-3050.

[Lu+22] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu. 'DPM-Solver: a fast ODE solver for diffusion probabilistic model sampling in around 10 steps'. In: Advances in Neural Information Processing Systems . Ed. by S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh. Vol. 35. Curran Associates, Inc., 2022, pp. 5775-5787.

[MW23] A. Montanari and Y. Wu. 'Posterior sampling from the spiked models via diffusion processes'. In: arXiv preprint 2304.11449 (2023).

[Pid22] J. Pidstrigach. 'Score-based generative models detect manifolds'. In: Advances in Neural Information Processing Systems . Ed. by S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh. Vol. 35. Curran Associates, Inc., 2022, pp. 35852-35865.

[PMM23] F. Pedrotti, J. Maas, and M. Mondelli. 'Improved convergence of score-based diffusion models via prediction-correction'. In: arXiv preprint 2305.14164 (2023).

[Ram+22] A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen. 'Hierarchical text-conditional image generation with CLIP latents'. In: arXiv preprint arXiv:2204.06125 (2022).

[Rom+22] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. 'High-resolution image synthesis with latent diffusion models'. In: 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) . Los Alamitos, CA, USA: IEEE Computer Society, June 2022, pp. 10674-10685.

[San15] F. Santambrogio. Optimal transport for applied mathematicians . Vol. 87. Progress in Nonlinear Differential Equations and their Applications. Calculus of variations, PDEs, and modeling. Birkhäuser/Springer, Cham, 2015, pp. xxvii+353.

[SE19] Y. Song and S. Ermon. 'Generative modeling by estimating gradients of the data distribution'. In: Advances in Neural Information Processing Systems . Ed. by H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett. Vol. 32. Curran Associates, Inc., 2019.

[Shi+21] C. Shi, S. Luo, M. Xu, and J. Tang. 'Learning gradient fields for molecular conformation generation'. In: Proceedings of the 38th International Conference on Machine Learning . Ed. by M. Meila and T. Zhang. Vol. 139. Proceedings of Machine Learning Research. PMLR, July 2021, pp. 9558-9568.

[Soh+15] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. 'Deep unsupervised learning using nonequilibrium thermodynamics'. In: Proceedings of the 32nd International Conference on Machine Learning . Ed. by F. Bach and D. Blei. Vol. 37. Proceedings of Machine Learning Research. Lille, France: PMLR, July 2015, pp. 22562265.

[Son+21a] Y. Song, C. Durkan, I. Murray, and S. Ermon. 'Maximum likelihood training of scorebased diffusion models'. In: Advances in Neural Information Processing Systems . Ed. by M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan. Vol. 34. Curran Associates, Inc., 2021, pp. 1415-1428.

[Son+21b] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. 'Scorebased generative modeling through stochastic differential equations'. In: International Conference on Learning Representations . 2021.

[Son+22] Y. Song, L. Shen, L. Xing, and S. Ermon. 'Solving inverse problems in medical imaging with score-based generative models'. In: International Conference on Learning Representations . 2022.

[TP18] M. K. Titsias and O. Papaspiliopoulos. 'Auxiliary gradient-based sampling algorithms'. In: J. R. Stat. Soc. Ser. B. Stat. Methodol. 80.4 (2018), pp. 749-767.

[Vin11] P. Vincent. 'A connection between score matching and denoising autoencoders'. In: Neural Comput. 23.7 (2011), pp. 1661-1674.

[VKK21] A. Vahdat, K. Kreis, and J. Kautz. 'Score-based generative modeling in latent space'. In: Advances in Neural Information Processing Systems . Ed. by M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan. Vol. 34. Curran Associates, Inc., 2021, pp. 11287-11302.

[WHZ23] Z. Wang, J. J. Hunt, and M. Zhou. 'Diffusion policies as an expressive policy class for offline reinforcement learning'. In: The Eleventh International Conference on Learning Representations . 2023.

[WY22] A. Wibisono and K. Y. Yang. 'Convergence in KL divergence of the inexact Langevin algorithm with application to score-based generative models'. In: arXiv preprint 2211.01512 (2022).

[ZC23] Q. Zhang and Y. Chen. 'Fast sampling of diffusion models with exponential integrator'. In: The Eleventh International Conference on Learning Representations . 2023.

A Notation and overview

In this section, we collect together the notation used throughout the proofs and provide a road map for the end-to-end analysis in §E.

Throughout the analysis, Assumptions 1-4 are in full force.

We will reserve q for the law of the reverse process (and denote the forward process by q → when needed). In §E, the law of the algorithm is denoted by p .

We use the following Markov kernels:

P t,h ODE is the output of running the ODE for time h , starting at (reverse) time t .

P LD (resp. P ULD ) is the output of running the continuous-time overdamped (resp. underdamped) Langevin diffusion for time h . In this notation, we have suppressed mention of the stationary distributions of the diffusion, which will be provided by context.

̂ P t,h ODE and ̂ P LMC (resp. ̂ P ULMC ) are the corresponding processes once discretized and using the estimated score.

For the ODE, we are more precise with the notation because even within a single epoch of predictor steps, the kernel for the probability flow ODE depends on time (as opposed to the kernels for the diffusions, which are constant within any epoch of corrector steps); moreover, for our analysis in §E, we also need to take time-varying step sizes for the predictor steps. We will omit the dependencies on t and h when clear from context. When P = P ODE or ̂ P ODE , we use P t,h 1 ,...,h N to denote P t,h 1 P t + h 1 ,h 2 · · · P t + h 1 + ··· + h N -1 ,h N (we compose kernels on the right).

We refer to §4 for a high-level description of the proof strategy. We begin in §B with our improved score perturbation lemma (Corollary 1); this is the only section of the analysis which is indexed by forward time (instead of reverse time). In Lemma 5 in §C, we establish our main result for the predictor steps, which combines together standard ODE discretization analysis with the score perturbation lemma of §B. Since Corollary 1 degrades near the end of the reverse process (or equivalently, near the start of the forward process, when the regularization has not yet kicked in), our analysis requires a geometrically decreasing step size schedule, which leads to the two-stage Algorithms 1 and 2.

In §D, we prove our main regularization results for the overdamped corrector (Theorem 4) and the underdamped corrector (Theorem 5). Finally, we put together the various constituent results in the end-to-end analysis in §E.

B Score perturbation

In this section, we prove a score perturbation lemma which refines that of [LLT22]. This improved lemma is necessary in order to obtain O ( √ d ) dependence for the probability flow ODE.

Lemma 1 (Score perturbation) . Suppose p t = p 0 ∗ N (0 , tI ) and x 0 ∼ p 0 , ˙ x t = -1 2 ∇ ln p t ( x t ) . Suppose that ∥∇ 2 ln p ( t -1 2 L ) ∨ 0 ( x ) ∥ op ≤ L for all x . Then

Proof. Without loss of generality, we may assume t ≤ 1 2 L , as otherwise, noting that p t = p t -1 2 L ∗ N (0 , 1 2 L I ) , we may replace p 0 with p t -1 2 L and t with 1 2 L . Suppose p 0 ( x ) = e -V ( x ) . Let P 0 ,t denote the joint distribution of ( X 0 , X t ) where X t = X 0 + √ t Z with Z ∼ N (0 , I ) independent of X 0 , and let P 0 | t ( · | x t ) denote the conditional distribution of X 0 given X t = x t . We first note that since

we have the following calculations:

Using ˙ x t = -1 2 ∇ ln p t ( x t ) , we calculate

Webound each term above separately. For the first term, a quick calculation shows that ∂ t ∇ ln p t ( y ) = -Cov P 0 | t ( ·| y ) ( ∥ y -·∥ 2 2 t 2 , ∇ V ) is finite a.s.: by Cauchy-Schwarz, it suffices to show E P 0 | t ( ·| y ) [ ∥ y -·∥ 4 ] and E P 0 | t ( ·| y ) [ ∥∇ V ∥ 2 ] are finite for all y , and this follows because for t ≤ 1 2 L the measure P 0 | t ( · | y ) is strongly log-concave and ∥∇ V ∥ 2 can be bounded by a quadratic. Because ∇ V is L -Lipschitz,

Now P 0 | t ( · | y ) has density p 0 | t ( x | y ) ∝ p 0 ( x ) e -∥ x -y ∥ 2 2 t so if ∥∇ 2 ln p 0 ∥ op ≤ L and t ≤ 1 2 L , then P 0 | t is 1 2 t -strongly log-concave. By Talagrand's transport cost inequality,

Plugging this back in (10) and using Fatou's lemma and the chain rule for KL,

Now

Plugging into (11) gives

For the second term, by assumption we have ∥∇ 2 ln p t ∥ op ≤ L . Then, since x t ∼ p t ,

using the fact E µ [ ∥∇ ln µ ∥ 2 ] ≤ Ld for any measure µ such that ln µ is L -smooth, which follows from integration by parts. From (9), (12), and (13), and the elementary inequality ⟨ a, b ⟩ ≤ ∥ a ∥ 2 +4 ∥ b ∥ 2 , we get

The above result holds for the dynamics ˙ x t = -1 2 ∇ ln p t ( x t ) for which ( p t ) t ≥ 0 follows the heat flow; this corresponds to the variance-exploding SGM. In this paper, since we wish to consider the SGM based on the variance-conserving Ornstein-Uhlenbeck (OU) process, we can apply the following reparameterization lemma.

Lemma 2 (Reparameterization) . Suppose that ( x t ) t ≥ 0 satisfies the probability flow ODE for Brownian motion starting at p 0 ; that is, letting p t = p 0 ∗ N (0 , tI ) , we have x 0 ∼ p 0 , ˙ x t = -1 2 ∇ ln p t ( x t ) . Then, if we set

then ( y t ) t ≥ 0 satisfies the probability flow ODE for the OU process starting at p 0 ; that is, letting q → t be the density of the OU process at time t , we have y 0 ∼ p 0 = q → 0 , ˙ y t = -y t -∇ ln q → t ( y t ) .

Proof. By direct calculation, one can check that for any y ∈ R d , it holds that q → t ( y ) ∝ p e 2 t -1 ( e t y ) . The claim follows from the chain rule.

Lemma 3 (Score perturbation for OU) . Suppose q → t is the density of the OU process at time t , started at q → 0 , and y 0 ∼ q → 0 , ˙ y t = -y t -∇ ln q → t ( y t ) . Suppose for all t and all x that ∥∇ 2 ln q → t ( x ) ∥ op ≤ L , where L ≥ 1 . Then,

Proof. Using the relationship q → t ( y ) ∝ p e 2 t -1 ( e t y ) ,

If ∥∇ 2 ln q → t ∥ op ≤ L , then ∥∇ 2 ln p e 2 t -1 ∥ op ≤ e -2 t L . By Lemma 1,

Hence

Next,

The result follows.

Finally, we use Lemma 3 to derive a bound on how much the score changes along the trajectory of the probability flow ODE.

Corollary 1. Consider the setting of Lemma 3, and suppose 0 < s < t , h = t -s .

In the first case, this is bounded by O ( L 3 dh 2 ) . In the second case, this is bounded by O ( L 2 dh ∫ t s 1 u d u ) = O ( L 2 dh ln( t/s )) = O ( L 2 dh 2 /t ) .

C Predictor step

Next, we need an ODE discretization analysis.

Lemma 4. Suppose the score function satisfies Assumption 2. Assume that L ≥ 1 , h ≲ 1 /L , and T -( t 0 + h ) ≥ T -t 0 2 . Then

Proof. We have the ODEs

for t 0 ≤ t ≤ t 0 + h , with x t 0 = ̂ x t 0 ∼ q , x t 0 + h ∼ qP ODE , and ̂ x t 0 + h ∼ q ̂ P ODE . Then,

By Grönwall's inequality, noting that h = O (1) ,

We split up the error term as

By Corollary 1, the expectation of the first term is bounded by

The second term is bounded in expectation by ε 2 sc . Plugging back into (14) gives

The Wasserstein distance is bounded by the square root of this quantity, and the lemma follows.

Lemma 4 suggests that focusing on the dependence on d , we will be able to take h ≍ d -1 / 2 (we need to keep one factor of h in the bound, as we need to sum up the bound over 1 /h iterations).

Remark 2 . Our improved score perturbation lemma is necessary to obtain this d 1 / 2 dependence. The original score perturbation lemma [LLT22, Lemma C.11-12] combined with a space discretization bound gives a bound of

in place of Corollary 1. Note this is a 1 2 -Hölder continuity bound rather than a Lipschitz bound. The bound in Lemma 4 then becomes

and we would only be able to take h ≍ d -1 . We also note that our bound has an extra factor of max { L 1 / 2 , ( T -t 0 ) -1 / 2 } ; we do not know if this extra factor is necessary.

We now iterate Lemma 4 to obtain the following result. Note that we now need to also assume that the score estimate is L -Lipschitz.

Lemma 5. Suppose that both Assumptions 2 and 3 hold. Let h 1 , . . . , h N > 0 be a sequence such that letting t N = h 1 + · · · + h N , we have t N ≤ 1 /L . Let h max = max 1 ≤ n ≤ N h n .

Proof. We abbreviate P N ODE := P t 0 ,h 1 ,...,h N ODE and ̂ P N ODE := ̂ P t 0 ,h 1 ,...,h N ODE . Using the triangle inequality,

where the bound on the first term is by Lemma 4. By induction,

By assumption, h n +1 + · · · + h N ≤ t N ≤ 1 /L . In the first case, we get

In the second case we get

by interpreting the summation as a Riemann sum, and noting that the condition h n +1 ≤ T -t 0 -t n 2 implies that this is a constant-factor approximation of the integral ∫ T -t 0 T -t 0 -t N 1 t 1 / 2 d t ≲ √ T -t 0 .

Choice of step sizes. In the first case, we can take all the step sizes to be equal, but in the second case, we may need to take decreasing step sizes. Given a target time T -t 0 -t N = δ , by taking h 1 = h max and then

we can reach the target time in

D Corrector step

In §D.1 (resp. §D.2), we will show that if p , q are close in Wasserstein distance, then running the corrector step based on the overdamped (resp. underdamped) Langevin diffusion starting from p and from q for some amount of time results in distributions which are close in total variation distance. In the end-to-end analysis in §E, we combine this 'total variation to Wasserstein' regularization with the Wasserstein discretization analysis of the predictor step in §C in order to establish our final results.

D.1 Corrector via overdamped Langevin

We will take the potential and score estimate defining the Markov kernels P LD and ̂ P LMC from §2.3 to be U and s respectively. Recall that these correspond respectively to running the overdamped Langevin diffusion with stationary distribution q ∝ exp( -U ) and running the discretized diffusion with score estimate s , both for time h .

The main result of this section is to show that p ̂ P N LMC and q are close in total variation if p and q are close in Wasserstein.

Theorem 4 (Overdamped corrector) . For any T corr := Nh ≲ 1 /L ,

In particular, for T corr ≍ 1 /L ,

We will bound TV ( pP N LD , q ) and TV ( p ̂ P N LMC , pP N LD ) separately. For the former, we use the following short-time regularization result:

Lemma 6 ([BGL01, Lemma 4.2]) . If T corr ≲ 1 /L , then

Proof. The first inequality is Pinsker's inequality. The second inequality is a consequence of [BGL01, Lemma 4.2], which gives a bound of KL ( pP N LD ∥ q ) ≲ L (1 + 1 / ( e 2 LT corr -1)) W 2 2 ( p, q ) . The claim then follows from simplifying by using T corr ≲ 1 /L .

For the latter term, we introduce notation for two stochastic processes.

Note that for any integer k ≥ 0 ,

Observe that marginally, qP k LD = q for any k ≥ 0 because q is the stationary distribution of the Langevin diffusion. The three processes are coupled by using the same Brownian motion and by coupling x 0 = ̂ x 0 ∼ p and x ◦ 0 ∼ q optimally.

Before we proceed to bound TV ( pP N LD , p ̂ P N LMC ) , we need the following simple lemma.

Lemma 7. If T corr ≲ 1 /L , then for all 0 ≤ t ≤ T corr .

Proof. By Itô's formula,

so that if we couple the two processes by coupling x 0 and x ◦ 0 optimally, we conclude that

recalling that t ≤ T ≲ 1 /L

It remains to bound TV ( p ̂ P N LMC , pP N LD ) .

Lemma 8. If T ≲ 1 /L

Proof. As x and ̂ x are driven by the same Brownian motion, by Girsanov's theorem 8 and the data processing inequality we have

We can decompose the integrand as follows:

where we used Lemma 7 to bound E [ ∥ x ◦ u -x u ∥ 2 ] .

It remains to bound E [ ∥ x ◦ kh -x ◦ u ∥ 2 ] . Note that

where in the last step we used that E [ ∥∇ U ( x ◦ u ) ∥ 2 ] ≤ Ld . Substituting this into (15), we obtain

The claimed bound on TV ( pP N LD , p ̂ P N LMC ) follows by Pinsker's inequality.

Proof of Theorem 4. This is immediate from Lemma 6 and Lemma 8, recalling that T corr ≲ 1 /L so that the bound in Lemma 6 dominates the W 2 ( p, q ) term in Lemma 8.

D.2 Corrector via underdamped Langevin

Throughout, we set the friction parameter to

We will take the potential and score estimate defining the Markov kernels P ULD and ̂ P ULMC from §2.3 to be U and s respectively. Recall that these correspond respectively to running the underdamped Langevin diffusion with stationary distribution q and running the discretized diffusion with score estimate s , both for time h .

Given probability measures p and q , we write p := p ⊗ γ d and q := q ⊗ γ d , where γ d is the standard Gaussian measure in R d .

The main result of this section is to show that p ̂ P N ULMC and q are close in total variation if p and q are close in Wasserstein. Compared to §D.1, the discretization error for the underdamped Langevin diffusion is smaller.

Theorem 5 (Underdamped corrector) . For T corr ≲ 1 / √ L ,

In particular, if we take T corr ≍ 1 / √ L , then

We will bound TV ( p P N ULD , q ) and TV ( p ̂ P N ULMC , p P N ULD ) separately. For the former, we use the short-time regularization result of [GW12]:

Lemma 9. If T corr ≲ 1 / √ L , then

Proof. This is a consequence of [GW12, Corollary 4.7 (1)]. The condition to check therein is their Eq. (3.6), which in our setting is satisfied by the constants K 1 = L and K 2 = γ . The Corollary then states that for the cost function

we have KL ( p P N ULD ∥ q ) ≤ W c T corr ( p ⊗ γ d , q ⊗ γ d ) . For v = v ′ and T corr ≲ 1 / √ L , note that

so the claim follows by Pinsker's inequality.

Next, we define the following processes: d z ◦ t = v ◦ t d t , d z t = v t d t ,

It follows that for any integer k ≥ 0 ,

We couple these processes by using the same Brownian motion and coupling q ⊗ γ d and p ⊗ γ d optimally (in particular, v 0 = v ◦ 0 ).

Before we proceed to bound TV ( pP N ULD , p ̂ P N ULMC ) , we start with the following lemma.

Lemma 10. If T corr ≲ 1 / L , then for all 0 ≤ t ≤ T corr ,

√

Proof. We have

and the operator H t satisfies

Let α := 2 / γ . For the vectors δ t := ( z t + αv t ) -( z ◦ + αv ◦ ) and η t := z t -z ◦ , we have

so if we couple the two processes by coupling z and z optimally and taking v = v , we obtain

E [ ∥ z t -z ◦ t ∥ 2 ] ≲ E [ ∥ δ t ∥ 2 + ∥ η t ∥ 2 ] ≤ e O ( Lt ) E [ ∥ δ 0 ∥ 2 + ∥ η 0 ∥ 2 ] ≲ e O ( Lt ) W 2 2 ( p, q ) ≲ W 2 2 ( p, q ) , recalling that t ≤ T corr ≲ 1 / L

0 ◦ 0 0 ◦ 0 √ √ √ by hypothesis.

It remains to bound TV ( p ̂ P N ULMC , p P N ULD ) .

Lemma 11. If T corr ≲ 1 / L , then

Proof. As ( z, v ) and ( z ◦ , v ◦ ) are driven by the same Brownian motion, by Girsanov's theorem 9 and the data processing inequality we have

We can decompose the integrand as follows:

It remains to bound E [ ∥ z ◦ kh -z ◦ u ∥ 2 ] . Note that

where in the last step we used the fact that v ◦ s ∼ γ d . Substituting this into (16), we conclude that

The claimed bound on TV ( p P N ULD , p ̂ P N ULMC ) follows by Pinsker's inequality.

Proof of Theorem 5. This is immediate from Lemma 9 and Lemma 11, recalling that T corr ≲ 1 / √ L so that the bound in Lemma 9 dominates the W 2 ( p, q ) term in Lemma 11.

Remark 3 . In all other sections of this paper, we abuse notation as follows. Given a distribution p on R d , we write p ̂ P ULMC to denote the projection onto the z -coordinate of p ̂ P ULMC , i.e., we view ̂ P ULMC as a Markov kernel on R d rather than on R d × R d (and similarly for P ULD ).

E End-to-end analysis

Lemma 12 (TV error after one round of predictor and corrector) . Choose predictor step sizes h 1 , . . . , h N pred as in Lemma 5 with T pred = h 1 + · · · + h N pred ≤ 1 /L . That is, if T -t 0 -T pred ≲ 1 /L , then we ensure that h n +1 ≤ T -t 0 -h 1 -···-h n 2 for all n , and if T -t 0 ≳ 1 /L , then we can take h 1 = · · · = h N . Let h pred := max 1 ≤ n ≤ N pred h n and abbreviate P ( N pred ) ODE := P t 0 ,h 1 ,...,h N pred ODE (and similarly for ̂ P ODE ).

Consider running the overdamped Langevin corrector for time T corr ≍ 1 /L , step size h corr , and stationary distribution q t 0 P ( N pred ) ODE = q t 0 + T pred ; set N corr = T corr /h corr . Then,

Consider running the underdamped Langevin corrector for time T corr ≍ 1 / √ L , step size h corr , and stationary distribution q t 0 P ( N pred ) ODE = q t 0 + T pred ; set N corr = T corr /h corr . Then,

√ √ ε

Proof. By the triangle inequality and the data-processing inequality,

For overdamped Langevin, applying Theorem 4,

For the Wasserstein term, Lemma 5 yields

Combining these bounds yields the result for the overdamped corrector. For the underdamped corrector, we modify (17) by replacing the use of Theorem 4 with Theorem 5.

We also need the following lemma on the convergence of the OU process.

Lemma 13. Let ( q → t ) t ≥ 0 denote the marginal law of the OU process started at q → 0 = q ⋆ . Then, for all T ≳ 1 , it holds that

Proof. This follows from [CLL23, Lemma C.4]. Alternatively, using the short-time regularization result of [BGL01, Lemma 4.2] for time t 0 ≍ 1 and the Wasserstein contraction of the OU process,

We now prove our main theorems.

Proof of Theorems 2 and 3. For t ∈ [0 , T ] , let p t := law( ̂ x t ) . From Lemma 13,

We divide our analysis according to the two stages of the algorithm. In the first stage, after iterating Lemma 12 for N 0 ≍ LT steps,

where p = 1 2 if we use the overdamped corrector and p = 1 if we use the underdamped corrector. Applying the second part of Lemma 12 for the second stage of the algorithm, we then conclude that

Finally, we note that if we take δ ≍ ε 2 L 2 ( d ∨ m 2 2 ) , then by [LLT23, Lemma 6.4], TV ( q T -δ , q T ) ≤ ε ; a triangle inequality thus finishes the proof.

Remark 4 . Alternatively, instead of taking geometrically decreasing step sizes and employing early stopping, we could split the algorithm into two stages: for time t < T -h pred , we take constant step size h pred , and for time t > T -h pred , we use a smaller constant step size h ′ as required if working with the original score perturbation lemma (see Remark 2).

F Numerical experiments

In this section, we provide preliminary numerical experiments to illustrate our theory. We implement DPUM on a toy model that is not log-concave (mixture of Gaussians).

Setup. The target distribution is a mixture of five Gaussians in dimension 5. On Figures 1 and 2, the red stars represent the means of the Gaussians and the red ellipses around the stars represent the variances of the Gaussians. We start by sampling 500 independent points (in blue) from a standard Gaussian. Then, we run DPUM from the blue dots over 300 iterations and plot the two first coordinates of the dots at iterations 0, 100, 200 and 300. This is a low-dimensional toy example so it does not illustrate our theory, rather we include it as a simple sanity check.

Parameters. We use a closed form formula for the score along the forward process. In other words, the score estimation error is equal to zero. The step size of the predictor is 0 . 01 and the step size of the corrector is 0 . 001 . The corrector consists in 3 steps of the underdamped Langevin algorithm. In the latter algorithm, we initialize the velocity as a centered Gaussian random variable with standard deviation 0 . 001 and set the parameter γ to 0 . 01 .

Observations. We observe the expected behavior: although the target distribution is highly non-logconcave, DPUM is able to provide samples from a distribution that is close to the target distribution. In particular, the initial Gaussian distribution splits in clusters that will fit each component of the target mixture of Gaussians. Recall that we experiment without score error but with discretization error: our numerical results illustrate the common wisdom that score knowledge along the forward process can replace convexity assumptions. In particular, we observe that even isolated, low probability components of the Gaussian mixture, are recovered by DPUM .

Iteration 0

x

-4

-2

0

2

4

x

-4

-2

0

2

4

y

p = 0.03

p = 0.32

p = 0.42

p = 0.03

p = 0.20

-4

-2

0

2

4

x

-4

-2

0

2

4

y

p = 0.03

p = 0.32

p = 0.42

p = 0.03

p = 0.20

-4

-2

0

2

4

-4

-2

0

2

4

y

p = 0.03

p = 0.32

p = 0.42

p = 0.03

p = 0.20

Iteration 200

-4

-2

0

2

4

-4

-2

0

2

4

y

p = 0.03

p = 0.32

p = 0.42

p = 0.03

p = 0.20

Iteration 300

Iteration 100

x

-3

-2

-1

0

1

2

3

4

5

x

-2

0

2

4

y

p = 0.46

p = 0.17 p = 0.08

p = 0.03

p = 0.25

Iteration 0

-3

-2

-1

0

1

2

3

4

5

x

-2

0

2

4

y

p = 0.46

p = 0.17 p = 0.08

p = 0.03

p = 0.25

Iteration 100

-3

-2

-1

0

1

2

3

4

5

-2

0

2

4

y

p = 0.46

p = 0.17 p = 0.08

p = 0.03

p = 0.25

Iteration 200

-3

-2

-1

0

1

2

3

4

5

-2

0

2

4

y

p = 0.46

p = 0.17 p = 0.08

p = 0.03

p = 0.25

Iteration 300

x

x

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, authors propose improved solution scheme with the specially chosen corrector step for the the score-based generative models. Within the denoising diffusion probabilistic modeling (DDPM) and the underdamped Langevin diffusion approach authors obtain better dimension dependence within certain restrictions on the smoothness of data distribution (Algorithm DPUM) than prior works.\n\nSTRENGTHS:\n1. Score-based generative models today are of great interest to the scientific community, while the theoretical establishment of their capabilities and limitations is a powerful impetus for the further development of new effective methods. The extensive analysis carried out in this work is certainly relevant and of direct interest to researchers.\n\n2. The work is well structured, written in a great style and easy to read.\n\n3. The provided approach with the polynomial-time convergence guarantees for the probability flow ODE seems to be essential if it will be successfully implemented in practice and applied to the real-world data.\n\nWEAKNESSES:\nGiven the high level of theoretical consideration carried out in the work, I am confused by such a weak study of the experimental part:\n\n1. It seems that a (at least brief) description of the performed numerical experiments should be included as a section of the main text (in the presented text, numerical experiments are contained exclusively in the appendix).\n\n2. I do not quite understand the conclusions that can be drawn from the only simple numerical experiment presented in the appendix for artificially generated data from a mixture of Gaussians. Does the proposed theoretical estimate of the convergence of the method hold in this case? How to interpret a word \"approximately\" in [Line 721, appendix]: \"DPUM is able to provide samples (approximately) from the target distribution\"? Can any quantitative conclusions be drawn from the experiment? Is it possible to do some form of convergence analysis and make a comparison with the baseline?\n\n3. Can numerical experiments be done on real-world datasets using a neural network ([Line 130]: \"the score estimates are produced via a deep neural network\"; [Line 723, appendix]: \"recall that we experiment without score error\")?\n\nQUESTIONS:\nPlease see the section \"Weaknesses\" above. I have only one more minor comment: [Line 123] for clarity, it may be worthwhile, in addition to the dimension of the noise, to also indicate the dimension of the x variable.\n\nLIMITATIONS:\nThe proposed approach is based on the assumption that the score functions along the forward process are Lipschitz continuous, and it does not cover non-smooth distributions. However, the authors note that this limitation can potentially be circumvented by using the early stopping technique.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
This paper provides polynomial-time convergence guarantees for the probability flow ODE implementation. This paper obtains a better better dimension dependence than prior works on DDPM. However, there is no experiments to prove the effectiveness of the proposed methods.

STRENGTHS:
1. The paper is well written.
2. The theoretical results are solid.

WEAKNESSES:
1. No experiments are performed to prove the proposed method. I think it is necessary to evaluate the method of diffusion models, such as Stable Diffusion or DeepFloyd. I admit that the theoretical results presented in this paper might be important, but that is not the reason of the lack of experiments.
2. The proposed method has an extra corrector step, which further introduces extra computational costs. 
3. Some comparisons with previous sampling methods (DPM-Solver++, DEIS, etc) are missing.

QUESTIONS:
Please see the weaknesses.

LIMITATIONS:
The limitations have been fully discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors proposed the convergence guarantees for the probability flow ODE implementation. Through the use of a specially chosen corrector step based on the underdamped Langevin diffusion, a better dimension dependence of $O(\\sqrt{d})$ than DDPM $O(d)$ is achieved.\n\nSTRENGTHS:\n1. the theoretical comparison between probability flow ODE and DDPM is less studied in the community. In this sense, analyzing the advantage of the ODE-based method brings new insight to the community. \n\n2. solid and clean derivations.\n\nWEAKNESSES:\n1. the term of fast looks a bit confusing to me. It is fast in inference or training? they have different meanings. As provided in the ranking: https://paperswithcode.com/sota/image-generation-on-cifar-10, **flow-based methods seem to perform quite mediocre in standard datasets**. So at least the claim of fast flow-ODE contradicts the fact of fast training in practice. Clarification is suggested.\n\n2. Flow-based methods are known to work well in **low dimensions** and I am even suspecting it suffers from the curse of dimensionality. As such. I appreciate the authors' solid math techniques, but I personally believe that the advantage of the well-known smaller discretization error $O(\\sqrt{d})$ v.s. $O(d)$ becomes **less interesting if it mainly adapts to low-dimensional problems** (see Q1 below). the authors can correct me if I am wrong.\n\n3. The authors should point out the **limitations of probability-flow ODE methods** compared to DDPM so as not to confuse readers.\n\nQUESTIONS:\nQ1. Do you know any popular real-world examples that probability-flow ODE works well in ultra-high dimensions in practice?\n\nQ2. Growall inequality does suffer from the exponential blowup. \nQ2.a) The authors propose to use the geometric contraction of the corrector (diffusion) step to control the exponential blowup. Is that right? \nQ2.b) Does it mean we can do lots of ODE steps but with only a few diffusion steps, the exponential blowup is still controlled in terms of complexity (ignore the constant) due to the contraction (closer to 1, e.g. alpha=0.999999999 also works, although the error / (1-alpha) is large)?\nQ2.c) Or do you need to maintain at least some amount of diffusion steps to make sure the bound hold?\n\nMinor: a relative work is missing and a short comment on this work would be appreciated: Provably Convergent Schrodinger Bridge with Applications to Probabilistic Time Series Imputation. ICML'23.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe current study offers a polynomial-time total variation convergence guarantee for score-based generative modeling, given that the generative process is a probabilistic Ordinary Differential Equation (ODE). This differs from previous work, which typically involves a reverse Stochastic Differential Equation (SDE).  To the best of my knowledge, this is the first paper to present such results for an ODE-implemented score-based generative model targeting a general distribution without assuming log-concavity. The proof technique includes analyzing the discretization error of both the ODE and SDE, and employs the Score Perturbation Lemma to limit the error introduced by using estimated score functions.\n\nIt's important to note that the convergence guarantee is provided for a modified version of the ODE-implemented score-based generative model. The authors have added an additional stochastic correction step to the deterministic dynamics, which is an unconventional approach.\n\nSTRENGTHS:\nThe results provided in the work is novel and important in my perspective. As mentioned in the summary, this is the first convergence guarantee for an ODE-based, score-based generative model for general target distributions that I am aware of (this could, of course, be due to my limited familiarity with the relevant literature). While some may contend that the additional stochastic correction step diminishes the impact of the work (given that the primary challenge of the proof is that the ODE transformation is purely deterministic, making Total Variation/Kullback-Leibler control nearly impossible), the theoretical framework employed remains significantly distinct from its SDE counterparts.\n\nThe clarity of the writing and the structure of the proof are commendable. Despite the complexity of the paper, the proof is not difficult to follow, largely owing to a comprehensive preliminary introduction to the proof techniques. I particularly value the fact that the authors dedicated a significant portion of the text to discussing the limitations of some existing proof concepts, and made comparisons with similar intermediate results from the literature. This approach further underscores the contributions of this work.\n\nWEAKNESSES:\nThe inclusion of an additional stochastic corrector in the generative process is somewhat unsatisfying. To the best of my knowledge regarding score-based generative modeling via ODE, this correction step is not typically included in practice. One aspect that I believe is not clearly addressed in the current results is the comparison between the corrected and uncorrected generative processes. Ideally, there would be a result showing that the corrected process converges to the target faster.\nWhile it is mentioned in the paper that we cannot expect such a comparison in terms of Total Variation (TV) or Kullback-Leibler (KL) divergence - as without the stochastic correction, the TV/KL bound would be vacuous - could we not show such a comparison in terms of weak distance or even Wasserstein distance?\n(I've also raised some questions regarding the corrector in the Questions section.)\n\nAnother minor shortcoming pertains to the toy numerical experiments presented in the study. It would have been more desirable if the authors had evaluated Algorithms 1 & 2 on actual image generation tasks and compared the performance of the corrected and uncorrected versions in these scenarios. While I admit it might seem somewhat unfair to expect this from a theory-oriented paper, such empirical evidence would undoubtedly render the study more comprehensive and impactful.\n\nQUESTIONS:\n* Step 5 for Alg 1 and 2: Can authors provide some intuition on why geometric decreasing stepsize is needed for the prediction at the end given that correction step is used eventually? (I understand from the proof that when $t>1/L$, decreasing step size is required to control the ODE discretization error + approximated score error, but cannot form a practical intuition on this.)\n\n* In practice, which one of  DPOM and DPUM is preferred? Or how does one choose the algorithm for practical applications? \n\n* Is it possible to get rid of the correction step for the convergence if the ODE discretization error is reduce by using more sphisticated numerical integrator (e.g., high order sympletic integrator)?  Based on my superfluous understanding of the proof, the answer is probably NO since the key of the corrector is to obtain TV bounds based on Wasserstein bounds, intead of \"correcting\" the introduced error during the ODE simulation. This leads to my next question.\n\n* If intead of TV distance, all we care is just wasserstein distance, can we get rid of the correction step? Or is the correction step beneficial for a faster convergence?\n\nLIMITATIONS:\nThe author did note limitations or boundaries of the adopted techniques through out the proof (e.g., line 565 in Appendix C), and raised practical questions base on the theoretical framework.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this article the authors propose polynomial-time convergence guarantees for probability flow ODEs related to score based generative models. \nPrevious analyses on score based diffusion models have resorted to the use of SDEs and relied on properties of stochastic flows to alleviate error accumulation due to discretisation. Here the authors study the ODE counterparts of diffusion models.\nSince the deterministic dynamics are not contractive without log-concave assumptions on the data distribution, and thereby \ndiscretisation errors accumulate, the authors derive bounds for a modified process which requires interleaving deterministic integration steps (predictor) with stochastic (corrector) steps that smooths the reverse process, drawing inspiration from the purely stochastic framework of [4,9]. They propose two variants of the algorithm, one relying on overdamped Langevin dynamics with dimension dependence that scales as $\\mathcal{O}(d)$ ,and one with underdamped Langevin dynamics that scales more favourably as $\\mathcal{O}(\\sqrt{d})$and thus has improved dimension dependence compared to previous approaches.\n\nThe topic of the paper is interesting and relevant for the audience of  NeurIPS for understanding denoising diffusion probabilistic models, however I find that the article could have been more accessible for the broader audience of NeurIPS, and importantly, I would expect that the authors would provide numerical evidence for their proposed guarantees.\n\nSTRENGTHS:\n- the authors provide improved convergence guarantees compared to existing approaches.\n- they relax log-concavity assumptions existing in previous frameworks that are rather restrictive.\n\nWEAKNESSES:\n- The work is purely theoretical and the authors do not provide numerical evidence for their claims excluding a simple experiment in the supplement. I find it a bit odd that they provide a convergence guarantee and propose a scaling with system dimension, but nevertheless they don't explore the said scaling numerically for increasing system dimension. The numerical example they present in the appendix concerns only a single experiment of dimension $d=5$.\n\n- the main theorem of the paper is formulated considering only linear dynamics (Ornstein-Uhlenbeck process). However according to the introduction of the\n paper a reader gets the impression that the provided results are relevant for general dynamics. I would propose to introduce this clarification both in the abstract and the introduction of the \npaper, otherwise the statements may be considered misleading.\n\n\n- the derivation of the probability flow ODE in Eq. 3 is a bit odd, and happens to be correct here because the considered dynamics is conservative. The whole point of [3]  first introduced the\nformalism of ODE dynamics for diffusion that got later named prob. flow ODE by [4], was that there is no need for a stationary solution for the prob. flow ODE and no need to derive the dynamics wrt to the stationary solution. \nFor general systems one can derive the dynamics of the deterministic process by rewriting the Fokker-Planck equation in the form of a Liouville equation\nas\n\\begin{align} \n\\partial_t p_t(\\mathbf{x})& = - \\nabla \\cdot \\left( f(\\mathbf{x}) p_t(\\mathbf{x}) - \\frac{\\sigma^2}{2} \\nabla p_t(\\mathbf{x})  \\right)   \\\\\\\\\n&=- \\nabla \\cdot \\left[p_t(\\mathbf{x})  \\left(  f(\\mathbf{x})  - \\frac{\\sigma^2}{2}  \\nabla \\log p_t(\\mathbf{x})  \\right) \\right] \n\\end{align} \nFrom this formulation one can apply the divergence theorem and obtain the deterministic dynamics similar to [5]. The explanation used in the paper uses the Wasserstein gradient flow structure.\n\n- Poorly written and limited conclusion without discussing the limitations and possible outcomes of the provided guarantees.\n\nMinor:\n\n- The notation of the stationary distribution of the overdamped Langevin dynamics is rather unfortunate, especially considering that then it is used in the same sentences/equations with the parameter $\\gamma$ indicating the damping of the underdamped process.\nI recommend adapting the notation to avoid confusing the readers unnecessarily.\n\nQUESTIONS:\n- How does the proposed approach relate to the flow matching methods, i.e. [8]?\n\n- How does the method compare with concurrent approaches [7] , [9] and [10]?\n\nLIMITATIONS:\n- The authors do not provide numerical experiments to support their claims.\n\n- They do not provide comparisons with already proposed bounds on either stochastic or deterministic implementations of diffusion based generative models.\n\n- They do not discuss limitations of their approach.\n\n\n\n\n---\n**References:**\n\n\n[1] Chen, Sitan, et al. \"Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions.\" arXiv preprint arXiv:2209.11215 (2022).\n\n[2] Lee, Holden, Jianfeng Lu, and Yixin Tan. \"Convergence of score-based generative modeling for general data distributions.\" International Conference on Algorithmic Learning Theory. PMLR, 2023.\n\n[3] Maoutsa, D., Reich, S., & Opper, M. (2020). Interacting particle solutions of Fokker–Planck equations through gradient–log–density estimation. Entropy, 22(8), 802.\n\n[4] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456.\n\n[5] Hermoso, A., Homar, V., & Yano, J. I. (2020). Exploring the limits of ensemble forecasting via solutions of the Liouville equation for realistic geophysical models. Atmospheric Research, 246, 105127. \n\n[6] Albergo, Michael S., and Eric Vanden-Eijnden. \"Building normalizing flows with stochastic interpolants.\" arXiv preprint arXiv:2209.15571 (2022).\n\n[7]  Li, Gen, et al. \"Towards Faster Non-Asymptotic Convergence for Diffusion-Based Generative Models.\" arXiv preprint arXiv:2306.09251 (2023).\n\n[8] Benton, Joe, George Deligiannidis, and Arnaud Doucet. \"Error Bounds for Flow Matching Methods.\" arXiv preprint arXiv:2305.16860 (2023).\n\n[9] Lee, Holden, Jianfeng Lu, and Yixin Tan. \"Convergence of score-based generative modeling for general data distributions.\" International Conference on Algorithmic Learning Theory. PMLR, 2023.\n\n[10] Albergo, Michael S., Nicholas M. Boffi, and Eric Vanden-Eijnden. \"Stochastic interpolants: A unifying framework for flows and diffusions.\" arXiv preprint arXiv:2303.08797 (2023).\n\n[9] Jain, Ajay, and Ben Poole. \"Journey to the BAOAB-limit: finding effective MCMC samplers for score-based models.\" NeurIPS 2022 Workshop on Score-Based Methods. 2022.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}