Optimizing Few-Step Sampler for Diffusion Probabilistic Model
This work addresses a practical bottleneck for users of DPMs in image generation by reducing inference time, though it appears incremental as it builds on existing methods.
The paper tackled the high computational cost of Diffusion Probabilistic Models (DPMs) during inference by optimizing the sampling schedule and fine-tuning the model, resulting in consistent improvements across various sampling steps on the ImageNet64 dataset.
Diffusion Probabilistic Models (DPMs) have demonstrated exceptional capability of generating high-quality and diverse images, but their practical application is hindered by the intensive computational cost during inference. The DPM generation process requires solving a Probability-Flow Ordinary Differential Equation (PF-ODE), which involves discretizing the integration domain into intervals for numerical approximation. This corresponds to the sampling schedule of a diffusion ODE solver, and we notice the solution from a first-order solver can be expressed as a convex combination of model outputs at all scheduled time-steps. We derive an upper bound for the discretization error of the sampling schedule, which can be efficiently optimized with Monte-Carlo estimation. Building on these theoretical results, we purpose a two-phase alternating optimization algorithm. In Phase-1, the sampling schedule is optimized for the pre-trained DPM; in Phase-2, the DPM further tuned on the selected time-steps. Experiments on a pre-trained DPM for ImageNet64 dataset demonstrate the purposed method consistently improves the baseline across various number of sampling steps.