Поиск дифференцируемых решателей для ускорения сэмплирования в диффузионных моделях

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 \). Мы показываем, что такая интерполяция является неоптимальной для диффузионных моделей, и раскрываем компактное пространство поиска, состоящее из временных шагов и коэффициентов решателя. На основе нашего анализа мы предлагаем новый дифференцируемый алгоритм поиска решателя, который позволяет находить более оптимальные решения. Оснащенные найденным решателем, модели с исправленным потоком, такие как SiT-XL/2 и FlowDCN-XL/2, достигают показателей FID 2.40 и 2.35 соответственно на ImageNet256 всего за 10 шагов. В то же время модель DDPM, DiT-XL/2, достигает показателя FID 2.33 также за 10 шагов. Примечательно, что наш найденный решатель значительно превосходит традиционные решатели. Более того, он демонстрирует универсальность, работая с различными архитектурами моделей, разрешениями и размерами моделей.

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.