Accelerating Diffusion Sampling with Optimized Time Steps

• Main Authors: Xue, Shuchen, Liu, Zhaoqiang, Chen, Fei, Zhang, Shifeng, Hu, Tianyang, Xie, Enze, Li, Zhenguo
• Format: Journal Article
• Language: English
• Published: 27.02.2024
• Subjects: Computer Science - Artificial Intelligence; Computer Science - Computer Vision and Pattern Recognition; Computer Science - Learning
• DOI: 10.48550/arxiv.2402.17376

Diffusion probabilistic models (DPMs) have shown remarkable performance in high-resolution image synthesis, but their sampling efficiency is still to be desired due to the typically large number of sampling steps. Recent advancements in high-order numerical ODE solvers for DPMs have enabled the generation of high-quality images with much fewer sampling steps. While this is a significant development, most sampling methods still employ uniform time steps, which is not optimal when using a small number of steps. To address this issue, we propose a general framework for designing an optimization problem that seeks more appropriate time steps for a specific numerical ODE solver for DPMs. This optimization problem aims to minimize the distance between the ground-truth solution to the ODE and an approximate solution corresponding to the numerical solver. It can be efficiently solved using the constrained trust region method, taking less than $15$ seconds. Our extensive experiments on both unconditional and conditional sampling using pixel- and latent-space DPMs demonstrate that, when combined with the state-of-the-art sampling method UniPC, our optimized time steps significantly improve image generation performance in terms of FID scores for datasets such as CIFAR-10 and ImageNet, compared to using uniform time steps.