用多级方法加速扩散模型采样,理论可实现多项式提速。
Polynomial Speedup in Diffusion Models with the Multilevel Euler-Maruyama Method
- 通过分级近似漂移项,减少高精度模型调用次数。
- 在γ≈2.5时,图像生成速度提升最高达4倍。
- 适合大模型扩散采样,尤其对超大规模网络效果更显著。
我们提出多级欧拉-马鲁雅米(ML-EM)方法,通过一系列精度递增、成本递增的漂移函数近似器 $f^1,\dots,f^k$ 求解随机微分方程(SDE)与常微分方程(ODE)。当漂移项处于“比蒙特卡洛更难”(HTMC)情形,即需 $ε^{-γ}$ 计算量才能达到 $ε$ 精度($γ>2$)时,ML-EM 仅需 $ε^{-γ}$ 计算即可实现 $ε$ 精度的解,优于传统欧拉-马鲁雅米方法的 $ε^{-γ-1}$ 复杂度。这意味着可在等价于单次高精度漂移评估的成本下求解整个 SDE。在扩散模型中,各层级 $f^1,\dots,f^k$ 对应逐步增大尺寸的 UNet,ML-EM 实现了等效于一次最大尺寸 UNet 的采样效率。数值实验验证了理论:在 64x64 缩放的 CelebA 数据集上,速度提升最高达 4 倍,此时 $γ\approx2.5$。由于这是多项式加速,实际应用中使用更大规模网络时预期提速更强。
原文摘要 · Abstract (English)
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 $ε^{-γ}$ compute to be $ε$-approximated for some $γ>2$, then ML-EM $ε$-approximates the solution of the SDE with $ε^{-γ}$ compute, improving over the traditional EM rate of $ε^{-γ-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 $γ\approx2.5$. Given that this is a polynomial speedup, we expect even stronger speedups in practical applications which involve orders of magnitude larger networks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。