Maximum Likelihood Training for Score-Based Diffusion ODEs by High-Order Denoising Score Matching

Introduction

Score-based generative models (SGMs) have shown promising performance in both sample quality and likelihood evaluation (Song et al., 2020b; Vahdat et al., 2021; Dockhorn et al., 2021; Kingma et al., 2021), with applications on many tasks, such as image generation (Dhariwal & Nichol, 2021; Meng et al., 2021b), speech synthesis (Chen et al., 2020; Kong et al., 2020), shape generation (Luo & Hu, 2021; Zhou et al., 2021) and lossless compression (Kingma et al., 2021). SGMs define a forward diffusion process by a stochastic differential equation (SDE) that gradually perturbs the data distribution to a simple noise distribution. The forward diffusion process has an equivalent reverse-time process with analytical form (Anderson, 1982), which depends on the score function (the gradient of the log probability density) of the forward process at each time step. SGMs learn a parameterized neural network (called “score model”) (Song et al., 2020a; Song & Ermon, 2019, 2020; Song et al., 2020b) to estimate the score functions, and further define probabilistic models by the score model.

Two types of probabilistic models are proposed with the learned score model. One is the score-based diffusion SDE (Song et al. (2020b), “ScoreSDE” for short), which is defined by approximately reversing the diffusion process from the noise distribution by the score model and it can generate high-quality samples. The other is the score-based diffusion ordinary differential equation (Song et al. (2020b, 2021), “ScoreODE” for short), which is defined by approximating the probability flow ODE (Song et al., 2020b) of the forward process whose marginal distribution is the same as the forward process at each time. ScoreODEs can be viewed as continuous normalizing flows (Chen et al., 2018). Thus, unlikely SDEs, they can compute exact likelihood by ODE solvers (Grathwohl et al., 2018).

In practice, the score models are trained by minimizing an objective of a time-weighted combination of score matching losses (Hyvärinen & Dayan, 2005; Vincent, 2011). It is proven that for a specific weighting function, minimizing the score matching objective is equivalent to maximizing the likelihood of the ScoreSDE (Song et al., 2021). However, minimizing the score matching objective does not necessarily maximize the likelihood of the ScoreODE. This is problematic since ScoreODEs are used for likelihood estimation. In fact, ScoreODEs trained with the score matching objective may even fail to estimate the likelihood of very simple data distributions, as shown in Fig. 1.

In this paper, we systematically analyze the relationship between the score matching objective and the maximum likelihood objective of ScoreODEs, and present a new algorithm to learn ScoreODEs via maximizing likelihood. Theoretically, we derive an equivalent form of the KL divergenceMaximum likelihood estimation is equivalent to minimizing a KL divergence. between the data distribution and the ScoreODE distribution, which explicitly reveals the gap between the KL divergence and the original first-order score matching objective. Computationally, as directly optimizing the equivalent form requires an ODE solver, which is time-consuming (Chen et al., 2018), we derive an upper bound of the KL divergence by controlling the approximation errors of the first, second, and third-order score matching. The optimal solution for this score model by minimizing the upper bound is still the data score function, so the learned score model can still be used for the sample methods in SGMs (such as PC samplers in (Song et al., 2020b)) to generate high-quality samples.

Based on the analyses, we propose a novel high-order denoising score matching algorithm to train the score models, which theoretically guarantees bounded approximation errors of high-order score matching. We further propose some scale-up techniques for practice. Our experimental results empirically show that the proposed training method can improve the likelihood of ScoreODEs on both synthetic data and CIFAR-10, while retaining the high generation quality.

Score-Based Generative Models

We first review the dynamics and probability distributions defined by SGMs. The relationship between these dynamics is illustrated in Fig. 2(a).

Let $q_{0}(\bm{x}_{0})$ denote an unknown $d$-dimensional data distribution. SGMs define an SDE from time to $T$ (called “forward process”) with $\bm{x}_{0}\sim q_{0}(\bm{x}_{0})$ and

where the marginal distribution of $\bm{x}_{t}$ is $p^{\text{SDE}}_{t}(\bm{x}_{t})$. By solving the parameterized reverse SDE, SGMs achieve excellent sample quality in many tasks (Song et al., 2020b; Dhariwal & Nichol, 2021). Moreover, the KL divergence between $q_{0}(\bm{x}_{0})$ and $p_{0}^{\text{SDE}}(\bm{x}_{0})$ can be bounded by the score matching error with the weight $g(\cdot)^{2}$ (Song et al., 2021):

Therefore, minimizing $\mathcal{J}_{\text{SM}}(\theta;g(\cdot)^{2})$ is equivalent to maximum likelihood training of $p^{\text{SDE}}_{0}$. In the rest of this work, we mainly consider the maximum likelihood perspective, namely $\lambda(\cdot)=g(\cdot)^{2}$, and denote $\mathcal{J}_{\text{SM}}(\theta)\coloneqq\mathcal{J}_{\text{SM}}(\theta;g(\cdot)^{2})$.

2 ScoreODEs and Exact Likelihood Evaluation

where the marginal distribution of $\bm{x}_{t}$ is $p_{t}^{\text{ODE}}(\bm{x}_{t})$. By the “Instantaneous Change of Variables” (Chen et al., 2018), $\log p_{t}^{\text{ODE}}(\bm{x}_{t})$ of $\bm{x}_{t}$ can be computed by integrating:

Relationship between Score Matching and KL Divergence of ScoreODEs

The likelihood evaluation needs the distribution $p_{0}^{\text{ODE}}$ of ScoreODEs, which is usually different from the distribution $p^{\text{SDE}}_{0}$ of ScoreSDEs (see Appendix B for detailed discussions). Although the score matching objective $\mathcal{J}_{\text{SM}}(\theta)$ can upper bound the KL divergence of ScoreSDEs (up to constants) according to Eqn. (4), it is not enough to upper bound the KL divergence of ScoreODEs. To the best of our knowledge, there is no theoretical analysis of the relationship between the score matching objectives and the KL divergence of ScoreODEs.

Let $q_{0}$ be the data distribution and $q_{t}$ be the marginal distribution at time $t$ through the forward diffusion process defined in Eqn. (1), and $p^{\text{ODE}}_{t}$ be the marginal distribution at time $t$ through the ScoreODE defined in Eqn. (6). Under some regularity conditions in Appendix A, we have

2 Bounding the KL Divergence of ScoreODEs by High-Order Score Matching Errors

By straightforwardly using Cauchy-Schwarz inequality for Eqn. (8), we have

Error-Bounded High-Order Denoising Score Matching (DSM)

We present some preliminaries on first-order denoising score matching (Vincent, 2011; Song et al., 2020b) in this section.

by optimizing the (first-order) DSM objective:

where $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$ and $\bm{x}_{t}=\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon}$.

2 High-Order Denoising Score Matching

In this section, we generalize the first-order DSM to second and third orders for matching the second-order and third-order score functions defined in Theorem 3.2. We firstly present the second-order method below.

Moreover, denote the first-order score matching error as $\delta_{1}(\bm{x}_{t},t)\coloneqq\|\hat{\bm{s}}_{1}(\bm{x}_{t},t)-\nabla_{\bm{x}}\log q_{t}(\bm{x}_{t})\|_{2}$, then $\forall\bm{x}_{t},\theta$, the score matching error for $\bm{s}_{2}(\bm{x}_{t},t;\theta)$ can be bounded by

Theorem 4.1 shows that the DSM objective in problem (13) is a valid surrogate for second-order score matching, because the difference between the model score $\bm{s}_{2}(\bm{x}_{t},t;\theta)$ and the true second-order score $\nabla^{2}_{\bm{x}}\log q_{t}(\bm{x}_{t})$ can be bounded by the training error $\|\bm{s}_{2}(\bm{x}_{t},t;\theta)-\bm{s}_{2}(\bm{x}_{t},t;\theta^{*})\|_{F}$ and the first-order score matching error $\delta_{1}(\bm{x}_{t},t)$. Note that previous second-order DSM method proposed in Meng et al. (2021a) does not have such error-bounded property, and we refer to Sec. 6 and Appendix F.4 for the detailed comparison. In addition, recent work (Bao et al., 2022, Theorem 3.3) about the optimal covariance of the diffusion models considers the estimation error of the mean of the diffusion models, which is equivalent to our proposed error-bounded second-order DSM with considering the first-order score matching error.

The DSM objective in problem (13) requires learning a matrix-valued function $\nabla^{2}_{\bm{x}}\log q_{t}(\bm{x}_{t})$, but sometimes we only need the trace $\tr(\nabla^{2}_{\bm{x}}\log q_{t}(\bm{x}_{t}))$ of the second-order score function. Below we present a corollary for only matching the trace of the second-order score function.

The estimation error for $\bm{s}_{2}^{\text{trace}}(\bm{x}_{t},t;\theta)$ can be bounded by:

Finally, we present the third-order DSM method. The score matching error can also be bounded by the training error and the first, second-order score matching errors.

Denote the first-order score matching error as $\delta_{1}(\bm{x}_{t},t)\coloneqq\|\hat{\bm{s}}_{1}(\bm{x}_{t},t)-\nabla_{\bm{x}}\log q_{t}(\bm{x}_{t})\|_{2}$ and the second-order score matching errors as $\delta_{2}(\bm{x}_{t},t)\coloneqq\|\hat{\bm{s}}_{2}(\bm{x}_{t},t)-\nabla_{\bm{x}}^{2}\log q_{t}(\bm{x}_{t})\|_{F}$ and $\delta_{2,\text{tr}}(\bm{x}_{t},t)\coloneqq|\tr(\hat{\bm{s}}_{2}(\bm{x}_{t},t))-\tr(\nabla_{\bm{x}}^{2}\log q_{t}(\bm{x}_{t}))|$. Then $\forall\bm{x}_{t},\theta$, the score matching error for $\bm{s}_{3}(\bm{x}_{t},t;\theta)$ can be bounded by:

Theorem 4.3 shows that the third-order score matching needs first-order score matching, second-order score matching and trace of second-order score matching. The third-order score matching error can also be bounded by the training error and the lower-order score matching errors. We believe that our construction for the error-bounded high-order DSM can be extended to even higher orders by carefully designing the training objectives. In this paper, we only focus on the second and third-order methods, and leave this extension for future work.

Training Score Models by High-Order DSM

Building upon the high-order DSM for a specific time $t$ (see Sec. 4), we present our algorithm to train ScoreODEs in the perspective of maximum likelihood by considering all timesteps $t\in[0,T]$. Practically, we further leverage the “noise-prediction” trick (Kingma et al., 2021; Ho et al., 2020) for variance reduction and the Skilling-Hutchinson trace estimator (Skilling, 1989; Hutchinson, 1989) to compute the involved high-order derivatives efficiently. The whole training algorithm is presented detailedly in Appendix. H.

Theoretically, to train score models for all $t\in[0,T]$, we need to integrate the DSM objectives in Eqn. (12)(13)(15)(16) from $t=0$ to $t=T$, which needs ODE solvers and is time-consuming. Instead, in practice, we follow the method in (Song et al., 2020b; Ho et al., 2020) which uses Monte-Carlo method to unbiasedly estimate the objectives by sample $t\in[0,T]$ from a proposal distribution $p(t)$, avoiding ODE solvers. The main problem for the Monte-Carlo method is that sometimes the sample variance of the objectives may be large due to the different value of $\frac{1}{\sigma_{t}}$. To reduce the variance, we take the time-reweighted objectives by multiplying $\sigma_{t}^{2}$,$\sigma_{t}^{4}$,$\sigma_{t}^{6}$ with the corresponding first, second, third-order score matching objectives at each time $t$, which is known as the “noise-prediction” trick (Kingma et al., 2021; Ho et al., 2020) for the first-order DSM objective. Specifically, assume that $\bm{x}_{0}\sim q_{0}(\bm{x}_{0})$, the time proposal distribution $p(t)=\mathcal{U}[0,T]$ is the uniform distribution in $[0,T]$, the random noise $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$ follows the standard Gaussian distribution, and let $\bm{x}_{t}=\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon}$. The training objective for the first-order DSM in Eqn. (12) through all $t\in[0,T]$ is

which empirically has low sample variance (Song et al., 2020b; Ho et al., 2020).

As for the second and third-order DSM objectives, we take $\hat{\bm{s}}_{1}(\bm{x}_{t},t)\coloneqq\bm{s}_{\theta}(\bm{x}_{t},t)$ for the first-order score function estimation, and $\hat{\bm{s}}_{2}(\bm{x}_{t},t)\coloneqq\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x}_{t},t)$ for the second-order score function estimation, and both disable the gradient computations for $\theta$ (which can be easily implemented by stop_gradient or detach). The reason for disabling gradients is because the high-order DSM only needs the estimation values of the lower-order score functions, as shown in Theorem. 4.1 and 4.3. The training objectives through all $t\in[0,T]$ for the second-order DSM in Eqn. (13), for the trace of second-order DSM in Eqn. (15) and for the third-order DSM in Eqn. (16) are

where $\lambda_{1},\lambda_{2}\geq 0$ are hyperparameters, and we discuss in Appendix. I.1 for details.

Although as shown in Eqn. (7), we can exactly compute $\log p_{0}^{\text{ODE}}(\bm{x}_{0})$ for a given data point $\bm{x}_{0}$ via solving the ODE (Chen et al., 2018; Grathwohl et al., 2018), this method is hard to scale up for the large neural networks used in SGMs. For example, it takes $2\sim 3$ minutes for evaluating $\log p^{\text{ODE}}_{0}$ for a single batch of the ScoreODE used in (Song et al., 2020b). Instead, our proposed method leverages Monte-Carlo methods to avoid the ODE solvers.

2 Scalability and Numerical Stability

As the second and third-order DSM objectives need to compute the high-order derivatives of a neural network and are expensive for high-dimensional data, we use the Skilling-Hutchinson trace estimator (Skilling, 1989; Hutchinson, 1989) to unbiasedly estimate the trace of the Jacobian (Grathwohl et al., 2018) and the Frobenius norm of the Jacobian (Finlay et al., 2020). The detailed objectives for high-dimensional data can be found in Appendix. G.

In practice, we often face numerical instability problems for $t$ near to . We follow Song et al. (2021) to choose a small starting time $\epsilon>0$, and both of the training and the evaluation are performed for $t\in[\epsilon,T]$ instead of $[0,T]$. Note that the likelihood evaluation is still exact, because we use $p_{\epsilon}^{\text{ODE}}(\bm{x}_{0})$ to compute the log-likelihood of $\bm{x}_{0}$, which is still a well-defined density model (Song et al., 2021).

Related Work

ScoreODEs are special formulations of Neural ODEs (NODEs) (Chen et al., 2018), and our proposed method can be viewed as maximum likelihood training of NODEs by high-order score matching. Traditional maximum likelihood training for NODEs aims to match the NODE distribution at $t=0$ with the data distribution, and it cannot control the distribution between and $T$. Finlay et al. (2020) show that training NODEs by simply maximizing likelihood could result in unnecessary complex dynamics, which is hard to solve. Instead, our high-order score matching objective is to match the distributions between the forward process distribution and the NODE distribution at each time $t$, and empirically the dynamics is kind of smooth. Moreover, our proposed algorithm uses the Monto-Carlo method to unbiasedly estimate the objectives, which does not need any black-box ODE solvers. Therefore, our algorithm is suitable for maximum likelihood training for large-scale NODEs.

Recently, Meng et al. (2021a) propose a high-order DSM method for estimating the second-order data score functions. However, the training objective of our proposed error-bounded high-order DSM is different from that of Meng et al. (2021a). Our proposed algorithm can guarantee a bounded error of the high-order score matching exactly by the lower-order estimation errors and the training error, while the score matching error raised by minimizing the objective in Meng et al. (2021a) may be unbounded, even when the lower-order score matching error is small and the training error of the high-order DSM is zero (see Appendix F.4 for detailed analysis). Moreover, our method can also be used to train a separate high-order score model in other applications, such as uncertainty quantification and Ozaki sampling presented in Meng et al. (2021a).

Experiments

In this section, we demonstrate that our proposed high-order DSM algorithm can improve the likelihood of ScoreODEs, while retaining the high sample quality of the corresponding ScoreSDEs. Particularly, we use the Variance Exploding (VE) (Song et al., 2020b) type diffusion models, which empirically have shown high sample quality by the ScoreSDE but poor likelihood by the ScoreODE (Song et al., 2021) when trained by minimizing the first-order score matching objective $\mathcal{J}_{\text{SM}}(\theta)$. We implement our experiments by JAX (Bradbury et al., 2018), which is efficient for computing the derivatives of the score models. In all experiments, we choose the start time $\epsilon=10^{-5}$, which follows the default settings in Song et al. (2020b). In this section, we refer to “first-order score matching” as minimizing the objective in Eqn. (19) with $\lambda_{1}=\lambda_{2}=0$, and “second-order score matching” as minimizing the one with $\lambda_{2}=0$. Please see Appendix. I.1 for detailed settings about $\lambda_{1}$ and $\lambda_{2}$. The released code can be found at https://github.com/LuChengTHU/mle_score_ode.

Song et al. (2021) use a proposal distribution $t\sim p(t)$ to adjust the weighting for different time $t$ to minimize $\mathcal{J}_{\text{SM}}(\theta)$. In our experiments, we mainly focus on the VE type, whose proposal distribution $p(t)=\mathcal{U}[0,T]$ is the same as that of our final objective in Eqn. (19).

We take an example to demonstrate how the high-order score matching training impacts the weighted Fisher divergence $\mathcal{J}_{\text{Fisher}}(\theta)$ and the model density of ScoreODEs. Let $q_{0}$ be a 1-D Gaussian mixture distribution as shown in Fig. 1 (a). We train a score model of VE type by minimizing the first-order score matching objective $\mathcal{J}_{\text{SM}}(\theta)$, and use the corresponding ScoreODE to evaluate the model density, as shown in Fig. 1 (b). The model density achieved by only first-order score matching is quite different from the data distribution at some data points. However, by third-order score matching, the model density in Fig. 1 (c) is quite similar to the data distribution, showing the effectiveness of our proposed algorithm.

We further analyze the difference of $\mathcal{J}_{\text{Fisher}}(\theta)$ between the first, second, third-order score matching training. As $q_{0}$ is a Gaussian mixture distribution, we can prove that each $q_{t}$ is also a Gaussian mixture distribution, and $\nabla_{\bm{x}}\log q_{t}$ can be analytically computed. Denote

2 Density Modeling on 2-D Checkerboard Data

We then train ScoreODEs (VE type) by the first, second, third-order score matchings on the checkerboard data whose density is multi-modal, as shown in Fig. 4. We use a simple MLP neural network with Swish activations (Ramachandran et al., 2017), and the detailed settings are in Appendix. I.3.

3 Density Modeling on Image Datasets

We also train ScoreODEs (VE type) on both the CIFAR-10 dataset (Krizhevsky, 2009) and the ImageNet 32x32 dataset (Deng et al., 2009), which are two of the most popular datasets for generative modeling and likelihood evaluation. We use the estimated training objectives by trace estimators, which are detailed in Appendix. G. We use the same neural networks and hyperparameters as the NCSN++ cont. model in (Song et al., 2020b) (denoted as VE) and the NCSN++ cont. deep model in (Song et al., 2020b) (denoted as VE (deep)), respectively (see detailed settings in Appendix. I.4 and Appendix. I.5). We evaluate the likelihood by the ScoreODE $p^{\text{ODE}}_{\epsilon}$, and sample from the ScoreSDE by the PC sampler (Song et al., 2020b). As shown in Table 1, our proposed high-order score matching can improve the likelihood performance of ScoreODEs, while retraining the high sample quality of the corresponding ScoreSDEs (see Appendix. J for samples). Moreover, the computation costs of the high-order DSM objectives are acceptable because of the trace estimators and the efficient “Jacobian-vector-product” computation in JAX. We list the detailed computation costs in Appendix. I.4.

We also compare the VE, VP and subVP types of ScoreODEs trained by the first-order DSM and the third-order DSM, and the detailed results are listed in Appendix. J.

Conclusion

We propose maximum likelihood training for ScoreODEs by novel high-order denoising score matching methods. We analyze the relationship between the score matching objectives and the KL divergence from the data distribution to the ScoreODE distribution. Based on it, we provide an upper bound of the KL divergence, which can be controlled by minimizing the first, second, and third-order score matching errors of score models. To minimize the high-order score matching errors, we further propose a high-order DSM algorithm, such that the higher-order score matching error can be bounded by exactly the training error and the lower-order score matching errors. The optimal solution for the score model is still the same as the original training objective of SGMs. Empirically, our method can greatly improve the model density of ScoreODEs of the Variance Exploding type on several density modeling benchmarks. Finally, we believe that our training method is also suitable for other SGMs, including the Variance Preserving (VP) type (Song et al., 2020b), the latent space type (Vahdat et al., 2021) and the critically-damped Langevin diffusion type (Dockhorn et al., 2021). Such extensions are left for future work.

Acknowledgements

This work was supported by National Key Research and Development Project of China (No. 2021ZD0110502); NSF of China Projects (Nos. 62061136001, 61620106010, 62076145, U19B2034, U1811461, U19A2081, 6197222, 62106120); Beijing NSF Project (No. JQ19016); Beijing Outstanding Young Scientist Program NO. BJJWZYJH012019100020098; a grant from Tsinghua Institute for Guo Qiang; the NVIDIA NVAIL Program with GPU/DGX Acceleration; the High Performance Computing Center, Tsinghua University; and Major Innovation & Planning Interdisciplinary Platform for the “Double-First Class” Initiative, Renmin University of China.

References

Appendix A Assumptions

We follow the regularity assumptions in (Song et al., 2021) to ensure the existence of reverse-time SDEs and probability flow ODEs and the correctness of the “integration by parts” tricks for the computation of KL divergence. And to ensure the existence of the third score function, we change the assumptions of differentiability. For completeness, we list all these assumptions in this section.

We make the following assumptions, most of which are presented in (Song et al., 2021):

$g\in\mathcal{C}$ and $\forall t\in[0,T],|g(t)|>0$.

$\forall t\in[0,T],\exists k>0:q_{t}(\bm{x})=O(e^{-\|\bm{x}\|_{2}^{k}})$, $p_{t}^{\text{SDE}}(\bm{x})=O(e^{-\|\bm{x}\|_{2}^{k}})$, $p_{t}^{\text{ODE}}(\bm{x})=O(e^{-\|\bm{x}\|_{2}^{k}})$ as $\|\bm{x}\|_{2}\rightarrow\infty$.

Appendix B Distribution gap between score-based diffusion SDEs and ODEs

We list the corresponding equivalent SDEs and ODEs of $q_{t}(\bm{x}_{t})$, $p_{t}^{\text{SDE}}(\bm{x}_{t})$ and $p_{t}^{\text{ODE}}(\bm{x}_{t})$ in Table 2. In most cases, $q_{t}(\bm{x}_{t})$, $p_{t}^{\text{SDE}}(\bm{x}_{t})$ and $p_{t}^{\text{ODE}}(\bm{x}_{t})$ are different distributions, which we will prove in this section.

According to the reverse SDE and probability flow ODE listed in the table, it is obvious that when $\bm{s}_{\theta}(\bm{x}_{t},t)\neq\nabla\log q_{t}(\bm{x}_{t})$, $q_{t}(\bm{x}_{t})$ and $p_{t}^{\text{SDE}}$ are different, and $q_{t}(\bm{x}_{t})$, $p_{t}^{\text{ODE}}$ are different. Below we show that in most cases, $p_{t}^{\text{SDE}}$ and $p_{t}^{\text{ODE}}$ are also different.

Firstly, we should notice that the probability flow ODE of the ScoreSDE in Eqn. (3) is

where the distribution of $\bm{x}_{t}$ during the trajectory is also $p_{t}^{\text{SDE}}(\bm{x}_{t})$. Below we show that for the commonly-used SGMs, in most cases, the probability flow ODE in Eqn. (20) of the ScoreSDE is different from the ScoreODE in Eqn. (6).

On the other hand, if we start the forward process in Eqn. (1) with $\bm{x}_{0}\sim p^{\text{SDE}}_{0}$, and the distribution of $\bm{x}_{t}$ at time $t$ during its trajectory follows the same equation as Eqn. (21). Therefore, the distribution of $\bm{x}_{t}$ is also $p_{t}^{\text{SDE}}$.

The ”variational gap” of ScoreSDEs in (Huang et al., 2021) is the gap between the “joint distribution” KL divergence of $p^{\text{SDE}}_{0:T}$ and “marginal distribution” KL divergence of $p_{0}^{\text{SDE}}$, which does not include $p^{\text{ODE}}_{t}$. We refer to Appendix C for further discussions.

When $\mathcal{J}_{\text{SM}}(\theta)=0$, we have $\bm{s}_{\theta}(\cdot,t)\equiv\nabla_{\bm{x}}\log q_{t}$. In this case, the ScoreSDE in Eqn. (3) becomes the reverse diffusion SDE in Eqn. (2), and the ScoreODE in Eqn. (6) becomes the probability flow ODE in Eqn. (5). However, if $\bm{f}(\bm{x}_{t},t)$ is linear to $\bm{x}_{t}$ and $q_{0}(\bm{x}_{0})$ is not Gaussian, we have $q_{T}$ is not Gaussian, so $q_{T}\neq p_{T}^{\text{SDE}}$ and $q_{T}\neq p_{T}^{\text{ODE}}$. Therefore, in this case, we still have $\bm{s}_{\theta}(\cdot,T)=\nabla_{\bm{x}}\log q_{T}\neq\nabla_{\bm{x}}\log p_{T}^{\text{SDE}}$, so the probability flow in Eqn. (20) of the ScoreSDE is still different from the ScoreODE in Eqn. (6), leading to the fact that $p_{0}^{\text{SDE}}\neq p_{0}^{\text{ODE}}$. (But the difference is extremely small, because they are both very similar to $q_{0}$).

Appendix C KL divergence and variational gap of ScoreSDEs

By Eqn. (20), we denote the probability flow ODE of ScoreSDE as

First we rewrite the KL divergence from $q_{0}$ to $p_{0}^{\text{SDE}}$ in an integral form

Given a fixed $\bm{x}$, by the special case of Fokker-Planck equation with zero diffusion term, we can derive the time-evolution of ODE’s associated probability density function by:

where Eqn. (27) is due to integration by parts under the Assumption. A.1(11), which shows that $\lim_{\bm{x}\rightarrow\infty}\bm{h}_{q}(\bm{x},t)q_{t}(\bm{x})=0$ and $\lim_{\bm{x}\rightarrow\infty}\bm{h}_{p^{\text{SDE}}}(\bm{x},t)p_{t}^{\text{SDE}}(\bm{x})=0$ for all $t\in[0,T]$. Combining with Eqn. (25), we can finish the proof. ∎

Thus by Eqn. (22) and Proposition C.1, the expectation of the variational gap for $p_{0}^{\text{SDE}}$ under data distribution $q_{0}$ is

Note that this is equivalent to the expectation of the variational gap in (Huang et al., 2021). This expression tells us that

Appendix D Instantaneous change of score function of ODEs

Given a fixed $\bm{x}$, by the special case of Fokker-Planck equation with zero diffusion term, we have

Appendix E KL divergence of ScoreODEs

In this section, we propose the proofs for Sec. 3.

First we rewrite the KL divergence from $q_{0}$ to $p_{0}^{\text{ODE}}$ in an integral form

Given a fixed $\bm{x}$, by the special case of Fokker-Planck equation with zero diffusion term, we can derive the time-evolution of ODE’s associated probability density function by:

where Eqn. (45) is due to integration by parts under the Assumption. A.1(11), which shows that $\lim_{\bm{x}\rightarrow\infty}\bm{h}_{q}(\bm{x},t)q_{t}(\bm{x})=0$ and $\lim_{\bm{x}\rightarrow\infty}\bm{h}_{p}(\bm{x},t)p_{t}^{\text{ODE}}(\bm{x})=0$ for all $t\in[0,T]$. Combining with Eqn. (40), we can conclude that

E.2 Proof of Theorem 3.2

Given a fixed $\bm{x}$, by the special case of Fokker-Planck equation with zero diffusion term, we have

Similarly, for $\nabla\log q_{t}(\bm{x}_{t})$ we have

where $\|\cdot\|_{2}$ is the $L_{2}$-norm for vectors and the induced $2$-norm for matrices.

For simplicity, we denote $\bm{h}_{p}(s)\coloneqq\bm{h}_{p}(\bm{x}_{s},s)$, $\bm{h}_{q}(s)\coloneqq\bm{h}_{q}(\bm{x}_{s},s)$, $p_{s}\coloneqq p_{s}^{\text{ODE}}(\bm{x}_{s})$ and $q_{s}\coloneqq q_{s}(\bm{x}_{s})$. We use $\|\cdot\|$ to denote the 2-norm for vectors and matrices. As

Then $\alpha(t)\geq 0,\beta(t)\geq 0$ are independent of $\theta$, and we have

E.3 Maximum likelihood training of ScoreODE by ODE solvers

We can directly call ODE solvers to compute $\log p_{0}^{\text{ODE}}(\bm{x}_{0})$ for a given data point $\bm{x}_{0}$, and thus do maximum likelihood training (Chen et al., 2018; Grathwohl et al., 2018). However, this needs to be done at every optimization step, which is hard to scale up for the large neural networks used in SGMs. For example, it takes $2\sim 3$ minutes for evaluating $\log p^{\text{ODE}}_{0}$ for a single batch of the ScoreODE used in (Song et al., 2020b). Besides, directly maximum likelihood training for $p^{\text{ODE}}_{0}$ cannot make sure that the optimal solution for $\bm{s}_{\theta}(\bm{x}_{t},t)$ is the data score function $\nabla_{\bm{x}}\log q_{t}(\bm{x}_{t})$, because the directly MLE for $p_{0}^{\text{ODE}}$ cannot ensure that the model distribution $p_{t}^{\text{ODE}}$ is also similar to the data distribution $q_{t}$ at each time $t\in(0,T)$, thus cannot ensure the model ODE function $\bm{h}_{p}(\bm{x}_{t},t)$ is similar to the data ODE function $\bm{h}_{q}(\bm{x}_{t},t)$. In fact, ScoreODE is a special formulation of Neural ODEs (Chen et al., 2018). Empirically, Finlay et al. (2020) find that directly maximum likelihood training of Neural ODEs may cause rather complex dynamics ($\bm{h}_{p}(\bm{x}_{t},t)$ here), which indicates that directly maximum likelihood training by ODE solvers cannot ensure $\bm{h}_{p}(\bm{x}_{t},t)$ be similar to $\bm{h}_{q}(\bm{x}_{t},t)$, and thus cannot ensure the score model $\bm{s}_{\theta}(\bm{x}_{t},t)$ be similar to the data score function $\nabla_{\bm{x}}\log q_{t}(\bm{x}_{t})$.

Appendix F Error-bounded High-Order Denoising Score Matching

In this section, we present all the lemmas and proofs for our proposed error-bounded high-order DSM algorithm.

Below we propose an expectation formulation of the first-order score function.

Assume $(\bm{x}_{t},\bm{x}_{0})\sim q(\bm{x}_{t},\bm{x}_{0})$, we have

We use $\nabla(\cdot)$ to denote the derivative of $\bm{x}_{t}$, namely $\nabla_{\bm{x}_{t}}(\cdot)$.

And below we propose an expectation formulation for the second-order score function.

Assume $(\bm{x}_{t},\bm{x}_{0})\sim q(\bm{x}_{t},\bm{x}_{0})$, we have

We use $\nabla(\cdot)$ to denote the derivative of $\bm{x}_{t}$, namely $\nabla_{\bm{x}_{t}}(\cdot)$. Firstly, the gradient of $q_{t0}$ w.r.t. $\bm{x}_{t}$ can be calculated as

Then using Lemma F.1 and the product rule, we have

Thus Eqn. (96) can be further transformed by subtracting $\bm{0}$ as

which completes the proof of second-order score matrix and its trace. ∎

Below we propose a corollary of an expectation formulation of the sum of the first-order and the second-order score functions, which can be used to design the third-order denoising score matching method.

Assume $(\bm{x}_{t},\bm{x}_{0})\sim q(\bm{x}_{t},\bm{x}_{0})$, we have

This is the direct corollary of Eqn. (96) by adding $\nabla_{\bm{x}_{t}}\log q_{t}(\bm{x}_{t})\nabla_{\bm{x}_{t}}\log q_{t}(\bm{x}_{t})^{\top}$ to both sides and using Lemma F.1. ∎

And below we propose an expectation formulation for the third-order score function.

Assume $(\bm{x}_{t},\bm{x}_{0})\sim q(\bm{x}_{t},\bm{x}_{0})$ and $\nabla_{\bm{x}_{t}}^{3}\log q_{0t}(\bm{x}_{t}|\bm{x}_{0})=\bm{0}$, we have

We use $\nabla(\cdot)$ to denote the derivative of $\bm{x}_{t}$, namely $\nabla_{\bm{x}_{t}}(\cdot)$. Firstly, according to the second-order score trace given by Lemma F.2, we have

Assume $(\bm{x}_{t},\bm{x}_{0})\sim q(\bm{x}_{t},\bm{x}_{0})$ and $\nabla_{\bm{x}_{t}}^{3}\log q_{0t}(\bm{x}_{t}|\bm{x}_{0})=\bm{0}$, we have

We use $\nabla(\cdot)$ to denote the derivative of $\bm{x}_{t}$, namely $\nabla_{\bm{x}_{t}}(\cdot)$. Firstly, according to the second-order score given by Lemma F.2, we have

Similar to the proof of Lemma F.4, we have

Also similar to the proof of Lemma F.4, we have

For simplicity, we denote $q_{0t}\coloneqq q_{0t}(\bm{x}_{t}|\bm{x}_{0})$, $q_{t0}\coloneqq q_{t0}(\bm{x}_{0}|\bm{x}_{t})$, $q_{t}\coloneqq q_{t}(\bm{x}_{t})$, $q_{0}\coloneqq q_{0}(\bm{x}_{0})$, $\hat{\bm{s}}_{1}\coloneqq\hat{\bm{s}}_{1}(\bm{x}_{t},t)$, $\bm{s}_{2}(\theta)\coloneqq\bm{s}_{2}(\bm{x}_{t},t;\theta)$.

As $\nabla\log q_{0t}=-\frac{\bm{\epsilon}}{\sigma_{t}}$ and $\nabla^{2}\log q_{0t}=-\frac{1}{\sigma_{t}^{2}}\bm{I}$, by rewriting the objective in Eqn. (13), the optimization is equivalent to

For fixed $t$ and $\bm{x}_{t}$, minimizing the inner expectation is a minimum mean square error problem for $\bm{s}_{2}(\theta)$, so the optimal $\theta^{*}$ satisfies

Moreover, by leveraging the property of minimum mean square error, we should notice that the training objective can be rewritten to

which shows that $\|\bm{s}_{2}(\bm{x}_{t},t;\theta)-\bm{s}_{2}(\bm{x}_{t},t;\theta^{*})\|_{F}$ can be viewed as the training error. ∎

F.2 Proof of Corollary 4.2

For simplicity, we denote $q_{0t}\coloneqq q_{0t}(\bm{x}_{t}|\bm{x}_{0})$, $q_{t0}\coloneqq q_{t0}(\bm{x}_{0}|\bm{x}_{t})$, $q_{t}\coloneqq q_{t}(\bm{x}_{t})$, $q_{0}\coloneqq q_{0}(\bm{x}_{0})$, $\hat{\bm{s}}_{1}\coloneqq\hat{\bm{s}}_{1}(\bm{x}_{t},t)$, $\bm{s}_{2}^{\text{trace}}(\theta)\coloneqq\bm{s}_{2}^{\text{trace}}(\bm{x}_{t},t;\theta)$.

As $\nabla\log q_{0t}=-\frac{\bm{\epsilon}}{\sigma_{t}}$ and $\nabla^{2}\log q_{0t}=-\frac{1}{\sigma_{t}^{2}}\bm{I}$, by rewriting the objective in Eqn. (15), the optimization is equivalent to

For fixed $t$ and $\bm{x}_{t}$, minimizing the inner expectation is a minimum mean square error problem for $\bm{s}_{2}^{\text{trace}}(\theta)$, so the optimal $\theta^{*}$ satisfies

By Lemma F.1 and Lemma F.2, similarly we have

Moreover, by leveraging the property of minimum mean square error, we should notice that the training objective can be rewritten to

which shows that $|\bm{s}_{2}^{\text{trace}}(\bm{x}_{t},t;\theta)-\bm{s}_{2}^{\text{trace}}(\bm{x}_{t},t;\theta^{*})|$ can be viewed as the training error. ∎

F.3 Proof of Theorem 4.3

For simplicity, we denote $q_{0t}\coloneqq q_{0t}(\bm{x}_{t}|\bm{x}_{0})$, $q_{t0}\coloneqq q_{t0}(\bm{x}_{0}|\bm{x}_{t})$, $q_{t}\coloneqq q_{t}(\bm{x}_{t})$, $q_{0}\coloneqq q_{0}(\bm{x}_{0})$, $\hat{\bm{s}}_{1}\coloneqq\hat{\bm{s}}_{1}(\bm{x}_{t},t)$, $\hat{\bm{s}}_{2}\coloneqq\hat{\bm{s}}_{2}(\bm{x}_{t},t)$, $\bm{s}_{3}(\theta)\coloneqq\bm{s}_{3}(\bm{x}_{t},t;\theta)$.

As $\nabla\log q_{0t}=-\frac{\bm{\epsilon}}{\sigma_{t}}$ and $\nabla^{2}\log q_{0t}=-\frac{1}{\sigma_{t}^{2}}\bm{I}$, by rewriting the objective in Eqn. (16), the optimization is equivalent to

For fixed $t$ and $\bm{x}_{t}$, minimizing the inner expectation is a minimum mean square error problem for $\bm{s}_{3}(\theta)$, so the optimal $\theta^{*}$ satisfies

As $\nabla^{2}\log q_{0t}=-\frac{1}{\sigma_{t}^{2}}\bm{I}$ is constant w.r.t. $\bm{x}_{0}$, by Lemma F.1, we have

Moreover, by leveraging the property of minimum mean square error, we should notice that the training objective can be rewritten to

which shows that $\|\bm{s}_{3}(\bm{x}_{t},t;\theta)-\bm{s}_{3}(\bm{x}_{t},t;\theta^{*})\|_{2}$ can be viewed as the training error. ∎

In this section, we analyze the DSM objective in Meng et al. (2021a), and show that the objective in Meng et al. (2021a) has the unbounded-error property, which means even if the training error is zero and the first-order score matching error is any small, the second-order score matching error may be arbitrarily large.

(Meng et al., 2021a) proposed an objective for estimating second-order score

For fixed $t$ and $\bm{x}_{t}$, minimizing the inner expectation is a minimum mean square error problem for $\bm{s}_{2}(\theta)$, so the optimal $\theta^{*}$ satisfies

However, below we show that even if $\theta$ achieves its optimal solution $\theta^{*}$ and the first-order score matching error $\|\hat{\bm{s}}_{1}-\nabla\log q_{t}\|$ is small, the second-order score matching error $\bm{s}_{2}(\theta^{*})-\nabla^{2}\log q_{t}$ may still be rather large.

We construct an example to demonstrate this point. Suppose the first-order score matching error $\hat{\bm{s}}_{1}-\nabla\log q_{t}=\delta_{1}\cdot\bm{1}$, where $\delta_{1}>0$ is small, then the first-order score matching error $\|\hat{\bm{s}}_{1}-\nabla\log q_{t}\|_{2}=\delta_{1}\sqrt{d}$, where $d$ is the dimension of $\bm{x}_{t}$. We have

where the second inequality is because $\|A\|_{F}\geq\|\mbox{diag}(A)\|_{2}$, where $\mbox{diag}(A)$ means the diagonal vector of the matrix $A$. Therefore, for any small $\delta_{1}$, the second-order score estimation error may be arbitrarily large.

In practice, Meng et al. (2021a) do not stop gradients for $\hat{\bm{s}}_{1}$, which makes the optimization of the second-order score model $\bm{s}_{2}(\theta)$ affecting the first-order score model $\bm{s}_{1}(\theta)$, and thus cannot theoretically guarantee the convergence. Empirically, we implement the second-order DSM objective in Meng et al. (2021a), and we find that if we stop gradients for $\hat{\bm{s}}_{1}$, the model quickly diverges and cannot work. We argue that this is because of the unbounded error of their method, as shown above.

Instead, our proposed high-order DSM method has the error-bounded property, which shows that $\|\bm{s}_{2}(\theta^{*})-\nabla\log q_{t}\|_{F}=\|\hat{\bm{s}}_{1}-\nabla\log q_{t}\|_{2}^{2}$. So our method does not have the problem mentioned above.

Appendix G Estimated objectives of high-order DSM for high-dimensional data

In this section, we propose the estimated high-order DSM objectives for high-dimensional data.

The second-order DSM objective in Eqn. (13) requires computing the Frobenius norm of the full Jacobian of the score model, i.e. $\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x},t)$, which typically has $\mathcal{O}(d^{2})$ time complexity and is unacceptable for high dimensional real data. Moreover, the high-order DSM objectives in Eqn. (15) and Eqn. (16) include computing the divergence of score network $\nabla_{\bm{x}}\cdot\bm{s}_{\theta}(\bm{x},t)$ i.e. the trace of Jacobian $\tr(\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x},t))$, and it also has $\mathcal{O}(d^{2})$ time complexity. Similar to Grathwohl et al. (2018) and Finlay et al. (2020), the cost can be reduced to $\mathcal{O}(d)$ using Hutchinson’s trace estimator and automatic diffentiation provided by general deep learning frameworks, which needs one-time backpropagation only.

For a $d$-by-$d$ matrix $A$, its trace can be unbiasedly estimated by (Hutchinson, 1989):

Let $A=\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x},t)$, The Jacobian-vector-product $\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x},t)\bm{v}$ can be efficiently computed by using once forward-mode automatic differentiation in JAX, making the evaluating of trace and Frobenius norm approximately the same cost as evaluating $\bm{s}_{\theta}(\bm{x},t)$. And we show the time costs for the high-order DSM training in Appendix. I.4.

By leveraging the unbiased estimator, our final objectives for second-order and third-order DSM are:

where $\bm{s}_{jvp}=\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x},t)\bm{v}$ is the Jacobian-vector-product as mentioned above, $\hat{\bm{s}}_{1}=\texttt{stop\_gradient}(\bm{s}_{\theta})$, $\hat{\bm{s}}_{jvp}=\texttt{stop\_gradient}(\bm{s}_{jvp})$, and $\bm{a}\cdot\bm{b}$ denotes the vector-inner-product for vectors $\bm{a}$ and $\bm{b}$. We compute $\bm{v}^{\top}\nabla_{\bm{x}}\bm{s}_{jvp}$ by using the auto-gradient functions in JAX (which can compute the ”vector-Jacobian-product” by one-time backpropagation”).

In this section, we show that the proposed estimated objectives are actually equivalent to or can upper bound the original objectives in Sec. 5.1 for high-order DSM. Specifically, we have

Firstly, the estimated second-order DSM objective in Eqn. (172) is equivalent to the original objective in Sec. 5.1, because

And the estimated trace of second-order DSM objective in Eqn. (173) can upper bound the original objective in Sec. 5.1, because

And the estimated third-order DSM objective in Eqn. (174) can also upper bound the original objective in Sec. 5.1. To prove that, we firstly propose the following theorem, which presents an equivalent form of the third-order DSM objective in Eqn. (16) for score models and the corresponding derivatives.

Denote the first-order score matching error as $\delta_{1}(\bm{x}_{t},t)\coloneqq\|\hat{\bm{s}}_{1}(\bm{x}_{t},t)-\nabla_{\bm{x}}\log q_{t}(\bm{x}_{t})\|_{2}$ and the second-order score matching errors as $\delta_{2}(\bm{x}_{t},t)\coloneqq\|\nabla_{\bm{x}}\hat{\bm{s}}_{1}(\bm{x}_{t},t)-\nabla_{\bm{x}}^{2}\log q_{t}(\bm{x}_{t})\|_{F}$ and $\delta_{2,\text{tr}}(\bm{x}_{t},t)\coloneqq|\tr(\nabla_{\bm{x}}\hat{\bm{s}}_{1}(\bm{x}_{t},t))-\tr(\nabla_{\bm{x}}^{2}\log q_{t}(\bm{x}_{t}))|$. Then $\forall\bm{x}_{t},\theta$, the score matching error for $\nabla_{\bm{x}}\tr(\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x}_{t},t))$ can be bounded by:

For simplicity, we denote $q_{0t}\coloneqq q_{0t}(\bm{x}_{t}|\bm{x}_{0})$, $q_{t0}\coloneqq q_{t0}(\bm{x}_{0}|\bm{x}_{t})$, $q_{t}\coloneqq q_{t}(\bm{x}_{t})$, $q_{0}\coloneqq q_{0}(\bm{x}_{0})$, $\hat{\bm{s}}_{1}\coloneqq\hat{\bm{s}}_{1}(\bm{x}_{t},t)$, $\hat{\bm{s}}_{jvp}\coloneqq\hat{\bm{s}}_{jvp}(\bm{x}_{t},t)$, $\nabla\tr(\nabla\bm{s}(\theta))\coloneqq\nabla_{\bm{x}}\tr(\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x}_{t},t))$.

As $\nabla\log q_{0t}=-\frac{\bm{\epsilon}}{\sigma_{t}}$ and $\nabla^{2}\log q_{0t}=-\frac{1}{\sigma_{t}^{2}}\bm{I}$, by rewriting the objective in Eqn. (188), the optimization is equivalent to:

For fixed $t$ and $\bm{x}_{t}$, minimizing the inner expectation is a minimum mean square error problem for $\nabla\tr(\nabla\bm{s}(\theta))$, so the optimal $\theta^{*}$ satisfies

Below we show that the objective in Eqn. (16) in Theorem 4.3 is equivalent to the objective in Eqn. (188) in Theorem G.1.

For simplicity, we denote $q_{0t}\coloneqq q_{0t}(\bm{x}_{t}|\bm{x}_{0})$, $q_{t0}\coloneqq q_{t0}(\bm{x}_{0}|\bm{x}_{t})$, $q_{t}\coloneqq q_{t}(\bm{x}_{t})$, $q_{0}\coloneqq q_{0}(\bm{x}_{0})$, $\hat{\bm{s}}_{1}\coloneqq\hat{\bm{s}}_{1}(\bm{x}_{t},t)$, $\hat{\bm{s}}_{jvp}\coloneqq\hat{\bm{s}}_{jvp}(\bm{x}_{t},t)$, $\nabla\tr(\nabla\bm{s}(\theta))\coloneqq\nabla_{\bm{x}}\tr(\nabla_{\bm{x}}\bm{s}_{\theta}(\bm{x}_{t},t))$.

On the one hand, by Eqn. (197), the optimal solution of the objective in Eqn. (188) is:

so by the property of least mean square error, the objective in Eqn. (188) w.r.t. the optimization of $\theta$ is equivalent to

On the other hand, by Eqn. (155), the optimal solution of the objective in Eqn. (16) is also:

so by the property of least mean square error, the objective in Eqn. (16) w.r.t. the optimization of $\theta$ is also equivalent to

Therefore, the two objectives in Eqn. (16) and Eqn. (188) are equivalent w.r.t. $\theta$. ∎

Therefore, by Corollary G.2, we can derive an equivalent formulation of $\mathcal{J}_{\text{DSM}}^{(3)}(\theta)$ by Theorem. G.1:

Appendix H Training algorithm

In this section, we propose our training algorithm for high-dimensional data, based on the high-order DSM objectives in Appendix. G.

Appendix I Experiment details

As we use Monto-Carlo method to unbiasedly estimate the expectations, we empirically find that we need to ensure the mean values of $\mathcal{J}_{\text{DSM}}^{(1)}(\theta),\lambda_{1}\left(\mathcal{J}_{\text{DSM}}^{(2)}(\theta)+\mathcal{J}_{\text{DSM}}^{(2,\text{tr})}(\theta)\right)$ and $\lambda_{2}\mathcal{J}_{\text{DSM}}^{(3)}(\theta)$ are in the same order of magnitude. Therefore, for synthesis data experiments, we choose $\lambda_{1}=0.5$ and $\lambda_{2}=0.1$ for our final objectives. And for CIFAR-10 experiments, we simply choose $\lambda_{1}=\lambda_{2}=1$ with no further tuning.

Specifically, for synthesis data, we choose $\lambda_{1}=0.5$, $\lambda_{2}=0$ as the second-order score matching objective, and $\lambda_{1}=0.5$, $\lambda_{2}=0.1$ as the third-order score matching objective. For CIFAR-10, we simply choose $\lambda_{1}=1$, $\lambda_{2}=0$ as the second-order score matching objective, and $\lambda_{1}=\lambda_{2}=1$ as the third-order score matching objective.

In practice, the implementation of the first-order DSM objective in (Song et al., 2020b, 2021) divides the loss function by the data dimension (i.e. use $\frac{\|\cdot\|_{2}^{2}}{d}$ instead of $\|\cdot\|_{2}^{2}$). We also divide the second-order DSM and the third-order DSM objectives by the data dimension for the image data to balance the magnitudes of the first, second and third-order objectives. Please refer to the implementation in our released code for details.

I.2 1-D mixture of Gaussians

We exactly compute the high-order score matching objectives in Sec. 5.1, and choose the starting time $\epsilon=10^{-5}$.

The density function of the data distribution is

We use the “noise-prediction” type model (Kingma et al., 2021), i.e. we use a neural network $\bm{\epsilon}_{\theta}(\bm{x}_{t},t)$ to model $\sigma_{t}\bm{s}_{\theta}(\bm{x}_{t},t)$. We use the time-embedding in Song et al. (2020b). For the score model, we use a two-layer MLP to encode $t$, and a two-layer MLP to encode the input $\bm{x}_{t}$, then concatenate them together to another two-layer MLP network to output the predicted noise $\bm{\epsilon}_{\theta}(\bm{x}_{t},t)$.

We use Adam (Kingma & Ba, 2014) optimizer with the default settings in JAX. The batch size is $5000$. We train the model for $50$k iterations by one NVIDIA GeForece RTX 2080 Ti GPU card.

As our proposed high-order DSM algorithm has the property of bounded errors, we directly train our model from the default initialized neural network, without any pre-training for the lower-order models.

I.3 2-D checkerboard data

We exactly compute the high-order score matching objectives in Sec. 5.1, and choose the starting time $\epsilon=10^{-5}$.

We use the same neural network and hyperparameters in Appendix. I.2, and we train for $100$k iterations for all experiments by one NVIDIA GeForece RTX 2080 Ti GPU card.

Similar to Appendix. I.2, we directly train our model from the default initialized neural network, without any pre-training for the lower-order models.

I.4 CIFAR-10 experiments

Our code of CIFAR-10 experiments are based on the released code of Song et al. (2020b) and Song et al. (2021). We choose the start time $\epsilon=10^{-5}$ for both training and evaluation.

We train the score model by the proposed high-order DSM objectives both from iteration and from the pre-trained checkpoints in Song et al. (2020b) to a fixed ending iteration, and we empirically find that the model performance by these two training procedure are nearly the same. Therefore, to save time, we simply use the pre-trained checkpoints to further train the second-order and third-order models by a few iteration steps and report the results of the fine-tuning models. Moreover, we find that further train the pre-trained checkpoints by first-order DSM cannot improve the likelihood of ScoreODE, so we simply use the pre-trained checkpoints to evaluate the first-order models.

The model architectures are the same as (Song et al., 2020b). For VE, we use NCSN++ cont. which has 4 residual blocks per resolution. For VE (deep), we use NCSN++ cont. deep which has 8 residual blocks per resolution. Also, our network is the “noise-prediction” type, same as the implementation in Song et al. (2020b).

Training

We follow the same training procedure and default settings for score-based models as (Song et al., 2020b), and set the exponential moving average (EMA) rate to 0.999 as advised. Also as in (Song et al., 2020b), the input images are pre-processed to be normalized to $$ for VESDEs. For all experiments, we set the “n_jitted_steps=1” in JAX code.

For the experiments of the VE model, we use 8 GPU cards of NVIDIA GeForece RTX 2080 Ti. We use the pre-trained checkpoints of $1200$k iterations of Song et al. (2020b), and further train $100$k iterations (for about half a day) and the model quickly converges. We use a batchsize of 128 of the second-order training, and a batchsize of 48 of the third-order training.

For the experiments of the VE(deep) model, we use 8 GPU cards of Tesla P100-SXM2-16GB. We use the pre-trained checkpoints of $600$k iterations of Song et al. (2020b), and further train $100$k iterations (for about half a day) and the model quickly converges. We use a batchsize of 128 of the second-order training, and a batchsize of 48 of the third-order training.

Likelihood and sample quality

We use the uniform dequantization for likelihood evaluation. For likelihood, we report the bpd on the test dataset with 5 times repeating (to reduce the variance of the trace estimator). For sampling, we find the ode sampler often produce low quality images and instead use the PC sampler (Song et al., 2020b) discretized at 1000 time steps to generate 50k samples and report the FIDs on them. We use the released pre-trained checkpoints of the VESDE in Song et al. (2020b) to evaluate the likelihood and sample quality of the first-order score matching models.

Computation time, iteration numbers and memory consumption

We list the computation time and memory consumption of the VE model (shallow model) on 8 NVIDIA GeForce RTX 2080 Ti GPU cards with batch size 48 in Table 3. The n_jitted_steps is set to 1. When training by second-order DSM, we only use score_jvp_fn in the training algorithm and remove grad_div_fn. The computation time is averaged over 10k iterations. We use jax.profiler to trace the GPU memory usage during training, and report the peak total memory on 8 GPUs.

We find that the costs of the second-order DSM training are close to the first-order DSM training, and the third-order DSM training costs less than twice of the first-order training. Nevertheless, our method can scale up to high-dimensional data and improve the likelihood of ScoreODEs.

I.5 ImageNet 32x32 experiments

We adopt the same start time and model architecture for both VE and VE (deep) as CIFAR-10 experiments. Note that the released code of Song et al. (2020b) and Song et al. (2021) provides no pretrained checkpoint of VE type for ImageNet 32x32 dataset, so we use their training of VP type as a reference, and train the first-order VE models from scratch. Specifically, we train the VE baseline for 1200k iterations, and the VE (deep) baseline for 950k iterations.

For the high-order experiments of both VE and VE (deep), we further train 100k iterations, using a batchsize of 128 for the second-order training and a batchsize of 48 for the third-order training. We use 8 GPU cards of NVIDIA GeForece RTX 2080 Ti for VE, and 8 GPU cards of NVIDIA A40 for VE (deep). We report the average bpd on the test dataset with 5 times repeating.

Appendix J Additional results for VE, VP and subVP types

We also train VP and subVP types of ScoreODEs by the proposed high-order DSM method, with the maximum likelihood weighting function for ScoreSDE in (Song et al., 2021) (the weighting functions are detailed in (Song et al., 2021, Table 1)). Note that for the VE type, the likelihood weighting is exactly the DSM weighting used in our experiments.

We firstly show that for the ScoreODE of VP type, even on the simple 1-D mixture-of-Gaussians, the model density trained by the first-order DSM is not good enough and can be improved by the third-order DSM, as shown in Fig. 5.

We then train ScoreODE of VP and subVP types on the CIFAR-10 dataset. We use the pretrained checkpoint by first-order DSM in (Song et al., 2021) (the checkpoint of “Baseline+LW+IS”), and further train 100k iterations. We use the same experiment settings as the VE experiments. We vary the start evaluation time $\epsilon$ and evaluate the model likelihood at each $\epsilon$, as shown in Fig. 6. For the ScoreODE trained by the first-order DSM, the model likelihood is poor when $\epsilon$ is slightly large. Note that Song et al. (2020b, 2021) uses $\epsilon=10^{-3}$ for sampling of VP type, and the corresponding likelihood is poor. Moreover, our proposed method can improve the likelihood for $\epsilon$ larger than a certain value (e.g. our method can improve the likelihood for the VP type with $\epsilon=10^{-3}$). However, for very small $\epsilon$, our method cannot improve the likelihood for VP and subVP. We suspect that it is because of the “unbounded score” problem (Dockhorn et al., 2021). The first-order score function suffers numerical issues for $t$ near and the first-order SM error is so large that it cannot provide useful information for the higher-order SM.

As the data score function $\nabla_{\bm{x}}\log q_{t}(\bm{x}_{t})$ is unknown for CIFAR-10 experiments, we cannot evaluate $\mathcal{J}_{\text{Fisher}}(\theta)$. To show the effectiveness of our method, we evaluate the difference between $\bm{s}_{\theta}(\bm{x}_{t},t)$ and $\nabla_{\bm{x}}\log p_{t}^{\text{ODE}}(\bm{x}_{t})$ in the $\mathcal{J}_{\text{Diff}}(\theta)$. Denote

Moreover, we randomly select a batch of generated samples by the PC sampler (Song et al., 2020b) of the same random seed by the VE model of first-order, second-order and third-order DSM training, as shown in Fig. 8. The samples are very close for human eyes, which shows that after our proposed training method, the score model can still be used for the sample methods of SGMs to generate high-quality samples.