SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shuchen Xue $^{1,4*}$ , Mingyang Yi $^{2\dagger}$ , Weijian Luo $^{3*}$ , Shifeng Zhang $^{2}$ , Jiacheng Sun $^{2}$
Zhenguo Li $^{2}$ , Zhi-Ming Ma $^{1,4}$
<sup>1</sup>University of Chinese Academy of Sciences <sup>2</sup> Huawei Noah's Ark Lab <sup>3</sup> Peking University <sup>4</sup> Academy of Mathematics and Systems Science, Chinese Academy of Sciences
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).
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]
*Work done during an internship at Huawei Noah's Ark Lab. Email: xueshuchen17@mails.ucas.ac.cn, luoweijian@stu.pku.edu.cn
†Corresponding authors: Mingyang Yi (yimingyang2@huawei.com)
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.
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 SDE-DPM-Solver++, SA-Solver has theoretical guarantees with proved convergence order; 3) SDE-DPM-Solver++ 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.
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:
$$
\mathrm {d} \boldsymbol {x} _ {t} = f (t) \boldsymbol {x} _ {t} \mathrm {d} t + g (t) \mathrm {d} \boldsymbol {w} _ {t}, \tag {1}
$$
where  $\pmb{w}_t\in \mathbb{R}^d$  represents the standard Wiener process,  $f(t)\pmb{x}_t$  and  $g(t)$  respectively denote the drift and diffusion coefficients. For each time  $t\in [0,T]$ ,  $\pmb{x}_t|\pmb{x}_0\sim \mathcal{N}(\alpha_t\pmb{x}_0,\sigma_t^2\pmb{I})$ .
Let  $p_t(\pmb{x})$  denotes the marginal distribution of  $\pmb{x}_t$ , the coefficients  $f(t)$  and  $g(t)$  are meticulously selected to guarantee that the marginal distribution  $p_T(\pmb{x}_T)$  closely approximates a Gaussian distribution, i.e.,  $\mathcal{N}(\mathbf{0},\pmb{I})$ , and the signal-to-noise-ratio (SNR)  $\alpha_t^2 /\sigma_t^2$  is strictly decreasing w.r.t.  $t$ . In the sequel, we follow the established notations in [33]:
$$
f (t) = \frac {\mathrm {d} \log \alpha_ {t}}{\mathrm {d} t}, \quad g ^ {2} (t) = \frac {\mathrm {d} \sigma_ {t} ^ {2}}{\mathrm {d} t} - 2 \frac {\mathrm {d} \log \alpha_ {t}}{\mathrm {d} t} \sigma_ {t} ^ {2}. \tag {2}
$$
Anderson [34] demonstrates a pivotal theorem that the forward process (1) has an equivalent reverse-time 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].
$$
\mathrm {d} \boldsymbol {x} _ {t} = \left[ f (t) \boldsymbol {x} _ {t} - g ^ {2} (t) \nabla_ {\boldsymbol {x}} \log p _ {t} (\boldsymbol {x} _ {t}) \right] \mathrm {d} t + g (t) \mathrm {d} \bar {\boldsymbol {w}} _ {t}, \quad \boldsymbol {x} _ {T} \sim p _ {T} (\boldsymbol {x} _ {T}) \tag {3}
$$
where  $\bar{w}_t$  represents the Wiener process in reverse time, and  $\nabla_{\pmb{x}}\log p_t(\pmb{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(\pmb{x})$  as (3), so that the ODE solver can be adopted for efficient sampling [23, 24]:
$$
\mathrm {d} \boldsymbol {x} _ {t} = \left[ f (t) \boldsymbol {x} _ {t} - \frac {1}{2} g ^ {2} (t) \nabla_ {\boldsymbol {x}} \log p _ {t} (\boldsymbol {x} _ {t}) \right] \mathrm {d} t, \quad \boldsymbol {x} _ {T} \sim p _ {T} (\boldsymbol {x} _ {T}) \tag {4}
$$
To get the score function  $\nabla_{\pmb{x}}\log p_t(\pmb {x}_t)$  in (3), we usually take neural network  $s_\theta (\pmb {x},t)$  parameterized by  $\pmb{\theta}$  to approximate it by optimizing the denoising score matching loss [3]:
$$
\boldsymbol {\theta} ^ {*} = \underset {\boldsymbol {\theta}} {\arg \min } \mathbb {E} _ {t} \left\{\lambda (t) \mathbb {E} _ {\boldsymbol {x} _ {0}} \mathbb {E} _ {\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0}} \left[ \left\| \boldsymbol {s} _ {\boldsymbol {\theta}} (\boldsymbol {x}, t) - \nabla_ {\boldsymbol {x} _ {t}} \log p _ {0 t} (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0}) \right\| _ {2} ^ {2} \right] \right\}. \tag {5}
$$
In practice, two methods are used to reparameterize the score-based model [35]. The first approach utilizes a noise prediction model such that  $\epsilon_{\pmb{\theta}}(\pmb{x}_t,t) = -\sigma_t s_{\pmb{\theta}}(\pmb{x}_t,t)$ , while the second employs a data prediction model, represented by  $\pmb{x}_{\pmb{\theta}}(\pmb{x}_t,t) = (\pmb{x}_t - \sigma_t \epsilon_{\pmb{\theta}}(\pmb{x}_t,t)) / \alpha_t$ . The reparameterized models are plugged into the sampling process (3) or (4) according to their relationship with  $s_{\pmb{\theta}}(\pmb{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 improves 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  $\tau(t): [0, T] \to \mathbb{R}$ , the following Reverse SDEs
$$
\mathrm {d} \boldsymbol {x} _ {t} = \left[ f (t) \boldsymbol {x} _ {t} - \left(\frac {1 + \tau^ {2} (t)}{2}\right) g ^ {2} (t) \nabla_ {\boldsymbol {x}} \log p _ {t} (\boldsymbol {x} _ {t}) \right] \mathrm {d} t + \tau (t) g (t) \mathrm {d} \bar {\boldsymbol {w}} _ {t}, \quad \boldsymbol {x} _ {T} \sim p _ {T} (\boldsymbol {x} _ {T}) \tag {6}
$$
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  $\tau(t)$ , and  $\tau(t) = 0$  or  $\tau(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  $\nabla_{\pmb{x}}\log p_t(\pmb{x}_t)$  in it with the "data prediction reparameterization model"  $\pmb{x}_{\theta}(\pmb{x}_t,t)$  according to  $\nabla_{\pmb{x}}\log p_t(\pmb{x}_t)\approx -(\pmb{x}_t - \alpha_t\pmb{x}_{\theta}(\pmb{x}_t,t)) / \sigma_t^2$  as pointed out in Section 3. Then diffusion SDEs to be solved become
$$
d \boldsymbol {x} _ {t} = \left[ f (t) \boldsymbol {x} _ {t} + \left(\frac {1 + \tau^ {2} (t)}{2 \sigma_ {t}}\right) g ^ {2} (t) \left(\frac {\boldsymbol {x} _ {t} - \alpha_ {t} \boldsymbol {x} _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {t} , t)}{\sigma_ {t}}\right) \right] d t + \tau (t) g (t) d \bar {\boldsymbol {w}} _ {t}. \tag {7}
$$
Remark 1. We reparameterize the score function in diffusion SDEs (6) with data prediction model  $\pmb{x}_{\theta}(\pmb{x}_t,t)$  to get Eq. (9). The equation can be also reparameterized by the "noise prediction model"  $\epsilon_{\theta}(\pmb{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  $\lambda_{t} = \log (\alpha_{t} / \sigma_{t})$ . Noting the following relationship in Eq. (2)
$$
f (t) = \frac {\mathrm {d} \log \alpha_ {t}}{\mathrm {d} t}, \quad g ^ {2} (t) = \frac {\mathrm {d} \sigma_ {t} ^ {2}}{\mathrm {d} t} - 2 \frac {\mathrm {d} \log \alpha_ {t}}{\mathrm {d} t} \sigma_ {t} ^ {2} = - 2 \sigma_ {t} ^ {2} \frac {\mathrm {d} \lambda_ {t}}{\mathrm {d} t}, \tag {8}
$$
and plugging them into (7), it becomes
$$
\mathrm {d} \boldsymbol {x} _ {t} = \left[ \frac {\mathrm {d} \log \alpha_ {t}}{\mathrm {d} t} \boldsymbol {x} _ {t} - (1 + \tau^ {2} (t)) (\boldsymbol {x} _ {t} - \alpha_ {t} \boldsymbol {x} _ {\theta} (\boldsymbol {x} _ {t}, t)) \frac {\mathrm {d} \lambda_ {t}}{\mathrm {d} t} \right] \mathrm {d} t + \tau (t) \sigma_ {t} \sqrt {- 2 \frac {\mathrm {d} \lambda_ {t}}{\mathrm {d} t}} \mathrm {d} \bar {\boldsymbol {w}} _ {t}. \tag {9}
$$
The above equation has an explicit solution owing to its semi-linear structure [37].
Proposition 4.2. Given  $\pmb{x}_s$  for any time  $s > 0$ , the solution  $\pmb{x}_t$  at time  $t \in [0, s]$  of (9) is
$$
\pmb {x} _ {t} = \frac {\sigma_ {t}}{\sigma_ {s}} e ^ {- \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} \pmb {x} _ {s} + \sigma_ {t} \pmb {F} _ {\pmb {\theta}} (s, t) + \sigma_ {t} \pmb {G} (s, t),
$$
$$
\boldsymbol {F} _ {\boldsymbol {\theta}} (s, t) = \int_ {\lambda_ {s}} ^ {\lambda_ {t}} e ^ {- \int_ {\lambda} ^ {\lambda_ {t}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} \left(1 + \tau^ {2} (\lambda)\right) e ^ {\lambda} \boldsymbol {x} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {\lambda}, \lambda\right) \mathrm {d} \lambda \tag {10}
$$
$$
\boldsymbol {G} (s, t) = \int_ {s} ^ {t} e ^ {- \int_ {\lambda_ {u}} ^ {\lambda_ {t}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} \tau (u) \sqrt {- 2 \frac {\mathrm {d} \lambda_ {u}}{\mathrm {d} u}} \mathrm {d} \bar {\boldsymbol {w}} _ {u},
$$
where  $G(s, t)$  is an Itô integral [38] with the special property
$$
\sigma_ {t} \boldsymbol {G} (s, t) \sim \mathcal {N} \left(\boldsymbol {0}, \sigma_ {t} ^ {2} \left(1 - e ^ {- 2 \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}}\right)\right). \tag {11}
$$
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  $\pmb{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  $\pmb{x}_{\theta}$ , timesteps  $\{t_i\}_{i=0}^{M}$ , initial value  $\pmb{x}_{t_0}$ , predictor step  $s_p$ , corrector step  $s_c$ , buffer  $B$  to store former evaluation of  $\pmb{x}_{\theta}$ ,  $\tau(t)$  to control variance.
1:  $B \xleftarrow{\text{buffer}} \boldsymbol{x}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t_0}, t_0)$
2: for  $i = 1$  to  $\max(s_p, s_c)$  do
3: sample  $\pmb {\xi}\sim \mathcal{N}(\pmb {0},\pmb {I})$
4: calculate steps for warm-up  $s_p^m = \min(i, s_p)$ ,  $s_c^m = \min(i, s_c)$ .
5:  $\pmb{x}_{t_i}^p \gets s_p^m$  -step SA-Predictor  $(\pmb{x}_{t_{i-1}}, B, \pmb{\xi})$  (Eq. (14))
6:  $B \xleftarrow{\text{buffer}} \boldsymbol{x}_{\theta}(x_{t_i}^p, t_i)$
7:  $\pmb{x}_{t_i} \gets s_c^m$ -step SA-Corrector( $\pmb{x}_{t_i}^p, \pmb{x}_{t_{i-1}}, B, \pmb{\xi}$ ) (Eq. (17))
8: for  $i = \max(s_p, s_c) + 1$  to  $M$  do
9: sample  $\xi \sim \mathcal{N}(\mathbf{0},\mathbf{I})$
0:  $\pmb{x}_{t_i}^p\gets s_p$  -step SA-Predictor  $(\pmb{x}_{t_{i - 1}},B,\pmb {\xi})$  (Eq. (14))
1:  $B \xleftarrow{\text{buffer}} \boldsymbol{x}_{\theta}(x_{t_i}^p, t_i)$
2:  $\pmb{x}_{t_i}\gets s_{c\text{-step}}SA\text{-Corrector} (\pmb{x}_{t_i}^p,\pmb{x}_{t_{i - 1}},B,\pmb {\xi})$  (Eq. (17))
return  $x_{t_M}$
$\triangleright$  Warm-up
$\triangleright$  Prediction Step
Evaluation Step
$\triangleright$  Correction Step
$\triangleright$  Prediction Step
$\triangleright$  Evaluation Step
$\triangleright$  Correction Step
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\}_{i = 0}^M$  which is strictly decreased from  $t_0 = T$  to  $t_M = 0$ . Then we can iteratively obtain the  $x_{t_i}$  (so that  $x_0$  approximates the required data) by the following relationship.
$$
\boldsymbol {x} _ {t _ {i + 1}} = \frac {\sigma_ {t _ {i + 1}}}{\sigma_ {t _ {i}}} e ^ {- \int_ {\lambda_ {t _ {i}}} ^ {\lambda_ {t _ {i + 1}}} \tau^ {2} (\lambda_ {u}) \mathrm {d} \lambda_ {u}} \boldsymbol {x} _ {t _ {i}} + \sigma_ {t _ {i + 1}} \boldsymbol {F} _ {\boldsymbol {\theta}} (t _ {i}, t _ {i + 1}) + \sigma_ {t _ {i + 1}} \boldsymbol {G} (t _ {i}, t _ {i + 1}) \tag {12}
$$
As pointed out in Proposition 4.2, the Ito 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_\theta(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  $\pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_i}, t_i), \pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_{i-1}}, t_{i-1}), \dots, \pmb{x}_{\pmb{\theta}}(\pmb{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_{\pmb{\theta}}(t_i, t_{i+1})$  involves the application of Lagrange interpolations [37] of these evaluations. Lagrange interpolation of  $s$  points  $\pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_i}, t_i), \pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_{i-1}}, t_{i-1}), \dots, \pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_{i-(s-1)}}, t_{i-(s-1)})$  is a polynomial  $L(t)$  with degrees  $s-1$ :
$$
L (t) = \sum_ {j = 0} ^ {s - 1} l _ {i - j} (t) \boldsymbol {x} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t _ {i - j}}, t _ {i - j}\right), \tag {13}
$$
where  $l_{i - j}(t):\mathbb{R}\to \mathbb{R}$  is the Lagrange basis. Lagrange interpolation is an excellent approximation of  $\pmb {x}_{\pmb{\theta}}(\pmb {x}_t,t)$  with the special property:  $L(t_{i - j}) = x_{\pmb{\theta}}(x_{t_{i - j}},t_{i - j}),\quad \forall 0\leq j\leq s - 1$  . Thus a natural
3The diffusion SDEs (7) are reverse-time SDEs, so that the  $t_i$  here is increased.
way to estimate  $F_{\pmb{\theta}}(t_i, t_{i+1})$  is to replace  $\pmb{x}_{\pmb{\theta}}(\pmb{x}_{\lambda_u}, \lambda_u)$  with  $L(\lambda)$ , which is just a change-of-variable of  $L(t)$ . The formula for  $s$ -step SA-Predictor is then derived.