Title: Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity
Venue: NeurIPS 2024
Accepted: yes
Keywords: diffusion model

parallel sampling

stochastic differential equations

probability flow ode
Abstract: Diffusion models have become a leading method for generative modeling of both image and scientific data.
As these models are costly to train and \emph{evaluate}, reducing the inference cost for diffusion models remains a major goal.
Inspired by the recent empirical success in accelerating diffusion models via the parallel sampling technique~\cite{shih2024parallel}, we propose to divide the sampling process into $\mathcal{O}(1)$ blocks with parallelizable Picard iterations within each block. Rigorous theoretical analysis reveals that our algorithm achieves $\widetilde{\mathcal{O}}(\mathrm{poly} \log d)$ overall time complexity, marking \emph{the first implementation with provable sub-linear complexity w.r.t. the data dimension $d$}. Our analysis is based on a generalized version of Girsanov's theorem and is compatible with both the SDE and probability flow ODE implementations. Our results shed light on the potential of fast and efficient sampling of high-dimensional data on fast-evolving modern large-memory GPU clusters.

Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity

Abstract

Diffusion models have become a leading method for generative modeling of both image and scientific data. As these models are costly to train and evaluate , reducing the inference cost for diffusion models remains a major goal. Inspired by the recent empirical success in accelerating diffusion models via the parallel sampling technique [1], we propose to divide the sampling process into O (1) blocks with parallelizable Picard iterations within each block. Rigorous theoretical analysis reveals that our algorithm achieves ˜ O (poly log d ) overall time complexity, marking the first implementation with provable sub-linear complexity w.r.t. the data dimension d . Our analysis is based on a generalized version of Girsanov's theorem and is compatible with both the SDE and probability flow ODE implementations. Our results shed light on the potential of fast and efficient sampling of high-dimensional data on fast-evolving modern large-memory GPU clusters.

1 Introduction

Diffusion and probability flow based models [2-11] are now state-of-the-art in many fields, such as computer vision and image generation [12-22], natural language processing [23, 24], audio and video generation [25-29], optimization [30, 31], sampling and learning of fixed classes of distributions [32-41], solving high-dimensional partial differential equations [42-46], and more recently several applications in physical, chemical and biological fields [47-63]. For a more comprehensive list of related work, one may refer to the following review papers [64-66]. While there are already many variants, such as denoising diffusion probabilistic models (DDPMs) [7], score-based generative models (SGMs) [9], diffusion schrödinger bridges [67], stochastic interpolants and flow matching [2-4], etc. , the recurring idea is to design a stochastic process that interpolates between the data distribution and some simple distribution, along which score functions or alike are learned by neural network-based estimators, and then perform inference guided by the learned score functions.

Due to the sequential nature of the sampling process, the inference of high-quality samples from diffusion models often requires a large number of iterations and, thus, evaluations of the neural network-based score function, which can be computationally expensive [68]. Efforts have been

Yinuo Ren ∗ †

ICME

Grant M. Rotskoff

̂

˜

made to accelerate this process by resorting to higher-order or randomized numerical schemes [6979], augmented dynamics [80], adaptive step sizes [81], operator learning [82], restart sampling [83], self-consistency [84-87] and knowledge distillation [88-90]. Recently, several empirical works [1, 91-94] leverage the Picard iteration and triangular Anderson acceleration to parallelize the sampling procedure of diffusion models and achieve empirical success in large-scale image generation tasks. Some other recent work [95, 96] also combine the parallel sampling technique with the randomized midpoint method [97] to accelerate the inference of diffusion models.

This efficiency issue is closely related to the problem of bounding the required number of steps and evaluations of score functions to approximate an arbitrary data distribution on R d to δ -accuracy, which has been analyzed extensively in the literature [98-115]. In terms of the dependency on the dimension d , the current state-of-the-art result for the SDE implementation of diffusion models is ˜ O ( d ) [107], improved from the previous ˜ O ( d 2 ) bound [104]. [111] gives a ˜ O ( √ d ) bound for the probability flow ODE implementation by considering a predictor-corrector scheme with the underdamped Langevin Monte Carlo (UMLC) algorithm.

In this work, we aim to provide parallelization strategies, rigorous analysis, and theoretical guarantees for accelerating the inference process of diffusion models. The time complexity of previous implementations of diffusion models has been largely hindered by the discretization error, which requires the step size to scale with ˜ O (1 /d ) for the SDE implementation and ˜ O (1 / √ d ) for the probability flow ODE implementation. We show that the inference process can be first divided into O (1) blocks with parallelizable evaluations of the score function within each, and thus reduce the overall time complexity to ˜ O (poly log d ) . We provide the first implementation of diffusion models with poly-logarithmic complexity , a significant improvement over the current state-of-the-art polynomial results that sheds light on the potential fast and efficient sampling of high-dimensional distributions with diffusion models on fast-developing memory-efficient modern GPU clusters.

1.1 Contributions

We propose parallelized inference algorithms for diffusion models in both the SDE and probability flow ODE implementations (PIADM-SDE/ODE) with exponential integrators, a shrinking step size scheme towards the data end, and the early stopping technique;

We provide a rigorous convergence analysis of PIADM-SDE, showing that our parallelization strategy yields a diffusion model with O (poly log d ) approximate time complexity;

˜ · Weshow that our strategy is also compatible with the probability flow ODE implementation, and PIADM-ODE could improve the space complexity from ˜ O ( d 2 ) to ˜ O ( d 3 / 2 ) while maintaining the poly-logarithmic time complexity.

2 Preliminaries

In this section, we briefly recapitulate the framework of score-based diffusion models, define notations, and discuss related work.

2.1 Diffusion Models

In score-based diffusion models, one considers a diffusion process ( x s ) s ≥ 0 in R d governed by the following stochastic differential equation (SDE):

where ( w s ) s ≥ 0 is a standard Brownian motion, and p 0 is the target distribution that we would like to sample from. The distribution of x s is denoted by p s . Once the drift β s ( · ) , the diffusion coefficient σ s , and a sufficiently large time horizon T are specified, (2.1) also corresponds to a backward process ( ⃗ x t ) 0 ≤ t ≤ T for another arbitrary diffusion coefficient ( υ s ) s ≥ 0 [116]:

where ⃗ ∗ t denotes ∗ T -t , with ⃗ p 0 = p T and ⃗ p T = p 0 .

For notational simplicity, we adopt a simple choice of the drift and the diffusion coefficients in what follows: β t ( x ) = -1 2 x , σ t = I d , and υ = υ I d , under which (2.1) is an Ornstein-Uhlenbeck (OU) process converging exponentially to its stationary distribution, i.e. p T ≈ ̂ p T := N (0 , I d ) , and (2.1) and (2.2) reduce to the following form:

In practice, the score function ∇ ⃗ p t ( ⃗ x t ) is often estimated by a neural network (NN) s θ t ( x t ) , where θ represents its parameters, by minimizing the denoising score-matching loss [117, 118]:

Implementations. Diffusion models admit multiple implementations depending on the choice of the parameter υ in the backward process (2.2). The SDE implementation with υ = 1 is widely used in the literature for its simplicity and efficiency [10], while recent studies [111] claim that the probability flow ODE implementation with υ = 0 may exhibit better time complexity. We refer to [111, 119] for theoretical and [120, 121] for empirical comparisons of different implementations.

2.2 Parallel Sampling

Parallel sampling algorithms have been actively explored in the literature, including the parallel tempering method [122-124] and several recent studies [125-127]. For diffusion models, the idea of parallel sampling is based on the Picard iteration [128, 129] for solving nonlinear ODEs. Suppose we have an ODE d x t = f t ( x t )d t and we would like to solve it for t ∈ [0 , T ] , then the Picard iteration is defined as follows:

Under assumptions on the Lipschitz continuity of f t , the Picard iteration converges to the true solution exponentially fast, in the sense that ‖‖ x ( k ) t -x t ‖‖ L ∞ ([0 ,T ]) ≤ δ with K = O (log δ -1 ) iterations. Unlike high-order ODE solvers, the Picard iteration is intrinsically parallelizable: for any t ∈ [0 , T ] , the computation of x ( k +1) t relies merely on the values of the most recent iteration x ( k ) t . With sufficient computational sources parallelizing the evaluations of f , the computational cost of solving the ODE no longer scales with T but with the number of iterations K .

Recently, this idea has been applied to both the Langevin Monte Carlo (LMC) and the underdamped Langevin Monte Carlo (UMLC) contexts [130]. Roughly speaking, it is proposed to simulate the Langevin diffusion process d x t = -∇ V ( x t )d t +d w t with the following iteration resembling (2.6):

where all iterations share a common Wiener process ( w t ) t ≥ 0 .

It is shown that for well-conditioned log-concave distributions, parallelized LMC would achieve an iteration depth of K = ˜ O (poly log d ) that matches the indispensable time horizon T = ˜ O (poly log d ) to achieve exponential ergodicity ( cf. [130, Theorem 13]). This promises a significant speedup in sampling high-dimensional distributions from the standard LMC of T = ˜ O ( d ) iterations, hindered by the o (1 /d ) step size as imposed by the discretization error and now evaded by the parallelization.

2.3 Approximate Time Complexity

A similar situation is expected in diffusion models, where the application bottleneck is largely the inference process with sequential iterations and expensive evaluations of the learned score function s θ t ( · ) , which is often parametrized by large-scale NNs. Despite several unavoidable costs involving pre- and post-processing, data storage and retrieval, and arithmetic operations, we define the following notion of the approximate time complexity of the inference process of diffusion models:

Definition 2.1 (Approximate time complexity) . For a specific implementation of diffusion models (2.5), we define the approximate time complexity of the sampling process as the number of unparallelizable evaluations of the learned NN-based score function s θ t ( · ) .

This definition coincides with the notion of the number of steps required to reach a certain accuracy in [104, 100], iteration complexity in [107, 111], etc. in the previous theoretical studies of diffusion models. We have adopted this notion in Table 1 for a comparison of the current state-of-the-art results and our bounds in this work. We will use the notion of space complexity likewise to denote the approximate required storage during the inference. Trivially, the space complexity of the sequential implementation is O ( d ) . Should no confusion occur, we omit the dependency of the complexities above on the accuracy threshold δ , etc. , during our discussion, as we focus on applications of diffusion models to high-dimensional data distributions, following the standard practice in the literature.

3 Main Results

Inspired by the acceleration achieved by the parallel sampling technique in LMC and ULMC, we aim to accommodate parallel sampling into the theoretical analysis framework of diffusion models. The benefit of the parallel sampling technique in this scenario has been recently confirmed by up to 14 × acceleration achieved by the ParaDiGMS algorithm [1] and ParaTAA [92], where several practical compromises are made to mitigate GPU memory constraints and theoretical guarantees are still lacking.

In this section, we will propose P arallelized I nference A lgorithms for D iffusion M odels with both the SDE and probability flow ODE implementations, namely the PIADM-SDE (Algorithm 1) and PIADM-ODE (Algorithm 2), and present theoretical guarantees of our algorithms, including the approximate time complexity and space complexity, for both implementations in Section 3.1 and Section 3.2, respectively. Due to the large number of notations used in the presentation, we give an overview of notations in Appendix A.1 for readers' convenience.

3.1 SDE Implementation

We first focus on the approximation, parallelization strategies, and error analysis of diffusion models with the SDE implementation, i.e. the forward and backward process (2.3) and its approximatation (2.5) with υ = 1 . We will show that PIADM-SDE achieves an ˜ O (poly log d ) approximate time complexity with ˜ O ( d 2 ) space complexity .

Outer Iterations

n

:

N

=

O

(log

d

)

blocks

̂

q

0

≈ N

(

0

,

I

d

)

̂

q

t

N

≈

p

data

h

0

h

1

h

n

-

1

h

N

-

1

η

O

(1)

k

= 0

k

= 1

k

=

K

(

d

-

1

)

or

(

d

-

1

/

2

)

.

.

.

.

.

.

M

n

=

˜

O

(

d

)

or

˜

O

(

√

d

)

parallalizable steps

̂

q

t

n

̂

q

t

n

+1

Inner

Iterations

k

:

K

=

˜

O

(log

d

)

depth

ϵ

n,

0

ϵ

n,

1

ϵ

n,M

-

1

3.1.1 Algorithm

PIADM-SDE is summarized in Algorithm 1 and illustrated in Figure 1. The main idea behind our algorithm is the fact that (2.5) can be efficiently solved by the Picard iteration within a period of O (1) length, transferring ˜ O ( d ) sequential computations to a parallelizable iteration of depth ˜ O (log d ) . In the following, we introduce the numerical discretization scheme of our algorithm and the implementation of the Picard iteration in detail.

Step Size Scheme. In our algorithm, the time horizon T is first segmented into N blocks of length ( h n ) N -1 n =0 , with each h n ≤ h := T/N = Ω(1) , forming a grid ( t n ) N n =0 with t n = ∑ n j =1 h j . For any n ∈ [0 : N -1] , the n -th block is further discretized into a grid ( τ n,m ) M n m =0 with τ n, 0 = 0 and τ n,M n = h n . We denote the step size of the m -th step in the n -th block as ϵ n,m = τ n,m +1 -τ n,m , and the total number of steps in the n -th block as M n .

For the first N -1 blocks, we simply use the unique discretization, i.e. h n = h , ϵ n,m = ϵ , and M n = M := h/ϵ , for n ∈ [0 : N -2] and m ∈ [0 : M -1] . Following [104, 107], to curb the potential blow-up of the score function as t → T , which is shown by [107] for 0 ≤ s < t < T to be of the order we apply early stopping at time t N = T -η , where η is chosen in a way such that the O ( √ η ) 2-Wasserstein distance between ⃗ p T and its smoothed version ⃗ p T -η that we aim to sample from alternatively, is tolerable for the downstream tasks. We also impose the exponential decay of the step size towards the data end in the last block. To be specific, we let h N -1 = h -δ , and discretize the interval [ t N -1 , t N ] = [( N -1) h, T -η ] into a grid ( τ N -1 ,m ) M N -1 m =0 with step sizes ( ϵ N -1 ,m ) M N -1 m =0 satisfying

As shown in Lemma B.7, this exponential decaying step size scheme towards the data end is crucial to bound the discretization error in the last block.

For the simplicity of notations, we introduce the following indexing function: for τ ∈ [ t n , t n +1 ] , we define I n ( τ ) to be the unique integer such that ∑ I n ( τ ) j =1 ϵ n,j ≤ τ < ∑ I n ( τ )+1 j =1 ϵ n,j . We also define

O

O

Algorithm 1: PIADM-SDE

a piecewise function g such that g n ( τ ) = ∑ I n ( τ ) j =1 ϵ n,j . It is easy to check that under the uniform discretization for n ∈ [1 : N -1] , we have I n ( τ ) = glyph[floorleft] τ/ϵ glyph[floorright] and g n ( τ ) = glyph[floorleft] τ/ϵ glyph[floorright] ϵ .

Exponential Integrator. For each step τ ∈ [ t n + τ n,m , t n + τ n,m +1 ] , we use the following exponential integrator scheme [77], as the numerical discretization of the SDE (2.5):

Remark 3.1. One could also implement a straightforward Euler-Maruyama scheme instead of the exponential integrator (3.4), where an additional high-order discretization error term would emerge [104, Theorem 1], which we believe would not affect the overall ˜ O (poly log d ) time complexity with parallel sampling.

Picard Iteration. Within each block, we apply Picard iteration of depth K . As shown by Lemma B.3, the discretized scheme (3.4) implements the following iteration for k ∈ [0 : K -1] :

We denote the distribution of ̂ y ( K ) t n ,τ by ̂ q t n + τ . As proved in Lemma B.6, the iteration above would converge to (3.2) in each block exponentially fast, which given a sufficiently accurate learned score estimation s θ t should be close to the true backward SDE (2.3). One should also notice that the Gaussians ξ m are only sampled once and used for all iterations.

The parallelization for (3.4) in Algorithm 1 should be understood as that for any k ∈ [0 : K -1] , each s θ t n + τ n,j ( ̂ y ( k ) t n ,τ n,j ) for j ∈ [0 : M n ] is evaluated in parallel, with subsequent floating-point operations comparably negligible, resulting in the overall O ( NK ) approximate time complexity.

3.1.2 Assumptions

Our theoretical analysis will be built on the following mild assumptions on the regularity of the data distribution and the numerical properties of the neural networks:

Assumption 3.1 ( L 2 ([0 , t N ]) δ -accurate learned score) . The learned NN-based score s θ t is δ 2 -accurate in the sense of

Assumption 3.2 (Regular and normalized data distribution) . The data distribution p 0 has finite second moments and is normalized such that cov p 0 ( x 0 ) = I d .

Assumption 3.3 (Bounded and Lipschitz learned NN-based score) . The learned NN-based score function s θ t has bounded C 1 norm, i.e. ‖‖ s θ t ( · ) ‖‖ L ∞ ([0 ,T ]) ≤ M s with Lipschitz constant L s .

Remark 3.2. Assumption 3.1 and the finite moment assumption in Assumption 3.2 are standard assumptions across previous theoretical works on diffusion models [100, 104, 111], while we adopt the normalization Assumption 3.2 from [107] to simplify true score function-related computations ( cf. Lemma A.8). Assumption 3.3 can be easily satisfied by truncation, ensuring computational stability. Notice that the exponential integrator, one actually applies Picard iteration to e -t/ 2 s θ t , a relaxation of Assumption 3.1 might be possible, which is left for future work.

3.1.3 Theoretical Guarantees

The following theorem summarizes our theoretical analysis for PIADM-SDE (Algorithm 1):

Theorem 3.3 (Theoretical Guarantees for PIADM-SDE) . Under Assumptions 3.1, 3.2, and 3.3, given the following choices of the order of the parameters

˜ and let L 2 s h n e 7 2 h n glyph[lessmuch] 1 , δ 2 ≲ δ , T ≲ log η -1 , the distribution ̂ q t N that PIADM-SDE (Algorithm 1) generates samples from satisfies the following error bound:

D KL ( p η ‖ ̂ q t N ) ≲ de -T + dϵT + δ 2 2 + dTe -K ≲ δ 2 , with a total of KN = ˜ O ( log 2 ( dδ -2 ) ) approximate time complexity and dM = ˜ O ( d 2 δ -2 ) space complexity for parallalizable δ -accurate score function computations. Remark 3.4. We would like to make the following remarks on the result above:

The acceleration from ˜ O ( d ) to ˜ O (poly log d ) is at the cost of a trade-off with extra memory cost of M = ˜ O ( d ) for computing and updating { s θ t n + τ n,j ( ̂ y ( k ) t n ,τ n,m ) } m ∈ [0: M n ] simultaneously during each Picard iteration;

As mentioned above, the scale of the step size ϵ within one block is still confined to Θ(1 /M ) = ˜ Θ(1 /d ) . The block length h , despite being required to be small compared to 1 /L s , is of order Θ(1) , resulting in only Θ(log d ) blocks and thus ˜ O (poly log d ) total iterations. 3.1.4 Proof Sketch

Compared with log-concave sampling [130], M being of order ˜ O ( d ) instead of ˜ O ( √ d ) therein is partly due to the time independence of the score function ∇ log p ( · ) in general sampling tasks. Besides, the scaling M = ˜ O ( d ) agrees with the current state-of-the-art dependency [107] for the SDE implementation of diffusion models;

The detailed proof of Theorem 3.3 is deferred to Section B. The pipeline of the proof is to (a) first decompose the error D KL ( ⃗ p t N ‖ ̂ q t N ) into blockwise errors using the chain rule of KL divergence; (b) bound the error in each block by invoking Girsanov's theorem; (c) sum up the errors in all blocks.

The key technical challenge lies in Step (b). Different from all previous theoretical works [100, 104, 111], the Picard iteration in our algorithm generates K paths recursively in each block using the learned score s θ t . And therefore the final path ( ̂ y ( K ) t n ,τ ) τ ∈ [0 ,h n ] depends on all previous paths ( ̂ y ( k ) t n ,τ ) τ ∈ [0 ,h n ] for k ∈ [0 : K -1] , ruling out a direct change of measure argument from the

naïve application of Girsanov's theorem. To this end, we need a more sophisticated mathematical framework of stochastic processes, as given in Appendix A.2. We define the measurable space (Ω , F ) with filtrations ( F t ) t ≥ 0 to specify the probability measures on (Ω , F ) of each Wiener process, and resort to one of the most general forms of Girsanov's theorem ( [131, Theorem 8.6.6]). For example, in the n -th block, we apply the following change of measure procedure:

Let q | F t n be the measure where w t ( ω ) is the shared Wiener process in the Picard iteration (3.3) for any k ∈ [0 : K -1] ;

Define another process d w t n + τ ( ω ) = d w t n + τ ( ω ) + δ t n ( τ, ω ) dτ , where

̂ ̂ 3. Invoke Girsanov's theorem, which yields that the Radon-Nikodym derivative of the measure ⃗ p | F t n with respect to q | F t n satisfies

Conclude that ( ˜ w t n + τ ) τ ≥ 0 is a Wiener process under the measure ⃗ p | F t n and thus (3.3) at iteration K satisfies the following SDE:

i.e. the true backward SDE (2.3) with the true score function for τ ∈ [ t n , t n +1 ] .

One should notice that this change of measure argument will cause an additional term in the bound of the discrepancy between the first iteration ̂ y (1) t n ,τ and the initial condition ̂ y (0) t n ,τ in Lemma B.5. However, due to the exponential convergence of the Picard iteration, this term does not affect the overall error bound.

3.2 Probability Flow ODE Implementation

In this section, we will show that our parallelization strategy is also compatible with the probability ODE implementation of diffusion models, i.e. the forward and backward process (2.3) and its approximatation (2.5) with υ = 0 . We will demonstrate that PIADM-ODE (Algorithm 2) further improves the space complexity from ˜ O ( d 2 ) to ˜ O ( d 3 / 2 ) while maintaining the same ˜ O (poly log d ) approximate time complexity .

3.2.1 Algorithm

Due to the space limit, we refer the readers to Section C.1 and Algorithm 2 for the details of our parallelization of the probability flow ODE formulation of diffusion models. PIADM-ODE keeps the discretization scheme detailed in Section 3.1.1 that divides the time horizon T into N blocks and uses exponential integrators for all updating rules. Notably, PIADM-ODE has the following distinctions compared with PIADM-SDE (Algorithm 1):

Instead of applying Picard iteration to the backward SDE as in (3.2), we apply Picard iteration to the probability flow ODE as in (C.3) within each block, which does not require sampling i.i.d. Gaussians to simulate a Wiener process;

The most significant difference is the adoption of an additional corrector step [111] after running the probability flow ODE with Picard iteration within one block. During the corrector step, one augments the state space with a Gaussian that represents the initial momentum and then simulates an underdamped Langevin dynamics for O (1) time with the learned NN-based score function at the time of the block end;

We then further parallelize the underdamped Langevin dynamics in the corrector step so that it can also be accomplished with O (log d ) approximate time complexity, as a naïve implementation would result in ˜ O ( √ d ) [130], which is incompatible with our desired poly-logarithmic guarantee.

3.2.2 Assumptions

Due to technicalities specific to this implementation, we need first to modify Assumption 3.1 and add assumption on the Lipschitzness of the true score functions ∇ log p t , which is a common practice in related literature [104, 111]. Recent work on the probability flow ODE implementation [112, 114] also adopts stronger assumptions compared to the SDE implementation.

Assumption 3.1' ( L ∞ ([0 , t N ]) δ -accurate learned score) . For any n ∈ [0 : N -1] and m ∈ [0 : M n -1] , the learned NN-based score s θ t n ,τ n,m is δ ∞ -accurate in the sense of

Assumption 3.4 (Bounded and Lipschitz true score) . The true score function ∇ log p t has bounded C 1 norm, i.e. ‖‖∇ log p t ( · ) ‖‖ L ∞ ([0 ,T ]) ≤ M p with Lipschitz constant L p .

Further relaxations on Assumption 3.4 to time-dependent assumptions accommodating the blow-up to the data end ( e.g. [103, Assumption 1.5]) are left for further work.

3.2.3 Theoretical Guarantees

Our results for PIADM-ODE are summarized in the following theorem:

Theorem 3.5 (Theoretical Guarantees for PIADM-ODE) . Under Assumptions 3.1', 3.2, 3.3, and 3.4, given the following choices of the order of the parameters

T = O (log( dδ -2 )) , h = Θ(1) , N = O (log( dδ -2 )) , ϵ = Θ ( d -1 / 2 δ log -1 ( d -1 / 2 δ -1 ) ) , M = O ( d 1 / 2 δ -1 log( d 1 / 2 δ -1 )) , K = ˜ O (log( dδ -2 )) , for the outer iteration and

for the inner iteration during the corrector step, and let L 2 s h 2 e h ∨ L 2 s h † 2 e h † /γ glyph[lessmuch] 1 , δ ∞ ≲ δ log -1 ( dδ -2 ) , and γ ≳ L 1 / 2 p , then the distribution ̂ q t N that PIADM-ODE (Algorithm 2) generates samples from satisfies the following error bound:

˜ The reduction of space complexity by the probability flow ODE implementation is intuitively owing to the fact that the probability flow ODE process is a deterministic process in time rather than a stochastic process as in the SDE implementation, getting rid of the O ( ϵ ) term derived by Itô's symmetry. This allows the discretization error to be bounded with O ( ϵ 2 ) instead ( cf. Lemma B.7 and C.5).

TV( p η , ̂ q t N ) 2 ≲ de -T + dϵ 2 T 2 +( T 2 + N 2 ) δ 2 ∞ + dN 2 e -K ≲ δ 2 , with a total of ( K + K † N † ) N = ˜ O ( log 2 ( dδ -2 ) ) approximate time complexity and d ( M ∨ M † ) = Θ ( d 3 / 2 δ -1 ) space complexity for parallalizable δ -accurate score function computations.

3.2.4 Proof Sketch

The details of the proof of Theorem 3.5 are provided in Section C. Along with the complexity benefits the deterministic nature of the probability flow ODE may bring, the analysis is technically more involved than that of Theorem 3.3 and requires an intricate interplay between statistical distances. Several major challenges and our corresponding solutions are summarized below:

The error of the parallelized algorithm within each block may now only be bounded by 2Wasserstein distance ( cf. Theorem C.7) instead of any f -divergence that admits data processing inequality as in the SDE case by Girsanov's theorem. The additional corrector step exactly handles this issue and would intuitively translate 2-Wasserstein proximity to TV distance proximity ( cf. Lemma C.18), allowing the decomposition of the overall error into each block;

For the corrector step, the underdamped Langevin dynamics as a second-order dynamics requires only O ( √ d ) steps to converge, instead of O ( d ) steps in its overdamped counterpart. We then adapt the parallelization technique mentioned in Section 2.2 to conclude that it can be accomplished with O (log d ) approximate time complexity ( cf. Theorem C.17). The error caused by the approximation to the true score and numerical discretization within this step is bounded in KL divergence by invoking Girsanov's theorem(Theorem A.4) as in the proof of Theorem 3.3;

Different from the SDE case, where the chain rule of KL divergence can easily decouple the initial distribution and the subsequent dynamics, we need several interpolating processes between the implementation and the true backward process in this case. The final guarantee is in TV distance as it connects with the KL divergence via Pinsker's inequality and admits data processing inequality. We refer the readers to Figure 2 for an overview of the proof pipeline, as well as the notations and intuitions of the auxiliary and interpolating processes appearing in the proof.

4 Discussion and Conclusion

In this work, we have proposed novel parallelization strategies for the inference of diffusion models in both the SDE and probability flow ODE implementations. Our algorithms, namely PIADMSDE and PIADM-ODE, are meticulously designed and rigorously proved to achieve ˜ O (poly log d ) approximate time complexity and ˜ O ( d 2 ) and ˜ O ( d 3 / 2 ) space complexity, respectively, marking the first inference algorithm of diffusion and probability flow based models with sub-linear approximate time complexity. Our algorithm intuitively divides the time horizon into several O (1) blocks and applies Picard iteration within each block in parallel, transferring the time complexity into space complexity. Our analysis is built on a sophisticated mathematical framework of stochastic processes and provides deeper insights into the mathematical theory of diffusion models.

Our findings echo and corroborate the recent empirical work [1, 91-94] that parallel sampling techniques significantly accelerate the inference process of diffusion models. Theoretical exploration of the adaptive block window scheme therein presents an interesting future research potential. Possible future work also includes the investigation of how to apply our parallelization framework to other variants of diffusion models, such as the discrete [23, 132-142] and multi-marginal [143] formulations. Although we anticipate implementing diffusion models in parallel may introduce engineering challenges, e.g. scalability, hardware compatibility, memory bandwidth, etc. , we believe that our theoretical contributions lay a solid foundation that not only supports but also motivates the empirical development of parallel inference algorithms for diffusion models since advancements continue in GPU power and memory efficiency.

Acknowledgments and Disclosure of Funding

LY acknowledges the support of the National Science Foundation under Grant No. DMS-2011699 and DMS-2208163. GMR is supported by a Google Research Scholar Award.

References

Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. Advances in Neural Information Processing Systems , 36, 2024.

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023.

Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571 , 2022.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 , 2022.

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

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems , 34: 1415-1428, 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Linfeng Zhang, Weinan E, and Lei Wang. Monge-ampère flow for generative modeling. arXiv preprint arXiv:1809.10188 , 2018.

Omer Bar-Tal, Lior Yariv, Yaron Lipman, and Tali Dekel. Multidiffusion: Fusing diffusion paths for controlled image generation. In International Conference on Machine Learning , pages 1737-1752. Pmlr, 2023.

Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. Deconstructing denoising diffusion models for self-supervised learning. arXiv preprint arXiv:2401.14404 , 2024.

Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research , 23(47):1-33, 2022.

Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint arXiv:2401.08740 , 2024.

Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2108.01073 , 2021.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning , pages 8821-8831. Pmlr, 2021.

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

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022.

Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint arXiv:2111.08005 , 2021.

Yu Sun, Zihui Wu, Yifan Chen, Berthy T Feng, and Katherine L Bouman. Provable probabilistic imaging using score-based generative priors. arXiv preprint arXiv:2310.10835 , 2023.

Xingyu Xu and Yuejie Chi. Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction. arXiv preprint arXiv:2403.17042 , 2024.

Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-lm improves controllable text generation. Advances in Neural Information Processing Systems , 35:4328-4343, 2022.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems , 35:8633-8646, 2022.

Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli S Sastry, Siddharth Gururani, Sageev Oore, and Yisong Yue. Symbolic music generation with non-differentiable rule guided diffusion. arXiv preprint arXiv:2402.14285 , 2024.

Gautam Mittal, Jesse Engel, Curtis Hawthorne, and Ian Simon. Symbolic music generation with diffusion models. arXiv preprint arXiv:2103.16091 , 2021.

Flavio Schneider. Archisound: Audio generation with diffusion. arXiv preprint arXiv:2301.13267 , 2023.

Ruihan Yang, Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for video generation. Entropy , 25(10):1469, 2023.

Zihao Li, Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Yinyu Ye, Minshuo Chen, and Mengdi Wang. Diffusion model for data-driven black-box optimization. arXiv preprint arXiv:2403.13219 , 2024.

Chen Xu, Jonghyeok Lee, Xiuyuan Cheng, and Yao Xie. Flow-based distributionally robust optimization. IEEE Journal on Selected Areas in Information Theory , 2024.

Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from mean-field gibbs measures via diffusion processes. arXiv preprint arXiv:2310.08912 , 2023.

Sitan Chen, Vasilis Kontonis, and Kulin Shah. Learning general gaussian mixtures with efficient score matching. arXiv preprint arXiv:2404.18893 , 2024.

Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from the sherringtonkirkpatrick gibbs measure via algorithmic stochastic localization. In 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS) , pages 323-334. IEEE, 2022.

Khashayar Gatmiry, Jonathan Kelner, and Holden Lee. Learning mixtures of gaussians using diffusion models. arXiv preprint arXiv:2404.18869 , 2024.

Ye He, Kevin Rojas, and Molei Tao. Zeroth-order sampling methods for non-log-concave distributions: Alleviating metastability by denoising diffusion. arXiv preprint arXiv:2402.17886 , 2024.

Xunpeng Huang, Hanze Dong, HAO Yifan, Yian Ma, and Tong Zhang. Reverse diffusion monte carlo. In The Twelfth International Conference on Learning Representations , 2023.

Brice Huang, Andrea Montanari, and Huy Tuan Pham. Sampling from spherical spin glasses in total variation via algorithmic stochastic localization. arXiv preprint arXiv:2404.15651 , 2024.

Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309.11420 , 2023.

Andrea Montanari. Sampling, diffusions, and stochastic localization. arXiv preprint arXiv:2305.10690 , 2023.

Andrea Montanari and Yuchen Wu. Posterior sampling from the spiked models via diffusion processes. arXiv preprint arXiv:2304.11449 , 2023.

Nicholas M Boffi and Eric Vanden-Eijnden. Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology , 4(3):035012, 2023.

Yan Huang and Li Wang. A score-based particle method for homogeneous landau equation. arXiv preprint arXiv:2405.05187 , 2024.

Lingxiao Li, Samuel Hurault, and Justin M Solomon. Self-consistent velocity matching of probability flows. Advances in Neural Information Processing Systems , 36, 2024.

Jianfeng Lu, Yue Wu, and Yang Xiang. Score-based transport modeling for mean-field fokkerplanck equations. Journal of Computational Physics , 503:112859, 2024.

Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of fokker-planck equations through gradient-log-density estimation. Entropy , 22(8):802, 2020.

Amira Alakhdar, Barnabas Poczos, and Newell Washburn. Diffusion models in de novo drug design. Journal of Chemical Information and Modeling , 2024.

Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04997 , 2024.

Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.05841 , 2024.

Pavel Avdeyev, Chenlai Shi, Yuhao Tan, Kseniia Dudnyk, and Jian Zhou. Dirichlet diffusion score model for biological sequence generation. In International Conference on Machine Learning , pages 1276-1301. PMLR, 2023.

Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X Lu, Nicolo Fusi, Ava P Amini, and Kevin K Yang. Protein generation with evolutionary diffusion: sequence is all you need. BioRxiv , pages 2023-09, 2023.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature , 620(7976):1089-1100, 2023.

Jordan Cotler and Semon Rezchikov. Renormalizing diffusion models. arXiv preprint arXiv:2308.12355 , 2023.

Florian Fürrutter, Gorka Muñoz-Gil, and Hans J Briegel. Quantum circuit synthesis with diffusion models. arXiv preprint arXiv:2311.02041 , 2023.

Zhiye Guo, Jian Liu, Yanli Wang, Mengrui Chen, Duolin Wang, Dong Xu, and Jianlin Cheng. Diffusion models in bioinformatics and computational biology. Nature reviews bioengineering , 2(2):136-154, 2024.

Artan Sheshmani, Yi-Zhuang You, Baturalp Buyukates, Amir Ziashahabi, and Salman Avestimehr. Renormalization group flow, optimal transport and diffusion-based generative model. arXiv preprint arXiv:2402.17090 , 2024.

Luke Triplett and Jianfeng Lu. Diffusion methods for generating transition paths. arXiv preprint arXiv:2309.10276 , 2023.

Lingxiao Wang, Gert Aarts, and Kai Zhou. Generative diffusion models for lattice field theory. arXiv preprint arXiv:2311.03578 , 2023.

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923 , 2022.

Mengjiao Yang, KwangHwan Cho, Amil Merchant, Pieter Abbeel, Dale Schuurmans, Igor Mordatch, and Ekin Dogus Cubuk. Scalable diffusion for materials generation. arXiv preprint arXiv:2311.09235 , 2023.

Jiahao Fan, Ziyao Li, Eric Alcaide, Guolin Ke, Huaqing Huang, and Weinan E. Accurate conformation sampling via protein structural diffusion. Journal of Chemical Information and Modeling , 2024.

Jason Yim, Hannes Stärk, Gabriele Corso, Bowen Jing, Regina Barzilay, and Tommi S Jaakkola. Diffusion models in protein structure and docking. Wiley Interdisciplinary Reviews: Computational Molecular Science , 14(2):e1711, 2024.

Yuchen Zhu, Tianrong Chen, Evangelos A Theodorou, Xie Chen, and Molei Tao. Quantum state generation with structure-preserving diffusion model. arXiv preprint arXiv:2404.06336 , 2024.

Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys , 56(4):1-39, 2023.

Stanley H Chan. Tutorial on diffusion models for imaging and vision. arXiv preprint arXiv:2403.18103 , 2024.

Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. Opportunities and challenges of diffusion models for generative ai. National Science Review , page nwae348, 2024.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. Advances in Neural Information Processing Systems , 35:30150-30166, 2022.

Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably. arXiv preprint arXiv:2403.03852 , 2024.

Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778 , 2022.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems , 35:5775-5787, 2022.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpmsolver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095 , 2022.

Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Dpm-solver-v3: Improved diffusion ode solver with empirical model statistics. Advances in Neural Information Processing Systems , 36:55502-55542, 2023.

Zhenyu Zhou, Defang Chen, Can Wang, and Chun Chen. Fast ode-based sampling for diffusion models in around 5 steps. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 7777-7786, 2024.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems , 35: 26565-26577, 2022.

Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204.13902 , 2022.

Saravanan Kandasamy and Dheeraj Nagaraj. The poisson midpoint method for langevin dynamics: Provably efficient discretization for diffusion models. arXiv preprint arXiv:2405.17068 , 2024.

Yuchen Wu, Yuxin Chen, and Yuting Wei. Stochastic runge-kutta methods: Provable acceleration of diffusion models. arXiv preprint arXiv:2410.04760 , 2024.

Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped langevin diffusion. arXiv preprint arXiv:2112.07068 , 2021.

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

Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning , pages 42390-42402. PMLR, 2023.

Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. Advances in Neural Information Processing Systems , 36:76806-76838, 2023.

Jonathan Heek, Emiel Hoogeboom, and Tim Salimans. Multistep consistency models. arXiv preprint arXiv:2403.06807 , 2024.

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

Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310.14189 , 2023.

Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081 , 2024.

Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 , 2021.

Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14297-14306, 2023.

Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512 , 2022.

Hyungjin Chung, Jeongsol Kim, Sehui Kim, and Jong Chul Ye. Parallel diffusion models of operator and image for blind inverse problems. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 6059-6069, 2023.

Zhiwei Tang, Jiasheng Tang, Hao Luo, Fan Wang, and Tsung-Hui Chang. Accelerating parallel sampling of diffusion models. arXiv preprint arXiv:2402.09970 , 2024.

Jiezhang Cao, Yue Shi, Kai Zhang, Yulun Zhang, Radu Timofte, and Luc Van Gool. Deep equilibrium diffusion restoration with parallel sampling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2824-2834, 2024.

Nikil Roashan Selvam, Amil Merchant, and Stefano Ermon. Self-refining diffusion samplers: Enabling parallelization via parareal iterations. arXiv preprint arXiv:2412.08292 , 2024.

Shivam Gupta, Linda Cai, and Sitan Chen. Faster diffusion-based sampling with randomized midpoints: Sequential and parallel. arXiv preprint arXiv:2406.00924 , 2024.

Gen Li and Yuchen Jiao. Improved convergence rate for diffusion probabilistic models. arXiv preprint arXiv:2410.13738 , 2024.

Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems , 32, 2019.

Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory , pages 30843114. PMLR, 2019.

Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto-encoders and langevin sampling. arXiv preprint arXiv:2002.00107 , 2020.

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

Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems , 35:2287022882, 2022.

Francesco Pedrotti, Jan Maas, and Marco Mondelli. Improved convergence of score-based diffusion models via prediction-correction. arXiv preprint arXiv:2305.14164 , 2023.

Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning , pages 4462-4484. PMLR, 2023.

Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning , pages 4735-4763. PMLR, 2023.

Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory , pages 946-985. PMLR, 2023.

Sokhna Diarra Mbacke and Omar Rivasplata. A note on the convergence of denoising diffusion probabilistic models. arXiv preprint arXiv:2312.05989 , 2023.

Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Linear convergence bounds for diffusion models via stochastic localization. arXiv preprint arXiv:2308.03686 , 2023.

Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv preprint arXiv:2306.09251 , 2023.

Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards non-asymptotic convergence for diffusion-based generative models. In The Twelfth International Conference on Learning Representations , 2024.

Gen Li, Yuting Wei, Yuejie Chi, and Yuxin Chen. A sharp convergence theory for the probability flow odes of diffusion models. arXiv preprint arXiv:2408.02320 , 2024.

Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast. Advances in Neural Information Processing Systems , 36, 2024.

Xuefeng Gao and Lingjiong Zhu. Convergence analysis for general probability flow odes of diffusion models in wasserstein distances. arXiv preprint arXiv:2401.17958 , 2024.

Yuchen Liang, Peizhong Ju, Yingbin Liang, and Ness Shroff. Non-asymptotic convergence of discrete-time diffusion models: New approach and improved rate. arXiv preprint arXiv:2402.13901 , 2024.

Daniel Zhengyu Huang, Jiaoyang Huang, and Zhengjiang Lin. Convergence analysis of probability flow ode for score-based generative models. arXiv preprint arXiv:2404.09730 , 2024.

Gen Li and Yuling Yan. o ( d/t ) convergence theory for diffusion probabilistic models under minimal assumptions. arXiv preprint arXiv:2409.18959 , 2024.

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

Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

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

Yu Cao, Jingrun Chen, Yixin Luo, and Xiang Zhou. Exploring the optimal choice for generative processes in diffusion models: Ordinary vs stochastic differential equations. Advances in Neural Information Processing Systems , 36, 2024.

Teo Deveney, Jan Stanczuk, Lisa Maria Kreusser, Chris Budd, and Carola-Bibiane Schönlieb. Closing the ode-sde gap in score-based diffusion models through the fokker-planck equation. arXiv preprint arXiv:2311.15996 , 2023.

Shen Nie, Hanzhong Allan Guo, Cheng Lu, Yuhao Zhou, Chenyu Zheng, and Chongxuan Li. The blessing of randomness: Sde beats ode in general diffusion-based image editing. arXiv preprint arXiv:2311.01410 , 2023.

Charles J Geyer. Markov chain monte carlo maximum likelihood. In E. M. Kernamidas, editor, Computing Science and Statistics: Proc. 23rd Symposium on the Interface , pages 156-163. Interface Foundation of North America, 1991.

Koji Hukushima and Koji Nemoto. Exchange monte carlo method and application to spin glass simulations. Journal of the Physical Society of Japan , 65(6):1604-1608, 1996.

Faming Liang. Use of sequential structure in simulation from high-dimensional systems. Physical Review E , 67(5):056101, 2003.

Nima Anari, Yizhi Huang, Tianyu Liu, Thuy-Duong Vuong, Brian Xu, and Katherine Yu. Parallel discrete sampling via continuous walks. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing , pages 103-116, 2023.

Holden Lee. Parallelising glauber dynamics. arXiv preprint arXiv:2307.07131 , 2023.

Lu Yu and Arnak Dalalyana. Parallelized midpoint randomization for langevin monte carlo. arXiv preprint arXiv:2402.14434 , 2024.

Ernest Lindelöf. Sur lapplication de la méthode des approximations successives aux équations différentielles ordinaires du premier ordre. Comptes rendus hebdomadaires des séances de lAcadémie des sciences , 116(3):454-457, 1894.

Emile Picard. Sur les méthodes d'approximations successives dans la théorie des équations différentielles. American Journal of Mathematics , pages 87-100, 1898.

Nima Anari, Sinho Chewi, and Thuy-Duong Vuong. Fast parallel sampling under isoperimetry. arXiv preprint arXiv:2401.09016 , 2024.

Bernt Oksendal. Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2110.02037 , 2021.

Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems , 34:12454-12465, 2021.

Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data. Advances in Neural Information Processing Systems , 35:34532-34545, 2022.

Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuoustime discrete diffusion models. arXiv preprint arXiv:2211.16750 , 2022.

Pierre H Richemond, Sander Dieleman, and Arnaud Doucet. Categorical sdes with simplex diffusion. arXiv preprint arXiv:2210.14784 , 2022.

Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834 , 2023.

Griffin Floto, Thorsteinn Jonsson, Mihai Nica, Scott Sanner, and Eric Zhengyu Zhu. Diffusion on the probability simplex. arXiv preprint arXiv:2309.02530 , 2023.

Javier E Santos, Zachary R Fox, Nicholas Lubbers, and Yen Ting Lin. Blackout diffusion: generative diffusion models in discrete-state spaces. In International Conference on Machine Learning , pages 9034-9059. PMLR, 2023.

Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising markov models. Journal of the Royal Statistical Society Series B: Statistical Methodology , 86(2):286-301, 2024.

Hongrui Chen and Lexing Ying. Convergence analysis of discrete diffusion model: Exact implementation through uniformization. arXiv preprint arXiv:2402.08095 , 2024.

Yinuo Ren, Haoxuan Chen, Grant M Rotskoff, and Lexing Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via a stochastic integral framework. arXiv preprint arXiv:2410.03601 , 2024.

Michael S Albergo, Nicholas M Boffi, Michael Lindsey, and Eric Vanden-Eijnden. Multimarginal generative modeling with stochastic interpolants. arXiv preprint arXiv:2310.03695 , 2023.

Jean-François Le Gall. Brownian motion, martingales, and stochastic calculus . Springer, 2016.

Arnaud Guillin and Feng-Yu Wang. Degenerate fokker-planck equations: Bismut formula, gradient estimate and harnack inequality. Journal of Differential Equations , 253(1):20-40, 2012.

A Mathematical Background

In this section, we will summarize used notations and rigorous mathematical framework of Itô processes as necessary in the proofs. We will also present several technical lemmas for later reference.

A.1 Notations

We adopt the following notations throughout the paper:

A.2 Preliminaries

Theorem A.1 (Properties of f -divergence) . Suppose p and q are two probability measures on a common measurable space (Ω , F ) with p glyph[lessmuch] q . The f -divergence between p and q is defined as where d p d q is the Radon-Nikodym derivative of p with respect to q , and f : R + → R is a convex function. In particular, D f ( · ‖ · ) coincides with the KullbackLeibler (KL) divergence when f ( x ) = x log x and D f ( · ‖ · ) = TV coincides with the total variation (TV) distance when f ( x ) = 1 2 | x -1 | . For the f -divergence defined above, we have the following properties:

(Data-processing inequality). Suppose H is a subσ -algebra of F , the following inequality holds

for any f -divergence D f ( ·‖· ) .

(Chain rule). Suppose X is a random variable generating a subσ -algebra F X of F , and p ( ·| X ) glyph[lessmuch] q ( ·| X ) holds for any value of X , then

In this paper, we consider a probability space (Ω , F , p ) on which ( w t ( ω )) t ≥ 0 is a Wiener process in R d . The Wiener process ( w t ( ω )) t ≥ 0 generates the filtration {F t } t ≥ 0 on the measurable space (Ω , F ) . For an Itô process z t ( ω ) with the following governing SDE:

for any time t , we denote the marginal distribution of z t by p t , i.e.

as well as the path measure of the process z t in the sense of

For the sake of simplicity, we define the following class of functions:

Definition A.2. For any 0 ≤ t 1 < t 2 , we define V ( t 1 , t 2 ) as the class of functions f ( t, ω ) : [0 , + ∞ ) × Ω → R such that

f ( t, ω ) is B × F -measurable, where B is the Borel σ -algebra on R d ;

f ( t, ω ) is F t -adapted for all t ≥ 0 ;

The following Novikov condition holds

and V = ∩ t> 0 V (0 , t ) . For vectors and matrices, we say it belongs to V n ( t, ω ) or V m × n ( t, ω ) if each component of the vector or each entry of the matrix belongs to V ( t, ω ) .

Remark A.3. Novikov's condition appeared in the third requirement is often relaxed to the squared integrability condition in the general definition of Itô processes, which requires

Here, we adopt the more restricted condition in the spirit of its necessity for Girsanov's theorem to hold, as we shall see later.

Similar to previous work [111], here we can avoid checking Novikov's condition throughout our proofs below by using the approximation argument presented in [100]. A review of Girsanov can be found in textbooks like in [131, 144]. We will present the following generalized version of Girsanov's theorem:

Theorem A.4 (Girsanov's Theorem [131, Theorem 8.6.6]) . Let α ( t, ω ) ∈ V m , Σ ( t, ω ) ∈ V m × n , and ( w t ( ω )) t ≥ 0 be a Wiener process on the probability space (Ω , F , q ) . For t ∈ [0 , T ] , suppose z t ( ω ) is an Itô process with the following SDE:

and there exist processes δ ( t, ω ) ∈ V n and β ( t, ω ) ∈ V m such that

Σ ( t, ω ) δ ( t, ω ) = α ( t, ω ) -β ( t, ω ) ;

The process M t ( ω ) as defined below is a martingale with respect to the filtration {F t } t ≥ 0 and probability measure q :

then there exists another probability measure p on (Ω , F ) such that

p glyph[lessmuch] q with the Radon-Nikodym derivative d p d q ( ω ) = M T ( ω ) ,

The process ˜ w t ( ω ) as defined below is a Wiener process on (Ω , F , p ) :

Any continuous path in C ([ t 1 , t 2 ] , R m ) generated by the process z t satisfies the following SDE under the probability measure p :

d ˜ z t ( ω ) = β ( t, ω )d t + Σ ( t, ω )d ˜ w t ( ω ) . (A.3) Corollary A.5. Suppose the conditions in Theorem A.4 hold, then for any t 1 , t 2 ∈ [0 , T ] with t 1 < t 2 , the path measure of the SDE (A.3) under the probability measure p in the sense of p t 1 : t 2 = p ( z -1 t 1 : t 2 ( · ) ) is absolutely continuous with respect to the path measure of the SDE (A.2) in the sense of q t 1 : t 2 = q ( z -1 t 1 : t 2 ( · ) ) . Moreover, the KL divergence between the two path measures is given by

Proof. First, by Theorem A.1, we have

From Girsanov's theorem (Theorem A.4), we have that the measure p | F t 1 is absolutely continuous with respect to q | F t 1 , which allows us to compute the second term above as follows:

and therefore

which completes the proof.

A.3 Helper Lemmas

Lemma A.6 ([107, Lemma 2]) . For the backward process (2.3), we have for 0 ≤ s < t < T ,

Lemma A.7 ([107, Lemma 3]) . For the forward process (2.3), we have for 0 ≤ t < T , where the posterior covariance matrix Σ t := cov p 0 | t ( x 0 ) and σ 2 t = 1 -e -t . Moreover, the posterior covariance matrix Σ t satisfies

Lemma A.8. For any n ∈ [0 : N -1] and τ ∈ [0 , h n ] , under the assumption cov p 0 ( x 0 ) = I d , we have and

Proof. Conditioned on x 0 , we have that and

for any τ ∈ [0 , h n ] . Therefore, we have

Taking the difference between them then implies that for any τ ∈ [0 , h n ] ,

Lemma A.9 (Lemma 9 in [104]) . For ̂ q 0 ∼ N (0 , I d ) and ⃗ p 0 = p T is the distribution of the solution to the forward process (2.3), we have

B Details of SDE Implementation

In this section, we will present the missing proofs for Theorem 3.3. For readers' convenience, we reiterate the backward process (2.3)

and its approximate version (2.5) with the learned score function

The filtration F t refers to the filtration of the SDE (B.1) up to time t .

B.1 Auxiliary Process

We would like first to consider the errors that Algorithm 1 may cause within one block of update. To this end, we consider the following auxiliary process for τ ∈ [0 , h n ] conditioned on the filtration F t n at time t n :

Definition B.1 (Auxiliary Process) . For any n ∈ [0 : N -1] , we define the auxiliary process ( ̂ y ( k ) t n ,τ ) τ ∈ [0 ,h n ] as the solution to the following SDE recursively for k ∈ [0 : K -1] :

with the initial condition

̂ ̂ ̂ The iteration should be perceived as a deterministic procedure to each event ω ∈ Ω , i.e. each realization of the Wiener process ( w t ) t ≥ 0 . The following lemma clarifies this fact and proves the well-definedness and parallelability of the iteration in (B.2).

Lemma B.2. The auxiliary process ( ̂ y ( k ) t n ,τ ( ω )) τ ∈ [0 ,h n ] is F t n + τ -adapted for any k ∈ [0 : K ] and n ∈ [0 : N -1] .

Proof. Since the initialization ̂ y (0) t n ,τ ( ω ) ≡ ̂ y t n ( ω ) for τ ∈ [0 , h n ] , where ̂ y t n ( ω ) is F t n -adapted, it is obvious that ̂ y (0) t n ,τ ( ω ) is F t n + τ -adapted. Now suppose that ( ̂ y ( k ) t n ,τ ( ω )) τ ∈ [0 ,h n ] is F t n + τ -adapted, since g n ( τ ) ≤ τ , we have the following Itô integral well-defined and F t n + τ -adapted:

and therefore (B.2) has a unique strong solution ( ̂ y ( k +1) t n ,τ ( ω )) τ ∈ [0 ,h n ] that is also F t n + τ -adapted. The lemma follows by induction.

LemmaB.3 (Equivalence between (3.4) and (B.2)) . For any n ∈ [0 : N -1] , the update rule (3.4) in Algorithm 1 is equivalent to the exact solution of the auxiliary process (B.2) for any k ∈ [0 : K -1] and τ ∈ [0 , h n ] .

Proof. The dependency on ω will be omitted in the proof below.

Rewriting (B.2) and multiplying e -τ 2 on both sides yield

d [ e -τ 2 ̂ y ( k +1) t n ,τ ] = e -τ 2 [ d ̂ y ( k +1) t n ,τ -1 2 ̂ y ( k +1) t n ,τ d τ ] = e -τ 2 [ s θ t n + g n ( τ ) ( ̂ y ( k ) t n ,g n ( τ ) ) d τ +d w t n + τ ] . Integrating on both sides from 0 to τ implies

where, by Itô isometry, we have

for τ > τ n,m and equals to 0 otherwise. Plugging in τ = τ j,m gives us (3.4), as desired.

B.2 Errors within Block

We shall invoke Girsanov's theorem (Theorem A.4) in the procedure as detailed below:

Setting (A.2) in Theorem A.4 as the auxiliary process (B.2) at iteration K , where w t ( ω ) is a Wiener process under the measure q | F t n ;

Defining another process w t n + τ ( ω ) governed by the following SDE:

where

and computing the Radon-Nikodym derivative of the measure ⃗ p | F t n with respect to q | F t n as

Concluding that (B.2) at iteration K under the measure q | F t n satisfies the following SDE:

with ( ˜ w t n + τ ) τ ≥ 0 being a Wiener process under the measure ⃗ p | F t n . If we replace ̂ y ( K ) t n ,g n ( τ ) ( ω ) by ⃗ x t n + τ ( ω ) , one should notice (B.5) is immediately the original backward SDE (2.3) with the true score function on t ∈ [ t n , t n +1 ] :

Remark B.4. The applicability of Girsanov's theorem here relies on the F τ -adaptivity of s θ t n + g n ( τ ) ( ̂ y ( K -1) t n ,g n ( τ ) ( ω ) ) established by Lemma B.2. One should notice the change of measure procedure above depends on the number of iterations K , and different K would lead to different transform (B.4).

Then Corollary A.5 provides the following computation

where the first inequality is by the data-processing inequality (Theorem A.1). Now, the problem remaining is to bound the discrepancy quantified by

Before we continue our proof, we would like first to provide the following lemma bounding the behavior of the auxiliary process (B.2) when k = 0 for τ ∈ [0 , h n ] .

Lemma B.5. For any n ∈ [0 : N -1] , suppose the initialization ̂ y t n in (B.3) of the auxiliary process (B.2) follows the distribution of ⃗ x t n ∼ ⃗ p t n , then the following estimate holds

Proof. Let z t n ,τ = y (1) t n ,τ y (0) t n ,τ . For k = 0 , we can rewrite (B.2) as

By applying Itô's lemma and plugging in the expression of w t n + τ given by Theorem A.4, we have

where δ t n ( τ, ω ) is defined in (B.4). Similar to (B.7), we may use triangle inequality to upper bound ‖ δ t n ( τ, ω ) ‖ 2 , which implies that for any τ ∈ [0 , h n ]

where in the second inequality above, we have used the fact that s θ t ( · ) is L s -Lipschitz for any t . By Grönwall's inequality, we have that for any τ ∈ [0 , h n ]

Substituting (A.5) into (B.10) then yields that for any τ ∈ [0 , h n ]

By assumption, ̂ y (0) t n ,τ = ̂ y t n follows the distribution of ⃗ x t n ∼ ⃗ p t n , which allows us to bound the second moment of ̂ y t n for any n ∈ [0 : N ] by Lemma A.8:

Taking supremum with respect to τ ∈ [0 , h n ] on both sides above completes our proof.

As utilized in the proof of the existence of solutions of SDEs, the following lemma demonstrates the exponential convergence of the iteration defined in (B.2).

LemmaB.6 (Exponential convergence of Picard iteration in PIADM-SDE) . For any n ∈ [0 , N ] , suppose the initialization ̂ y t n in (B.3) of the auxiliary process (B.2) follows the distribution of ⃗ x t n ∼ ⃗ p t n , then the two ending terms ̂ y ( K ) t n ,τ and ̂ y ( K -1) t n ,τ of the sequence { ̂ y ( k ) t n ,τ } k ∈ [0: K -1] satisfy the following exponential convergence rate

Proof. For each ω ∈ Ω conditioned on the filtration F t n , subtracting (B.2) from the process as defined by

we have

where the diffusion term d w t n + τ cancels each other out. Now we may use the formula above to compute derivative d d τ ′ ∥ ∥ ∥̂ y ( k +1) t n ,τ ′ ( ω ) -̂ y ( k ) t n ,τ ′ ( ω ) ∥ ∥ ∥ 2 explicitly, integrate it from τ ′ = 0 to τ , and obtain the following inequality

Taking expectation on both sides above further implies that for any τ ∈ [0 , h n ] ,

Furthermore, we take supremum over τ ∈ [0 , h n ] on both sides above and iterate (B.12) over k ∈ N , which indicates

where the last inequality follows from Lemma B.5. By rearranging the inequality above, setting k = K -1 and using the assumption that L 2 s h n e 2 h n glyph[lessmuch] 1 , we obtain

as desired.

The following lemma from [107] bounds the expectation of the term A t n ( ω ) in (B.7): Lemma B.7 ([107, Section 3.1]) . We have

where η

Proof. Notice that

where for the last equality, we use the fact that the process ̂ y ( K ) t n ,τ ( ω ) follows the backward SDE with the true score function under the measure ⃗ p . In the following, we drop the superscript ω ∼ ⃗ p | F t n of the expectation for simplicity.

Now noticing that we further have

and thus

For n ∈ [0 , N -2] , we have M n ϵ = h n and thus E ω ∼ ⃗ p | F t n [ A t n ( ω )] ≲ ϵdh n , and for n = N -1 , we have

and thus E ω ∼ ⃗ p | F t n [ A t N -1 ( ω ) ] ≲ ϵ 2 dM n ≲ ϵd log η -1 .

B.3 Overall Error Bound

Proof of Theorem 3.3. We first continue the computation in (B.6) and (B.7):

Then plugging in the result of Lemma B.6, we have

where we used the assumption that L 2 s h n e 7 2 h n glyph[lessmuch] 1 .

The term ∑ N -1 n =0 E ω ∼ ⃗ p | F t n [ B t n ( ω )] is bounded by Assumption 3.1 as where the last equality is because the process ̂ y ( K ) t n ,τ ( ω ) under measure ⃗ p follows the backward SDE (B.5). Thus, by Theorem A.1 and plugging in the iteration relations above

as T ≳ log η -1 , h n ≲ 1 , and δ 2 ≲ δ , and then it is straightforward to see that the following choices of parameters would yield an overall error of O ( δ 2 ) .

C Details of Probability Flow ODE Implementation

In this section, we provide the details of the parallelized algorithm for the probability flow ODE formulation of diffusion models. We first introduce the algorithm and define the necessary notations, then discuss the error analysis during the predictor and corrector steps, respectively, and finally provide the proof of Theorem 3.5.

C.1 Algorithm

In the parallelized inference algorithm for diffusion models in the probability flow ODE formulation, we adopt the same discretization scheme as in Section 3.1.1 and the exponential integrator for all updating rules. For each block, we first run a predictor step , which consists of running the probability flow ODE in parallel. Then we run a corrector step , which runs an underdamped Langevin dynamics in parallel to correct the distribution of the samples. The algorithm is summarized In Algorithm 2.

Parallelized Predictor Step The parallelization strategies in the predictor step are similar to those in the SDE algorithm (Algorithm 1). The only difference here is that instead of applying Picard iteration to the backward SDE as in (3.2), we apply Picard iteration to the probability flow ODE as in (C.3), which does not require i.i.d. samples from standard Gaussian distribution. As shown in Lemma C.3, the update rule in the predictor step (C.1) in Algorithm 2 is equivalent to running

Algorithm 2: PIADM-ODE

the auxiliary predictor process (C.3). The auxiliary predictor process takes in the result from the previous corrector step (or the initialization if n = 0 ) and outputs ̂ y ( K ) t n ,h n as the initialization for the next corrector step.

We run the underdamped Langevin dynamics for time T † , which is set to be of order Ω(1) so that it is large enough to correct the distribution of the samples ( cf. Lemma C.18) while being comparably short to ensure numerical stability ( cf. Theorem C.17). Following a similar strategy as in Section 2.2 and in Algorithm 1, we further divide the time horizon T † into N † blocks with step size h † , and for each block the block length h † into M † steps with step size ϵ † . Within each block, we run the underdamped Langevin dynamics in parallel for K † iterations. As shown in Lemma C.9, the update rule in the corrector step (C.2) in Algorithm 2 is equivalent to running the auxiliary corrector process (C.11).

Parallelized Corrector Step The parallelization of the underdamped Langevin dynamics is similar to that mentioned in Section 2.2. Given a sample resulting from the predictor step, we initialize the auxiliary corrector process (Definition C.8) which is an underdamped Langevin dynamics with the initialization ̂ u t n , 0 = y ( K ) t n ,h n and the augmented variable ̂ v t n , 0 ∼ N (0 , I d ) representing the momentum.

In the following subsections, we proceed to provide theoretical guarantees for the algorithm.

C.2 Parallelized Predictor Step

Definition C.1 (Auxiliary Predictor Process) . For any n ∈ [0 : N -1] , we define the auxiliary predictor process ( ̂ y ( k ) t n ,τ ) τ ∈ [0 ,h n ] as the solution to the following ODE recursively for k ∈ [0 : K -1] :

with the initial condition

̂ y (0) t n ,τ ≡ ̂ y t n for τ ∈ [0 , h n ] , and y ( k ) t n , 0 ≡ ̂ y t n for k ∈ [1 : K ] (C.4) where ̂ y t n = ̂ u t n -1 ,N † h † if n ∈ [1 : N -1] and ̂ y t 0 ∼ N (0 , I d ) . We will also denote the probability distribution of ̂ y ( K ) t n ,τ as ̂ q t n ,τ .

Definition C.2 (Interpolating Process) . For any n ∈ [0 : N -1] , we define the interpolating process ( ˜ y ( k ) t n ,τ ) τ ∈ [0 ,h n ] as the solution to the following ODE recursively for k ∈ [0 : K -1] :

with initial condition

˜ y (0) t n ,τ ≡ ˜ y (0) t n , 0 for τ ∈ [0 , h n ] , and ˜ y ( k ) t n , 0 ≡ ˜ y (0) t n , 0 for k ∈ [1 : K ] , where ˜ y (0) t n , 0 ∼ ⃗ p t n . We will also denote the probability distribution of ˜ y ( K ) t n ,τ as ˜ q t n ,τ . Similar to the equivalence between (3.4) and (B.2), we have the following lemma:

Lemma C.3 (Equivalence between (C.1) and (C.3)) . For any n ∈ [0 : N -1] , the update rule (C.1) in Algorithm 2 is equivalent to the exact solution of (C.3) for any k ∈ [0 : K -1] and τ ∈ [0 , h n ] .

Proof. Rewriting (C.3) and multiplying e -τ 2 on both sides yield d [ e -τ 2 ̂ y ( k +1) t n ,τ ] = e -τ 2 [ d ̂ y ( k +1) t n ,τ -1 2 ̂ y ( k +1) t n ,τ d τ ] = e -τ 2 2 s θ t n + g n ( τ ) ( ̂ y ( k ) t n ,g n ( τ ) ) d τ Integrating on both sides from 0 to τ implies

and then multiplying e τ 2 on both sides above yields

Lemma C.4 (Error between the interpolating process and the true process) . Under the Picard iteration, we have that the ending process { ̂ y ( K ) t n ,τ } τ ∈ [0 ,h n ] satisfies the following exponential convergence rate

Proof. Recall that the backward true process { ⃗ x t n + τ } τ ∈ [0 ,h n ] satisfies the following backward SDE within one block

By subtracting (C.6) from (C.5), we obtain that

Then by and integrating for τ ′ ∈ [0 , h n ] , we have

Using AM-GM inequality and taking expectations on both sides, we further upper bound the summation above as

where the last equality is by Assumption 3.1'.

Applying Grönwall's inequality, we have

and by taking supremum

Given that constant h n is sufficiently small, which ensures L 2 s h n e 5 2 h n glyph[lessmuch] 1 , iterating the above inequality for k ∈ [0 : K -1] gives us that

substituting which into (C.9) then gives us that

Now it remains to bound C t n and D t n in Lemma C.4. We first bound D t n using the following lemma:

Lemma C.5. For any n ∈ [0 : N -1] , we have that

For n = N -1 , notice that by the step size schedule ( cf. Section 3.1.1) and suppose ϵ ≤ 1 / 2 , we have

and then again [111, Corollary 1] states and thus

Remark C.6. The above lemma is able to achieve a better dependency on ϵ compared to Lemma B.7, because the backward process ( ⃗ x t ) t ∈ [0 ,T ] is now a deterministic process in the probability flow ODE formulation, instead of a stochastic process as in the SDE formulation as in Lemma B.7. Thus, intuitively applying Cauchy-Schwarz rather than Itô symmetry gives us a O ( ϵ 2 ) -dependency rather than O ( ϵ ) -dependency.

Theorem C.7. Under Assumptions 3.1', 3.2, 3.3, and 3.4, then the distribution ˜ q t n ,h n that the parallelized predictor step generates samples from satisfies the following error bound:

for n ∈ [0 : N -1] .

Proof. By the definition of 2-Wasserstein distance, we have for any coupling of ˜ y ( K ) t n ,h n and ⃗ x t n + h n , and therefore

where for the second to last inequality we used Lemma C.4, the last inequality is due to Lemma C.5 and the assumption h 2 n e h n L 2 s glyph[lessmuch] 1 .

C.3 Parallelized Corrector Step

After each predictor step, we run the corrector step for O (1) time to reduce the error. Particularly, we apply the Parallelized underdamped Langevin dynamics algorithm [130] to the corrector step, which yields O (1) approximate time complexity compared to the ordinary implementation of the ULMC dynamics as in [111]. In the following, we will drop the dependency on ω for notational simplicity, and we refer readers to Appendix A.2 and B.2 to review the change of measure arguments and the application of Girsanov's theorem A.4. We will also use a general notation ∗ † to distinguish the time in the backward process and the inner time in the corrector step of the n -th block.

We first define the true underdamped Langevin dynamics ( u t n ,t † , v t n ,t † ) t ≥ 0 :

with initial condition u t n , 0 ≡ ˜ y ( K † ) t n ,h n from the predictor step and v t n , 0 ∼ N (0 , I d ) , where ( b t n ,t † ) t ≥ 0 is a Wiener process. We may also write the system of SDEs above in the following matrix form:

We run this underdamped Langevin dynamics until the pre-determined time horizon T † . We also define the joint probability distribution of ( u t n ,t † , v t n ,t † ) at time t as π t n ,t † ( u t n ,t † , v t n ,t † ) and its marginal on u t n ,t † as π u t n ,t † ( u t n ,t † ) .

Similar to the parallelizing strategy in Section 3.1.1, we discretize the time interval [0 , T † ] into N † blocks with length h † = T † /N † . Within the n -th block, we further divide the block [ n † h † , ( n +1) h † ] into M † steps, each with step size ϵ † = h † /M † .

Definition C.8 (Auxiliary corrector process) . For any n † ∈ [0 : N † -1] , we define the auxiliary corrector process ( ̂ u ( k † ) t n ,n † h † ,τ † ) τ † ∈ [0 ,h † ] as the solution to the following SDE recursively for k † ∈ [0 : K † -1] :

where

for n † ∈ [1 : N † -1] , and

̂ u t n , 0 = y ( K ) t n ,h n , ̂ v t n , 0 ∼ N (0 , I d ) . We define the joint probability distribution of ( ̂ u t n ,t † , ̂ v t n ,t † ) at time t as ̂ π t n ,t † ( ̂ u t n ,t † , ̂ v t n ,t † ) and its marginal on ̂ u t n ,t † as ̂ π ̂ u t n ,t † ( ̂ u t n ,t † ) . We will also denote the resulting probability distribution of π ̂ u t n ,T † as q t n +1 .

̂ ̂ Lemma C.9 (Equivalence between (C.2) and (C.11)) . For any n † ∈ [0 : N † -1] , the update rule in Algorithm 2 is equivalent to the exact solution of the auxiliary process (C.11) for any k † ∈ [0 : K † -1] and τ † ∈ [0 , h † ] .

Proof. Without loss of generality, we will prove the lemma for m † = M † . The proof for m † ∈ [0 : M † -1] can be done similarly.

We first rewrite (C.2) into the matrix form:

Define the time-dependent matrix G ( · ) as

satisfying that

Now we multiply G ( -τ † ) on both sides of (C.13) to obtain:

Integrating on both sides from 0 to h † and multiplying G ( h † ) on both sides, we have

By Itô isometry, we have

as desired

Definition C.10 (Interpolating corrector process) . For any n † ∈ [0 : N † -1] , we define the interpolating corrector process ( ̂ u ( k † ) t n ,n † h † ,τ † ) τ † ∈ [0 ,h † ] as the solution to the following SDE recursively for k † ∈ [0 : K † -1] :

where

for n † ∈ [1 : N † -1] , and

We define the joint probability distribution of ( ˜ u t n ,t † , ˜ v t n ,t † ) at time t as ˜ π t n ,t † ( ˜ u t n ,t † , ˜ v t n ,t † ) and its marginal on ˜ u t n ,t † as ˜ π ˜ u t n ,t † ( ˜ u t n ,t † ) . We invoke Girsanov's theorem (Theorem A.4) again by the following procedure

Setting (A.2) as the auxiliary process (C.15) at iteration K † , where b t n ,t † ( ω ) is a Wiener process under the measure Q ;

Defining another process b t n ,n † h † + τ † governed by the following SDE:

where

ϕ t n ,n † h † ( τ † ) = 1 √ 2 γ ( s t n +1 ( ˜ u ( K † -1) t n ,n † h † , ⌊ τ † ϵ † ⌋ ϵ † ) -∇ log ⃗ p t n +1 ( ˜ u ( K † ) t n ,n † h † ,τ † ) ) (C.18) and computing the Radon-Nikodym derivative of the measure P with respect to Q as

Concluding that (C.15) at iteration K † under the measure Q satisfies the following SDE:

with ( ˜ b t n ,n † h † + τ † ) τ † ≥ 0 being a Wiener process under the measure P . If we replace ( ˜ u ( K † ) n † h † ,τ † , ˜ v ( K † ) n † h † ,τ † ) by ( u t n ,n † h † + τ † , v t n ,n † h † + τ † ) , one should notice (C.20) is immediately the original backward SDE (C.10) with the true score function on t ∈ [ n † h † , ( n +1) h † ] :

We further define the joint probability distribution of ( u t n ,t † , v t n ,t † ) at time t as π t n ,t † ( u t n ,t † , v t n ,t † ) and its marginal on u t n ,t † as π u t n ,t † ( u t n ,t † ) .

Remark C.11. The application of Girsanov's theorem A.4 is by writing the system of SDEs in the matrix form.

Definition C.12 (Stationary process) . Under the P -measure that is defined by the Radon-Nikodym derivative (C.19), we may define a stationary underdamped Langevin process for n † ∈ [0 : N † -1] and τ † ∈ [0 , h † ] as with the initial condition u ∗ t n ,n † h † ∼ ⃗ p t n +1 and v ∗ t n ,n † h † ∼ N (0 , I d ) . We define the joint probability distribution of ( u ∗ t n ,t † , v ∗ t n ,t † ) at time t as π ∗ t n ,t † ( u ∗ t n ,t † , v ∗ t n ,t † ) and its marginal on u ∗ t n ,t † as π ∗ , u ∗ t n ,t † ( u ∗ t n ,t † ) .

Thus, from Corollary A.5, we have that

By triangle inequality, we have

where we used the Lipschitz continuity of the learned score function (Assumption 3.3) and the true score function (Assumption 3.4), and the δ ∞ -accuracy of the learned score function at each time step (Assumption 3.1').

Now we proceed to bound the terms in the error decomposition (C.24). We first bound the F t n ,n † h † term by the following lemma:

Lemma C.13. For any n ∈ [0 : N -1] and τ † ∈ [0 , h † ] , we have and therefore

Proof. By the definition of ( u ∗ t n .n † h † + τ , v ∗ t n ,n † h † + τ ) as the stationary underdamped Langevin dynamics (C.22), we have where the first inequality follows from Cauchy-Schwarz inequality and the last inequality is by the fact that

Consequently, we have

The term E t n ,n † h † can be bounded with the following lemma:

Lemma C.14. For any n † ∈ [0 : N † -1] , suppose that γ ≲ L -1 / 2 p and T † ≲ L -1 / 2 p , then we have the following inequality for any τ † ∈ [0 , h † ]

and therefore

Proof. Recall that under the measure P , ˜ u ( K † ) t n ,n † h † ,τ † follows the dynamics of u t n ,n † h † ,τ † (C.21) for τ † ∈ [0 , h † ] , which coincides with that of u ∗ t n ,n † h † + τ † . As the only difference between the two processes u t n ,n † h † ,τ † and u ∗ t n ,n † h † + τ † is the initial condition, we can invoke Lemma 10 proved in [111] to deduce that

where the assumption that γ ≲ L -1 / 2 p and T † ≲ L -1 / 2 p is required.

Now notice that u ∗ t n ,n † h † + τ † and u t n ,n † h † ,τ † also follow the same dynamics with the true score function for τ † ∈ [0 , n † h † ] , for any coupling of u ∗ t n ,n † h † and u t n ,n † h † , we have where the last equality is again by [111, Lemma 10].

Therefore, we have

Now, we provide lemmas that are used to bound the first term in (C.24).

Lemma C.15. For any n † ∈ [0 : N † -1] , we have the following estimate:

Proof. Let µ t n ,n † h † ,τ † := ˜ u (1) t n ,n † h † ,τ † -˜ u (0) t n ,n † h † ,τ † and ν t n ,n † h † ,τ † := ˜ v (1) t n ,n † h † ,τ † -˜ v (0) t n ,n † h † ,τ † . Then for k = 0 , we may rewrite (C.15) as follows

and integrate it for τ † ′ ∈ [0 , τ † ] , which yields

We then take expectation with respect to the path measure P and then the supremum with respect to τ † ∈ [0 , h † ] , implying that

On the other hand, by applying Itô's lemma and plugging in the expression of b t n ,n † h † + τ † given by (C.17), we have

Then similarly, we may compute the derivative of ‖ ν t n ,n † h † ,τ † ‖ 2 , integrate it for τ † ∈ [0 , h † ] , and take the supremum with respect to τ † to obtain

By Itô's lemma, this equals to

where in the last equality, we used the initialization of the auxiliary corrector process ˜ v (0) t n ,n † h † , 0 ∼ N (0 , I d ) .

Again, we apply Gronwall's inequality to the above inequality and take the supremum with respect to τ † ∈ [0 , h † ] to obtain and for the difference term within the expectation, we decompose it again by the triangle inequality in (C.24), i.e.

substituting which into (C.26) completes our proof of this Lemma.

Lemma C.16 (Exponential convergence of Picard iteration in the corrector step of PIADM-ODE) . For any n † ∈ [0 , N † -1] , then the two ending terms ˜ u ( K † ) n † h † ,τ † and ˜ u ( K † ) n † h † ,τ † of the sequence { u ( k † ) n † h † ,τ † } k † ∈ [0: K † -1] satisfy the following exponential convergence rate

where the coefficient

Proof. We subtract the dynamics of u ( k +1) n † h † ,τ † and u ( k ) n † h † ,τ † in (C.15) to obtain

Then, we use the formula above to compute the derivative

and integrate for τ † ′ ∈ [0 , τ † ] to obtain

Applying Grönwall's inequality gives us that

and taking the supremum with respect to τ † ∈ [0 , h † ] on both sides above implies

integrate which for τ † ∈ [0 , τ † ] to obtain

Substituting (C.31) into (C.30) and iterating over k ∈ [1 : K † -1] , we obtain that where we plug in the results from Lemma C.15 in the last inequality. Rearranging the inequality above completes our proof.

Theorem C.17. Under Assumptions 3.1', 3.2, 3.3, and 3.4, given the following choices of the order of the parameters

and let

˜ with a total of K † N † = O ( log( dδ -2 ) ) approximate time complexity and M = Θ ( d 1 / 2 δ -2 ) space complexity for parallalizable δ -accurate score function computations.

Proof. Now, we continue the computation by plugging the decomposition in (C.24) and all the error bounds derived above into the equation. First for the last term in (C.23)

where the last inequality is by Lemma C.16. We further substitute Lemma C.14 and C.13 into (C.23) to obtain

D KL ( π t n ,T † ‖ ˜ π t n ,T † ) = D KL ( π t n ,N † h † ‖ ˜ π t n ,N † h † ) ≲ D KL ( π t n , 0 ‖ ˜ π t n , 0 ) + N † h † W 2 2 ( ˜ q t n ,h n , ⃗ p t n +1 ) + N † h † δ 2 ∞ + dN † h † ϵ † 2 + e -K † N † h † 2 d = T † W 2 2 ( ˜ q t n ,h n , ⃗ p t n +1 ) + T † δ 2 ∞ + dT † ϵ † 2 + e -K † T † h † d. Then, it is straightforward to see that when the following order of the parameters holds

and δ ∞ ≤ δ , we have

t

n

,T

†

t

n

+1

Lemma C.18. Suppose T † ≲ L -1 / 2 p , then we have

TV(

π

,

⃗ p

)

√

‖

L

1

4

p

1

(

T

†

)

3

2

˜

Proof. A complete proof of the Lemma above is presented in [111, Lemma 9], which is derived based on [145, Corollary 4.7 (1) ].

D

(

π

⃗ p

)

≲

W

(

π

,

⃗ p

KL

t

n

,T

†

t

n

+1

2

t

n

,

0

t

n

+1

)

≲

W

2

(

q

t

n

,h

n

,

⃗ p

t

n

+1

)

.

C.4 Overall Error Bound

We are now ready to prove Theorem 3.5.

Proof of Theorem 3.5. Notice that the interpolating corrector process ( ˜ u t n ,n † h † ,τ † , ˜ v t n ,n † h † ,τ † ) is constructed to follow the same dynamics as the auxiliary corrector process ( ̂ u t n ,n † h † ,τ † , ̂ v t n ,n † h † ,τ † ) in the corrector step. Therefore, we have by data processing inequality that

TV( ̂ q t n ,h n , ˜ q t n ,h n ) ≤ TV( ̂ q t n , 0 , ˜ q t n , 0 ) = TV( ̂ q t n , ⃗ p t n ) . (C.33) Furthermore, applying triangle inequality, Pinsker's inequality along with Theorem C.17 and Theorem C.7 proved above, we may upper bound the second term above as follows

TV( ̂ π ̂ u t n ,T † , ˜ π ˜ u t n ,T † ) ≤ TV( ̂ π ̂ u t n , 0 , ˜ π ˜ u t n , 0 ) = TV( ̂ q t n ,h n , ˜ q t n ,h n ) , (C.32) and again, since the interpolating predictor process ˜ y t n ,n † h † is constructed to follow the same dynamics as the auxiliary predictor process ̂ y t n ,n † h † in the predictor step, we further have by data processing inequality that

Summarizing the above inequalities, we have

where the second last inequality is deduced from Theorem C.17 and Lemma C.18 and the last inequality is derived via Theorem C.7. Therefore, for any n ∈ [0 : N -1] , applying triangle inequality along with data processing inequality ( cf. Theorem A.1) yields

where the last inequality is derived by plugging in (C.32), (C.33) and (C.36). Applying Lemma A.9 and summing the inequalities above further give us that

By setting the parameters

as desired.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We have carefully reviewed the abstract and introduction to ensure that they accurately reflect the paper's contributions and scope.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We have discussed the limitations of our work in Section 4 (Discussion and Conclusion).

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We have provided the full set of assumptions and complete proofs for all theoretical results in the paper.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification: As a theoretical paper, we do not include experiments.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: As a theoretical paper, we do not include experiments.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification: As a theoretical paper, we do not include experiments.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: As a theoretical paper, we do not include experiments.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification: As a theoretical paper, we do not include experiments.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the NeurIPS Code of Ethics and have ensured that our research conforms to it.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: As a theoretical paper, we do not expect direct societal impacts of our work.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: As a theoretical paper, we do not release data or models that have a high risk for misuse.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: As a theoretical paper, we do not use existing assets.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: As a theoretical paper, we do not introduce new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: As a theoretical paper, we do not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: As a theoretical paper, we do not involve crowdsourcing nor research with human subjects.

Peer reviews:
comments: SUMMARY:
This work improves the sample complexity of diffusion models by using the parallel sampling technique and achieves $\tilde{O}(\text{poly} \log d)$ results for reverse SDE and PFODE settings at the same time. To achieve these results, this work provides a general version of Girsanov’s theorem to deal with the additional dependence introduced by the Picard iteration.

STRENGTHS:
1.	The general version of Girsanov’s theorem is novel and suitable for the parallel sampling setting, which would raise independent interest.

WEAKNESSES:
1. For the PIADM type algorithm, the framework is similar to Algorithm 2 of [1], which is used to analyze ULMC. It seems that this work replaces the ULMC process with the  denoised process. Then, this work relaxes the LSI assumption of data by using a similar technique to [2]. It would be helpful to discuss the technique challenge except for the general Girsanov’s theorem.

2. For the PFODE setting, the PFODE-ULMC type algorithm is introduced and analyzed by [3]. Using this technique, [3] achieve a similar $\sqrt{d}$ improvement. It would be better to discuss the technique novelty compared to [3].

[1] Anari, N., Chewi, S., & Vuong, T. D. (2024). Fast parallel sampling under isoperimetry. arXiv preprint arXiv:2401.09016.

[2] Benton, J., Bortoli, V. D., Doucet, A., & Deligiannidis, G. (2024). Nearly d-linear convergence bounds for diffusion models via stochastic localization.

[3] Chen, S., Chewi, S., Lee, H., Li, Y., Lu, J., & Salim, A. (2024). The probability flow ode is provably fast. Advances in Neural Information Processing Systems, 36.

QUESTIONS:
Please see the Weakness part.

LIMITATIONS:
The authors adequately addressed the limitations and societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this article, the authors propose parallel methods for diffusion models, achieving a poly-logarithmic complexity for the first time under commonly used assumptions. By applying existing Picard Iterations to both the SDE and probability flow ODE of diffusion models, the backward process can be efficiently solved within a period of $O(1)$ length, effectively transferring sequential computations of $O(d)$ into parallelizable iterations with a depth of $O(\\log d)$. However, the author's analysis is not trivial. They employ a more intricate mathematical framework for stochastic processes and general forms of Girsanov's theorem to bound the error in each block.\n\nSTRENGTHS:\n- I believe this paper is technically solid and addresses an important open question regarding parallel methods for diffusion models.\n- This paper is well-written, and the proof is clear and comprehensive.\n- I think that the sketches and ideas presented in the proof will be useful for future research, although it is similar to the parallel method used in sampling.\n\nWEAKNESSES:\n- The authors apply Picard iterations to diffusion models without providing a comparison to other parallel methods, such as randomized midpoint methods. Including such comparisons would enhance the evaluation of their approach.\n- Please compare the total query complexity of the proposed method to that of sequential methods.\n\nMinor:\n- Line 85: Please add a reference for the Picard iteration.\n- Line 89: Please include a reference for the claim of exponentially fast convergence.\n- Last equation on page 4: Specify the distribution that is being assumed for this equation.\n- Lemma A.6: Ensure that $_F$ is included.\n- Provide an explanation regarding the different parameter choices for the predictor and corrector in Theorem 3.5.\n\nQUESTIONS:\nSee Weaknesses\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDenoising diffusion models generate samples from a complex target distribution by solving stochastic/ordinary differential initialized at the Gaussian noise. Solving these differential equations is typically done by simulating diffusion solutions. This is inherently sequential since simulating the solutions at time t_i requires the solution at time t_{i-1}.\n\nAn alternative approach to solve these diffusions is to use Picard Iterations (PI), which involve constructing a contraction mapping where the solution is the fixed point. Notably, a PI iteration updates the entire path, making parallelization amenable.\n\nThis paper introduces how to sample from a diffusion model incorporating PI into the pipeline, allowing practitioners to utilize their parallel processes efficiently. They show that a PI can speed up the run time in O(poly(log d)) in dimension compared to O(d) runtime of serial implementations of diffusion models.\n\nSTRENGTHS:\nI think the topic is well-motivated, and the authors do a great job of presenting the ideas in the paper. The idea feels quite natural and the algorithms developed have practical utility, which is uncommon for a theory paper. The theory is also strong, I can tell a lot of care has gone into the proofs.\n\nWEAKNESSES:\nI understand it is a theory paper, but I was disappointed to see no experiments, or empirical verification of the theory. The main result is the poly(log(d)) run-time, which would have been nice to see how this hold in practice. The algorithms presented in this paper seem fairly easy to implement. Experimenting with even a Gaussian toy model where analytical solutions are known would have made this paper more convincing to practitioners.\n\nQUESTIONS:\nCan you provide some intuition on how the constants scale, and how long does it take for the asymptotic to kick in?\n\nCan you provide some intuition on how to tune hyperparameters, including a number of blocks, discretization steps per block, and depth?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose to analyse the error in parallel sampling for diffusion models. This represents the first theoretical analysis of parallel sampling for diffusion models. The initial methodology was proposed in ParaDiGMS [1]. In this paper, the authors propose some incremental improvements on the samplers used in ParaDiGMS. The main contribution resides in the control of the error in parallel sampling. Notably, instead of deriving a global time complexity which depends on $O(d)$ they derive rates that are of order $O(\\mathrm{poly log}(d))$. This improvement comes from the parallel sampling nature of the scheme they analyse. The authors not only deal with the SDE (Stochastic Differential Equation) sampler but also with ODE (Ordinary Differential Equation) samplers. Parallel sampling for ULA (Unadjusted Langevin Algorithm) was already theoretically analysed in [2] but to the best of my knowledge this is the first rigorous work on the complexity of parallel sampling for diffusion models.\n\n[1] Shih et al. (2024) -- Parallel Sampling of diffusion models\n\n[2] Anari et al. (2024) -- Fast parallel sampling under isoperimetry\n\nSTRENGTHS:\n* This paper is very well-written. I found the flow of the paper to be quite easy to follow. \n\n* I have read the proofs in the case of the SDE samplers and they are correct. I also commend the authors on adding a proof sketch in the main paper. This makes the whole paper much easier to read and the proofs become more intuitive.\n\n* The rates obtained by the authors are quite compelling. In particular the $O(\\mathrm{poly log}(d))$. time complexity echoes the rates obtained in [1].\n\n* The authors not only analyse the SDE sampler but also the ODE sampler which is usually more complicated than the SDE one. I am pleasantly surprised by the gains they obtain in terms of memory complexity.\n\n[1] Anari et al. (2024) -- Fast parallel sampling under isoperimetry\n\nWEAKNESSES:\n* My main complain is that the results of Theorem 3.3 and Theorem 3.5 are given in terms of a given level of optimisation $\\delta^$. This gives a specific choice of constant for the discretisation stepsize $\\epsilon$, the time $T$, the block size $h$, the number of blocks $N$, the number of parallel iterations $K$ and $M$ the number of steps in each block. However, it would also be interesting to give a bound that depends on those quantities. What I am thinking about is that then one could optimize the bound under time complexity and space complexity constraints (for example with a requirement that $d M \\leq M_0$ where $M_0$ is the available memory of the system). Such a bound would be more interesting to gain insights on what we really gain by moving to parallel sampling. \n\n* While I'm interested in the ODE framework I am a bit confused by the explanation on how the authors have managed to reduce the $O(d^2)$ memory complexity to $O(d^{3/2})$. The explanation of l.281 \"The reduction of space complexity by the probability flow ODE implementation is intuitively owing to the fact that the probability flow ODE process is a deterministic process in time rather than a \nstochastic process as in the SDE implementation.\" is not very intuitive to me. It would be nice to point out where the improvement comes from in the proof sketch (3.2.4).\n\n* Regarding the ODE sampler, I am not sure that the bounds with SDE are directly comparable since the proof techniques use different metrics (even though the KL can be related to the TV using Pinsker's inequality). To  get comparable rates it would have been better to provide a unified framework in which both the SDE and the ODE samplers are analysed under the same techniques (and assumptions). \n\n* Assumption 3.4 is extremely strong and not satisfied in most applications. It seems that compared to [1] which analyses ODE samplers, Assumption 3.4 is not required there. Similarly for Assumption 3.3. These assumptions prevent the model to explode near the data which is something that is observed in practice (and would happen theoretically). I would have liked a more detailed comparisons between the differences of assumptions required for the analysis of sequential and parallel samplers of diffusion models.\n\n[1] Chen et al. (2023) -- \"Restoration-Degradation Beyond Linear Diffusions: A Non-Asymptotic Analysis For DDIM-Type Samplers\"\n\nQUESTIONS:\n* \"We propose new parallel inference algorithms for diffusion models using parallel sampling\" (in the TLDR). It seems that one of the claim of the authors is that they also propose a new sampler. Compared to ParaDiGMS [1], the methodological innovation seems pretty incremental (exponential integrator, shrinking step size, early stopping, see Section 1.1 \"Contributions\"). Is there something more? If there is a key methodological contribution, this should be assessed rigorously against ParaDiGMS in experimental settings. I realize that this is a theoretical paper so I would suggest to tame down the claims of methodological novelty and instead focus on the theoretical analysis. \n\n* I would be interested to see if the improvements obtained by the authors can be translated to the manifold hypothesis setting of [2, 3]. In particular, can the $O(\\mathrm{poly log}(d))$ rate be kept for this setting and/or improved to replace the dimension of the space by the dimension of the implicit manifold? \n\n* See \"Weaknesses\" for other questions and concerns. \n\n[1] Shih et al. (2024) -- Parallel Sampling of diffusion models\n\n[2] De Bortoli (2022) -- Convergence of denoising diffusion models under the manifold hypothesis\n\n[3] Yang et al. (2023) -- The Convergence of Variance Exploding Diffusion Models under the Manifold Hypothesis\n\nLIMITATIONS:\nIn the paper the limitations are discussed in Section 4, i.e. \"Although we anticipate implementing diffusion models in parallel may introduce engineering challenges, e.g. scalability, hardware compatibility, memory bandwidth, etc., we believe that our theoretical contributions lay a solid foundation that not only supports but also motivates the empirical development of parallel inference algorithms for diffusion models since advancements continue in GPU power and memory efficiency.\" I think there are other limitations that need to be discussed regarding the current work and the theoretical assumptions that the authors needed to make in order to derive their results. While I don't think these limitations hinder the results of the paper they should be clearly laid out.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}