SlowGuess/NeurIPS_2023_OCR huggingface dataset:
s-step SA-Predictor Given the initial value  $\pmb{x}_{t_i}$  at time  $t_i$ , a total of  $s$  former model evaluations  $\pmb{x}_{\theta}(\pmb{x}_{t_i}, t_i), \pmb{x}_{\theta}(\pmb{x}_{t_{i-1}}, t_{i-1}), \dots, \pmb{x}_{\theta}(\pmb{x}_{t_{i-(s-1)}}, t_{i-(s-1)})$ , our  $s$ -step SA-Predictor is defined as:
$$
\boldsymbol {x} _ {t _ {i + 1}} = \frac {\sigma_ {t _ {i + 1}}}{\sigma_ {t _ {i}}} e ^ {- \int_ {\lambda_ {t _ {i}}} ^ {\lambda_ {t _ {i + 1}}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} \boldsymbol {x} _ {t _ {i}} + \sum_ {j = 0} ^ {s - 1} b _ {i - j} \boldsymbol {x} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t _ {i - j}}, t _ {i - j}\right) + \tilde {\sigma} _ {i} \boldsymbol {\xi}, \quad \boldsymbol {\xi} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}), \tag {14}
$$
where  $\tilde{\sigma}_i = \sigma_{t_{i + 1}}\sqrt{1 - e^{-2\int_{\lambda t_i}^{\lambda t_{i + 1}}\tau^2(\tilde{\lambda})\mathrm{d}\tilde{\lambda}}}$  according to Proposition 4.2 and  $b_{i - j}$  is given by:
$$
b _ {i - j} = \sigma_ {t _ {i + 1}} \int_ {\lambda_ {t _ {i}}} ^ {\lambda_ {t _ {i + 1}}} e ^ {- \int_ {\lambda} ^ {\lambda_ {t _ {i + 1}}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} (1 + \tau^ {2} (\lambda)) e ^ {\lambda} l _ {i - j} (\lambda) \mathrm {d} \lambda , \quad \forall 0 \leq j \leq s - 1 \tag {15}
$$
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  $\mathcal{O}(\sup_{0\leq t\leq T}\tau (t)h + h^s)$ ,
where  $h = \max_{1\leq i\leq M}(t_i - t_{i - 1})$
5.2 SA-Corrector
Eq. (14) offers a "prediction"  $\pmb{x}_{t_{i+1}}^p$  that relies on information preceding or coinciding with the time step  $t_i$  since we only use  $\pmb{x}_{\pmb{\theta}}(\pmb{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_{\pmb{\theta}}(t_i, t_{i+1})$  in Eq. (12). We perform a model evaluation  $\pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_{i+1}}^p, t_{i+1})$  and construct the Lagrange interpolations of  $\pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_{i+1}}^p, t_{i+1}), \pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_i}, t_i), \dots, \pmb{x}_{\pmb{\theta}}(\pmb{x}_{t_{i-(s'-1)}}, t_{i-(\hat{s}-1)})$ :
$$
\hat {L} (t) = \hat {l} _ {i + 1} (t) \boldsymbol {\mathcal {X}} _ {\boldsymbol {\theta}} \left(\boldsymbol {\mathcal {X}} _ {t _ {i + 1}} ^ {p}, t _ {i + 1}\right) + \sum_ {j = 0} ^ {\hat {s} - 1} \hat {l} _ {i - j} (t) \boldsymbol {\mathcal {X}} _ {\boldsymbol {\theta}} \left(\boldsymbol {\mathcal {X}} _ {t _ {i - j}}, t _ {i - j}\right), \tag {16}
$$
where  $\hat{l}_{i-j}(t): \mathbb{R} \to \mathbb{R}$  is the Lagrange basis and  $\hat{s}$  can be different with  $s$  in Eq. (13). The  $\hat{s}$ -step SA-Corrector is derived by replacing  $\boldsymbol{x}_{\theta}(\boldsymbol{x}_{\lambda_u}, \lambda_u)$  with  $\hat{L}(\lambda)$  which is a change-of-variable of  $\hat{L}(t)$ .
$\hat{s}$ -step SA-Corrector Given the initial value  $\pmb{x}_{t_i}$  at time  $t_i$ , a total of  $\hat{s}$  former model evaluations  $\pmb{x}_{\theta}(\pmb{x}_{t_i}, t_i)$ ,  $\pmb{x}_{\theta}(\pmb{x}_{t_{i-1}}, t_{i-1})$ ,  $\dots$ ,  $\pmb{x}_{\theta}(\pmb{x}_{t_{i-(\hat{s}-1)}}, t_{i-(\hat{s}-1)})$ , model evaluation of "prediction"  $\pmb{x}_{\theta}(\pmb{x}_{t_{i+1}}^p, t_{i+1})$ , our  $\hat{s}$ -step SA-Corrector is defined as:
$$
\boldsymbol {x} _ {t _ {i + 1}} = \frac {\sigma_ {t _ {i + 1}}}{\sigma_ {t _ {i}}} e ^ {- \int_ {\lambda_ {t _ {i}}} ^ {\lambda_ {t _ {i + 1}}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} \boldsymbol {x} _ {t _ {i}} + \hat {b} _ {i + 1} \boldsymbol {x} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t _ {i + 1}} ^ {p}, t _ {i + 1}\right) + \sum_ {j = 0} ^ {\hat {s} - 1} \hat {b} _ {i - j} \boldsymbol {x} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t _ {i - j}}, t _ {i - j}\right) + \tilde {\sigma} _ {i} \boldsymbol {\xi}, \tag {17}
$$
where  $\xi \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ ,  $\tilde{\sigma}_i = \sigma_{t_{i + 1}}\sqrt{1 - e^{-2\int_{\lambda t_i}^{\lambda t_{i + 1}}\tau^2(\tilde{\lambda})\mathrm{d}\tilde{\lambda}}}$  according to Proposition 4.2 and the coefficients  $\hat{b}_{i + 1},\hat{b}_{i - j}$  is given by:
$$
\hat {b} _ {i - j} = \sigma_ {t _ {i + 1}} \int_ {\lambda_ {t _ {i}}} ^ {\lambda_ {t _ {i + 1}}} e ^ {- \int_ {\lambda} ^ {\lambda_ {t _ {i + 1}}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} (1 + \tau^ {2} (\lambda)) e ^ {\lambda} \hat {l} _ {i - j} (\lambda) \mathrm {d} \lambda , \quad \forall 0 \leq j \leq s - 1 \tag {18}
$$
$$
\hat {b} _ {i + 1} = \sigma_ {t _ {i + 1}} \int_ {\lambda_ {t _ {i}}} ^ {\lambda_ {t _ {i + 1}}} e ^ {- \int_ {\lambda} ^ {\lambda_ {t _ {i + 1}}} \tau^ {2} (\tilde {\lambda}) \mathrm {d} \tilde {\lambda}} \left(1 + \tau^ {2} (\lambda)\right) e ^ {\lambda} \hat {l} _ {i + 1} (\lambda) \mathrm {d} \lambda
$$
Theorem 5.2 (Strong Convergence of  $\hat{s}$ -step SA-Corrector). Under mild regularity conditions, our  $\hat{s}$ -step SA-Corrector (Eq. (17)) has a global error in strong convergence sense of  $\mathcal{O}(\sup_{0\leq t\leq T}\tau (t)h + h^{\hat{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:
$$
\boldsymbol {x} _ {t _ {i + 1}} = \alpha_ {t _ {i + 1}} \left(\frac {\boldsymbol {x} _ {t _ {i}} - \sigma_ {t _ {i}} \boldsymbol {\epsilon} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t _ {i}} , t _ {i}\right)}{\alpha_ {t _ {i}}}\right) + \sqrt {1 - \alpha_ {t _ {i + 1}} ^ {2} - \hat {\sigma} _ {t _ {i}} ^ {2}} \boldsymbol {\epsilon} _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {t _ {i}}, t _ {i}\right) + \hat {\sigma} _ {t _ {i}} \boldsymbol {\xi}, \tag {19}
$$
where  $\xi \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ ,  $\hat{\sigma}_{t_i}$  is a variable parameter. In practice, DDIM introduces a parameter  $\eta$  such that when  $\eta = 0$ , the sampling process becomes deterministic and when  $\eta = 1$ , the sampling process coincides with original DDPM [2]. Specifically,  $\hat{\sigma}_{t_i} = \eta \sqrt{\frac{1 - \alpha_{t_i + 1}^2}{1 - \alpha_{t_i}^2}\left(1 - \frac{\alpha_{t_i}^2}{\alpha_{t_i + 1}^2}\right)}$ .
Corollary 5.3 (Relationship with DDIM). For any  $\eta$  in DDIM, there exists a  $\tau_{\eta}(t): \mathbb{R} \to \mathbb{R}$  which is a piecewise constant function such that DDIM- $\eta$  coincides with our 1-step SA-Predictor when  $\tau(t) = \tau_{\eta}(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  $\tau(t) \equiv 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  $\tau(t) \equiv 0$  with predictor step  $p$ , corrector step  $p$  in Algorithm 1.
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  $\tau(t) \equiv \tau$  as a constant function or a piecewise constant function. We leave the detailed settings for  $\tau(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  $\tau$  to demonstrate the principles for selecting  $\tau$  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  $\tau(t) \equiv 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 Predictor-Corrector 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
Table 1: Compared results by FID  $\downarrow$  under data-prediction and noise-prediction models, measured by different NFEs. The latent diffusion model in ImageNet 256x256 is used for evaluation.
NFEs	Noise-prediction	Data-prediction
20	310.5	3.88
40	5.85	3.47
60	3.54	3.41
80	3.41	3.38
Table 2: Compared results by FID  $\downarrow$  under different predictor steps and corrector steps, measured by different NFEs. The VE-baseline model [27] in CIFAR10 32x32 is used for evaluation.
method \ setting (NFE, τ)	15,0.4	23,0.8	31,1.0	47,1.4
Predictor 1-steps only	13.76	12.44	11.72	14.67
Predictor 1-steps, Corrector 1-step	8.49	6.87	6.13	6.75
Predictor 3-steps only	5.30	3.93	3.52	2.98
Predictor 3-steps, Corrector 3-steps	4.91	3.77	3.40	2.92
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  $\tau$  values for SA-Solver, namely  $\{0.0, 0.2, 0.4, \dots, 1.6\}$ , where larger value of  $\tau$  correspond to larger magnitude of stochasticity. The FID results under different NFE and  $\tau$  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  $\tau$  values can achieve better FID results; (2) under a considerable number of steps (20-100), large  $\tau$  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  $\tau$  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  $\tau(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], DPM-Solver [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.
Table 3: Sampling quality measured by FID of different sampling methods on DiT, Min-SNR ImageNet [41, 42] models. DiT-XL/2-G and ViT-XL-patch2-32 with  $s = 1.5$  are used.
Model	FID (↓)
DiT ImageNet 256x256	DDPM (NFE=250)	SA-Solver (Ours) (NFE=60)
2.27	2.02
Min-SNR ImageNet 256x256	Heun (NFE=50)	SA-Solver (Ours) (NFE=20)
2.06	1.93
DiT ImageNet 512x512	DDPM (NFE=250)	SA-Solver (Ours) (NFE=60)
3.04	2.80
Figure 1: Sampling quality measured by FID  $\downarrow$  of SA-Solver under a different number of function evaluations (NFE), varying the stochastic noise scale  $\tau$ . For LSUN Bedroom, 10K samples are used to evaluate FID.
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 improves 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  $\tau$  and existing state-of-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  $\tau$  value results in better performance. We also conduct a theoretical analysis that stochasticity can mitigate the error of estimation (see Appendix C).
Figure 2: Sampling quality measured by FID  $\downarrow$  of different sampling methods of DPMs under different NFEs.
Figure 3: Qualitative comparisons between our SA-Solver and previous state-of-the-art methods. All images are generated by Stable Diffusion v1.5 with the same random seed. The main part of the prompt is "portrait of curly orange-haired mad scientist man". We set the guidance scale as 7.5. The proposed SA-Solver is able to generate images with more details.
Figure 4: Sampling quality measured by FID  $\downarrow$  of different sampling methods of DPMs under different training epochs.
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  $\tau(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-of-the-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  $\tau(t)$  selection under different NFEs needs further research. The paper proposes empirical criteria for selecting  $\tau(t)$ , more in-depth theoretical analysis is still needed.
[38] B. Oksendal, Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
[39] P. E. Kloeden and E. Platen, Numerical Solution of Stochastic Differential Equations. Springer, 1992.
[40] A. Roberts, "Modify the improved euler scheme to integrate stochastic differential equations," 2012.
[41] W. Peebles and S. Xie, “Scalable diffusion models with transformers,” arXiv preprint arXiv:2212.09748, 2022.
[42] 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.
[43] 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.
[44] 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.
[45] 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
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 {20}
$$
has the same marginal probability distributions with (4) and (3).
Proof. Denote  $p(\pmb{x}, t): \mathbb{R}^d \times [0, T] \to \mathbb{R}^+$  as the probability density function of  $\pmb{x}_t$  at time  $t$ , thus  $p(\pmb{x}, T) = p_T(\pmb{x})$ . Fokker-Planck equation [38] determines a Partial Differential Equation (PDE) that  $p_t(\pmb{x})$  satisfies:
$$
\begin{array}{l} - \frac {\partial p _ {t} (\boldsymbol {x})}{\partial t} = - \sum_ {i = 1} ^ {d} \frac {\partial}{\partial x _ {i}} \left[ - \left[ f (t) p _ {t} (\boldsymbol {x}) x _ {i} - \left(\frac {1 + \tau^ {2} (t)}{2}\right) g ^ {2} (t) p _ {t} (\boldsymbol {x}) \frac {\partial \log p _ {t} (\boldsymbol {x})}{\partial x _ {i}} \right] \right] \\ + \frac {1}{2} \sum_ {i = 1} ^ {d} \sum_ {j = 1} ^ {d} \frac {\partial^ {2}}{\partial x _ {i} \partial x _ {j}} \left[ \tau^ {2} (t) g ^ {2} (t) \delta_ {i j} p _ {t} (\boldsymbol {x}) \right] \tag {21} \\ = \sum_ {i = 1} ^ {d} \frac {\partial}{\partial x _ {i}} \left[ f (t) p _ {t} (\pmb {x}) x _ {i} - \left(\frac {1 + \tau^ {2} (t)}{2}\right) g ^ {2} (t) p _ {t} (\pmb {x}) \frac {\partial \log p _ {t} (\pmb {x})}{\partial x _ {i}} \right] \\ + \frac {1}{2} \sum_ {i = 1} ^ {d} \frac {\partial}{\partial x _ {i}} \left[ \frac {\partial}{\partial x _ {i}} \left[ \tau^ {2} (t) g ^ {2} (t) p _ {t} (\boldsymbol {x}) \right] \right]. \\ \end{array}
$$
Eq. (20) is a reverse-time SDE running from  $T$  to 0, thus there are two additional minus signs in Eq. (21) before term  $\frac{\partial p_t(\pmb{x})}{\partial t}$  and term  $\left[f(t)p_t(\pmb {x})x_i - \left(\frac{1 + \tau^2(t)}{2}\right)g^2 (t)p_t(\pmb {x})\frac{\partial\log p_t(\pmb{x})}{\partial x_i}\right]$  compared with vanilla Fokker-Planck equation in general cases. Here  $\delta_{ij}$  is the Dirac symbol satisfies  $\delta_{ij} = 1$  when  $i = j$ , otherwise,  $\delta_{ij} = 0$ . Notice that
$$
\frac {\partial}{\partial x _ {i}} \left[ \tau^ {2} (t) g ^ {2} (t) p _ {t} (\boldsymbol {x}) \right] = \tau^ {2} (t) g ^ {2} (t) \frac {\partial}{\partial x _ {i}} p _ {t} (\boldsymbol {x}) = \tau^ {2} (t) g ^ {2} (t) p _ {t} (\boldsymbol {x}) \frac {\partial \log p _ {t} (\boldsymbol {x})}{\partial x _ {i}}. \tag {22}
$$
Substituting Eq. (22) into Eq. (21), we obtain that
$$
\begin{array}{l} - \frac {\partial p _ {t} (\pmb {x})}{\partial t} = \sum_ {i = 1} ^ {d} \frac {\partial}{\partial x _ {i}} \left[ f (t) p _ {t} (\pmb {x}) x _ {i} - \left(\frac {1 + \tau^ {2} (t)}{2}\right) g ^ {2} (t) p _ {t} (\pmb {x}) \frac {\partial \log p _ {t} (\pmb {x})}{\partial x _ {i}} \right] \\ + \frac {1}{2} \sum_ {i = 1} ^ {d} \frac {\partial}{\partial x _ {i}} \left[ \tau^ {2} (t) g ^ {2} (t) p _ {t} (\boldsymbol {x}) \frac {\partial \log p _ {t} (\boldsymbol {x})}{\partial x _ {i}} \right] \tag {23} \\ = \sum_ {i = 1} ^ {d} \frac {\partial}{\partial x _ {i}} \left[ f (t) p _ {t} (\pmb {x}) x _ {i} - \frac {1}{2} g ^ {2} (t) p _ {t} (\pmb {x}) \frac {\partial \log p _ {t} (\pmb {x})}{\partial x _ {i}} \right], \\ \end{array}
$$
which is independent of  $\tau(t)$ . With the same initial condition  $p(\boldsymbol{x}, T) = p_T(\boldsymbol{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  $\tau(t) = 0$ , Eq. (20) degenerates to diffusion ODEs and when  $\tau(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  $\nabla_{\pmb{x}}\log p_t(\pmb{x}_t)$  with  $s_\theta (\pmb {x}_t,t)$  and reparameterizing  $s_{\pmb{\theta}}(\pmb {x}_t,t)$  with  $-(x_{t} - \alpha_{t}\pmb{x}_{\pmb{\theta}}(\pmb{x}_{t},t)) / \sigma_{t}^{2}$ , Eq. (20) becomes
$$
\mathrm {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] \mathrm {d} t + \tau (t) g (t) \mathrm {d} \bar {\boldsymbol {w}} _ {t}. \tag {24}
$$