可微求解器搜索用於快速擴散採樣

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.