[Paper Review] Parallel Sampling of Diffusion Models

Andy Y. Shih, Suneel Belkhale+3 more|arXiv (Cornell University)|May 25, 2023

Advanced Neuroimaging Techniques and Applications9 citations

TL;DR

ParaDiGMS introduces parallelized denoising via Picard iterations to accelerate diffusion model sampling, achieving 2–4x speedups with no measurable quality loss across robotics and image-generation tasks.

ABSTRACT

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

Motivation & Objective

• Motivate faster diffusion-model sampling without sacrificing sample quality by trading compute for speed.
• Develop a general parallel sampling framework compatible with pretrained diffusion models and existing fast samplers.
• Demonstrate speedups across robotics policies and image-generation models while preserving metrics like reward, FID, and CLIP.
• Show compatibility with prior fast-sampling methods (DDIM, DPMSolver) and classifier-free guidance.

Proposed method

• Use Picard iterations to solve the reverse-time diffusion ODE by iteratively refining a guess of the full denoising trajectory.
• Discretize the Picard update to enable parallel computation of drifts s(x_i, t_i) across time steps, followed by prefix-sum aggregation.
• Introduce a sliding-window minibatch strategy to fit GPU memory and a stopping criterion ensuring total-variation distance guarantees.
• Absorb fixed forward-noise realizations into a deterministic ODE when extending to reverse SDEs.
• Demonstrate compatibility by substituting Euler or higher-order solvers within ParaDiGMS and integrating with ParaDDPM, ParaDDIM, and ParaDPMSolver.
• Employ tolerance-based stopping criteria to bound distributional distance from the sequential baseline.

Experimental results

Research questions

• RQ1 Can diffusion-model sampling be accelerated by parallelizing denoising steps without sacrificing sample quality?
• RQ2 How does Picard-iteration-based parallel sampling compare with existing fast samplers (DDIM, DPMSolver) in terms of speed and quality?
• RQ3 Is the approach compatible with classifier-free guidance and prior speedups, and can it generalize to both image and robotic-policy tasks?

Key findings

• ParaDiGMS achieves about 2–4x speedups across robotics and image-generation models without measurable degradation in task reward, FID, or CLIP score.
• On DiffusionPolicy with 100-step sampling, time reduced from 0.74s to 0.2s per sample for a 100-step action generation task (example provided).
• On StableDiffusion-v2 with 1000-step sampling, time reduced from 50.0s to 14.6s per sample with no loss in quality.
• ParaDiGMS is compatible with DDPM, DDIM, and DPMSolver, enabling combined speedups (ParaDDPM, ParaDDIM, ParaDPMSolver) across tasks.
• In latent-space diffusion for StableDiffusion-v2, net speedups exceed 3x when using multi-GPU setups, with speedups depending on batch window size and hardware.
• ParaDiGMS maintains sample quality as measured by CLIP score on COCO captions, across several configurations.