OMS-DPM: Optimizing the Model Schedule for Diffusion Probabilistic
  Models
Enshu Liu Xuefei Ning Zinan Lin Huazhong Yang Yu Wang
Diffusion probabilistic models (DPMs) are a new class of generative models
that have achieved state-of-the-art generation quality in various domains.
Despite the promise, one major drawback of DPMs is the slow generation speed
due to the large number of neural network evaluations required in the
generation process. In this paper, we reveal an overlooked dimension -- model
schedule -- for optimizing the trade-off between generation quality and speed.
More specifically, we observe that small models, though having worse generation
quality when used alone, could outperform large models in certain generation
steps. Therefore, unlike the traditional way of using a single model, using
different models in different generation steps in a carefully designed
\emph{model schedule} could potentially improve generation quality and speed
\emph{simultaneously}. We design OMS-DPM, a predictor-based search algorithm,
to optimize the model schedule given an arbitrary generation time budget and a
set of pre-trained models. We demonstrate that OMS-DPM can find model schedules
that improve generation quality and speed than prior state-of-the-art methods
across CIFAR-10, CelebA, ImageNet, and LSUN datasets. When applied to the
public checkpoints of the Stable Diffusion model, we are able to accelerate the
sampling by 2$\times$ while maintaining the generation quality.