Differentiable Solver Search for Fast Diffusion Sampling

TL;DR

This paper introduces a novel differentiable solver search algorithm for diffusion models, significantly improving sampling efficiency and quality with fewer steps by optimizing solver parameters beyond traditional methods.

Contribution

The paper proposes a new differentiable solver search method that finds more optimal diffusion solvers, outperforming traditional approaches in speed and accuracy.

Findings

Achieved state-of-the-art FID scores with only 10 steps on ImageNet256.

Searched solver outperforms traditional solvers by a large margin.

Demonstrated solver's generality across models, resolutions, and sizes.

Abstract

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.