Polynomial Speedup in Diffusion Models with the Multilevel Euler-Maruyama Method

TL;DR

This paper introduces the Multilevel Euler-Maruyama method, which significantly accelerates the solution of SDEs and diffusion models by leveraging multiple approximators, achieving polynomial speedups especially in high-accuracy regimes.

Contribution

The paper presents a novel multilevel approach for SDEs and diffusion models that reduces computational cost by efficiently combining multiple approximators, enabling faster sampling with large neural networks.

Findings

Up to fourfold speedup in image generation on CelebA dataset.

The method achieves polynomial speedup in the Harder than Monte Carlo regime.

Numerical experiments confirm theoretical speedup predictions.

Abstract

We introduce the Multilevel Euler-Maruyama (ML-EM) method compute solutions of SDEs and ODEs using a range of approximators $f^1,\dots,f^k$ to the drift $f$ with increasing accuracy and computational cost, only requiring a few evaluations of the most accurate $f^k$ and many evaluations of the less costly $f^1,\dots,f^{k-1}$. If the drift lies in the so-called Harder than Monte Carlo (HTMC) regime, i.e. it requires $\epsilon^{-\gamma}$ compute to be $\epsilon$-approximated for some $\gamma>2$, then ML-EM $\epsilon$-approximates the solution of the SDE with $\epsilon^{-\gamma}$ compute, improving over the traditional EM rate of $\epsilon^{-\gamma-1}$. In other terms it allows us to solve the SDE at the same cost as a single evaluation of the drift. In the context of diffusion models, the different levels $f^{1},\dots,f^{k}$ are obtained by training UNets of increasing sizes, and ML-EM allows us to perform sampling with the equivalent of a single evaluation of the largest UNet. Our numerical experiments confirm our theory: we obtain up to fourfold speedups for image generation on the CelebA dataset downscaled to 64x64, where we measure a $\gamma\approx2.5$. Given that this is a polynomial speedup, we expect even stronger speedups in practical applications which involve orders of magnitude larger networks.