高速拡散サンプリングのための微分可能ソルバー探索

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ベースのソルバーが開発されました。しかし、これらのソルバーはAdams型多段階法に強く影響を受けており、t関連のラグランジュ補間にのみ依存しています。本研究では、t関連のラグランジュ補間が拡散モデルにとって最適ではないことを示し、時間ステップとソルバー係数から構成されるコンパクトな探索空間を明らかにします。この分析に基づき、より最適なソルバーを特定するための新しい微分可能なソルバー探索アルゴリズムを提案します。探索されたソルバーを装備した修正フローモデル（例：SiT-XL/2およびFlowDCN-XL/2）は、ImageNet256においてわずか10ステップでそれぞれ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.