SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, Nima Anari  
Computer Science, Stanford University  
{andyshih, belkhale, ermon, dorsa, anari}@cs.stanford.edu
Diffusion models are powerful generative models but suffer from slow sampling, often taking 1000 sequential denoising steps for one sample. As a result, considerable efforts have been directed toward reducing the number of denoising steps, but these methods hurt sample quality. Instead of reducing the number of denoising steps (trading quality for speed), in this paper we explore an orthogonal approach: can we run the denoising steps in parallel (trading compute for speed)? In spite of the sequential nature of the denoising steps, we show that surprisingly it is possible to parallelize sampling via Picard iterations, by guessing the solution of future denoising steps and iteratively refining until convergence. With this insight, we present ParaDiGMS, a novel method to accelerate the sampling of pretrained diffusion models by denoising multiple steps in parallel. ParaDiGMS is the first diffusion sampling method that enables trading compute for speed and is even compatible with existing fast sampling techniques such as DDIM and DPM-Solver. Using ParaDiGMS, we improve sampling speed by 2-4x across a range of robotics and image generation models, giving state-of-the-art sampling speeds of 0.2s on 100-step DiffusionPolicy and 14.6s on 1000-step StableDiffusion-v2 with no measurable degradation of task reward, FID score, or CLIP score. $^{1}$
Diffusion models [28, 9, 32] have demonstrated powerful modeling capabilities for image generation [34, 15, 24], molecular generation [38], robotic policies [12, 4], and other applications. The main limitation of diffusion models, however, is that sampling can be inconveniently slow. For example, the widely-used Denoising Diffusion Probabilistic Models (DDPMs) [9] can take 1000 denoising steps to generate one sample. In light of this, many works like DDIM [29] and DPMSolver [18] have proposed to improve sampling speed by reducing the number of denoising steps. Unfortunately, reducing the number of steps can come at the cost of sample quality.
We are interested in accelerating sampling of pretrained diffusion models without sacrificing sample quality. We ask the following question: rather than trading quality for speed, can we instead trade compute for speed? That is, could we leverage additional (parallel) compute to perform the same number of denoising steps faster? At first, this proposal seems unlikely to work, since denoising proceeds sequentially. Indeed, naive parallelization can let us generate multiple samples at once (improve throughput), but generating a single sample with faster wall-clock time (improving latency) appears much more difficult.
We show that, surprisingly, it is possible to improve the sample latency of diffusion models by computing denoising steps in parallel. Our method builds on the idea of Picard iterations to guess the full denoising trajectory and iteratively refine until convergence. Empirically, we find that the number of iterations for convergence is much smaller than the number of steps. Therefore, by computing each iteration quickly via parallelization, we sample from the diffusion model much faster.
$^{1}$ Code for our paper can be found at https://github.com/AndyShih12/paradigms
Our method ParaDiGMS (Parallel Diffusion Generative Model Sampling) is the first general method that allows for the tradeoff between compute and sampling speed of pretrained diffusion models. Remarkably, ParaDiGMS is compatible with classifier-free guidance [10] and with prior fast sampling methods [29, 18] that reduce the number of denoising steps. In other words, we present an orthogonal solution that can form combinations with prior methods (which we call ParaDDPM, ParaDDIM, ParaDPMSolver) to trade both compute and quality for speed.
We experiment with ParaDiGMS across a large range of robotics and image generation models, including Robosuite Square, PushT, Robosuite Kitchen, StableDiffusion-v2, and LSUN. Our method is strikingly consistent, providing an improvement across all tasks and all samplers (ParaDDPM, ParaDDIM, ParaDPMSolver) of around 2-4x speedup with no measurable decrease in quality on task reward, FID score, or CLIP score. For example, we improve the sample time of the 100-step action generation of DiffusionPolicy from 0.74s to 0.2s, and the 1000-step image generation of StableDiffusion-v2 on A100 GPUs from 50.0s to 14.6s. Our improvements also extend to few-step image generation, showing speedups for as low as 50-step DDIM. By enabling these faster sampling speeds without quality degradation, ParaDiGMS can enhance exciting applications of diffusion models such as real-time execution of diffusion policies or interactive generation of images.
Diffusion models [28, 9] such as Denoising Diffusion Probabilistic Models (DDPM) were introduced as latent-variable models with a discrete-time forward diffusion process where  $q(\pmb{x}_0)$  is the data distribution,  $\alpha$  is a scalar function, with latent variables  $\{\pmb{x}_t : t \leq T\}$  defined as
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {0}\right) = \mathcal {N} \left(\boldsymbol {x} _ {t}; \sqrt {\alpha (t)} \boldsymbol {x} _ {0}, (1 - \alpha (t)) \boldsymbol {I}\right).
$$
By setting  $\alpha (T)$  close to  $0$ ,  $q(\pmb{x}_T)$  converges to  $\mathcal{N}(\mathbf{0},\pmb{I})$ , allowing us to sample data  $\pmb{x}_0$  by using a standard Gaussian prior and a learned inference model  $p_{\theta}(\pmb{x}_{t - 1}\mid \pmb{x}_t)$ . The inference model  $p_{\theta}$  is parameterized as a Gaussian with predicted mean and time-dependent variance  $\sigma_t^2$ , and can be used to sample data by sequential denoising, i.e.,  $p_{\theta}(\pmb{x}_0) = \prod_{t = 1}^{T}p_{\theta}(\pmb{x}_{t - 1}\mid \pmb{x}_t)$ ,
$$
p _ {\theta} \left(\boldsymbol {x} _ {t - 1} \mid \boldsymbol {x} _ {t}\right) = \mathcal {N} \left(\boldsymbol {x} _ {t - 1}; \mu_ {\theta} \left(\boldsymbol {x} _ {t}\right), \sigma_ {t} ^ {2} \boldsymbol {I}\right). \tag {1}
$$
Many works [32, 18] alternatively formulate diffusion models as a Stochastic Differential Equation (SDE) by writing the forward diffusion process in the form
$$
\mathrm {d} \boldsymbol {x} _ {t} = f (t) \boldsymbol {x} _ {t} \mathrm {d} t + g (t) \mathrm {d} \boldsymbol {w} _ {t}, \quad \boldsymbol {x} _ {0} \sim q (\boldsymbol {x} _ {0}), \tag {2}
$$
with the standard Wiener process  $\boldsymbol{w}_t$ , where  $f$  and  $g$  are position-independent functions that can be appropriately chosen to match the transition distribution  $q(\boldsymbol{x}_t \mid \boldsymbol{x}_0)$  [18, 15]. These works use an important result from [2] that the reverse process of Eq. (2) takes on the form
$$
\mathrm {d} \boldsymbol {x} _ {t} = \underbrace {\left(f (t) \boldsymbol {x} _ {t} - g ^ {2} (t) \nabla_ {\boldsymbol {x}} \log q _ {t} (\boldsymbol {x})\right)} _ {\text {d r i f t} s} \mathrm {d} t + \underbrace {g (t)} _ {\sigma_ {t}} \mathrm {d} \bar {\boldsymbol {w}} _ {t}, \quad \boldsymbol {x} _ {T} \sim q (\boldsymbol {x} _ {T}), \tag {3}
$$
where  $\bar{\boldsymbol{w}}_t$  is the standard Wiener process in reverse time. This perspective allows us to treat the sampling process of DDPM as solving a discretization of the SDE where the DDPM inference model  $p_{\theta}$  can be used to compute an approximation  $p_{\theta}(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t) - \boldsymbol{x}_t$  of the drift term in Eq. (3).
Since the focus of this paper is on sampling from a pretrained diffusion model, we can assume  $p_{\theta}$  is given. For our purposes, we only need two observations about sampling from the reverse process in Eq. (3): we have access to an oracle that computes the drift at any given point, and the SDE has position-independent noise. We will use the latter observation in Section 3.
2.1 Reducing the number of denoising steps
DDPM typically uses a  $T = 1000$  step discretization of the SDE. These denoising steps are computed sequentially and require a full pass through the neural network  $p_{\theta}$  each step, so sampling can be extremely slow. As a result, popular works such as DDIM [29] and DPMSolver [18] have explored the possibility of reducing the number of denoising steps, which amounts to using a coarser discretization with the goal of trading sample quality for speed.
Empirically, directly reducing the number of steps of the stochastic sampling process of DDPM hurts sample quality significantly. Therefore many works [29, 32, 18] propose using an Ordinary Differential Equation (ODE) to make the sampling process more amenable to low-step methods. These works appeal to the probability flow ODE [20], a deterministic process with the property that the marginal distribution  $p(\pmb{x}_t)$  at each time  $t$  matches that of the SDE, so in theory sampling from the probability flow ODE is equivalent to sampling from the SDE:
$$
\mathrm {d} \boldsymbol {x} _ {t} = \underbrace {\left(f (t) \boldsymbol {x} _ {t} - \frac {1}{2} g ^ {2} (t) \nabla_ {\boldsymbol {x}} \log q _ {t} (\boldsymbol {x})\right)} _ {\text {d r i f t} s} \mathrm {d} t, \quad \boldsymbol {x} _ {T} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}).
$$
By sampling from the ODE instead of the SDE, works such as DDIM and DPMSolver (which have connections to numerical methods such as Euler and Heun) can reduce the quality degradation of few-step sampling (e.g., 50 steps).
As a summary, the current landscape of sampling from pretrained diffusion models is comprised of full-step DDPM or accelerated sampling techniques such as DDIM and DPMSolver that trade quality for speed by reducing the number of denoising steps.
Notation We write  $[a, b]$  to denote the set  $\{a, a + 1, \ldots, b\}$  and  $[a, b)$  to denote the set  $\{a, a + 1, \ldots, b - 1\}$  for  $b > a$ . We write  $\pmb{x}_{a:b}$  to denote the set  $\{\pmb{x}_i : i \in [a, b)\}$ . Since our focus is on sampling, in the rest of the paper we denote time as increasing for the reverse process.
3 Parallel computation of denoising steps
Rather than investigating additional techniques for reducing the number of denoising steps, which can lead to quality degradation, we look towards other approaches for accelerating sampling. In particular, we investigate the idea of trading compute for speed: can we accelerate sampling by taking denoising steps in parallel? We clarify that our goal is not to improve sample throughput – that can be done with naive parallelization, producing multiple samples at the same time. Our goal is to improve sample latency – minimize the wall-clock time required for generating a single sample by solving the denoising steps for a single sample in parallel. Lowering sample latency without sacrificing quality can greatly improve the experience of using diffusion models, and enable more interactive and real-time generation applications.
Parallelizing the denoising steps, however, seems challenging due to the sequential nature of existing sampling methods. The computation graph has a chain structure (Fig. 1), so it is difficult to propagate information quickly down the graph. To make headway, we present the method of Picard iteration, a technique for solving ODEs through fixed-point iteration. An ODE is defined by a drift function  $s(\boldsymbol{x}, t)$  with position and time arguments, and initial value  $\boldsymbol{x}_0$ . In the integral form, the value at time  $t$  can be written as
$$
\boldsymbol {x} _ {t} = \boldsymbol {x} _ {0} + \int_ {0} ^ {t} s (\boldsymbol {x} _ {u}, u) d u.
$$
In other words, the value at time  $t$  must be the initial value plus the integral of the derivative along the path of the solution. This formula suggests a natural way of solving the ODE by starting with a guess of the solution  $\{x_{t}^{k}:0\leq t\leq 1\}$  at initial iteration  $k = 0$ , and iteratively refining by updating the value at every time  $t$  until convergence:
$$
\text {(P i c a r d I t e r a t i o n)} \quad \boldsymbol {x} _ {t} ^ {k + 1} = \boldsymbol {x} _ {0} ^ {k} + \int_ {0} ^ {t} s \left(\boldsymbol {x} _ {u} ^ {k}, u\right) d u. \tag {4}
$$
Under mild conditions on  $s$ , such as continuity in time and Lipschitz continuity in position as in the well-known Picard-Lindelöf theorem, the iterates form a convergent sequence, and by the Banach fixed-point theorem, they converge to the unique solution of the ODE with initial value  $x_0$  [cf. 5]. To perform Picard iterations numerically, we can write the discretized form of Eq. (4) with step size  $1 / T$ , for  $t \in [0,T]$ :
$$
\boldsymbol {x} _ {t} ^ {k + 1} = \boldsymbol {x} _ {0} ^ {k} + \frac {1}{T} \sum_ {i = 0} ^ {t - 1} s \left(\boldsymbol {x} _ {i} ^ {k}, i / T\right). \tag {5}
$$
Figure 1: Computation graph of sequential sam- Figure 2: Computation graph of Picard iterations, pling by evaluating  $p_{\theta}(\boldsymbol{x}_{t + 1}\mid \boldsymbol{x}_t)$  , from the per- which introduces skip dependencies.   
spective of reverse time.
Examining the iterative update rule in Eq. (5), we see that an update at time  $t$  depends on all previous timesteps instead of just the previous timestep  $t - 1$ . This amounts to introducing skip dependencies in the computation graph (Fig. 2), which enables information to propagate quickly down the chain and accelerate sampling.
The key property of interest is that each Picard iteration can be parallelized by performing the expensive computations  $\{s(\pmb{x}_i^k,\frac{i}{T}):i\in [0,T)\}$  in parallel and then, with negligible cost, collecting their outputs into prefix sums. Given enough parallel processing power, the sampling time then scales with the number of iterations  $K$  until convergence, instead of the number of denoising steps  $T$ .
The number of iterations until convergence depends on the drift function  $s$ . More concretely, sequential evaluation can be written as a nested evaluation of functions  $\boldsymbol{x}_{t+1}^{\star} = h_t(\ldots h_2(h_1(\boldsymbol{x}_0)))$  on the initial value  $\boldsymbol{x}_0$  where  $h_i(\boldsymbol{x}) = \boldsymbol{x} + s(\boldsymbol{x}, i/T)/T$ . If, for all timesteps, the drift at the true solution can be accurately obtained using the drift at the current guess, then the parallel evaluation will converge in one step.
Proposition 1. (Proof in Appendix A)
$$
s \left(\boldsymbol {x} _ {i} ^ {k}, i / T\right) = s \left(h _ {i - 1} \left(\dots h _ {2} \left(h _ {1} \left(\boldsymbol {x} _ {0}\right)\right)\right), i / T\right) \quad \forall i \leq t \Longrightarrow \boldsymbol {x} _ {t + 1} ^ {k + 1} = \boldsymbol {x} _ {t + 1} ^ {\star}
$$
It is also easy to see that even in the worst case, exact convergence happens in  $K \leq T$  iterations since the first  $k$  points  $x_{0:k}$  must equal the sequential solution  $x_{0:k}^{\star}$  after  $k$  iterations. In practice, the number of iterations until (approximate) convergence is typically much smaller than  $T$ , leading to a large empirical speedup.
The idea of Picard iterations is powerful because it enables the parallelization of denoising steps. Remarkably, Picard iterations are also fully compatible with prior methods for reducing the number of denoising steps. Recall that the drift term  $s(\pmb{x}_t, t / T) / T$  can be written as  $h_t(\pmb{x}_t) - \pmb{x}_t$  and approximated using Euler discretization as  $p_{\theta}(\pmb{x}_{t + 1} \mid \pmb{x}_t) - \pmb{x}_t$ , but it can also be readily approximated using higher-order methods on  $p_{\theta}$ . In our experiments, we demonstrate the combination of the two axis of speedups to both reduce the number of denoising steps and compute the steps in parallel.
3.1 Practical considerations
Implementing Picard iteration on diffusion models presents a few practical challenges, the most important being that of GPU memory. Performing an iteration requires maintaining the entire array of points  $\boldsymbol{x}_{0:T}$  over time, which can be prohibitively large to fit into GPU memory. To address this, we devise the technique of (mini-)batching which performs Picard iteration only on points  $\boldsymbol{x}_{t:t+p}$  inside a window of size  $p$  that can be chosen appropriately to satisfy memory constraints. Moreover, instead of iterating on  $\boldsymbol{x}_{t:t+p}$  until convergence of the full window before advancing to the next window, we use a sliding window approach to aggressively shift the window forward in time as soon as the starting timesteps in the window converge.
One other issue is the problem of extending Picard iteration to SDEs, since we rely on the determinism of ODEs to converge to a fixed point. Fortunately, since the reverse SDE (Eq. (3)) has position-independent noise, we can sample the noise up-front and absorb these fixed noises into the drift of the (now deterministic) differential equation. Note that the resulting ODE is still Lipschitz continuous in position and continuous in time, guaranteeing the convergence of Picard iteration.
Figure 3: ParaDiGMS algorithm: accelerating an ODE solver by computing the drift at multiple timesteps in parallel. During iteration  $k$ , we process in parallel a batch window of size  $p$  spanning timesteps  $[t, t + p)$ . The new values at a point  $x_{t + j}^{k + 1}$  are updated based on the value  $x_{j}^{k}$  at the left end of the window plus the cumulative drift  $1 / T\sum_{i = t}^{t + j - 1}s(\pmb{x}_i^k,i / T)$  of points in the window. We then slide the window forward until the error is greater than our tolerance, and repeat for the next iteration.
Finally, we need to choose a stopping criterion for the fixed-point iteration, picking a low tolerance to avoid degradations of sample quality. A low enough tolerance ensures that the outcome of parallel sampling will be close to the outcome of the sequential sampling process in total variation distance.
Proposition 2. (Proof in Appendix B) Assuming the iteration rule in Eq. (5) has a linear convergence rate with a factor  $\geq 2$ , using the tolerance  $\| \pmb{x}_t^K - \pmb{x}_t^{K-1} \|^2 \leq 4\epsilon^2 \sigma_t^2 / T^2$  ensures that samples of  $\pmb{x}_T^K$  are drawn from a distribution with total variation distance at most  $\epsilon$  from the DDPM model distribution of Eq. (1).
The above is based on a worst-case analysis, and in our experiments, we find that using a much more relaxed tolerance such as  $\frac{1}{D} \| \pmb{x}_i^{k+1} - \pmb{x}_i^k \|^2 \leq \tau^2 \sigma_i^2$ , with  $\tau = 0.1$  and  $D$  being the dimensionality of data, gives reliable speedups without any measurable degradation in sample quality.
In Algorithm 1 we present the complete procedure of ParaDiGMS, incorporating sliding window over a batch, up-front sampling of noise, and tolerance of Picard iterations (Fig. 3). The loop starting on Line 4 performs a sliding window over the batch of timesteps  $[t, t + p)$  in each iteration. Line 5 computes the drifts, which is the most compute-intensive part of the algorithm, but can be efficiently parallelized. Line 6 obtains their prefix sums in parallel to run the discretized Picard iteration update, and Lines 7-8 check the error values to determine how far forward we can shift the sliding window.
Algorithm 1: ParaDiGMS: parallel sampling via Picard iteration over a sliding window
Input: Diffusion model  $p_{\theta}$  with variances  $\sigma_t^2$  , tolerance  $\tau$  , batch window size  $p$  , dimension  $D$  Output: A sample from  $p_{\theta}$    
1  $t,k\gets 0,0$    
2  $\mathbf{z}_i\sim \mathcal{N}(\mathbf{0},\sigma_i^2\mathbf{I})\quad \forall i\in [0,T)$  // Up-front sampling of noise (for SDE)   
3  $\pmb{x}_0^k\sim \mathcal{N}(\pmb {0},\pmb {I}),\qquad \pmb{x}_i^k\gets \pmb{x}_0^k\quad \forall i\in [1,p]$  // Sample initial condition from prior   
4 while  $t <   T$  do   
5  $\pmb{y}_{t + j}\gets p_{\theta}(\pmb{x}_{t + j}^{k},t + j) - \pmb{x}_{t + j}^{k}\quad \forall j\in [0,p)$  // Compute drifts in parallel   
6  $\pmb{x}_{t + j + 1}^{k + 1}\leftarrow \pmb{x}_t^k +\sum_{i = t}^{t + j}\pmb {y}_i + \sum_{i = t}^{t + j}\pmb {z}_i\quad \forall j\in [0,p)$  // Discretized Picard iteration   
7 error  $\leftarrow \{\frac{1}{D}\| \pmb{x}_{t + j}^{k + 1} - \pmb{x}_{t + j}^{k}\|^{2}:\forall j\in [1,p)\}$  // Store error value for each timestep   
8 stride  $\leftarrow \min \left(\{j:error_j > \tau^2\sigma_j^2\} \cup \{p\}\right)$  // Slide forward until we reach tolerance   
9  $\pmb{x}_{t + p + j}^{k + 1}\leftarrow \pmb{x}_{t + p}^{k + 1}\quad \forall j\in [1,\text{stride} ]$  // Initialize new points that the window now covers   
10  $t\gets t +$  stride,  $k\gets k + 1$    
11  $p\gets \min (p,T - t)$    
12 return  $x_{T}^{k}$
The ParaDiGMS algorithm is directly compatible with existing fast sequential sampling techniques such as DDIM and DPMSolver, by swapping out the Euler discretization in Lines 5-6 for other
2For ODE methods (DDIM, DPMSolver) we still pick a tolerance value relative to the noise variance of the corresponding SDE of DDPM.
solvers, such as higher-order methods like Heun. As we see in our experiments, the combination of reducing the number of steps and solving the steps in parallel leads to even faster sample generation.
We experiment with our method ParaDiGMS on a suite of robotic control tasks [4] including Square [41], PushT, Franka Kitchen [7], and high-dimensional image generation models including StableDiffusion-v2 [24] and LSUN Church and Bedroom [39]. We observe a consistent improvement of around 2-4x speedup relative to the sequential baselines without measurable degradation in sample quality as measured by task reward, FID score, or CLIP score.
4.1 Diffusion policy
Recently, a number of works have demonstrated the advantages of using diffusion models in robotic control tasks for flexible behavior synthesis or robust imitation learning on multimodal behavior [12, 4, 22, 36]. We follow the setup of DiffusionPolicy [4], which is an imitation learning framework that models action sequences. More specifically, DiffusionPolicy first specifies a prediction horizon  $h$  and a replanning horizon  $r$ . At each environment step  $l$ , DiffusionPolicy conditions on a history of observations and predicts a sequence of actions  $\{a_{l;l + h}\}$ . Then, the policy executes the first  $r$  actions  $\{a_{l;l + r}\}$  of the prediction. Therefore, for an episode of length  $L$  and scheduler with  $T$  steps, executing a full trajectory can take  $T\times L / r$  denoising steps over a dimension of  $|a|\times h$ , which can be inconveniently slow.
We examine our method on the Robosuite Square, PushT, and Robosuite Kitchen tasks. Each environment uses a prediction horizon of 16, and replanning horizon 8. The Square task uses state-based observations with a maximum trajectory length of 400 and a position-based action space of dimensionality 7. This means the diffusion policy takes 50 samples per episode, with each sample being a series of denoising steps over a joint action sequence of dimension 112. The PushT task also uses state-based observations and has a maximum trajectory length of 300 and action space of 2, which results in 38 samples with denoising steps over a joint action sequence of dimension 32. Lastly, the Kitchen task uses vision-based observations and has a maximum trajectory length of 1200 with an action space of 7, giving 150 samples per episode and denoising steps over a joint action sequence of dimensionality 112. For all three tasks, we use a convolution-based architecture.
The DDPM scheduler in DiffusionPolicy [4] uses 100 step discretization, and the DDIM/DPMSolver schedulers use 15 step discretization. For example, a trajectory in the Kitchen task requires  $1200 / 8 = 150$  samples, which amounts to  $150 \times 100 = 15000$  denoising steps over an action sequence of dimensionality 112 with the DDPM scheduler.