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

Authors: haoxuan  chen
yinuo  ren
lexing  ying
grant m rotskoff
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 23:25:49

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.

Main paper: 
 introduction: 
 Diffusion and probability flow based models [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] are now state-of-the-art in many fields, such as computer vision and image generation [12] [13] [14] [15] [16] [17] [18] [19] [20] [21] [22] , natural language processing [23, 24] , audio and video generation [25] [26] [27] [28] [29] , optimization [30, 31] , sampling and learning of fixed classes of distributions [32] [33] [34] [35] [36] [37] [38] [39] [40] [41] , solving high-dimensional partial differential equations [42] [43] [44] [45] [46] , and more recently several applications in physical, chemical and biological fields [47] [48] [49] [50] [51] [52] [53] [54] [55] [56] [57] [58] [59] [60] [61] [62] [63] . For a more comprehensive list of related work, one may refer to the following review papers [64] [65] [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] [3] [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 made to accelerate this process by resorting to higher-order or randomized numerical schemes [69] [70] [71] [72] [73] [74] [75] , augmented dynamics [76] , adaptive step sizes [77] , operator learning [78] , restart sampling [79] , self-consistency [80, 81] and knowledge distillation [82] [83] [84] . Recently, several empirical works [1, 85] 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 [86, 87] also combine the parallel sampling technique with the randomized midpoint method [88] 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 [89] [90] [91] [92] [93] [94] [95] [96] [97] [98] [99] [100] [101] [102] [103] [104] [105] [106] . 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) [98] , improved from the previous O(d 2 ) bound [95] . [102] 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.• 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;• We show 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. 
preliminaries: 
 In this section, we briefly recapitulate the framework of score-based diffusion models, define notations, and discuss related work.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):dx s = β s (x s )ds + σ s dw s , with x 0 ∼ p 0 , (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 [107] :EQUATIONwhere ⃗ * 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:dx s = - 1 2x s ds + dw s , and d ⃗x t = 1 2 ⃗ x t + 1 + υ 2 2 ∇ log ⃗ p t ( ⃗ x t ) dt + υdw t . (2.3)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 [108, 109] :EQUATIONand the backward process in (2.3) is approximated by the following SDE thereafter:dy t = 1 2 y t + 1 + υ 2 2s θ t (y t ) dt + υdw t , with y 0 ∼ N (0, I d ).(2.5)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 [102] claim that the probability flow ODE implementation with υ = 0 may exhibit better time complexity. We refer to [102, 110] for theoretical and [111, 112] for empirical comparisons of different implementations.Parallel sampling algorithms have been actively explored in the literature, including the parallel tempering method [113] [114] [115] and several recent studies [116] [117] [118] . For diffusion models, the idea of parallel sampling is based on the Picard iteration [119, 120] for solving nonlinear ODEs. Suppose we have an ODE dx t = f t (x t )dt and we would like to solve it for t ∈ [0, T ], then the Picard iteration is defined as follows:x (0) t ≡ x 0 , and x (k+1) tEQUATIONUnder 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 [121] . Roughly speaking, it is proposed to simulate the Langevin diffusion process dx t = -∇V (x t )dt + dw t with the following iteration resembling (2.6):EQUATIONwhere 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. [121, 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.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 [95, 91] , iteration complexity in [98, 102] , 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.. Suppose X is a random variable generating a sub-σ-algebra F X of F, and p(•|X) q(•|X) holds for any value of X, thenD KL (p q) = D KL (p| F X q| F X ) + E F X [D KL (p(•|X) q(•|X))].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:dz t (ω) = α(t, ω)dt + Σ(t, ω)dw t (ω),for any time t, we denote the marginal distribution of z t by p t , i.e.p t := p z -1 t (•), where z t : Ω → R m , ω → z t (ω), as well as the path measure of the process z t in the sense ofp t1:t2 := p z -1 t1:t2 (•) , where z t1:t2 : Ω → C([t 1 , t 2 ], R m ), ω → (z t (ω)) t∈[t1,t2] .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 1. f (t, ω) is B × F -measurable,E t2 t1 f 2 (t, ω)dt < +∞.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 [102] , here we can avoid checking Novikov's condition throughout our proofs below by using the approximation argument presented in [91] . A review of Girsanov can be found in textbooks like in [122, 135] . We will present the following generalized version of Girsanov's theorem: Theorem A.4 (Girsanov's Theorem [122, 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:dz t (ω) = α(t, ω)dt + Σ(t, ω)dw t (ω), (A.2)and there exist processesδ(t, ω) ∈ V n and β(t, ω) ∈ V m such that 1. Σ(t, ω)δ(t, ω) = α(t, ω) -β(t, ω);2. The process M t (ω) as defined below is a martingale with respect to the filtration {F t } t≥0 and probability measure q: 2) in the sense of q t1:t2 = q z -1 t1:t2 (•) . Moreover, the KL divergence between the two path measures is given byM t (ω) = exp - t 0 δ(s, ω) ⊤ dw s (ω) - 1 2 t 0 δ(s, ω)D KL (p t1:t2 q t1:t2 ) = D KL (p t1 q t1 ) + E ω∼p| F t 1 1 2 t2 t1 δ(t, ω) 2 dt (A.4)Proof. First, by Theorem A.1, we haveD KL (p t1:t2 q t1:t2 ) = D KL (p| Ft 1 q| Ft 1 )+E z∼p| F t 1 D KL p( z -1 t1:t2 (•))| z t1 = z) q( z -1 t1:t2 (•))| z t1 = z) .From Girsanov's theorem (Theorem A.4), we have that the measure p| Ft 1 is absolutely continuous with respect to q| Ft 1 , which allows us to compute the second term above as follows:D KL p( z -1 t1:t2 (•)| z t1 = z) q( z -1 t1:t2 (•)| z t1 = z) = E zt 1 :t 2 log dp( z -1 t1:t2 (•)| z t1 = z) dq( z -1 t1:t2 (•)| z t1 = z) = E ω∼p| F t 1 log dp| Ft 1 dq| Ft 1 = E ω∼p| F t 1 - t2 t1 δ(t, ω) ⊤ dw t (ω) - 1 2 t2 t1 δ(t, ω) 2 dt = E ω∼p| F t 1 - t2 t1 δ(t, ω) ⊤ (d w t (ω) -δ(t, ω)dt) - 1 2 t2 t1 δ(t, ω) 2 dt = E ω∼p| F t 1 1 2 t2 t1 δ(t, ω) 2 dt ,and thereforeD KL (p t1:t2 q t1:t2 ) = D KL (p t1 q t1 ) + E ω∼p| F t 1 1 2 t2 t1 δ(t, ω) 2 dt ,which completes the proof. 
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 [85] , where several practical compromises are made to mitigate GPU memory constraints and theoretical guarantees are still lacking.In this section, we will propose Parallelized Inference Algorithms for Diffusion Models 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.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. n: N = O(log d) blocks q 0 ≈ N (0, I d ) q t N ≈ p data h0 h1 hn-1 hN-1 η O(1) k = 0 k = 1 k = K O(d -1 ) or O(d -1/2 ) . . . . . . M n = O(d) or O( √ d) parallalizable steps q tn q tn+1 Inner Iterations k: K = O(log d) depth ϵn,0 ϵn,1 ϵn,M-1t = t n + τ n,m for m = 1, 2, M n -1, M n .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 ) Mn m=0 with τ n,0 = 0 and τ n,Mn = 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 [95, 98] , to curb the potential blow-up of the score function as t → T , which is shown by [98] for 0 ≤ s < t < T to be of the orderM n = M := h/ϵ, for n ∈ [0 : N -2] and m ∈ [0 : M -1]. FollowingE t s ∇ log ⃗ p τ ( ⃗ x τ ) -∇ log ⃗ p s ( ⃗ x s ) 2 dτ ≲ d t -s T -t 2 ,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 ϵ N -1,m ≤ ϵ ∧ ϵ (h -τ N -1,m+1 ) . (3.1)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 thatIn(τ ) j=1 ϵ n,j ≤ τ < In(τ )+1 j=1ϵ n,j . We also define 2 (e ϵn,j -1) s θ tn+τn,j ( yAlgorithm 1: PIADM-SDE Input: y 0 ∼ q 0 = N (0, I d ), a discretization scheme (T , (h n ) N n=1 and (τ n,m ) n∈[1:N ],m∈[0:M ] ) satisfying (3.1), the depth of iteration K, the learned NN-based score function s θ t (•). Output: A sample y t N ∼ q t N ≈ ⃗ p T . 1 for n = 0 to N -1 do 2 y (0) tn,τn,m ← y tn , ξ m ∼ N (0, I d ) for m ∈ [0 : M n ] in parallel;(k)tn,τn,j ) + √ e ϵn,j -1ξ j ;(3.4) EQUATIONWe denote the distribution of y(K)tn,τ by q tn+τ . 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 θ tn+τn,j ( y(k)tn,τn,j ) for j ∈ [0 : M n ] is evaluated in parallel, with subsequent floating-point operations comparably negligible, resulting in the overall O(N K) approximate time complexity.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: [91, 95, 102] , while we adopt the normalization Assumption 3.2 from [98] 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.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 parametersT = O(log(dδ -2 )), h = Θ(1), N = O log(dδ -2 ) , ϵ = Θ d -1 δ 2 log -1 (dδ -2 ) , M = O dδ -2 log(dδ -2 ) , K = O(log(dδ -2 )),and let L 2 s h n e 7 2 hn1, δ 2 ≲ δ, T ≲ log η -1 , the distribution q t N that PIADM-SDE (Algorithm 1) generates samples from satisfies the following error bound: • Compared with log-concave sampling [121] , 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 [98] for the SDE implementation of diffusion models;D KL (p η q t N ) ≲ de -T + dϵT + δ 2 2 + dT e -K ≲ δ 2 ,• 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.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 [91, 95, 102] , 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)tn,τ ) τ ∈[0,hn] depends on all previous paths ( y(k) tn,τ ) τ ∈[0,hn] 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 ( [122, Theorem 8.6.6] ). For example, in the n-th block, we apply the following change of measure procedure: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) tn,τ and the initial condition y (0) tn,τ in Lemma B.5. However, due to the exponential convergence of the Picard iteration, this term does not affect the overall error bound.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.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 [102] 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) [121] , which is incompatible with our desired poly-logarithmic guarantee.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 [95, 102] . Recent work on the probability flow ODE implementation [103, 105] 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 θ tn,τn,m is δ ∞ -accurate in the sense of E ⃗ p tn +τn,m s θ tn+τn,m ⃗ x tn+τn,m -∇ log ⃗ p tn+τn,m ⃗ x tn+τn,m 2 ≤ δ 2 ∞ .Assumption 3.4 (Bounded and Lipschitz true score). The true score function ∇ log p t has boundedC 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. [94, Assumption 1.5]) are left for further work.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 parametersT = 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 andT † = O(1) ≲ L -1/2 p ∧ L -1/2 s , h † = Θ(1), N † = O(1), ϵ † = Θ(d -1/2 δ), M † = O(d 1/2 δ -1 ), K † = O(log(dδ -2 )),for the inner iteration during the corrector step, and letL 2 s h 2 e h ∨ L 2 s h † 2 e h † /γ 1, δ ∞ ≲ δ log -1 (dδ -2 ), and γ ≳ L 1/2p , then the distribution q t N that PIADM-ODE (Algorithm 2) generates samples from satisfies the following error bound: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.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).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 2-Wasserstein 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. 
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 PIADM-SDE 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, 85] 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, [123] [124] [125] [126] [127] [128] [129] [130] [131] [132] [133] and multi-marginal [134] 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. 
 
 Appendix: 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.We adopt the following notations throughout the paper:Notation DescriptionThe set {a, a + 1, . . . , b}I d Identity matrix in R d×d ⃗ * t * T -t *Used to denote quantities produced by the algorithm * Used to denote quantities related to the auxiliary processes * †Used to denote quantities related to the corrector step•The Euclidean norm of a vector ≲ or ≳The inequality holds up to a constant factor Absolute continuity (for measures)/ much less than (for quantities)(x t ) t≥0The forward process of the diffusion model (2.3)( ⃗ x t ) t∈[0,T ]The backward process of the diffusion model (2.3) (y t ) t∈[0,T ] The approximate backward process of the diffusion model (2.5)y (k) tn,τn,mThe approximate value of the approximate process y t at time t n +τ n,m after k iterations in the (n + 1)-th block y tnThe value of the approximate process y t at time t n q tnThe distribution of y tn z t1:t2The path(z t ) t∈[t1,t2] of the process z t D f (• •) The f -divergence between two distributions D KL (• •) The KL divergence between two distributions T V (• , •) The total variation distance between two distributions W 2 (•, •)The 2-Wasserstein distance between two distributions Table 2 : Summary of notationsTheorem A.1 (Properties of f -divergence). Suppose p and q are two probability measures on a common measurable space (Ω, F) with p q. The f -divergence between p and q is defined asD f (p q) = E X∼q f dp dq , (A.1)where dp dq 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:1. (Data-processing inequality). Suppose H is a sub-σ-algebra of F, the following inequality holdsD f (p| H q| H ) ≤ D f (p q); for any f -divergence D f (• •).Lemma A.6 ([98, Lemma 2] ). For the backward process (2.3), we have for 0 ≤ s < t < T ,d dt E ∇ log ⃗ p t ( ⃗ x t ) -∇ log ⃗ p s ( ⃗ x s ) 2 ≤ 1 2 E ∇ log ⃗ p s ( ⃗ x s ) 2 + E ∇ 2 log ⃗ p t ( ⃗ x t ) 2 F .Lemma A.7 ([98, Lemma 3] ). For the forward process (2.3), we have for 0 ≤ t < T ,E [∇ log p t (x t )] ≤ dσ -2 t , and E ∇ 2 log p t (x t ) 2 F ≤ dσ -4 t + 2 d dt σ -4 t E [tr Σ t ] ,where the posterior covariance matrix Σ t := cov p 0|t (x 0 ) and σ Proof. Conditioned on x 0 , we have that⃗ x tn = x T -tn ∼ N e -1 2 (T -tn) x 0 , (1 -e -(T -tn) )I d ,and⃗ x tn+τ = x T -tn-τ ∼ N e -1 2 (T -tn-τ ) x 0 , (1 -e -(T -tn-τ ) )I d for any τ ∈ [0, h n ]. Therefore, we haveE ⃗ x tn 2 = E E x T -tn 2 x 0 ≤ E E x T -tn -e -1 2 (T -tn) x 0 2 x 0 + e -1 2 (T -tn) x 0 2 ≤ d(1 -e -(T -tn) ) + e -(T -tn) E[ x 0 2 ] ≤ 2d.Taking the difference between them then implies that for any τ ∈ [0, h n ],E ⃗ x tn -⃗ x tn+τ 2 = E E x T -tn -x T -tn-τ 2 |x 0 ≤ d(2 -e -(T -tn) -e -(T -tn-τ ) ) + e -1 2 (T -tn) -e -1 2 (T -tn-τ ) 2 E[ x 0 2 ] ≤ 2d + e -(T -tn-τ ) (1 -e -1 2 τ ) 2 E[ x 0 2 ] ≤ 3d.Lemma A.9 (Lemma 9 in [95] ). 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 haveTV( ⃗ p 0 , q 0 ) 2 ≤ D KL ( ⃗ p 0 q 0 ) ≲ de -T .In this section, we will present the missing proofs for Theorem 3.3. For readers' convenience, we reiterate the backward process (2.3)d ⃗ x t = 1 2 ⃗ x t + ∇ log ⃗ p t ( ⃗ x t ) dt + dw t , with ⃗ x 0 ∼ p T , (B.1)and its approximate version (2.5) with the learned score functiondy t = 1 2 y t + s θ t (y t ) dt + dw t , with y 0 ∼ N (0, I d ).The filtration F t refers to the filtration of the SDE (B.1) up to time t.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 with the initial condition yFEQUATIONy (k) tn,0 (ω) ≡ y tn (ω) for k ∈ [1 : K] (B.3)where y tn (ω) = y(K) tn-1,τ n-1,M n-1 (ω) if n ∈ [1 : N -1] and y t0 (ω) ∼ N (0, I d ).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) tn,τ (ω)) τ ∈[0,hn] is F tn+τ -adapted for any k ∈ [0 : K] and n ∈ [0 : N -1].Proof. Since the initialization y Proof. The dependency on ω will be omitted in the proof below. (0) tn,τ (ω) ≡ y tn (ω) for τ ∈ [0, h n ], where y tn (ω) is F tn -adapted, it is obvious that y (0) tn,τ (ω) is F tn+τ -adapted. Now suppose that ( y (k) tn,τ (ω)) τ ∈[0,hn] is F tn+τ -adapted, since g n (τ ) ≤ τ ,We shall invoke Girsanov's theorem (Theorem A.4) in the procedure as detailed below:1. 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| Ft n ;2. Defining another process w tn+τ (ω) governed by the following SDE:d w tn+τ (ω) = dw tn+τ (ω) + δ tn (τ, ω)dτ,where EQUATIONThen Corollary A.5 provides the following computationD KL ( ⃗ p tn+1 q tn+1 ) ≤ D KL ( ⃗ p tn:tn+1 q tn:tn+1 ) = D KL ( ⃗ p tn q tn ) + E ω∼q| F tn 1 2 hn 0 δ tn (τ, ω) 2 dτ , (B.6)where the first inequality is by the data-processing inequality (Theorem A.1). Now, the problem remaining is to bound the discrepancy quantified byhn 0 δ tn (τ, ω) 2 dτ = hn 0 s θ tn+gn(τ ) ( y (K-1) tn,gn(τ ) (ω)) -∇ log ⃗ p tn+τ ( y (K) tn,τ (ω)) 2 dτ ≤ 3 hn 0 ∇ log ⃗ p tn+gn(τ ) y (K) tn,gn(τ ) (ω) -∇ log ⃗ p tn+τ y (K)tn,τ (ω)2 dτ :=At n (ω) + hn 0 s θ tn+gn(τ ) y (K) tn,gn(τ ) (ω) -∇ log ⃗ p tn+gn(τ ) y (K)tn,gn(τ ) (ω)2 dτ :=Bt n (ω) + hn 0 s θ tn+gn(τ ) y (K) tn,gn(τ ) (ω) -s θ tn+gn(τ ) y (K-1)tn,gn(τ ) (ω) By integrating from 0 to τ and taking the expectation on both sides of (B.9), we obtain thatEQUATIONE ω∼ ⃗ p| F tn z tn,τ 2 = E ω∼ ⃗ p| F tn τ 0 z tn,τ ′ 2 + z ⊤ tn,τ ′ y (0) tn,τ ′ + 2z ⊤ tn,τ ′ s θ tn+gn(τ ′ ) y (0) tn,gn(τ ′ ) + d dτ ′ + 2E ω∼ ⃗ p| F tn τ 0 z ⊤ tn,τ ′ d w tn+τ ′ (ω) -δ tn (τ ′ , ω)dτ ′ ,and by AM-GM, we further haveE ω∼ ⃗ p| F tn z tn,τ 2 ≤ E ω∼ ⃗ p| F tn τ 0 7 2 z tn,τ ′ 2 + 1 2 y (0) tn,τ ′ 2 + s θ tn+gn(τ ′ ) y (0) tn,gn(τ ′ ) 2 + d + δ tn (τ, ω) 2 dτ ′ ≤ τ 0 E ω∼ ⃗ p| F tn 7 2 z tn,τ ′ 2 + δ tn (τ, ω) 2 dτ ′ + 1 2 E y (0) tn,τ 2 + M 2 s + d τ,where δ tn (τ, ω) is defined in (B.4). Similar to (B.7), we may use triangle inequality to upper bound δ tn (τ, ω) 2 , which implies that for any τ ∈ [0, h n ]E ω∼ ⃗ p| F tn z tn,τ 2 ≤ 7 2 τ 0 E ω∼ ⃗ p| F tn z tn,τ ′ 2 dτ ′ + 1 2 E y (0) tn,τ 2 + M 2 s + d τ + 3E ω∼ ⃗ p| F tn τ 0 s θ tn+gn(τ ) ( y (K-1) tn,gn(τ ) (ω)) -s θ tn+gn(τ ) ( y (K) tn,gn(τ ) (ω)) 2 dτ ′ + 3E ω∼ ⃗ p| F tn τ 0 s θ tn+gn(τ ) ( y (K) tn,gn(τ ) (ω)) -∇ log ⃗ p tn,gn(τ ) ( y (K) tn,gn(τ ) (ω)) 2 dτ ′ + 3E ω∼ ⃗ p| F tn τ 0 ∇ log ⃗ p tn+gn(τ ) ( y (K) tn,gn(τ ) (ω)) -∇ log ⃗ p tn+τ ( y (K) tn+τ (ω)) 2 dτ ′ ≤ 7 2 τ 0 E ω∼ ⃗ p| F tn z tn,τ ′ 2 dτ ′ + 1 2 E y (0) tn,τ 2 + M 2 s + d τ + 3L 2 s τ 0 E ω∼ ⃗ p| F tn y (K) tn,gn(τ ′ ) (ω) -y (K-1)tn,gn(τ ′ ) (ω)2 dτ ′ + 3E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] ,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 ]E ω∼ ⃗ p| F tn z tn,τ 2 ≤ e 7 2 τ 1 2 E y (0) tn,τ 2 + M 2 s + d τ + 3e 7 2 τ E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] + 3e 7 2 τ L 2 s τ 0 E ω∼ ⃗ p| F tn y (K) tn,gn(τ ′ ) (ω) -y (K-1)tn,gn(τ ′ ) (ω) E ω∼ ⃗ p| F tn y (K) tn,τ ′ (ω) -y (K-1) tn,τ ′ (ω) 2 .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. explicitly, integrate it from τ ′ = 0 to τ , and obtain the following inequalityE ω∼ ⃗ p| F tn y (K) tn,τ (ω) -y (K-1) tn,τ (ω) 2 2 ≤ L 2 s h n e 2hn K-1 he 7 2 hn M 2 s + 2d 1 -3 (L 2 s h n e 2hn ) K-1 e 7 2 hn h n L 2 s + 3 L 2 s h n e 2hn K-1 e 7 2 hn E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] 1 -3 (L 2 s h n e 2hn ) K-y (k+1) tn,τ (ω) -y (k) tn,τ (ω) 2 = τ 0 2 y (k+1) tn,τ ′ (ω) -y (k) tn,τ ′ (ω) ⊤ s θ tn+gn(τ ′ ) y (k) tn,gn(τ ′ ) (ω) -s θ tn+gn(τ ′ ) y (k-1) tn,gn(τ ′ ) (ω) dτ ′ + τ 0 y (k+1) tn,τ ′ (ω) -y (k) tn,τ ′ (ω) 2 dτ ′ ≤ 2 τ 0 y (k+1) tn,τ ′ (ω) -y (k) tn,τ ′ (ω) 2 dτ ′ + τ 0 s θ tn+gn(τ ′ ) y (k) tn,gn(τ ′ ) (ω) -s θ tn+gn(τ ′ ) y (k-1)tn,gn(τ ′ ) (ω)2 dτ ′ ≤ 2 τ 0 y (k+1) tn,τ ′ (ω) -y (k) tn,τ ′ (ω) 2 dτ ′ + L 2 s τ 0 y (k) tn,gn(τ ′ ) (ω) -y (k-1)tn,gn(τ ′ ) (ω) Taking expectation on both sides above further implies that for any τ ∈ [0, h n ],E ω∼ ⃗ p| F tn y (k+1) tn,τ (ω) -y (k) tn,τ (ω) 2 ≤ L 2 s e 2τ τ 0 E ω∼ ⃗ p| F tn y (k) tn,gn(τ ′ ) (ω) -y (k-1) tn,gn(τ ′ ) (ω) 2 dτ ′ ≤ L 2 s τ e 2τ sup τ ′ ∈[0,τ ] E ω∼ ⃗ p| F tn y (k) tn,τ ′ (ω) -y (k-1) tn,τ ′ (ω) 2 . (B.13)Furthermore, we take supremum over τ ∈ [0, h n ] on both sides above and iterate (B.12) over k ∈ N, which indicates The following lemma from [98] bounds the expectation of the term A tn (ω) in (B.7):sup τ ∈[0,hn] E ω∼ ⃗ p| F tn y (k+1) tn,τ (ω) -y (k) tn,τ (ω) 2 ≤ L 2 s h n e 2hn sup τ ∈[0,hn] E ω∼ ⃗ p| F tn y (k) tn,τ (ω) -y (k-1) tn,τ ′ (ω) 2 ≤ L 2 s h n e 2hn k sup τ ∈[0,hn] E y (1) tn,τ (ω) -y (0) tn,τ (ω) 2 ≤ L 2 s h n e 2hnLemma B.7 ([98, Section 3.1]). We haveE ω∼ ⃗ p| F tn [A tn (ω)] ≲ ϵdh n , for n ∈ [0 : N -2], andE ω∼ ⃗ p| F tn A t N -1 (ω) ≲ ϵd log η -1 ,where η is the parameter for early stopping.Proof. Notice thatE ω∼ ⃗ p| F tn [A tn (ω)] =E ω∼ ⃗ p| F tn hn 0 ∇ log ⃗ p tn+gn(τ ) y (K) tn,gn(τ ) (ω) -∇ log ⃗ p tn+τ y (K) tn,gn(τ ) (ω) 2 dτ =E ω∼ ⃗ p| F tn Mn m=0 τn,m+1 τn,m ∇ log ⃗ p tn+τn,m y (K) tn,τn,m (ω) -∇ log ⃗ p tn+τ y (K) tn,τ (ω) 2 dτ , = Mn m=0 τn,m+1 τn,m E ω∼ ⃗ p| F tn ∇ log ⃗ p tn+τn,m ⃗ x tn+τ (ω) -∇ log ⃗ p tn+τ ⃗ x tn+τ (ω) 2 dτ,where for the last equality, we use the fact that the process y(K)tn,τ (ω) follows the backward SDE with the true score function under the measure ⃗ p. In the following, we drop the superscript ω ∼ ⃗ p| Ft n of the expectation for simplicity. By Lemma A.6 and A.7, we have E ∇ log ⃗ p tn+τn,m ⃗ x tn+τ (ω) -∇ log ⃗ p tn+τ ⃗ x tn+τ (ω) 2 ≤ τ 0 1 2 E ∇ log ⃗ p tn+τn,m ⃗ x tn+τn,m (ω) 2 + E ∇ 2 log ⃗ p tn+τ ′ ⃗ x tn+τ ′ (ω) 2 F dτ ′ ≤ τ 0 1 2 d ⃗ σ -2 τ ′ + d ⃗ σ -4 τ ′ dτ ′ + ⃗ σ -4 tn+τn,m E tr ⃗ Σ tn+τn,m -⃗ σ -4 tn+τ E tr ⃗ Σ tn+τ , Now noticing that ⃗ σ 2 t = σ 2 T -t ≲ T -t, we further have τn,m+1 τn,m E ∇ log ⃗ p tn+τn,m ⃗ x tn+τ (ω) -∇ log ⃗ p tn+τ ⃗ x tn+τ (ω) 2 dτ ≲ τn,m+1 τn,m τ ′ 0 d (T -t n -τ n,m+1 ) 2 dτ ′ dτ + ϵ n,m E tr ⃗ Σ tn+τn,m -E tr ⃗ Σ tn+τn,m+1 (T -t n -τ n,m ) 2 ≲d ϵ 2 n,m (T -t n -τ n,m+1 ) 2 + ϵ E tr ⃗ Σ tn+τn,m -E tr ⃗ Σ tn+τn,m+1(T -t n -τ n,m+1 ) 2 + Mn m=0 ϵ T -t n -τ n,m E tr ⃗ Σ tn+τn,m -E tr ⃗ Σ tn+τn,m+1 ≤dϵ 2 M n + ϵE tr ⃗ Σ tn+τn,0 T -t n -τ n,0 + Mn m=0 ϵϵ n,m E tr ⃗ Σ tn+τn,m (T -t n -τ n,m+1 )(T -t n -τ n,m ) ≤dϵ 2 M n + ϵd + dϵ 2 M n ≲ dϵ 2 M n .For n ∈ [0, N -2], we have M n ϵ = h n and thus E ω∼ ⃗ p| F tn [A tn (ω)] ≲ ϵdh n , and for n = N -1, we haveM N ≲ h η 1 ϵτ dτ = log η -1 ϵ -1and thusE ω∼ ⃗ p| F tn A t N -1 (ω) ≲ ϵ 2 dM n ≲ ϵd log η -1 .Proof of Theorem 3.3. We first continue the computation in (B.6) and (B.7):D KL ( ⃗ p tn+1 q tn+1 ) ≤ D KL ( ⃗ p tn q tn ) + E ω∼ ⃗ p| F tn 1 2 hn 0 δ tn (τ, ω) 2 dτ ≤D KL ( ⃗ p tn q tn ) + 3E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] +3E ω∼ ⃗ p| F tn hn 0 s θ tn+gn(τ ) y (K) tn,gn(τ ) (ω) -s θ tn+gn(τ ) y (K-1)tn,gn(τ ) (ω)2 dτ ≤D KL ( ⃗ p tn q tn ) + 3E ω∼ ⃗ p| F tn A tn (ω) + B tn (ω) + L 2 s hn 0 y (K) tn,gn(τ ) (ω) -y (K-1)tn,gn(τ ) (ω)2 dτ ≤D KL ( ⃗ p tn q tn ) + 3E ω∼ ⃗ p| F tn A tn (ω) + B tn (ω) + h n L 2 s sup τ ∈[0,hn] y (K) tn,τ (ω) -y (K-1) tn,τ (ω) 2 .Then plugging in the result of Lemma B.6, we haveD KL ( ⃗ p tn+1 q tn+1 ) ≤D KL ( ⃗ p tn q tn ) + 3E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] + 3h n L 2 s L 2 s h n e 2hn K-1 he 7 2 hn M 2 s + 2d 1 -3 (L 2 s h n e 2hn ) K-1 e 7 2 hn h n L 2 s +h n L 2 s 9 L 2 s h n e 2hn K-1 e 7 2 hn E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] 1 -3 (L 2 s h n e 2hn ) K-1 e 7 2 hn h n L 2 s ≲D KL ( ⃗ p tn q tn ) + 1 + e -K h n e hn 1 -e -K h n e hn E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] + e -K h 2 n e hn d ≲D KL ( ⃗ p tn q tn ) + E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] + e -K h 2 n e hn d,where we used the assumption that L2 s h n eThe termN -1 n=0 E ω∼ ⃗ p| F tn [B tn (ω)] is bounded by Assumption 3.1 asN -1 n=0 E ω∼ ⃗ p| F tn [B tn (ω)] ≤E ω∼ ⃗ p| F tn N -1 n=0 hn 0 s θ tn+gn(τ ) y (K) tn,gn(τ ) (ω) -∇ log ⃗ p tn+gn(τ ) y (K) tn,gn(τ ) (ω) 2 dτ =E ω∼ ⃗ p| F tn N -1 n=0 Mn-1 m=0 ϵ n,m s θ tn+τn,m y (K) tn,τn,m (ω) -∇ log ⃗ p tn+τn,m y (K) tn,τn,m (ω) 2 =E ω∼ ⃗ p| F tn N -1 n=0 Mn-1 m=0 ϵ n,m s θ tn+τn,m ⃗ x tn+τ (ω) -∇ log ⃗ p tn+τn,m ⃗ x tn+τ (ω) 2 ≤ δ 2 2 ,where the last equality is because the process y(K)tn,τ (ω) under measure ⃗ p follows the backward SDE (B.5). Thus, by Theorem A.1 and plugging in the iteration relations aboveD KL (p η q t N ) = D KL ( ⃗ p t N q t N ) ≤D KL ( ⃗ p 0 q 0 ) + N -1 n=0 E ω∼ ⃗ p| F tn [A tn (ω) + B tn (ω)] + e -K h 2 n e hn d ≤D KL ( ⃗ p 0 q 0 ) + N -2 n=0 ϵdh n + ϵd log η -1 + N -1 n=0 E ω∼ ⃗ p| F tn [B tn (ω)] + e -K h 2 n e hn dN ≤de -T + ϵd(T + log η -1 ) + δ 2 2 + e -K dT ≤ de -T + ϵdT + δ 2 + e -K dT,as T ≳ log η -1 , h n ≲ 1, and δ 2 ≲ δ, and then it is straightforward to see that the following choices of parametersT = O(log(dδ -2 )), h = Θ(1), N = O log(dδ -2 ) , ϵ = Θ d -1 δ 2 log -1 (dδ -2 ) , M = O dδ -2 log(dδ -2 ) , K = O(log(dδ -2 )),would yield an overall error of O(δ 2 ).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.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.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 tn,0 ∼ N (0, I d );12 for n † = 0 to N † -1 do 13 ( u (0) tn,n † h † ,m † ϵ † , v (0) tn,n † h † ,m † ϵ † ) ← ( u tn,n † h † , v tn,n † h † ) for m † ∈ [0 : M † ]; 14 for k † = 1 to K † do 15 ( u (k † ) tn,n † h † ,0 , v (k † ) tn,n † h † ,0 ) ← ( u tn,n † h † , v tn,n † h † ); 16 ξ j † ∼ N 0, 2γ(1 + γ -2 )(1 -e -γϵ † ) 2 e -2γ((M † -j † +1)ϵ † ) I d for j † ∈ [0 : M † ]; 17 for m † = 1 to M † in parallel do 18   u (k † ) tn,n † h † ,m † ϵ † v (k † ) tn,n † h † ,m † ϵ †   ← G(m † ϵ † )   u (k † -1) tn,n † h † ,0 v (k † -1) tn,n † h † ,0   + m-1 j † =0 G((m † -j † -1)ϵ † ) I d -G(ϵ † ) 0 s θ tn+1 ( u (k † -1) tn,n † h † ,j † ϵ † ) + m-1 j † =0 G((m † -j † -1)ϵ † ) 0 ξ j † ; (C.2) 19 end 20 end 21 ( u tn,(n † +1)h † , v tn,(n † +1)h † ) ← ( u (K † ) tn,n † h † ,h † , v (K † ) tn,n † h † ,h † ); 22 end 23 y tn+1 ← u tn,T † ; 24 end q tn = q tn,0 q tn,hn = π u tn,0 π u tn,T † = q tn+1 y (k) tn,τ u (k) tn,n † h † ,τ † , v (k) tn,n † h † ,τ † ⃗ p tn q tn,hn = π u tn,0 = π u tn,0 π u tn,T † π u tn,T † y (k) tn,τ u (k) tn,n † h † ,τ † , v (k) tn,n † h † ,τ † u tn,n † h † +τ † , v tn,n † h † +τ † ⃗ p tn+1 = π * ,u * tn,0 π * ,u * tn,T † = ⃗ p tn+1 ⃗ x tn+τ u * tn,n † h † +τ † , v * tn,n † h † +τ † W 2 ( q tn,hn , ⃗ p tn+1 ) (Theorem C.7) D KL (π tn,T † π tn,T † ) (Theorem C.17) TV(π tn,T † , ⃗ p tn+1 ) (Lemma C.18) TV( π u tn,T † , π u tn,T † ) (Theorem A.1)Figure 2 : Illustration of the proof pipeline of Theorem 3.5 for PIADM-ODE within the n-th block.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) tn,hn as the initialization for the next corrector step.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 tn,0 = y (K) tn,hn and the augmented variable v tn,0 ∼ N (0, I d ) representing the momentum.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).In the following subsections, we proceed to provide theoretical guarantees for the algorithm. wherey tn = u tn-1,N † h † if n ∈ [1 : N -1]and y t0 ∼ N (0, I d ). We will also denote the probability distribution of y(K)tn,τ as q tn,τ . for k ∈ [1 : K],where y(0)tn,0 ∼ ⃗ p tn . We will also denote the probability distribution of y(K)tn,τ as q tn,τ . Similar to the equivalence between (3.4) Plugging in τ = τ n,m gives us (C.1), as desired.Lemma C.4 (Error between the interpolating process and the true process). Under the Picard iteration, we have that the ending process { y(K) tn,τ } τ ∈[0,hn] satisfies the following exponential convergence rate sup τ ∈[0,hn] E y (K) tn,τ -⃗ x tn+τ 2 ≤ 3d h 2 n e hn+ 3 2 L 2 s 2 K + e hn+ 3 2 h n /2 1 -h 2 n e hn+ 3 2 L 2 s /2 h n δ 2 ∞ + E[D tn ] ,whereD tn := hn 0 s θ tn+gn(τ ′ ) ⃗ x tn+gn(τ ′ ) -s θ tn+τ ′ ( ⃗ x tn+τ ′ ) 2 dτ ′ .Proof. Recall that the backward true process { ⃗ x tn+τ } τ ∈[0,hn] satisfies the following backward SDE within one blockd ⃗ x tn+τ = 1 2 ⃗ x tn+τ + 1 2 ∇ log ⃗ p tn+τ ( ⃗ x tn+τ ) dτ. (C.6)By subtracting (C.6) from (C.5), we obtain that d dτ y(k+1) tn,τ -⃗ x tn+τ = 1 2 y (k+1) tn,τ -⃗ x tn+τ + 1 2 s θ tn+gn(τ ) y (k) tn,gn(τ ) -s θ tn+gn(τ ) ( ⃗ x tn+gn(τ ) ) + 1 2 s θ tn+gn(τ ) ( ⃗ x tn+gn(τ ) ) -∇ log ⃗ p tn+gn(τ ) ( ⃗ x tn+gn(τ ) ) + 1 2 ∇ log ⃗ p tn+gn(τ ) ( ⃗ x tn+gn(τ ) ) -∇ log ⃗ p tn+τ ( ⃗ x tn+τ ) . (C.7)Then byd y (k+1) tn,τ ′ -⃗ x tn+τ ′ 2 = 2 y (k+1) tn,τ ′ -⃗ x tn+τ ′ ⊤ d y (k+1) tn,τ ′ -⃗ x tn+τ ′ ,and integrating for τ ′ ∈ [0, h n ], we havey (k+1) tn,τ -⃗ x tn+τ 2 = τ 0 y (k+1) tn,τ ′ -⃗ x tn+τ ′ ⊤ s θ tn+gn(τ ′ ) y (k) tn,gn(τ ′ ) -s θ tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) dτ ′ + τ 0 y (k+1) tn,τ ′ -⃗ x tn+τ ′ ⊤ s θ tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) -∇ log ⃗ p tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) dτ ′ + τ 0 y (k+1) tn,τ ′ -⃗ x tn+τ ′ ⊤ ∇ log ⃗ p tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) -∇ log ⃗ p tn+τ ′ ( ⃗ x tn+τ ′ ) dτ ′ + τ 0 y (k+1) tn,τ ′ -⃗ x tn+τ ′ 2 dτ ′ .Using AM-GM inequality and taking expectations on both sides, we further upper bound the summation above asE y (k+1) tn,τ -⃗ x tn+τ 2 ≤ 1 + 3 2h n τ 0 E y (k+1) tn,τ ′ -⃗ x tn+τ ′ 2 dτ ′ + h n 2 τ 0 E s θ tn+gn(τ ′ ) y (k) tn,gn(τ ′ ) -s θ tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) 2 dτ ′ + h n 2 τ 0 E s θ tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) -∇ log ⃗ p tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) 2 dτ ′ + h n 2 E τ 0 ∇ log ⃗ p tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) -∇ log ⃗ p tn+τ ′ ( ⃗ x tn+τ ′ ) 2 dτ ′ ≤Dt n ≤ 1 + 3 2h n τ 0 E y (k+1) tn,τ ′ -⃗ x tn+τ ′ 2 dτ ′ + h n 2 τ δ 2 ∞ + E [D tn ] + L 2 s h n 2 τ 0 E y (k) tn,gn(τ ′ ) -⃗ x tn+gn(τ ′ ) 2 dτ ′ ,where the last equality is by Assumption 3.1'.Applying Grönwall's inequality, we haveE y (k+1) tn,τ -⃗ x tn+τ 2 ≤ e (1+ 3 2hn )τ L 2 s h n 2 τ 0 E y (k) tn,gn(τ ′ ) -⃗ x tn+gn(τ ′ ) 2 dτ ′ + e (1+ 3 2hn )τ h n 2 τ δ 2 ∞ + E[D tn ] ≤ τ e (1+ 3 2hn )τ L 2 s h n 2 sup τ ′ ∈[0,τ ] E y (k) tn,τ ′ -⃗ x tn+τ ′ 2 + e (1+ 3 2hn )τ h n 2 τ δ 2 ∞ + E[D tn ] ,(C.8) and by taking supremumsup τ ∈[0,hn] E y (k+1) tn,τ -⃗ x tn+τ 2 ≤ h 2 n e hn+ 3 2 L 2 s 2 sup τ ′ ∈[0,τ ] E y (k) tn,τ ′ -⃗ x tn+τ ′ 2 + e hn+ 3 2 h n 2 h n δ 2 ∞ + E[D tn ] (C.9)Given that constant h n is sufficiently small, which ensures L 2 s h n e 5 2 hn1, iterating the above inequality for k ∈ [0 :K -1] gives us that sup τ ∈[0,hn] E y (K) tn,τ -⃗ x tn+τ 2 ≤ h 2 n e hn+ 3 2 L 2 s 2 K sup τ ∈[0,hn] E y (0) tn,τ -⃗ x tn+τ 2 + e hn+ 3 2 h n /2 1 -h 2 n e hn+ 3 2 L 2 s /2 h n δ 2 ∞ + E[D tn ] ,Notice that by Lemma A.8, we haveE y (0) tn,τ -⃗ x tn+τ 2 = E ⃗ x tn -⃗ x tn+τ 2 ≤ 3d,substituting which into (C.9) then gives us thatsup τ ∈[0,hn] E y (K) tn,τ -⃗ x tn+τ 2 ≤ 3d h 2 n e hn+ 3 2 L 2 s 2 K + e hn+ 3 2 h n /2 1 -h 2 n e hn+ 3 2 L 2 s /2 h n δ 2 ∞ + E[D tn ] ,as desired.Now it remains to bound C tn and D tn in Lemma C.4. We first bound D tn using the following lemma: 1) and thus by [102, Corollary 1] thatLemma C.5. For any n ∈ [0 : N -1], we have that E [D tn ] ≲ dϵ 2 h n . Proof. For any n ∈ [0 : N -2], we have T -t n+1 ≳ O(E ∇ log ⃗ p t N -1 +τn,m ( ⃗ x t N -1 +τn,m ) -∇ log ⃗ p t N -1 +τ ′ ( ⃗ x t N -1 +τ ′ ) 2 ≲ dϵ 2 n,m , for any τ ′ ∈ [τ n,m , τ n,m+1 ],and thusE [D tn ] = hn 0 E ∇ log ⃗ p tn+gn(τ ′ ) ( ⃗ x tn+gn(τ ′ ) ) -∇ log ⃗ p tn+τ ′ ( ⃗ x tn+τ ′ ) 2 dτ ′ = Mn m=0 τn,m+1 τn,m E ∇ log ⃗ p tn+τn,m ( ⃗ x tn+τn,m ) -∇ log ⃗ p tn+τ ′ ( ⃗ x tn+τ ′ ) 2 dτ ′ ≲ Mn m=0 dϵ 2 n,m ϵ n,m ≤ dϵ 2 h n .For n = N -1, notice that by the step size schedule (cf. Section 3.1.1) and suppose ϵ ≤ 1/2, we haveT -τ 2 ≤ T -g n (τ ) ≤ T -τ,and then again [102, Corollary 1] statesE ∇ log ⃗ p tn+ϵn,m ( ⃗ x tn+ϵn,m ) -∇ log ⃗ p tn+τ ′ ( ⃗ x tn+τ ′ ) 2 ≲ dϵ 2 n,m T -τ n,m ,and thusE D t N -1 = h N -1 0 E ∇ log ⃗ p t N -1 +gn(τ ′ ) ( ⃗ x t N -1 +gn(τ ′ ) ) -∇ log ⃗ p t N -1 +τ ′ ( ⃗ x t N -1 +τ ′ ) 2 dτ = M N -1 m=0 τn,m+1 τn,m E ∇ log ⃗ p t N -1 +τn,m ( ⃗ x t N -1 +gn(τ ′ ) ) -∇ log ⃗ p t N -1 +τ ′ ( ⃗ x t N -1 +τ ′ ) 2 dτ ′ ≲ M N -1 m=0 dϵ 2 n,m T -τ n,m ϵ n,m ≤ M N -1 m=0 dϵ 2 n,m ϵ ≲ T -t N -1 δ∞ dτ dτ ≲ dϵ 2 h N -1 .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 tn,hn that the parallelized predictor step generates samples from satisfies the following error bound:W 2 ( q tn,hn , ⃗ p tn+1 ) 2 ≲ de -K + h 2 n δ 2 ∞ + dϵ 2 h 2 n , for n ∈ [0 : N -1].Proof. By the definition of 2-Wasserstein distance, we have for any coupling of y(K)tn,hn and ⃗ x tn+hn ,W 2 ( q tn,hn , ⃗ p tn+1 ) 2 ≤ E y (K) tn,hn -⃗ x tn+hn 2 ,and thereforeW 2 ( q tn,hn , ⃗ p tn+1 ) 2 ≤ E y (K) tn,hn -⃗ x tn+hn 2 ≤ sup τ ∈[0,hn] E y (K) tn,τ -⃗ x tn+τ 2 ≤3d h 2 n e hn+ 3 2 L 2 s 2 K + e hn+ 3 2 h n /2 1 -h 2 n e hn+ 3 2 L 2 s /2 h n δ 2 ∞ + E[D tn ] ≲de -K + h 2 n δ 2 ∞ + dϵ 2 h 2 n, 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 hn L 2 s 1.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 [121] to the corrector step, which yields O(1) approximate time complexity compared to the ordinary implementation of the ULMC dynamics as in [102] . 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 tn,t † , v tn,t † ) t≥0 :du tn,t † = v tn,t † dt † dv tn,t † = -γv tn,t † dt † -∇ log ⃗ p tn+1 (u tn,t † )dt † + √ 2γdb tn,t † , (C.10)with initial condition u tn,0 ≡ y(K † )tn,hn from the predictor step and v tn,0 ∼ N (0, I d ), where (b tn,t † ) t≥0 is a Wiener process. We may also write the system of SDEs above in the following matrix form:d u tn,t † v tn,t † = 0 I d 0 -γI d u tn,t † v tn,t † - 0 ∇ log ⃗ p tn+1 (u tn,t † ) dt † + 0 0 0 √ 2γI d d b ′ tn,t † b tn,t † .We run this underdamped Langevin dynamics until the pre-determined time horizon T † . We also define the joint probability distribution of (u tn,t † , v tn,t † ) at time t as π tn,t † (u tn,t † , v tn,t † ) and its marginal on u tn,t † as π u tn,t † (u tn,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 † ) tn,n † h † ,τ † ) τ † ∈[0,h † ] as the solution to the following SDE recursively for k † ∈ [0 : K † -1]: d u (k+1) tn,n † h † ,τ † = v (k+1) tn,n † h † ,τ † dτ † , d v (k+1) tn,n † h † ,τ † = -γ v (k+1) tn,n † h † ,τ † dτ † -s tn+1 u (k † ) tn,n † h † ,gn(τ † ) dτ † + √ 2γdb tn,n † h † +τ † (C.11)with the initial conditionEQUATION   u (k † ) tn,n † h † ,τ † ≡ u tn,n † h † v (k † ) tn,n † h † ,0 ≡ v tn,n † h † for k ∈ [1 : K † ],(C.12) where u tn,n † h † := u(K † ) tn,(n † -1)h † ,h † , v tn,n † h † := v (K † ) tn,(n † -1)h † ,h † for n † ∈ [1 : N † -1],andu tn,0 = y (K)tn,hn , v tn,0 ∼ N (0, I d ).We define the joint probability distribution of ( u tn,t † , v tn,t † ) at time t as π tn,t † ( u tn,t † , v tn,t † ) and its marginal on u tn,t † as π u tn,t † ( u tn,t † ). We will also denote the resulting probability distribution of π u tn,T † as q tn+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:d   u (k † ) tn,n † h † ,τ † v (k † ) tn,n † h † ,τ †   =   0 I d 0 -γI d   u (k † ) tn,n † h † ,τ † v (k † ) tn,n † h † ,τ †   - 0 s tn+1 u (k † ) tn,n † h † ,gn(τ † )   dτ † + 0 0 0 √ 2γI d d b ′ tn,n † h † +τ † b tn,n † h † +τ † . (C.13)Define the time-dependent matrix G(•) asG(t † ) := I d 1-e -γt † γ I d 0 e -γt † I d = exp 0 I d 0 -γI d t † , (C.14) satisfying that d dt † G(t † ) = 0 I d 0 -γI d G(t † ) = G(t † ) 0 I d 0 -γI d .Now we multiply G(-τ † ) on both sides of (C.13) to obtain:d   G(-τ † )   u (k † ) tn,n † h † ,τ † v (k † ) tn,n † h † ,τ †     = -G(-τ † ) 0 s tn+1 u (k † ) tn,n † h † ,gn(τ † ) dτ † + G(-τ † ) 0 0 0 √ 2γI d d b ′ tn,n † h † +τ † b tn,n † h † +τ † .Integrating on both sides from 0 to h † and multiplying G(h † ) on both sides, we have  u (k † ) tn,n † h † ,τ v (k † ) tn,n † h † ,τ   -G(h † )   u (k † ) tn,n † h † ,0 v (k † ) tn,n † h † ,0   = - h † 0 G(h † -τ † ′ ) 0 s tn+1 u (k † ) tn,n † h † ,g(τ † ′ ) dτ † ′ + h † 0 G(h † -τ † ′ ) 0 0 0 √ 2γI d d b ′ tn,n † h † +τ † ′ b tn,n † h † +τ † ′ = - M † -1 m † =0 (m † +1)ϵ † m † ϵ † G(h † -τ † ′ )dτ † ′ 0 s tn+1 u (k † ) tn,n † h † ,m † ϵ † + M † -1 m † =0 (m † +1)ϵ † m † ϵ † G(h † -τ † ′ ) 0 0 0 √ 2γI d d b ′ tn,n † h † +τ † ′ b tn,n † h † +τ † ′ = - M † -1 m † =0 G(ϵ † ) -I d G((M † -m † -1)ϵ † ) 0 s tn+1 u (k † ) tn,n † h † ,m † ϵ † + M † -1 m † =0 (m † +1)ϵ † m † ϵ † G(h † -τ † ′ ) 0 0 0 √ 2γI d d b ′ tn,n † h † +τ † ′ b tn,n † h † +τ † ′ .By Itô isometry, we have(m † +1)ϵ † m † ϵ † G(h † -τ † ′ ) 0 0 0 √ 2γI d d b ′ tn,n † h † +τ † ′ b tn,n † h † +τ † ′ ∼N 0, 0 0 0 √ 2γI d G((M † -m † -1)ϵ † ) ⊤ G(ϵ † ) -I d ⊤ G(ϵ † ) -I d G((M † -m † -1)ϵ † ) 0 0 0 √ 2γI d ∼ 0 N 0, 2γ(1 + γ -2 )(1 -e -γϵ † ) 2 e -2γ(M † -m † +1)ϵ † ) I d ,as desired Definition C.10 (Interpolating corrector process). For any n † ∈ [0 : N † -1], we define the interpolating corrector process ( u(k † ) tn,n † h † ,τ † ) τ † ∈[0,h † ]as the solution to the following SDE recursively fork † ∈ [0 : K † -1]: d u (k+1) tn,n † h † ,τ † = v (k+1) tn,n † h † ,τ † dτ † , d v (k+1) tn,n † h † ,τ † = -γ v (k+1) tn,n † h † ,τ † dτ † -s tn+1 u (k † ) tn,n † h † ,gn(τ † ) dτ † + √ 2γdb tn,n † h † +τ † (C.15)with the initial conditionEQUATIONu tn,0 = y (K)tn,hn , v tn,0 ∼ N (0, I d ).We define the joint probability distribution of ( u tn,t d b tn,n † h † +τ † = db tn,n † h † +τ † -ϕ tn,n † h † (τ † )dτ † , (C.17)whereϕ tn,n † h † (τ † ) = 1 √ 2γ s tn+1 ( u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 ( u (K † ) tn,n † h † ,τ † ) (C.18)and computing the Radon-Nikodym derivative of the measure P with respect to Q asdP dQ = exp h † 0 ϕ tn,n † h † (τ † ) ⊤ db tn,n † h † +τ † - 1 2 h † 0 ϕ nh (τ † ) 2 dτ † ; (C.19)3. Concluding that (C.15) at iteration K † under the measure Q satisfies the following SDE:   d u (K † ) n † h † ,τ † = v (K † ) n † h † ,τ † dτ † d v (K † ) tn,n † h † ,τ † = -γ v (K † ) tn,n † h † ,τ † dτ † -∇ log ⃗ p tn+1 ( u (K † ) n † h † ,τ † )dτ † + √ 2γd b tn,n † h † +τ † , (C.20) with ( b tn,n † h † +τ † ) τ † ≥0being a Wiener process under the measure P . If we replace ( u (K † ) n † h † ,τ † , v (K † ) n † h † ,τ † ) by (u tn,n † h † +τ † , v tn,n † h † +τ † ),t ∈ [n † h † , (n + 1)h † ]: du tn,n † h † +τ † = v tn,n † h † +τ † dτ † dv tn,n † h † +τ † = -γv tn,n † h † +τ † dτ † -∇ log ⃗ p tn+1 (u tn,n † h † +τ † )dτ † + √ 2γd b tn,n † h † +τ † .(C.21) We further define the joint probability distribution of (u tn,t † , v tn,t † ) at time t as π tn,t † (u tn,t † , v tn,t † ) and its marginal on u tn,t † as π u tn,t † (u tn,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 du * tn,n † h † +τ † = v * tn,n † h † +τ † dτ † , dv * tn,n † h † +τ † = -γv * tn,n † h † +τ † dτ † -∇ log ⃗ p tn+1 (u * n † h † +τ † )dτ † + √ 2γd b tn,n † h † +τ † , (C.22) with the initial condition u * tn,n † h † ∼ ⃗ p tn+1 and v * tn,n † h † ∼ N (0, I d ).We define the joint probability distribution of (u h † π tn,n † h † ) ≤D KL (π tn,(n-1)h † π tn,(n-1)h † ) + N † -1 n=0 D KL (π tn,n † h † :(n+1)h † π tn,n † h † :(n+1)h † ) ≤D KL (π tn,(n-1)h † π tn,(n-1)h † ) + 1 4γ E P h † 0 s tn+1 ( u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 ( u (K † ) tn,n † h † ,τ † ) 2 dτ † . (C.23)By triangle inequality, we have Now we proceed to bound the terms in the error decomposition (C.24). We first bound the F tn,n † h † term by the following lemma: Lemma C.13. For any n ∈ [0 : N -1] and τ † ∈ [0, h † ], we haveh † 0 s tn+1 ( u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 ( u (K † ) tn,n † h † ,τ † ) 2 dτ † ≤5 h † 0 s tn+1 ( u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -s tn+1 ( u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) 2 dτ † +5 h † 0 s tn+1 ( u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -s tn+1 (u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † ) 2 dτ † +5 h † 0 s tn+1 (u * tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 (u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † ) 2 dτ † +5 h † 0 ∇ log ⃗ p tn+1 (u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 (u * tn,n † h † ,τ † ) 2 dτ † +5 h † 0 ∇ log ⃗ p tn+1 (u * tn,n † h † +τ † ) -∇ log ⃗ p tn+1 ( u (K † ) tn,n † h † ,τ † ) 2 dτ † ≤ 5L 2 s h † 0 u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † -u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † 2 dτ † +5 L 2 s h † 0 u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † -u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † 2 dτ † + L 2 p h † 0 u (K † ) tn,n † h † ,τ † -u * tn,n † h † +τ † 2 dτ † :=E tn ,n † h † +5h † δ 2 ∞ + 5L 2 p h † 0 u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † -u * tn,n † h † +τ †E P u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † -u * tn,n † h † +τ † 2 ≤ dϵ † 2 ,and thereforeE P F tn,n † h † ≤ dh † ϵ † 2 .Proof. By the definition of (u * tn.n † h † +τ , v * tn,n † h † +τ ) as the stationary underdamped Langevin dynamics (C.22), we haveE P u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † -u * tn,n † h † +τ † 2 = E P   τ † ⌊ τ † ϵ † ⌋ϵ † v * tn,n † h † +τ † ′ dτ † ′ 2   ≤ ϵ † τ † ⌊ τ † ϵ † ⌋ϵ † E P v * tn,n † h † +τ † ′ 2 dτ † ′ ≤ dϵ † 2 ,where the first inequality follows from Cauchy-Schwarz inequality and the last inequality is by the fact thatv * tn,n † h † +τ † ′ ∼ N (0, I d ), for any τ † ′ ∈ [0, h † ].Consequently, we haveE P F tn,n † h † = h † 0 E P u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † -u * tn,n † h † +τ † 2 dτ † ≤ dh † ϵ † 2 .The term E tn,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 † ] E P u (K † ) tn,n † h † ,τ † -u * tn,n † h † +τ † 2 ≲ W 2 2 ( q tn,hn , ⃗ p tn+1 ),and thereforeE P E tn,n † h † ≲ h † (L 2 s + L 2 p )W 2 2 ( q tn,hn , ⃗ p tn+1 ).Proof. Recall that under the measure P , u(K † ) tn,n † h † ,τ † follows the dynamics of u tn,n † h † ,τ † (C.21) for τ † ∈ [0, h † ], which coincides with that of u * tn,n † h † +τ † .As the only difference between the two processes u tn,n † h † ,τ † and u * tn,n † h † +τ † is the initial condition, we can invoke Lemma 10 proved in [102] to deduce thatE P u (K † ) tn,n † h † ,τ † -u * tn,n † h † +τ † 2 ≲ W 2 2 (π tn,n † h † , ⃗ p tn+1 ),where the assumption that γ ≲ L W 2 2 (π tn,n † h † , ⃗ p tn+1 ) ≤ E u tn,n † h † -u * tn,n † h † 2 ≤W 2 2 (π tn,0 , ⃗ p tn+1 ) = W 2 2 ( q tn,hn , ⃗ p tn+1 ),where the last equality is again by [102, Lemma 10] .Therefore, we haveE P E tn,n † h † = h † 0 E P L 2 s u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † -u * tn,n † h † +⌊ τ † ϵ † ⌋ϵ † 2 + L 2 p u (K † ) tn,n † h † ,τ † -u * tn,n † h † +τ † 2 dτ † ≤h † (L 2 s + L 2 p )W 2 2 ( q tn,hn , ⃗ p tn+1 ).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:sup τ † ∈[0,h † ] E P u (1) tn,n † h † ,τ † -u (0) tn,n † h † ,τ † 2 ≤ 5L 2 s h † e (3+γ)h † 2γ sup τ † ∈[0,h † ] E P u (K † -1) tn,n † h † ,τ † -u (K † ) tn,n † h † ,τ † 2 + 5h † e (3+γ)h † 2γ E P E tn,n † h † + h † δ 2 ∞ + L 2 p F tn,n † h † + h † 2 e (3+γ)h † 3γd + M 2 s + h † e 2h † d. Proof. Let µ tn,n † h † ,τ † := u (1) tn,n † h † ,τ † -u (0) tn,n † h † ,τ † and ν tn,n † h † ,τ † := v (1) tn,n † h † ,τ † -v (0) tn,n † h † ,τ † . Then for k = 0, we may rewrite (C.15) as follows    dµ tn,n † h † ,τ † = ν tn,n † h † ,τ † + v (0) tn,n † h † ,τ † dτ † dν tn,n † h † ,τ † = -γ(ν tn,n † h † ,τ † + v (0) tn,n † h † ,τ † )dτ † -s tn+1 ( u (0) tn,n † h † ,τ † )dτ † + √ 2γdb tn,n † h † +τ † (C.25) On the one hand, by using the first equation in (C.25), we may compute the derivatived dτ † ′ µ tn,n † h † ,τ † ′ 2 = 2µ ⊤ tn,n † h † ,τ † ′ ν tn,n † h † ,τ † ′ + v (0) tn,n † h † ,τ † ′ and integrate it for τ † ′ ∈ [0, τ † ], which yields µ tn,n † h † ,τ † 2 = 2 τ † 0 µ ⊤ tn,n † h † ,τ † ′ (ν tn,n † h † ,τ † ′ + v (0) tn,n † h † ,τ † ′ )dτ † ′ ≤2 τ † 0 µ tn,n † h † ,τ † ′ 2 dτ † ′ + τ † 0 ν tn,n † h † ,τ † ′ 2 dτ † ′ + τ † 0 v (0) tn,n † h † ,τ † ′ 2 dτ † ′ .Applying Gronwall's inequality, we haveµ tn,n † h † ,τ † 2 ≤ e 2τ † τ † 0 ν tn,n † h † ,τ † ′ 2 dτ † ′ + τ † 0 v (0) tn,n † h † ,τ † ′ 2 dτ † ′ .We then take expectation with respect to the path measure P and then the supremum with respect toτ † ∈ [0, h † ], implying that sup τ † ∈[0,h † ] E P µ tn,n † h † ,τ † 2 ≤ sup τ † ∈[0,h † ] e 2τ † τ † 0 E P ν tn,n † h † ,τ † ′ 2 dτ † ′ + e 2τ † τ † 0 E P v (0) tn,n † h † ,τ † ′ 2 dτ † ′ ≤ h † e 2h † sup τ † ∈[0,h † ] E P ν tn,n † h † ,τ † ′ 2 + h † e 2h † d. (C.26)On the other hand, by applying Itô's lemma and plugging in the expression of b tn,n † h † +τ † given by (C.17), we haved ν tn,n † h † ,τ † 2 = -2γ ν tn,n † h † ,τ † 2 + 2γν ⊤ tn,n † h † ,τ † v (0) tn,n † h † ,τ † + 2ν ⊤ tn,n † h † ,τ † s tn+1 u (0) tn,n † h † ,τ † -2γd dτ † +2ν ⊤ tn,n † h † ,τ † 2γ d b tn,n † h † +τ † + ϕ tn,n † h † (τ † )dτ † , (C.27)Then similarly, we may compute the derivative of ν tn,n † h † ,τ † 2 , integrate it for τ † ∈ [0, h † ], and take the supremum with respect to τ † to obtainE P ν tn,n † h † ,τ † 2 =E P - τ † 0 2γ ν tn,n † h † ,τ † ′ 2 + 2γν ⊤ tn,n † h † ,τ † ′ v (0)tn,n † h † ,τ † ′ -2γd dτ + M 2 s + 2γd τ † =(1 + γ) τ † 0 E P ν tn,n † h † ,τ † ′ 2 dτ † ′ + τ † 0 E P ϕ tn,n † h † (τ † ′ ) 2 dτ † ′ + τ † 3γd + M 2 s ,where in the last equality, we used the initialization of the auxiliary corrector process v (0) tn,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 sup τ † ∈[0,h † ] E P ν tn,n † h † ,(K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 ( u (K † ) tn,n † h † ,τ † ) 2 dτ † ≤5L 2 s h † 0 u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † -u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † 2 dτ † + 5E tn,n † h † + 5h † δ 2 ∞ + 5L 2 p F tn,n † h † ,to obtain that sup EQUATIONn † h † ,τ † -v (k † ) n † h † ,τ † 2 = - τ † 0 2γ v (k+1) n † h † ,τ † ′ -v (k † ) n † h † ,τ † ′ 2 dτ † ′ -2 τ † 0 v (k+1) n † h † ,τ † ′ -v (k † ) n † h † ,τ † ′ ⊤ s tn+1 ( u (k † ) tn,n † h † ,⌊ τ † ′ ϵ † ⌋ϵ † ) -s tn+1 ( u (k-1) tn,n † h † ,⌊ τ † ′ ϵ † ⌋ϵ † ) dτ † ′ ≤ 1 2γ τ † 0 s tn+1 ( u (k † ) tn,n † h † ,⌊ τ † ′ ϵ † ⌋ϵ † ) -s tn+1 ( u (k-1) tn,n † h † ,⌊ τ † ′ ϵ † ⌋ϵ † ) 2 dτ † ′ ≤ L 2 s 2γ τ † 0 u (k † ) tn,n † h † ,⌊ τ † ′ ϵ † ⌋ϵ † -u (k-1) tn,n † h † ,⌊ τ † ′ ϵ † ⌋ϵ † 2 dτ † ′ .And then taking the supremum with respect to τ † ∈ [0, h † ] on both sides above implies sup τ † ∈[0,h † ] E P v (k+1) n † h † ,τ † -v (k † ) n † h † ,τ † 2 ≤ h † L 2 s 2γ sup τ † ∈[0,h † ] E P u (k † ) tn,n † h † ,τ † -u (k-1) tn,n † h † ,τ †E P u (K † ) n † h † ,τ † -u (K † -1) n † h † ,τ † 2 ≤ L 2 s h † 2 e h † 2γ sup τ † ∈[0,h † ] E P u (K † -1) tn,n † h † ,τ † -u (K † -2) tn,n † h † ,τ † 2 ≤ L 2 s h † 2 e h † 2γ K † -1 sup τ † ∈[0,h † ] E P u (1) tn,n † h † ,τ † -u (0) tn,n † h † ,τ † 2 ≤ L 2 s h † 2 e h † 2γ K † -1 5h † e (3+γ)h † 2γ E P E tn,n † h † + h † δ 2 ∞ + L 2 p F tn,n † h † + L 2 s h † 2 e h † 2γ K † -1 h † 2 e (3+γ)h † 3γd + M 2 s + h † e 2h † d + L 2 s h † 2 e h † 2γ K † -1 5L 2 s h † e (3+γ)h † 2γ sup τ † ∈[0,h † ] E P u (K † -1) tn,n † h † ,τ † -u (K † ) tn,n † h † ,τ † 2 ,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 parametersT † = O(1), N † = O(1), h † = Θ(1)M † = Θ(d 1/2 δ -1 ), ϵ † = Θ(d -1/2 δ), K † = O(log(dδ -2 ))and letL 2 s h † 2 e h † 2γ 1, γ ≲ L -1/2 p , T † ≲ L -1/2 p ∧ L -1/2 s , δ ∞ ≲ δthen the distribution π tn,T † satisfies the following error bound:EQUATION2 ( q tn,hn , ⃗ p tn+1 ) + δ 2 , 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)E P h † 0 s tn+1 ( u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † ) -∇ log ⃗ p tn+1 ( u (K † ) tn,n † h † ,τ † ) 2 dτ † ≤5L 2 s h † sup τ † ∈[0,h † ] E P u (K † -1) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † -u (K † ) tn,n † h † ,⌊ τ † ϵ † ⌋ϵ † 2 + 5E P E tn,n † h † + h † δ 2 ∞ + L 2 p F tn,n † h † ≤5 1 + L 2 s h † C K † 5h † e (3+γ)h † 2γ E P E tn,n † h † + h † δ 2 ∞ + L 2 p F tn,n † h † +5L 2 s h † C K † h † 2 e (3+γ)h † 3γd + M 2 s + h † e 2h † d ,where the last inequality is by Lemma C. 16 . We further substitute Lemma C. 14 ) + N † h † δ 2 ∞ + dN † h † ϵ † 2 + e -K † N † h † 2 d =T † W 2 2 ( q tn,hn , ⃗ p tn+1 ) + T † δ 2 ∞ + dT † ϵ † 2 + e -K † T † h † d.Then, it is straightforward to see that when the following order of the parameters holdsT † = O(1), h † = Θ(1), N † = O(1), ϵ † = Θ(d -1/2 δ), M † = O(d 1/2 δ -1 ), K † = O(log(dδ -2 ))and δ ∞ ≤ δ, we have Dhn 1. 


Decision: This paper presents a novel parallel sampling method for diffusion models, achieving a poly-logarithmic time complexity. The authors leverage Picard Iterations (PI) to solve both the SDE and probability flow ODE associated with diffusion models, enabling efficient parallelization of the sampling process.

The reviews highlight the paper's strengths in addressing a significant problem, presenting a novel approach, and providing a solid theoretical analysis.

However, reviewers suggest that a more detailed comparison with other parallel methods like the randomized midpoint method and ParaDiGMS should be included. Additionally, they were also disappointed with the absence of empirical verification of the theoretical results. While a theory paper, they strongly encourage the inclusion of experiments to demonstrate the practical utility of the proposed methods. I concur with them and strongly recommend the authors include such experiments in the final version of the paper.

Overall, the paper presents an excellent contribution to the field of diffusion models. However, the authors should address carefully the concerns raised by reviewers regarding clarity, empirical validation, and comparison with existing methods while preparing the final version of the paper. This will significantly enhance the paper's impact.

Reviews: review: paper_summary: paper_summary: In 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.

[1] Shih et al. (2024) -- Parallel Sampling of diffusion models

[2] Anari et al. (2024) -- Fast parallel sampling under isoperimetry
strength_weakness: strength_weakness: * 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. 

* 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 
stochastic 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).

* 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). 

* 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.

[1] Chen et al. (2023) -- "Restoration-Degradation Beyond Linear Diffusions: A Non-Asymptotic Analysis For DDIM-Type Samplers"
questions: questions: * "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. 

* 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? 

* See "Weaknesses" for other questions and concerns. 

[1] Shih et al. (2024) -- Parallel Sampling of diffusion models

[2] De Bortoli (2022) -- Convergence of denoising diffusion models under the manifold hypothesis

[3] Yang et al. (2023) -- The Convergence of Variance Exploding Diffusion Models under the Manifold Hypothesis
limitations: limitations: In 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.
paper_summary: paper_summary: Denoising 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}.

An 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.

This 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.
strength_weakness: strength_weakness: I 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.
questions: questions: Can you provide some intuition on how the constants scale, and how long does it take for the asymptotic to kick in?

Can you provide some intuition on how to tune hyperparameters, including a number of blocks, discretization steps per block, and depth?
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: paper_summary: In 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.
strength_weakness: strength_weakness: - 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.
- Please compare the total query complexity of the proposed method to that of sequential methods.

Minor:
- Line 85: Please add a reference for the Picard iteration.
- Line 89: Please include a reference for the claim of exponentially fast convergence.
- Last equation on page 4: Specify the distribution that is being assumed for this equation.
- Lemma A.6: Ensure that $_F$ is included.
- Provide an explanation regarding the different parameter choices for the predictor and corrector in Theorem 3.5.
questions: questions: See Weaknesses
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Please see the Weakness part.
limitations: limitations: The authors adequately addressed the limitations and societal impact.

Comments: title: Respone to Reviewer GP1D
Thank you so much!
comment: We would like to thank the reviewer for the response and appreciation for our work. As for now, we can only obtain the memory reduction for the ODE case via mathematical proof. To the best of our knowledge, there has been no empirical work comparing the ODE implementation with the SDE implementation under the parallel sampling framework. We think it would be interesting to compare them to see if such memory reduction can be achieved in future work.
Thanks a lot for your rebuttal. I would like to keep my score. I think the paper is good. Again, I am still a bit confused by the ODE memory reduction (and the provided answer does not really clarify things for me). I guess my question is the following. Do you think that the memory reduction is an artifact of the proof (i.e. some of the bounds are loose in the SDE case) or do you think this is something that could be actually verified experimentally?
We would like to thank the reviewer again for your time and efforts! Your insightful suggestions and comments have greatly helped improve the quality of our manuscript. In case you have any other question, please don't hesitate to let us know.
Thanks for the detailed discussion on the previous works. It would be helpful to add these discussions and highlight the contribution of the technique in the next version. Since this discussion addresses my concerns, I will raise my score to $6$.