OMS-DPM: Optimizing the Model Schedule for Diffusion Probabilistic Models

TL;DR

This paper introduces OMS-DPM, a method to optimize the sequence of models used during diffusion probabilistic model generation, significantly improving speed without sacrificing quality by leveraging a model schedule approach.

Contribution

It proposes a predictor-based algorithm to optimize model schedules for DPMs, enabling faster generation with maintained quality across multiple datasets.

Findings

OMS-DPM outperforms prior methods in quality-speed trade-offs.

Achieves 2x acceleration on Stable Diffusion checkpoints.

Effective across diverse datasets like CIFAR-10, CelebA, ImageNet, and LSUN.

Abstract

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.