Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity
Authors: Haoxuan Chen
Yinuo Ren
Lexing Ying
Grant M. Rotskoff
Venue: NeurIPS 2024 spotlight (2024)
Keywords: diffusion model
parallel sampling
stochastic differential equations
probability flow ode
Primary area: diffusion_based_models
TL;DR: We propose new parallel inference algorithms for diffusion models using parallel sampling and rigorously prove that our algorithms enjoy sub-linear inference cost w.r.t. data dimension for both SDE and probability flow ODE implementations.

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
Haoxuan Chen" ICME
Yinuo Ren* ICME
Stanford University haoxuancdstanf ord. edu
Stanford University
yinuorendstanforr edu
Lexing Ying
Grant M. Rotskoff
Department of Mathematics and ICME Stanford University
Department of Chemistry and ICME
Stanford University
1 texing@stanford...
rotskoff@stanford edu
Abstract
Diffusion models have become leading method for generative modeling of both image and scientific data. As these models are costly train and evaluate, re- ducing the inference cost for diffusion models remains major goal. Inspired by the recent empirical success in accelerating diffusion models via the parallel sam- pling technique [M]), propose divide the sampling process into O(1) blocks with parallelizable Picard iterations within each block. Rigorous theoretical anal-
ysis reveals that our algorithm achieves (pooyy log d) overall time complexity, marking the first implementation with provable sub-linear complexity w.r.f. the data dimension d. Our analysis is based on a generalized version of Girsanov's theorem and compatible with both the SDE and probability flow ODE imple- mentations. Our results shed light on the potential of fast and efficient sampling of high digeenimeniional data on fast-evolving modern large-memov memory GPU clusters.
Introduction
Diffusion and probability flow based models [2-10]] are now state-of-the-art in many fields, such computer vision and image generation [[22-22] natural language processing [23, 23]], audio and video generation 255-299,, optimization [60 B1I], sampling and learning of fixed classes of distribu- tions [32] [.] solving high high-dimensional partial differential equations [21 266], and more recently several applications in physical, chemical and biological fields [47 633.. For a more comprehen-
sive 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) [[], score-based generative models (SGMs) [S]], diffusion schrödinger bridges [672], stochastic interpolants and flow matching [2- etc., the recurring idea is to design 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 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
"qqual contribution, alphabetical order. Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Work
Implementation
Measure
Approx. Time Complexity
Theorem 2]
SDE
TV(po, Tr)
(dd--1
Theorem 2]
SDE
DKL (Py (T
(
Corollary 1]
SDE
DKL (pr) (191--)
O(d8
Theorem 3]
ODE w/MLCC correction
TV 9r-
(Vd8-1)
Theorem 1
SDE w/parallel sampling
DKL. (P [I
(poly log(d6--))
Theorem --
ODE w/parallel sampling
(P Qr-n)²
(poly log (dd (og(dd-)) (655-))
Table I: Comparison of the approximate time complexity (ef. Definition 2.II) of different implemen-
tations of diffusion models. 17 is small parameter that controls the smooth approximation of the data distribution (cf. Section
made accelerate this process by resorting to higher-order randomized numerical schemes [69]-
79], augmented dynamics [8Q], adaptive step sizes [KT]. operator learning [82], restart sampling [8
self-consistency [84-87]] and knowledge distillation [88 [0]. Recently, several empirical works [M]
91- leverage the Picard iteration and triangular Anderson acceleration to parallelize the sampling
procedure of diffusion models and achieve empirical success large-scale image generation tasks.
Some other recent work [95], |6] 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 arbitrary data distribution on IRd to 6-accuracy, which has been analyzed extensively in the literature (98-1155). In terms of the dependency on the dimension d, the current state-of-the-arr result for the SDE implementation of diffusion models is (dd [IOM] improved from the previous (d²) bound [_0+]] || gives ((vd) bound for the probability flow ODE implementation by considering predictor-corrector scheme with the under- damped Langevin Monte Carlo (UMLC) algorithm.
In this work, we aim to provide parallelization strategies, rigorous analysis, and theoretical guaran- tees 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 scale with ((1d)) for the SDE implementation and (1//Vd) for the prob- ability flow ODE implementation. We show that the inference process can be first divided into ((1) blocks with parallelizable evaluations of the score function within each, and thus reduce the overall time complexity O(poly log d). We provide the first implementation of diffusion mod- els with poly- lobaritgaii complexity, significant improvement over the current state-of-the-arr polynomial results that sheds light on the potential fast and efficient sampling of high- nigh-dimensional distributions with diffusion models on fast-developing memory- efficiiet modern GPU clusters.
1.1 Contributions
We propose parallelized inference algorithms for diffusion models in both the SDE and proba- bility flow ODE implementations (PIADM-SDE/ODE) with exponential integrators, shrinking step size scheme towards the data end, and the early stopping technique;
We provide a rigoros convergence analysis of PIADM-SDE, showing that our parallelization strategy yields diffusion model with Opoly 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 (cl) to (dd3/2) while maintaining the logatthgai time complexity.
2 Preliminaries
In this section, we briefly recapitulate the framework of score-based diffusion models, define nota- tions, and discuss related work.

---

2.1 Diffusion Models
In score-based based diffusion models, one considers a diffusion process (2x)520 IRd governed by the following stochastic differential equation (SDE):
dæ, = (3(xx. + a,dws, with Do ~ po+
(2.1)
where (w,) (w2)s20 is standard Brownian motion, and po is the target distribution that would like to sample from. The distribution of æ, is denoted by Ps. Once the drift (-). the diffusion coefficient
ass and sufficiently large time horizon T are specified, (2.II) also corresponds backward process (oe)o for another arbitrary diffusion coefficient (Us)s20 [[]6]]:
dē. (z,)) + ,, + J v (. (T+)
dt + Ü,dw;,
(2.2)
2
where F2 denotes -r- with Po = pr and Pr = Po-
For notational simplicity, we adopt simple choice of the drift and the diffusion coefficients in what follows: (x(x) = j o; = La, and v = v.a. under which (2.I) is an Manstein-Uhlenbede (OU)
process converging exponentially to its stationary distribution, i.e. pr se pr == NO, Ia), and (2.
and (2 reduce the following form:
1 dæ, - ds ls + dww;; and dē.
1 Er
1 0²
log (i.)) de + edwer
(2.3)
2
In practice, the score function ( is often estimated by neural network (NN) 8 (x.). where 0 represents its parameters, by minimizing the denoising score-matching loss [117 I18]:
((0)= = Excepe
V Vlogpt(z;) s (xx))
1']
(2.4)
20e -,
Bao
E B (ejo
1 1-c
and the backward process in (2 is approximated by the following SDE thereafter:
2 2 of (8r) dt + udw,, with yo ~ NVO, Ia).
dy: = yt +
(2.5)
2
Implementations. Diffusion models admit multiple implementations depending on the choice of the parameter V in the backward process (222). The SDE implementation with v = widely used in the literature for its simplicity and efficiency [[Q]. while recent studies [للل]] claim that the probability flow ODE implementation with v = may exhibit better time complexity. We refer to [ [
for theoretical and [[20) [21]] 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 - and several recent studies [125 [27]. 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æ, = fx ) and we would like to solve it for [0, TJ. then the Picard iteration is defined as follows:
xD + I' (a( for k [0: K 1]. K-1]
Fo, and x, (k+1)
x
(2.6)
Under assumptions on the Lipschitz continuity of ft, the Picard iteration converges to the true so- lution exponentially fast, in the sense that |||| x; ((0,77) 8 with K Oloo 8-1) iterations. Unlike high-order ODE solvers, the Picard itcration is intrinsically parallelizable: for any 1. E [0, T), the computation of at (k+1) relies merely on the values of the most recent iteration a; (k) With sufficient computational sources parallelizing the evaluations of f. the computational cost of solving the ODE no longer scales with 7 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 [[30]] Roughly speaking, it is proposed to simulate the Langevin diffusion process dæ. = -V )dt + dur, with the following iteration resembling (2.66):
o0,
(1+1)) and F = 20
(1) )ds + wi, for k E [0: VV(x (
1}, (2.7)
where all iterations share common Wiener process (wi))20-
It is shown that for well-conditioned log-concave distributions, parallelized LMC would achieve an iteration depth of K Öppoly log d) that matches the indispensable time horizon T (poly log d) to achieve exponential ergodicity (ef. [[L3], Theorem 13]). This promises a signif- icant speedup in sampling high-dimensions 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 (-). which is often parametrized by large-scale NNs. Despite several unavoidable costs involv- ing pre- and post-processing following notion of the approximate time complexity of the inference process of diffusion models:
data storage and retrieval, and arithmetic operations, we define the
Definition 2.1 (Approximate time complexity). For a specific implementation of diffusion mod- els (2.S), we define the approximate time complexity of the sampling process as the number of unparallelizable evaluations of the learned NN-bused score function (.).
This definition coincides with the notion of the number of steps required to reach a certain accuracy in (IOA, лΛоI), iteration complexity in [I00]] [[]. elc. in the previous theoretical studies of diffusion models. We have adopted this notion in Table MI for a comparison of the current state-of-the-ait re- sults 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 (d). Should confusion occur, we omit the dependency of the complexities above on the accuracy threshold 6, etc., during our discussion, as we focus on applications of diffu- sion models to dimenssoal 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 henefit of the parallel sampling technique in this scenario has been recently confirmed by up to 14x acceleration achieved by the ParaDiGMS algorithm [M]] 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 Parallelized Inference Algorithms for Diffusion Models with both the SDE and probability flow ODE implementations, namely the PIADM-SDE (Algorithm MI] 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 B.ll and Section B. respectively. Due to the large number notations used in the presentation, we give an overview of notations in Appendix A. for readers' convenience.
3.1 SDE Implementation
We first focus on the approximation, parallelization strategies, and error analysis of diffusion mod- els with the SDE implementation, iee. the forward and backward process (2.3) and its approximata- tion (255) with e = 1. We will show that PIADM-SDE achieves an Ö(poly log d) approximate time complexity with (dd) space complexity.

---

Outer Iterations N (log d) blocks
g0 ~ (O,Iaa)
F. Pdlata
hn
hN-1
(1)
Mn (d) ((Vd) parallalizable steps
Cn,D
Cn,M-1
k k=
Inner Iterations kt K Ollog d) depth
k=1
== K
O(d-1) or (dd-1/2)
Figure 1: Illustration of WRADM-SDE/ON The outer iterations are divided into O d) blocks of O1)) length. Within each block, the inner iterations are parallelized with (d) steps for SDE (ef. Theorem B.3), or Ö (d for probability flow ODE implementation (ef. Theorem D. The overall approximate time complexity is KN (poly log d). brown, green, blue, and red curves represent Mn,
the computation graph at f in + Ta,m for m 1,2, M,
3.1.1 Algorithm
PIADM-SDE is summarized in Algorithm D and illustrated in Figure M. The main idea behind algorithm is the fact that (2.5) can be efficiently solved by the Picard iteration within period of O(1) length, transferring (d) sequential computations to parallelizable iteration of depth Ö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 7 is first segmented into N blocks of length (h,, W with each hn < h := T/N = [((1,, forming a grid (t,))_ with t, = Σド= hj. For Ma
any n E [0: N:N- 1]. the n-th block is further discretized into a grid (Tmm))m with Tn,0 = 0 and Tn,M, hn- We denote the step size of the m-th step in the n-th block En.m -n,m+1 T
and the total number of steps in the n-th block as M,.
For the first N - blocks, we simply use the unique discretization, ie. hn = h, 6a.m = and
Mn = M h/c, for n E [0 N 2) and m (0 M - 1]. Following [[004 [00I], curb the potential blow-up of the score function of the order
T, which is shown by [0107] for 0 < T to be
T V log (ū,) V log ( S d
2
T
apply early stopping time iN = T 7, where n is chosen in way such that the O(VT) Wasserstein distance between Pr and its smoothed version Pr-, 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 hy-1
ö, and discretize the interval [IN=1, [[N - 1)h, T v] into grid (mm- with step sizes (EN-1, ) My satisfying
Mx-1
6N-1m S A (h ----NN-1111
(3.1)
As shown in Lemma B this exponential decaying step size scheme towards the data end is crucial bound the discretization error in the last block.
For the simplicity of notations, we introduce the following indexing function: for r E [fns tn+ we
define (T) to be the unique integer such that E/-(1 enjj S (r))+1 En.j- We also define

---

Algorithm 1: PIADM-SDE
Input: Yo ~ 9o = (0, Ia). a discretization scheme (T, (ha) and (Tu.m)nE(La [I:N].me satisfying (),,,, the depth of iteration K, the learned NN-based score function s (-). Output: sample Vee qeN Pr.
for n = 0 to N -1 do
(0)
Tm..Tn.m yt.. 5. ~O, Ia) for m E [0 M,] parallet;
3
for h = 0 to 1o
(R)
Vtn.0 f- yr,,
for = 0 to Mn in parallel do
(k+1)
To (k)
Dt 7m.m t
0l..0
(3.4)
+ 0 1
Tmoill
2 (E
1) t
(K) tnan
( ) Veend -oo-i-14 16:
end
end
Dtni1 f Van.ta.Ma
end
a piecewise function g such that (T) = -Eaat) en.j- It is easy to check that under the uniform discretization for n E [1 N -N-11 1], we have (T) = -Tr/c] and (T) [t/elc.
Exponential Integrator. For each step T E [t,, + Tn.m: in + mn. we use the following expo- nential integrator scheme [7]] as the numerical discretization of the SDE
(n..// 1 8 +T (it, +rm. Vetn.m 1g.
where & (O0, Ia). Lemma B shows its equivalence approximating (2.55) as
dft... -
1 + 8 mm (Dinsto..m ) dr + dw.,+r for T E [Tnms Tn.m+1]-
(3.2)
Remark 3.1. One could also implement a straightforward Euler-Maruyana scheme instead of the exponential integrator (B.4, where additional high-order discretization error term would emerge [НA Theorem 11, which believe would not affect the overall Oppoly log d) time com- plexity with parallel sampling.
Picard Iteration. Within each block, we apply Pieard iteration of depth K. As shown by Lemma B.3, the discretized scheme (3.4) implements the following iteration for k E [0: K -1]
1]:
(1+11) dÿ
1 (k+++) 2 Vino
+ 6 (r)
( ( dr + dw,, tr for T 00, h,].
(3.3)
(K) We denote the distribution of yi..r by +7, As proved in Lemma B.d. the iteration above would converge to ()=22) in each block exponentially fast, which given a sufficiently accurate learned score estimation se should be close the true backward SDE (2 One should also notice that the Gaussians E.. are only sampled once and used for all iterations.
The parallelization for (6.4) in Algorithm 0 should be understood as that for any le [0 K - 1), each 8f, + (0) (0())).)) for j E [0 : M,]] is evaluated parallel, with subsequent doating-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 ([0,tv] 8-accurate learned score). The learned NN-based score 8 is 82- accurate in the sense of
Ma-1 N-
Ep
En.m 66 +Tamm (Zom+ramm) log Pt., Fn.m ( mT.m))))))))
82.
(3.5)
===
m==0
Assumption 3.2 (Regular and normalized data distribution). The data distribution Po has finite second moments and is normalized such that co (Zu) Ia-
Assumption 3.3 (Bounded and Lipschitz learned NN-based score). The learned NN-based score function s; has bounded C' norm, i.e. ()lll, Σ M, with Lipschitz constant L
Remark 3.2. Assumption 15l and the finite moment assumption in Assumption 3.2 are standard assumptions across previous theoretical works on diffusion models (TO0, WI4 LTJ, while we adopt the normalization Assumption 3.A from [MDA] to simplify true score imnction-related compulations (ef. Lemma A.S).. Assumption 5 can easily satisfied by truncation, ensuring computational stability. Notice that the exponential integrator, one actually applies Picard tteration to e e sg. relaxation of Assumption 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 I]:
Theorem 3.3 (Theoretical Guarantees for PIADM-SDE), Under Assumptions given the following choices of the order of the parameters
H and
5
T = (log(dd622) )), h @()), = (log(dd8
é c= O log 2 (d)- = (do (dā-2 1)) K = (log(dd--)
and let Lihn in < 1, 82 8, T S logn the distribution qNN that PIADM-SDE (Algorithm M generates samples from satisfies the following error bound:
(PL(Ppl((e)) E de + dcT + 82 + dee K 28"
with a total of KN Ō ( - (1og (d8 (d)-2)) approximate time complexity and dM = Ō (d space complexity for parallalizable d-accurate score functionn computations.
Remark 3.4. We would like to make the following remarks on the result above:
The acceleration from ((d) to Ö(poyy log d) is at the cost of a trade-off with extra memory cost
of M = (d)) for computing and updating {sf. tTa, (1 (Dtn. mme():M..) simultaneously during each Picard iteration:
((k)
Compared with log- log-concave sampling [L30]] M being of order ((d) instead of Ö(Vd) therein is partly due to the time independence of the score function V log p(-) in general sampling tasks. Besides, the scaling M (d) agrees with the current state-of-the-ee dependency (UOZ) for the SDE implementation of diffusion models;
As mentioned above, the scale of the step size within one block is still confined to O(L/M) = Ô (1/d). The block length h, despite being required to be small compared to 1/L., is of order (1), resulting in only Ollog d) blocks and thus Öppoyy log d) total iterations.
3.1.4 Proof Sketch
The detailed proof of Theorem decompose the error DKL (P,,, x bound the error in each block by invoking Girsanov's theorem; (c) sum up the errors in all blocks.
+ is deferred to Section H. The pipeline of the proof is to (a) first 1 into blockwise errors using the chain rule of KL divergence; (b)
The key technical challenge lies in Step (b). Different from all previous theoretical works [[оτ ma4 IT], the Picard iteration in our algorithm generates K paths recursively in each block using the learned score sQ. And therefore the final path (K) (o!"))) depends on all previous paths (t)Ak) (Hrnr) )re(o.n.)) for A e [0 K 1). ruling out direct change of measure argument from the
7

---

naive application of Girsanov's theorem. To this end, we need a more sophisticated mathematical framework of stochastic processes, as given Appendix A.2 We define the measurable space
(R,F) with filtrations (Fe) to specify the probability measures on (S, FF) of each Wiener process, and resort to one of the most general forms of Girsanov's theorem (ا Theorem 8.6.6)). For example, in the n-th block, we apply the following change of measure procedure:
1. Let glz,,, be the measure where wr(w) is the shared Wiener process the Picard iteration (3. for any k [0: K -1];
2. Define another process dlii (w()) +- + Se, (T,W)dT, where
6.n (T,W) o (g)(t) ( (w)) V log (Den (w)):
3. Invoke Girsanov's 's theorem, which yields that the Radon-Nikodym derivative of the measure Pl., with respect to glzs, satisfies
dpl (w) - Inm 8c, (T,(T,W)) (T,w) dw 1/m, 16,,, (T, log dglF
4. Conclude that (+),+))77 is a Wiener process under the measure Pl.,, and thus at iteration K satisfies the following SDE:
(K) dg (w)
(K) Vin.r (w) V log logir,, (Dlm) (w)) dr tr
i.e. the true backward SDE (2)) with the true score function for T E [tn, tn+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 =(1) and the initial condition yi,. in Lemma B 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, will show that our parallelization strategy is also compatible with the probabil- ity 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 D further improves the space complexity from (dd) to O(d//2) while maintaining the same Opooly log d) approximate time complexity.
3.2.1 Algorithm
Due to the space limit, we refer the readers to Section [ and Algorithm for the details of our parallelization of the probability flow ODE formulation of diffusion models. PIADM-ODE keeps the discretization scheme detailed in Section س that divides the time horizon into blocks and uses exponential integrators for all updating rules. Notably, PIADM-ODE has the following distinctions compared with PIADM-SDE (Algorithm MI:
Instead of applying Picard iteration to backward SDE as (B.2), we apply Picard iteration to the probability flow ODE as in (CK) within each block, which does not require sampling i.i.d.
Gaussians to simulate a Wiener process;
The most significant difference the adoption of an additional corrector step [|| 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 ((1) time with the Icarned NN-based score function at the time 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, naive implemen- tation would result in ((vi) [I30]] 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 B.J add assumption on the Lipschitzness of the true score functions log Pt. which is common practice in related literature [m m. Recent work on the probability flow ODE implementation [л22 [14] also adopts stronger assumptions compared the SDE implementation.
Assumption 3.1' (L ([0, (x]) d-accurate learned score). For any n [0 N 1] and m E [0
Mn 11. the learned NN-based score sy ft..omm is 8x acccrrate in the sense of
||
Epen Teym 8 plu (Zintrn.n) - V log Ptn (Tt, +
S 880
Assumption 3.4 (Bounded and Lipschitz true score). The true score function V log Pt has bounded I Mp with Lipschitz constant T.p-
C¹ norm, i.e. |II log pe( (olggp())llllz=>>>/ [[..T=
Further relaxations on Assumption B. to time-dependent assumptions accommodating the blow-up to the data end (e.g. [|0,) 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 given the following choices of the order of the parameters
and 14
T = (log(dd
h = ((), N = -1og((d6
o ( -1 -d (d=1122-1)) for the outer iteration and
M = 1o K = (log(dd 3),
T = 0 x 1,p 1/2 A 1-11/2 h² 0(1, N =1),(1,
d= O
Mt= = Odd Kt = 0(lgg(dō ?))
for the inner iteration during the corrector step, and let Lightt v Li eh // < 1, 8~ 2 ö log (d6-2), and 7 2 D then the distribution quN that PIADM-ODE (Algorithm D gener- ates samples from satisfies the following error bound:
TV )2 S de de- + de + (T + N2) + eN ~80
with total of (K + KINT)N = ō (log2 approximate time complexity and d(M M³) G (d//28-1) space complexity for parallalizable 8-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 deterministic process in time rather than stochastic process as in the SDE implementation, getting rid of the O(e) term derived by Itô's symmetry. This allows the discrctization error to be bounded with OEe2) instead (cf. Lemma B. and
3.2.4 Proof Sketch
The details of the proof of Theorem 6x are provided in Section Q Along with the complexity bene- fits the deterministic nature of the probability flow ODE may bring, the analysis is technically more involved than that of Theorem GM 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 (ef. Theorem (и) instead of any f-divergence that admits data processing inequality as in the SDE case by Girsanov's theorem. The additional corrector step exactly han- dles this issue and would intuitively translate 2-Wasserstein proximity to TV distance proximity (cf. Lemma N), allowing the decomposition of the overall error into each block;

---

For the corrector step, the underdamped Langevin dynamics second -oderr dynamics requires only (Vvd) steps to converge, instead of ((d) steps in its overdamped counterpart. We then adapt the parallelization technique mentioned in Section 2.2 to conclude that it can be accom- plished with (log d) approximate time complexity (ef. Theorem C.I.). 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 's mneorem(Theoree A.4) as in the proof of Theorem BSC
Different from the SDE case, where the chain rule of KL divergence can easily decouple the ini- tial 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 dis- tance as it connects with the KI. divergence via Pinsker's inequality and admits data processing inequality. We refer the readers to Figure @ 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, 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 PLADM-ODE, are meticulously designed and rigorously proved to achieve (poly log d) approximate time complexity and (dd) and ((d//2) space complexity. respectively, marking the first inference algorithm 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 applics Picard iteration within each block in parallel, transferring the time complexity into space complexity. Our analysis is built 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 [M. 21 [E]] that parallel sampling tech- niques 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 the discrete [MK [I3 [EW] and multi- marginal [|<T]] formula- tions. Although we anticipate implementing diffusion models in parallel may introduce engineering challenges, e.g. scalability, hardware compatibility, memory bandwidth, etc., we believe that theoretical contributions lay a solid foundation that not only supports but also motivates the empiri- cal 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 Google Research Scholar Award.
References
[1] Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sam- pling of diffuusion models. Advances in Neural Information Processing Systems, 36, 2024.
[2] Michael Albergo, Nicholas M Boffi, and Eric Vanden Rijnen- Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08897,, 2023.
[3] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15577 2022.
[4] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXivv2210.02747, 2022.
[5] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint wtXXi:2209.030000 2022.
10

---

[6] Jascha Sohl- Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynamiis. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[7] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Ad- vances in neural information processing systems, 33:6840-666 685000668811188 2020.
[8] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34: 1415-1428, 2021.
[9] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[10] 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:201 1. 2020.
[1I] Linfeng Zhang, Weinan E, and Lei Wang. Monge-ampten flow for generative modeling. arXiv preprint arXiv: 10/88, 2018.
[12] Omer Bar- 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.
[13] Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. Deconstructing denoising diffusion models for self-supervised learning. arXiv preprint arXiv:2401. 144044 2024.
[14] Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 2344))11-33, 2022.
[15] Nanye Ma, Mark Goldstein, Michael Albergo, Nicholas M Boffi, Eric Vanden-Eijndenn and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint arXiv:2401 087440 2024.
[16] 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.
[17] 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 Pmir, 2021.
[18] 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.
[19] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Om- mer. 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.
[20] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint w:Xi::2/1.00000 2021.
[21] Yu Sun, Zihui Wu, Yifan Chen, Berthy T Feng, and Katherine L Bouman. Provable proba- bilistic imaging using score-basee based generative priors. arXiv preprint (rXiv:2310.10835 2023.
[22] Xingyu Xu and Yuejie Chi. Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction. arXiv preprint 1iXi:22403.17444 2024.
[23] Jacob Austin, Daniel D Johnson, Jonathan Ho. Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models discrete state-spaces. Advances Neural Informa- tion Processing Systems,
2021.

---

[24] Xiang Li, John Thickstun, Ishaan Guirajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Pro- cessing Systems, 35:4328 4343, 2022.
[25] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 118633--6468 2022.
[26] Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli Sastry, Siddharth Gururani, Sageev Oore, and Yisong Yue. Symbolic music generation with non-differentiable rule guided diffusion. arXiv preprint arXiv:2402.14285, 2024.
[27] Gautam Mittal, Jesse Engel, Curtis Hawthorne, and lan Simon. Symbolic music generation with diffusion models. arXiv preprint arXiv:2103. 2021.
[28] Flavio Schneider. Archisound: Audio generation with diffusion. arXiv preprint
vrXiv:2301.33267 2023.
[29] Ruihan Yang, Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for video generation. Entropy, 25(10):1469, 2023.
[30] 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.
[31] Chen Xu, Jonghyeok Lee. Xiuyuan Cheng, and Yao Xie. Flow-based distributionally robust optimization. IEEE Journal on Selected Areas in Information Theory, 2024.
[32] Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from mean-field gibbs measures via diffusion processes. arXiv preprint arXiv:2310.08912 2023.
[33] Sitan Chen, Vasilis Kontonis, and Kulin Shah. Learning general gaussian mixtures with efficient score matching. arviv preprint arXiv:2404. 1...22444..889,, 2024.
[34] Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from the sherrington- kirkpatrick gibbs measure via algorithmic stochastic localization. In 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS), pages 323-334. IEEE, 2022.
[35] Khashayar Gatmiry, Jonathan Kelner, and Holden Lee. Learning mixtures of gaussians using diffusion models. arXiv preprint arXiv:2404.18869, 2024.
[36] Ye He, Kevin Rojas, and Molei Tao. Zeroth-order sampling methods for non-log-concave dis- tributions: Alleviating metastability by denoising diffusion. arXiv preprint arXiv:2402.17886, 2024.
[37] Xunpeng Huang, Hanze Dong, HAO Yifan, Yian Ma, and Tong Zhang. Reverse diffusion monte carlo. The Twelfth International Conference on Learning Representations, 2023.
[38] 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.
[39] 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.
[40] Andrea Montanari. Sampling, diffusions, and stochastic localization. UrXiv:2305.10690, 2023.
arXiv preprint
[41] Andrea Montanari and Yuchen Wu. Posterior sampling from the spiked models via diffusion processes. arXiv preprint arXiv:2304.11449, 2023.
[42] Nicholas M Boffi and Eric Vanden-Fijnden. Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology, 4(3):035012, 2023.
12

---

[43] Yan Huang and Li Wang. A score-based particle method for homogeneous landau equation. arXiv preprint arXiv:2405.05187, 2024.
[44] Lingxiao Li, Samuel Hurault, and Justin M Solomon. Self-consistent velocity matching of probability flows. Advances in Neural Information Processing Systems, 36, 2024.
[45] Jianfeng Lu, Yue Wu, and Yang Xiang. Score-based transport modeling for mean-field fokker- planck equations. Journal of Computational Physics, 503:1 12859, 2024.
[46] Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of fokker-planck equations through gradient-log-densee estimation. Entropy, 22(8):802, 2020.
[47] Amira Alakhdar, Barnabas Poczos, and Newell Washburn. Diffusion models in de novo drug design. Journal of Chemical Information and Modeling, 2024.
[48] Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gener- ative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04999 2024.
[49] Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications dna sequence design. arXiv preprint 2402::20000. 2024.
[50] 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.
[51] 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 all you need. BioRxiv, pages 2023-09, 2023.
[52] Joseph L Watson, David Juergens, Nathaniel Bennett, Brian L. Trippe, Jason Yim, Helen Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo de- sign of protein structure and function with rfdiffusion. Nature, 620(7976) 2023.
[53] Jordan Cotler and Semon Rezchikov. Renormalizing diffusion models. arXiv preprint arXiv:2308. 12355, 2023.
[54] Florian Fürrutter, Gorka Muñoz-Gil, and Hans J Briegel. Quantum circuit synthesis with diffusion models. arXiv preprint arXiv:2311.0244/ 2023.
[55] Zhiye Guo, Jian Liu, Yanli Wang, Mengrui Chen, Duolin Wang, Dong Xu, and Jianlin Cheng- Diffusion models in bioinformatics and computational biology. Nature reviews bioengineer- ing, 2(2): 154, 2024.
[56] Artan Sheshmani, Yi-Zhuang You, Baturalp Buyukates, Amir Ziashahabi, and Salman Aves- timehr. Renormalization group flow, optimal transport and diffusion-based generative model. arXiv preprint arXiv:2402. 10.002...700 7090, 2024.
[57] Luke Triplett and Jianfeng Lu. Diffusion methods for generating transition paths. arXiv preprint 2009 2023.
[58] Lingxiao Wang, Gert Aarts, and Kai Zhou. Generative diffusion models for lattice field theory. arXiv preprint arXiv:2311.03377 2023.
[59] Minkai Xu, Lantao Yu, Yang Song. Chence Shi, Stefano Ermon, and Jian Tang. Geod- iff: A geometric diffusion model for molecular conformation generation. arXiv preprint NXNi:220330092233 2022.
[60] Mengjiao Yang, KwangHwan Cho, Amil Merchant, Pieter Abbeel, Dale Schuurmans, Igor Mordatch, and Ekin Dogus Cubuk, Scalable diffusion for materials generation, arXiv preprint 0rXiv:2311.09235, 2023.
13

---

[61] 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.
[62] 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.
[63] Yuchen Zhu, Tianrong Chen, Evangelos A Theodorou, Xie Chen, and Molei Tao. Quantum state generation with structure-pressrviing diffusion model. arXiv preprint arXiv:2404.06336, 2024.
[64] 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 meth- ods and applications. ACM Computing Surveys, 56(4):1-39, 2023.
[65] Stanley H Chan. Tutorial on diffusion models for imaging and vision. arXiv preprint arXiv:2403. 18103, 2024.
[66] Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. Opportunities and challenges of diffusion models for generative ai. National Science Review, page nwae348, 2024.
[67] Valentin De Bortoli, James Thorton,, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to basee generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
[68] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020.
[69] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. Advances Neural Information Processing Systems, 35:30150-30166, 2022.
[70] Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably. arXiv preprint MxXiv:2403.33522 2024.
[71] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09777. 2022.
[72] 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.
[73] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm- solver+4: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095, 2022.
[74] 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 Sys- tems, 36:55502-55542, 2023.
[75] Zhenyu Zhou, Defang Chen, Can Wang, and Chun Chen. Fast ode-based sampling for dif- fusion models in around 5 steps. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7777-7786, 2024.
[76] 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.
[77] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint 2rXiv:2204.13902. 2022.
[78] Saravanan Kandasamy and Dheeraj Nagaraj. The poisson midpoint method for langevin dynamics: Provably efficient discretization for diffusion models. arXiv preprint arXiv:2405.17068, 2024.
14

---

[79] Yuchen Wu, Yuxin Chen, and Yuting Wei. Stochastic runge-kutta methods: Provable acceler- ation of diffusion models. arXiv preprint arXiv:2410.04760, 2024.
[80] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-damped langevin diffusion. arXiv preprint urXiv:2112.07066 2021.
[81] Alexia Jolicoeur-Martinent Ke Li, Rémi Piché-Taillefer. Tal Kachman, and loannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2105. 140000 2021.
[82] 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 22390-42408 PMLR, 2023.
[83] Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. Advances in Neural Information Pro- cessing Systems, 36:76806-78888, 2023.
[84] Jonathan Heek, Emiel Hoogeboom, and Tim Salimans. Multistep consistency models. arXiv preprint arXiv:2403.06807, 2024.
[85] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.014000 2023.
[86] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310.141500 2023.
[87] Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410. 110081110. 2024.
[88] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388. 2021.
[89] 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 2023.
[90] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005/1 00512, 2022.
[91] 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 Computer Vision and Pattern Recognition, pages 6059 6069, 2023.
[92] Zhiwei Tang, Jiasheng Tang. Hao Luo, Fan Wang, and Tsung-Hur Hui Chang. Accelcrating par- allel sampling of diffusion models. arXiv preprint arXiv:2402.0999 09970, 2024.
[93] Jiezhang Cao, Yue Shi, Kai Zhang, Yulun Zhang, Radu Timofte, and Lue 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.
[94] Nikil Roashan Selvam, Amil Merchant, and Stefano Ermon. Self-refining diffusson samplers: Enabling parallelization via parareal iterations. arXiv preprint arXiv:24/2.08899 2024.
[95] Shivam Gupta, Linda Cai, and Sitan Chen. Faster diffusion- sampling with randomized midpoints: Sequential and parallel. arXiv preprint arXiv:2406.000 00924, 2024.
[96] Gen Li and Yuchen Jiao. Improved convergence rate for diffusion probabilistic models. arXiv preprint arxiv:24 13738, 2024.
[97] Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems, 32, 2019.
15

---

[98] Belinda Tzen and Maxim Raginsky, Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pages 3084- 3114. PMLR, 2019.
[99] Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto-encoders and langevin sampling. arXiv preprini arXiv:2002.00107 2020.
[100] 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 2rXiv:2209.11213 2022.
[101] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 35:22870- 22882, 2022.
[102] Francesco Pedrotti, Jan Maas, and Marco Mondelli. Improved convergence of ccore-based diffusion models via prediction-correatte arXiv preprint arXiv:2305.14164 2023.
[103] Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degraaattton beyond linear diffu- sions: A non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning, pages 4462-4484. PMLR, 2023.
[104] Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User- frernffiendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning, pages 4735-4766 PMLR, 2023.
[105] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference Algorithmic Learning Theory, pages 946-985. PMLR, 2023.
[106] Sokhna Diarra Mbacke and Omar Rivasplata. A note on the convergence of denoising diffu- sion probabilistic models. arXiv preprini arXiv:2312.05996 2023.
[107] Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Linear convergence bounds for diffusion models via stochastic localization. arXiv preprint TrXiv:2308.03660 2023.
[108] Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv preprint arXiv:2306.09255 2023.
[109] Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards non-asymptotic convergence for diffusion-basees generative models. In The Twelfth International Conference on Learning Representations, 2024.
[110] Gen Li, Yuting Wei, Yuejie Chi, and Yuxin Chen. A sharp convergence theory for the proba- bility flow odes of diffusion models. arXiv preprint arXiv:2408. 20:::22288.22200 2024.
[111] Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The prob- ability flow ode is provably fast. Advances in Neural Information Processing Systems, 36, 2024.
[112] Xuefeng Gao and Lingjiong Zhu. Convergence analysis for general probability flow odes of diffusion models in wasserstein distances. arXiv preprint arXi::2401. 2024.
[113] Yuchen Liang, Peizhong Ju. Yingbin Liang, and Ness Shroff. Non-asymptotic conver- gence of discrete-time diffusion models: New approach and improved rate. arXiv preprint arXiv:2402.13901 2024.
[114] Daniel Zhengyu Huang, Jiaoyang Huang, and Zhengjiang Lin. Convergence analysis of prob- ability flow ode for score-based generative models. arXiv preprint arXiv:2404.097300 2024.
[115] Gen Li and Yuling Yan. o(d/t) convergence theory for diffusion probabilistic models under minimal assumptions. arXiv preprint arXiv:2409.18955 18959, 2024.
16

---

[116] Brian DO Anderson. Reverse-timg time diffusion equation models. Stochastic Processes and their Applications, 12(3):11--266 1982.
[117] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[118] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674 2011.
[119] Yu Cao, Jingrun Chen, Yixin Luo, and Xiang Zhou. Exploring the optimal choice for gener- ative processes in diffusion models: Ordinary stochastic differential equations. Advances in Neural Information Processing Systems, 36, 2024.
[120] 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.
[121] 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 wxXv:2331.004100 2023.
[122] 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.
[123] Koji Hukushima and Koji Nemoto. Exchange monte carlo method and application to spin glass simulations. Journal of the Physical Society of Japan, 6566):1604-1608 1996.
[124] Faming Liang. Use of sequential structure in simulation from high-dimennional systems. Physical Review E, 67(5):056101, 2003.
[125] 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.
[126] Holden Lee. Parallelising glauber dynamics. arXiv preprint
2023.
[127] Lu Yu and Arak Dalalyana. Parallelized midpoint randomization for langevin monte carlo. arXiv preprint arXiv:2402.14434, 2024.
[128] 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.
[129] Emile Picard. Sur les méthodes d' d'approximations successives dans la théoric des équations différenticlles. American Journal of Mathematics, pages 87-100, 1898.
[130] Nima Anari, Sinho Chewi, and Thuy-Duong Vuong. Fast parallel sampling under isoperime- try. arXiv preprint arXiv:2401.0900 2024.
[131] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
[132] Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2/ 10.02037 2021.
[133] 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.
[134] Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data. Advances Neural Information Processing Systems, 35:34532-34545, 2022.
17

---

[135] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous- time discrete diffusion models. arXiv preprint arXiv:2211.16750, 2022.
[136] Pierre H Richemond, Sander Dieleman, and Arnaud Doucet. Categorical sdes with simplex diffusion. arXiv preprint arXiv:2210.14788, 2022.
[137] 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.
[138] Griffin Floto, Thorsteinn Jonsson, Mihai Nica, Scott Sanner, and Eric Zhengyu Zhu. Diffusion on the probability simplex. arXiv preprint wrXiv:2309.025300 2023.
[139] Javier 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.
[140] Joe Benton, Yuyang Shi, Valentin 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-30) 2024.
[141] Hongrui Chen and Lexing Ying. Convergence analysis discrete diffusion model: Exact implementation through uniformization. arXiv preprint arXiv:2402. 08095, 2024.
[142] Yinuo Ren, Haoxuan Chen, Grant M Rotskoff, and Lexing Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via stochastic integral framework. arXiv preprint 2rXiv:2410.03601, 2024.
[143] Michael Albergo, Nicholas M Boffi, Michael Lindsey, and Eric Vanden Fijden- Multi- marginal generative modeling with stochastic interpolants. arXiv preprint arXiv: 10000395, 2023.
[144] Gean-François Le Gall. Brownian motion, martingales, and stochastic calculus. Springer, 2016.
[145] Arnaud Guillin and Feng- Yu Wang. Degenerate fokker-planck equations: Bismut formula,
gradient estimate and harnack inequality. Journal of Differential Equations, 253(1):20-44) 2012.
18

---

A Mathematical Background
In this section, we will summarize used notations and rigorous mathematical framework of Itō pro- cesses 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:
Notation
Description
[a b]
The set {a, a + 1, 1.....b) Identity matrix in Rdxd
Ia
*r-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 vector The inequality holds up to a constant factor
or
Absolute continuity (for measures\/ much less than (for quantities)
(m) 1:20
The forward process of the diffusion model (2.3)
( [E
The backward process of the diffusion model (L
JeE[o,11
The approximate backward process of the diffusion model
(k)
The approximate value of the approximate process y: at time bn + Tn.m after & iterations in the (n + 1)-th block
Utn
The value of the approximate process y; at time tn
The distribution of yt..
The path (z)) of the process ž₁ The J-divergence between two distributions The KL. divergence between two distributions
Dy(- II
TV(.
The total variation distance between two distributions The 2- W-Wasersteinn distance between two distributions
:2(-))
Table Summary of notations
A.2 Preliminaries
Theorem A.1 (Properties of f-divergence). Suppose p and q are two probability measures on common measurable space (2,, with p q. The f-divergence between p and q defined
( dq
Dy(pl)g) Ex
(A.1)
dp is the Radon-Nikodym derivative of with respect to q. and f R+ + RR is comvex where dg function. In particular, Dy(- || -) coincides with the KullbackLeibler (KL) divergence when f(x) TV coincides with the total variation (TV) distance when f(x) 1 1/ 1/.
log and ( II
For the f-divergence defined above, we have the following properties:
1. (Data-processing inequality). Suppose H. is a sub-o-algebra of F, the following inequality holds
Dy (pbu (b)) S Dy(p q):
for any f-divergence D-(-1l)).
2. (Chain rule). Suppose is a random variable generating a sub-o-algebra Fx of F, and x(IX)) & q(1lX) holds for any value of X, then
(xl(pllg) = DRL (pl Fx 1xabxx) + Exx
19

---

In this paper, we consider a probability space F,, on which (wy(w)))2o is a Wiener process in R.. The Wiener process (w: (w) 1:20 generates the filtration {FeJe20 on the measurable space (S2 F). For Itô process z (w) with the following governing SDE:
dæ, (w) att, w)dt + [t, (w),
for any time i, we denote the marginal distribution of 2₁ by Pt. i.e.
Pr (z. (.)) where Ω + R., + (w),
as well as the path measure of the process ž₁ in the sense of
Patta == p=2p((z))) (z22 (-) where zt,t, f2 + C( t2). R.), w + (za(W))ee(hsta))
For the sake of simplicity, we define the following class of functions:
Definition A.2. For any S t i2 we define V(ti,t2) as the class of functions f(t,w) [0, +00) x Ω + R such that
1. f(t,w) is B x F F-measurable, where B is the Borel a-algebra on R.::
2. f(t,w) is F₂ Faaadeee for all t 2 0;
3. The following Novikov condition holds
t exp I"
P((,))de < +0c,
E
and V = N V(0,). For vectors and matrices, we say it belongs to V V* w) or Vmxx (t,w)) if each component of the vector or each entry of the matrix belongs to ((I, w)
Remark A.3. Novikov's condition appeared in the third requirement is often relaxed the squared integrability condition in the general definition of lt processes, which requires
E L,,0 P(l,w)dee < +00.
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 [НЛЛ]], here avoid checking Novikov's condition throughout proofs below by using the approximation argument presented in [[C]] A review of Girsanov can be found in textbooks like in [[КI] [44]. We will present the following generalized version of
Girsanov's 's theorem:
Theorem A.4 (Girsanov's Theorem [B1 Theorem 8.6.6)). Let a<l,w) Vm, E(1,w) vmxn, and (w,(ar (w2) be Wiener process the probability space (Ω, (.)). For [0, 71, suppose ze z((w) is an Itô process with the following SDE:
dz)(w) (tt, + E(t, (ww,
(A.2)
and there exist processes 6(t,w) E v and B(t, E v such that
1. E(1,w))(1, = all,, - 3(1,w)
2. The process Me(w) as defined below is a martingale with respect to the filtration (Felezo and probability measure q:
M((w) =cxp exp (-/ ((s,w)Tdw,(() dw,(ww S. | [ w1)
then there exists another probability measure p (), such that
dp (w) = MT(w),
I. < with the Radon-Nikodym derivative
dg
20

---

2. The process (.( defined below is a Wiener process on (2, F, p):
(i(w) ww(w) + L. 8(s,w)ds,
3. Any continuous path in C( 12]. )) generated by the process 2t satisfies the following SDE under the probability measure p:
dī((w) 3(, + E(t,w)dū,(w).
(A.3)
Corollary A.5. Suppose the conditions in Theorem A hold, then for any f1, [0. with t < 12, the path measure of the SDE (A.5) under the probability measure in the sense of Ph = (Zirida (-) is absolutely continuous with respect to the path measure of the SDE (A in the sense of Q1att = (Z ()). Moreover, the KL divergence between the two path measures is given by
1 [8(t,
DKL (P (1 (KL(Ps, (1t;,) + Beowpl F11
2 Jr
(A.4)
Proof. First, by Theorem have
DKL(Pt,a (14)))2) Dku llas,, 1lgl,,,), Ben- [DKL (p( ())E,, (1 = =])]
From Girsanov's theorem (Theorem A 4) we have that the measure pla,, is absolutely continuous with respect to glz,,, which allows us to compute the second term above as follows:
DKL (p( (-),,, ")119( 1)) =2))
log dg (Z to (-) = =x) Euu-plen log dpl
dp(2 (-) = =2)
dglF,,
8ll w) (T,,))Tder,, ddwr(w) - 1/ (10(1, (((,,w))
Ewwxp Fi
t2
5(t, (T (d( - wttt / 11ō(t, w/ (
Ew-p Fi
-
1 pla
Ewnp Fi
2 (12
and therefore
[ f
(kL(Phsslll0...)) = DKL 114,,) + Eu Fi 2 (F)ttt
(15(t,
which completes the proof.
A.3 Helper Lemmas
Lemma A.6 (1107) Lemma 21). For the backward process (2.3), we have for 0 S <<<1<) <TT,
d
V log (ā.) log D. (t.)) <
[[[ logVo (1.)
|| 2 (*,
dt
Lemma A.7 (1077 Lemma 3J). For the forward process (2.3) we have for < < <T
E V (o))( < 2 and E [|[ logp, < do, + 2 d
dl. (o, (o E [tr Edl)
where the posterior covariance matrix E, COV (m:) and of 1- Moreover, the posterior covariance matrix E, satisfies
E ΣxJ dd do?
21

---

Lemma A.8. For any n E [0: N 1] and T E [0, hnl, under the assumption covpo (21) = Id. we have
E [e, [1
S 2d.
(A.5)
and
B [it. t
S 3d.
(A.6)
Proof. Conditioned 20, have that
Tt. = xT ~ N (c-t(T-In)0o, 20, (1 ((T--()
and
3 WT ~ N (e -(T 20. (1 (T tn- -n-- {))
for any T E [0, An). Therefore, we have
Lt, [2] = E Dr tm [0]]
S E E | e (T xo |00 + lle 3T 2001?
d(1 (T-) (( F1 æn [| < 2d.
Taking the difference between them then implies that for any E [0, hn],
Lt, It, 22,+T 11 = | - xT 2T-4--7 2 x00
< d(2 c- -(T - (T-
+ e ((T - -2 E[| zeo ||
< 2d - (T - (1 t B - [13] S 3d.
Lemma A.9 (Lemma 9 in (124)) For Co ~ NN00 and Po pr is the distribution of the solution to the forward process we have
TV T0)² < DKI. (follo)) S de-
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
dæ.
2 E V log (z.) dt dw,, with To Pr.
(B.1)
and its approximate version
(2.5)) with the learned score function
2 yr + a (vr)) dt + dwr, with Yo N0 (0) Ia).
dy:
The filtration F, refers to the filtration of the SDE up to time t.
22

---

B.1 Auxiliary Process
We would like first consider the errors that Algorithm may cause within one block of update. To this end, we consider the following auxiliary process for T E [0, hn) conditioned on the filtration Ft, time ty:
Definition B.1 (Auxiliary Process). For any n E [0 N - 11, we define the auxiliary process (ot (0. the solution to the following SDE recursively for & 10: K 1]:
dy ( (w)
(k+ (w) + s stnt (o dr + den,,+ Dini
(B.2)
with the initial condition
(w) = ÿtn (W) for E [0, Nnl, and g Flw = Drn (w) for E (1: K]
(B.3)
[1: e where Dr., (w) = (KK) D Dta-155e-IMMM11 (w)
and Tro (w) ~ NO Iz).
The iteration should be perceived as a deterministic procedure to each event w E Ω i.e. each realization of the Wiener process (w,) The following lemma clarifies this fact and proves the well-definedness and parallelability of the iteration in (B.22).
Lemma B.2. The auxiliary process E 00: N-1 1).
(k) (w) (E(Oo))) is Fe.tt-adapted for any lo E [0: K] and
Proof. Since the initialization T. (w) Dr., (w) for [0, where Tr. (w) is Fr,, aadatedd, it is obvious that (0) (w) ) is + Now suppose that (o!")) (e))) is Fr, +r adapted, plo!
since 9n (T) T, we have the following Itô integral well-defined and Fr, +r aadatee::
(19 (w)(w))
and therefore (1:32) has a unique strong solution (Di.r (w))rejo.nnl) that is also Fe.+r--dapted The lemma follows by induction.
Lemma B.3 (Equivalence between G4 and
For any n E [0: N the update rule (3.4 in
Algorithm 0 is equivalent to the exact solution of the auxiliary process (B. for any k E [0: K and T E [0, hn..
Proof. The dependency w will be omitted in the proof below.
Rewriting (B.2) and multiplying eー on both sides yield
d
(k+1) Yiner e
d (k+1)
( dr -
50 ( ( lr deve
Integrating on both sides from 0 to implies
8 (k+1) e
Dtnit
Di,
(k+1)
8 (T')
(16)) (Wrnddenntr + deoetrr
Mn
AATmmm-
Tonn )dr' + 1' e +dde,
e 8 to
m==0
T/tn.m
M
2
*AT
2
(kk)) 8 tiuttoi
D Vlasta.m) ) + (" e
- dww,
m=0
and then multiplying ei both sides above yields
M,
((+1)
=
e -220
(1+)1)
Σ
TATa
AAram
1)
outr Ta
[kk Fn.m Vinitam
Vin.
M
A+11
dwi,
m==
23

---

where, by Itō isometry, we have
e dw... ~ N (o, (e -hrmmm ovel Faom+
F
for Tn.m and equals to otherwise. Plugging Tj.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:
1. Setting 2 in Theorem as the auxiliary process at iteration K, where we (w) is Wiener process under the measure Fl.,,.
2. Defining another process +r (w) governed by the following SDE:
dr dio,,, (w(w) + di, w),
where
6t, (T.W) = 8 +n,t (r) (9) (0t (r) +. (w)).
(B.4)
and computing the Radon-Nikodym derivative of the measure p s., with respect to gf gl
Fil (w) = exp (-f dglsr.
de,, (T, (f
1 h.. dawe - 2 f/ 11ō,,, (I
3. Concluding that (B.2) at iteration K under the measure qlF.,, satisfies the following SDE:
(K) dÿ ty (w)
(K) Din. (w) logā (D) (w)) dr +r
with (7) being a Wiener process under the measure Pl.,, If we replace (K) De (w)
by t- (w), should notice (B.5) is immediately the original backward SDE (2.3) with the true score function on E [tn, +1):
tr (w)
1 1
(w) (w)) logggi Vo (w)) dr deū,,
(B.S)
Remark B.4. The applicability of Girsanov's theorem here relies on the F--adaptivity of (r,ggo(7)(w)) established by Lemma B. 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.A).
80 (..tg (K)
Then Corollary A. provides the following computation
DKL ( (1Ten1)) S DKL (Pinitae) (
1 tm ||,,, (T, (T,
(B.6)
= D (P.,, (1 + Ewigg Feen
2 J
24

---

where the first inequality is by the data-procest processing inequality (Theorem A D). Now, the problem remaining bound the discrepancy quantified by
16,, (T,
An
1 bn tg) (onnsn (w)) V loog (0 dr
hn
V log Pr.,, ta (r) (r) (w)) - V Vlogār, +r ((K) (w)) | dr
<
(B.7)
Aen (w)
s +an (Drnn (T) log (
(r) (w))) (-)
dr
:Be (w)
8 tontonl)) (D (r) (w)) 8 (r) (o (w))|/ J
hn
( dr
Before continue proof, would like first provide the following lemma bounding the behavior of the auxiliary process (B.2) when = for E [0, hn]-
Lemma B.S. For any 10 N - 11, suppose the initialization Dt. in (B.3)) of the auxiliary process B follows the distribution of ti. ~ Prn.' then the following estimate holds
sup Ewwilren [I (w) - T [w))?
rE(0,ha]
< h mn 3mm (M² + 2d) + 30 BBnnBll [Ae, (w) + Bi,, (w)
(B.8)
3e hm hn I sup rE(o.hn]
Eur
(w) (w)
<KK 1) Din.t (w)
(1)
Proof. Let Zt,,. = Vt...
Dtn.r For k = O, we can rewrite
B22) as
dz
2 Zin.r + vi..)) + S (9)())
V (m dr + dw
By applying Itô's lemma and plugging in the expression of wr given by Theorem A4 we have
d d| A||2 11
|| 122 2T TT D..0 22, - o tom(7)
(10)) (o + d dr
(B.9)
22,,, (dwi, ++ (w) 8,, (F.
By integrating from 7 and taking the expectation on both sides of (B. we obtain that
Bwnplen [| []]]
= -ww
[|
T D.. + 22 T2,
8 tnn(rr)
( + d dr'
Z (d(0 8 (7, (F)(w)))r)))
+ 2E 2E Fr..
and by AM-GM, we further have
BwwBlren [| [|3]
< EEuwplre
2
(0)
+
(0))) 12 +d+ d+ 1|8,, (T, (T.))1 dr'
2
< / Bumnle
2 ||2 11,,, (T, 2|
dr' +
E
(0)
27
+ M2 + +d)
25

---

where Sen (T) is defined in (BE). Similar to (B 2).
we may use triangle inequality to upper bound
1, (T,w)] w) which implies that for any T E [0, hm
Bumple [1]
7
|| |1 dr' +
(1 (0)
Vin.T 1"] M² + ta)
2
+ 3.u--ple
6 (r) (r) (1 (K-1)
(oonion (w))
ttnt (B fr (wW)]] dr'
8 a (r) (K)
(T) (w)) V log Dt. (1) (0( (r) (w)) / dr'
[K]
3.B
+ 3.B--Bb
V log Pt,, (T) V
(7)) (w)) (w))
(K)
V Vlog po ++ Vi. Vi.tt (1 (K)
(w)) dr'
7 / Bunnben [1E2a07111]] d' dr'
12/
J
R
M + +d)
2
+ 3L2 ('" Bu-i Fin
|| ( (1t) (w)
(w))
|||
dr' 3B si., [Ae, (w) Be, (w)),
where in the second inequality above, we have used the fact that s.) is Li-Liipchitz for any l. By Grönwall's inequality, we have that for any [0, km]
E Bw- [1]=,_.-11]] S 3rr [
+ M² + d )d)r r|
+ 3 Buuibr [Ar, (w) + Be, (w))
(B.10)
+ 30 L3 / Bumplien
I (
y (in
(K-1)
(w) dtt
dr'.
By assumption, Vi = V.. follows the distribution of "z,, Pe,,, which allows us to bound the second moment of y: for any n E [0: N] by Lemma
E (.. 12] = E [ 1| - 2d.
Substituting (A5) into (B 100 then yields that for any E [0, hr
E E Fin 1t...r [13] S = (M² + 2d) + 3E S Eu-Flii [A + Bi,, (w)]
| (w) - " (w) 2
3re 3 L3 sup Bawely r'E[U.A..]
Taking supremum with respect to 7 E [0, h,,] 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).
Lemma B.6 (Exponential convergence of Picard iteration in PIADM-SDE), For any E 10, NJ. sup- pose the initialization Dt., in (B.5) of the auxiliary process (B.2) follows the distribution ot.. Pr..'
then the two ending terms y and
_(K) <(- yiu.r
of sequence (o(%) ] satisfy the following
exponential convergence rate
(w))
sup Euw-Blrn re) e2gh K-1 hežhn (M2 + 2d)
3 (I 2Rr)
K-1 c [/hn Es Euw [Ai, (w) B., (w)
+
1-3 (L3 (2 K-1 Tb h 12
1- 3-3 (L2tin k-1 eggn h L2
Proof. For each Ω conditioned the filtration Ft,., subtracting defined by
B from the process as
[1 (k) Vt (w) + S (a (1 (w)) dr + dww,,r(w), (B.11)
(k) (w) = d
26

---

we have
(k) d 1k++1) (w)) 1l (w)
(+( 2 Vinst (w)) Vinit (w) + o tga(t) (T) ot on (0 (1) (ww) dr,
(k)
where the diffusion term dw,,, +r cancels each other out. Now we may use the formula above to compute derivative dr. | (w) sl (w)] explicitly, integrate it from 7' 0 T, and
obtain the following inequality
(litt) Vinir (w) Vin.. (w) /
(k)
y (k-1) (w) 0 (w))
( (w)) s tnn) (D (r)))(())))))) dr'
Den.T! (w) D (w) dr'
(k+1)
< (ottt) (w) D.n. (w) || dr'
+ ! (lutou(t) (Fin.g (w)) 50 ten()))))) (oth-nir (wi) || dr' (k) | (w) - ( (w) || dr' + L2 / ( (w) / (w/)) dr'
<
2
By Grönwall's inequality, we have
(1++1) (w) (w) || S 1.ge2r
[
) (r') [w)) dr'.
(B.12)
Taking expectation both sides above further implies that for any t E hm),
1 | (w)
(A) Diat
) (1))
Bw-:lr
1-¹ Eunplenn 11 Venign (or))s((T((]]))))))))))) ((w) -1)
/|² dr'
1. 27
Vin.ga (7)) (w)
(B.13)
(k) 1 (w) - (w) |
< TT 2r sup Buw r'E(0.-]
Furthermore, we take supremum over T E [0, hn] on both sides above and iterate ( over l e N, which indicates
Venit (((ktt) (w) 0/ ( [w))
sup rE(0,hm]
Bu-iliin
((k)) Vinr (w) -  (w)) 22
< I. e 2hn
sup re(
< e2
E ( - 71001 [w]-(ww) [12
(B.14)
sup
re(o,
(L k the (M² + 2d) + + (L3 2hn, t 2² Bumplen AA (w) + Br, (w))
EuwFlenn ( - (K- (w))) 12]
+ e2 ezhn L
sup
rE(0,Am.
27

---

where the last inequality follows from Lemma B5
By rearranging the inequality above, setting
== K and using the assumption that L ha 2hn
we obtain
sup Eu-Blsen [ (w) g(K-1) (w) re)0,aal
(22 )k-1 he (M² + 2d) + Light 2h..) )H-1
K-1 e Bumplr. [A,, (w) + Bt,, (w)
1 - 3 cn K --1 LinmnnIL2
(B.15)
as desired.
The following lemma from [1ол]] bounds the expectation of the term Ar., (w) in (BB.)):
Lemma B.7 (1107, Section 3.1]). We have
Blr [A, (w)) S edlnn for n E : No:N-22 2], and Eumpler [Aex (w)) S ed log
where 1i is the parameter for early stopping.
Proof. Notice that
EunEber
[Ae, (w))
hn
V log Pr., tan(r) (r) (K( (t) - V loog logār..+ +r (1 (w)) || dr
To.m+t
V log Vlogie
+na. (Venrta (w)) Vlog Pe. ( (wm) || dr
m==
Tnam
Ma
Tn,ma
(stt V log Prg (Tt +- (wi)) /| dr;
m==0-
where for the last equality, we use the fact that the process y ( the true score function under the measure P. In the following, we drop the superscript ~ pl.,, of the expectation for simplicity.
(w) follows the backward SDE with
By Lemma A.G and A.A, we have
V log Dt., + (we, (w)) V log logPr.+> (#.., + (ww)) (1
E
E
V log Frn (Bi, tram (w)) [12
E. [[[ log De +o (Bt, ((()) dr'
Jo
<
/
2 dō + -ā dr' +
E te E See J b tr Den,tr
Now noticing that
82 = -T
we further have
Tm.
V log Prgtr (wi, + (w)) | dr
E v log Per.tm (Tr,++(w)) (w))
6n E tr 21 to + Tm.m
B tr Er, Tnmm-1
Tnuml1
12
dr' d dr +
(T' - iu Tnm+1)²
(T bn Tn,m)"
Fm.m
2
c
E
+ Tm,
E tr Tmmmt1
+
TT im Tn,m+1 3
tp Tn.m
28

---

and thus
Mr
Fn.m++
V log Pan, trim (w))+r(w)) - V log Pt (Zintr (w]) | dr
Σ
E
m==
Fn.m
Mn
2 n,m
Mn
<d
EB
tr E
tr E t tTo, m+1
(T tn
Tn,m-1
T tn Tn,m
in Tumm
c tr Et +T.
Mn
6mmm B ir
Mn + T- in Tn.0
(T in Tn, 1I to Tn.m)
m==
Mn + ed + de de2M,, S de Mn.
For n E 10, - 2]). we have Mnt = ln and thus Bunilren [Ar, (w)) S edht, and for n = N 1,
we have
" 1 dr = =log" logn lc-
MN S
CT
and thus Blr,, [wi) S edMM, S edlogn
B.3 Overall Error Bound
Proof of Theorem 3. We first continue the computation in (B.6) and (B
DRL. (Pr,z (1 DKL (Pr,, ll@e..) + Eu
11 (T.)))
2 2
<DKL (p,, ll@m.) + 3E.mF
[Ar + Be, (w]]
+3EWwwBl (%am) thn
s |slutoa(t))
( (w)) - 1
to (r)()) -1) ((,)) /| dr
1 pnn
( (w)) ) (w)) dr
(DKL(P,,1100 + 3Ew~Fl
Ar, (w) + Bin (w) + 13
31 [w) - - (w) |2]
(Fr, + 3Ew-ili A.,(w) + Bi, (w) + L sup
[k(
rE(O.An)
Then plugging in the result of Lemma B.G. we have
DKI ll@em.))
(Laht 22 K=1 he (M² + 2d)
< <P,,(P,, (1 + 3B.-BIF (w) + (w)] + 3h, L2
1 3 e23 K e
3hm 1, L3
(L3 )k-1 B.hnBunpee
[A (w) + (w)(w))
1 3 (I3 R-1 e 2i I?2
1 + c- - h
Drr (1qen))
+
E ehn BuniiZe, [Ae (w) + Bi, (w)) + "o
SDKI (1Tt..)) + Euwelrr [Ae,, (w) + (w)) + e- 13,.
where we used the assumption that 12 The < 1.
29

---

The term En=o Bw--lE [Be (w)) is bounded by Assumption B as
N-1
Bu sen
BB, (w)]
mm
< <l <luwplr
Σ
8 tox(t) ( (Din.ga (w)) V log Pr, (T) (r (w)) || /d+
N M
En.m 8 to (Blntt (w)) log Pen (Dtnita,, (wi) ||
N
M
s trum (Zen+r(ww) (w)) V log Pt., +ra.m (Ze (
S S
Cn.m
==0
m==
where the last equality is because the process y - (w) under measure p follows the backward SDE (B.5).
Thus, by Theorem A and plugging in the iteration relations above
(pylee DKL I
N-1
(Fo1(bl1
+
Buu-Bls, (Ar((w) + Br, (w)[ + e-
n=0
N
N-1
< (Pol(Pol00) +
Σ
cdhn + edlogn Σ
Bu- [Be, (w)) + c K e
n=
==
-T + cd(( log
82 -c- Ko-KTT
dT S de + cdT + 5² + ee- -dT
as T 2 log hn $ 1, and 82 S 8, and then it straightforward to see that the following choices
of parameters
T = (log(dd 2))d h= ((1, = (log(d6-2)) θ (d 82 log (do 3) M = O (dō log (dō ()))
K = (log (d))
would yield an overall error of (8²))
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 (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 [ 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 parallel correct the distribution of the samples. The algorithm is summarized In Algorithm Z.
Parallelized Predictor Step The parallelization strategies in the predictor step are similar to those the SDE algorithm (Algorithm M.. The only difference here that instead of applying Picard iteration to the backward SDE as in (3.2), we apply Picard iteration to the probability flow ODE as (OK),, which does not require i.i.d. samples from standard Gaussian distribution. As shown in Lemma the update rule in the predictor step (Cal) in Algorithm Ø is equivalent to running
30

---

Algorithm 2: PIADM-ODE
Input: Do 0o (0, Ia), a discretization scheme (T' (h ) and
satisfying (G.JI), parameters for the corrector step (7 Nt, ht, Mt, ct), the depth of iteration K and Ki, the learned NN-based score of (-).
Output: sample ýr qr Pr-
1 for n = to N - do
Predictor Step (Section (T2)
Tm Dt, for E [0 Mn];
(0)
for k = to K do
(R) t Dr.: Vt..0
5
for m = 1 to Mn in parallel do
1 Iu.m (k-1))
y Fn.m
e
Vtn.
2
(C.1)
+ 2
mm
2==0
Tmim Tmim-Fagt1
(eens - 1) 8 trass (Fen
end end
D Corrector Step (Section a... Hen, and (00)
(K)
Vi (O,
for nt to N do (7)0)
12
tn.ntht.mtet Lnnthh (n.nthtmmte mtee) f- (itttt (tn, for m E [0 Mt); for =1 to K do
(0)
13
14
(a(k))
((k)
15
tn 1n.nthhp tt. or in.ntht. -- (Utmintht ktt)
Eg¹ M 0, 1, + + e1 )ee -f +1))))))+))))))))))))))))))))))) Ia) for
16
j 0 M'];
for ml to Mt in parallel do
17
18
(Rt) tn attett (kt) nn htmme
u -1)
G(mtet)
ta
4 (k--1) -1)
G((m! ((mll-jt- 1-1-1)cl) -1))) (Ia-
(ā -11-13
(C.2)
8
In jtct))
0
+ Σ
G((mt - 1)et) -1ett)
Est
===
19
end
20
end
(ū (un (nt +1)ht; Ur (nt+jjat) < (n) (K(
(K))
end
23
Yints f- Ti;
24
end
31

---

(k)
tnnttht rt: (matthtrr)
(R)
(K)
ylast
inn 0
kn 0n.
9ma-1
TV (#t
L,,Tt L r+)
(Theorem A.
π tn.rt
DKL O (T... IT (Theorem
T TT)
ha t..
(R)
Vt...7
T t,,
We(Dt, tie Pen (Theorem (W))
Ptn
TV TV (t,,,. Tt, Pt...) (Lemma C.
Ptnt1
t..00
n 2n,2t Ptwtt
Figure Illustration of the proof pipeline of Theorem 3.5 for PIADM-ODE within the n-th block.
the auxiliary predictor process (MCC The auxiliary predictor process takes in the result from the previous corrector step (or the initialization if 0) and outputs V.,,y, as the initialization for the next corrector step.
Parallelized Corrector Step The parallelization of the underdamped Langevin dynamics is simi- lar to that mentioned in Section 2 Given sample resulting from the predictor step, initialize the auxiliary corrector process (Definition C.S) which is an underdamped Langevin dynamics with (K)
the initialization üt,, Din and the augmented variable it,,0 NO Ia) representing the momentum.
We run the underdamped Langevin dynamics for time T which is set to be of order Ω(11) so that it large enough to correct the distribution the samples (ef. Lemma CT while being comparably short to ensure numerical stability (cf. Theorem C. Following similar strategy as in Section 2 and Algorithm III we further divide the time horizon T into Nt blocks with step size h1, and for each block the block length ht into Mt steps with step size el. Within each block, the underdamped Langevin dynamics in parallel for K iterations. As shown Lemma eaL the update rule in the corrector step (C.2) in Algorithm is equivalent to running the auxiliary corrector process (T
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 E [0 N - 1), we define the auxiliary
predictor process ) the solution to the following ODE recursively for k. E [0
(k)
K 1):
dy
(k+1)
2 Dini
1 (_(k))
1 8/0 - (T(tt (D +
2
dr,
(C.3)
with the initial condition
-00) Pt.. Vi, for [0, hn], and yr.. It., for E [1:K]
(k)
(C.4)
where It., E NIA! in E [1:N and St.
N
NO, Ia). We will also denote the probability
(K) distribution V qc...-
32

---

Definition C.2 (Interpolating Process). For any E [0: N -11,
11. we define the interpolating process
(re,o)) the solution to the following ODE recursively 10 K 1]:
(A-1))
1 (k++) Utmit
1 + 2 (rn(() ((insnett)
(k) dr:
(C.5)
with initial condition (1)
Vinst
(R) (0)) Vtn for [0, Anl, and Vi, Vi for k [1 K],
(0)
Vino
-(0) where Vin.0 Pe.. We will also denote the probubility distribution of Dins B...T
(KK
Similar to the equivalence between (5.4) and we have the following lemma: Lemma C.3 (Equivalence between and
For any n E [0 N -1,, the update rule (C..I in Algorithm is equivalent to the exact solution of (C.5) ) for any k E [0 K 1] and T E [0, kn].
Proof. Rewriting and multiplying e & on both sides yield
d 2 Via.
(k+1)
-
do (k+1) (k+1)
2 ymn
dr
e
(k) 8 ton(t) (inien(t) dr
Integrating on both sides from 0 to r implies
/-- e i
1 Ytns?
(k+1)
(k+1)
Ven.o
J 2
In (ntaa(r))
M
A
(k) Ylaitomm dr
e lo tta,m
m=U
TAbmm
Mn
FAre.m m
o taats
m==
and then multiplying e on both sides above yields
Mn
(+++t) PP... = e e
Vr.,
(R+11
PhTmm 1
1)
OVI ovir Tomm
St tham
m=D
Plugging in T Tn,m gives us as desired.
Lemma C.4 (Error between the interpolating process and the true process). Under the Picard itera- tion, have that the ending process (o!.)) (i...rrrejo.k. satisfies the following exponential convergence rate
K No. 3 12/22
sup rE k. where
E
Vin.t
(R)
12
< 3d
1
hn, + E(Dr..)
2
eg L2/2
An
Dr,
(te
86 & (Te, }| dr' dr'.
(T)e,
Proof. Recall that the backward true process (āt, +r Feeoom satisfies the following backward SDE within one block
1
dā
E +r + v log (o dr.
(C.6)
By subtracting
(C.. from
we obtain that
d
(k+1)
(k++)
Vi... Blntr
Dtun Bintt]
dr
1
to (g) (Ulmlem(tt) o (n+on(7) ( (E,, (()
+
2
(C.7)
1
toutan (r) (Ten+on(ts) V log Pr.. + (r) (% tn (on)t))
+
2
+
1
[ log P.,. t.. (r) (Er, +om(t))
Vlog (Z).+))
2
33

---

Then by
3 Binnr
2 = 2
((+t))
Ltn ( ) T
(hitt! )
and integrating for E [0, Mnl. we have
v(htt) Tentrl ||
(iunir tentrr)
8 ga (r') (Winlen(rr)) ( (+ (o.n))))))) dr'
+
(olhtt ton
hh to (Ern
log ton (T), (Tin dr'
+
(ilitt )entrr)
V log (0 Vlogpr +r ( dr'
+
y Pinnt'
(k+
E +
2 || dr'.
Using AM-GM inequality and taking expectations on both sides, we further upper bound the sum- mation above
(k+1)
Btntt
3
S. IF
(((+)) Ve M
Trn + t=1] dt'
+
2hn
+
hn
E
8 la tr (mn))a(T))) o (t (tt dr
2
hn
11, (g ( ( (2) - V llog Dr,, toa toa(r)) do'
E
2
Jo
log Pt., (79 (oa())n--)(())))))))))))))))))) log p (Tr... + || dr'
+
E
2
DDin
<
1 +
3
E
(++1)
dr'
hn (T8% + +E(D...)
2hm
2
12
E
(k) en De (T')
+gn())(7))
dr',
where the last equality is by Assumption M Applying Grönwall's inequality, we have
E
Vinit
(k+1)
B.,,
2
+ Zunz
L2gh
27
cl1+ ) (782 + +E(Di..)
Vin.an(t)) Dr,, tom
dr' +
2
2
re
L2hn
E r' E(0,7]
sup
Vin.T #entr 19/
27
c (2+ thn ) hn
T + F.DD..))
2
2
(C.8)
and by taking supremum
(++1)
x B1, 11.tT/ 11²
sup E
Vin.r
(C.9)
h h
2
['E(0.r]
sup
E
Di,
+'
2
+
pn chn+thn
2
(hno2 + B(Dr..))
34

---

Given that constant hn is sufficiently small, which ensures L < 1, iterating the above inequality for [0: K gives that
(K) Vin..
sup
E
I +
re(O,bm]
h
K
110) Vinit
|3
chn 2 h/
(nno + +B(D..)).
sup
E
Be,,
1 e 1/2/2
TE km
Notice that by Lemma
A.N. we have
2t 1']
IE
Vtn.
=
S 3d, [ ttn+t 11
substituting which into (C.9) then gives us that
K
sup E Te 0, as desired.
[K) Vinir
x to t0
19
s 3d
ch 2 hn/2
(npo + +((D..)).
2
e 12/2
Now it remains to bound Cr, and Di, in Lemma C4
We first bound Dr, using the following
lemma:
Lemma C.5. For any n E [0: N 11, we have that
E Dt,, S de hn.
Proof. For any [0: N 2]. have T' bx+1 2 (1) and thus by [ Corollary 1] that
E v log PeN tm (ity ++mmm V loggiN 1tr ( 1
S de n,m
for any r' E [= mn, il. and thus
hn
v log Po MogPentonnt)) (Fentarrr)) V log D DogDi.tr ( tr:) [r))|/ dr'
E (Dimll
=
E
0
Ma
Ptg (Foutramm) V log De, + ( ( 11" | dr'
Σ
E
m=0
Tn.m
M
2
S de h
n.m
m.
m=U
For N 1, notice that by the step size schedule (ef. Section k.. and suppose S 1/2,
have
T- ST (T( 5T-7,
and then again [ Corollary 1] states
27
de
den,m
v log Prn (Zeg +emm) log Pen+or (Te, +r,)
x
Tn
and thus
E
[DNN-]]
E
Pey (B log Pry + (D (BeN- dr
My-1
Thim-1
Poy Pay it (#ex_iton(r))) V loggey (wt stri dr
E
V log
m
Tmmm
My 1
de2
My
( I-
<
Σ
En.m
Σ m=0
de2 de S
drdr S de de
T
Tn.m
===
35

---

Remark C.6. The above lemma is able to achieve a better dependency on compared to Lemma B because the backward process (3I)) is now a deterministic process in the probability flow ODE formulation, instead of stochastic process as in the SDE formulation as in Lemma B. Thus, intuitively applying Cauchy-Schwar rather than Itô symmetry gives us a O (eee))deeeneeccy rather than O deeeendene.
Theorem C.7. Under Assumptions
B.L
B.A B and B.A.A then the distribution ht that the
parallelized predictor step generates samples from satisfies the following error bound:
We(@tn Ptars S de- K + -12,00 - -de
for n E [0: N -I]. 1].
((K) Proof. By the definition of 2- 2-Wasserstein distance, we have for any coupling of Dinch. and Tt, tha
We(@tnahnn Pents) )2
E
(K) Vi..
Tr. the
|
and therefore
W2 We(Utn.hn Penes )2 $ E
Pin.hi It,, the thn 1"] 1
sup
E
Vinst
(K)
Il, 1
27
r0(0...))
12 L2
an 2
<3d
(In0's + :B(D
2
1ig 1.2/2
Sde K
12
de
h2
where for the second to last inequality we used Lemma C.4, the last inequality is due to Lemma and the assumption e², L
1.
C.3 Parallelized Corrector Step
After each predictor step, we run the corrector step for (1)) time to reduce the error. Particularly, we apply the Parallelized underdamped Langevin dynamics algorithm [[3]]] to the corrector step. which yields (11) approximate time complexity compared to the ordinary implementation of the ULMC dynamics as in [[[ In the following, we will drop the dependency
w for notational simplicity, and we refer readers to Appendix A.2 and B. to review the change of measure arguments and the application of Girsanov's theorem A We will also use general notation *1 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 (t Dt, )1z0i
dd Utn, at det ddv
(C.10)
- et det - V logg (u, )dtt /2ydbt...tt
(Kt) with initial condition ut,.0 Vinth from the predictor step and vt,,0 N(O,, Ia), where (b hzz Wiener process. matrix form:
We may also write the system of SDEs above in the following
Ue Ur Vinnet
[
Ia -Ia
utmnee Vennet
) det
T O
b bl_,et
V Vogg Proes (Me,,
2
Id
bin
We run this underdamped Langevin dynamics until the pre-determined time horizon T We also define the joint probability distribution of (u,,,,t, (t,, ) at time Te,,,t (Ut, (t, and its marginal on ut,.tt as ",,,, (ut,,et).
Similar to the parallelizing strategy in Section M we discretize the time interval [0, T!] into All blocks with length ht = =T /Nt. Within the n-th block, we further divide the block [ntht, (n+1)) into Mt steps, each with step size et = ht ht/Mt.
36

---

Definition C.8 (Auxiliary conrector process). For any nt E [0: Nt - 1}, we define the auxiliary tn,mtht, (rrte as the solution to the following SDE recursively for kt E
corrector process (a(k))
kt -1]]
du 1)
bott 1)
((++1)
tnithtt drt,
(C.11)
do faitthi,r
(x+k+1)
-no (k++1
8tut1 bayei g. + an
(kt)
with the initial condition
(0) nt .rt mith
(ke) to h ni for k E [1: K*),
for f [0,,10 and
((())
ta,thto = ntttt
(C.12)
where
(Kt)
Dt x((k))
ū (n.(nt-1)
for [1: 11. and
Y Vin.hn.'
(K)
Ton,0 NO La).
We define the joint probability distribution of (u,,, th De,,tt) at time as (t,, (ūt,,tt (,,, and its marginal on Ut as Tt tn (1).. We will also denote the resulting probability distribution of
t.rrt as Ua+t1"
Lemma C.9 (Equivalence between (22) and II.. For any nt [0 N - the update rule in Algorithm B equivalent to the exact solution of the auxiliary process (C. for any kt E K 1 and r E [0,
Proof. Without loss of generality, we will prove the lemma for mf Mt. The proof for mt [0:
M can be done similarly.
We first rewrite into the matrix form:
(kt) bnt v tnathtt
0
0
Ia
(kT)
0
| drt
=
7la
(kl)
St.
(kt) tn ht..ttttt....)) (T))))
(C.13)
0
by
+
d
Define the time-dependent matrix G(.) as
G(tt)
=
Id
1-
Id
Ia = exp (0:
(1. )+)
(C.14)
0
e
1t La
satisfying that
d
0 Ia
G( G(tt) O
Ia -yla
G(I
di
- Il
Now multiply G(-rt) on both sides of (C (C.33 to obtain:
(ktt))
) = G(-r!)
G(-r))
drt
(kk) tt
Sinti Commiss
0
b'
+ G((rt)
d
0
2
37

---

Integrating on both sides from 0 to ht and multiplying G( on both sides, we have
(k))
(kt)
(k i
G(hl)
to tt,
((k)) mnntttto
L G(hl ((hl->l')
h1
drl'
Stnl1 eaintht (ttr)))
/ G(ht
-
0
b
0
V27la
M (m)+1)e
G(hl - =l')drt
Stara (ox (
mnett
tn mtte
(m (m)+
G(Mtt -+t) |
b
Σ
V2
ml mtet
M
Σ
(G(c) Ia) G((MT - m -1)ct)
Senti (zi (zilkl) (kt) bannth
m
MI
(1)+
G(hl G(Ml- -0
b bminttttrr binnnthi++
/2
ml =
By Itō isometry, we have
(m(+1)e
-0 I2 d
by
G(hl
b b
~V
0,
0
0
G((M - ml - - (G(e!)-IA)
0
V2
(G(ct) Ia) G((M -m² m - 1)ct) [0 v
0 N (0.2 2) - 27 faattmm+1yy +111))))
as desired
Definition C.10 (Interpolating corrector process). For any nt E [0 N 1). we define the interpo- lating corrector process (ū (ke)
-
nthtt Dreeeo as the solution to the following SDE recursively for
kt [0: Kt 1):
du (k+1) bi 1) (k+1)
(C.15)
do (k+1)
-yo (k++)
(k))
hn Stats ta (t + 22 ht
"nattht,=t
with the initial condition
u (0)
(k+) tyynthtrr t ((k)
Sinnt .rt ū,_. nhh
for E 10, and
for k E [1= Kt],
tnathht = Ut
(C.16)
where
U ntt == 5 (Kt)
1 u = (KK)
fornt E [ Nt - 1), and
U,,,0 V Diniiinn
-(K)
ūe,,o ~ I2).
38

---

We define the joint probability distribution of (i,,,,t Dt,, at time as " (Utn 0 and its marginal tt
We invoke Girsanov's s theorem (Theorem A.4) again by the following procedure
1. Setting (A.2) as the auxiliary process 1 iteration Kt, where bt,,ct (w) is Wiener pro- cess under the measure Q;
2. Defining another process b mithttrr governed by the following SDE:
db nti,t db, nf.t.ttrrr Pt ht (5 )d
(C.17)
where
tt (r+) V2y St (u (k) f.
(ū) (T(Kt) en mlbttrt))
V log Pinpi
(C.18)
and computing the Radon-Nikodyn derivative of the measure P with respect to Q as
=exp (1. sat (Tt) nttttt
1 ph I1 (()m((t))((
dp
/
(C.19)
dQ
2
3. Concluding that (C.T5) at iteration Kt under the measure Q satisfies the following SDE:
ntt U ((( dē dv:(kt) nibtt drt
dk(k))
fn,ntht,rt - tt drt V log Pe.es (ū (Kt)) nt )rr + /2 hhl
(Kt)
(C.20)
with (b, (ttnttttrtrrt being a Wiener process under the measure P. If we replace o(k))
(u(K)) nint,rt .ti wthirrt) (trrrt) by (Utm,m/ht4r n ht+tt) one should notice (2.0 is immediately the original backward SDE (CCII) with the true score function on E [ntht, (n + 1)h7]:
du,, attitrrt U1 a dn ( at log Poo.P (M + /2rdbb thttert ht4tt
(C.21) We further define the joint probability distribution of (u v.,))) at time as a (Ut (t,,et and its marginal ot as T (H),,,t))
Remark C.11. The application of Girsanov's 's theorem A 4.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-Nikodyy derivative (K.T9),, we may define a stationary underdamped Langevin process for nt E [0 N - 1] and E [0, hi] as
duu mtt vi ht dr dv tn +y -M0 h14,,drt V log Pr., (u ( + V2ydb ht+rt;
(C.22)
with the initial condition u ntt bility distribution of (u at (n ar) at time I T (1 (t, and its marginal on u. as
htt ~ Penes and oi whitt ~ NO, Ia). We define the joint proba-
r. ( (ut,,,))
Thus, from Corollary A S. we have that
DRI(TT (tn))
N1
<DKL (tm (n- 1
+
Σ
DEL (RL(TTe (T on (1+1)) ((+()))
(C.23)
bomi (1 /² drt
Stats (It-
log Pr Proae (x(kt)
+ Ep
47
39

---

By triangle inequality, we have
2
8t 1 ④ (x(K-11)
ta bt, b 11 log Pin-.
(ā(ht)
tn rt)
drt
<5
h1
11 ④ (1)
montht 1iftc - 8tati (a(
2
drt
mnontht lct
2
-5
montht, 1
stati (u
let foonth + 1
drt
2
5
Stat1 1n
log Pents (u
ba ba lct
dr
ht
5
log Penpe (u
la Moathh+11 tet - llog (w) m utht
drt
ht
V log Pongt Ptnti m V log Prop (K))
Pan || drt
-5
< 5
② (1k-
inattt,L - ((*))
Annutht 1 dr1
dr + Li / h1
(k()
(k)
2
-5
inittttrtt - 2 a
drl
0
-5ht82 + +,,
Jo
u
drt
(C.24)
where we used the Lipschitz continuity of the learned score function (Assumption (.3 and the true score function (Assumption B.4),, and the -accuray of the learned score function at each time step (Assumption D.
Now we proceed to bound the terms in the error decomposition (.24
We first bound the Fi
term by the following lemma:
Lemma C.13. For any E [0: N and - [0, have
S d
Ep
u
and therefore
Ep [F. & dhict?
Proof. By the definition of (u: namics (C.22). we have
nthtt D nlh4 as the stationary underdamped Langevin dy-
1
27
M
Ep
ui atnhhttr
Ep
||7 dr S det2
Set
Ep
1ett
where the first inequality follows from Cauchy-Schward inequality and the last inequality is by the fact that
0 ht ~ NO.. Ia). for any p [0, kt..
40

---

Consequently, we have
Ep [ = //²²
pht Ep
" miht +rl drl & dli dhice?
u (nntth +1
+L
The term Et bt can be bounded with the following lemma:
Lemma C.14. For any nt :0:Nt -1], the following inequality for any - E [0,
11, suppose that y S Lp 1/2 and T & Lp 1/2 then we have
Ep
(K))
to mlbltr
S W2 (ttmihe P(ns1):
and therefore
Ep [Er S 11 12) (Ct.. hn Pen41)
Proof. Recall that under the measure P. u (k())
nt follows the dynamics of wn I| for T1 E [0, which coincides with that of u, mt.htt++" As the only difference between the two processes m!!t,r and u," in is the initial condition, we can invoke Lemma 10 proved
in [| to deduce that
(Kt) ht u min [|]
Ep u
S
Pent1);
where the assumption that y S Lp 1/2 and T S I,p required.
Now notice that "h nthttrr and wt also follow the same dynamics with the true score function for 7 n0,nthl for any coupling of u and u, nttt we have
W2 (T ntht Ptn.s S EE
[ ( nt 112
WW o. Pt. ( Pewer)s
where the last equality is again by [ΑTл, Lemma 10). Therefore, we have
Ep [Ei..ttht] ht
Ep 1 ④
(K*)
u t
13 (K
tninthtrr t ht +-1
/ drt
In
<h (L2 + 17)) WT Mo Pr...).
Now, we provide lemmas that are used to bound the first term in (C.24)
Lemma C.15. For any n E [0 N - 11. we have the following estimate:
sup
Ep
u
r1E(0,b1]
5L hi (3+7)A)
ū (K1-1)
(KK)
27
sup
Ep
27
5h S (3+7))+)
E + hi02 + (ntt + 1it2 (3+ts)n (37d + M²) + 2ht
Ep
2-
41

---

Proof. Let ntt,t == u t..ntht,r1 tt.tt and U thitt ((1)
(0)
((())
Then for k may rewrite (C.T5) follows
dupe (lunn (0) de demmathht = -V( + tn.ntht, (t )dT
(n drt
6(0)
8t, (u)
(0) miih )dr + /2ydbt tt
(C.25)
On the one hand, by using the first equation in (CC25)), we may compute the derivative
dr | 2 miit (Mi,,
((0) tn,nt
and integrate it for 7 E [0, which yields
M 112 L- Monntaatr (um + (:(0) tn dot' x2 L / drt' + Lo [1 dot + S / Part'
drtt'
Applying Gronwall's inequality, we have
| st 1|2 < 2r
I | drt +
//" n or l dr")
We then take expectation with respect to the path measure P and then the supremum with respect to 7t [0, ht], implying that
[|
Ep +tE[0.A1]
sup
2xr f"r Ep /| drt' + + ( Ep
(0)) At
1 dr
sup -lg[0.A1]
< h htee2h 2h
sup rte 00,0
Ep
| |2] + 2te2h" d.
(C.26)
On the other hand, by applying Ité's lemma and plugging in the expression of b httttrt given by (CT,)), have
d (lnatthtzzl
27 U 1ttt,, 2yy, 2 tnattht
(0) ( katntht + Tt tt (Th.nttttss) 2nd drt
(0)
+2u 22. nt.hthrr v 27 (db, (db + tt (F
(C.27)
Then similarly, we may compute the derivative of (1 IP, integrate it for 7 E [0, and take the supremum with respect to ! to obtain
Ep
-Ep
2
" (0) ( mnntth U tn - 2nd drtt
1" 2v tnnith .: 8tats
(0(0)
tnintht drtt
+Ep
50" v.
V. (ab (dbmmntht+tt" + Φ aat (rt) )art))
/2/Ep
42

---

By Ito's lemma, this equals to
Ep
|| nitht [1?]
-Ep
2711 syllummnnttt
+ 27 in mi rrt 1(0)
inottht 2nd drl
u (0)
tmmolhi m0)) + 2 27 mnnt u Pen (7 (rt))r))))))
+Ep
2v mn tn P 88,11
Applying AM-GM gives
Ep
Pn.nh',r [12]
(1 + wot , + + (T")) drl'
Ep
? b (0)
tnintht t 1| + stats (z) (+t |1 2nd drt'
Ep
/ Ep (1 1)
1| nohht Nth mtrl Pen (T |13] drt' + 7B bniitat ai 1'] +M3 2nd 77
(0)
=(1 + v) ( Ep (I [1*] dr f." Ep
|1 (I
(7 ||3]
dr + rl (3yd + M²)
where in the last equality, we used the initialization of the auxiliary corrector process 6. (( Ia).
Again, apply Gronwall's inequality to the above inequality and take the supremum with respect to rt E [0, h to obtain
sup Ep [1 on [12] rE(00At1]
/ ntt
(e(1+))
Ep
[ll (ntt drt + + Nte0+si (3nd + M2)
(C.28)
c(I+))ht
2
Ep
stree (ū (1))
V log Prmps (ū) (Kt) ( rt
drt
2y
+ ht (1+)()))) (3yd + M²).
and for the difference term within the expectation, we decompose it again by the triangle inequality i.e.
8 St
W (1)
11
log P
(KT) (t
/|
drt
rhtt S572
-1)
(KK)
dr + 5 5F + 5h' + 5T.2 mi
i taintht,L#eee
Ju to obtain that
in Ao Je
Ep [| [13]
sup =tE(o,At]
Ep [
((Kt)
drf
27
to,ntht,12f1 lct
5e (1+7)1
Ep EL + hto2
+ at + (37d + M2)
27
5L26
ū (k)
((tt)
27
ht
Ep rTE(0,A1]
sup
tountat,
5e (1+)))h
Ep nin + h102 + L2 F ntan + (++()+))) (3rd + +M?).
27
43

---

For any n! E [0, - 11, then the two ending terms i x(K)) ntt,, ,t and T(K") ntid of the sequence substituting which into (C26 completes our proof of this Lemma.
Lemma C.16 (Exponential convergence of Picard iteration in the corrector step of PIADM-ODE).
(u(k') niht ( satisfy the following exponential convergence rate
20 20(t-11)1
sup
Ep
r'E(0.A)]
5h c(1c(3+))a
Ep Et + 12202 + Ft + + (3+r))h (37d M2) + h 2h
<Ox1
2y
(C.29)
where the coefficient
L
kt-1
((+-)))
Lihh Eh?
Kt-1
Ckt
2
2y
27
Proof. We subtract the dynamics of it (k+1)
atht,rt and ū (k)
wthttrt in (C.55) to obtain
(k+1]
u
(k")
(x(k+]
((kt
n' drt.
Then, the formula above compute the derivative
d
(k+1)
(
(k(+1) atht,rt
T (((+1)
drt
at
ntht
and integrate for [0,7!] to obtain
(k+1) ((k))
-2 /-" (w(k+t) (kt)
(((+))) witt
((A1)
wibttrt)) drt
< 1
(k++1) u att
(ka) nt
||
2
drt'
/"
((k+1)
-0(1)) (1k))
ntht drt'
//2
Applying Grönwall's inequality gives us that
n (k+1) =(kt) wihttrt nnthtrr 1| S er
((()) ntht
dr
and taking the supremum with respect r [0, on both sides above implies
sup Ep
1 (k- x(k))
S hhth"
sup Ep rle(o.t
[100et1) ((k)) atht,rt
ntht .71
rte(0.11]
(C.30)
We then apply similar argument for σ (k+1)) pthirtt - v ntht as well
(kt)
(k+1)
(i(kt)
v ta,nttt,tt ((.
(w(( go(ket
ti (nottttt drt
81mns (+)
tn 1 lot) - 8tats (ū (t-1)
twinth aot,, ((t) dr!
-
44

---

integrate which for for E [0, to obtain
b((+1) Nt nt ||
(((k)) blet1)
(k0) will mr
drl''
27
1
-2 !,"
(((tt) 2(tt)
Stati (2)
bnjuthht 1 let) - Sents (eilk-1)
booetht 1 ( drt'
nibitrt!
ft (T(
(2)
baittht 1 )ot) /
drt
mnnttal 1 jot
8t
27
2
(k-1)
drl'.
2y
tn.nth 1
And then taking the supremum with respect to rt E [0, on both sides above implies
1k+1) atttert Wo
g(k))
htT22 sup
(kt) M
((K-1) Vn he A
sup Ep
Ep
+ =lE[o,A1]
2y r E(U,A
(C.31)
Substituting into 1.50 and iterating over & E [1
Kt 1]. we obtain that
T.ght2 2
(k') nthirtt
(Kt- 11 nthttrt
I (K*-1) ((K--2)
sup Ep rtE(
sup teE(0
Ep
2
KT-1
Lihh? E
sup Ep
(0) mntntt .rt
2-y
-te(o,At]
L 12 E
K1-1
5h c3+
Ep E + 10 + T Pen
2-
27
1.2 2 el
K!-1
(1t3 (3+T)))) (3yd M2) + ht222 a)
2y
kt-1
1 c1
(kt (KR)) buattht
sup Ep rte[o.Mt]
2y
24
where we plug in the results from Lemma C.I the last inequality. Rearranging the inequality above completes our proof.
Theorem C.17. Under Assumptions of the parameters
and 3.4 given the following choices of the order
T =1)) Nt =1),(1), hi h²= 0(1) M = - et (dd-1/2))) Kt = (log(d8-2)
and let
L2ht2 ht
1 754,112
T 800 28
2
then the distribution T satisfies the following error bound:
(RI. (T. ( m1) ST W2 VT . P(...) + + + dT -
SW? h. Dro-.) 82
with total of KtNt Olog(ad approximate time complexity and M - (1-128-2) space complexity for parallalizable d-accurate score function computations.
(log(dā-2))
45

---

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.
(1)
(ū(K')
Ep
Stat1
V log Pt.. ton
dr
551.gh
sup
Ep
i bnntthi, 1
(((1)
+ +5Ep [Fmmntit + h 11822 + 1 (mntt
rtE]U,kt]
5hie(3+1)h
5h
<5
1
C
Cxt
Ep [Emmnntt + hh122 + [n
27
+5LShOCr Cki
(h12 (3+-)n (3yd + M2) + hte d)
where the last inequality is by Lemma C. We further substitute Lemma C.4 and CT into (C.233 obtain
DKL (xL,
1)nt [ 11) + 27 + k22Kxth 5hh ((3+)) (
<DKL. (TKI(Te.. (Tem.
Ep E + h182 + (
4
5LghtKKt (hT2c(3ts))))
(3+ (3yd + M2)) + + 2h
ent d)
27
DKI..TT.. (Ttn.(n-1)111 (1-),,)))1))
27 +52 L22Cx1 5h
(L' + (Tm Penn + 16 ddile!?)
+5
472
5L2 (1t2 (1t22(3t)))) (3yd + M2) + ht 2h d)
27
SDKL. (Tem.(n-1)A1 (1))))))))))))) h W (Te Menhnn An Pr., + +ho02 + ddhlel2 -e-K
and then sum over n. to obtain
(RI ( (TI Vli
(DKL(Fen (em..) .) N N (t hn Penta) + Nt1t82 dNthtet te-k Nt
=IT (Tt Powee) T ditet -e-k'tthd
Then, it straightforward to see that when the following order of the parameters holds
T 0(1) ht = 0(1), Nt = 0(1) el = (d 1/2),, M = (28 Kt = Olog((dō 3))
and 6n S 6, we have
DKL [T....)) T+) S W2 (qtn. Pen+i) + 82.
Lemma C.18. Suppose T S Lp 1/2 then we have
TV P S v
S
W Wa ) Wn W2(Ceni Pro..):
L (TT)
Proof. A complete proof of the Lemma above is presented in (Лл]) Lemma 9], which derived based on [155 Corollary 4.7 (1)
D
46

---

C.4 Overall Error Bound
We are now ready prove Theorem
Proof of Theorem B.N. Notice that the interpolating corrector process (ūt; (ii ,t 8 b is constructed to follow the same dynamics as the auxiliary corrector process (Ut, nt ott D atttt in the corrector step. Therefore, we have by data processing inequality that
TV ET., ((.,,.)) S TV(T) (t,.0) TV h., H...
(C.32)
and again, since the interpolating predictor process nttt is constructed follow the same dy- namics as the auxiliary predictor process nt in the predictor step, we further have by data processing inequality that
(V Itn.h..) & TV (o..o = Tr. Pt..)
(C.33)
Furthermore, applying triangle inequality, Pinsker's 's inequality along with Theorem CA i and Theo- rem C.7 proved above, we may upper bound the second term above as follows
TV T.) S D (.....+) S W2 (atn,hn Pente) 82
(C.34)
Summarizing the above inequalities, we have
TV T Pt.. - T (..Ttt )²
rt.
TV T Th Tt )² TV t.. T T
(T
u.u"
TV (R Ne. rt) )²² TV (V 27 Ten
(C.35)
TV TV Te
T rt )2 + TV (T,, Tt Pt..
< W (Tt P(,+s) + 8 + W Mn Ptwes)
2 de
12, + de 82,
where the second last inequality is deduced from Theorem M and Lemma CAS and the last in-
equality derived via Theorem Owl. Therefore, for any n E [0: 1]. applying triangle inequality along with data processing inequality (cf. Theorem A. DI yields
TV Pro..) TV( rt..pan))) zu < TV tn TI (.. TV (T,,,
(C.36)
S TV .... ) d-/2
8 + /22/ch, +
where the last inequality is derived by plugging in (C32), (OKE) and (M and summing the inequalities above further give us that
Applying Lemma A
TV Pv) TV P(N)
N-1
$ TV(Go:P)))
Σ
12 hn, + d22 +6)
(C.37)
$ / + 1/2 e KK22 + +6o0 + /2 +
By setting the parameters
= (log(dd- )), h = 0(1, N = 0log(d88 )));
e -=000-10000000-1----1) M = Odd log(d//2-11),
(dd11/28-1)
K = (log(dd
and letting dn S 8T S 8 log (do-2), finally obtained the upper bound
TV( pon) de de + Nd + + + < 8²
as desired.
47

---

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 the work performed by the authors?
Answer: [Yes]
Justification: We have discussed the limitations of our work in Section R (Discussion and Conclusion).
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and 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 conclu- sions of the paper (regardless of whether the code and data provided or not)?
Answer: [NA]
Justification: As theoretical paper, we do not include experiments.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: As 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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: As theoretical paper, we do not include experiments.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined other appropri- ate information about the statistical significance of the experiments?
Answer: [NA]
Justification: As theoretical paper, we do not include experiments.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: As theoretical paper, we do not include experiments.
48

---

9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http://///arrp...cooa/a//////aaa
Question:
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and have ensured that our research conforms 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 theoretical paper, do not expect direct societal impacts of our work.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (c.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: As theoretical paper, we do not release data or models that have 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 theoretical paper, we do existing assets.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: As 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 pa- per 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 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: As theoretical paper, do not involve crowdsourcing research with human subjects.
49


Reviews:
[{"review_id": "Cnuvre074S", "reviewer": "Reviewer_GP1D", "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.\n\n[1] Shih et al. (2024) -- Parallel Sampling of diffusion models\n\n[2] Anari et al. (2024) -- Fast parallel sampling under isoperimetry", "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. \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", "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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* 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", "weaknesses": "* 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\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed feedback, constructive suggestions, and kind affirmation of our work. We would like to address the reviewer’s comments in a one-by-one manner below.\n\n---\n\n## Weaknesses\n### Regarding the presentation of results\nWe thank the reviewer’s suggestions regarding the presentation of results in Theorem 3.3 and Theorem 3.5. Acknowledging your suggestion, we will revise our paper to include such a bound with explicit dependencies on the step size $\\epsilon$, the number of iterations $K$, etc., enhancing the utility and applicability of our results.\n\n### Regarding memory reduction with ODE formulation\nThank you for pointing out this confusion. Here is an example to illustrate the intuition behind the acceleration of ODE-based samplers: Consider a 1D Itô process $dx_t = b(x_t) \\ dt + \\sqrt{2\\sigma} \\ dw_t,\n$ and by Itô's formula for $f \\in C^2$, we obtain\n$$f(x_t) - f(x_0) = \\int \\left( f'(x_t) b(x_t) + \\sigma f''(x_t) \\right) \\ dt + \\int f'(x_t) \\sqrt{2\\sigma} \\ dw_t,$$\nand thus\n$$\\mathbb{E} \\left[ f(x_t) - f(x_0) \\right]^2 \\lesssim \\mathbb{E} \\left[ \\int \\left( f'(x_t) b(x_t) + \\sigma f''(x_t) \\right) \\ dt \\right]^2 + \\mathbb{E} \\left[ \\int f'(x_t) \\sqrt{2\\sigma} \\ dw_t \\right]^2 \\leq t \\mathbb{E} \\left[ \\int \\left( f'(x_t) b(x_t) + \\sigma f''(x_t) \\right)^2 \\ dt \\right] + 2\\sigma \\mathbb{E} \\left[ \\int f'(x_t)^2 \\ dt \\right] \\sim O(t^2) + \\sigma O(t),$$\nwhere the second-to-last inequality utilizes the Cauchy-Schwarz inequality and Itô's isometry. This derivation suggests that when $\\sigma > 0$ (in the case of SDEs), the discretization error with step size $\\epsilon$ ($\\mathbb{E}[f(x_\\epsilon) - f(x_0)]^2$) is of order $O(\\epsilon)$, whereas it is of order $O(\\epsilon^2)$ when $\\sigma = 0$ (for ODEs). The improved step size dependency thus allows for better storage complexity.\n\nFor more rigorous arguments, please refer to Lemma B.7 and C.5. We will expand the corresponding discussions in the paper to improve clarity.\n\n### Regarding the comparability of SDE and ODE formulations\nWe appreciate the reviewer’s insightful observation. While ODEs do not permit the application of Girsanov’s theorem, we have to make more restrictive assumptions and some compromises in the algorithm. Notwithstanding, we have tried our best to align the results as closely as possible by giving KL bounds for SDE and TV$^2$ bound for ODE. Indeed, as the reviewer suggests, the bound in Theorem 3.3 could be relaxed to a TV$^2$ bound using just one step of Pinsker’s inequality. Creating a unified framework to analyze both samplers under the same techniques and assumptions would be both desirable and challenging, and we intend to pursue this direction in future work.\n\n###  Regarding Assumptions\nWe acknowledge the reviewer’s comments on the assumptions. We agree with the reviewer that Assumption 3.4 is less favorable when it comes to wider applicability and is possibly due to technicality as we only require it for the ODE formulation. We regard an improvement on our current bound, by replacing Assumption 3.4 with Assumption 1.5 in [1] (arxiv version), as viable. Assumption 3.3, which is implementable by simply adding penalization terms based on the weights of the neural network, is often adopted for numerical stability. We believe the possible blow-up near the data end could be partially averted with an appropriate early stopping time $\\eta$ and is of distinct research interest that we would like to study in future work. We will expand the discussions on the assumptions regarding the reviewer’s comments.\n\n---\n\n## Questions\n### Regarding the algorithm compared to ParaDiGMS\nWe thank the reviewer for the feedback. We will adjust our paper to more accurately emphasize the theoretical nature of our work, focusing on the specific analytical benefits our methods provide. We believe this will ensure our contributions are properly positioned within the existing literature, and we plan to rigorously assess their practical impacts in future empirical studies.\n\n### Regarding the manifold hypothesis\nWe thank the reviewer for raising this potential improvement with the manifold hypothesis. We believe this hypothesis can be integrated into our framework harmoniously, whereby the $O(\\mathrm{poly}\\log d)$ rate would be improved to $O(\\mathrm{poly}\\log \\mathrm{diam}(M))$, with $\\mathrm{diam}(M)$ being the diameter of the data manifold that only depends on the intrinsic dimension $p$ of the manifold. This would be a significant result, especially for applications in computer vision, where we have $p\\ll d$. We will add related discussions in the paper and possibly conduct further theoretical analysis on this aspect in future works.\n\n---\n\n## Limitations\n### Regarding the discussion of limitations\nWe appreciate the reviewer for pointing out the necessity of a more comprehensive discussion on the theoretical assumptions underlying our results. We will expand our discussion in Section 3, focusing on how the assumptions might affect the generalizability and applicability of our findings. Noting that the performance of parallel algorithms can be greatly influenced by practical implementation and hardware constraints, we recognize that scaling may deviate from the theory, particularly when unexpected communication and synchronization issues introduce considerable delays. To further elaborate on these issues, we will enhance our discussion in Section 2.3 to more explicitly outline the approximations and assumptions made during our analysis of the time complexity of parallel algorithms. \n\n---\n\nWe thank again for the reviewer’s detailed comments and valuable suggestions. We hope our responses have resolved their concern and our modifications to the paper are up to their expectations.\n\n### References\n[1] Chen, S., Daras, G., & Dimakis, A. (2023, July). Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning (pp. 4462-4484). PMLR.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0M7kjYtz0d", "reviewer": "Reviewer_sp2J", "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}.\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.", "questions": "Can 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?", "limitations": "The authors adequately addressed the limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I 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.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the appreciation of our work and the valuable feedback. Below are our responses to the questions raised in the review:  \n\n---\n## Weaknesses\n\n### Regarding numerical experiments\nWe appreciate the reviewer’s suggestion and recognize the value of empirical verification of our theoretical results. \n\nOur decision to focus primarily on theoretical analysis was influenced by existing literature, particularly recent studies in references [1, 2], which have extensively explored and positively demonstrated the acceleration of diffusion model inference through parallel sampling. Therefore, we chose to concentrate on advancing the theoretical discussion to avoid repetition and build directly on these established empirical findings. \n\nMoreover, as mentioned in Section 2.3, the verification of the sublinear acceleration achieved by our proposed parallel algorithms might require large-scale experiments with data dimensions $d \\gg \\log d$ and thus substantial engineering efforts that could be of independent research interest. In light of this, we will expand Section 2.3 to further discuss the applicability of our results and to clarify the assumptions required for implementation.\n\n---\n\n## Questions\n### Regarding asymptotics\nWe thank the reviewer for the insightful question regarding the constant scalings and asymptotic behavior.\nIn our work, the scaling of constants in the two main theorems (Theorem 3.3 and 3.5) is influenced by the constants ($L_{\\boldsymbol{s}}$ and $M_{\\boldsymbol{s}}$) specified in Assumption 3.3., These constants are in turn derived from the properties of the ground-truth density from which we aim to generate samples.\n\nIt is important to note that the actual onset of asymptotic behavior may depend not only on these constants, but also the specific implementation and the empirical setup. As such, further empirical studies are needed to precisely determine when the asymptotics will take effect under various conditions. We plan to explore this aspect in future work to provide a clearer understanding of the dynamics involved.\n\n### Regarding hyperparameters\nWe appreciate the reviewer’s question concerning the tuning of the hyperparameters. We acknowledge the tuning of hyperparameters, although involving considerable engineering effort, is crucial for the actual performance of the algorithms and may vary significantly depending on specific problems and the desired accuracy of the model. In this context, our theoretical work aims to provide insights into the optimal scaling and relative magnitudes of these hyperparameters, potentially facilitating the tuning process through a more structured approach.\n\nAdditionally, several empirical studies [1, 2] have implemented various kinds of parallel algorithms, incorporating advanced tuning techniques. For instance, the algorithm in [1] employs a striding window of adaptive length to tune the step sizes. These implementations serve as excellent references for observing the practical impacts of hyperparameter settings in real-world scenarios. We are eager to further explore and validate our theoretical findings by experimenting with different combinations of these hyperparameters in empirical settings in subsequent research efforts.\n\n---\n\nFinally, we would like to thank the reviewer for all the helpful comments and suggestions, which have greatly helped us improve the quality of our manuscript. \n\n### References\n[1] Shih, A., Belkhale, S., Ermon, S., Sadigh, D., & Anari, N. (2024). Parallel sampling of diffusion models. Advances in Neural Information Processing Systems, 36.\n\n[2] Tang, Z., Tang, J., Luo, H., Wang, F., & Chang, T. H. (2024, January). Accelerating parallel sampling of diffusion models. In Forty-first International Conference on Machine Learning.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pyYyN6zMwR", "reviewer": "Reviewer_jWja", "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.", "questions": "See Weaknesses", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the reviewer’s insightful comments and appreciation for our work. We have listed our answers to the questions raised in the review as follows:\n\n---\n\n### Regarding comparison with other parallel sampling methods\nWe appreciate the reviewer’s suggestion regarding the necessity of contrasting the Picard iteration and other parallel methods. To address this, we will revise our paper accordingly, particularly enhancing the introduction and expanding the discussion in Section 2.2 (Parallel Sampling). Our aim is to provide a more comprehensive outline of parallel methods developed for various tasks and to elucidate the rationale behind selecting the Picard iteration for our study. \n\nSpecifically for the randomized midpoint method [1, 2], we acknowledge its potential for achieving better query complexity under additional isoperimetric assumptions on the target density, such as the logarithmic Sobolev inequality. This observation will guide our future research efforts. We will incorporate these insights and recent advances along this research direction into our revised paper, ensuring a thorough evaluation and justification of the methodologies employed in our work.\n### Regarding total query complexity\nWe thank the reviewer for the suggestion. In response, we have included a remark in Section 3 that compares the total query complexity of our proposed methods (PIADM-SDE and PIADM-ODE) with their sequential counterparts. We wish to clarify that the core idea behind our acceleration approach is to offset the number of sequential evaluations with parallel iterations. Leveraging the exponential convergence properties of Picard iterations, our algorithm manages to achieve $O(\\mathrm{poly} \\log d)$ time complexity, at the cost of a potentially higher total number of score queries. We will make additional revisions to our paper to further elucidate this aspect.\n### Regarding minor issues\nWe appreciate the reviewer’s meticulous attention in identifying the typos and missing references (on line 85, on line 89, Lemma A.6) and have accordingly revised our paper. \n- For the last equation on page 4, we have inserted an explanatory line below equation (2.2) to more clearly define the true density associated with the reversed SDE. We have also expanded the “Step Size Scheme” paragraph in Section 3.1.1 to enhance understanding of Lemma A.6, A.7, A.8, along with their corollary, Lemma B.7 on which the equation in the reviewer’s equestion is based.\n- Regarding the two sets of parameter choices for the predictor and corrector presented in Theorem 3.5, we have added a paragraph dedicated to elaborate on the reasoning behind our parameter selections, including $h$, $\\epsilon$, $M$, $K$, etc.. We specifically address the conditions that allow for the block size $h$ to be of constant order, how the magnitude of the step size $\\epsilon$ aligns with that in previous work [3], and the relationships between the scaling of $T$ in the predictor step and $T^\\dagger$ in the corrector, etc., within the ODE formulation.\n\nWe believe these revisions will significantly enhance the clarity and coherence of the presentation of our paper.\n\n---\n\nFinally, we sincerely thank the reviewer for all your time and efforts. Your suggestions have tremendously helped us improve the presentation of our paper.\n\n### References\n[1] Shen, R., & Lee, Y. T. (2019). The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems, 32.\n\n[2] Yu, L., & Dalalyana, A. (2024). Parallelized midpoint randomization for langevin monte carlo. arXiv preprint arXiv:2402.14434.\n\n[3] 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DEftbkwVRr", "reviewer": "Reviewer_FgeD", "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.", "questions": "Please see the Weakness part.", "limitations": "The authors adequately addressed the limitations and societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tThe 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.\n\n2. 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].\n\n[1] Anari, N., Chewi, S., & Vuong, T. D. (2024). Fast parallel sampling under isoperimetry. arXiv preprint arXiv:2401.09016.\n\n[2] Benton, J., Bortoli, V. D., Doucet, A., & Deligiannidis, G. (2024). Nearly d-linear convergence bounds for diffusion models via stochastic localization.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer’s comments and suggestions on our work. In the following paragraphs, we address the reviewer’s concerns.\n\n---\n\n### Regarding comparison with [1, 2]\nWe appreciate the reviewer’s suggestions to deepen the discussions on our results relative to those in [1, 2]. Along with the comparisons already detailed in the second bullet point of Remark 3.4, we wish to offer the following additional remarks:\n\nIn contrast to [1], which limits parallel sampling techniques to log-concave sampling within Langevin dynamics, our approach significantly applies these techniques to denoising diffusion models - a topic of widespread interest within the ML community. To the best of our knowledge, our work is the first to provide parallel algorithms for diffusion models that are rigorously analyzed and proven to achieve $O( \\mathrm{poly} \\log d)$ time complexity. \n\nAs aptly pointed out by the reviewer, unlike the strategies in [1], our analysis does not rely on LSI-type assumptions on the data distribution, making it applicable to almost any distributions of interest, given a sufficiently accurate score approximation. This flexibility underpins both the empirical success of diffusion models [8] and the theoretical advancements in [2, 3] - further validating our results under the parallel sampling framework.\n\nWe would like to emphasize that the aim of our work is not merely the sophisticated generalization and application of the Girsanov’s theorem, which the reviewer recognized as “novel and suitable”, and “raises independent interest”, but also its coordinated integration into the algorithmic design of denoising diffusion models that answers the open question of the benefit of parallel sampling theoretically in this context. \n\nAlthough primarily theoretical, our PIADM-SDE algorithm remains practically feasible, and our analysis incorporates practical inference strategies such as exponential integrator, early stopping, etc., which hold promise for real-world applications, a potential recognized by Reviewer sp2J. Further discussion on PIADM-ODE provides possible directions for reducing storage complexity while preserving a sub-linear rate. We will refine our paper to more clearly articulate these contributions.\n\n### Regarding comparison with [3]\nWe appreciate the reviewer’s question regarding our work’s comparison with [3]. As aforementioned, our exploration of PIADM-ODE, building upon our PIADM-SDE findings, aims to investigate the potential reduction in storage requirements by replacing the SDE with the probability flow ODE formulation, which has been demonstrated as “provably faster” in terms of $d$ in [3].\n\nAs discussed in Section 3.2.1 (Algorithm) and Section 3.2.4 (Proof Sketch), our analysis addresses not only the absence of a data processing inequality for 2-Wasserstein distance as in [3], but also the $O(\\sqrt{d})$ time complexity introduced by a naive implementation of the underdamped Langevin corrector, which would compromise our sublinear time efficiency. Leveraging techniques, such as the generalized Girsanov’s Theorem initially developed for our PIADM-SDE, we successfully mitigate these challenges. This allows us to show that the improvement of time efficiency from $O(d)$ to $O(\\sqrt{d})$ with probability flow ODE in [3] can be adapted to reduce storage complexity from $O(d^2)$ to $O(d^{3/2})$ in our parallel algorithms, as detailed in Theorem 3.5. \n\nWe believe our results not only echo those found in [3] but also introduce significant additional technical contributions. We will refine our paper to further elaborate on these points in the revised version.\n\n---\n\nFinally, we sincerely thank the reviewer for their insightful comments. Should our clarifications and modifications address their concerns, we would be grateful if the reviewer could possibly consider reevaluation of our work based on our theoretical contributions and we are more than happy to answer any follow-up questions.\n\n### References\n[1] Anari, N., Chewi, S., & Vuong, T. D. (2024). Fast parallel sampling under isoperimetry. arXiv preprint arXiv:2401.09016.\n\n[2] Benton, J., Bortoli, V. D., Doucet, A., & Deligiannidis, G. (2024). Nearly d-linear convergence bounds for diffusion models via stochastic localization.\n\n[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.\n\n[4] Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. Advances in Neural Information Processing Systems, 36, 2024.\n\n[5] Chen, H., Lee, H., & Lu, J. (2023, July). Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning (pp. 4735-4763). PMLR.\n\n[6] Lee, H., Lu, J., & Tan, Y. (2022). Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 35, 22870-22882.\n\n[7] Lee, H., Lu, J., & Tan, Y. (2023, February). Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory (pp. 946-985). PMLR.\n\n[8] Yang, Ling, et al. \"Diffusion models: A comprehensive survey of methods and applications. arXiv 2022.\" arXiv preprint arXiv:2209.00796 (2022).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (spotlight)