Recherche de Solveurs Différentiables pour un Échantillonnage Rapide de Diffusion

Differentiable Solver Search for Fast Diffusion Sampling

May 27, 2025

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

cs.AI

Résumé

Les modèles de diffusion ont démontré une qualité de génération remarquable, mais au prix d'un nombre élevé d'évaluations de fonctions. Récemment, des solveurs avancés basés sur les équations différentielles ordinaires (EDO) ont été développés pour atténuer les exigences computationnelles substantielles de la résolution de la diffusion inverse avec un nombre limité d'étapes d'échantillonnage. Cependant, ces solveurs, fortement inspirés des méthodes multi-étapes de type Adams, reposent uniquement sur l'interpolation de Lagrange liée au temps t. Nous montrons que l'interpolation de Lagrange liée au temps t est sous-optimale pour les modèles de diffusion et révélons un espace de recherche compact composé des pas de temps et des coefficients des solveurs. Sur la base de notre analyse, nous proposons un nouvel algorithme de recherche de solveur différentiable pour identifier un solveur plus optimal. Équipés du solveur recherché, les modèles de flux rectifiés, par exemple, SiT-XL/2 et FlowDCN-XL/2, atteignent des scores FID de 2,40 et 2,35, respectivement, sur ImageNet256 avec seulement 10 étapes. Parallèlement, le modèle DDPM, DiT-XL/2, atteint un score FID de 2,33 avec seulement 10 étapes. Notamment, notre solveur recherché surpasse les solveurs traditionnels de manière significative. De plus, notre solveur recherché démontre une généralité à travers diverses architectures de modèles, résolutions et tailles de modèles.

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.