UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models

Wenliang Zhao* Lujia Bai* Yongming Rao Jie Zhou Jiwen Lu

Tsinghua University

Figure 1: The main idea of UniPC. We provide an illustration of the multistep variant of UniPC with noise prediction. As a predictor-corrector style fast sampler for DPMs, UniPC consists of UniP and UniC which share the same analytical form. Apart from only considering the output in previous timesteps, UniC can increase the order of accuracy by also leveraging the model output of the current timestep.

Abstract

Diffusion probabilistic models (DPMs) have demonstrated a very promising ability in high-resolution image synthesis. However, sampling from a pre-trained DPM usually requires hundreds of model evaluations, which is computationally expensive. Despite recent progress in designing high-order solvers for DPMs, there still exists room for further speedup, especially in extremely few steps (e.g., 5~10 steps). Inspired by the predictor-corrector for ODE solvers, we develop a unified corrector (UniC) that can be applied after any existing DPM sampler to increase the order of accuracy without extra model evaluations, and derive a unified predictor (UniP) that supports arbitrary order as a byproduct. Combining UniP and UniC, we propose a unified predictor-corrector framework called UniPC for the fast sampling of DPMs, which has a unified analytical form for any order and can significantly improve the sampling quality over previous methods. We evaluate our methods through extensive experiments including both unconditional and conditional sampling using pixel-space and latent-space DPMs. Our UniPC can achieve 3.87 FID on CIFAR10 (unconditional) and 7.51 FID on ImageNet 256x256 (conditional) with only 10 function evaluations.

A Unified Predictor-Corrector Solver

We propose a unified predictor-corrector solver of DPMs called UniPC, consisting of UniP and UniC. Our UniPC is unified in mainly two aspects: 1) the predictor (UniP) and the corrector (UniC) share the same analytical form; 2) UniP supports arbitrary order and UniC can be applied after off-the-shelf fast samplers of DPMs to increase the order of accuracy.

Results

We show UniPC outperforms previous methods in both unconditional and conditional sampling tasks.
We demonstrate that UniC can be applied after off-the-shelf fast solvers of DPMs to further accelerate the sampling.

Figure 2: Unconditional sampling results. We compare our UniPC with DDIM and DPM-Solver++ on CIFAR10, LSUN Bedroom, and FFHQ. We report the FID↓ of the methods with different numbers of function evaluations (NFE). Experimental results demonstrate that our method is consistently better than previous ones on both pixel-space DPMs and latent-space DPMs, especially with extremely few steps.

Figure 3: Conditional sampling results. (a)(b) We compare the sample quality measured by FID↓ on ImageNet 256x256 with guidance scale s=8.0/4.0; (c) We adopt the text-to-image model provided by stable-diffusion to compare the convergence error, which is measured by the l2 distance between the results of different methods and 1000-step DDIM. We show that our method outperforms previous ones with various guidance scales and NFE.

Table 1: Ablation on the choice of B(h). We investigate how the choice of B(h) affects the performance of our UniPC.

Table 2: Applying UniC to any solvers. We demonstrate that our UniC can be a plug-and-play component to boost the performance of a wide range of solvers, including both singlestep and multistep solvers with different orders. The sampling quality is measured by FID↓ on the CIFAR10 dataset.

Figure 4: Qualitative comparisons between our UniPC and previous methods. All images are generated by sampling from a DPM trained on ImageNet 256x256 with only 7 number of function evaluations (NFE) and a classifier scale of 8.0. We show that our proposed UniPC can generate more plausible samples with more visual details compared with previous samplers.

BibTeX

@article{zhao2023unipc,

title={UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models},

author={Zhao, Wenliang and Bai, Lujia and Rao, Yongming and Zhou, Jie and Lu, Jiwen},

journal={arXiv preprint arXiv:2302.04867},

year={2023}

}