Fast ODE-based Sampling for Diffusion Models in Around 5 Steps
This work addresses the computational bottleneck of slow sampling in diffusion models for image synthesis, offering a plugin method to improve efficiency, though it is incremental as it builds on existing ODE-based samplers.
The paper tackles the problem of sampling from diffusion models with very few function evaluations (around 5 steps), which degrades sample quality due to approximation errors in existing ODE-based methods; it proposes AMED-Solver, which learns the mean direction to eliminate truncation errors, achieving FID scores of 6.61 on CIFAR-10, 10.74 on ImageNet 64x64, and 13.20 on LSUN Bedroom with only 5 NFE.
Sampling from diffusion models can be treated as solving the corresponding ordinary differential equations (ODEs), with the aim of obtaining an accurate solution with as few number of function evaluations (NFE) as possible. Recently, various fast samplers utilizing higher-order ODE solvers have emerged and achieved better performance than the initial first-order one. However, these numerical methods inherently result in certain approximation errors, which significantly degrades sample quality with extremely small NFE (e.g., around 5). In contrast, based on the geometric observation that each sampling trajectory almost lies in a two-dimensional subspace embedded in the ambient space, we propose Approximate MEan-Direction Solver (AMED-Solver) that eliminates truncation errors by directly learning the mean direction for fast diffusion sampling. Besides, our method can be easily used as a plugin to further improve existing ODE-based samplers. Extensive experiments on image synthesis with the resolution ranging from 32 to 512 demonstrate the effectiveness of our method. With only 5 NFE, we achieve 6.61 FID on CIFAR-10, 10.74 FID on ImageNet 64$\times$64, and 13.20 FID on LSUN Bedroom. Our code is available at https://github.com/zju-pi/diff-sampler.