How Differentiable Solver Search Accelerates Diffusion Model Sampling

This article shares a fast sampling method for diffusion models proposed by Alibaba's Intelligent Creation and AI Application team, based on a differentiable solver search that discovers high‑quality sampling paths for a given number of steps.

Paper: Differentiable Solver Search for Fast Diffusion Sampling

Authors: Shuai Wang, Zexian Li, Qipeng Zhang, Tianhui Song, Xubin Li, Tiezheng Ge, Bo Zheng, Limin Wang

Paper URL: https://arxiv.org/abs/2505.21114

Paper Quick Read

The paper introduces a solver search algorithm: given a model and a target number of inference steps, it efficiently finds an excellent sampling trajectory.

Current state‑of‑the‑art samplers rely on linear multistep methods with time‑only Lagrange interpolation, which is suboptimal for a specific diffusion model. By pre‑integrating the interpolation function into solver coefficients, the sampling error depends only on the timestep and coefficients, enabling a compact search space of only dozens of parameters.

Extensive experiments on Rectified‑Flow, DDPM/VP, and various text‑to‑image models validate the method’s superiority.

1. Background

Image generation aims to learn the data distribution of real images and sample from it. Diffusion models have recently outperformed GANs and autoregressive models, but they require many denoising steps, leading to high computational cost.

Two main acceleration routes exist:

Training‑based methods : distill fast sampling paths into model parameters, but they require retraining and cannot fully exploit pretrained models.

Solver‑based methods : design ODE solvers for the reverse diffusion process, preserving the pretrained model’s performance and being applicable to any model with a similar noise schedule.

This work proposes a data‑driven solver search that customizes solvers for each pretrained model without altering its parameters.

2. Task Definition

The focus is on the Rectified‑Flow model as a representative diffusion family, though the algorithm also applies to DDPM/VP and other models. The reverse diffusion ODE is treated as an integral problem, and the goal is to find a solver that minimizes integration error within a limited number of function evaluations.

3. Method

We show that the specific form of the interpolation function is irrelevant; after pre‑integration it reduces to a set of coefficients. These coefficients, together with timesteps, form a compact search space.

3.1 Revisiting Multistep Methods

Classic multistep solvers such as Adams–Bashforth use Lagrange polynomials, which can be pre‑computed into constant coefficients, reducing ODE solving to simple summations.

3.2 From Interpolation Functions to Solver Coefficients

By pre‑integrating the interpolation function, we replace it with coefficient vectors that depend on the current state, avoiding over‑fitting and preserving generalization.

3.3 Differentiable Solver Search

We parameterize timesteps with unbounded variables transformed via Softmax to ensure they sum to one, and we re‑parameterize the diagonal of the coefficient matrix while initializing off‑diagonal entries to zero.

Mono‑alignment Supervision

A high‑step Euler solver provides reference ODE trajectories; we minimize mean‑squared error between the N‑step solver’s trajectory and the reference, supplemented by a Huber loss on the final image.

4. Experiments

We evaluate on open‑source diffusion models using DiT‑XL/2 (DDPM), SiT‑XL/2 and FlowDCN‑XL/2 (Rectified‑Flow). Training uses the Lion optimizer with a fixed learning rate of 0.01 and no weight decay. Searching 50,000 samples for FlowDCN‑B/2 takes about half an hour.

4.1 Rectified‑Flow Models

On ImageNet‑256×256, the searched solver achieves FID 2.40 for SiT‑XL/2 and 2.35 for FlowDCN‑XL/2 within 10 steps, far surpassing traditional solvers. Similar gains are observed on ImageNet‑512×512 and in text‑to‑image generation with Flux.1‑dev and SD3.

4.2 DDPM/VP Models

Using DiT‑XL/2 on ImageNet‑256×256, the searched solver reaches FID 2.33 in 10 steps (CFG = 1.5). On ImageNet‑512×512, it attains FID 3.64, again outperforming DPM‑Solver++ and UniPC. The solver also improves PixArt‑Σ results under CFG = 2.0.

4.3 Solver Parameter Visualization

Coefficient matrices and timestep distributions are visualized, showing diagonal concentration for DDPM/VP and smoother patterns for Rectified‑Flow.

5. Conclusion

By converting diffusion ODE interpolation functions into coefficient vectors via pre‑integration, we define a compact solver search space and introduce a differentiable, data‑driven search algorithm that yields optimal solvers. Experiments on multiple diffusion frameworks confirm substantial image quality improvements, highlighting the method’s low training cost and high practical value.

6. References

[1] Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. NeurIPS, 2020.

[2] Song, Y., et al. Score‑based generative modeling through stochastic differential equations. arXiv, 2020.

[3] Karras, T., et al. Elucidating the design space of diffusion‑based generative models. NeurIPS, 2022.

[4] Brock, A., et al. Large scale GAN training for high fidelity natural image synthesis. arXiv, 2018.

[5] Sauer, A., et al. StyleGAN‑XL: Scaling StyleGAN to large diverse datasets. SIGGRAPH, 2022.

[6] Chang, H., et al. MaskGit: Masked generative image transformer. CVPR, 2022.

[7] Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. arXiv, 2022.

[8] Song, Y., et al. Consistency models. arXiv, 2023.

[9] Zheng, J., et al. Trajectory consistency distillation. arXiv, 2024.

[10] Song, T., et al. DMM: Building a versatile image generation model via distillation‑based model merging. arXiv, 2025.

[11] Lin, S., et al. SDXL‑Lightning: Progressive adversarial diffusion distillation. arXiv, 2024.

[12] Zhou, M., et al. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one‑step generation. ICML, 2024.

[13] Lu, C., et al. DPM‑Solver++: Fast solver for guided sampling of diffusion probabilistic models, 2023.

[14] Zhao, W., et al. UniPC: A unified predictor‑corrector framework for fast sampling of diffusion models. arXiv, 2023.