中文

S4S:为扩散模型求解器寻解

机器学习 2025-02-25 v1

摘要

扩散模型(DM)通过从随机噪声出发并迭代求解逆时间常微分方程(ODE),从数据分布中生成样本。由于迭代求解中的每一步都需要昂贵的神经函数评估(NFE),人们对于在不修改底层模型的情况下仅用少量 NFE 近似求解这些扩散 ODE 产生了浓厚的兴趣。然而,在少量 NFE 机制下,我们观察到使用传统 ODE 求解器根本无法追踪真实的 ODE 演化。在本工作中,我们提出了一种为 DM 学习良好求解器的新方法,我们称之为为求解器寻解(S4S)。S4S 通过学习匹配强教师求解器的输出,直接优化求解器以获得良好的生成质量。我们在六个不同的预训练 DM 上评估了 S4S,包括用于条件采样和无条件采样的像素空间和潜在空间 DM。在所有设置中,相对于传统 ODE 求解器,S4S 均匀地提高了样本质量。此外,我们的方法轻量级、无需数据,并且可以作为黑盒即插即用地应用于任何离散化调度或架构之上以提升性能。在此基础上,我们还提出了 S4S-Alt,它同时优化求解器和离散化调度。通过充分利用 DM 求解器的完整设计空间,在 5 次 NFE 下,我们在 CIFAR10 上实现了 3.73 的 FID,在 MS-COCO 上实现了 13.26 的 FID,相较于以往免训练 ODE 方法代表了 1.5×1.5\times 的提升。

关键词

引用

@article{arxiv.2502.17423,
  title  = {S4S: Solving for a Diffusion Model Solver},
  author = {Eric Frankel and Sitan Chen and Jerry Li and Pang Wei Koh and Lillian J. Ratliff and Sewoong Oh},
  journal= {arXiv preprint arXiv:2502.17423},
  year   = {2025}
}