Faster Diffusion Sampling with Randomized Midpoints: Sequential and Parallel
Authors: Shivam Gupta
Linda Cai
Sitan Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion Sampling
Generative Model
Statistical Theory
Primary area: generative models

Abstract: Sampling algorithms play an important role in controlling the quality and runtime of diffusion model inference. In recent years, a number of works (Chen et al., 2023c;b; Benton et al., 2023; Lee et al., 2022) have analyzed algorithms for diffusion sampling with provable guarantees; these works show that for essentially any data distribution, one can approximately sample in polynomial time given a sufficiently accurate estimate of its score functions at different noise levels. 

In this work, we propose a new scheme inspired by Shen and Lee's randomized midpoint method for log-concave sampling  (Shen & Lee, 2019). We prove that this approach achieves the best known dimension dependence for sampling from arbitrary smooth distributions in total variation distance ($\widetilde O(d^{5/12})$ compared to $\widetilde O(\sqrt{d})$ from prior work). We also show that our algorithm can be parallelized to run in only $\widetilde O(\log^2 d)$ parallel rounds, constituting the first provable guarantees for parallel sampling with diffusion models.
    
As a byproduct of our methods, for the well-studied problem of log-concave sampling in total variation distance, we give an algorithm and simple analysis achieving dimension dependence $\widetilde O(d^{5/12})$ compared to $\widetilde O(\sqrt{d})$ from prior work.

Published as a conference paper at ICLR 2025
FASTER DIFFUSION SAMPLING WITH RANDOMIZED MIDPOINTS: SEQUENTIAL AND PARALLEL
Shivam ,upta', Linda Caf,, Sitan Chen
UT Austin "CC Berkeley Harvard SEAS
thivamuutaa@utekee edu, tc (0brreleeyy..........e.eeerrrlee edu, sitanêseas www.seeaa....arr.....oo harvard.
ABSTRACT
Sampling algorithms play an important role in controlling the quality and runtime of diffusion model inference. In recent years, number of works (Chen al., 2023c;b; Benton et al., 2023; Lee et al., 2022) have analyzed algorithms for dif- fusion sampling with provable guarantees; these works show that for essentially any data distribution, one can approximately sample in polynomial time given a sufficiently accurate estimate of its score functions at different noise levels.
In this work, we propose new scheme inspired by Shen and Lee's 's randomized midpoint method for log-concave sampling (Shen & Lee, 2019). We prove that this approach achieves the best known dimension dependence for sampling from arbitrary smooth distributions in total variation distance (015/12) compared to O(Vd) from prior work). We also show that our algorithm can be parallelized to run in only Ö(log" d) parallel rounds, constituting the first provable guarantees for parallel sampling with diffusion models.
As a byproduct of our methods, for the well-studied problem of log- log-concave sam- pling in total variation distance, give an algorithm and simple analysis achiev- ing dimension dependence (dd/112 compared to O(vd) from prior work.
INTRODUCTION
Diffusion models (Sohl-Dicksted) al., 2015; Song & Ermon, 2019; Ho et al., 2020; Dhariwal & Nichol, 2021; Song et al., 2021a;b; Vahdat et al., 2021) have emerged as the de facto approach to generative modeling across range of data modalities like images (Betker al., 2023; Esser et al., 2024), audio (Kong et al., 2020), video (Brooks et al., 2024), and molecules (Wu et al., 2024). In recent years a slew of theoretical works have established surprisingly general convergence guarantees for this method (Chen et al., 2023c; Lee et al., 2023; Chen et al.. 2023a;b;d; Benton et al., 2024: Gupta et al., 2023b; Liet al., 2023; 2024). They show that for essentially any data distribution, assuming one has a sufficiently accurate estimate for its score function, one can approximately sample from in polynomial time.
While these results offer some theoretical justification for the empirical successes of diffusion mod- els, the upper bounds they furnish for the number of iterations needed to generate a single sample are quite loose relative to what is done in practice. The best known provable bounds scale as O ((vd/c), c). where d is the dimension of the space in which the diffusion is taking place (e.g. d = 16384 for Stable Diffusion) (Chen et al., 2023b), and is the target error. Even ignoring the dependence on e and the hidden constant factor, this is at least 2 - 3x larger than the default value of 50 inference steps in Stable Diffusion.
In this work consider needed to sample with diffusion models. Our approach is rooted in the randomized midpoint method, originally introduced by Shen and Lee (Shen & Lee, 2019) in the context of Langevin Monte Carlo for log- log-concave sampling. At a high level, this is a method for numerically solving differential equations where within every discrete window of time, one forms an unbiased estimate for the drift by evaluating it at a random "midpoint" (see Section 2.2 for formal treatment). For sampling from log-co concave densities, the number of iterations needed by their method seales with /33 and this remains the best known bound in the "low-accuracy" regime.
approach for driving down the amount of compute that is provably

---

Published as a conference paper at ICLR 2025
While this method is well-studied in the log-concave setting (He et al., 2020; Yu et al., 2023; Yu & Dalalyana, 2024; Shen & Lee, 2019), its applicability to diffusion models has been unexplored both theoretically and empirically. Our first result uses the randomized midpoint method to obtain im- provement the prior best known bound of O(Vá/c) for sampling arbitrary smooth distributions with diffusion models:
Theorem 1.1 (Informal, see Theorem B.10). Suppose that the data distribution has bounded second moment, its score functions V qu along the forward process are L-Lipschitz, and we are given score estimates which are L-Lipschitz and Ö( T172/ "close to V In for all t. Then there is feffiasion-based sampler using these score estimates (see Algorithm 1) which oulputs sample whose law is c-close in total variation distance to q using ( (25/2///12//// /c) iterations.
Our algorithm is based on the ODE-based predictor-corrector algorithm introduced in (Chen et al., 2023b), but in place of the standard exponential integrator discretization in the predictor step, we employ randomized midpoint discretization. We note that in the domain of log-concave sampling. the result of Shen and Lee only achieves recovery in Wasserstein distance. Prior to our work, was actually open whether one can achieve the same dimension dependence in total variation or KL divergence, for which the best known bound was Ö(vd) (Ma et al., 2021; Zhang et al., 2023; Altschuler & Chewi, 2023). In contrast, our result cireumvents this barrier by carefully trading off time spent in the corrector phase of the algorithm for time spent in the predictor phase. We defer the details of this, as well other important technical hurdles, to Section 3.1.
Next, we turn to a different computational model: instead of quantifying the cost of an algorithm in terms of the total number of iterations, we consider the parallel setting where one has access to multiple processors and wishes to minimize the total number of parallel rounds needed to generate a single sample. This perspective has been explored in a recent empirical work (Shih et al., 2024), but to our knowledge,
provable guarantees were known for parallel sampling with diffusion models (see Section 1 for discussion of concurrent and independent work). Our second result provides the first such guarantee:
Theorem 1.2 (Informal, see Theorem C.13), Under the same assumptions on q as in Theorem I.I, and assuming that we are given score estimates which are Õ vT. ) to V In qe for all t, there is diffusion-baser sampler using these score estimates (see Algorithm 9) which outputs sample whose law is e-close in total variation distance to q using O(L polyloolyylllllo parallel rounds.
This result follows in the wake of several recent theoretical works on parallel sampling of log- concave densities using Langevin Monte Carlo (Anari et al., 2023; 2024; Shen & Lee, 2019). A common thread among these works is the observation that differential equations can be numerically solved via fixed point iteration (see Section 2.3 for details), and we adopt similar perspective in the context of diffusions. To our knowledge this is the first provable guarantee for parallel sampling beyond the log-concave setting.
Finally, we show that, as a byproduct of our methods, we can actually obtain a similar dimension de- pendence of (d5/112 as in Theorem 1.1 for log-concave sampling in TV, superseding the previously best known bound of ((v)) mentioned above.
Theorem 1.3 (Informal, see Theorem D.3). Suppose distribution q is bestrongyyloo--onccae and its score function V In is L-Lipschitz. Then, there is underdamped- Laadd--annn-n-aee based sampler that uses this score (Algorithm 11) and outputs a sample whose law is c-close in total variation to q using LA/3
(1 277mm473 + +))) iterations.
1.1 RELATED WORK
Our discretization scheme is based on the randomized midpoint method of (Shen & Lee, 2019), which has been studied at length in the domain of log-concave sampling (He et al., 2020; Yu et al., 2023; Yu & Dalalyana, 2024).
(-)-) hides polylogarithmic factors in d, I, and R [1=[1]]]
2

---

Published as a conference paper at ICLR 2025
The proof of our parallel sampling result builds on the ideas of (Shen & Lee, 2019; Anari et al., 2023; 2024) on parallelizing the collocation method. These prior results were focused on Langevin Monte Carlo, rather than crfffusion-basee sampling. We review these ideas in Section 2.3.
In (Chen et al., 2023b), the authors proposed the predictor-corrector framework that we also use for analysing convergence guarantee of the probability flow ODE and which achieved iteration com- plexity scaling with d(vd). addition this, there have been many works recent years giving general convergence guarantees for diffusion models (De Bortoli et al.. 2021; Block et al., 2022; De Bortoli, 2022: Lee et al., 2022; Liu et al., 2022; Pidstrigach, 2022; Wibisono Yang. 2022; Chen et al., 2023c;d; Lee et al., 2023; Li et al., 2023; Benton et al., 2023; Chen et al., 2023b; Benton et al., 2024; Chen al., 2023a; Gupta et al., 2023b). Of these, one line of work (Chen et al., 2023e; Lee et al., 2023; Chen et al., 2023a; Benton al., 2024) analyzed DDPM, the stochastic analogue of the probability flow ODE, and showed (d) iteration complexity bounds. Another set of works (Chen et al., 2023b;d; Li et al., 2023; 2024) studied the probability flow ODE, for which our work provides a new discretization scheme for the probability flow ODE, that achieves a state-of-the-a (d5/12) dimension dependence for sampling from diffusion model.
Concurrent work. Here we discuss the independent works of (Chen et al., 2024) and (Kandasamy & Nagaraj. 2024). (Chen et al., 2024) gave an analysis for parallel sampling with diffusion models that also achieves polylog(d) number of parallel rounds like in the present work. (Kandasamy & Nagaraj, 2024) showed an improved dimension dependence of (dd5/12) for log-concave sampling in total variation, similar to our analogous result, but via a different proof techmique. In addition to this, they show a similar result when the distribution only satisfies log-Sobolev inequality. They also show empirical results for diffusion models, showing that an algorithm inspired by the randomized midpoint method outperforms ODE based methods with similar compute. While their work builds on the randomized midpoint method, they do not theoretically analyze the diffusion setting and do not study parallel sampling.
2
PRELIMINARIES
2.1 PROBABILITY FLOW ODE
In this section we review basics about deterministic liffusion-basee samplers; we refer the reader to (Chen et al., 2023b) for more thorough exposition.
Let q" denote the data distribution over Rd. We consider the standard Ornstein-Ullenbed (OU) forward process, i.e. the "VP SDE," given by
dxpt = - dt V2dB.
25 q",
(1)
where (Bt)220 denotes standard Brownian motion in Rd This process converges exponentially quickly to its stationary distribution, the Gaussian distribution (N(ittiooonnOO (N Id).
Suppose the OU process is run until terminal time T > O, and for any E [0, let of A law( ) i.c. the law of the forward process at time We will consider the reverse process given by the probability flow ODE
dx. (x+ + V Vln di.
This is a time-reversal of the forward process, so that xip ~ yr, then law(x.) -r- In practice, one initializes at xo (O, and instead of using the exact score function V In -r one uses estimates -r x V In gr which are learned from data. Additionally, the ODE is solved numerically using any of a number of discretization schemes. The theoretical literature on diffusion models has focused primarily on exponential integration, which we review next before turning to the discretization scheme, the randomized midpoint method used in the present work.
2.2 DISCRETIZATION SCHEMES
Suppose we wish to discretize the following semilinear ODE:
dx. (Ite + f((x))) dt.
(3)
3

---

Published as a conference paper at ICLR 2025
For our application will eventually take fi _ S1 but use ft in this section to condense notation.
Suppose want discretize Equation (3) over time window [to, to + h).. The starting point is the integral formulation for this ODE:
+lo = 2to + foth
goth-*felxd) dt.
(4)
Under the standard exponential integrator discretization, would approximate the integrand by eooth- (Ita) and obtain the approximation
21oth R eta + 10
(5)
The drawback of this discretization is that it uses an inherently biased estimate for the integral Eq. equation 4. The key insight of (Shen & Lee, 2019) was to replace this with the following unbiased estimate
10th forr (ooth--f((e)) dt = hell-a)h ftotah (Ztg+ah))
where is uniformly random sample from [0, 1]. While this alone does not suffice the estimate depends on taoh naturally we could iterate the above procedure again to obtain an approxima- tion +oh- It turns though that even if we simply approximate Itotoh using exponential integrator discretization, we can obtain nontrivial improvements in discretization error (e.g. our Theorem 1.1). In this case, the above sequence of approximations takes the following form:
tth se e Tto + (e ft (1tt))
Ftoth ra
Fto + he(l-a)h ftotah (Pto+ah).
(8)
Note that a similar idea can be used to discretize stochastic differential equations, but in this work only use it discretize the probability flow ODE.
Peffitto-CCorecttrr For important technical reasons, in our analysis we actually consider a slightly different algorithm than simply running the probability flow ODE with approximate score, Gaussian initialization, and randomized midpoint discretization. Specifically, we interleave the ODE with corrector steps that periodically inject noise into the sampling trajectory. We refer to the phases in which we are running the probability flow ODE predictor steps.
The corrector step will be given by running underdamped Langevin dynamics. As our analysis of this will borrow black-box from bounds proven in (Chen et al., 2023b), we refer to Section B.2 for details.
2.3 PARALLEL SAMPLING
The scheme outlined in the previous section is simple special of the collocation method. In the context of the semilinar ODE from Eq. equation 3, the idea behind the collocation method is to solve the integral formulation of the ODE in Eq. equation via fixed point iteration. For our parallel sampling guarantees, instead of choosing a single randomized midpoint a, we break up the window [to, to + h) into R sub-windows select randomized midpoints a1...., OR for these sub- window,, and approximate the trajectory of the ODE any time to + i5, where & h/ R. by
thg RS aadh Fip
Σ caub- max(c .1) (Ftotaah))
(9)
One can show that as R 0o, this approximation tends equality. For sufficiently large R, Eq. equation naturally suggests fixed point iteration that can be used to approximate each to h. i.e. we can maintain a sequence of estimates =(k)
tottag defined by the iteration
(k) to + ( + Σ
aaah- (1-1 nh- 11) -iotagh (iotas((((kt))) (=to (10)
4

---

Published as a conference paper at ICLR 2025
for k ranging from 1 up to some sufficiently large K. Finally, analogously to Eq. equation 8, estimate #1 via
+toth 2 cla(()
8
Σ i=1
(1-aa)h Sto
(F( tatagh)
(11)
The key observation, made (Shen Lee, 2019) and also related works of (Anari et al., 2024; Shih et al., 2024; Anari et al., 2023), is that for any fixed round k, all of the iterations Eq. equation 10 for different choices of = 1,..., R can be computed in parallel. With R parallel processors, one can thus compute the estimate for tto in K parallel rounds, with O(KR) total work.
2.4 ASSUMPTIONS
Throughout the paper, for our diffusion results, will make the following standard assumptions on the data distribution and score estimates.
Assumption 2.1 (Bounded Second Moment).
m
E x-
[||=|1]] < co.
Assumption 2.2 (Lipschitz Score). For all I, the score V qt is L-Lipschitz.
Assumption 2.3 (Lipschitz Score estimates). For all I. for which we need to estimate the score function in algorithms, the score estimate 8t is T-lipschitz.
Assumption 2.4 (Score Estimation Error). For all & ffror which need estimate the score function in our algorithms.
E ar
[z,,
V In ( Sec
TECHNICAL OVERVIEW
Here we provide an overview of our sequential and parallel algorithms, along with the analysis of iteration complexity bounds. We begin with description of the sequential algorithm.
3.1 SEQUENTIAL ALGORITHM
Following the framework of (Chen et al., 2023b), our algorithm consists of "predictor" steps in- terspersed with "corrector" steps, with the time spent on each carefully tuned to obtain our final (d5/122 dimension dependence. We first describe our predictor step this is the piece of our algorithm that makes use of the Shen and Lee's 'e's randomized midpoint method (Shen Lee, 2019).
Algorithm PREDICTORSTEP (SEQUENTIAL)
Input parameters:
Starting sample To. Starting time lo. Number of steps N, Step sizes hnejo.....N Score estimates 8,
1. For n O
(a) Let ln = to En-1 h;
(b) Randomly sample uniformly from [].
(c) Let F = cahn + (eat. - 1) St, (ss
(d) Let Zn+1 hin + hn ell-a)lin ali (2 kn+g)
2. Let EN to NN11 hi
3. Return CN IN-
5

---

Published as a conference paper at ICLR 2025
The main difference between the above and the predictor step of (Chen et al., 2023b) are steps 1(b) [(d). 1(b) and (cc) together compute a randomized midpoint, and ((d) uses this midpoint to obtain approximate solution to the integral of the ODE. We describe these steps in more detail in Section 3.3.
Next, describe the "corrector" step, introduced in (Chen al., 2023b). First, recall the under- damped Langevin ODE:
dē 00 dt
(12)
dē.
76;) di /27 c.B.
Here is our 2 accurate score estimate for fixed time (say t). Then, the corrector step is described below.
Algorithm CORRECTORSTEP (SEQUENTIAL)
Input parameters:
Starting sample To. Total time Toorr Step size heom Score estimate
1. Run underdamped Langevin Monte Carlo in equation 12 for total time Toorr using step size heore and let the result be Ny.
2. Return Zy-
Finally, Algorithm 3 below puts the predictor and corrector steps together to give our final sequential algorithm.
Algorithm SEQUENTIALALLGORTTHH
Input parameters:
Start time T, End time 8, Corrector steps time Teorr S I/VI, Number of predictor- corrector steps No. Predictor step size hpred, Corrector step size heorr, Score estimates &r
1. Draw D0 (o).
2. For n = ......No - 1:
(a) Starting from En. run Algorithm with starting time T - nt L using step sizes /pprdd for all steps, with "hepred so that the total time is 1 L. Let the result be F
(b) Starting from E L run Algorithm for total time Toor with step size heore and score estimate 1/-(++1)L obtain Fa+1-
3. Starting from NN., run Algorithm with starting time T - -No/L using step sizes hpped/22 hpred /4, /ppea/8...... to obtain Z Noo+1-
4. Starting from No run Algorithm for total time Teorr with step size heors and score estimate 85 to obtain NN,+1-
5. Return No
For the final setting of parameters in Algorithm 3, see Theorem 10 Now, we describe the analysis of the above algorithm detail.
3.2
COORECTR FRAMEWORK
The general framework of our algorithm closely follows that of (Chen et al., 2023b), which proposed to run the (discretized) reverse ODE but interspersed with "corrеccor" steps given by rumning un- derdamped Langevin dynamics. The idea is that the "predictor" steps where the discretized reverse

---

Published as a conference paper at ICLR 2025
ODE is being run keep the sampler close to the true reverse process in Wasserstein distance, but they cannot be run for too long before potentially incurring exponential blowups. The main purpose of the corrector steps is then to inject stochasticity into the trajectory of the sampler in order to convert closeness in Wasserstein to closeness in KL, divergence. This effectively allows one to "restart the coupling" used to control the predictor steps. For technical reasons that are inherited from (Chen et al., 2023b), for most of the reverse process the predictor steps (Step 2(a)) are run with fixed step size, but the end of the reverse process (Step 3), they are run with exponentially decaying step sizes.
We follow the same framework, and the core of our result lies in refining the algorithm and analysis for the predictor steps by using the randomized midpoint method. Below, we highlight our key technical steps.
3.3 PREDICTOR STEP- IMPROVED DISCRETIZATION ERROR WITH RANDOMIZED MIDPOINTS
Here, we explain the main idea behind why randomized midpoint allows us to achieve improved dimension dependence. We first focus on the analysis of the predictor (A (Algorithm 1) and restrict our attention to running the reverse process for a small amount of time h < I/L.
We begin by recalling the dimension dependence achieved by the standard exponential integrator scheme. One can show (see e.g. Lemma 4 in (Chen et al., 2023b)) that if the true reverse process and the discretized reverse process are both run for small time h starting from the same initialization, the two processes drift by a distance of )2 By iterating this coupling ((1hh) times, we conclude that in an O(1) window of time, the processes drift by a distance of O(d)/2). To ensure this is not too large, one would take the step size h to be O(1/Vd), thus obtaining an iteration complexity of ((Z)) as in (Chen et al., 2023b).
The starting point in the analysis of randomized midpoint is to instead track the squared displace- ment between the two processes instead. Given two neighboring time steps t - h and t in the algorithm, let Vg denote the true reverse process time I, and let T. denote the algorithm (in the notation of Algorithm 2, this is T. for some nt, but we use / in the discussion here to make the comparison to the true reverse process clearer). Note that क. depends on the choice of randomized midpoint a (see Step 1(b)). One can bound the squared displacement E ||æ, - F2ll as follows. Let Ve be the result of running the reverse process for time h starting from Tz-h. Then by writing I2- - Et
time t
as (xx - y.) - (ae- - yt) and applying Young's 's inequality, we obtain
Lh 2, hi 1| -.11 < (1 (hh
2
E
E | Se
11
E JET. 1ell + E E 2o11.. (13)
2
Lh 2.
For the first term, because Iy and y are the result of rumning the same ODE on initializations zt-h and Zi-h- the first term is close to E 1|= - 12 provided h < 1/L. The upshot is that the squared displacement at time is at most the squared displacement at time l. - h plus the remaining two terms on the right of Equation (13).
The main part of the proof lies in bounding these two terms, which can be thought of as "bias" and "variance" terms respectively. The variance term can be shown to scale with the square of the aforementioned (2) displacement bound that arises in the exponential integrator analysis,
giving Oddh)): Lemma 3.1 (Informal, see Lemma B.4 for formal statement). If S 7, and T - 2 (T - ---h)))/22 12) then
E 110-h 12. well S I2 [LV (LV T- 1222 + 12222 Bi-k
1
E E 0 12
E
Note that in this bound, in addition to the Oddh*)) term and a term for the score estimation error, there is an additional term which depends the squared displacement from the previous time step. Because the prefactor 1.21,2 is sufficiently small, this will ultimately be negligible.
The upshot of the above Lemma is that if the bias term is of lower order, then this means that the squared displacement essentially increases by O(dh*) with every time step of length h. Over O(1/h) such steps, the total squared displacement is Oddh),, so if we take the step size h to be 0(1/d1/3 this suggests improved iteration complexity of /d.3
7

---

namely Published as a conference paper at ICLR 2025
Arguing that the bias term LA Ez.-, |E,, 3t 1l is dominated by the variance term is where it is crucial that we use randomized midpoint instead of exponential integrator. But recall that the randomized midpoint method was engineered that it would give unbiased estimate for the true solution to the reverse ODE if the estimate of the trajectory at the randomized midpoint were exact. In reality we only have approximation to the latter, but this is indeed of lower order (see Lemma B.3). One technical complication that arises here is that the relevant quantity to bound is the distance between the true process at the randomized midpoint versus the algorithm, when both are initialized at an intermediate point in the algorithm's 's trajectory. Bounding such quantities in expectation over the randomness of the algorithm's trajectory can be difficult, but our proof identifies way of "offloading" some of this difficulty by absorbing some excess terms into term of the form || - æ |22 i.e. the squared displacement from the previous time step. Concretely, we obtain the following bound on the bias term:
2
show, the error incurred by
Lemma 3.2 (Informal, see Lemma B.2 for formal statement).
|| | - -oll S d' (LV T TT + +122 + +L .I 2111
E Be-h
3.4 SHORTENING THE CORRECTOR STEPS
While we have outlined how to improve the predictor step in the framewok of (Chen et al., 2023b), is quite unclear whether the same can be achieved for the corrector step. Whereas the the former is geared towards closeness in Wasserstein distance, the latter is geared towards closeness in KL divergence, and it is a well-known open question in the log-concave sampling literature to obtain analogous discretization bounds in KL for the randomized midpoint method (Chewi, 2023).
We will sidestep this issue and argue that even using exponential integrator discretization of the underdamped Langevin dynamics will suffice for purposes, by simply shortening the amount of time for which each corrector step is run.
First, let us briefly recall what was shown in (Chen et al., 2023b) for the corrector step. If one runs underdamped Langevin dynamics with stationary distribution for time T and exponential integrator discretization with step size h starting from two distributions p and q. then the resulting distributions p' and g satisfy
We(p.g)
TV 4) S L1/173/2 + 3/////711 1/2
(14)
where L is the Lipschitzness of V In (see Theorem B.6). At first glance this appears insufficient for our purposes: because of the / term coming from the discretization error, we would need to take step size h 1/vd, which would suggest that the number of iterations must scale with Vd.
To improve the dimension dependence for our overall predictor-correcto algorithm, we observe that if we take T itself to be smaller, then we can take h to be larger while keeping the discretization error in Equation (14) sufficiently small. Of course, this comes at a cost, as T also appears in the term Wa(p-g) in Equation (14). But in our overall proof, the (2(P, g) term is bounded by the predictor T4/475/2
analysis. There, we had quite bit of slack: even with step size as large as 1/d'/3, we could achieve small Wasserstein error. By balancing appropriately, we get our improved dimension dependence.
3.5 PARALLEL ALGORITHM
Now, we summarize the main proof ideas for our parallel sampling result. In Section 2.3, we de- scribed how approximately solve the reverse ODE over time h by running K rounds of the iter- ation in Equation (10). In our final algorithm, we will take h to be ddimension-independeer
h = ((/VL), so that the main part of the proof is to bound the discretization error incurred over
each of these time windows of length h. As in the sequential analysis, we will interleave these "predictor" steps with corrector steps given by (parallelized) underdamped Langevin dynamics.
We begin by describing the parallel predictor step. Suppose we have produced an estimate for the reverse process at to and now wish to solve the ODE from time to to lq + f. We initialize at
(1) (otta, via exponential integrator steps starting from the beginning of the window see Line
8

---

Published as a conference paper at ICLR 2025
I(c) in Algorithm 9 (this can be thought of as the analogue of Equation (7) used in the sequential algorithm). The key difference relative to the sequential algorithm is that here, because the length of the window is dimension-free the discretization error incurred by this initialization is too large and must be refined using the fixed point iteration in Equation (10). The main step is then to show that with each iteration of Equation (10), the distance the true reverse process contracts:
Lemma 3.3 (Informal, see Lemma C.2 for formal statement), Suppose S L/L. If yi denotes the solution of the true ODE starting at Ttt and running unfil time to + agh, then for all l E (1,-.. K} and E ,1,...R R).
R Σ
I lotach Mooach 1|
(8h2
(00) botash
- Notayhh l')
Big O1
R j=1
L'dh²
+ 1/2
60c
+
(LV
+ T.2. E Peg
1
Fig
112
R²
T
+
(15)
where Tto is the iterate of the algorithm from the previous time window, and Eto is the corresponding iterate in the true ODE.
In particular, because h is at most a small multiple of 1/L, the prefactor (8h²L2) is exponentially decaying in k so that the error incurred by the estimate a (k) to is contracting with each fixed point iteration. Because the initialization is distance poly(d) from the true process, Olog d) rounds of contraction thus suffice, which translates to (log d) parallel rounds for the sampler. The rest of the analysis of the predictor step is quite similar the analogous proofs for the sequential algorithm (i.e. Lemma C.4 and Lemma C.5 give the corresponding bias and variance bounds).
One shortcoming of the predictor analysis is that the contraction achieved by fixed point iteration ultimately bottoms out at error which scales with d/ Rt² (see the second term in Equation (15)). In order for the discretization error be sufficiently small, we thus have take and thus the total work of the algorithm, scale with O d). So this case do not improve over the dimension dependence of (Chen et al., 2023b), and instead the improvement is in obtaining parallel algorithm.
For the corrector analysis, we mostly draw upon the recent work of (Anari et al., 2024) which analyzed parallel implementation of the underdamped Langevin dynamics. While their guarantee focuses on sampling from log-concave distributions, implicit their analysis a bound for general smooth distributions on how much the law of the algorithm and the law of the true process drift apart bounded time window (see Lemma C.8). This bound suffices for our analysis of the corrector step, and we can conclude the following:
Theorem 3.4 (Informal, see Theorem C.12). Let 2 be an adjustable parameter. Let p' denote the law of the output of running the parallel corrector (see Algorithm 8) for total time 1/ L and step size h, using epproppre estimate for V In and starting from sample from another distribution p.
TV S VKL(V', S
W2(p.4)
VL
Bvd
Furthermore, this algorithm uses @(BVa/c) score evaluations over (llg((2d/22) parallel rounds.
Overall, parallel algorithm somewhat different from the parallel sampler developed in the empiri- cal work of (Shih et al., 2024), even apart from the fact that we use randomized midpoint discretiza- tion and corrector steps. The reason is that our algorithm applies collocation fixed windows of time, whereas the algorithm of (Shih et al., 2024) utilizes sliding window approach that proac- tively shifts the window forward as soon as the iterates the start of the previous window begin to converge. We leave rigorously analyzing the benefits of this approach as an interesting future direction.
3.6 LOG-CONCAVE SAMPLING IN TOTAL VARIATION
Finally, we briefly summarize the simple proof for our result log-concave sampling in TV, which achieves the best known dimension dependence of (d5/12). Our main observation is that Shen

---

Published as a conference paper at ICLR 2025
and Lee's randomized midpoint method (Shen & Lee, 2019) applied to the underdamped Langevin process gives a Wasserstein guarantee for log-concave sampling, while the corrector step of (Chen et al., 2023b) can convert Wasserstein guarantee to closeness in TV. Thus, randomized midpoint method, followed by the corrector step to achieve closeness in TV. Carefully tuning the amount of time spend and step sizes for each phase of this algorithm yields our improved dimension dependence see Appendix D for the full proof.
simply run the
4 DISCUSSION AND FUTURE WORK
In this work, we showed that it is possible to leverage Shen and Lee's randomized midpoint method (Shen & Lee. 2019) to achieve the best known dimension dependence for sampling from arbitrary smooth distributions in TV using diffusion. We also showed how to parallelize our algo- rithm, and showed that Ölog" d) parallel rounds suffice for sampling. These constitute the first provable guarantees for parallel sampling with diffusion models. Finally, we showed that our tech- niques can be used to obtain an improved dimension dependence for log-concave sampling in TV.
We note that relative to (Chen et al., 2023b), our result requires a slightly stronger guarantee on the score estimation error, by a dd/12 factor; we believe this is an artifact of our analysis, and it would be interesting to remove this dependence in future work. Importantly however, it was not known how to achieve an improvement over the O(vd) dependence shown in that paper even in case that the scores are known exactly prior the present work. Moreover, another line of work (Li et al., 2023; 2024; Dou et al., 2024) analyzing diffusion sampling makes the stronger assumption that the score estimation error Ō va an assumption stronger than ours by 5/12 factor, this does not detract from the importance of these works, and we feel the same is true in case.
We also note that our diffusion results require smoothness assumptions we assume that the true score, well our score estimates are L-Lipschitz. Although this assumption is standard in the literature, recent work (Chen et al., 2023a; Benton et al., 2024) has analyzed DDPM in the absence of these assumptions, culminating in (dd) dependence for sampling using a discretization of the reverse SDE. However, unlike in the smooth case, it is not known whether even a sublinear in d dependence is possible without smoothness assumptions via any algorithm. We leave this as an interesting open question for future work.
REFERENCES
Jason M Altschuler and Sinho Chewi. Faster high-accuracy log-concave sampling via algorithmic warm starts. In 2023 IEEE 64th Annual Symposium on Foundations of Computer Science (FOCS), pp. 2169-2176. IEEE, 2023.
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, pp. 103-116, 2023.
Nima Anari, Sinho Chewi, and Thuy-Duong Vuong. Fast parallel sampling under isoperimetry. CoRR, 0bbs/2401.09010 2024. doi: 10448500/ARXIV.2400..0001 URL https: htt///o..o//// 10. 884488550/40 arriv. 2401. 09016.
Joe Benton, George Deligiannidis, and Arnaud Doucet. Error bounds for flow matching methods. arXiv preprint 2.X.::2305.16660 2023.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear con- vergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learning Representations, 2024.
James Betker, Gabriel Goh, Li Jing, Tim Brooks, Jianfeng Wang, Linjie Li, Long Ouyang, Juntang Zhuang, Joyce Lee, Yufei Guo, al. Improving image generation with better captions. Computer Science. https://cdn. openai. com/papers/dallle-.. pdf, 2(3):8, 2023.
Adam Block, Youssef Mrouch, and Alexander Rakhlin. Generative modeling with denoising auto- encoders and Langevin sampling. arXiv preprint 2002.00107, 2022.
10

---

Published as a conference paper at ICLR 2025
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL https: httpp::///paaa.../////// wwt. coa////eerer..ee...... -- 1 d- simulators.
widdeo-generaticn-mme
Haoxuan Chen, Yinuo Ren, Lexing Ying. and Grant M. Rotskoff. Accelerating diffusion models with parallel sampling: Inference sub-linear time complexity, 2024.
Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learming, pp. 4735 4763. PMLR, 2023a.
Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probabil- ity flow ODE provably fast. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b. URL https: httpapp::///ppeeeeeeee openreview net/ http://wwwneee/uuutmieeem ?immmiiddii LD 6MF eWSAd.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is easy as learning the score: theory for diffusion models with minimal data assumptions. International Conference on Learning Representations, 2023c.
Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degradaee beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning, pp. 4462-4484. PMLR, 2023d.
Siinho Chewi. Log-concave sampling. Book draft available at https://chewisinho. github. io, 2023.
Valentin De Bortoli. Convergence denoising diffusion models under the manifold hypothesis.
Transactions on Machine Learning Research, 2022.
Valentin De Bortoli, James Thornoon, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Pro- Curran Associates, Inc., 2021.
cessing Systems, volume 34, pp. 17695-17709.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 8780-8794. Curran Associates, Inc., 2021.
Zehao Dou, Minshuo Chen, Mengdi Wang, and Zhuoran Yang. Theory of consistency diffusion models: Distribution estimation meets fast sampling. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Pro- ceedings of Machine Learning Research, pp. 11592-11612. PMLR, 21-27 Jul 2024. URL https::/proceet wttpooeeeiinssm.......o.... mlr htw.aainss.....esssss..ee v235/22355//// wotpess//23//oo24.aa..oo 2...3//..o2aaa.m.m.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.03226, 2024.
Shivam Gupta, Jasper C.H. Lee, and Eric Price. High-dimensional location estimation via norm concentration for subgamma vectors. In Proceedings of the 40th International Conference on Machine Learning, ICML'2) 23. JMLR.org. 2023a.
Shivam Gupta, Aditya Parulekar, Eric Price, and Zhiyang Xun. Sample-efficient training for diffu- sion. arXiv preprint arXiv:2311.13745. 2023b.
Ye He, Krishnakumar Balasubramanian, and Murat A Erdogdu. On the ergodicity, bias and asymp- totic normality of randomized midpoint sampling method. Advances in Neural Information Pro- cessing Systems, 33:7366-7376, 2020.
11

---

Published as a conference paper at ICLR 2025
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. In Advances in Neural Information Processing Systems (NeurIPS), volume 30, 2017. URL https: arxiv. org/ hrg//bbs// 1706. 08500.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilisticc models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Saravanan Kandasamy and Dheeraj Nagaraj. The poisson midpoint method for langevin dynamies: Provably efficient discretization for diffusion models, 2024.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint WxXi::200.009766 2020.
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.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference Algorithmic Learning Theory, pp- 946-985, PMLR, 2023.
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.
Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating conver- gence of score-based diffusion models, provably. arXiv preprint arXiv:2403.03852, 2024.
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208. /4699, 2022.
Yi-An Ma. Niladri S Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L Bartlett, and Michael I Jordan. there an analog of nesterov acceleration for gradient-based meme? 2021.
Jakiw Pidstrigach. Score-based generative models detect manifolds. In S. Koyejo. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 35852-35865 Curran Associates, Inc., 2022.
Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence Fox, and Roman Garnett (eds.), Advances in Neural Infor- Neural Information Process- ing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 2098 2109, 2019. URL htttps::/proceeediii neuripp wttpeuuppps.ppppeeeeeeeeeeeeeeee htc/pappr//219/baaeee (w1.66613661112255111eelclCc wtt.://bb.....ccccbbb..tttteett. html.
d' d'Alché-Buc, Emily mation Processing Systems 32: Annual Conference
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.), Proceed- ings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp.
Lille, France, 2015. PMLR.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- A'A,hé,BBc,, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 1415-1428. Curran Associates, Inc., 2021a
12

---

Published as a conference paper at ICLR 2025
Yang Song, Jascha Sohl-Dicksteii Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021b.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 11287-1 1302. Curran Associates, Inc., 2021.
Andre Wibisono and Kaylee Y. Yang. Convergence in KL divergence of the inexact Langevin algorithm with application to score-based generative models. arXiv preprint 2211.01512, 2022.
Kevin E Wu, Kevin K Yang, Rianne van den Berg. Sarah Alamdari, James Y Zou, Alex X Lu, and Ava Amini. Protein structure generation via folding diffusion. Nature Communicattons, 15(1): 1059, 2024.
Lu Yu and Amak Dalalyana. Parallelized midpoint randomization for langevin monte carlo. arXiv preprint arXiv:2402. 14434, 2024.
Lu Yu, Avetik Karagulyan, and Arnak Dalalyan. Langevin monte carlo for strongly log-concave distributions: Randomized midpoint revisited. arXiv preprint arXiv:2306.08494 2023.
Shunshi Zhang, Sinho Chewi, Mufan Li, Krishna Balasubramanian, and Murat A Erdogdu. Im- proved discretization analysis for underdamped langevin monte carlo. In The Thirty Sixth Annual Conference on Learning Theory, pp. 36-71. PMLR, 2023.
Roadmap. In Section A, we provide empirical evaluation of the randomized midpoint algorithm. In Section B, we give the proof of Theorem 1.1, our main result sequential sampling with diffu- sions. In Section C, we give the proof of Theorem 1.2, our main result on parallel sampling with diffusions. In Section D. we give the proof of Theorem 1.3 on log-concave sampling.
As a notational remark, in the proofs follow we will sometimes use the notation KL(x || y), (2, y). and TV(z, y) for random variables and to denote the distance between their associated probability distributions. Also, throughout the Appendix, we use t to denote time in the forward process.
A EXPERIMENTS
Experimental setup. For all of our experiments, we use one NVIDIA A100 GPU. We evalu- ate our (sequential) randomized midpoint scheduler (predictor step only), deterministic midpoint scheduler (where the midpoint the average of the start and end times of a step), the default DDIM scheduler, and the default DDPM scheduler on the following datasets: CIFAR-10 (generated image dimension: 32 x 32), and CelebAHQ (generated image dimension: 256 x 256). To obtain score estimations for both of the schedulers, 2020) for the two d and generate 50k sample images for each dataset.
public pretrained DDPM models release by (Ho al.,
Evaluation. The performance of our scheduler and the default DDIM scheduler is evaluated by comparing the Fréchet Inception Distance (FID) scores (Heusel et al., 2017), which measure the quality of generated samples relative to the target distribution. Specifically. we use pytorch-fid?. The most expensive part of diffusion scheduling algorithm is the number of evaluations of the score estimation function, as it involves calling the pre-trained neural network. To obtain a fair comparison, we maintain the same number of function calls (NFE) between the schedulers. Since randomized and deterministic midpoint takes two NFE per step, while DDIM and DDPM takes one NFE per step, this means that randomized and deterministic midpoint schedulers takes twice the step size of DDIM and DDPM in our experiments. Our results can be found in Figure 1.
These pretrained models can found respectively at https://heibxx.n uni- http://///g///oooooooooooo/o///a hhaps//mmpi.agg//mocc/yyhhhhhhdd
13

---

Published as a conference paper at ICLR 2025
Discussion. In our experiments, randomized midpoint consistently outperforms DDIM and DDPM on FID scores, which matches our theoretical analysis. The randomized midpoint appears to perform better than the deterministic midpoint the CIFAR-10 dataset, while showing compa- rable performance on the HH dataset. We suspect that this relative performance difference is related to the empirical bias of the deterministic midpoint estimator.
FD Scores vs MF for 01M0-10
70 Scones us WFF for 25e
Miljorr
42-
22-
Namber actio Drak 1N77
Wumber of Funccten Eallartoo Der1
Figure 1: Number of Function Calls vs. FID score for CIFAR-10 and CelebAHQ Datasets (both axes are log scale).
B SEQUENTIAL ALGORITHM
In this section, we describe our sequential celtoomieeddiippinnttt--e based algorithm in detail. Follow- ing the framework of (Chen et al., 2023b), we begin by describing the predictor Step and show in Lemma B.5 that in (dd//) steps (ignoring other dependencies), when for time t at most O (.) starting from fn, it produces a sample that is close to the true distribution at time f, - tt Then, we show that the corrector step be used to convert our W2 error error in TV distance by running the underdamped Langevin Monte Carlo algorithm, as described in (Chen et al., 2023b). We show in B.8 that if we run our predictor and corrector steps in succession for careful choice of times, we obtain a sample that is close to the true distribution in TV using just ((25/12) steps, but covering a time O (+) Finally, in Theorem 10 we iterate this bound Ölog' d) times to obtain our final iteration complexity of ((d/12).
B. PREDICTOR STEP
To show the 0(d dependence on dimension for the predictor step, we will, roughly speaking, show that its bias after one step is bounded by z Oz (dh") in Lemma B.3, and that the variance is bounded by OL (dh+) in Lemma B.4. Then, iterating these bounds x to times as shown in Lemma B.5 will give error Or. (dht + dhi²) in squared Wasserstein Distance.
Lemma B.1 (Naive ODE Coupling). Consider two variables Io, zn starting at time to, and consider the result of running the true ODE for time h, and let the results be xI, 24, For L 2 1, h < 1/ L. we have
| 1/11 < exp(O(Lh))]loo =6612
Proof. Recall that the true ODE is given by
dr, = (xx + V In (
So,
0:112 - x/612 = 2(xx( 2 0, 0,2²)
2(1xe 2/, It- L' + v (ln V (lngr-((z1))
S L| 1/112
14

---

Published as a conference paper at ICLR 2025
Algorithm 4 PREDICTORSTED (SEQUENTIAL)
Input parameters:
Starting sample To. Starting time to. Number of steps N, Step sizes Aneio.....N Score estimates $6
1. For n .,...,.--- 1:
(a) Let in = 1o Ei=0" hi
(b) Randomly sample uniformly from [0, 1].
(c) Let Fn = Tahnnn + (eatn - 1) St., (F ds
(d) Let ēn+1 = = + bn el1-a)hn athn (2 (x+tz)
2. Le = lo -i h4
3. Return TN IN-
So,
1 --1111 S exp (O(Lh)) ||:0 1'611
Lemma B.2. Suppose L 1. In Algorithm for all (O...., N 1). let (t be the solution
of the true ODE starting at Fn at time Ins running until time ln - L. If hn S I and in - hn 2 In 12/2, we have
E 1 a)hnn St,, -ahn (Zn+z) - Wne el (l-a)hha V Qt., ahn (ath))))
S h + L Lhg LV bn
LL 12, 2t, 11?
where E refers to the expectation over the initial choice To ~ Qto
Proof. For the proof, will let h hn. It suffices to show that
1 1 L2 E tn
E 11th--ah (fx+3) V In qt., (ah)))l I 2 + L'dh"
(LV
- x2, 11? (16)
We can separate the above quantity into three parts: the expected difference in score estimation
function s(-) evaluated at Rn+1/2 and -ahh the expected difference in the true score function g())
evaluated ar (oh) and aah and the expected difference between s(-) and q(-) both evaluated It, -ah. By our Lipschitz assumption on both the score and score estimation function (Theorem 2.2 and Theorem 2.3), the first two terms can be bounded by L times the difference in estimated and true sample By Theorem 2.4, the last term can be bounded by 2 Formally, we have
St., -nh an (fn+z)
In -an (x, (ah))
E [t, on ( (Tn+t) St,, on (Ft, - )I E E In lt,, - ob (x (ah)) V In Ht,, -ah (I (vtt, -ah)) 11
+ E | - (2 (ah) - V In qt, aah (Zr, )1
2 E | An+ Ft,, -ah 1 + 1,2 E Ixt,- -ah x (ah))l 600
2 [ - x", (ah))l + 2 E || (ah) It. -chl 1|2 1 2
(17)
By Theorem B.1, for any E h0,
|I, -.,-s 2 S exp (O(L)) |,,, Ft,, 112 I || It,
(18)
thus the second term in Equation (17) is bounded by L E ||= 21,-ahl 1| S 2 It., 11.
Now, we just need to bound E.. | xx (1
Note that 2; (ah) is the solution to the following ODE run for time añ, starting at 2, at time t,
dx: (xx V (lne( dt
15

---

Published as a conference paper at ICLR 2025
Similarly, En4 1 is the solution to the following ODE run for time ah, starting at Tr time t,
d3: = (T+ + Re,, (Rn)) di
So, we have
0, || 8,112 ==11== 2(x², t1, 0,0, 0,5,)
||-4 - Tz 122 + (5ee - Tc. (2 5r, (f.)))
2 1 1t) || - 2,11 + h In ( (#,)
where the last line is by Young's inequality. So. by Grönwall's inequality,
(2+ T.) ahh) I'" hll 2. Se,, (fnn)l? ds
(1th) Fx++112 S exp
2
h S h
IV I (1 Se,, (Fn))l2 ds
Ii - (x V In Sr, (It,,) 112 + V In qt, (Itm.) Se, (I )11
+ 11, (Gn) 3r, (Ir, )1( ds.
Taking expectation of the above quantity (as always the intial choice of fo Ita), we now use a similar computation as the one in Equation (17):
E 1 (ah)
3 2
x
h
E V
In Cc (1) V In gt. (3re..) 11 ds
V qt, (3tt,) St, (Itt,, )11 122 E 1 (Fn) St, (112 1 E
h
V In ann (x) - V ge.. (wt...) 1| ds + 12222 + 122 1
Ita 112
Now we just need to bound
E V In qtn (s)( - V In Sr. (It..) 112. By Corollary E.1, gt.
1 2 Fm
E V qt (Tt, V (It,,) 112 S d??hh T.
By Lipschitzness of V In and Lemma B.1, for S hn $ 1/L,
|| IV - (s) In ft ( 1] L (s) - It, -s1l
$ L exp (O(Lii)) || Ve,, 112
S L nn Iz, 11?
Combining the above two equations together, we get
1
E In (t.,, (x V In gt., (1=,,) S d2d
L.
+ 112115 xt., 11.
mm
and thus
LV 1 tm
E (ah) - 2 Z f 112
h2 L2dha
+ 12hh 16, - 22,
Plugging the above bound into Equation (17) yields the claim.
Lemma B.3 (Sequential Predictor Bias). Suppose 1. In Algorithm 4, for all n E .O....,N
1}. let x (t) be the solution of the true ODE starting at En at time by and runming until time to t, and let xa qi be the solution of the true ODE, starting at Ty Qto hf E 1 and ty hi 2 t /2, E.
we have
+12 1 (
EE E 1 (h,)) < h
L'dh
It,,
where E, is the expectation with respect to the a chosen in the n step, and is the expectation with respect to the choice of the initial To Stg
16

---

Published as a conference paper at ICLR 2025
Proof. For the proof, wil fix and let f = hn By the integral formulation of the true ODE,
(h) ex + /'00 e--(tnhh Ving (ln - 8)) ds)ds.
-n
Thus, we have
/'h, e--lln Ving
EE 32 -
/1 E c1 (1-0)h Stn (fu+t)
-
1n-h
2 IE | (1lheelloli 1/ at (3n- 1) hell-all In q6, -n
1
fm
l1 Be(l-a)ii v In -ah aah (ah))
(c-(n- Vnng - 5)) ds|i?
tn-h
The second term is 0 since
h (1 Mol--o)// In In -a = h ('' c(1-a)h Ur an (ah)) da
to
c=-(t. (0-(tn-h) In g ())
For the first term, have, by Lemma B.2 (he(l-))) St, St.- an
(Fn+z) ( he(1-1o0 aa) v qr, aan (z)
h + L'dhh
LLV
_ bn
+ L 1h 12,
2t, 11.
The claimed bound follows.
Lemma B.4 (Sequential Predictor Variance). Suppose 1. In Algorithm 4, for all n (O...., N - 1,. let x, (t) be the solution of the true ODE starting at En at time ta and running until time tn and let Ie qi be the solution of the true ODE starting at To qo If hn L and In hn 2 ln/2, we huve
1
E 1 I (hw))l + L'dh (LV))
1))
L' Rnll2
where E refers to the expectation wrt the random a in the nh step, along with the initial choice To to
Proof. Fix n and let ;= hn. We have
E E 1 x²,, (1)ll2
add ( (fntz) < E lh (1 -ayh St., - all (Gn+g)
= E |h e(I-a)) -
J (c-(tn Vng 8)) ds|
hell (1-)aa In Ut -at (2
+ E (1-a)/ In -ah (2 (ah)) (hv (l 5)) ds|12
Vn ds - phn
+ E
-- -- - 8)) doll
j
The first term was bounded in Lemma B.2:
E h e (1-a)hs, st -ad (Fn+z) - he aah In qi. a0l
2 12 [ tx
1
+ 12 1T. Ft.,
17

---

Published as a conference paper at ICLR 2025
For the second term,
fr" -hallhhh - d))
E El (1-a)h/ In st, -o -n (all))
1n-h
E1| 1'h
fm cll-ajh ( Van (1ng o))) dol?
$ h
E -o - Ving - s)l
ln
Now,
E V In St, al (2, (ah)) - In gn - 8) s))1
V In lt,, -ah - Vin qr,))(
+ E V In It, -a - V In (o.
E ( - (1n =s)]l
The first these terms is bounded in Corollary E.1:
LV - in
v -ta (zt, (ah) - v (..g((r)))
2 d'dh?
For the remaining two terms, note that by the Lipschitzness of In qi and Lemma B.1, E V In qt, an (It, aah)
V In -x < L E || -ah (ah))l
S La exp(O(LA)) E| 11
2 1 E Blsse - Fr 11
and similarly, for tn h S tu.
E V In (x.)) V In (1 - s])1| z L' - 122
(19)
Thus, we have shown that the second term in our bound on E [int1 2, (h) 1 is bounded as follows:
E Ellh- e(l-a)/ In an (x, (ah)) - ffr frr (lnga\Vngg +))
-h-h
SL'd Lv
(Lv) - bm + 2 | - 2..12.
For the third term.
Inqo - 5)) ,ss
cx- VV -
8)) dssl
jr
fr h
In (ell-a)h (c(-aahhe"* Ell
V Vinga - -s) s)),
in
$ hn
-
E
[Ving - (s))l
Now, we have
E
s))1 < V
b 3)) (zng(
d + LL - 12
where the last step follows by Lemma E.4 and equation 19. So.
E
Inn
s)) ds J c- -- Vi s)) ds|?
(tn 2
S
E. (c(l-a)h
+ I.2 - - ds
+ I2 11
22.
tn
Thus, noting that h < J. we obtain the claimed bound on E E x3, (h..)).
18

---

Published as a conference paper at ICLR 2025
Finally, we put together the bias and variance bounds above to obtain bound the Wasserstein error at the end of the Predictor Step.
Lemma B.5 (Sequential Predictor Wasserstein Guarantee). Suppose that L 2 1, and that for our sequence of step sizes ho...., hN-1- E. hi I I/L. Let hmax max; h2s Then, the end of Algorithm 4,
1. IftN 1/L,
W 2ix) S
Sc
dh max + dd max
L
2. Ift S 1/L and hn S bo for each
W2 NGN,Tay S
Egc
(L'sh + L33 (dchiaas) N
1?
Here, Van Nt is the solution of the true ODE beginning at Tto To Qta
Proof. For all n E [I,..., NJ, let yn be the solution of the exact one step ODE starting from Tn-1- Let the operator E, be the expectation over the random choice of in the n iteration. Note that only En depends on a. We have
E |||e 2n ||
E [[ UN) - (ãN -UN)ll
| yN 112 2(TEN UN EyN
uN)
+ E
EN
yN yN 1|
L
2
- + 2 ||TeN - UN 1 LhN
BEE Vn 112 + LE 1 - UN 1?
S exp (O(LhN 1|TENN1 1 Lh
| |N + E H UN 113
where the third line is by Young's incquality, and the fourth line is by Lemma B.I. Taking the expectation wrt To to by Lemmas B.3 and B.4,
E ||te
Ex|
expp (O(LhN-1))
TiN-1 - 22-1 + LhN-1
EI - yN - E EISN yN
S exp eepp(O(Lh))) (1(
- ((0x--x1))
+0
1 12 + L'd
LV 1
1) + 1 ||aee-, - Ex-11²))
Lh
tN-1
1 -1
+0
+ +1 LV tN-1
+ 1
1, 1lleen_ - 2N-111?)
S exp (O((NN--))
- - 1112
+0 1
(LV 1 2)
+ - + (I + I'd
L
By induction, noting that Ito
= To. we have
E B|
EN 2 $
N-1
Σ
2
+ 2 (IA + 1.2 (LV. 11)
tn 1 exp
(0
Σ 11
N-1
))
i=n+1
By assumption, Σ, /4 S }. In the first case, LV t. S L for all n, so
2
+ 13 d max + Ldd max
1,2
In the second case.
N-1 hn
WG, I(N) S
+
tn
p=0
2 S
+ (L dh
+ La ahhh
.N.
1
19

---

Published as a conference paper at ICLR 2025
B.2 CORRECTOR STEP
For the sequential algorithm, we make use of the underdamped Langevin corrector step and analysis from (Chen al., 2023b). We reproduce the same here for convenience.
The underdamped Langevin Monte Carlo process with step size is given by:
d 0g dt
(20)
dē; (5( (1/njN)) 70s dt 2 dB.
where B, Brownian motion, and & satisfies
E
[x(x) Vin(<)]l S Sc
(21)
for some target measure q, Here, we set the friction parameter 7 = O
Then, our corrector step is follows.
Algorithm CORRECTORSTEP (SEQUENTIAL)
Input parameters:
Starting sample To. Total time Toore. Step size hcoms Score estimate
1. Run underdamped Langevin Monte Carlo in equation 20 for total time Toorr using step size heore. and let the result be nN.
2. Return nN-
Theorem B.6 (Theorem 5 of (Chen al., 2023b), restated). Suppose Eq. equation 21 holds. For any distribution p over Rd, and total time Teor S 1/vL, if we let pn be the distribution of ã resulting from running Algorithm 5 initialized at Ty ~ P, then we have
TV(pN. 4) $ W2(p.q)
+ +13/441 2/2 /2 hoarr
T./473//2
L
corr
Corollary B.7 (Underdamped Corrector). For Teorr = = Viav/1
TV(px, $ Wa(p. g) d/12, VL+
Esc
+ Ld
heort
V
B.3 END-TO-END ANALYSTS
Finally, we put together the analysis of the predictor and corrector step to obtain our final ((d5/12) dependence sampling time. We first show that carefully choosing the amount of time the corrector results in small TV error after successive rounds of the predictor and corrector steps in Lemma B.8, Finally, we iterate this bound to obtain our final guarantee, given by Theorem 10
Lemma B.8 (TV error after one round of predictor and corrector). Let
be sample from the
true distribution time l. Let bn - -n/L for E [O.....No). If we set Teorr = O VLaizss have,
1. For n E [0,...,N, ift Z 1/L,
Goodl//1
TV(@441)0m-1) TV Ft,, +0 L2 Apred ppeed 173 heorr +
VL
2. If 11/L.
TV($No+1, < TV( IAN,
+0 (
hpred + 3/366/ /oorr + E 4//2
(L 2 + 1//2
pprodd ) log
VT.
20

---

Published as a conference paper at ICLR 2025
Algorithm 6 SEQUEENIALALGGRITTIN
Input parameters:
Start time T. End time 8, Corrector steps time Teors X 1/VL, Number of predictor- corrector steps No. Predictor step size Uprod, Corrector step size hears Score estimates &t
1. Draw To ~ NO0 la)-
2. For n 0,...,
(a) Starting from 2n. run Algorithm 4 with starting time T n using step sizes hpred for all N steps, with N "iipraa so that the total time is 1/ Let the result be 44+1-
-
(b) Starting from T1+1= run Algorithm for total time Teorr with step size heorr and score estimate (1/(x+1)/L obtain Tn+1-
3. Starting from No. Algorithm 4 with starting time T - -No using step sizes /4, 11prea/...... to obtain Noo+1
4. Starting from NNo+1 run Algorithm 5 for total time Teorr with step size hearr and score estimate 88 to obtain +I.
5. Return No-11
Proof. For [O....,.No)) No). let Un+1 be the result of single predictor-corrector sequence de- scribed in step 2 of Algorithm 6, but starting from It, ~ qt. instead of T.. Additionally, let 7N,+1 be the result of running steps and starting from Tin, 9tx, instead of 2 N., Similarly, let Uz+1 be the result of only applying the predictor step starting from Ft,, 9t,, analogous to F defined in step 2a.
We have, by the triangle inequality and the data-processing inequality, for n E [O,..., No 1].
TV Lt, 1) S TV Fo+1) + TV 1t.11)
< TV (...) + (V
By Corollary B.7,
TV st S W2 Ttn4s) d2/12 VL
/ v 17/30 heort
Egc
Now, for in 1/L, by Lemma B.5,
Wa W2 Pta-i) S + 1/12 vdh2
pred + LV dh 3/2
Combining the above gives the first claim.
For the second claim, similar to above, we have
TV( S TV( +1, + TV(
< TV( (eNo) + TV((N++1,
By Corollary B.7,
TV( S W2((NNo+11s)) / VL+
Eac
+ Ld hcorr
VLd1/36
For Nno 1/L, by Lemma B.5, noting that the number of predictor steps in this case is O log Apoed
T. We (7 26) S Cse + (13/2Van2ze
pred + L
Ldd13/2 pred)
log h pred
The second claim follows by combining the above.
21

---

Published as a conference paper at ICLR 2025
We recall the following lemma the convergence of the OU process from (Chen al., 2023b)
Lemma B.9 (Lemma 13 of (Chen et al., 2023b)), Let qe denote the marginal law of the OU process started qo q. Then, for all T Z
TV(ar,NV Ia)) S (Vd+ma) exp(-T)
Finally, we prove our main theorem on the convergence of our sequential algorithm.
Theorem B.10 (Convergence bound for sequential algorithm). Suppose Assumptions 2.1-2.4 hold.
If x denotes the output of Algoritiu 6, for T = θ log
dvm
Teorr
θ
VL
and d θ (T>(ÍVmLT) if we set lppred-
+ min 1-73
2/2
26/12
Tog(mz]
hoor
6 217/36 La72 log (ma)) and the score estimation satisfies Esc S Ō
have
VLd1/12 log m,
that
TV( Sc
with iteration complexity ē (1AS)g/11 ( loog (ogg(m2)
Proof. We will let ln = n/ First, note that by Lemma B.9
TV( vip)) S (vd + m2) exp(-T)
We divide our analysis into two steps. For the first No = O(LT) steps, we iterate the first part of Lemma B.8 obtain
faod21122
TV( TeND) ) S TV(22, Izo.) + (1 122 pred + 1/3/22/112 3/2 'ree + L7
heorr
No
VL
$ (Vd+mm) exp(-T) + 7/12 pred T + L 5/2 )/ 3/2
prea
T + 87/26) hearT
T
esod 2/12
Applying the second part of Lemma B.8 for the second stage of the algorithm, we have TV
2 VVd+ m2)
exp(-IT+ + (L' k/2 ppred + 1/2 /22
prod)
T log prod
+ 3//2 heonT T + esod 1/12T/I
θ (log dVm3 hpred θ
1/2 min
Setting T
and
heor
L
25732
log(ma)
if the score estimation satisfies 6 S
VTal/12 log(m2)) with
2,77
[mg]
iteration complexity ④
L//3 20/12 log" m²) ). we obtain TV( S
C PARALLEL ALGORITHM
C.I PREDICTOR STEP
In this section, we will apply a parallel version of randomized midpoint for the predictor step, where only (log" (dd) iteration complexity will be required to attain our desired error bound for one predictor step.
In each iteration n, we will first sample randomized midpoints that are expectation evenly spaced with 8, = R. time intervals between consecutive midpoints. Next, our step (c), we provide an initial estimate on the x value of midpoints using our estimate of position on at time t, provided by iteration n 1. This step is analogous step (c) in Algorithm for the sequential predictor step. Then, in step (d) we refine our initial estimates by using a discrete version of Picard iteration, where for round k compute new estimate of It,- aahh based the estimates of
hns
22

---

Published as a conference paper at ICLR 2025
ho, for j S i in round k 1. Note that trajectory z(t) that starts from time to and follows the true ODE is fix point of operator r that maps continuous function to continuous function, where
h,,
= ( + [ e" e Vg.(x(s))ds.
By smoothness of the true ODE, the continuous Picard iteration converges exponentially the true trajectory, and we will show that discretization error for Picard iteration controlled. After the refinements have sufficiently reduced the estimation error for randomized midpoints, make final calculation, estimating the value of +h, based on the estimated value at all the randomized midpoints.
Algorithm PREDICTORSTEP (PARALLEL)
Input parameters:
Starting sample To. Starting time lo. Number of steps N, Step size (h ) Number of midpoint estimates {Rn J.y Number of parallel iteration {Kn Score estimates 3.
For all n 0, ... N 1: let 8n Ra
1. For n .,..., N 1: 1:
(a) Let in lo Ew=o/t Aw
(b) Randomly sample a, uniformly from [(i ((i-1) (0)-1)/RB 1), 1/ for all E f1,... Rnd
(c) F = 1,.., R2 in parallel: Let x!! = eannn + (cashn 1) St, (Zn)
(d) For 1,... Kn:
For = 1.... Rr in parallel:
x(k) n,3
Co Fn
+ Ej= (carha- -(--16- max -Jon
=
St. kyn (3
(e) T Tn = cn + Sn E/= el aailen Se, Bn (s((K))
2. Let In = to - nn
3. Return EN. IN.
In our analysis, we follow the same notation as in Appendix B.1. We first establish a poly L,, d) bound on the initial estimation error incurred in step (c) of each iteration in Algorithm 7.
Claim C.1. Suppose L L Assume hn $1/L. For any O.... .N suppose we draw Tn from an arbitrary distribution Pw, then run step (a) (e) in Algorithm 7. Then for any i = -.... R
E 10 z (a.h.))
ha dL tn hn
12222,12 11.
where x", (t) is solution of the true ODE starting at Fn at time tn and running until time t, t.
Proof. Notice that in step (c) of Algorithm 7, the initial estimate of the randomized midpoint is done with the exact same formula as in step (c) of Algorithm 6, except we calculate this initial estimate for R, different randomized midpoints. Notice also that the bound for discretization error in Theorem B.2 is not dependent on specific value of a, as long as the randomed value is at most 1. Hence
identical calculation yield the claim.
Next, we show how to drive the initialization error from Lemma C.I down (exponentially) using the Picard iterations step (d) Algorithm 7.
Lemma C.2. Suppose L 2 1. Assume hn S 1/LL. For all iterations n E {0...., N 1}, suppose draw Tn from arbitrary distribution Pu, then step (a) (e) in Algorithm 7. Then for all
23

---

Published as a conference paper at ICLR 2025
k E {1,... Kn) and {1, R
R - 2 (ay)ta))
E [(h) - 2, (a:ttt)) 2 (8)2
Σ
R
==1
+ L2 Bl - Den 1), (22)
+ 11
2
Ssc
+
R2
(LV
in
hn.
where 2, (t) is solution of the true ODE starting at En at time tn and running until time t, t.
Proof. Fixing iteration n, we will let h == hn. R = Rn and 5 := 6n The formula of ((a) and 2 (ash) has the same coefficient for f.. thus we can bound the difference as follows:
(k) Fn x (agh)
E
f'h
In
E
ce-(n--shh ds St, -ayh (2
- 8)) ds
-ach
11-1)6
$
cx (t-(tn ds 8t -a.-aa (o
V V At. Vigg:-oaalth(aatt))
(23)
6n mii((j.,..)))
to -0 118
+ E
- (1 s))) ds
In minnjjhaahh)
(24)
The first second line is by definition and the second third line is by Young's inequality. Now, we will bound Equation (23) and Equation (24) separately. By Theorem 2.2 and Theorem 2.4,
E
h (#p3
(x V In St, (x
2
S E
8 nn aajh (k-))) 1)) St, oyh (It, -0,n mn)
2
+ V In (x - V In qt. oy 1n (
+ E 8t, 0z (It, -0gh) - Vlngt -oyh (It, (0t,--,,i01
Se2
_
[k- 1)
1.2 F (a.h) 2it,, -a,h ?
x n.
It.,, -1
<2
I2
E
(k-1)
x (ayh)
||
2
+ +I2 E 1E Irn 3en-ajhhll
By Theorem B.I, for any E [0, h).
E -. (s) =2,-=112 < exp (O(Lh) E SE =1,
(25)
24

---

Published as a conference paper at ICLR 2025
thus L E
E (ayh) Tt, -am-a,h11 L2 L Iea 112. The term in Equation (23) can now be
bounded as follows
2
11-166
(0-(haaiih) ds (2 V In in o
E
min(j6,a.A)
ln 11-16
(tn aaihi ds (ygh-( - - bo -min(((),
R
/=1
SR. 82
2a,h
E
St, aayh(nnn (1/
52R 88 20ch Σ
2 + 2
ing - (ayh)))
+L2 E 11 2x,1!
==1
R
<2R² 88. .2,th 1
Eec + 2 1 - -aa)))
+L2 13 1| - Iea Ip)
R
j=1
<4h² 2. +4h2
Iz,, 112 + 222L2
1
EE
Tn.j
(k-11)
x (a,hi))
The first second line is by inequality E a;)² S n Σ at the second to third line by the fact that -e. 1n--111166 1))6
m.in(j,,aa.))) ( ds 8. eacb the fifth to sixth line is by R5 h and that e 2a,h
at most when < 1/4.
Similarly, the term in Equation (24) can be bounded follows
tn 13-1)6
E
e" 5-((,, -b,b) (V In St., (( - (1n ds
min(j4...))
11-188
<R.
E
e- ca-tn--ath In ft. (V
Vnnglo 8))) ds
ln -min(jj66,
j=1
1)5
<R 8. 20,h,
E
V In It -, - (1 6))]1 ds
min((5.auc))
Since | (ayhi - 51 5.
Vtt. (iag - s)]l|
<3 V qt
Vq -a,n -V
+3E (x.(x.) 3
Vgol= s))11
+ 3 91.
=3 R V , (ayh)) V (F,, 33 [ - (t
=))l
+3E
It,, =0n_- 31.2 || (t,, -s) -x.112 3 bn- On v In gn
2
$31. E
25

---

Published as a conference paper at ICLR 2025
The second to third line is by Young's inequality, and the fourth to fifth line is by Theorem 2.2. By Lemma 3 in (Chen et al., 2023b),
a
ta an-a
E
Ou V la (Ing(
S6.
E V du
2 pnn phn
L L'd _ S (L
1
u
tn
Now, by Theorem B.1 and the fact that h S t,
E Ve - (ingo(= s]){|² 12L exp( [ xt, do"
1
tn
h
36L 2 E 13 2t, 11 d2 (LV
1
tn
We conclude that the term in Equation (24) can be bounded by
15-16 1)6
R 8 2ach
Σ
36L E E| It, 112 112 33L (L
dd ds
In h
min(jj,aah)))
<2R28²
366.2 |s
Ft,, 11 36.226
tn -
=4h² L²
18E j - Iz,, 2| + 2
do
(L
In
Combining the bounds for Equation (23) and Equation (24), we get
- - (add))) S 12.
1 R
T n. - xa
(k-1)
2 + 12 1
dō
1 + + B| It,, 12
In - h
Given sufficiently small constant h, eash < 2. Moreover, by the definition of 6, Rō h. By unrolling the recursion, we get
R
E
T:
(k) n,i
Ti, a,h
I (8h2112)
R
Ø=1
812
+ 1 812
+
do²
In - + El - se.ll??
7,2
1 R
$ (3)*
zn(ashill
R
+ 12 ( + L2d8
(LV (LV
In , E xr.,, 11²)
As a consequence of Lemma C.2, we take the number Kn, of Picard iterations sufficiently large, the error incurred in our estimate for Ft,, -a, is dominated by the terms in the second line of Eq. equation 22.
Corollary C.3. Assume L 2 L. For all n E {O..... ,N 1}. suppose we draw En from an arbitrary distribution Pu. then step (a) (e) in Algorithm 7. In addition, suppose hn 37 and Kn, Z log( RR). Then for any i E (1,... Ra).
((c (( ||
E 2 +
L'dh4
LV
1 1
1.212 1 E It.,, 11?
R2
fn
hn
where 57, (1) is solution of the true ODE starting Tn time in and running until time t.
26

---

Published as a conference paper at ICLR 2025
Proof. Fixing iteration n, will let h := hns R := Rn, := 6n and K := Kn Notice that when
is at most Tea Now by plugging Theorem C.1 into Theorem C.2, K3 log sx,2, -o)g(()) (2) K
get
(K) - n (acl)) E
2 (81212) .1.. ( +
I'd
In -n- -,=-,)+++ 1,2 1-1 Ftm! 11
+ h 2e 1I
1
+ 1,2 J0 - Ien 11?
in - h
x 13) +1) hr (2 + L²
E 1 - Iznll 11
+ ((8h2 12)K
+ R
L'dhh
[LV tn h
d'
$ h22 50
1.213 E T Iz,,l 112
L.
R²
t,
-
The first second inequality by rearrangement of terms, while the second to third inequality is by the fact that the terms
1 Ra dominates (
We can now prove the parallel analogue of Theorem B.3 and Theorem B.4. Note that the bounds in Theorem C.4 and Theorem C.5 are identical to the bounds in Theorem B.3 and Theorem B.4, except from additional Ry factor for the middle term. This additional factor stems from using Rn midpoints in each iteration (compared to using one midpoint each iteration in Algorithm 6).
Lemma C.4 (Parallel Predictor Bias), Assume L > 1. For all n E fo..... .N N-13, 1}, suppose we draw En from an arbitrary distribution pas then run step (a) (e) in Algorithm In addition, suppose bn < 3L and Kn Z log( R.). Then we have
Lnhed R
1
E E (hn)) S hn 2
(LV
})) + L B It. 11?
in hn
where xn (1) is solution of the true ODE starting at Fn at time in and running until time ln L.
Proof. Fixing iteration n, we will let = hns = Rns 5 = 8n and K = Kn, We have
E E E. nt1 x,, 11)
da c- - Vin 8)) ds
h a St. hih (2111))
E
E
8
<2
E
Σ de h-ah acb 1an-aak (2(11) (K)) V Vinge aith (ach (e
(26)
( v Ling 8)) ds Vlng..-a - tn
+2
E
òe h ach
-n
(27)
Since o; is drawn uniformly from [(i - 1)6, 11), i0).
(1a-61-1)6
E. Te k-sh,, V In gr., (a, (osh))) =
e c -1) In qo 5)) ds,
27

---

Published as a conference paper at ICLR 2025
and the term in Equation (27) is equal to O. Now need to bound Equation (26).
E 8/, a. (1 V In qc, oa
5 E
-anh (2 (2)1)) - St, -ahh anh)
+ E In -ah V In qt -n
EE ||tt ast ( -ash)
Vin qt, -ach (Ttn
SL2 E
[ - (agh))))
+ + 18 || (agh) Ita -a,h- Exc
SIZ
122 +
d'
LV R² [LV,,'
1 -m) + 1.212
E 1 Itn 12)
12 Ft,, 112 2
in
The first to second line by inequality (Σ o:)² < n Ei= a₂ the second to third line is by Theorem 2.4. Theorem 2.2 and Theorem 2.3, and the third to fourth line by Theorem C.3 (K 2 log( R), which satisfies the condition in Theorem C.3) and Theorem B.1. Hence
R
2
E
de he
V In St (Ft oihh)
R
<2R(26)²
2
E2
St (an-a In qt. -ach (Ft, (B
i=1
5852R2.
(2e
(1222 +
dh
LVV In - + 1.22 B -Tt,, 122
+L2BIIT, - -xem 11
R2
L'' R²
1
5h2 22 +
(LV
+
T.2 2
E
1 Ie., 11?
tn -
The first to second step by inequality (E:-1 a:) 2 Σラ at and Young's inequality, the second to third line is by plugging in our previous calculation, and the third to forth line is by = 8Rt and that h S 1/L.
Lemma C.5 (Parallel Predictor Variance), Assume L L. For all E {0, NN 1.. suppose we draw En from an arbitrary distribution Pro then run step (a) - (e) in Algorithm addition, 1 and Kn 2 log(Rn). Then have
suppose hn
L'dh R2
(LV
1 -V. + L2hh, 1E Zz,, 11?
$ x (hn))l S . 22.22-
In In
where (1)( is solution of the true ODE starting at Fy at time in and running until time ln L.
Proof. Fixing iteration n, will let h hns R = Ra, 5 = 5n and K Kn, We will separate Eg [ x,, (h)l² [2 into several terms and bound each term separately.
E [ (h)1
R Σ c a-ash
f hc 8)) ds
St (1
E
ta-h
R
<3
8e h-ahh
h
(2(((((())))
V In qt -ach (It, (Fth
(28)
In - (f-1)8
+3E
h-ash
(V -oth) (ing ss)) ds
-ach
(29)
Mu (b-16
+3E
_ (h-i)) (lng 8)) ds
(30)
i=1 -i-
28

---

Published as a conference paper at ICLR 2025
Equation (28) is identical to Equation (26) in Theorem C.4, and can be bounded by
L'dh
1 ))) In
Equation (28) S h22 < + 1...2-B B It,
(L
R²
Next will bound Equation (29). By Theorem E.1 and Theorem B.I,
1 in - h
E Vgt, a - v (1 s))[| s L d'
1 It, 1l.
hence Equation (29) can be bounded with similar calculations as for Equation (24) in Theorem C.2, by the following term:
1 - + O(L3). - 82 ( (LUBLLV in
d2d R2
E 22,.13)
(LV
tn - -
( x1,1l
Finally will bound Equation (30) Since both ag and belong to the range [(i - 1)6, i8].
a-ah x- (tn hh) h) S (e (1-11) e -5) S e 5 S 25.
Moreover, by the fact that V In qo )112 & Ld (by integration by parts), Theorem E.1, Theo- rem B.I and the fact that Lb o(1), we have
E V ln s))1 < E V In (I,, )l + V
s)) V (z S Ld L expl L5) 1 It., 11 z Ld + I,2 1/1 It,, 11.
Hence Equation (30) can be bounded by
R
t, 116
(cmaach - (lngs (3) - d)
E
18
Mn
136
<3R6
2| In (n -
In (1-1)8
<3R6 482
Ld + L2 12n II,. 11?
Ldh
h
SR80" LLd + L2 1 2t,11
It, 112
122
R²
By adding together Equation (28), Equation (29) and Equation (30), and combining terms, we con- clude that
LA
E Bn41
(1)11
S 122 2 + 122
E E Ft,, 112
(L
V
tn h )))
R²
a
+ 12 E T Dz, 1| + Ldh
I
+ R²
(LV
11 Fz, 11?
tn - h
R2
R
d2
1 tn
$
2 +
R2 (L.V
+ 122 1T Ft., 11
We can now prove our main guarantee for the parallel predictor step, which states that with logarith- mically many parallel rounds and Õ( /d) score estimate queries over a short time interval (of length O(ILL)) of the reverse process, the algorithm does not drift too far from the true ODE. Our proof follows a similar flow as Theorem B.5.
Note that due the existence of R., midpoints in each step, can set h. O(1). We will set h, to be @ O unless t, is close to the end time 8 (see Algorithm 9 for the global algorithm and timeline). If L, is close to 6, will repeated half h, as increases, until we reach the end time.
29

---

Published as a conference paper at ICLR 2025
Theorem C.6. Assume L 2 1. Let 3 1 be adjustable parameter. When hn Lvd
min TL 1n/2,t2 8}. K Z log(
Sva and Ra - ho .3 33
O Bva the Wasserstein
distance between the true ODE process and the process in Algorithm 7, both starting from to no and run for total time TN S 1 is bounded by
L
CCc L BVT.
Proof. To avoid confusion, we will reserve En as the result of running Algorithm 7 for n steps, starting at To go- We will use Jn denote the result from running the true ODE process, starting at Tn-1. Then by an identical calculation as in Theorem B.5,
2
2 + E yn
TEx Gxl1 S exp (O((hN-))) || - 1N-1 + L EêN
uNll
Now we do a similar calculation as in Theorem B.5, but utilizes the bias and variance of one step in the parallel algorithm instead of sequential. Taking the expectation wrt To Uta by Lem- mas C.4 and C.5,
E 1I
ENl 12
exp
F - 2N-1112 +
E EI - yN - E -UN
LhN-1
$ exp (O(LAN
- 1|
+0
112-160 T
1
11
L L'dh
LV IN-1
+L2 E 1| :N- 11²)
R2
1 1)
+0
1
2
+
LV
+ 1
Elaexx_ - 2N-11?
1R
IN-1
5 exp (O(LAN- () 11112E - -8N -11
IN /
Lh N-1
+ dh N-1
1 L
+0
20 +
fN-1
L
R²
Since hN-1 1 the term L'dh is dominated by the term L'dhhy- and the term h cec os
dominated by I By induction, noting that Tto To. have
E ||
EN 12 S
N-1
hn s0 +
L'd (IV ) -exp O L N
(0)1)
Σ h1
m)
R
n=
i=n+1
Since N- t S t. exp
(0 L Dinnt kx)) is a constant. Moreover, by our choice of h it
is always the case that hn < to < tn - hn and that hn S 1. Therefore
L dh
1
dh
dh
hn
LV
R2 in hmn
R
LLd
f2
2
and thus
E Blaae - 2x x
N-1
Σ
hn
+
h
6sc
+
2
8c
32
I
LB
n=0
We conclude that
W W2
/B
11xex
S
Esc
+
T,
BVL
C.2
CORRECTOR STEP
In this step we will be using the parallel algorithm in (Anari et al., 2024) to estimate the underdamped Langevin diffusion process. Since we will be fixing the score function in time, will use V In g to
30

---

Published as a conference paper at ICLR 2025
denote the true score function for the diffusion process, and 8 to denote the estimated score function. We will choose the friction parameter to be y x VL.
Algorithm CORRECTOR STEP (PARALLEL) (Anari et al., 2024)
Input parameters:
Starting sample (0o. 00) ~ p X |0, Ia), Number of steps N, Step size h, Score estimates h
3 V qlgg Number of midpoint estimates R, == R
1. For n = 0,...,N-1.
(a) Let En = nh
(b) Let (# o(h)) represent the algorithmic estimate of (ft, RR: +sh//)) at iteration
k.
(c) Let (C", (t) be a correlated gaussian vector corresponding to change caused by the Brownian motion term in h/ R time (see more detail in (Anari et al., 2024))
(d) For i 0, R in parallel: Let (7 (00))
(Ezz Fx)
(e) For = -,,.,, ,...,, K:
For = 1.... ...... .R in parallel:
=!(t = - + 1-expl exp(->h/BB) - A/R (1 expp( 7/4 ////R//1 ((ann ) +
on exp(- (->h/ ((k-1)
nniil 1-epp(-ahi 3(1,1,11 + 5"
(k-1)
( (In+1, 1n+1) (*t T
2. Let IN Nh
3. Return CN IN.
Let Tv denote the total time the parallel corrector step is run (namely, IN = nh). Consider two continuous underdamped Langevin diffusion processes u (t) = (t), v= (t)) and Utot = (1t Uto+e) with coupled brownian motions. The first one start from position x: z²(0) = Eo ~ and the second one start from position Zt, ~ q. Both processes start with velocity V* 00)
Uto
N(O, Ia). We will bound both the distance measure between (t) and the true sample t and the distance measure between x (lv) and outputs Algorithm 8. First, (Chen et al., 2023b) gives the following bound on the total variation error between x" (TN) and Tex-
Lemma C.7 (CChen et al., 2023b), Lemma 9). If h S VI
1 then
TV(2"(Tx); W2(p.g) L
Next, (Anari et al., 2024) bounds the discretization error in Algorithm in terms of quantities that relates to the supremum E || qa= (f)) and E [|ee (t) where t E [o, Tw]
Lemma C.8 (Annri et al., 2024), Theorem 20, Implicit). Assume L 2 1. In Algorithm 8, assume K log(d) (for sufficiently large constant), K S log R and h. S VL Then
VT. (2c + 2dh2
TN
122
h RI
0)
KL x ((Tx) S
+
P
+
R
where P = SUP1e(0.Tk)
[[]]=(()]]" and Q = supre(o,Tx]
| V In ((xx()))))))
To reason about the value of P and Q, we will use the following lemma in (Chen et al., 2023b).
Lemma C.9 ((Chen et al., 2023b), Lemma 10). For any S VI
E ||uu" (()())
tel S W (p. q)-
31

---

Published as a conference paper at ICLR 2025
Lemma C.10. Assume L 1. For any TN E VI¹
Elll In ( L Ld
sup te[o,Tx]
and
E |o00 (2 + d.
sup 1e
Proof. Note that (0, Ia)) is stationary distribution of the underdamped Langevin diffusion process, hence 3y g and 1/2 (0, Ta)- Hence E. V/ g(x))l s Ld by integration by parts. Similarly, E
for any E [o, Tx]. we can now bound || /| = [O) Ia) I4)ll & d. Since Tv 2 VI
E In ((= (1) 2 and E [(o=(()) by Theorem C.9 as follows:
E V In (( (0)] < 2- E
+ 2- B V In (x (1)) - Vin (xx)))
Ld + 1.2 (t) xall s Ld + 1.2 Elu well Ld 1.23(
and
E Bl(x((t)1
2 110" [10)())))) - v111 E (t)ll 2. 12
E +E (t) - u, d = (2(
Theorem C.11. Let 2 be an adjustable parameter. Algorithm with parameter 8. 0 Vi K 4 1.log(8 log(R) and Tn 2 + has discretization error
R
VST
TV(@N, (Tx)) S V /L x² (Tx)) S
Esc
+ +
W2(p.g).
VL
B 3Vd
Proof. Since TN S VL we get that
and h 0 ( 0(1). Plugging Theorem C0 into Theorem C.8,
TN
1dhh3
h
+ ha Rea (LL + LL) )
KL x² (Tx))
S
2 + 2
+
(d+ + w2(p.
VL.
R1
R2
Ld
Ld
L
L2
<
+
+
+
w2(x,))
T
8c
R4
R2
Rd
R2
R
2
2
W2
(p.
L
Bdd
The first to second line is by combining terms and setting / = - O(VL), and the second Taking the square root of KL( x= (Tv)) yields the third line is by setting = 3 1 v claim.
Theorem C.12. Let 3 2 1 be adjustable parameter. When Algorithm 8 is initialized at (Zo, Do) p @ (0, Ia), there exists parameters h = VST. R = B ( yd K = Olog( ( and TN S + such that the total variation distance between the final output of Algorithm 8 and the true distribution can be bounded as
TV(
Esc VL
+ W)(.g)).
3
Proof. By triangle inequality, TV( < TV( x" (Tx)) + TV (Tx), T(x). Combining
Theorem C.7 and Theorem C.11 yields the claim.
32

---

Published as a conference paper at ICLR 2025
C.3 END-TO-ENT ANALYSIS
Algorithm 9 PARALLELALGORITHN
Input parameters:
Start time T, End time 8, Corrector steps time Teors X 1/VI, Number of predictor- corrector steps No. Score estimates St
1. Draw To (N0
2. For n .,.
(a) Starting from Tn.. run Algorithm with starting time T P-nL with total time min(1//., n/ 6). Let the result be H 2+1
(b) Starting from 2 run Algorithm for total time Teorr and score estimate (x+(++)) to obtain Tat1.
3. Return NN++1
Theorem C.13 (Parallel End to End Error), By setting T = O (log (dvm)
Teorr
VT. 8
+
A(dvm3) and 81 $2 θ Log ((vm))
in Algorithum 7 and Algorithm 8, when Eac S the total variation distance between the output of Algorithm and the target distribution
VT.
Fo ~ q" is
TV( Sc.
with iteration complexity Õ(L log² [dvmm]
Proof. Let Tt, be the result of running the true ODE for time T starting from 2T ~ q"' Let v/2 be the result of running the predictor step in step 1 of Algorithm 9, starting from Tt 9ta-1 and start time 1- In addition, let Dn be the result of the corrector step in step n of Algorithm 9, starting from vna
We will first bound the error in one predictor + corrector step that starts at fn-1 = T (n - 1)/La By triangle inequality of TV distance and data processing inequality (applied to T, and ÿ..)
II, TV In) TV( Zt..)
S TV( + TV( II.))
(31)
By Theorem C.6 parametrized by 31 and Theorem C.12 parametrized by B2.
TV(0n, It,)) 2
Esc
+ VL. W2(Wh I...)
VT.
32
Scc
VL
Esc
VL
32
BIVI.
Ssc
c
VL
min 32)
The first line is by Theorem C.12, and the first to second line is by Theorem C.12. Next, note that the beginning of the process, to T, and the end of the process, N 5. By induction Equation (31).
TV(@N++1. 20) & TV(x0, I1N + TV FoNo+1)
No+1
TV( 2s) TV(2T. +
Σ
TV Fim)
n=1
TV(20, ((rr,N(0 Ta)) No
VI min( 32)
33

---

Published as a conference paper at ICLR 2025
By Theorem B.9, TV(rr, x (vd m2) exp(-T). By (Lee et al., 2023, Lemma 6.4),
TV rz) - 6. Therefore by setting T = e log dvms
No = θ L
dvm3
and
e
31 32 θ -6(logg T0)) e.
AVmá ))
in Algorithm 7 and Algorithm 8, when E x O (V7))
obtain
The iteration complexity of Algorithm 9 given above parameters is roughly number of predictor- corrector steps times the iteration complexity in one predictor-correctto step. Note that in any cor- rector step and any predictor step except the last one, only N O(1) number of sub-steps are taken, therefore the iteration complexity of one predictor step (except the last step) is O(log( 9 V))
and iteration complexity of one corrector step is (loog( 22)). In the last predictor step, the number of steps taken is O (log (1z) (log(L) + T), and thus the iteration complexity is
((log(L)) +T) loog( (LvZ)) )). We conclude that the total iteration complexity of Algorithm is
LT. (log( Biva )) + (log(
Bid e²
+0 (log(L) + +T) log Bivd )) = (logg Ld [d3
m²
1)
c
D LOG-CONCAVE SAMPLING IN TOTAL VARIATION
In this section, we give simple proof, using our observation about trading off the time spent on the predictor and corrector steps, of an improved bound for sampling from log-concave distribution in total variation.
Definition D.1. distribution with probability density p is a log-concave distribution if log p is concave. Formally, for any I, y in the domain of and for any X e (0, log f(Ar + (1- A)y) 2 Alog fx) + (1 X) log fly)
Note that for this section, assume that is the true score of the distribution and is known, as standard in the log-concave sampling literature. We begin by recalling Shen and Lee's random- ized midpoint method applied approximate the underdamped Langevin process, for log-concave sampling in the Wasserstein metric (Shen & Lee, 2019) in Algorithm 10.
Algorithm 10 RANDOMIZDDMIPPONTEETHHO (Shen & Lee, 2019)
Input parameters:
Starting sample To. Starting to. Number of steps N, Step size h, Score function 8, u = 1
1. For n = 0,.,,
(a) Randomly sample a uniformly from [0, 1].
(b) Generate Gaussian random variable (w("), w:("), w(()) E 3ddd as in Appendix A of (Shen & Lee, 2019).
Dn (1 c 22ah) Un fuu (ah - (1-c e e(2(h)))))) 8(zn) + (d) Let 2n+1 = In + 1/2 (1- e -2h) un - uu (1 c 2(2(h-ah)) ((Entt) + vuw.(m)
(c) Let Ent] (ww((a)
(e) Let Un+1 = UnE 22h uhe 2(h 2(h-aa)) ( + 2 2//W((
2. Return TN-
Theorem D.2 (Theorem 3 of (Shen & Lee, 2019), restated). Let š V lnp, the score function of a log-concave distribution be such that 0 & m Ia $ Jz(x)) & L La. for the Jacobian Jg of 88. Let To be the root of's s. and ui = 0. Let Æ = be the condition number. For any 0 < c < 1, if we
set the step size of Algorithm 10 as = C Cinn 2/9m 1/6
2:/ log -1/6
213 1/3
log - 1/3
4178
34

---

Published as a conference paper at ICLR 2025
for some small constant and run the algorithm for N = log iterations, then Algorithm 10 after N iterations can generate Fy such that
20d
<ō
//0/// 23/3m
+
kd
h
2m
22/3m
W2(2N. x) S
where I ~ p.
Now, we make the following simple observation if run the corrector step from Section B.2 for a short time, we can convert the above Wasserstein guarantee to TV guarantee. We carefully trade off the time spent on the Randomized Midpoint step above and the corrector step obtain the improved dimension dependence. Our final algorithm is given in Algorithm 11.
Algorithm 11 LooCONCAVRSAMPIING (Shen & Lee, 2019)
Input parameters:
Number of Randomized Midpoint steps Nrand, Corrector steps Time Teorr 2 1 ized Midpoint Step size hrand Corrector step size heorrs Score function S.
Random-
VI*
1. Let To be the root of s, and let = O.
2. Run Algorithm 10 with Nrand steps and step size haads using Tu, uo and let the result be
T Nama"
3. Run Algorithm starting from Nen.. for time Teoms using step size heor. Let the result be
T Naws
4. Return 'Nas
We obtain the following guarantee with our improved dimension dependence of (d
Theorem D.3 (Log-Concave Sampling in Total Variation). Let 3 = V Inp be the score function of Ia Jz(x) & L for the Jacobian J3 Let
log-concave distribution p such that
be the condition number. For any < 1, we set hrand C 20/12/1/3 log 1/3 ddc for
2/3
small constant C, Nrand log c
< Ö (5(~"00 1/12 )
hoom Ō (a:7/OVT.) and Toore
62/3
O ( we have that Algorithm 11 returns &iae with
TV( 1)) S
61/3 for x ~ p. Furthemore, the total iteration complexity is Ō (ds/12
2313 + 2))
Proof. By Theorem D.2, we have, for our setting of Nrand and hyand that, at the end of step 2 of Algorithm 11,
W2 x) d
Then, by the first part of Corollary B.7,
TV x) S + V 27/36
E Cc.
(17/366 /L
Our iteration complexity is bounded by Nman
Teamz
δ
2612
claimed.
Heare
E
HELPER LEMMAS
Lemma E.1 (Corollary 1 of (Chen et al., 2023b)). For the ODE
dr. (x, V In (xx))) dt,
L 2 and 2 (x:1 < we have for < and h =t
LV 1
(xng(x)))) || S a2
35

---

Published as a conference paper at ICLR 2025
Lemma E.2 (Implicit in Lemma of (Chen al., 2023b)), Suppose L 2 h S and tu/2. For ODEs starting at zt, = Tta where
h
dx. = (It + Vln dt
dē; = (Ix + 8to (Tto dl,
we have
Eto 111 S 1i2 ( 12ddh² L
- to
E
+ 600
Lemma E.3 (Lemma B.1. of (Gupta et al., 2023a), restated). Let po be distribution over Rd. For To Pa, let It = + Pt for ~ ~N(,, (O,)Ta) independent of vo. Then,
Pe( +c)
Tx 1412
E
[c
(x))
and
v Iipp((x:)=
E
2t
Re ie
Lemma E.4. For (e)))) x xp p NO0, (e2r T)) we have
VIng(y))= = (h = e¹
=t
=eletu
e
Furthermore,
d
E Mi-a
[I
S
Proof. The first claim is immediate consequence of the definition of and Lemma E.3, For the second claim, note that
[[] In (v:)11
26
E
2t
Seruge
We ruge
2e yt
21
||ee| 1|
2t
E
(e2t - 1)3
||2
28
E
2e
1)2
2t
since ~ ~N(0 (e -1)4a)
21
36


Reviews:
[{"review_id": "RcOuiSD4sc", "reviewer": "Reviewer_s9Bu", "summary": "The authors explore sampling for diffusion with mild assumptions and introduce both sequential and parallel algorithms, which extend the ODE-based predictor-corrector framework initially presented in [CCL+23a]. In their sequential approach, they adeptly utilize the randomized midpoint method in[SL19] to discretize the probability flow ODE during the predictor phase, and use the corrector method outlined in [CCL+23a] with a short time $T$. Impressively, this method achieves a complexity of $O(L^{5/3}d^{4/9}/\\varepsilon)$, which marks a substantial improvement over the previous best complexity of $O(L^2d^{1/2}/\\varepsilon)$. For the parallel algorithm, the authors also apply a similar parallel randomized midpoint method from [SL19] for the predictor phase and use the parallel underdamped Langevin method from [ACV24] as the corrector. This innovative combination gives the first provable guarantees for parallel sampling within diffusion models. Furthemore, this paper gives an improved bound for log-concave sampling.", "questions": "- is it possible to improve the total query complexity of paralell version to $d^{4/9}$?\n\n- What are the primary challenges associated with removing the smoothness assumption?\n\n- In the parallel version of sampling for diffusion models and the log-concave sampling algorithm, the predictor-corrector framework runs only once, whereas in the sequential version, the predictor-corrector loop is repeated $N_0$ times. Why does this difference occur?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Their method achieves significantly better sequential complexity and parallel complexity.\n\n- Overall I found the paper well-written and clear.", "weaknesses": "- Their sequential result requires a slightly stronger assumption on the score estimation error, but I think it is fine.\n\n- The technical contributions are limited to combining existing ODE-based predictor-corrector framework and discrete schemes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "027c0DEsBt", "reviewer": "Reviewer_u9VU", "summary": "The manuscript introduces a novel approach for enhancing sampling efficiency in diffusion models by incorporating a randomized midpoint method. This method, inspired by Shen and Lee's work on log-concave sampling, is adapted to diffusion processes to reduce sample complexity and improve iteration bounds. The authors present theoretical guarantees for both sequential and parallel sampling, addressing the limitations of previous diffusion model sampling techniques.", "questions": "**Space Complexity and Scalability.** Can the authors elaborate on how the proposed method would scale with large data distributions in practical applications? How feasible is this method for scenarios requiring substantial memory or computation resources?\n\n**Empirical Benchmarks.** Could the authors provide insights on the types of empirical tests they recommend or any preliminary results? Testing on standard datasets or benchmarks in score-based generative models would be particularly valuable for evaluating this method’s efficiency.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Theoretical Innovation.** The manuscript extends a randomized midpoint technique, traditionally used in other domains, to diffusion sampling, presenting significant advancements in sample complexity for smooth distributions.\n\n**Parallel Sampling Framework.** Including a parallel sampling approach with provable efficiency bounds is particularly notable, potentially paving the way for scalable implementations in large-scale generative modeling tasks.\n\n**Solid Analysis.** The authors offer a rigorous mathematical foundation for their method, with thorough proofs and well-defined assumptions on score functions and dimensionality dependence.", "weaknesses": "**Assumptions and Practical Applicability.** While the theoretical advancements are clear, the reliance on strong Lipschitz assumptions and high accuracy in score estimates might limit practical applicability. The authors could consider discussing the feasibility of these assumptions in real-world diffusion models.\n\n**Lack of Experimental Validation.** Despite the paper's focus on theoretical results, experimental validation would have been beneficial. Empirical benchmarks could highlight practical gains and offer insights into scenarios where the proposed method outperforms existing approaches.\n\n**Parallelization Overhead.** While the approach achieves polylogarithmic round complexity, it requires significant parallel resources to fully realize these benefits. The paper could improve by addressing potential bottlenecks and clarifying hardware implications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "O1tq5V9LAx", "reviewer": "Reviewer_henE", "summary": "This paper uses the randomized midpoint method to solve the reverse-diffusion SDE in generative diffusion models, to generate samples.  The randomized midpoint method is a more efficient method of solving SDEs than methods which incur a first-order error, such as the Euler-Maruyama integrator, as long as the coefficients are smooth functions.   The randomized midpoint method was originally developed for implementing the Hamiltonian Monte Carlo Markov Chain algorithm, but the here the authors apply the same technique in the setting of generative diffusion models.  This allows the authors to obtain faster runtime bounds for sampling (in the total variation norm) from the reverse-diffusion SDE in diffusion models, provided the certain smoothness assumptions are satisfied.  As a by-product of their analysis the authors also obtain improved bounds for sampling from log-concave distributions using the Langevin MCMC algorithm.", "questions": "Can the authors provide simulations of their methods? It would be interesting to see if their method works in practice.  This concerns both the runtime and the quality of the generated samples.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper obtains guarantees, in the TV norm, for solving the reverse SDE in diffusion models, with improved dimension dependence on the runtime.\n\nThe paper obtains improved guarantees, which are logarithmic in dimension, for parallelized solving of reverse SDE in diffusion models.\n\nAs a by-product of their analysis the authors also obtain improved bounds for sampling from log-concave distributions using the Langevin MCMC algorithm.\n\nThe paper introduces a new technique (the randomized midpoint method, originally developed for log-concave sampling via MCMC algorithms) to the area of generative diffusion models.", "weaknesses": "The paper does not include any numerical simulations.  Numerical simulations would help evaluate whether their algorithm works well in practice (both in terms of runtime and in terms of sample generation quality).\n\nAlso, while the authors adapt a new technique to the generative diffusion models, this is still an existing technique, but was used in a different area (MCMC for log-concave sampling).  That being said, I still feel the technical contribution is good since the authors introduce an existing technique to a new area.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "7d4W80AoGa", "reviewer": "Reviewer_g31r", "summary": "This paper studies diffusion models and introduces a faster sampling algorithm utilizing the randomized midpoint method within a predictor-corrector framework. In the predictor step, the paper introduces the randomized midpoint method, which allows for bounding the squared displacement into “bias” and“variance”terms. The corrector step follows the analysis in Chen et al.,2023b. This approach reduces discretization error, leading to an algorithm with improved theoretical dependency on the dimension $d$. Similar analysis can also be used in the setting of log-concave sampling in total variation distance, achieving a better dependency in $d$.", "questions": "1. As claimed in [1], the linear dependency on $d$ is optimal for KL divergence, thus $\\sqrt{d}$ for TV. Can you remark on this compared with the better than $\\sqrt{d}$ result in this paper?\n\n2. In the proof of Theorem A.10, given the value of the hyperparameters in Line 1109, can you explain more about how you can get the iteration complexity? Is it max $\\\\{1/h_{pred}, 1/h_{corr}\\\\}$? Then why the dependency on $d$ is not $d^{17/36}$, which is larger than $d^{5/12}$?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The organization is clear, highlighting the main contribution.\n\n2. The analysis method is very novel, and it’s an interesting result to see that the dimensional dependency for sampling can be improved to be better than $\\tilde O(\\sqrt{d})$.\n\n3. The paper shows that the analysis can also be extended to parallel algorithms.", "weaknesses": "1.The assumptions are a little strong.  It has been shown in [1][2] that the assumption of Lipschitz score can be relaxed. Moreover, it is often enough to assume that the true score is Lipschitz (Assumption 2.2). However, in this paper, the estimated score is also required to be Lipschitz. Although I don’t think it’s a big problem, can the authors remark about where this additional assumption is used?\n\n2.There are issues with the proof of Lemma A.2. The argument in (Line 736 equation 17) holds only under pointwise estimation (the estimated score is accurate for any $x$). However, in Assumption 2.4, it is defined as the expectation. Similar conditions occur in Line 755. As all the lemmas after this depend on Lemma A.2, the main result is correct only with a stronger assumption, which is not reflected in the paper. \n\n3.It misses some background knowledge and discussion:\n\n(1) The paper discusses log-concave sampling in section 3.6. However, there is no background knowledge in this paper. It is unclear what is the exact setting. Thus, the analysis in Appendix C is confusing. \n\n(2) In the main content, all the theorems are informal. Therefore, in the proof technique sections, like Sections 3.3, and 3.4, the target of the argument is ambiguous. I think at least a formal version of the theorem should be provided, and more details of the proof should be provided to help readers understand.\n\n4.Some typos: I think Algorithms 1, 2, and 3 are exactly the same as Algorithms 4,5,6 in the appendix. Thus, the reference in Algorithm 3 Line 3 should be Algorithm 1 instead of Algorithm 4, right?\n\nI'm willing to increase my score if the authors can solve my concerns, especially weakness 2.\n\n----\n[1] Benton et al. 2024, Nearly $d$-Linear Convergence Bounds for Diffusion Models via Stochastic Localization, ICLR2024\n\n[2] Chen et al. 2023, Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions, ICML2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)