Differentiable Solver Search for Fast Diffusion Sampling

TL;DR

This work tackles the computational bottleneck of diffusion sampling by introducing a differentiable solver search (DS-Solver) that learns an optimal, compact set of timesteps and solver coefficients tailored to pre-trained diffusion models. By shifting focus from interpolation forms to solver parameters and leveraging pre-integral analysis, the method achieves large gains in sampling efficiency, delivering high-quality images with as few as 10 function evaluations and demonstrating transferability across Rectified Flow and DDPM/VP models. The approach yields 2.33 FID on ImageNet-256 with 10 steps for DiT-XL/2 and 2.40–2.35 FID for Rectified Flow variants, highlighting practical impact for fast diffusion-based generation. Overall, the paper provides a data-driven, solver-centric pathway to near-state-of-the-art diffusion performance under stringent inference budgets, with broad applicability to different architectures and resolutions.

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.

Figures (6)

• Figure 1: Generated images from Flux.1-dev with Guidance=2.0 and our solver (searched on SiT-XL/2).Euler-Shift3 is the default solver provided by diffusers and Flux community. Our solver(DS-Solver) achieves better visual quality from 5 to 10 steps(NFE).
• Figure 2: Ablations studies of Differentiable Solver Search.We evaluate the searched solver on SiT-XL/2, and report the FID performance curve of searched solvers.
• Figure 3: The same searched solver on different Rectified-Flow Models.R256 and R512 indicate the generation resolution of given model. We search solver with FlowDCN-B/2 on ImageNet-$256\times256$ and evaluate it with SiT-XL/2 and FlowDCN-XL/2 on different resolution datasets. Our searched solver outperforms traditional solvers by a significant margin. More metrics(sFID, IS, Precision, Recall) are places at Appendix
• Figure 4: The images generated from PixArt-$\Sigma$ with CFG=2.0 equipped with Our DS-Solver ( searched on DiT-XL/2-R256 ). In comparison to DPM-Solver++ and UniPC, our results consistently exhibit greater clarity and possess more details. Our solver achieves better quality from 5 to 10 steps(NFE).
• Figure 5: The same searched solver on different Rectified-Flow Models.R256 and R512 indicate the generation resolution of given model. We search solver with FlowDCN-B/2 on ImageNet-$256\times256$ and evaluate it with SiT-XL/2 and FlowDCN-XL/2 on different resolution datasets. Our searched solver outperforms traditional solvers by a significant margin.

Theorems & Definitions (7)

• Theorem 4.2
• Theorem 4.4
• Theorem 4.5
• Theorem 8.1
• proof
• Theorem 9.1
• proof