[论文解读] UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models

Wenliang Zhao, Lujia Bai+3 more|arXiv (Cornell University)|Feb 9, 2023

Cell Image Analysis Techniques被引用 29

一句话总结

UniPC 引入一个训练无关的统一预测-校正框架，以加速扩散模型采样，在任意阶数下实现更高阶精度且无需额外模型评估，在无条件与有条件任务中用极少的函数评估进行演示。

ABSTRACT

Diffusion probabilistic models (DPMs) have demonstrated a very promising ability in high-resolution image synthesis. However, sampling from a pre-trained DPM is time-consuming due to the multiple evaluations of the denoising network, making it more and more important to accelerate the sampling of DPMs. Despite recent progress in designing fast samplers, existing methods still cannot generate satisfying images in many applications where fewer steps (e.g., $<$10) are favored. In this paper, 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, especially in extremely few steps. 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 256$ imes$256 (conditional) with only 10 function evaluations. Code is available at https://github.com/wl-zhao/UniPC.

研究动机与目标

• 推动扩散概率模型（DPMs）在高分辨率图像合成中的快速采样。
• 开发一个训练无关的框架，在不增加额外模型评估的情况下提升采样精度。
• 提供统一预测器（UniP）和统一校正器（UniC），从而给出可任意阶数的灵活 UniPC 框架。

提出的方法

• 引入 UniC-p 作为校正器，利用当前输出和先前模型输出在不进行额外评估的情况下对当前估计进行细化。
• 推导 UniP-p 作为预测器，利用先前点来实现 p 阶精度。
• 将 UniP 与 UniC 结合为 UniPC，一种具有闭式、阶数统一框架的统一预测-校正家族，适用于任意阶数（p+1 精度）。
• 在正则性条件下证明 UniC-p 可达到(p+1)阶精度，且 UniP-p 能达到 p 阶精度，从而实现更高阶的 UniPC 变体。
• 实现多步更新以重复利用过去的结果以提高效率，并提供与现有求解器配合使用、提升精度阶数的 UniPC-p 变体。

实验结果

研究问题

• RQ1 一种训练无关的预测-校正框架是否能够在不增加函数评估次数的情况下提升扩散模型采样器的精度阶数？
• RQ2 UniPC 是否提供任意阶数的统一解析形式？在无条件与有条件设置下与现有快速采样器相比如何？
• RQ3 将 UniC 作为后处理校正器应用时，对现有采样器的改进程度有多大？
• RQ4 在像素空间和潜在空间的 DPMs 中，UniPC 的实际折衷（速度、内存、稳定性）有哪些？

主要发现

• 在 5–10 次函数评估（NFE）下，UniPC 的采样质量优于先前方法，其中在 CIFAR10（无条件）上达到 3.87 的 FID，在 ImageNet 256×256（有条件）上达到 7.51 的 FID，且使用 10 NFEs。
• UniC 可以在现有采样器之后应用以在不增加额外模型评估的情况下提高精度，UniP 提供与阶数兼容的预测器，从而实现统一的(p+1) 阶框架。
• 该统一框架支持带解析形式的任意阶数，多步 UniPC 更新复用过去的结果，在速度和内存方面的效率与 DPM-Solver++ 相当。
• UniPC 同时改进像素空间和潜在空间的扩散模型，并增强在分类器引导和无分类器引导下的引导采样。
• 实证结果表明，在极低的 NFEs 下，UniPC 在无条件和有条件任务以及不同数据集（CIFAR10、FFHQ、LSUN、ImageNet）上均优于之前的快速采样器。
• 在 UniPC 内可自定义的阶数调度可以产生比简单提升阶数更好的结果，且 UniPC 维持具有竞争力的推理速度和内存使用。