SA-Solver: Stochastic Adams Solver for Fast Sampling of Diffusion Models

TL;DR

This work tackles the slow sampling of diffusion probabilistic models by introducing SA-Solver, a stochastic Adams-based solver for diffusion SDEs that employs variance-controlled SDEs via a tunable noise scale tau(t) and an analytic exponential-integrator framework. It provides convergence guarantees for both predictor and corrector stages and establishes strong connections to existing samplers (DDIM, DPM-Solver++, UniPC) as special cases. Through extensive experiments, SA-Solver achieves state-of-the-art or competitive Fréchet Inception Distance (FID) on multiple benchmarks with few to moderate NFEs, and demonstrates robustness when score estimates are imperfect. The study also highlights the benefit of data-prediction reparameterization and offers practical guidance on choosing stochasticity magnitudes. Overall, SA-Solver advances fast, high-quality diffusion sampling with theoretical guarantees and practical efficacy across vision tasks and beyond.

Abstract

Diffusion Probabilistic Models (DPMs) have achieved considerable success in generation tasks. As sampling from DPMs is equivalent to solving diffusion SDE or ODE which is time-consuming, numerous fast sampling methods built upon improved differential equation solvers are proposed. The majority of such techniques consider solving the diffusion ODE due to its superior efficiency. However, stochastic sampling could offer additional advantages in generating diverse and high-quality data. In this work, we engage in a comprehensive analysis of stochastic sampling from two aspects: variance-controlled diffusion SDE and linear multi-step SDE solver. Based on our analysis, we propose \textit{SA-Solver}, which is an improved efficient stochastic Adams method for solving diffusion SDE to generate data with high quality. Our experiments show that \textit{SA-Solver} achieves: 1) improved or comparable performance compared with the existing state-of-the-art (SOTA) sampling methods for few-step sampling; 2) SOTA FID on substantial benchmark datasets under a suitable number of function evaluations (NFEs). Code is available at https://github.com/scxue/SA-Solver.

Figures (10)

• Figure 1: Sampling quality measured by FID $\downarrow$ of SA-Solver under a different number of function evaluations (NFE), varying the stochastic noise scale $\tau$. For LSUN Bedroom, 10K samples are used to evaluate FID.
• Figure 2: Sampling quality measured by FID $\downarrow$ of different sampling methods of DPMs under different NFEs.
• Figure 3: Qualitative comparisons between our SA-Solver and previous state-of-the-art methods. All images are generated by Stable Diffusion v1.5 with the same random seed. The main part of the prompt is "portrait of curly orange haired mad scientist man". We set the guidance scale as 7.5. The proposed SA-Solver is able to generate images with more details.
• Figure 4: Sampling quality measured by FID $\downarrow$ of different sampling methods of DPMs under different training epochs.
• Figure 5: Samples by DDIM, DPM-Solver, UniPC, EDM(ODE), EDM(SDE) and our SA-Solver with 15, 23, 47, 95 NFEs with the same random seed from CIFAR10 32x32 VE baseline model Karras2022edm

Theorems & Definitions (32)

• Proposition 4.1
• Remark 1
• Proposition 4.2
• Theorem 5.1: Strong Convergence of $s$-step SA-Predictor
• Theorem 5.2: Strong Convergence of $\hat{s}$-step SA-Corrector
• Corollary 5.3: Relationship with DDIM
• proof
• proof
• Proposition A.1
• proof