Self-Refining Diffusion Samplers: Enabling Parallelization via Parareal Iterations

TL;DR

The paper introduces Self-Refining Diffusion Samplers (SRDS), a parallelizable method inspired by Parareal, that accelerates diffusion model sampling without sacrificing quality by iteratively refining initial rough estimates in parallel.

Contribution

SRDS leverages Parareal iterations to enable parallel diffusion sampling, reducing generation time while maintaining sample quality, unlike prior approximation methods.

Findings

Speedup of up to 1.7x on 25-step models

Speedup of up to 4.3x on longer trajectories

Guarantees convergence to serial solution

Abstract

In diffusion models, samples are generated through an iterative refinement process, requiring hundreds of sequential model evaluations. Several recent methods have introduced approximations (fewer discretization steps or distillation) to trade off speed at the cost of sample quality. In contrast, we introduce Self-Refining Diffusion Samplers (SRDS) that retain sample quality and can improve latency at the cost of additional parallel compute. We take inspiration from the Parareal algorithm, a popular numerical method for parallel-in-time integration of differential equations. In SRDS, a quick but rough estimate of a sample is first created and then iteratively refined in parallel through Parareal iterations. SRDS is not only guaranteed to accurately solve the ODE and converge to the serial solution but also benefits from parallelization across the diffusion trajectory, enabling batched inference and pipelining. As we demonstrate for pre-trained diffusion models, the early convergence of this refinement procedure drastically reduces the number of steps required to produce a sample, speeding up generation for instance by up to 1.7x on a 25-step StableDiffusion-v2 benchmark and up to 4.3x on longer trajectories.