可微求解器搜索加速扩散采样

Differentiable Solver Search for Fast Diffusion Sampling

May 27, 2025

作者: Shuai Wang, Zexian Li, Qipeng zhang, Tianhui Song, Xubin Li, Tiezheng Ge, Bo Zheng, Limin Wang

cs.AI

摘要

扩散模型已展现出卓越的生成质量，但代价是大量的函数评估。近期，基于ODE的高级求解器被开发出来，以缓解在有限采样步数下逆向扩散求解的巨大计算需求。然而，这些深受亚当斯类多步方法启发的求解器，仅依赖于与时间t相关的拉格朗日插值。我们证明，对于扩散模型而言，t相关的拉格朗日插值并非最优，并揭示了一个由时间步长和求解器系数构成的紧凑搜索空间。基于此分析，我们提出了一种新颖的可微分求解器搜索算法，以识别更优的求解器。配备所搜索到的求解器后，修正流模型，如SiT-XL/2和FlowDCN-XL/2，在仅10步的情况下，于ImageNet256上分别取得了2.40和2.35的FID分数。同时，DDPM模型DiT-XL/2在仅10步时也达到了2.33的FID分数。值得注意的是，我们搜索到的求解器显著优于传统求解器。此外，我们的求解器在不同模型架构、分辨率及模型大小上均展现出良好的通用性。

English

Diffusion models have demonstrated remarkable generation quality but at the cost of numerous function evaluations. Recently, advanced ODE-based solvers have been developed to mitigate the substantial computational demands of reverse-diffusion solving under limited sampling steps. However, these solvers, heavily inspired by Adams-like multistep methods, rely solely on t-related Lagrange interpolation. We show that t-related Lagrange interpolation is suboptimal for diffusion model and reveal a compact search space comprised of time steps and solver coefficients. Building on our analysis, we propose a novel differentiable solver search algorithm to identify more optimal solver. Equipped with the searched solver, rectified-flow models, e.g., SiT-XL/2 and FlowDCN-XL/2, achieve FID scores of 2.40 and 2.35, respectively, on ImageNet256 with only 10 steps. Meanwhile, DDPM model, DiT-XL/2, reaches a FID score of 2.33 with only 10 steps. Notably, our searched solver outperforms traditional solvers by a significant margin. Moreover, our searched solver demonstrates generality across various model architectures, resolutions, and model sizes.