[论文解读] OMS-DPM: Optimizing the Model Schedule for Diffusion Probabilistic Models

Enshu Liu, Xuefei Ning+3 more|arXiv (Cornell University)|Jun 15, 2023

Advanced Neural Network Applications被引用 6

一句话总结

OMS-DPM 提出了一种新颖的方法，通过在去噪步骤中动态选择不同的预训练模型来优化扩散概率模型的模型调度，显著提升了生成质量和速度。采用基于预测器的进化搜索方法，其采样速度比 Stable Diffusion 快 2 倍，同时在 CIFAR-10、CelebA、ImageNet 和 LSUN-Church 上保持了最先进水平的 FID 分数。

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$ imes$ while maintaining the generation quality.

研究动机与目标

• 解决扩散概率模型（DPMs）采样速度慢的问题，其生成过程需要数百至数千次神经网络评估。
• 识别并利用被忽视的优化维度——模型调度，即在不同去噪步骤中使用不同模型，以同时提升质量和速度。
• 开发一种方法，在仅使用预训练模型且无需微调的情况下，基于任意时间预算找到最优模型调度。
• 实现在多种数据集和预训练模型（包括 Stable Diffusion 等公开检查点）上的高效、可迁移的调度策略。
• 通过利用去噪步骤中模型异质性，实现优于现有方法的质量-速度权衡。

提出的方法

• 提出 OMS-DPM，一种基于预测器的搜索算法，通过少量训练数据学习从模型调度预测生成质量。
• 在少量模型调度子集上训练性能预测器，以实现对未见调度的泛化，从而减少对昂贵完整评估的依赖。
• 使用进化算法在给定时间预算下高效探索指数级庞大的模型调度搜索空间。
• 通过将 DDIM 步骤跳过和 DPM-Solver 顺序选择等现有加速技术视为搜索中的超参数，支持灵活集成。
• 设计搜索过程，根据预测性能将不同大小和速度的预训练模型分配至不同去噪步骤。
• 通过在每个数据集上仅使用少量数据点进行预测器训练，实现可迁移性，使新数据集仅需极少微调即可复用。

实验结果

研究问题

• RQ1 在扩散概率模型中，是否可以通过在不同去噪步骤中动态分配不同的预训练模型，同时提升生成质量和采样速度？
• RQ2 与仅关注噪声调度、离散化或求解器设计的传统加速方法相比，模型调度优化有何优势？
• RQ3 在时间预算受限的情况下，模型大小和放置位置（早期 vs. 晚期步骤）对最终生成质量有何影响？
• RQ4 轻量级预测器是否能泛化至多样化的模型调度和数据集，从而实现在无需完整评估情况下的高效搜索？
• RQ5 最优模型调度的结构在不同数据集和采样预算下如何变化？

主要发现

• OMS-DPM 在保持或提升 FID 分数的同时，使 Stable Diffusion 模型的采样速度提升 2 倍，优于基线方法。
• 在 CIFAR-10 上，OMS-DPM 在 1400ms 延迟预算下将 FID 从基线的 3.56 降低至 3.19，显著改善了质量-速度权衡。
• 在时间预算紧张时，最优调度仅使用最小、最快的模型（例如 2–3 个模型），因为大模型在步骤过少时会引入误差。
• 在 LSUN-Church 上，最优策略是在去噪过程早期使用大模型；而在 ImageNet-64 和 CIFAR-10 上，大模型在后期更有效，表明调度策略具有数据集依赖性。
• 在仅使用 100 个数据点进行训练的情况下，预测器仍具有良好泛化能力，使评估成本极低的高效搜索成为可能。
• 对于 DPM-Solver，在时间预算紧张时优先使用一阶和二阶求解器，而三阶求解器仅在时间充足时使用，表明最优调度中存在自适应求解器选择。