Busca de Solucionadores Diferenciáveis para Amostragem Rápida em Difusão

Differentiable Solver Search for Fast Diffusion Sampling

May 27, 2025

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

cs.AI

Resumo

Os modelos de difusão demonstraram uma qualidade de geração notável, mas ao custo de inúmeras avaliações de função. Recentemente, solucionadores avançados baseados em EDOs foram desenvolvidos para mitigar as demandas computacionais substanciais da resolução de difusão reversa com um número limitado de passos de amostragem. No entanto, esses solucionadores, fortemente inspirados em métodos multistep do tipo Adams, dependem exclusivamente da interpolação de Lagrange relacionada a t. Mostramos que a interpolação de Lagrange relacionada a t é subótima para modelos de difusão e revelamos um espaço de busca compacto composto por passos de tempo e coeficientes do solucionador. Com base em nossa análise, propomos um novo algoritmo de busca de solucionador diferenciável para identificar um solucionador mais ótimo. Equipados com o solucionador encontrado, modelos de fluxo retificado, como SiT-XL/2 e FlowDCN-XL/2, alcançam pontuações FID de 2,40 e 2,35, respectivamente, no ImageNet256 com apenas 10 passos. Enquanto isso, o modelo DDPM, DiT-XL/2, atinge uma pontuação FID de 2,33 com apenas 10 passos. Notavelmente, nosso solucionador encontrado supera os solucionadores tradicionais por uma margem significativa. Além disso, nosso solucionador encontrado demonstra generalidade em várias arquiteturas de modelo, resoluções e tamanhos de modelo.

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.