Faster Diffusion Sampling with Randomized Midpoints: Sequential and Parallel

Sampling algorithms play an important role in controlling the quality and runtime of diffusion model inference. In recent years, a number of works~citechen2023sampling,chen2023ode,benton2023error,lee2022convergence have proposed schemes 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~citeShenL19. We prove that this approach achieves the best known dimension dependence for sampling from arbitrary smooth distributions in total variation distance (widetilde O(d5/12) compared to widetilde O(√d) from prior work). We also show that our algorithm can be parallelized to run in only widetilde O(log2 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(d5/12) compared to widetilde O(√d) from prior work.