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

1 Introduction

• DPM sampling can demand hundreds of network function evaluations, creating a substantial computational burden for practical use.Source: §1 · p. 1
• Training-free samplers improve efficiency by numerically solving diffusion SDEs or ODEs, while ODE sampling empirically offers better efficiency than stochastic SDE sampling.Source: §1 · p. 1
• Prior empirical observations suggest stochastic sampling may produce higher-quality data as sampling steps increase, motivating further work on efficient stochastic samplers.Source: §1 · p. 2
• The paper studies noise-scaled diffusion SDEs with shared marginals and linear multi-step SDE solvers as a route to efficient stochastic sampling.It proposes SA-Solver based on the stochastic Adams method, adapted to an exponentially weighted integral and analytical variance.Source: §1 · p. 2
• SA-Solver reports improved or comparable FID under small NFEs, state-of-the-art FID across benchmark datasets at suitable NFEs, and superior performance to deterministic samplers with incompletely trained models.The experiments evaluate image-generation benchmarks using FID across different numbers of function evaluations.Source: §1 · p. 2

2 Related Works

• Training-based acceleration includes distillation, learning-to-sample, and GAN integration, while this paper focuses on training-free methods.The passage characterizes the lack of plug-and-play applicability as limiting the broad use of training-based approaches.Source: §2 · p. 2
• ODE samplers such as DDIM have attracted attention because ODEs are easier to solve than SDEs, and later methods use semi-linear structure or predictor-corrector mechanisms [24].Source: §2 · p. 2
• SDE sampling has potential to generate higher-quality data, but existing samplers are described as limited by low efficiency or hyperparameter sensitivity.The cited literature includes first-order discretization, a higher-order stochastic sampler [25], and a stochastic sampler tailored to diffusion problems.Source: §2 · p. 2
• Compared with SDE-DPM-Solver++ [31], SA-Solver includes predictor-corrector support, proves convergence order, and supports arbitrary predictor orders with analytical forms.The passage describes SDE-DPM-Solver++ as a special case with predictor step 2 and no corrector.Source: §2 · p. 2

3 Preliminary

• Continuous DPMs construct noisy data through a linear SDE, with a Wiener process and drift and diffusion coefficients described in the setup.The coefficients are selected so the terminal distribution approximates a Gaussian and the signal-to-noise ratio decreases over time.Source: §3 · p. 3
• The forward process has an equivalent reverse-time diffusion process, making generation equivalent to numerically solving the diffusion SDE.The reverse-time Wiener process and score function are identified in the setup.Source: §3 · p. 3
• A corresponding deterministic process shares the diffusion process's marginal probability densities, allowing ODE solvers to be used for sampling.Source: §3 · p. 3
• A neural network sθ(x,t) is trained to approximate the score function, and the paper describes noise-prediction and data-prediction reparameterizations for use in sampling.Source: §3 · p. 3

4 Variance Controlled Diffusion SDEs

• Under 20 NFEs, DDIM improves FID over DDPM on CIFAR10 from 133.37 to 6.84, but under 1000 NFEs DDPM beats DDIM, with FID 3.17 versus 4.04.These observations motivate the authors' conjecture that appropriately scaled noise may improve generated-data quality with few NFEs.Source: §4 · p. 3
• For bounded τ(t), the proposed variance-controlled reverse SDEs share marginal probability distributions with the diffusion ODE and SDE, while τ(t) controls the magnitude of added noise.τ(t)=0 and τ(t)=1 correspond respectively to the diffusion ODE and the SDE in.Source: §4 · pp. 3–4
• The authors use data-prediction reparameterization of the score in the SDEs, and report that this choice significantly improves generated-data quality for SA-Solver.The paper focuses on data reparameterization for the remainder of its analysis.Source: §4 · p. 4
• Changing the time variable to log-SNR λ_t gives the SDE an explicit solution because of its semi-linear structure.The solution is described for a given state at time s and a target time t in [0,s].Source: §4 · p. 4
• The numerical sampling procedure solves the resulting diffusion SDE starting from x_T approximated by a Gaussian distribution.Source: §4 · p. 4

5 SA-Solver: Stochastic Adams Method to Solve Diffusion SDEs

• SA-Solver combines predictor and corrector steps, samples the Gaussian stochastic integral directly, and uses stochastic Adams methods to reduce evaluations relative to Stochastic Runge-Kutta schemes.The method estimates the deterministic integral using preceding model evaluations.Source: §5 · p. 5, §5.1 · p. 5

5.1 SA-Predictor

• Under mild regularity conditions, the s-step SA-Predictor has a global strong-convergence result, estimating the integral with Lagrange interpolation of preceding model evaluations.The interpolation uses s prior model evaluations.Source: §5.1 · pp. 5–6

5.2 SA-Corrector

• The ŝ-step SA-Corrector incorporates a model evaluation at the predicted next state and has a global strong-convergence result under mild regularity conditions.It uses Lagrange interpolation, whose order may differ from the predictor's.Source: §5.2 · p. 6

5.3 Connection with other samplers

• DDIM-η coincides with the 1-step SA-Predictor when the variance-control function τ(t) is the corresponding piecewise-constant function in the data-parameterized SDE.Source: §5.3 · p. 7
• With τ(t) ≡ 0, DPM-Solver++(2M) is the 2-step SA-Predictor, while UniPC-p is SA-Solver with equal predictor and corrector steps.Source: §5.3 · p. 7

6.1 Comparison between Data-Prediction Model and Noise-Prediction Model

• Data-prediction achieves better sampling-quality values than noise-prediction across tested NFEs, so the remaining experiments use data-prediction.The comparison uses the ImageNet 256x256 latent diffusion model with τ(t) ≡ 1.Source: §6.1 · p. 7, Table 1 · p. 8

6.2 Ablation Study on Predictor/Corrector Steps and Predictor-Corrector Method

• On CIFAR10, both multi-step predictors and predictor-corrector configurations improve sampler performance over predictor-only configurations.The ablation varies predictor steps and includes comparisons with and without corrector steps.Source: §6.2 · p. 8, Table 2 · p. 8

6.3 Effect on Magnitude of Stochasticity

• The best stochastic-noise scale depends on NFE: smaller nonzero τ performs better at low NFEs, while larger τ helps at 20–100 steps.The experiments vary constant τ from 0.0 to 1.6; Figure 1 reports FID for SA-Solver across NFEs and noise scales, using 10K LSUN Bedroom samples.Source: §6.3 · p. 8, Figure 1 · p. 9

6.4 Comparison with State-of-the-Art

• Across CIFAR10 and ImageNet benchmarks, SA-Solver outperforms existing samplers and reports state-of-the-art results at suitable NFEs.At 95 NFEs, it reaches FID 2.63 on CIFAR10 and 1.81 on ImageNet 64x64; on ImageNet 256x256 it reaches FID 3.87 at 20 NFEs and below 3.5 at 40 NFEs or more. Reported state-of-the-art FIDs include 1.93 on ImageNet 256x256 and 2.80 on ImageNet 512x512.Source: §6.4 · pp. 8–9, Table 3 · p. 9, Figure 2 · p. 10
• For the same Stable Diffusion prompt and random seed, SA-Solver generates text-to-image outputs described as more reasonable and detailed.The comparison uses guidance scale 7.5.Source: §6.4 · p. 9, Figure 3 · p. 10

6.5 Effect of Stochasticity for Inaccurate Score Estimation

• With inaccurate score estimates from limited training data or compute, SA-Solver outperforms deterministic methods, especially early in training, and larger τ gives better performance.Figure 4 compares sampling-method FID across training epochs.Source: §6.5 · p. 9, Figure 4 · p. 10

7 Conclusions

• SA-Solver achieves state-of-the-art sampling performance across various pretrained DPMs and performs well with inaccurate score estimation.The conclusion also reports high sampling performance in both minimal and suitable numbers of steps.Source: §7 · p. 10
• Variance-controlled diffusion SDEs use τ(t) to control injected noise, and SA-Solver adapts the stochastic Adams method using an exponentially weighted integral and analytical variance.The conclusion reports an optimal theoretical convergence bound.Source: §7 · p. 10
• Selecting τ(t) for different NFEs needs further research, including deeper theoretical analysis beyond the paper’s empirical selection criteria.Source: §7 · p. 10

A.1 Proof of Proposition 4.1

• For bounded measurable τ(t), the reverse SDE family shares marginal distributions with Eqs. (4) and (3), with τ(t)=0 and τ(t)=1 recovering the diffusion ODE and SDE.Source: §A.1 · p. 14

A.2 Two Reparameterizations and Exact Solution under Exponential Integrator

• The exponential-integrator analysis compares data and noise prediction reparameterizations and states that the data-prediction noise term has smaller variance.Source: §A.2 · p. 14

A.2.2 Proof of Proposition 4.2

• The Itô integral in the data-prediction solution is Gaussian; its mean is zero by the martingale property, and its variance is computed using Itô isometry.Source: §A.2.2 · pp. 16–17

A.2.3 Noise Prediction Reparameterization

• The noise-prediction formulation also has an analytic solution, presented as Proposition A.1.Source: §A.2.3 · pp. 17–18

A.2.4 Comparison between Data and Noise Reparameterizations

• Under equal stochasticity, data reparameterization gives better convergence and FID; its numerical updates inject smaller noise, which the authors suggest explains the difference.Source: §A.2.4 · p. 19

B.1 Preliminary

• The preliminaries define strong and mean-square convergence and state linear-growth and uniform-Lipschitz assumptions for the data-prediction model.Source: §B.1 · p. 20

B.2 Outline of the Proof

• The convergence proof reduces the s-step corrector rate to predictor and corrector lemmas, targeting O(max τ(t)h + h^(s+1)).Source: §B.2 · p. 21

B.3 Lemmas for the Proof; Then the following s+1 equalities hold

• The proof bounds multiple Wiener integrals under a linear-growth condition and verifies Lagrange-interpolation identities using constant, linear, and higher-degree polynomials.Source: §B.3 · p. 22, Then the following s+1 equalities hold · p. 22

B.4 Proof of Lemma B.7 (for Theorem. 5.1) and Lemma B.8 (for Theorem. 5.2)

• The s-step SA-Predictor has deterministic local error O(h^(s+1)), while the s-step SA-Corrector has deterministic local error O(h^(s+2)).The proof attributes the difference to the corrector using s+1 points, versus s points for the predictor.Source: §B.4 · p. 29

B.5.1 Relationship with DDIM

• For any DDIM η, a piecewise-constant τ(t) makes DDIM-η coincide with the 1-step SA-Predictor under data parameterization.Source: §B.5.1 · p. 29

B.5.2 Relationship with DPM-Solver++(2M)

• DPM-Solver++(2M) is a special case of the 2-step SA-Predictor when τ(t)≡0.Source: §B.5.2 · p. 30

B.5.3 Relationship with UniPC

• UniPC-p is a special case of SA-Solver when τ(t)≡0, with predictor and corrector steps both set to p.Source: §B.5.3 · p. 31

C Selection on the Magnitude of Stochasticity

• Under stated conditions, the appendix examines choosing τ(t)≡1 because it can minimize an upper bound on KL divergence to the true distribution.Source: C Selection on the Magnitude of… · p. 31

D Implementation Details

• The implementation uses a modified coefficient whose O(h^3) difference preserves the 2-step predictor’s convergence order; the 1-step corrector detail is analogous.Source: D Implementation Details · p. 32

E.1 Details on τ(t), Predictor Steps and Corrector Steps

• CIFAR10 uses piecewise-constant τ(t); predictor and corrector step counts vary in Section 6.2, but are three each in Section 6.5; other experiments use constant τ(t).For other experiments, the solver uses three predictor and three corrector steps under 20 NFEs, then two predictor and one corrector step beyond 20 NFEs.Source: §E.1 · p. 33

E.2 Details on Pretrained Models and Settings

• CIFAR10 uses the pretrained unconditional VE model from with σmin=0.02 and σmax=80; EDM reports its lowest observed FID, while SA-Solver and UniPC report averaged FID.Source: §E.2 · p. 33
• ImageNet 64x64 uses a conditional VP-cosine model from with the cited schedule; EDM reports its lowest observed FID, while compared samplers report averaged FID.Source: §E.2 · p. 33
• ImageNet 256x256 evaluates LDM, DiT-XL/29, and Min-SNR with classifier-free guidance scale 1.5; ImageNet 512x512 uses DiT-XL/2 with the same guidance scale.Source: §E.2 · p. 34
• LSUN Bedroom 256x256 uses the pretrained model from, unconditional sampling, and a uniform lambda step schedule from.Source: §E.2 · p. 34

• 3.87 FID with 20 NFEs; SA-Solver outperforms all ODE samplers even with far more steps.Source: F Additional Results · p. 34
• Stochasticity improves sample quality; small magnitudes suit small NFEs, while larger magnitudes help at large NFEs.Source: F Additional Results · p. 34
• In latent space, SDE converges faster; at 10 NFEs, τ = 0.6 is better than τ = 0.Source: F Additional Results · p. 34
• Supplementary tables report FID across NFE settings on CIFAR10 VE-baseline and ImageNet 64x64, with separate stochasticity-magnitude ablations on both datasets.Source: Table 4 · p. 34, Table 5 · p. 34, Table 6 · p. 35

• SA-Solver generates more diverse samples with more details in the additional examples.Source: G Additional Samples · p. 35
• Supplementary FID tables cover CIFAR10 and ImageNet 256x256, including sampling-method and training-epoch comparisons and NFE variation.Source: Table 8 · p. 36, Table 9 · p. 36, Table 10 · p. 36
• Separate ablation tables vary stochasticity magnitude and NFEs on CIFAR10, ImageNet 64x64, ImageNet 256x256, and LSUN Bedroom.Source: Table 11 · p. 36, Table 12 · p. 37, Table 13 · p. 37
• Figures 5–6 compare samplers at 15, 23, 47, and 95 NFEs using the same seed for CIFAR10 VE-baseline and conditional ImageNet 64x64.Source: Figure 5 · p. 38, Figure 6 · p. 39
• Figures 7–8 vary stochasticity for LSUN Bedroom at 20, 40, 60, and 100 NFEs and for ImageNet 512x512 DiT at 60 NFEs.Source: Figure 7 · p. 40, Figure 8 · p. 41
• Figures 9–10 compare solvers and NFEs on Stable-Diffusion v1.5 text-to-image samples using landscape and portrait prompts.Source: Figure 9 · p. 42, Figure 10 · p. 43