Faster Diffusion Sampling with Randomized Midpoints: Sequential and Parallel

Our discretization scheme is based on the randomized midpoint method of [SL19], which has been studied at length in the domain of log-concave sampling [HBE20, YKD23, YD24].

The proof of our parallel sampling result builds on the ideas of [SL19, AHL⁺23, ACV24] on parallelizing the collocation method. These prior results were focused on Langevin Monte Carlo, rather than diffusion-based sampling. We review these ideas in Section 2.3.

In [CCL⁺23a], the authors proposed the predictor-corrector framework that we also use for analysing convergence guarantee of the probability flow ODE and which achieved iteration complexity scaling with $\widetilde{O}(\sqrt{d})$. In addition to this, there have been many works in recent years giving general convergence guarantees for diffusion models [DBTHD21, BMR22, DB22, LLT22, LWYL22, Pid22, WY22, CCL⁺23b, CDD23, LLT23, LWCC23, BDD23, CCL⁺23a, BDBDD24, CLL23, GPPX23]. Of these, one line of work [CCL⁺23b, LLT23, CLL23, BDBDD24] analyzed DDPM, the stochastic analogue of the probability flow ODE, and showed $\tilde{O}(d)$ iteration complexity bounds. Another set of works [CCL⁺23a, CDD23, LWCC23, LHE⁺24] 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-art $\widetilde{O}(d^{5/12})$ dimension dependence for sampling from a diffusion model.

In this section we review basics about deterministic diffusion-based samplers; we refer the reader to [CCL⁺23a] for a more thorough exposition.

Let $q^{*}$ denote the data distribution over $\mathbb{R}^{d}$. We consider the standard Ornstein-Uhlenbeck (OU) forward process, i.e. the “VP SDE,” given by

where $(B_{t})_{t\geq 0}$ denotes a standard Brownian motion in $\mathbb{R}^{d}$. This process converges exponentially quickly to its stationary distribution, the Gaussian distribution $\mathcal{N}(0,\mathrm{Id})$.

Suppose the OU process is run until terminal time $T>0$, and for any $t\in[0,T]$, let $q^{*}_{t}\triangleq\text{law}(x^{\rightarrow}_{t})$, i.e. the law of the forward process at time $t$. We will consider the reverse process given by the probability flow ODE

This is a time-reversal of the forward process, so that if $x_{0}\sim q_{T}$, then $\text{law}(x_{t})=q^{*}_{T-t}$. In practice, one initializes at $x_{0}\sim\mathcal{N}(0,\mathrm{Id})$, and instead of using the exact score function $\nabla\ln q_{T-t}$, one uses estimates $\widehat{s}_{T-t}\approx\nabla\ln q_{T-t}$ 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.

Suppose we wish to discretize the following semilinear ODE:

For our application we will eventually take $f_{t}\triangleq\widehat{s}_{T-t}$, but we use $f_{t}$ in this section to condense notation.

Suppose we want to discretize Equation 3 over a time window $[t_{0},t_{0}+h]$. The starting point is the integral formulation for this ODE:

Under the standard exponential integrator discretization, one would approximate the integrand by $e^{t_{0}+h-t}f_{t_{0}}(x_{t_{0}})$ and obtain the approximation

The drawback of this discretization is that it uses an inherently biased estimate for the integral in Eq. (4). The key insight of [SL19] was to replace this with the following unbiased estimate

where $\alpha$ is a uniformly random sample from $[0,1]$. While this alone does not suffice as the estimate depends on $x_{t_{0}+\alpha h}$, naturally we could iterate the above procedure again to obtain an approximation to $x_{t_{0}+\alpha h}$. It turns out though that even if we simply approximate $x_{t_{0}+\alpha h}$ 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:

Note that a similar idea can be used to discretize stochastic differential equations, but in this work we only use it to discretize the probability flow ODE.

The scheme outlined in the previous section is a simple special case of the collocation method. In the context of the semilinar ODE from Eq. (3), the idea behind the collocation method is to solve the integral formulation of the ODE in Eq. (4) via fixed point iteration. For our parallel sampling guarantees, instead of choosing a single randomized midpoint $\alpha$, we break up the window $[t_{0},t_{0}+h]$ into $R$ sub-windows, select randomized midpoints $\alpha_{1},\ldots,\alpha_{R}$ for these sub-windows, and approximate the trajectory of the ODE at any time $t_{0}+i\delta$, where $\delta\triangleq h/R$, by

One can show that as $R\to\infty$, this approximation tends to an equality. For sufficiently large $R$, Eq. (9) naturally suggests a fixed point iteration that can be used to approximate each $x_{t_{0}+\alpha_{i}h}$, i.e. we can maintain a sequence of estimates $\widehat{x}^{(k)}_{t_{0}+\alpha_{i}h}$ defined by the iteration

for $k$ ranging from $1$ up to some sufficiently large $K$. Finally, analogously to Eq. (8), we can estimate $x_{t_{0}+h}$ via

The key observation, made in [SL19] and also in related works of [ACV24, SBE⁺24, AHL⁺23], is that for any fixed round $k$, all of the iterations Eq. (10) for different choices of $i=1,\ldots,R$ can be computed in parallel. With $R$ parallel processors, one can thus compute the estimate for $x_{t_{0}+h}$ in $K$ parallel rounds, with $O(KR)$ total work.

Following the framework of [CCL⁺23a], our algorithm consists of “predictor” steps interspersed with “corrector” steps, with the time spent on each carefully tuned to obtain our final $\widetilde{O}(d^{5/12})$ dimension dependence. We first describe our predictor step – this is the piece of our algorithm that makes use of the Shen and Lee’s randomized midpoint method [SL19].

The main difference between the above and the predictor step of [CCL⁺23a] are steps $1(b)$ – $1(d)$. $1(b)$ and $1(c)$ together compute a randomized midpoint, and $1(d)$ uses this midpoint to obtain an approximate solution to the integral of the ODE. We describe these steps in more detail in Section 3.3.

Next, we describe the “corrector” step, introduced in [CCL⁺23a]. First, recall the underdamped Langevin ODE:

Here $\widehat{s}$ is our $L^{2}$ accurate score estimate for a fixed time (say $t$). Then, the corrector step is described below.

Finally, Algorithm 3 below puts the predictor and corrector steps together to give our final sequential algorithm.

For the final setting of parameters in Algorithm 3, see Theorem A.10. Now, we describe the analysis of the above algorithm in detail.

The general framework of our algorithm closely follows that of [CCL⁺23a], which proposed to run the (discretized) reverse ODE but interspersed with “corrector” steps given by running underdamped Langevin dynamics. The idea is that the “predictor” steps where the discretized reverse 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 [CCL⁺23a], for most of the reverse process the predictor steps (Step 2(a)) are run with a fixed step size, but at 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.

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 (Algorithm 1) and restrict our attention to running the reverse process for a small amount of time $h\ll 1/L$.

We begin by recalling the dimension dependence achieved by the standard exponential integrator scheme. One can show (see e.g. Lemma 4 in [CCL⁺23a]) 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 $O(d^{1/2}h^{2})$. By iterating this coupling $O(1/h)$ times, we conclude that in an $O(1)$ window of time, the processes drift by a distance of $O(d^{1/2}h)$. To ensure this is not too large, one would take the step size $h$ to be $O(1/\sqrt{d})$, thus obtaining an iteration complexity of $O(\sqrt{d})$ as in [CCL⁺23a].

The starting point in the analysis of randomized midpoint is to instead track the squared displacement between the two processes instead. Given two neighboring time steps $t-h$ and $t$ in the algorithm, let $x_{t}$ denote the true reverse process at time $t$, and let $\widehat{x}_{t}$ denote the algorithm at time $t$ (in the notation of Algorithm 2, this is $\widehat{x}_{n}$ for some $n$, but we use $t$ in the discussion here to make the comparison to the true reverse process clearer). Note that $\widehat{x}_{t}$ depends on the choice of randomized midpoint $\alpha$ (see Step 1(b)). One can bound the squared displacement $\operatorname*{\mathbb{E}}\,\left\|x_{t}-\widehat{x}_{t}\right\|^{2}$ as follows. Let $y_{t}$ be the result of running the reverse process for time $h$ starting from $\widehat{x}_{t-h}$. Then by writing $x_{t}-\widehat{x}_{t}$ as $(x_{t}-y_{t})-(\widehat{x}_{t}-y_{t})$ and applying Young’s inequality, we obtain

For the first term, because $x_{t}$ and $y_{t}$ are the result of running the same ODE on initializations $x_{t-h}$ and $\widehat{x}_{t-h}$ , the first term is close to $\operatorname*{\mathbb{E}}\|x_{t-h}-\widehat{x}_{t-h}\|^{2}$ provided $h\ll 1/L$. The upshot is that the squared displacement at time $t$ is at most the squared displacement at time $t-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 $O(d^{1/2}h^{2})$ displacement bound that arises in the exponential integrator analysis, giving $O(dh^{4})$:

Note that in this bound, in addition to the $O(dh^{4})$ term and a term for the score estimation error, there is an additional term which depends on the squared displacement from the previous time step. Because the prefactor $L^{2}h^{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^{4})$ with every time step of length $h$. Over $O(1/h)$ such steps, the total squared displacement is $O(dh^{3})$, so if we take the step size $h$ to be $O(1/d^{1/3})$, this suggests an improved iteration complexity of $O(d^{1/3})$.

Arguing that the bias term $\frac{2}{Lh}\operatorname*{\mathbb{E}}_{\widehat{x}_{t-h}}\left\|\operatorname*{\mathbb{E}}_{\alpha}\widehat{x}_{t}-y_{t}\right\|^{2}$ 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 so that it would give an 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 an approximation to the latter, but as we show, the error incurred by this is indeed of lower order (see Lemma A.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 trajectory. Bounding such quantities in expectation over the randomness of the algorithm’s trajectory can be difficult, but our proof identifies a way of “offloading” some of this difficulty by absorbing some excess terms into a term of the form $\|x_{t-h}-\widehat{x}_{t-h}\|^{2}$, i.e. the squared displacement from the previous time step. Concretely, we obtain the following bound on the bias term:

While we have outlined how to improve the predictor step in the framewok of [CCL⁺23a], it 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 [Che23].

We will sidestep this issue and argue that even using exponential integrator discretization of the underdamped Langevin dynamics will suffice for our purposes, by simply shortening the amount of time for which each corrector step is run.

First, let us briefly recall what was shown in [CCL⁺23a] for the corrector step. If one runs underdamped Langevin dynamics with stationary distribution $q$ for time $T$ and exponential integrator discretization with step size $h$ starting from two distributions $p$ and $q$, then the resulting distributions $p^{\prime}$ and $q$ satisfy

where $L$ is the Lipschitzness of $\nabla\ln q$ (see Theorem A.6). At first glance this appears insufficient for our purposes: because of the $d^{1/2}h$ term coming from the discretization error, we would need to take step size $h=1/\sqrt{d}$, which would suggest that the number of iterations must scale with $\sqrt{d}$.

To improve the dimension dependence for our overall predictor-corrector 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 $\frac{W_{2}(p,q)}{L^{1/4}T^{3/2}}$ in Equation 14. But in our overall proof, the $W_{2}(p,q)$ term is bounded by the predictor analysis. There, we had quite a bit of slack: even with step size as large as $1/d^{1/3}$, we could achieve small Wasserstein error. By balancing appropriately, we get our improved dimension dependence.

Now, we summarize the main proof ideas for our parallel sampling result. In Section 2.3, we described how to approximately solve the reverse ODE over time $h$ by running $K$ rounds of the iteration in Equation 10. In our final algorithm, we will take $h$ to be dimension-independent, namely $h=\Theta(1/\sqrt{L})$, 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 $t_{0}$ and now wish to solve the ODE from time $t_{0}$ to $t_{0}+h$. We initialize at $\{\widehat{x}^{(0)}_{t_{0}+\alpha_{i}h}\}_{i\in[R]}$ via exponential integrator steps starting from the beginning of the window – see Line 1(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 to the true reverse process contracts:

One shortcoming of the predictor analysis is that the contraction achieved by fixed point iteration ultimately bottoms out at error which scales with $d/R^{2}$ (see the second term in Equation 15). In order for the discretization error to be sufficiently small, we thus have to take $R$, and thus the total work of the algorithm, to scale with $O(\sqrt{d})$. So in this case we do not improve over the dimension dependence of [CCL⁺23a], and instead the improvement is in obtaining a parallel algorithm.

For the corrector analysis, we mostly draw upon the recent work of [ACV24] which analyzed a parallel implementation of the underdamped Langevin dynamics. While their guarantee focuses on sampling from log-concave distributions, implicit in their analysis is a bound for general smooth distributions on how much the law of the algorithm and the law of the true process drift apart in a bounded time window (see Lemma B.8). This bound suffices for our analysis of the corrector step, and we can conclude the following:

Overall, parallel algorithm is somewhat different from the parallel sampler developed in the empirical work of [SBE⁺24], even apart from the fact that we use randomized midpoint discretization and corrector steps. The reason is that our algorithm applies collocation to fixed windows of time, whereas the algorithm of [SBE⁺24] utilizes a sliding window approach that proactively shifts the window forward as soon as the iterates at the start of the previous window begin to converge. We leave rigorously analyzing the benefits of this approach as an interesting future direction.

Finally, we briefly summarize the simple proof for our result on log-concave sampling in TV, which achieves the best known dimension dependence of $\widetilde{O}(d^{5/12})$. Our main observation is that Shen and Lee’s randomized midpoint method [SL19] applied to the underdamped Langevin process gives a Wasserstein guarantee for log-concave sampling, while the corrector step of [CCL⁺23a] can convert a Wasserstein guarantee to closeness in TV. Thus, we can simply run the 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 C for the full proof.