Optimizing Few-Step Sampler for Diffusion Probabilistic Model

TL;DR

This paper introduces an optimized sampling schedule for diffusion probabilistic models that reduces computational cost and improves image generation quality by theoretically analyzing and empirically validating a two-phase optimization approach.

Contribution

It derives an upper bound for discretization error in diffusion sampling and proposes a two-phase optimization algorithm to enhance sampling efficiency and quality.

Findings

Improved image quality with fewer sampling steps.

Consistent performance gains across different sampling step counts.

Effective optimization of sampling schedule for pre-trained models.

Abstract

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.