步调一致：优化扩散模型中的采样计划

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

作者: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

摘要

扩散模型（DMs）已经成为视觉领域及其他领域中最先进的生成建模方法。DMs的一个关键缺点是其较慢的采样速度，依赖于通过大型神经网络进行许多顺序函数评估。从DMs中进行采样可以被视为通过离散化的噪声水平集合解决微分方程。虽然过去的研究主要集中在推导高效求解器上，但对于寻找最佳采样计划却鲜有关注，整个文献都依赖于手工制定的启发式方法。在这项工作中，我们首次提出了一种通用且有原则的方法来优化DMs的采样计划，以获得高质量的输出，称为“调整您的步骤”。我们利用随机微积分方法，并找到了针对不同求解器、训练过的DMs和数据集的最佳计划。我们在几个图像、视频以及2D玩具数据合成基准上评估了我们的新方法，使用了各种不同的采样器，并观察到我们优化的计划在几乎所有实验中均优于先前手工制定的计划。我们的方法展示了采样计划优化的潜力，特别是在少步骤合成方案中。

English

Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. Sampling from DMs can be seen as solving a differential equation through a discretized set of noise levels known as the sampling schedule. While past works primarily focused on deriving efficient solvers, little attention has been given to finding optimal sampling schedules, and the entire literature relies on hand-crafted heuristics. In this work, for the first time, we propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs, called Align Your Steps. We leverage methods from stochastic calculus and find optimal schedules specific to different solvers, trained DMs and datasets. We evaluate our novel approach on several image, video as well as 2D toy data synthesis benchmarks, using a variety of different samplers, and observe that our optimized schedules outperform previous hand-crafted schedules in almost all experiments. Our method demonstrates the untapped potential of sampling schedule optimization, especially in the few-step synthesis regime.