Title: SA-Solver: Stochastic Adams Solver for Fast Sampling of Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion Model Sampler

Multi-step SDE Solver
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 SA-Solver, which is an improved efficient stochastic Adams method for solving diffusion SDE to generate data with high quality. Our experiments show that 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).

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

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 SA-Solver , which is an improved efficient stochastic Adams method for solving diffusion SDE to generate data with high quality. Our experiments show that 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).

1 Introduction

Diffusion Probabilistic Models (DPMs) [1-3] have demonstrated substantial success across a broad spectrum of generative tasks such as image synthesis [4-6], video generation [7, 8], text-to-image generation [9-11], speech synthesis [12, 13], etc . The primary mechanism of DPMs involves a forward diffusion process that incrementally introduces noise into data. Simultaneously, a reverse diffusion process is learned to generate data from this noise. Despite DPMs demonstrating enhanced generation performance in comparison to alternative methods such as Generative Adversarial Networks (GAN) [14] or Variational Autoencoders (VAE) [15], the sampling process of DPMs demand hundreds of evaluations of network function evaluations (NFE) [2]. The substantial computation requirement poses a significant limitation to their wider application in practice.

The existing literature on improving the sampling efficacy of DPMs can be categorized into two ways, depending on whether conducting extra training on the DPMs. The first category necessitates supplementary training [16-21], which often emerges as a bottleneck, thereby limiting their practical application. Due to this, we focus on exploring the second category, which consists training-free methods to improve the sampling efficiency of DPMs in this paper. Current training-free samplers employ efficient numerical schemes to solve the diffusion SDE/ODE[22-26]. Compared with solving diffusion SDE (stochastic sampler) [25-27], solving diffusion ODE (deterministic sampler) [22-24]

empirically exhibits better sampling efficiency. Existing stochastic samplers typically exhibit slower convergence speed. However, empirical observations in [3, 27] indicate that the stochastic sampler has the potential to generate higher-quality data when increasing the sampling steps. This empirical observation motivates us to further explore the efficient stochastic sampler.

Owing to the observed superior performance of stochastic sampler [3, 27], we speculate that adding properly scaled noise in the diffusion SDE may facilitate the quality of generated data. Thus, instead of solving the vanilla diffusion SDE in [22], we propose to consider a family of diffusion SDEs which shares the same marginal distribution [28, 27] with different noise scales. Meanwhile, efficient stochastic solvers are not carefully studied, which could be the reason that diffusion ODE exhibits better sampling efficiency. To overcome this problem, we study the linear multi-step SDE solvers [29] and incorporate them in the sampling.

Based on these studies, we propose SA-Solver with theoretical convergence order to solve the proposed diffusion SDEs. Our SA-Solver is based on the stochastic Adams method in [29], by adapting it to the exponentially weighted integral and analytical variance. With the proposed diffusion SDEs and SA-Solver , we can efficiently generate data with controllable noise scales. We empirically evaluate our SA-Solver on plenty of benchmark datasets of image generation. The evaluation criterion is the Fréchet Inception Distance (FID) score [30] under different number of function evaluations (NFEs). The experimental results can be summarized as three folds: 1) Under small NFEs, our SA-Solver has improved or comparable FID scores, compared with baseline methods; 2) Under suitable NFEs our SA-Solver achieves the State-of-the-Art FID scores over all benchmark datasets; 3) SA-Solver achieves superior performance over deterministic samplers when the model is not fully trained.

2 Related Works

The DPMs originate from the milestone work [1], and are further developed by [2] and [3] to successfully generate high-quality data, under the framework of discrete and continuous diffusion SDEs respectively. In this paper, we mainly focus on the latter framework. As mentioned in Section 1, plenty of papers are working on accelerating the sampling of DPMs due to their low efficiency, distinguished by whether conducting a supplementary training stage. Training-based methods, e.g., knowledge distillation [16-18], learning-to-sample [19], and integration with GANs [20, 21], have the potential to sampling for one or very few steps to enhance the efficiency, but their applicability is limited by the lack of a plug-and-play nature, thereby constraining their broad applicability across diverse tasks. Thus we mainly focus on the training-free methods in this paper.

Solving Diffusion ODE. Since the sampling process is equivalent to solving diffusion SDE (ODE), the training-free methods are mainly built on solving the differential equations via high-efficiency numerical methods. As ODEs are easier to solve compared with SDEs, the ODE sampler has attracted great attention. For example, Song et al. [22] provides an empirically efficient solver DDIM. Zhang and Chen [28] and Lu et al. [23] point out the semi-linear structure of diffusion ODEs, and develop higher-order ODE samplers based on it. Zhao et al. [24] further improve these samplers in terms of NFEs by integrating the mechanism of predictor-corrector method.

Solving Diffusion SDE. Though less explored than the ODE sampler, the SDE sampler exhibits the potential of generating higher-quality data [27]. Thus developing an efficient SDE sampler as we did in this paper is a meaningful topic. In the existing literature, researchers [2, 26, 3] solve the diffusion SDE by first-order discretization numerical method. The higher-order stochastic sampler of diffusion SDE has also been discussed in [25]. Karras et al. [27] proposes another stochastic sampler (which is not a general SDE numerical solver) tailored for diffusion problems. However, in contrast to our proposed SA-Solver , the existing SDE samplers are limited due to their low efficiency [2, 26, 3] or sensitivity to hyperparameters [27].

We found a concurrent paper proposing an SDE sampler SDE-DPM-Solver++ [31] which is similar to our SA-Solver . Though both methods develop multi-step diffusion SDE samplers, our SA-Solver is different from SDE-DPM-Solver++ as follows: 1) SA-Solver incorporates the predictor-corrector method, which helps improve the quality of generated data [3, 32, 24]; 2) In contrast to SDEDPM-Solver++, SA-Solver has theoretical guarantees with proved convergence order; 3) SDE-DPMSolver++ is a special case of SA-Solver when the predictor step equals 2 with no corrector in our predictor-corrector method, while our solver supports arbitrary orders with analytical forms.

3 Preliminary

In the regime of the continuous stochastic differential equation (SDE), Diffusion Probabilistic Models (DPMs) [1-3, 33] construct noisy data through the following linear SDE:

where w t ∈ R d represents the standard Wiener process, f ( t ) x t and g ( t ) respectively denote the drift and diffusion coefficients. For each time t ∈ [0 , T ] , x t | x 0 ∼ N ( α t x 0 , σ 2 t I ) .

Let p t ( x ) denotes the marginal distribution of x t , the coefficients f ( t ) and g ( t ) are meticulously selected to guarantee that the marginal distribution p T ( x T ) closely approximates a Gaussian distribution, i.e., N ( 0 , I ) , and the signal-to-noise-ratio (SNR) α 2 t /σ 2 t is strictly decreasing w.r.t. t . In the sequel, we follow the established notations in [33]:

Anderson [34] demonstrates a pivotal theorem that the forward process (1) has an equivalent reversetime diffusion process (from T to 0 ) as the following equation, so that generating process can be equivalent to numerically solve the diffusion SDE [2, 3].

where ¯ w t represents the Wiener process in reverse time, and ∇ x log p t ( x ) is the score function.

Moreover, Song et al. [3] also prove that there exists a corresponding deterministic process whose trajectories share the same marginal probability densities p t ( x ) as (3), so that the ODE solver can be adopted for efficient sampling [23, 24]:

To get the score function ∇ x log p t ( x t ) in (3), we usually take neural network s θ ( x , t ) parameterized by θ to approximate it by optimizing the denoising score matching loss [3]:

In practice, two methods are used to reparameterize the score-based model [35]. The first approach utilizes a noise prediction model such that ϵ θ ( x t , t ) = -σ t s θ ( x t , t ) , while the second employs a data prediction model , represented by x θ ( x t , t ) = ( x t -σ t ϵ θ ( x t , t )) /α t . The reparameterized models are plugged into the sampling process (3) or (4) according to their relationship with s θ ( x t , t ) .

4 Variance Controlled Diffusion SDEs

As mentioned in Section 1, most of the existing training-free efficient samplers are based on solving diffusion ODE (4), e.g., [23, 22, 24], because of their improved efficiency compared with the solvers of diffusion SDE (3). However, the empirical observations in [27, 22] exhibit that the quality of data generated by solving diffusion SDE outperforms diffusion ODE given sufficient computational budgets. For example, in [22], the diffusion ODE sampler DDIM [22] significantly improve the FID score of diffusion SDE sampler DDPM [2] (from 133.37 to 6.84) on CIFAR10 dataset [36] under 20 NFEs. However, under 1000 NFEs, the DDPM beats the DDIM in terms of FID score (3.17 v.s. 4.04). There may be a trade-off between stochasticity and efficiency. Thus, we conjecture that adding proper scale noise during the generating process may improve the quality of generated data with few NFEs.

In this section, we explore a family of variance-controlled diffusion SDEs, so that we can use proper noise scales during the sampling stage. Inspired by Proposition 1 in [28] and Eq. (6) in [27], we propose the following proposition to construct the aforementioned diffusion SDEs.

Proposition 4.1. For any bounded measurable function τ ( t ) : [0 , T ] → R , the following Reverse SDEs

share the same marginal probability distributions with (4) and (3) .

The proof can be found in Appendix A.1. The proposition indicates that by solving any of the diffusion SDEs in (6), we can sample from the target distribution. It is worth noting that the magnitude of noise varies with τ ( t ) , and τ ( t ) = 0 or τ ( t ) = 1 respectively correspond to the diffusion ODE and SDE in [3]. Thus we can control the magnitude of added noise during the sampling process by varying it.

In practice, we numerically solve the diffusion SDEs (6) by substituting score function ∇ x log p t ( x t ) in it with the 'data prediction reparameterization model' x θ ( x t , t ) according to ∇ x log p t ( x t ) ≈ -( x t -α t x θ ( x t , t )) /σ 2 t as pointed out in Section 3. Then diffusion SDEs to be solved become

Remark 1. We reparameterize the score function in diffusion SDEs (6) with data prediction model x θ ( x t , t ) to get Eq. (9) . The equation can be also reparameterized by the 'noise prediction model' ϵ θ ( x t , t ) as discussed in Section 3. Though the obtained diffusion ODEs e.g., Eq. (9) are equivalent, the numerical solver applied to them will result in different solutions. For our proposed SA-Solver, we find the diffusion SDEs reparameterized data prediction model significantly improves the quality of generated data. More details and theoretical explanations are in Sec. 6 and Appendix A.2.4. For the remaining part of the paper, we focus on data reparameterization.

We then solve the diffusion SDEs (9) with change-of-variable applying to it, i.e., changing time variable t to log-SNR λ t = log ( α t /σ t ) . Noting the following relationship in Eq. (2)

and plugging them into (7), it becomes

The above equation has an explicit solution owing to its semi-linear structure [37].

Proposition 4.2. Given x s for any time s > 0 , the solution x t at time t ∈ [0 , s ] of (9) is

where G ( s, t ) is an Itô integral [38] with the special property

The proof can be seen in Appendix A.2.2. With this proposition, we can sample from the diffusion model via numerically solving Eq. (10) starting from x T approximated by a Gaussian distribution.

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

Stochastic training-free samplers for solving diffusion SDEs have not been studied as systematically as their deterministic ODE counterparts. This stems from the inherent challenges associated with designing numerical schemes for SDEs compared to ODEs [39]. Existing stochastic sampling methods either use only variant of one-step discretization of diffusion SDEs [2, 26, 3], or are specifically designed sampling procedures for diffusion processes [27] which are not general purpose SDE solvers. Jolicoeur-Martineau et al. [25] uses stochastic Improved Euler's method [40] with adaptive step sizes. However, it still necessitates hundreds of steps to yield a high-quality sample. As observed by [25], off-the-shelf SDE solvers are generally ill-suited for diffusion models, often exhibiting inferior qualities or even failing to converge. We postulate that the current dearth of fast stochastic samplers is principally due to factor that existing methodologies predominantly tend to rely

Algorithm 1 SA-Solver

Require: data prediction model x θ , timesteps { t i } M i =0 , initial value x t 0 , predictor step s p , corrector step s c , buffer B to store former evaluation of x θ , τ ( t ) to control variance.

on one-step discretization or its variants, or alternatively, on heuristic designs of stochastic samplers. To address this factor, we leverage advanced contemporary tools in numerical solutions for SDEs, specifically, stochastic Adams methods [29]. It necessitates fewer evaluations compared to Stochastic Runge-Kutta schemes, making it a more suitable choice for problems which are computationally expensive - a characteristic that diffusion sampling certainly exemplifies.

Next, we formally present our Stochastic Adams Solver ( SA-Solver ). To solve Eq. (9), we first take M +1 time steps { t i } M i =0 which is strictly decreased from t 0 = T to t M = 0 . 3 Then we can iteratively obtain the x t i (so that x 0 approximates the required data) by the following relationship.

As pointed out in Proposition 4.2, the Itô integral term G ( t i , t i +1 ) in above equation follows a Gaussian that can be directly sampled so we need to solve the deterministic integral term F θ ( t i , t i +1 ) .

We further combine Eq. (12) with the predictor-corrector method, which is a widely used numerical method. It works in two main steps. First, a predictor step is taken to make an initial approximation of the solution. Second, a corrector step will refine the predictor's approximation by taking the predicted value into account. It has been proven successful in the wide application of numerical analysis [37]. Especially, there are some attempts to use the predictor-corrector method to help sample diffusion models [3, 32, 24]. In the subsequent Section 5.1 and Section 5.2, we will separately derive our SA-Predictor and SA-Corrector using Eq. (12). Our algorithm is outlined in Algorithm 1.

5.1 SA-Predictor

The fundamental idea behind stochastic Adams methods is to leverage preceding model evaluations like x θ ( x t i , t i ) , x θ ( x t i -1 , t i -1 ) , · · · , x θ ( x t i -( s -1) , t i -( s -1) ) . These evaluations can be retained with negligible cost implications. Given these preceding model evaluations, a natural strategy for estimating F θ ( t i , t i +1 ) involves the application of Lagrange interpolations [37] of these evaluations. Lagrange interpolation of s points x θ ( x t i , t i ) , x θ ( x t i -1 , t i -1 ) , · · · , x θ ( x t i -( s -1) , t i -( s -1) ) is a polynomial L ( t ) with degrees s -1 :

where l i -j ( t ) : R → R is the Lagrange basis. Lagrange interpolation is an excellent approximation of x θ ( x t , t ) with the special property: L ( t i -j ) = x θ ( x t i -j , t i -j ) , ∀ 0 ≤ j ≤ s -1 . Thus a natural

way to estimate F θ ( t i , t i +1 ) is to replace x θ ( x λ u , λ u ) with L ( λ ) , which is just a change-of-variable of L ( t ) . The formula for s -step SA-Predictor is then derived.

s -step SA-Predictor Given the initial value x t i at time t i , a total of s former model evaluations x θ ( x t i , t i ) , x θ ( x t i -1 , t i -1 ) , · · · , x θ ( x t i -( s -1) , t i -( s -1) ) , our s -step SA-Predictor is defined as:

where ˜ σ i = σ t i +1 √ 1 -e -2 ∫ λt i +1 λt i τ 2 ( ˜ λ )d ˜ λ according to Proposition 4.2 and b i -j is given by:

We show the convergence result in the following theorem. The proof can be found in Appendix B.

Theorem 5.1 (Strong Convergence of s -step SA-Predictor ) . Under mild regularity conditions, our s -step SA-Predictor (Eq. (14) ) has a global error in strong convergence sense of O ( sup 0 ≤ t ≤ T τ ( t ) h + h s ) ,

5.2 SA-Corrector

Eq. (14) offers a 'prediction' x p t i +1 that relies on information preceding or coinciding with the time step t i since we only use x θ ( x t i , t i ) along with other model evaluations antecedent to it, while the integration is over time [ t i , t i +1 ] . Then predictor-corrector method can be incorporated to better estimate F θ ( t i , t i +1 ) in Eq. (12). We perform a model evaluation x θ ( x p t i +1 , t i +1 ) and construct the Lagrange interpolations of x θ ( x p t i +1 , t i +1 ) , x θ ( x t i , t i ) , · · · , x θ ( x t i -( s ′ -1) , t i -(ˆ s -1) ) :

where ˆ l i -j ( t ) : R → R is the Lagrange basis and ˆ s can be different with s in Eq. (13). The ˆ s -step SA-Corrector is derived by replacing x θ ( x λ u , λ u ) with ˆ L ( λ ) which is a change-of-variable of ˆ L ( t ) .

ˆ s -step SA-Corrector Given the initial value x t i at time t i , a total of ˆ s former model evaluations x θ ( x t i , t i ) , x θ ( x t i -1 , t i -1 ) , · · · , x θ ( x t i -(ˆ s -1) , t i -(ˆ s -1) ) , model evaluation of 'prediction' x θ ( x p t i +1 , t i +1 ) , our ˆ s -step SA-Corrector is defined as:

where ξ ∼ N ( 0 , I ) , ˜ σ i = σ t i +1 √ 1 -e -2 ∫ λt i +1 λt i τ 2 ( ˜ λ )d ˜ λ according to Proposition 4.2 and the coefficients ˆ b i +1 , ˆ b i -j is given by:

We show the convergence result in the following theorem. The proof can be found in Appendix B. Theorem 5.2 (Strong Convergence of ˆ s -step SA-Corrector ) . Under mild regularity conditions, our ˆ s -step SA-Corrector (Eq. (17) ) has a global error in strong convergence sense of O ( sup 0 ≤ t ≤ T τ ( t ) h + h ˆ s +1 ) ,

5.3 Connection with other samplers

We briefly discuss the relationship between our SA-Solver and other existing solvers for sampling diffusion ODEs or diffusion SDEs.

Relationship with DDIM [22] DDIM generate samples through the following process:

where ξ ∼ N ( 0 , I ) , ˆ σ t i is a variable parameter. In practice, DDIM introduces a parameter η such that when η = 0 , the sampling process becomes deterministic and when η = 1 , the sampling process

Corollary 5.3 (Relationship with DDIM) . For any η in DDIM, there exists a τ η ( t ) : R → R which is a piecewise constant function such that DDIMη coincides with our 1 -step SA-Predictor when τ ( t ) = τ η ( t ) with data parameterization of our variance-controlled diffusion SDE.

The proof can be found in Appendix B.5.1.

Relationship with DPM-Solver++(2M) [31] DPM-Solver++ is a high-order solver which solves diffusion ODEs for guided sampling. DPM-Solver++(2M) is a special case of our 2 -step SA-Predictor when τ ( t ) ≡ 0 .

Relationship with UniPC [24] UniPC is a unified predictor-corrector framework for solving diffusion ODEs. UniPC-p is a special case of our SA-Solver when τ ( t ) ≡ 0 with predictor step p , corrector step p in Algorithm 1.

6 Experiments

In this section, we demonstrate the effectiveness of SA-Solver over the existing sampling methods on both a small number of function evaluations (NFEs) settings and a considerable number of NFEs settings, with extensive experiments. We use Fenchel Inception Distance (FID) [30] as the evaluation metric to show the effectiveness of SA-Solver . Unless otherwise specified, 50K images are sampled for evaluation. The experiments are conducted on various datasets, with image sizes ranging from 32x32 to 256x256. We also evaluate the performance of various models, including ADM [4], EDM [27], Latent Diffusion [5], and DiT [41].

For ease of computation, we take τ ( t ) ≡ τ as a constant function or a piecewise constant function. We leave the detailed settings for τ ( t ) , predictor step, and corrector step in Appendix E. For the following experiments, we first discuss the effectiveness of the data-prediction model. Then we evaluate the performance of SA-Solver under different random noise scales τ to demonstrate the principles for selecting τ under few-steps and a considerable number of steps. Finally, we compare SA-Solver with the existing solver to demonstrate its effectiveness.

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

We first discuss the necessity of using a data-prediction model for SA-Solver . We test on ImageNet 256x256 (latent diffusion model) with τ ( t ) ≡ 1 . Results of the data-prediction and noise-prediction model are shown in Table 1. It can be seen that the data-prediction model can achieve better sampling quality values under different NFEs, thus we use the data-prediction model in the rest of the experiments. More detailed discussions and theoretical analysis can be seen in Appendix A.2.4.

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

To verify the effectiveness of our proposed Stochastic Linear Multi-step Methods and PredictorCorrector Method, we conduct an ablation study on the CIFAR10 dataset as follows. We use EDM [27] baseline-VE pretrained checkpoint. Concretely, we vary the number of predictor steps

and meanwhile conduct them with/without corrector to separately explore the effect of the two components. As can be seen in Table 2, both Stochastic Linear Multi-step Methods (Predictor 1-steps only v.s. Predictor 3-steps only) and Predictor-Corrector Method (Predictor 1-steps only v.s. Predictor 1-steps, Corrector 1-step, and Predictor 3-steps only v.s. Predictor 3-steps, Corrector 3-steps) improve the performance of our sampler.

6.3 Effect on Magnitude of Stochasticity

The proposed SA-Solver is evaluated on various types of datasets and models, including ImageNet 256x256 [43] (latent diffusion model [5]), LSUN Bedroom 256x256 [44] (pixel diffusion model [4]), ImageNet 64x64 (pixel diffusion model [4]), and CIFAR10 32x32 (pixel diffusion model [27]). The models corresponding to these datasets cover pixel-space and latent-space diffusion models, with unconditional, conditional, and classifier-free guidance settings ( s = 1 . 5 in ImageNet 256x256).

We used different constant τ values for SA-Solver , namely { 0 . 0 , 0 . 2 , 0 . 4 , ..., 1 . 6 } , where larger value of τ correspond to larger magnitude of stochasticity. The FID results under different NFE and τ values are shown in Fig. 1. Note that for LSUN Bedroom, 10K images are sampled for evaluation. The experiments indicate that (1) under relatively small NFEs, smaller nonzero τ values can achieve better FID results; (2) under a considerable number of steps (20-100), large τ can achieve better FID. This phenomenon is consistent with the theoretical analysis we conducted in Appendix B and Appendix C, in which the sampling error with stochasticity is dominated under small NFE, while larger τ can significantly improve the quality of generated samples as the number of steps increases. In subsequent experiments, unless otherwise specified, we will report the results of a proper τ ( t ) value. Details can be found in Appendix E.

6.4 Comparison with State-of-the-Art

We compare SA-Solver with existing state-of-the-art sampling methods, including DDIM [22], DPMSolver [23], UniPC [24], Heun sampler and stochastic sampler in EDM [27]. Unless otherwise specified, the methods are tested using the default hyper-parameters in the original papers or code.

Results on CIFAR10 32x32 and ImageNet 64x64. We use the EDM [27] baseline-VE model for the CIFAR10 32x32 experiments and the ADM [4] model for the ImageNet 64x64 experiments. We use EDM's timesteps selection for all samplers for fair comparisons. EDM introduces a certain type of SDE and a corresponding stochastic sampler, which is used for comparison. The experimental results are shown in Fig. 2(a-b). It can be seen that the proposed SA-Solver consistently outperforms other samplers and achieves state-of-the-art FID results. It should be noticed for EDM samplers, we report its optimal result which is searched over four hyper-parameters. In fact, at 95 NFEs, SA-Solver can achieve the best FID value of 2.63 in CIFAR10 and 1.81 in ImageNet 64x64 which outperforms all other samplers.

0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

3.0

3.5

4.0

5.0

6.0

FID

NFE = 23

NFE = 31

NFE = 47

NFE = 63

NFE = 95

0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

2.0

2.5

3.0

4.0

5.0

FID

(b) ImageNet 64x64 (Pixel DPM)

NFE = 20

NFE = 40

NFE = 60

NFE = 80

NFE = 100

0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

2.5

3.0

3.5

4.0

5.0

6.0

FID-10K

(c) LSUN Bedroom 256x256 (Pixel DPM)

NFE = 20

NFE = 40

NFE = 60

NFE = 80

NFE = 100

0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

3.0

3.5

4.0

4.5

5.0

6.0

7.0

FID

(d) ImageNet 256x256 (Latent DPM)

NFE = 10

NFE = 20

NFE = 40

NFE = 60

NFE = 80

NFE = 100

Results on ImageNet 256x256 and 512x512. We evaluate with two classifier-free guidance models, one is the UNet-based latent diffusion model [5] in which the VQ-4 encoder-decoder model is adopted, and the other is the DiT [41] model using Vision Transformer based model with KL-8 encoder-decoder. The corresponding classifier-free guidance scale, namely s = 1 . 5 , is adopted for evaluation. For ImageNet 256x256 dataset with UNet based latent diffusion model, the results of different samplers are shown in Fig. 2(c). Under a considerable number of steps, SA-Solver achieves the best sample quality, in which the FID value is 3.87 with only 20 NFEs and below 3.5 with 40 NFEs or more. While for ODE solvers, the FID values cannot reach below 4, which shows the superiority of the proposed SDE solver.

Table 3 consists results of current SOTA models in ImageNet 256x256 and 512x512. Note that the (Min-SNR) DiT-XL/2-G models are adopted [41, 42]. It can be seen clearly that better FID results are achieved compared with baseline solvers used by corresponding methods. We achieve 1.93 FID value in Min-SNR DiT model at ImageNet 256x256, and 2.80 in DiT model at ImageNet 512x512, both of which are state-of-the-art results under existing DPMs.

Results of text-to-image generation Fig. 3 shows the qualitative results on text-to-image generation. It can be seen that both UniPC and SA-Solver can generate images with more details. Our SA-Solver is able to generate more reasonable images with better details.

6.5 Effect of Stochasticity for Inaccurate Score Estimation

When the training data is not enough or the computational budget is limited, the estimated score is inaccurate. We empirically observed that the stochasticity significantly improve the sample quality under the circumstance. To further investigate this effect, we reproduce the early training stage of EDM [27] baseline-VE model for the CIFAR10 32x32 dataset and DiT-XL/2 [41] model for the ImageNet 256x256 dataset. We compare SA-Solver with different stochastic level τ and existing stateof-the-art deterministic sampling methods. We use the same hyper-parameters as the corresponding experiment in section 6.4.

Figure 4 shows that SA-Solver outperforms deterministic sampling methods, especially in the early stage of the training process. Moreover, larger τ value results in better performance. We also conduct a theoretical analysis that stochasticity can mitigate the error of estimation (see Appendix C).

1216

24

32

48

64

96

NFE

2.5

3.0

3.5

4.0

5.0

6.0

10.0

FID

SA-Solver (ours)

UniPC

EDM (SDE)

EDM (ODE)

DPM-Solver

DDIM(

=0)

16

24

32

48

64

96

NFE

1.8

2.0

2.5

3.0

4.0

6.0

10.0

FID

(b) ImageNet 64x64 (Pixel DPM)

SA-Solver (ours)

DDIM(

=0)

DPM-Solver

UniPC

EDM (SDE)

EDM (ODE)

5 10

20

40

60

80

100

NFE

3.0

3.5

4.0

6.0

10.0

FID

(c) ImageNet 256x256 (Latent DPM)

SA-Solver (ours)

DDIM(

=0)

DDIM(

=1)

DPM-Solver

UniPC

NFE=50

NFE=50

NFE=50

DDIM

UniPC

SA-Solver(Ours)

1250

1300

1350

1400

1450

1500

Epoch

3

5

8

10

20

40

FID

SA-Solver(

=0.6)

SA-Solver(

=1.0)

DDIM

DPM-Solver

EDM(ODE)

50

100

150

200

250

Epoch

4

5

6

8

10

20

FID

(b) ImageNet 256x256 (Latent DPM)(NFE=40)

SA-Solver(

=0.4)

SA-Solver(

=0.8)

DDIM

DPM-Solver

7 Conclusions

In this paper, we propose an efficient solver named SA-Solver for solving Diffusion SDEs, achieving high sampling performance in both minimal steps and a suitable number of steps. To better control the scale of injected noise, we propose Variance Controlled Diffusion SDEs based on noise scale function τ ( t ) and propose the analytic form of the SDEs. Based on Variance Controlled Diffusion SDE, we propose SA-Solver , which is derived from the stochastic Adams method and uses exponentially weighted integral and analytical variance to achieve efficient SDE sampling. Meanwhile, SA-Solver has the optimal theoretical convergence bound. Experiments show that SA-Solver achieves state-ofthe-art sampling performance in various pre-trained DPMs models. Moreover, SA-Solver achieves superior performance when the score estimation is inaccurate.

Although SA-Solver achieves optimal sampling performance, the noise scale τ ( t ) selection under different NFEs needs further research. The paper proposes empirical criteria for selecting τ ( t ) , more in-depth theoretical analysis is still needed.

References

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli, 'Deep unsupervised learning using nonequilibrium thermodynamics,' in International Conference on Machine Learning . PMLR, 2015, pp. 2256-2265.

J. Ho, A. Jain, and P. Abbeel, 'Denoising diffusion probabilistic models,' in Advances in Neural Information Processing Systems , vol. 33, 2020, pp. 6840-6851.

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, 'Score-based generative modeling through stochastic differential equations,' in International Conference on Learning Representations , 2021.

P. Dhariwal and A. Q. Nichol, 'Diffusion models beat GANs on image synthesis,' in Advances in Neural Information Processing Systems , vol. 34, 2021, pp. 8780-8794.

R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, 'High-resolution image synthesis with latent diffusion models,' in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , June 2022, pp. 10 684-10 695.

J. Ho, C. Saharia, W. Chan, D. J. Fleet, M. Norouzi, and T. Salimans, 'Cascaded diffusion models for high fidelity image generation,' Journal of Machine Learning Research , vol. 23, no. 47, pp. 1-33, 2022.

J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet, 'Video diffusion models,' in Advances in Neural Information Processing Systems , 2022.

A. Blattmann, R. Rombach, H. Ling, T. Dockhorn, S. W. Kim, S. Fidler, and K. Kreis, 'Align your latents: High-resolution video synthesis with latent diffusion models,' in IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2023.

A. Q. Nichol, P. Dhariwal, A. Ramesh, P. Shyam, P. Mishkin, B. McGrew, I. Sutskever, and M. Chen, 'GLIDE: towards photorealistic image generation and editing with text-guided diffusion models,' in International Conference on Machine Learning (ICML), 2022 .

A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen, 'Hierarchical text-conditional image generation with clip latents,' 2022.

C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, J. Ho, D. J. Fleet, and M. Norouzi, 'Photorealistic text-to-image diffusion models with deep language understanding,' in Advances in Neural Information Processing Systems , 2022.

M. W. Lam, J. Wang, D. Su, and D. Yu, 'Bddm: Bilateral denoising diffusion models for fast and high-quality speech synthesis,' in International Conference on Learning Representations , 2022.

K. Song, Y. Leng, X. Tan, Y. Zou, T. Qin, and D. Li, 'Transcormer: Transformer for sentence scoring with sliding language modeling,' in Advances in Neural Information Processing Systems , 2022.

I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. C. Courville, and Y. Bengio, 'Generative adversarial nets,' in Advances in Neural Information Processing Systems , vol. 27, 2014, pp. 2672-2680.

D. P. Kingma and M. Welling, 'Auto-encoding variational bayes,' in International Conference on Learning Representations , 2014.

E. Luhman and T. Luhman, 'Knowledge distillation in iterative generative models for improved sampling speed,' arXiv preprint arXiv:2101.02388 , 2021.

T. Salimans and J. Ho, 'Progressive distillation for fast sampling of diffusion models,' in International Conference on Learning Representations , 2022.

C. Meng, R. Gao, D. P. Kingma, S. Ermon, J. Ho, and T. Salimans, 'On distillation of guided diffusion models,' in NeurIPS 2022 Workshop on Score-Based Methods , 2022.

D. Watson, W. Chan, J. Ho, and M. Norouzi, 'Learning fast samplers for diffusion models by differentiating through sample quality,' in International Conference on Learning Representations , 2022.

Z. Xiao, K. Kreis, and A. Vahdat, 'Tackling the generative learning trilemma with denoising diffusion GANs,' in International Conference on Learning Representations , 2022.

Z. Wang, H. Zheng, P. He, W. Chen, and M. Zhou, 'Diffusion-GAN: Training GANs with diffusion,' in The Eleventh International Conference on Learning Representations , 2023.

J. Song, C. Meng, and S. Ermon, 'Denoising diffusion implicit models,' in International Conference on Learning Representations , 2021.

C. Lu, Y. Zhou, F. Bao, J. Chen, C. LI, and J. Zhu, 'Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps,' in Advances in Neural Information Processing Systems , S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, Eds., vol. 35, 2022, pp. 5775-5787.

W. Zhao, L. Bai, Y. Rao, J. Zhou, and J. Lu, 'Unipc: A unified predictor-corrector framework for fast sampling of diffusion models,' arXiv preprint arXiv:2302.04867 , 2023.

A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, 'Gotta go fast when generating data with score-based models,' arXiv preprint arXiv:2105.14080 , 2021.

F. Bao, C. Li, J. Zhu, and B. Zhang, 'Analytic-DPM: An analytic estimate of the optimal reverse variance in diffusion probabilistic models,' in International Conference on Learning Representations , 2022.

T. Karras, M. Aittala, T. Aila, and S. Laine, 'Elucidating the design space of diffusion-based generative models,' in Proc. NeurIPS , 2022.

Q. Zhang and Y. Chen, 'Fast sampling of diffusion models with exponential integrator,' in The Eleventh International Conference on Learning Representations , 2023.

E. Buckwar and R. Winkler, 'Multistep methods for sdes and their application to problems with small noise,' SIAM Journal on Numerical Analysis , 2006.

M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, 'GANs trained by a two time-scale update rule converge to a local Nash equilibrium,' in Advances in Neural Information Processing Systems , I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, Eds., vol. 30, 2017, pp. 6626-6637.

C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu, 'Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models,' 2023.

S. Li, L. Liu, Z. Chai, R. Li, and X. Tan, 'Era-solver: Error-robust adams solver for fast sampling of diffusion probabilistic models,' 2023.

D. P. Kingma, T. Salimans, B. Poole, and J. Ho, 'Variational diffusion models,' in Advances in Neural Information Processing Systems , 2021.

B. D. Anderson, 'Reverse-time diffusion equation models,' Stochastic Processes and their Applications , vol. 12, no. 3, pp. 313-326, 1982.

A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen, 'Hierarchical text-conditional image generation with CLIP latents,' arXiv preprint arXiv:2204.06125 , 2022.

A. Krizhevsky, 'Learning multiple layers of features from tiny images,' Tech. Rep., 2009.

K. Atkinson, W. Han, and D. E. Stewart, Numerical solution of ordinary differential equations . John Wiley & Sons, 2011, vol. 108.

B. Oksendal, Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

P. E. Kloeden and E. Platen, Numerical Solution of Stochastic Differential Equations . Springer, 1992.

A. Roberts, 'Modify the improved euler scheme to integrate stochastic differential equations,' 2012.

W. Peebles and S. Xie, 'Scalable diffusion models with transformers,' arXiv preprint arXiv:2212.09748 , 2022.

T. Hang, S. Gu, C. Li, J. Bao, D. Chen, H. Hu, X. Geng, and B. Guo, 'Efficient diffusion training via min-snr weighting strategy,' arXiv preprint arXiv:2303.09556 , 2023.

J. Deng, W. Dong, R. Socher, L. Li, K. Li, and L. Fei-Fei, 'ImageNet: A large-scale hierarchical image database,' in 2009 IEEE Conference on Computer Vision and Pattern Recognition . IEEE, 2009, pp. 248-255.

F. Yu, Y. Zhang, S. Song, A. Seff, and J. Xiao, 'Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop,' arXiv preprint arXiv:1506.03365 , 2015.

E. Buckwar and R. Winkler, 'Improved linear multi-step methods for stochastic ordinary differential equations,' Journal of Computational and Applied Mathematics , 2007.

A Derivations of Variance Controlled Diffusion SDEs

A.1 Proof of Proposition 4.1

Proposition 4.1 For any bounded measurable function τ ( t ) : [0 , T ] → R , the following Reverse SDEs

has the same marginal probability distributions with (4) and (3) .

Proof. Denote p ( x , t ) : R d × [0 , T ] → R + as the probability density function of x t at time t , thus p ( x , T ) = p T ( x ) . Fokker-Planck equation [38] determines a Partial Differential Equation (PDE) that p t ( x ) satisfies:

Eq. (20) is a reverse-time SDE running from T to 0 , thus there are two additional minus signs in Eq. (21) before term ∂p t ( x ) ∂t and term [ f ( t ) p t ( x ) x i -( 1+ τ 2 ( t ) 2 ) g 2 ( t ) p t ( x ) ∂ log p t ( x ) ∂x i ] compared with vanilla Fokker-Planck equation in general cases. Here δ ij is the Dirac symbol satisfies δ ij = 1 when i = j , otherwise, δ ij = 0 . Notice that

Substituting Eq. (22) into Eq. (21), we obtain that

which is independent of τ ( t ) . With the same initial condition p ( x , T ) = p T ( x ) , the family of Reverse SDEs in Eq. (20) have exactly the same evolutions of probability density function because they share the same Fokker-Planck equation. Especially, when τ ( t ) = 0 , Eq. (20) degenerates to diffusion ODEs and when τ ( t ) = 1 , Eq. (20) degenerates to diffusion SDEs.

A.2 Two Reparameterizations and Exact Solution under Exponential Integrator

In this subsection, we will show the exact solution of SDE in both data prediction reparameterization and noise prediction reparameterization. The noise term in data prediction has smaller variance than noise prediction ones, implying the necessity of adopting data prediction reparameterization for the SDE sampler.

A.2.1 Data Prediction Reparameterization

After approximating ∇ x log p t ( x t ) with s θ ( x t , t ) and reparameterizing s θ ( x t , t ) with -( x t -α t x θ ( x t , t )) /σ 2 t , Eq. (20) becomes

Applying change-of-variable with log-SNR λ t = log( α t /σ t ) and substituting the following relationship

Eq. (24) becomes

A.2.2 Proof of Proposition 4.2

Proposition 4.2 Given x s for any time s > 0 , the solution x t at time t ∈ [0 , s ] of Eq. (9) is

where G ( s, t ) is an Itô integral [38] with the special property

Proof. Define y t = e -∫ t t 0 ( dlog αv d v -( 1+ τ 2 ( v ) ) d λv d v ) d v x t , where t 0 ∈ [0 , T ] is a constant. Differentiate y t with respect to t , we get

Integrating both sides from s to t

Substituting the definition of y t , y s into Eq. (30), we obtain Eq. (27)

The last term σ t ∫ t s e -∫ t u τ 2 ( v ) d λv d v d v τ ( u ) √ -2 d λ u d u d¯ w u of Eq. (31) is the Itô integral term. It follows a Gaussian distribution, which can be directly derived from two basic facts [38]: first, the definition of Itô integral is the limitation in L 2 space; second, the limit of Gaussian Process in L 2 space is still Gaussian. Then we can compute the mean as:

and the variance is

The expectation equals zero because the Itô integral is a martingale [38]. The computation of variance uses the Itô Isometry , which is a crucial fact of Itô integral. We can further simplify the result by using the change of variable P ( λ ) = e ∫ λ λt 2 τ 2 ( ˜ λ )d ˜ λ .

A.2.3 Noise Prediction Reparameterization

After approximating ∇ x log p t ( x t ) with s θ ( x t , t ) and reparameterizing s θ ( x t , t ) with -ϵ θ ( x t , t )) /σ t , Eq. (20) becomes

Applying change-of-variable with log-SNR λ t = log( α t /σ t ) and substituting the following relationship

Eq. (35) becomes

Eq (35) is the formulation of noist prediction model . Similar with Proposition 4.2, Eq. (37) can be solved analytically, which is shown in the following propositon:

Proposition A.1. Given x s for any time s > 0 , the solution x t at time t ∈ [0 , s ] of (37) is

where G ( s, t ) is an Itô integral [38] with the special property

Proof. Define y t = e -∫ t t 0 dlog αv d v d v x t , where t 0 ∈ [0 , T ] is a constant. Differentiate y t with respect to t , we get

Integrating both sides from s to t

Substituting the definition of y t , y s into Eq. (41), we obtain

The Itô integral term α t ∫ t s e -λ u τ ( u ) √ -2 d λ u d u d¯ w u follows a Gaussian distribution. Following the derivation in Proposition 4.2, the mean of the Itô integral term is:

and the expectation is

A.2.4 Comparison between Data and Noise Reparameterizations

In Table 1 we perform an ablation study on data and noise reparameterizations, the experiment results show that under the same magnitude of stochasticity, the proposed SA-Solver in data reparameterization has a better convergence which leads to better FID results under the same NFEs. In this subsection, we provide a theoretical view of this phenomenon.

Corollary A.2. For any bounded measurable function τ ( t ) , the following inequality holds

Proof. It's equivalent to show that

From the basic inequality 1 -e -x ≤ x , we have

Thus it's sufficient to show that

which is true because

This corollary indicates that the same SDE under two different reparameterizations has different properties under the effect of the exponential integrator. Specifically, in the numerical scheme, the data reparameterization will inject smaller noise in each step's updation. We speculate that this is the reason that the data reparameterization has a better convergence, shown as in Table 1.

B Derivations and Proofs for SA-Solver

B.1 Preliminary

We will first review some basic concepts and formulas in the numerical solutions of SDEs [39]. Suppose we have an Itô SDE d x t = f ( x t , t )d t + g ( x t , t )d w t and time steps { t i } M i =0 , t i ∈ [0 , T ] to numerically solve the SDE. For a random variable Z , we define the L 1 norm ∥ Z ∥ L 1 = E [ | Z | ] , the

Definition B.1. We shall say that a time-discrete approximation x 0 , · · · , x M , where x i is a numerical approximation of x t i , converges strongly with order γ > 0 , if there exists a positive constant C , which does not depend on h and a h 0 > 0 such that

Definition B.2. We say it is mean-square convergent with order γ > 0 , if there exists a positive constant C , which does not depend on h and a h 0 > 0 such that

Remark 2. To prove the strong convergence order γ of a numerical scheme, it's sufficient to show the mean-square convergence order γ . This is from Hölder Inequality E [ | Z | ] ≤ E [ | Z | 2 ] 1 2 E [ | 1 | 2 ] 1 2 ≤

We also need the following definition and assumptions, which usually holds in practical diffusion models.

Definition B.3. A function h : R d × [0 , T ] → R d satisfies a linear growth condition if there exists a constant K such that

Assumption B.4. The data prediction model x θ and its derivatives such as ∂ t x θ , ∇ x x θ and ∆ x θ satisfy the linear growth condition.

Assumption B.5. The data prediction model x θ satisfies a uniform Lipschitz condition with respect to x

B.2 Outline of the Proof

In the remaining part of this section, we will focus on our variance controlled SDE

Consider the general case of the numerical scheme as follows:

in which Eq. (17) and Eq (14) are the special case of this scheme. We will provide proof of the mean-square convergence order of the numerical scheme max 0 ≤ i ≤ M ∥ x t i -x i ∥ L 2 . We define the local

error of the numerical scheme Eq. (55) for the approximation of the SDE Eq. (54) as

L i +1 can be decomposed into R i +1 and S i +1 . Then the mean-square convergence can be derived, which is summarized in the following theorem proved by [29]:

Theorem B.6 ([29], Theorem 1) . The mean-square convergent of x i is bounded by

In Eq. (57), D i , i = 0 , · · · , s -1 are the initial error which we do not consider. Given Theorem B.6, to show the convergence order O ( max 0 ≤ t ≤ T τ ( t ) h + h s ) of our s -step SA-Predictor and the convergence order O ( max τ ( t ) h + h s +1 ) of our s -step SA-Corrector

0 ≤ t ≤ T , we just need to prove the following lemmas.

Lemma B.7 (Convergence rate of s -step SA-Predictor ) . For

There exists an decomposition of local error L i such that L i = R i + S i and

Lemma B.8 (Convergence rate of s -step SA-Corrector ) . For

There exists an decomposition of local error L i such that L i = R i + S i and

Lemma B.7 and B.8 will be proved in Sec. B.4.

B.3 Lemmas for the Proof

To better analyze the local error here, we state the following definitions and results from [45]. For a continuous function y : R d × [0 , T ] → R d , a general multiple Wiener integral over the subinterval [ t, t + h ] ⊂ [0 , T ] is given by

where r i ∈ { 0 , 1 , · · · , d } and d w 0 ( s ) = d s . Then we have the following lemma.

Lemma B.9 (Bound of Wiener Integral) . For any function y : R d × [0 , T ] → R d that satisfies a growth condition in the form | y ( x , t ) | ≤ K (1 + | x | 2 ) 1 2 , for any x ∈ R d , and any t ∈ [0 , T ] , h > 0 such that t + h ∈ [0 , T ] , we have that

̸

where l 1 is the number of zero indices and l 2 is the number of non-zero indices r i .

Lemma B.10 (Property of Lagrange interpolation polynomial) . For s + 1 points ( t i +1 , y i +1 ) , ( t i , y i ) , · · · , ( t i -( s -1) , y i -( s -1) ) , the Lagrange interpolation polynomial is

Then the following s+1 equalities hold

Proof. For the first equality, consider y k ≡ 1 for i -( s -1) ≤ k ≤ i +1 . The Lagrange interpolation polynomial for these y k s is a constant function L ( t ) ≡ 1 . We have L ( u ) = ∑ i +1 k = i -( s -1) l k ( u ) = 1 .

For the second equality, consider y k = ∫ t k t i -( s -1) d u 2 . The Lagrange interpolation polynomial for these y k s is a polynomial of degree 1 L ( t ) = t -t i -( s -1) . We have L ( u ) = ∑ i +1 k = i -( s -1) l k ( u ) ∫ t k t i -( s -1) d u 2 = u -t i -( s -1) = ∫ u t i -( s -1) d u 2 .

For equalities from the third to the last, without loss of generality, we prove the p -th equality, where 3 ≤ p ≤ s + 1 . Consider y k = ∫ t k t i -( s -1) ∫ u 2 t i -( s -1) · · · ∫ u p -1 t i -( s -1) d u p · · · d u 3 d u 2 . The Lagrange interpolation polynomial for these y k s is a polynomial of degree p -1 L ( t ) = ∫ t t i -( s -1) ∫ u 2 t i -( s -1) · · · ∫ u p -1 t i -( s -1) d u p · · · d u 3 d u 2 We have L ( u ) = ∫ u t i -( s -1) ∫ u 2 t i -( s -1) · · · ∫ u p -1 t i -( s -1) d u p · · · d u 3 d u 2 = ∫ u t i -( s -1) ∫ u 2 t i -( s -1) · · · ∫ u p -1 t i -( s -1) d u p · · · d u 3 d u 2 .

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

To simplify the notation, we will introduce two operators which will appear in the Itô formula. Suppose we have an Itô SDE d x t = f ( x t , t )d t + g ( x t , t )d w t and h ( x , t ) is a twice continuously differentiable function. Let Γ 0 ( · ) = ∂ t ( · ) + ∇ x ( · ) f , Γ 1 ( · ) = g 2 2 ∆( · ) and Γ 2 ( · ) = ∇ x ( · ) g in which

∇ x is the Jacobian matrix, and ∆ is the Laplacian operator. With the notation here, we can express the Itô formula for h ( x , t ) as

Given the above lemmas, we will analyze the local error L i +1 step by step. Inspired by Theorem B.6, for data-prediction reparameterization model, L i + i can be estimated by decomposing the terms step by step. The first step of decomposition is summarized as the following lemma:

Lemma B.11 (First step of estimating local error L i +1 in data-prediction reparameterization model) . Given the exact solution of data prediction model

With proper b k , k ∈ [ i -( s -1) , i +1] , The local error L i +1 in Eq. (56) is

where

Proof. The difference between Eq. (55) and Eq. (56) is that x j is our numerical approximation, while x t j is the exact solution of SDE Eq. (54) at time t = t j . Substitute the exact solution Eq. (31) of x t i +1 , we have

Let f ( x , t ) = ( dlog α t d t -( 1 + τ 2 ( t ) ) d λ t d t ) x + ( 1 + τ 2 ( t ) ) α t x θ ( x , t ) d λ t d t , g ( t ) = τ ( t ) σ t √ -2 d λ t d t . By Itô's formula [38], we have

Substituting Eq. (72) and Eq. (73) into Eq. (71), we have

We will divide the local error L i +1 into distinct terms. The first term has a coefficient

By Lemma B.10, b k constructed by the integral of Lagrange polynomial in Eq. (58) and Eq. (60) satisfies b k = O ( h ) and the coefficient (75) is zero. Furthermore, we have g ( t ) = τ ( t ) σ t √ -2 d λ t d t =

( )

The remaining terms of local error are

(77)

which completes the proof.

The remaining problem is to estimate the R (1) i +1 in Eq. (70). We can further expand the term Γ 0 ( x θ ) as following

Substituting the expansion of Γ 0 ( x θ ) , we perform the approximation of ˜ L i +1 , which is summarized with the following lemma:

Lemma B.12 (Second step of estimating L i +1 in data-prediction reparameterization model) . R (1) i +1 in Eq. (70) can be decomposed as

where

Proof. We start with decomposing the term R (1) i +1

We further estimate the terms with Γ 1 and Γ 2 .

The summation of the above terms is S (2) = O ( max 0 ≤ t ≤ T τ ( t ) h 5 2 ) . Compared with S (1) , this term can be omitted.

The remaining local error is

which completes the proof.

Remark 3. With Lemma B.11 and B.12, the local error L i +1 can be decomposed to the term S i +1 = S (1) i +1 + S (2) i +1 and the term R (2) i +1 . It is clear that S i +1 = O ( max 0 ≤ t ≤ T τ ( t ) h 3 2 ) , and we will show that given b i -j constructed by integral of Lagrange polynomial in Eq. (58) and Eq. (60) , R (2) i +1 = O ( h 3 ) .

By Lemma B.10, b k constructed by the integral of Lagrange polynomial in Eq. (58) and Eq. (60) satisfies that the coefficient for Γ 0 ( x θ )( x t i -( s -1) , t i -( s -1) )

equals zero. And the remaining term in R (2) i +1 is O ( h 3 ) .

Remark 4. We will show that the local error can be further decomposed such that L i +1 = R ( s ) i +1 + ∑ s j =1 S ( j ) i +1 . In this case S i +1 = ∑ s j =1 S ( j ) i +1 is the term such that S i +1 = O ( max 0 ≤ t ≤ T τ ( t ) h 3 2 ) , and we will show that by our constructed b i -j , R ( s ) i +1 = O ( h s +1 ) .

Lemma B.13 ( j -th step of estimating L i +1 in data-prediction reparameterization model) . For j ≤ s +1 , R ( j -1) i +1 in Eq. (70) can be decomposed as

where

Furthermore, given that b k is constructed by the integral of Lagrange polynomial in Eq. (58) and Eq. (60) , R ( j ) i +1 = O ( h j +1 )

Sketch of the proof (1) Use the Itô formula Eq. (67) to expand j -1 ︷ ︸︸ ︷ Γ 0 · · · Γ 0 ( x θ ) . (2) Use Lemma B.9 to estimate the stochastic term S ( j ) . For the remaining term R ( j ) , by Lemma B.10, b k constructed by the integral of Lagrange polynomial in Eq. (58) and Eq. (60) satisfies that the coefficient before

equals zero. And the remaining term in R ( j ) i +1 is O ( h j +1 ) .

The process can be repeated until the coefficient before s ︷ ︸︸ ︷ Γ 0 · · · Γ 0 ( x θ )( x t i -( s -1) , t i -( s -1) ) is

which equals zero. And the remaining term R s +1 i +1 is O ( h s +2 ) .

We conclude with the proof of Lemma B.7 and B.8.

Proof for Lemma B.8 (Convergence for s -step SA-Corrector ) The stochastic term S i +1 can be estimated as O ( max 0 ≤ t ≤ T τ ( t ) h 3 2 ) . Lemma B.10 prove that with b i -j defined in Theorem 5.2, the coefficients of Eq. (75), Eq. (84), Eq. (87) and Eq. (88) equal zero. Thus the deterministic term R i +1 can be estimated as O ( h s +2 ) . The proof is completed.

Proof for Lemma B.7 (Convergence for s -step SA-Predictor ) The stochastic term S i +1 can be estimated as O ( max 0 ≤ t ≤ T τ ( t ) h 3 2 ) from Eq. (76) and Eq. (82). Lemma B.10 prove that with b i -j defined in Theorem 5.1, the coefficients of Eq. (75), Eq. (84), Eq. (87) and Eq. (88) equal zero except for the last term. This is because in s -step SA-Predictor we only have s points in contrast to s +1 points in s -step SA-Corrector , for which we can only obtain the first s equalities in Lemma B.10. Thus the deterministic term R i +1 can be estimated as O ( h s +1 ) . The proof is completed.

B.5 Relationship with Existing Samplers

B.5.1 Relationship with DDIM

DDIM [22] generates samples through the following process:

where ξ ∼ N ( 0 , I ) , ˆ σ t i is a variable parameter. In practice, DDIM introduces a parameter η such that when η = 0 , the sampling process becomes deterministic and when η = 1 , the sampling process

Corollary 5.3 For any η in DDIM, there exists a τ η ( t ) : R → R which is a piecewise constant function such that DDIMη coincides with our 1 -step SA-Predictor when τ ( t ) = τ η ( t ) with data parameterization of our variance-controlled diffusion SDE.

Proof. Our 1 -step SA-Predictor when τ ( t ) = τ, t ∈ [ t i , t i +1 ] with data parameterization of our variance-controlled diffusion SDE is

DDIMη generates samples through the following process

If we substitute ˆ σ t i with σ t i +1 √ 1 -e -2 τ 2 ( λ t i +1 -λ t i ) , we can verify that √ 1 -α 2 t i +1 -ˆ σ 2 t i = σ t i +1 e -τ 2 ( λ t i +1 -λ t i ) . The DDIMη then becomes

which is exactly the same with our 1 -step SA-Predictor . To find the τ η , we solve the relationship

The relationship between τ and η is

In a concurrent paper [31], Lu et al. prove the result that their SDE-DPM-Solver++1 coincides with DDIM with a special η . Their result is a special case of Corollary 5.3 when τ η ≡ 1 and η take a special value, while our result holds for arbitrary η .

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

DPM-Solver++ [31] is a high-order solver which solves diffusion ODEs for guided sampling. DPMSolver++(2M) is equivalent to the 2-step Adams-Bashforth scheme combined with the exponential integrator. While our 2-step SA-Predictor is also equivalent to the 2-step Adams-Bashforth scheme combined with the exponential integrator when τ ( t ) ≡ 0 . Thus DPM-Solver++(2M) is a special case of our 2 -step SA-Predictor when τ ( t ) ≡ 0 .

B.5.3 Relationship with UniPC

UniPC [24] is a unified predictor-corrector framework for solving diffusion ODEs. Specifically, UniPC-p uses a p-step Adams-Bashforth scheme combined with the exponential integrator as a predictor and a p-step Adams-Moulton scheme combined with the exponential integrator as a corrector. While our p-step SA-Predictor is also equivalent to the p-step Adams-Bashforth scheme combined with the exponential integrator when τ ( t ) ≡ 0 and our p-step SA-Corrector is also equivalent to the p-step Adams-Moulton scheme combined with the exponential integrator when τ ( t ) ≡ 0 . Thus UniPC-p is a special case of our SA-Solver when τ ( t ) ≡ 0 with predictor step p , corrector step p in Algorithm 1.

C Selection on the Magnitude of Stochasticity

In this section, we will show that we choose τ ( t ) ≡ 1 in a number of NFEs. We will show that under certain conditions, the upper bound of KL divergence between the marginal distribution and the true distribution can be minimized when τ ( t ) ≡ 1 .

Let p t ( x ) denotes the marginal distribution of x t , by Proposition 4.1, we know that for any bounded measurable function τ ( t ) : [0 , T ] → R , the following Reverse SDEs

have the same marginal probability distributions. In practice, we substitute ∇ x log p t ( x t ) with s θ ( x t , t ) and substitute p T ( x T ) wiht π to sample the reverse SDE.

where π is a known distribution, specifically here a Gaussian. We have the following theorem under the Assumption in Appendix A in [3].

Theorem C.1. Let p = p 0 be the data distribution, which is the distribution if we sample from the ground truth reverse SDE (54) at time 0 . Let p τ ( t ) θ be the distribution if we sample from the practical reverse SDE (96) at time 0 . Under the assumptions above, we have

This evidence lower bound (ELBO) is minimized when τ ( t ) ≡ 1 .

Proof. Denote the path measure of Eq. (95) and Eq. (96) as µ and ν respectively. Both µ and ν are uniquely determined by the corresponding SDEs due to assumptions. Consider a Markov kernel K ( { z t } t ∈ [0 ,T ] , y ) = δ ( z 0 = y ) . Thus we have the following result

By data processing inequality for KL divergence

By the chain rule of KL divergence, we have

By Girsanov Thoerem, D KL ( µ ( ·| x T = z ) ∥ ν ( ·| x θ T = z ) ) can be computed as

D Implementation Details

For our 2-step SA-Predictor and 1-step SA-Corrector, we find that the coefficient will degenerate to a simple case.

For 2-step SA-Predictor, assume on [ t i , t i +1 ] , τ ( t ) = τ is a constant,

we have

Let h = λ t i +1 -λ t i , we have

Thus we implement ˆ b i -1 as α t i +1 λ t i -λ t i -1 1 2 (1 + τ 2 ) h 2 and ˆ b i as α t i +1 (1 -e -h (1+ τ 2 ) ) -ˆ b i -1 . Note that substituting b i , b i -1 as ˆ b i , ˆ b i -1 will maintain the convergence order result of 2-step SA-Predictor since the modified term is O ( h 3 ) . The implementation detail for 1-step SA-Corrector is technically the same.

E Experiment Details

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

CIFAR10 32x32 For the CIFAR10 experiment in Section 6.4, we use the pretrained baselineunconditional-VE model 4 from [27]. It's an unconditional model with VE noise schedule. To fairly compare with results in [27], we use a piecewise constant function τ ( t ) inspired by [27]. Concretely, denoting σ t EDM = σ t α t , our τ ( t ) is set to be a constant τ in the interval [( σ EDM ) -1 (0 . 05) , ( σ EDM ) -1 (1)] and to be zero outside the interval. We find empirically that this piecewise constant function setting makes our SA-Solver converge better, especially in large noise scale cases. We use a 3-step SA-Predictor and a 3-step SA-Corrector. For the CIFAR10 experiment in Section 6.2 and 6.5, we also use the piecewise constant function τ ( t ) as above. The predictor steps and corrector steps vary to verify the effectiveness of our proposed method in Section 6.2, while they are both set to be 3-steps in Section 6.5.

ImageNet 64x64 For the ImageNet 64x64 experiment in Section 6.4, we use the pretrained model 5 from [4]. It's a conditional model with VP cosine noise schedule. To fairly compare with results in [27], we use a piecewise constant function τ ( t ) inspired by [27]. Concretely, denoting σ t EDM = σ t α t , our τ ( t ) is set to be a constant τ in the interval [( σ EDM ) -1 (0 . 05) , ( σ EDM ) -1 (50)] and to be zero outside the interval. We find empirically that this piecewise constant function setting makes our SA-Solver converge better, especially in large noise scale cases. We use a 3-step SA-Predictor and a 3-step SA-Corrector.

Other experiments For other experiments, we use a constant function τ ( t ) ≡ τ . It's generally not the optimal choice for each individual task, thus further fine-grained tuning has the potential to improve the results. We aim to report the result of our SA-Solver without extra hyperparameter tuning. We use a 3-step SA-Predictor and a 3-step SA-Corrector under 20 NFEs and 2-step SA-Predictor and a 1-step SA-Corrector beyond 20 NFEs.

E.2 Details on Pretrained Models and Settings

CIFAR10 32x32 For the CIFAR10 experiment, we use the pretrained baseline-unconditional-VE model 6 from [27]. It's an unconditional model with VE noise schedule. To fairly compare with results in [27], we follow the time step schedule in it. Specifically, we set σ min = 0 . 02 and σ max = 80 and select the step by σ i = ( σ 1 7 max + i N -1 ( σ 1 7 min -σ 1 7 max )) 7 for SA-Solver and UniPC. We directly report the results of the deterministic sampler and stochastic sampler of EDM. To make it a strong baseline, we report the results of the optimal setting for 4 hyper-parameters { S churn , S tmin , S tmax , S noise } and report its lowest observed FID. While for SA-Solver and UniPC, we report the averaged observed FID.

ImageNet 64x64 For the ImageNet 64x64 experiment, we use the pretrained model 7 from [4]. It's a conditional model with VP cosine noise schedule. To fairly compare with results in [27], we follow the time step schedule in it and use conditional sampling. Specifically, we set σ min = 0 . 0064 and σ max = 80 and select the step by σ i = ( σ 1 7 max + i N -1 ( σ 1 7 min -σ 1 7 max )) 7 for SA-Solver , UniPC, DPM-Solver and DDIM. We directly report the results of the deterministic sampler and stochastic sampler of EDM. To make it a strong baseline, we report the results of the optimal setting for 4 hyper-parameters { S churn , S tmin , S tmax , S noise } and report its lowest observed FID. While for SA-Solver , UniPC, DPM-Solver and DDIM, we report the averaged observed FID.

ImageNet 256x256 For the ImageNet 256x256 experiment, we use three different pretrained models:LDM 8 (VP, handcrafted noise schedule) from [5], DiT-XL/2 9 (VP, linear noise schedule) from [41], Min-SNR 10 (VP, cosine noise schedule) from [42]. We use classifier-free guidance of scale s = 1 . 5 and a uniform time step schedule because it's the most common setting for guided sampling for ImageNet 256x256.

ImageNet 512x512 For the ImageNet 256x256 experiment, we use the pre-trained model: DiTXL/2 11 from [41]. We use classifier-free guidance of scale s = 1 . 5 and a uniform time step schedule following the settings of DiT [41].

LSUN Bedroom 256x256 For the LSUN Bedroom 256x256 experiment, we use the pretrained model 12 from [4]. We use unconditional sampling and a uniform lambda step schedule from [23].

F Additional Results

We include the detailed FID results in Figure 1, Figure 2 and Figure 4 in the tables 4 to 14. The ablation study shows that stochasticity indeed helps improve sample quality. We find that for small NFEs, the magnitude of stochasticity should be small while for large NFEs, large magnitude of stochasticity helps improve sample quality. It can also be observed that in latent space, SDE converges faster as in Table 13. With only 10 NFEs, τ = 0 . 6 is better than τ = 0 . With 20 NFEs, our SA-Solver can achieve 3.87 FID, which outperforms all ODE samplers even with far more steps.

G Additional Samples

We include additional samples in this section. In Figure 5 and Figure 6 we compare samples of our proposed SA-Solver with other diffusion samplers. In Figure 7 and Figure 8, we compare samples of our proposed SA-Solver under different NFEs and τ . In Figure 9 and Figure 10, we compare samples of our proposed SA-Solver with other diffusion samplers on text-to-image tasks. Our SA-Solver can generate more diverse samples with more details.

NFE = 15

NFE = 23

NFE = 47

NFE = 95

DDIM(

η

= 0

)

DPM-Solver

UniPC

EDM(ODE)

EDM(SDE)

SA-Solver(ours)

NFE = 15

NFE = 23

NFE = 47

NFE = 95

DDIM(

η

= 0

)

DPM-Solver

UniPC

SA-Solver(ours)

NFE = 20

NFE = 40

NFE = 60

NFE = 100

SA-Solver(

η

= 0

)

SA-Solver(

η

= 0

.

2

)

SA-Solver(

η

= 0

.

4

)

SA-Solver(

η

= 0

.

6

)

SA-Solver(

η

= 0

.

8

)

SA-Solver(

η

= 1

.

0

)

SA-Solver(

η

= 1

.

2

)

山

NFE = 60

SA-Solver( τ = 0 )

SA-Solver( τ = 0 . 4 )

SA-Solver(

τ

= 0

.

8

)

SA-Solver(

τ

= 1

.

0

)

CRA

OLA

NFE = 20

NFE = 50

NFE = 100

DDIM(

η

= 0

)

UniPC

SA-Solver(Ours)

NFE = 20

NFE = 50

NFE = 100

DDIM(

η

= 0

)

UniPC

SA-Solver(Ours)

Peer reviews:
comments: SUMMARY:
The authors propose a sampling method for diffusion probabilistic models by solving an alternative SDE with the same marginal distribution. Approximation techniques are applied in constructing the computationally efficient solver and corrector . The results show a superior FID with less function evaluations.

STRENGTHS:
1. The construction of the alternative SDE and the approximation techniques, including score function approximation, change of variable and Lagrange interpolation, provide a complete framework for efficient DFM generation.
2. The proposed method can achieve lower FID scores with less function evaluations which is promising.

WEAKNESSES:
1. The paper content seems not well organized yet.
2. In the appendix, for other methods, with different NFEs, the image content is consistent. But for the proposed method, the image content varies a lot. It might be caused by the inaccurate approximations.

QUESTIONS:
1. Could you elaborate on why the images are not consistent for different NFEs in your method?

LIMITATIONS:
The authors did not address the limitations.

There is no potential negative societal impact of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-17

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends UniPC to a stochastic manner to obtain a faster sampler called SA-Solver for diffusion models. Specifically, the authors start from the formulation of diffusion SDE and derive the SA-Predictor and SA-Corrector. Extensive experiments on CIFAR10, ImageNet, LSUN, etc demonstrate the effectiveness of the proposed SA-Solver.\n\nSTRENGTHS:\n1. It is interesting to consider sampling from the SDE instead of ODE. This opens up a new direction for the fast sampling of DPMs. I am surprised that this simple practice can improve the sampling quality so much.\n2. The writing is clear and the derivation of SA-Solver is easy to understand. I especially appreciate Section 5.3, where the relationship with UniPC helps me to understand the proposed method at a glance.\n3. The experiments are thorough and convincing. Most required experiments are included, such as unconditional/conditional sampling with data/noise prediction with different image resolutions.\n\nWEAKNESSES:\n1. It is questionable whether the novelty of SA-Solver is enough. As said in L228, SA-Solver can be viewed as an extension of UniPC with a non-zero $\\tau(t)$. However, I think the idea to introduce stochasticity back to the sampling of the diffusion models also contributes to the novelty. I just feel worried that this issue might be raised by other reviewers.\n2. Some experiments are lacking, such as comparing the convergence of different methods on COCO using a pre-trained text-to-image diffusion model. For now, only some qualitative results are provided. It would be better to also demonstrate some quantitative comparisons.\n3. Some minor format issues: too much space below Table 1. Maybe the authors can try to make it looks better.\n\nQUESTIONS:\nSee the weaknesses.\n\nLIMITATIONS:\nThe limitations have been fully discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: SUMMARY:
The paper proposed a stochastic Adam solver for solving diffusion SDEs in an efficient way with a convergence guarantee. Authors adapt the stochastic Adam from numerical literature and use Lagrange interpolation to predict unknown terms. They show strong convergence for both predictor and corrector. Numerical experiments are done to verify their claims.

STRENGTHS:
1. By deriving an explicit form of the solution of the SDE, authors are able to utilize stochastic Adam with Lagrange interpolation to find an efficient diffusion SDE solver, 
2. Theoretical analysis is well-written.
3. Numerical experiments show improvements compared to other solvers, especially the diffusion ODE solvers.
4. It sort of unifies previous methods in some sense.

WEAKNESSES:
1. My main concern is, from Figure 2, it seems SA-solver only outperforms other methods within a certain range of NFE. This "optimal" range looks very different for different data sets. In reality, if I use SA-solver for the sake of doing fewer function evals, how do I know when it outperforms other methods? 

2. I don't think I fully understand how to choose the parameter tau(t) in the experiments. Please comment on this.

QUESTIONS:
See weaknesses.

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new solver for diffusion SDEs, termed SA-Solver, combining the ideas of a predictor-corrector scheme and the stochastic Adam solvers. The predictor and corrector utilize the Lagrange polynomials for extrapolation to lower the approximation error at future time stamps. Experimentally, SA -Solver improves over previous ODE and SDE across a wide range of image generation benchmarks.\n\nSTRENGTHS:\n- The considered problem of accelerating SDE solvers is of great practical value in the field, as SDE solver often delivers better sample quality than ODE solvers, but are hindered by their slow sampling speed.\n\n- Adapting the idea in Stochastic Adams is interesting, offering higher-order convergence with Lagrange polynomial extrapolation.\n\n- The fusion of the PC sampler and Lagrange polynomial yields superior empirical results on diffusion SDE, tested on a range of dataset resolutions from 32x32 to 256x256. It's nice to see the methods scale to different resolutions and architectures.\n\nWEAKNESSES:\n- Throughout the paper, the authors mention several times that empirically the quality of data generated by SDE has a better upper limit. I think some theoretical analysis in the main text to make the paper more self-contained. A concurrent [1] provides some theoretical arguments and comparisons between SDE and ODE, showing that the stochasticity in SDE can reduce overall sample errors. I think it would be beneficial to discuss them to better motivate the idea of the paper.\n\n- The authors didn't provide detailed discussions of the benefit of introducing $\\tau(t)$ (at the end the author set $\\tau(t)=1$ in experiments). Since the authors observe the benefit of larger $\\tau$ on larger NFE, why not use a varying $\\tau$ in Fig.2 based on NFE? In addition, the benefit of a larger $\\tau$ on larger NFE seems predictable based on the theory in [1].\n\n - The authors did some comparisons of epsilon-prediction and data prediction models. It seems that EDM [2] uses interpolation between these two by pre-conditioning. Does it offer better results?\n\n- Unlike the deterministic path in ODE, the randomness in SDE can change the sample trajectories quite a bit. For example, the trajectories of generating a dog can stray into generating a cat. I wonder if this is an issue for using Lagrange polynomials in SDE, since the preceding predictions could be misleading. If not, could the author provide some intuitions? In addition, the authors mention that \"predictor step 2 and corrector step 1 is the most stable setup\". such small steps could be caused by the issue I raise.\n\n- The paper separately gives the convergence order of predictor and corrector. Is it possible to give the theoretical order after combing the predictor and corrector?\n\n- I think the author didn't use the \"EDM VE\" model for CIFAR-10. In Table 2 of EDM [2], the unconditional VE (in config F, EDM VE) can achieve an FID of 1.98 using only 35 NFE, as opposed to the >3 FID in Fig 2 in the current paper. So I guess the author use their VE model in config A (baseline VE). Could the author compare different methods using the VP model (config A), as in [1] and [2], because the VP model is more commonly used in practice and offers better results?\n\n- The stochastic sampler in the concurrent work [1] obtains an FID ~ 1.8 using less than 100 NFE on ImageNet-64 with Pixel DPM, as opposed to the > 2.1 FID in Fig 2 of the current paper. I wonder if one could combine the multi-step idea in this paper with [1] to obtain further improvements.\n\n- Could the author also provide quantitative results on Stable Diffusion experiments, to better showcase the advantage of SA-Solver? Current visualized images do not seem compelling.\n\n\n### Minors\n\n- line 225, DPM-solver++ and UniPC are *ODE* solvers which would not naturally be special cases of the proposed *SDE* solver.\n\n\n\n[1] Restart Sampling for Improving Generative Processes, Xu et al, https://arxiv.org/abs/2306.14878\n\n[2] Elucidating the Design Space of Diffusion-Based Generative Models, Karras et al, https://arxiv.org/abs/2206.00364\n\n\n# Post rebuttal\n\nThanks the authors for the rebuttal. After another pass of the paper, it occurred to me that that paper is a direct application of the ideas (variation-of-parameter, or exponential integrator, and Langraian polynomial for approximation) in [1] to SDE samplers. The same formula of the exponential integrated version of SDE (Eq.5) has been proposed in prior works [1] (see their eq.17, arXiv version 1). In addition, [1] also uses a Lagrange polynomial to extrapolate. Due to it's straightforward extension of prior works, I will keep my score unchanged.\n\n[1] Zhang, Qinsheng, and Yongxin Chen. \"Fast sampling of diffusion models with exponential integrator.\" arXiv preprint arXiv:2204.13902 (2022).\n\nQUESTIONS:\nI copy-and-paste some questions from above:\n\n- Since the authors observe the benefit of larger $\\tau$ on larger NFE, why not use a varying $\\tau$ in Fig.2 based on NFE? \n\n- I wonder if this is an issue for using Lagrange polynomials in SDE, since the preceding predictions could be misleading. If not, could the author provide some intuitions? In addition, the authors mention that \"predictor step 2 and corrector step 1 is the most stable setup\". such small steps could be caused by the issue I raise.\n\n- Is it possible to give the theoretical order after combing the predictor and corrector?\n\n- Could the author compare different methods using the VP model (config A), as in [1] and [2], because the VP model is more commonly used in practice and offers better results?\n\n- Could the author also provide quantitative results on Stable Diffusion experiments, to better showcase the advantage of SA-Solver?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the problem of fast sampling of diffusion models. Since standard DDPM sampling is slow, this is a very important topic of late, with many competing methods. Many methods reformulate as a ODE solving problem, which makes it easier to do few-step sampling. However, it has been noted that sampling from the original SDE formulation (as opposed to ODE) can lead to better samples if there is budget for taking many diffusion steps. Therefore, this paper studies the problem of accelerating sampling of the SDE. \n\nThe method, named SA-Solver is a predictor corrector method. It incorporates some recent findings, such as semi-linearity, but mainly uses the method of Lagrange interpolations. More specifically, the predictor step involves (uniquely) fitting a polynomial of degree s-1 using s predictions and reading off the polynomial value at desired points. Finally, a correction step is incorporated that plug in the predicted value at the new timestep (t+1) back into the formula to choose a new predicted value (similar in spirit to heun's method). The paper draws connections to DDIM (deterministic version of 1-step SA-Predictor), DPM-Solver (deterministic version of 2-step SA-Predictor), and UniPC (deterministic version of SA_Solver)\n\nThe paper shows sampling results on cifar10, imagenet64, and imagenet256 (latent), demonstrating best FID scores (conditioned on the same number of NFEs) when using more than 30 NFEs.\n\nSTRENGTHS:\nThe problem being studied is an important one, as sampling speed is arguably the biggest weakness of diffusion models (and diffusion models are being used in a huge number of applications right now).\n\nThe experimental results seem quite promising, showing better FID scores than DDIM/DPMSolver, which are very popular baselines.\n\nWEAKNESSES:\nSeems to be missing quantitative evaluation (e.g. CLIP score on stable diffusion) for text-to-image tasks. Also, evaluating on different domains besides image would improve the empirical results.\n\nQUESTIONS:\nHow much of the performance is due to the predictor/corrector approach, versus the lagrange interpolation? An ablation study in this regard would improve the paper and our understanding of the method.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a multistep SDE solver for diffusion models instead of ODE solvers. The main goal is to have diverse and high-quality samples while reducing the number of solver step required. To do this, the paper proposes a new SDE that includes an additional term, ensuring the marginal distribution unchanged. Based on this new SDE, the paper adopts stochastic Adam method, introducing two key component: SA-predictor, and SA-corrector. The convergences of these two are provided. Experiments are conducted in different scenario including varying stochastic noise scale, different models. The results support the claims on the reduction in NFEs.\n\nSTRENGTHS:\nThe paper has the following strong points:\n\n- The experiment shows state-of-the-art result in term of FID given a limited NFEs. This highlights the effectiveness of the proposed method in achieving impressive results with efficient resource utilization.\n- It provides theoretical studies on the convergence of proposed method, offering valuable insights\n\nWEAKNESSES:\nThe main weaknesses I find in this paper are:\n\n- In term of the contribution to difussion model research, the proposed method may be considered incremental. This is because the key factor leading to good result, such as predictor-corrector, multi-step have been well-established in the existing literature.\n- The paper presentation can be improved, in particular the exposure of stochastic Adam methods. I would be good to dedicate a small paragraph or section explicitly discussing ths method to provide readers with clear understand of how it is employed in the proposed approach. \n- The paper lacks clarity in establishing the connection between the variance-control SDE and the use of stochastic Adam methods. The discussion should be more thorough and elaborate, especially in addressing why setting $\\tau=1$ (falling back to DDPM) is the best choice in many cases. Then, stochastic Adam significantly contributes to the model's overall performance.\n\nQUESTIONS:\nHow many of predictor step $s_p$ and corrector step $s_c$ is used?\n\nMinor points,\n- Line 199, please add the notation stepsize $h$\n- Eq (13), please add the description for the Lagrange basis $l(t)$.\n- Caption of Fig. 1, \"scholastic\" -> \"stochastic\"\n- There are some typos in Appendix\n\nLIMITATIONS:\nThe paper mentioned a current limitation is that the optimal $\\tau(t)$ is unknown, which I agree. There many factors affecting this parameters including NFE, dataset and approximated score function.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-01"}