Ricerca di Solver Differenziabili per il Campionamento Rapido nella Diffusione

Differentiable Solver Search for Fast Diffusion Sampling

May 27, 2025

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

cs.AI

Abstract

I modelli di diffusione hanno dimostrato una qualità di generazione notevole, ma al costo di numerose valutazioni di funzioni. Recentemente, sono stati sviluppati risolutori avanzati basati su ODE per mitigare le sostanziali richieste computazionali della risoluzione della diffusione inversa con un numero limitato di passi di campionamento. Tuttavia, questi risolutori, fortemente ispirati da metodi multistep di tipo Adams, si basano esclusivamente sull'interpolazione di Lagrange relativa a t. Dimostriamo che l'interpolazione di Lagrange relativa a t è subottimale per i modelli di diffusione e riveliamo uno spazio di ricerca compatto composto da passi temporali e coefficienti del risolutore. Basandoci sulla nostra analisi, proponiamo un nuovo algoritmo di ricerca differenziabile per identificare un risolutore più ottimale. Dotati del risolutore trovato, i modelli di flusso rettificato, ad esempio SiT-XL/2 e FlowDCN-XL/2, raggiungono punteggi FID di 2.40 e 2.35, rispettivamente, su ImageNet256 con soli 10 passi. Nel frattempo, il modello DDPM, DiT-XL/2, raggiunge un punteggio FID di 2.33 con soli 10 passi. È importante notare che il nostro risolutore trovato supera i risolutori tradizionali con un margine significativo. Inoltre, il nostro risolutore trovato dimostra generalità su varie architetture di modelli, risoluzioni e dimensioni del modello.

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.