用于快速扩散采样的可微分求解器搜索

扩散模型展示了卓越的生成质量，但代价是大量的函数评估。最近，开发了基于 ODE 的先进求解器，以减轻在有限采样步数下进行逆扩散求解的巨大计算需求。然而，这些求解器在很大程度上受 Adams 类多步方法的启发，仅依赖于与 t 相关的拉格朗日插值。我们表明，与 t 相关的拉格朗日插值对于扩散模型来说并非最优，并揭示了一个由时间步长和求解器系数组成的紧凑搜索空间。基于我们的分析，我们提出了一种新颖的可微分求解器搜索算法，以识别更优的求解器。配备了搜索到的求解器后，经过校正流的模型，例如 SiT-XL/2 和 FlowDCN-XL/2，在 ImageNet256 上仅需 10 步就分别达到了 2.40 和 2.35 的 FID 分数。同时，DDPM 模型 DiT-XL/2 在仅需 10 步的情况下达到了 2.33 的 FID 分数。值得注意的是，我们搜索到的求解器显著优于传统求解器。此外，我们搜索到的求解器在各种模型架构、分辨率和模型尺寸上都表现出了通用性。