Efficient Learning of Generative Models via Finite-Difference Score Matching

Introduction

Deep generative models have achieved impressive progress on learning data distributions, with either an explicit density function or an implicit generative process . Among explicit models, energy-based models (EBMs) define the probability density as $p_{\theta}(x)=\widetilde{p}_{\theta}(x)/Z_{\theta}$, where $\widetilde{p}_{\theta}(x)$ denotes the unnormalized probability and $Z_{\theta}=\int\widetilde{p}_{\theta}(x)dx$ is the partition function. EBMs allow more flexible architectures with simpler compositionality compared to other explicit generative models , and have better stability and mode coverage in training compared to implicit generative models . Although EBMs are appealing, training them with maximum likelihood estimate (MLE), i.e., minimizing the KL divergence between data and model distributions, is challenging because of the intractable partition function .

Score matching (SM) is an alternative objective that circumvents the intractable partition function by training unnormalized models with the Fisher divergence , which depends on the Hessian trace and (Stein) score function of the log-density function. SM eliminates the dependence of the log-likelihood on $Z_{\theta}$ by taking derivatives w.r.t. $x$, using the fact that $\nabla_{x}\log p_{\theta}(x)=\nabla_{x}\log\widetilde{p}_{\theta}(x)$. Different variants of SM have been proposed, including approximate back-propagation , curvature propagation , denoising score matching (DSM) , a bi-level formulation for latent variable models and nonparametric estimators , but they may suffer from high computational cost, biased parameter estimation, large variance, or complex implementations. Sliced score matching (SSM) alleviates these problems by providing a scalable and unbiased estimator with a simple implementation. However, most of these score matching methods optimize (high-order) derivatives of the density function, e.g., the gradient of a Hessian trace w.r.t. parameters, which are several times more computationally expensive compared to a typical end-to-end propagation, even when using reverse-mode automatic differentiation . These extra computations need to be performed in sequential order and cannot be easily accelerated by parallel computing (as discussed in Appendix B.1). Besides, the induced repetitive usage of the same intermediate results could magnify the stochastic variance and lead to numerical instability .

To improve efficiency and stability, we first observe that existing scalable SM objectives (e.g., DSM and SSM) can be rewritten in terms of (second-order) directional derivatives. We then propose a generic finite-difference (FD) decomposition for any-order directional derivative in Sec. 3, and show an application to SM methods in Sec. 4, eliminating the need for optimizing on higher-order gradients. Specifically, our FD approach only requires independent (unnormalized) likelihood function evaluations, which can be efficiently and synchronously executed in parallel with a simple implementation (detailed in Sec. 3.3). This approach reduces the computational complexity of any $T$-th order directional derivative to $\mathcal{O}(T)$, and improves numerical stability because it involves a shallower computational graph. As we exemplify in Fig. 1, the FD reformulations decompose the inherently sequential high-order gradient computations in SSM (left panel) into simpler, independent routines (right panel). Mathematically, in Sec. 5 we show that even under stochastic optimization , our new FD objectives are asymptotically consistent with their gradient-based counterparts under mild conditions. When the generative models are unnormalized, the intractable partition function can be eliminated by the linear combinations of log-density in the FD-form objectives. In experiments, we demonstrate the speed-up ratios of our FD reformulations with more than $2.5\times$ for SSM and $1.5\times$ for DSM on different generative models and datasets, as well as the comparable performance of the learned models.

Background

As an alternative to KL divergence, score matching (SM) minimizes the Fisher divergence between $p_{\theta}(x)$ and $p_{\textup{data}}(x)$, which is equivalent to

up to a constant and $\text{tr}(\cdot)$ is the matrix trace. Note that the derivatives w.r.t. $x$ eliminate the dependence on the partition function, i.e., $\nabla_{x}\log p_{\theta}(x)=\nabla_{x}\log\widetilde{p}_{\theta}(x)$, making the objective function tractable. However, the calculation of the trace of Hessian matrix is expensive, requiring the number of back-propagations proportional to the data dimension . To circumvent this computational difficulty, two scalable variants of SM have been developed, to which we will apply our methods.

Denoising score matching (DSM). Vincent circumvents the Hessian trace by perturbing $x$ with a noise distribution $p_{\sigma}(\widetilde{x}|x)$ and then estimating the score of the perturbed data distribution $p_{\sigma}(\widetilde{x})=\int p_{\sigma}(\widetilde{x}|x)p_{\textup{data}}(x)dx$. When using Gaussian noise, we obtain the DSM objective as

The model obtained by DSM only matches the true data distribution when the noise scale $\sigma$ is small enough. However, when $\sigma\rightarrow 0$, the variance of DSM could be large or even tend to infinity , requiring grid search or heuristics for choosing $\sigma$ .

Sliced score matching (SSM). Song et al. use random projections to avoid explicitly calculating the Hessian trace, so that the training objective only involves Hessian-vector products as follows:

2 Computational cost of gradient-based SM methods

Although SM methods can bypass the intractable partition function $Z_{\theta}$, they have to optimize an objective function involving higher-order derivatives of the log-likelihood density. Even if reverse mode automatic differentiation is used , existing SM methods like DSM and SSM can be computationally expensive during training when calculating the Hessian-vector products.

Complexity of the Hessian-vector products. Let $\mathcal{L}$ be any loss function, and let $\textup{Cal}(\nabla\mathcal{L})$ and $\textup{Mem}(\nabla\mathcal{L})$ denote the time and memory required to compute $\nabla\mathcal{L}$, respectively. Then if the reverse mode of automatic differentiation is used, the Hessian-vector product can be computed with up to five times more time and two times more memory compared to $\nabla\mathcal{L}$, i.e., $5\!\times\textup{Cal}(\nabla\mathcal{L})$ time and $2\!\times\textup{Mem}(\nabla\mathcal{L})$ memory . When we instantiate $\mathcal{L}=\log p_{\theta}(x)$, we can derive that the computations of optimizing DSM and SSM are separately dominated by the sequential operations of $\nabla_{\theta}(\|\nabla_{x}\mathcal{L}\|)$ and $\nabla_{\theta}(v^{\top}\nabla_{x}(v^{\top}\nabla_{x}\mathcal{L}))$, as illustrated in Fig. 1 for SSM. The operations of $\nabla_{\theta}$ and $\nabla_{x}$ require comparable computing resources, so we can conclude that compared to directly optimizing the log-likelihood, DSM requires up to $5\times$ computing time and $2\times$ memory, while SSM requires up to $25\times$ computing time and $4\times$ memory . For higher-order derivatives, we empirically observe that the computing time and memory usage grow exponentially w.r.t. the order of derivatives, i.e., the times of executing the operator $v^{\top}\nabla$, as detailed in Sec. 3.3.

Approximating directional derivatives via finite difference

In this section, we first rewrite the most expensive terms in the SM objectives in terms of directional derivatives, then we provide generic and efficient formulas to approximate any $T$-th order directional derivative using finite difference (FD). The proposed FD approximations decompose the sequential and dependent computations of high-order derivatives into independent and parallelizable computing routines, reducing the computational complexity to $\mathcal{O}(T)$ and improving numerical stability.

Note that the objectives of SM, DSM, and SSM described in Sec. 2.1 can all be abstracted in terms of $v^{\top}\nabla_{x}\mathcal{L}_{\theta}(x)$ and $v^{\top}\nabla_{x}^{2}\mathcal{L}_{\theta}(x)v$. Specifically, as to SM or DSM, $v$ is the basis vector $\bm{e}_{i}$ along the $i$-th coordinate to constitute the squared norm term $\|\nabla_{x}\mathcal{L}_{\theta}(x)\|_{2}^{2}=\sum_{i=1}^{d}(\bm{e}_{i}^{\top}\nabla_{x}\mathcal{L}_{\theta}(x))^{2}$ or the Hessian trace term $\textup{tr}(\nabla_{x}^{2}\mathcal{L}_{\theta}(x))=\sum_{i=1}^{d}\bm{e}_{i}^{\top}\nabla_{x}^{2}\mathcal{L}_{\theta}(x)\bm{e}_{i}$. As to SSM, $v$ denotes the random direction.

We regard the gradient operator $\nabla_{x}$ as a $d$-dimensional vector $\nabla_{x}=(\frac{\partial}{\partial x_{1}},\cdots,\frac{\partial}{\partial x_{d}})$, and $v^{\top}\nabla_{x}$ is an operator that first executes $\nabla_{x}$ and then projects onto the vector $v$. For notation simplicity, we denote $\|v\|_{2}=\epsilon$ and rewrite the above terms as (higher-order) directional derivatives as follows:

Here $\frac{\partial}{\partial v}$ is the directional derivative along $v$, and $\left(v^{\top}\nabla_{x}\right)^{2}$ means executing $v^{\top}\nabla_{x}$ twice.

2 FD decomposition for directional derivatives

We propose to adopt the FD approach, a popular tool in numerical analysis to approximate differential operations , to efficiently estimate the terms in Eq. (4). Taking the first-order case as an example, the key idea is that we can approximate $\frac{\partial}{\partial v}\mathcal{L}_{\theta}(x)\!=\!\frac{1}{2\epsilon}(\mathcal{L}_{\theta}(x\!+\!v)\!-\!\mathcal{L}_{\theta}(x\!-\!v))\!+\!o(\epsilon)$, where the right-hand side does not involve derivatives, just function evaluations. In FD, $\|v\|_{2}=\epsilon$ is assumed to be a small value, but this does not affect the optimization of SM objectives. For instance, the SSM objective in Eq. (3) can be adaptively rescaled by $C_{v}$ (generally explained in Appendix B.2).

In general, to estimate the $T$-th order directional derivative of $\mathcal{L}_{\theta}$, which is assumed to be $T$ times differentiable, we first apply the multivariate Taylor’s expansion with Peano’s remainder as

(Existence of $o(1)$ estimator) If $\mathcal{L}_{\theta}(x)$ is $T$-times-differentiable at $x$, then given any set of $T+1$ different real values $\{\gamma_{i}\}_{i=1}^{T+1}$, there exist corresponding coefficients $\{\beta_{i}\}_{i=1}^{T+1}$, such that

Lemma 1 states that it is possible to approximate the $T$-th order directional derivative as to an $o(1)$ error with $T\!+\!1$ function evaluations. In fact, as long as $\mathcal{L}_{\theta}(x)$ is $(T\!+\!1)$-times-differentiable at $x$, we can construct a special kind of linear combination of $T\!+\!1$ function evaluations to reduce the approximation error to $o(\epsilon)$, as stated below:

It is easy to generalize Theorem 1 to achieve approximation error $o(\epsilon^{N})$ for any $N\geq 1$ with $T\!+\!N$ function evaluations, and we can show that the error rate $o(\epsilon)$ is optimal when evaluating $T\!+\!1$ functions. So far we have proposed generic formulas for the FD decomposition of any-order directional derivative. As to the application to SM objectives (detailed in Sec. 4), we can instantiate the decomposition in Theorem 1 with $K=1$, $\alpha_{1}=1$, and solve for $\beta_{1}=1$, which leads to

In addition to generative modeling, the decomposition in Theorem 1 can potentially be used in other settings involving higher-order derivatives, e.g., extracting local patterns with high-order directional derivatives , training GANs with gradient penalty , or optimizing the Fisher information . We leave these interesting explorations to future work.

Remark. When $\mathcal{L}_{\theta}(x)$ is modeled by a neural network, we can employ the average pooling layer and the non-linear activation of, e.g., Softplus to have an infinitely differentiable model to meet the condition in Theorem 1. Note that Theorem 1 promises a point-wise approximation error $o(\epsilon)$. To validate the error rate under expectation for training objectives, we only need to assume that $p_{\textup{data}}(x)$ and $\mathcal{L}_{\theta}(x)$ satisfy mild regularity conditions beyond the one in Theorem 1, which can be easily met in practice, as detailed in Appendix B.3. Conceptually, these mild regularity conditions enable us to substitute the Peano’s remainders with Lagrange’s ones. Moreover, this substitution results in a better approximation error of $\mathcal{O}(\epsilon^{2})$ for our FD decomposition, while we still use $o(\epsilon)$ for convenience.

3 Computational efficiency of the FD decomposition

Theorem 1 provides a generic approach to approximate any $T$-th order directional derivative by decomposing the sequential and dependent order-by-order computations into independent function evaluations. This decomposition reduces the computational complexity to $\mathcal{O}(T)$, while the complexity of explicitly computing high-order derivatives usually grows exponentially w.r.t. $T$ , as we verify in Fig. 2. Furthermore, due to the mutual independence among the function terms $\mathcal{L}_{\theta}(x+\gamma_{i}v)$, they can be efficiently and synchronously executed in parallel via simple implementation (pseudo code is in Appendix C.1). Since this parallelization acts on the level of operations for each data point $x$, it is compatible with data or model parallelism to further accelerate the calculations.

To empirically demonstrate the computational efficiency of our FD decomposition, we report the computing time and memory usage in Fig. 2 for calculating the $T$-th order directional derivative, i.e., $\frac{\partial^{T}}{\partial v^{T}}$ or $(v^{\top}\nabla_{x})^{T}$, either exactly or by the FD decomposition. The function $\mathcal{L}_{\theta}(x)$ is the log-density modeled by a deep EBM and trained on MNIST, while we use PyTorch for automatic differentiation. As shown in the results, our FD decomposition significantly promotes efficiency in respect of both speed and memory usage, while the empirical approximation error rates are kept within $1\%$. When we parallelize the FD decomposition, the computing time is almost a constant w.r.t. the order $T$, as long as there is enough GPU memory. In our experiments in Sec. 6, the computational efficiency is additionally validated on the FD-reformulated SM methods.

Application to score matching methods

Finite-difference SSM. For SSM, the scale factor is $C_{v}=\epsilon^{2}$ in Eq. (3). By instantiating $\mathcal{L}_{\theta}=\log p_{\theta}(x)$ in Eq. (8), we propose the finite-difference SSM (FD-SSM) objective as

In Fig. 1, we intuitively illustrate the computational graph to better highlight the difference between the gradient-based objectives and their FD reformations, taking SSM as an example.

Finite-difference DSM. To construct the FD instantiation for DSM, we first cast the original objective in Eq. (2) into sliced Wasserstein distance with random projection $v$ (detailed in Appendix B.4). Then we can propose the finite-difference DSM (FD-DSM) objective as

It is easy to verify that \mathcal{J}_{\text{FD-DSM}}(\theta)=\mathcal{J}_{\text{DSM}}(\theta)+{\color[rgb]{0,0,1}o(\epsilon)}, and we can generalize FD-DSM to the cases with other noise distributions of $p_{\sigma}(\widetilde{x}|x)$ using similar instantiations of Eq. (8).

If $s_{\theta}(x)$ is (element-wisely) twice-differentiable at $x$, we have the expansion that $s_{\theta}(x+v)+s_{\theta}(x-v)=2s_{\theta}(x)+o(\epsilon)$ and $s_{\theta}(x+v)-s_{\theta}(x-v)=2\nabla_{x}s_{\theta}(x)v+o(\epsilon^{2})$. Then we can construct the finite-difference SSMVR (FD-SSMVR) for the score-based models as

We can verify that \mathcal{J}_{\text{FD-SSMVR}}(\theta)=\mathcal{J}_{\text{SSMVR}}(\theta)+{\color[rgb]{0,0,1}o(\epsilon)}. Compared to the FD-SSM objective on the likelihood-based models, we only use two counterparts $s_{\theta}(x+v)$ and $s_{\theta}(x-v)$ in this instantiation.

Consistency under stochastic optimization

In practice, we usually apply mini-batch stochastic gradient descent (SGD) to update the model parameters $\theta$. Thus beyond the expected $o(\epsilon)$ approximation error derived in Sec. 4, it is critical to formally verify the consistency between the FD-form objectives and their gradient-based counterparts under stochastic optimization. To this end, we establish a uniform convergence theorem for FD-SSM as an example, while similar proofs can be applied to other FD instantiations as detailed in Appendix B.5. A key insight is to show that the directions of $\nabla_{\theta}\mathcal{J}_{\textup{FD-SSM}}(\theta)$ and $\nabla_{\theta}\mathcal{J}_{\textup{SSM}}(\theta)$ are sufficiently aligned under SGD, as stated in Lemma 2:

(Consistency under SGD) Optimizing $\nabla_{\theta}\mathcal{J}_{\textup{FD-SSM}}(\theta)$ with stochastic gradient descent, then the model parameters $\theta$ will converge to the stationary point of $\mathcal{J}_{\textup{SSM}}(\theta)$ under the conditions including: (i) the assumptions for general stochastic optimization in Bottou et al. hold; (ii) the differentiability assumptions in Lemma 2 hold; (iii) $\epsilon$ decays to zero during training.

In the proof, we further show that the conditions (i) and (ii) largely overlap, and these assumptions are satisfied by the models described in the remark of Sec. 3.2. As to the condition (iii), we observe that in practice it is enough to set $\epsilon$ be a small constant during training, as shown in our experiments.

Experiments

In this section, we experiment on a diverse set of generative models, following the default settings in previous work .Our code is provided in https://github.com/taufikxu/FD-ScoreMatching. It is worth clarifying that we use the same number of training iterations for our FD methods as their gradient-based counterparts, while we report the time per iteration to exclude the compiling time. More implementation and definition details are in Appendix C.2.

Deep EBMs utilize the capacity of neural networks to define unnormalized models. The backbone we use is an 18-layer ResNet following Li et al. . We validate our methods on six datasets including MNIST , Fashion-MNIST , CelebA , CIFAR-10 , SVHN , and ImageNet . For CelebA and ImageNet, we adopt the officially cropped images and respectively resize to $32\times 32$ and $128\times 128$. The quantitative results on MNIST are given in Table 6.1. As shown, our FD formulations result in $2.9\times$ and $1.7\times$ speedup compared to the gradient-based SSM and DSM, respectively, with consistent SM losses. We simply set $\epsilon=0.1$ to be a constant during training, since we find that the performance of our FD reformulations is insensitive to a wide value range of $\epsilon$. In Fig. 4 (a) and (b), we provide the loss curve of DSM / FD-DSM and SSM / FD-SSM w.r.t. time. As seen, FD-DSM can achieve the best model (lowest SM loss) faster, but eventually converges to higher loss compared to DSM. In contrast, when applying FD on SSM-based methods, the improvements are much more significant. This indicates that the random projection trick required by the FD formula is its main downside, which may outweigh the gain on efficiency for low-order computations.

As an additional evaluation of the learned model’s performance, we consider two tasks using deep EBMs: the first one is out-of-distribution detection, where we follow previous work to use typicality as the detection metric (details in Appendix C.3), and report the AUC scores and the training time per iteration in Table 6.4; the second one is image generation, where we apply annealed Langevin dynamics for inference and show the generated samples in the left of Fig. 3.

2 Flow-based generative models

In addition to the unnormalized density estimators, SM methods can also be applied to flow-based models, whose log-likelihood functions are tractable and can be directly trained with MLE. Following Song et al. , we adopt the NICE model and train it by minimizing the Fisher divergence using different approaches including approximate back-propagation (Approx BP) and curvature propagation (CP) . As in Table 6.1, FD-SSM achieves consistent results compared to SSM, while the training time is nearly comparable with the direct MLE, due to parallelization. The results are averaged over 5 runs except the SM based methods which are averaged over 10 runs. Howev the variance is still large. We hypothesis that it is because the numerical stability of the baseline methods are relatively poor. In contrast, the variance of FD-SSM on the SM loss is much smaller, which shows better numerical stability of the shallower computational graphs induced by the FD decomposition.

3 Latent variable models with implicit encoders

SM methods can be also used in score estimation . One particular application is on VAE / WAE with implicit encoders, where the gradient of the entropy term in the ELBO w.r.t. model parameters can be estimated (more details can be found in Song et al. and Shi et al. ). We follow Song et al. to evaluate VAE / WAE on both the MNIST and CelebA datasets using both SSMVR and FD-SSMVR. We report the results in Table 6. The reported training time only consists of the score estimation part, i.e., training the score model. As expected, the FD reformulation can improve computational efficiency without sacrificing the performance. The discussions concerned with other applications on the latent variable models can be found in Appendix B.6.

4 Score-based generative models

The noise conditional score network (NCSN) trains a single score network $s_{\theta}(x,\sigma)$ to estimate the scores corresponding to all noise levels of $\sigma$. The noise level $\{\sigma_{i}\}_{i\in}$ is a geometric sequence with $\sigma_{1}=1$ and $\sigma_{10}=0.01$. When using the annealed Langevin dynamics for image generation, the number of iterations under each noise level is $100$ with a uniform noise as the initial sample. As to the training approach of NCSN, Song and Ermon mainly use DSM to pursue state-of-the-art performance, while we use SSMVR to demonstrate the efficiency of our FD reformulation. We train the models on the CIFAR-10 dataset with the batch size of $128$ and compute the FID scores on $50,000$ generated samples. We report the results in Table 6.4 and provide the generated samples in the right panel of Fig. 3. We also provide a curve in Fig. 4 (c) showing the FID scores (on 1,000 samples) during training. As seen, our FD methods can effectively learn different generative models.

Related work

As to the more implicit connections to FD, the minimum probability flow (MPF) is a method for parameter estimation in probabilistic models. It is demonstrated that MPF can be connected to SM by a FD reformulation, where we provide a concise derivation in Appendix B.7. The noise-contrastive estimation (NCE) train the unnormalized models by comparing the model distribution $p_{\theta}(x)$ with a noise distribution $p_{n}(x)$. It is proven that when we choose $p_{n}(x)=p_{\text{data}}(x+v)$ with a small vector $v$, i.e., $\|v\|=\epsilon$, the NCE objective can be equivalent to a FD approximation for the SSM objective as to an $o(1)$ approximation error rate after scaling . In contrast, our FD-SSM method can achieve $o(\epsilon)$ approximation error with the same computational cost as NCE.

Conclusion

We propose to reformulate existing gradient-based SM methods using finite difference (FD), and theoretically and empirically demonstrate the consistency and computational efficiency of the FD-based training objectives. In addition to generative modeling, our generic FD decomposition can potentially be used in other applications involving higher-order derivatives. However, the price paid for this significant efficiency is that we need to work on the projected function in a certain direction, e.g., in DSM we need to first convert it into the slice Wasserstein distance and then apply the FD reformulation. This raises a trade-off between efficiency and variance in some cases.

Broader Impact

This work proposes an efficient way to learn generative models and does not have a direct impact on society. However, by reducing the computation required for training unnormalized models, it may facilitate large-scale applications of, e.g., EBMs to real-world problems, which could have both positive (e.g., anomaly detection and denoising) and negative (e.g., deepfakes) consequences.

Acknowledgements

This work was supported by the National Key Research and Development Program of China (No.2017YFA0700904), NSFC Projects (Nos. 61620106010, 62076145, U19B2034, U1811461), Beijing Academy of Artificial Intelligence (BAAI), Tsinghua-Huawei Joint Research Program, a grant from Tsinghua Institute for Guo Qiang, Tiangong Institute for Intelligent Computing, and the NVIDIA NVAIL Program with GPU/DGX Acceleration. C. Li was supported by the Chinese postdoctoral innovative talent support program and Shuimu Tsinghua Scholar.

References

Appendix A Proofs

In this section we provide proofs for the conclusions in the main text.

If $\mathcal{L}_{\theta}(x)$ is $T$-times-differentiable at $x$, then according to the general form of multivariate Taylor’s theorem , there is

In order to extract the $T$-th component $G_{\theta}^{t}(x,v,\epsilon)$, we arbitrarily select a set of $T+1$ different real values as $\{\gamma_{i}\}_{i\in[T+1]}$, and denote the induced Vandermonde matrix $V$ as

where $\bm{\beta}$ is the vector of coefficients. The determinant of the Vandermonde matrix $V$ is $\det(V)=\prod_{i<j}(\gamma_{j}-\gamma_{i})\neq 0$ since the values $\gamma_{i}$ are distinct. We consider the linear combination

To eliminate the term $G_{\theta}^{t}(x,v,\epsilon)$ for any $t<T$ and keep the $T$-th order term $G_{\theta}^{T}(x,v,\epsilon)$, we just need to set the coefficient vector $\bm{\beta}$ be solution of $V^{\top}\bm{\beta}=\bm{e_{(T+1)}}$, where $\bm{e_{(T+1)}}$ is the one-hot vector of the $(T\!+\!1)$-th element. Then we have

A.2 Proof of Theorem 1

where the second equation holds because $\left(1+(-1)^{t}\right)=0$ for any odd value of $t$. Note that there is $G_{\theta}^{0}(x,v,\epsilon)=\mathcal{L}_{\theta}(x)$, thus in order to eliminate the zero-order term, we let

and then we can rewrite Eq. (LABEL:even) as

Now in Eq. (LABEL:even_2) we only need to eliminate the term $G_{\theta}^{2t+2}(x,v,\epsilon)$ for $t<K-1$ and keep the term $G_{\theta}^{2K}(x,v,\epsilon)$, i.e., the $T$-th order term. we define the Vandermonde matrix $V$ generated by $\{\alpha_{1}^{2},\cdots,\alpha_{K}^{2}\}$ as

It is easy to know that $V$ is non-singular as long as $\alpha_{k}$ are positive and different. Then if $\bm{\beta}$ is the solution of $V^{\top}\bm{\beta}=\bm{e}_{K}$ is the one-hot vector of the $K$-th element. Then we have

Similarly when $T=2K-1$ is an odd number, we can construct the linear combination

where the second equation holds because $\left(1-(-1)^{t}\right)=0$ for any even value of $t$. Now we only need to eliminate the term $G_{\theta}^{2t+1}(x,v,\epsilon)$ for $t<K-1$ and keep the term $G_{\theta}^{2K-1}(x,v,\epsilon)$, i.e., the $T$-th order term. Then if we still let $\bm{\beta}$ be the solution of $V_{\text{even}}^{\top}\bm{\beta}=\bm{e}_{K}$, we will have

A.3 Proof of Lemma 2

where $|\bm{\alpha}|=\alpha_{1}+\cdots+\alpha_{d}$, $\bm{\alpha}!=\alpha_{1}!\cdots\alpha_{d}!$, $v^{\bm{\alpha}}=v_{1}^{\alpha_{1}}\cdots v_{d}^{\alpha_{d}}$, and

So the remainder term in Eq. (24) has a upper bound as

where similar results also hold for $x-v$ and we represent the corresponding upper bound as $U_{\omega}^{-}$. Then we further have

Similar for the expansion of $\log p_{\theta}(x+v)$, the remainder is

and we can obtain the uniform upper bound on the compact set $\overline{B}_{\epsilon_{0}}$ as

We denote the bound for $R_{\bm{\alpha}}(x-v)$ as $U^{-}$ and further have

We denote $\Delta R_{\bm{\alpha}}=R_{\bm{\alpha}}(x+v)-R_{\bm{\alpha}}(x-v)$ and $\Delta R_{\bm{\alpha}}^{\omega,+}=R_{\bm{\alpha}}^{\omega}(x+v)+R_{\bm{\alpha}}^{\omega}(x-v)$ and $\Delta R_{\bm{\alpha}}^{\omega,-}=R_{\bm{\alpha}}^{\omega}(x+v)-R_{\bm{\alpha}}^{\omega}(x-v)$ for notation compactness. Thus for $\forall(x,\theta)\in B$ and $\|v\|_{2}=\epsilon$, we obtain the partial derivative of $\mathcal{J}_{\text{FD-SSM}}(x,v;\theta)$ as

Note that the first term in the above equals to $\frac{\partial}{\partial\omega}\mathcal{J}_{\text{SSM}}(x,v;\theta)$. Due to the continuity of the norm functions $\|\nabla_{x}\log p_{\theta}(x)\|_{2}$ and $\|\nabla_{x}\frac{\partial}{\partial\omega}\log p_{\theta}(x)\|_{2}$ on the compact set $\overline{B}_{\epsilon_{0}}$, we denote their upper bound as $G$ and $G_{\omega}$, respectively. Then we have $|v^{\top}\nabla_{x}\log p_{\theta}(x)|\leq\epsilon G$ and $\frac{\partial}{\partial\omega}v^{\top}\nabla_{x}\log p_{\theta}(x)=v^{\top}\nabla_{x}\frac{\partial}{\partial\omega}\log p_{\theta}(x)\leq\epsilon G_{\omega}$. Now we can derive the bound between the partial derivatives of FD-SSM and SSM as

where we denote $\Delta U=U^{+}+U^{-}$ and $\Delta U_{\omega}=U_{\omega}^{+}+U_{\omega}^{-}$. By setting $\epsilon_{0}<\frac{1}{d}$, we can omit the condition $\epsilon<\frac{1}{d}$ since $\epsilon<\epsilon_{0}=\min(\epsilon_{0},\frac{1}{d})$. Note that the condition $\epsilon<\frac{1}{d}$ can be generalize to, e.g., $\epsilon<1$ without changing our conclusions. Then it is easy to show that

where $\min_{(x,\theta)\in B,\|v\|_{2}<\epsilon_{0}}\left\|\nabla_{\theta}\mathcal{J}_{\text{SSM}}(x,v;\theta)\right\|_{2}$ must exist and larger than due to the continuity of $\nabla_{\theta}\mathcal{J}_{\text{SSM}}(x,v;\theta)$ and the condition that $\|\nabla_{\theta}\mathcal{J}_{\text{SSM}}(x,v;\theta)\|_{2}>0$ on the compact set. So we only need to choose $\xi$ as

When we choose $\|v\|_{2}=\epsilon<\min(\epsilon_{0},\xi)$, we can guarantee the angle between $\nabla_{\theta}\mathcal{J}_{\text{FD-SSM}}(x,v;\theta)$ and $\nabla_{\theta}\mathcal{J}_{\text{SSM}}(x,v;\theta)$ to be uniformly less than $\eta$ on $B$. ∎

A.4 Proof of Theorem 2

We consider in the compact set $\overline{B}_{\epsilon_{0}}$ defined in Lemma 2. The assumptions for general stochastic optimization include:

(i) The condition of Corollary 4.12 in Bottou et al. : $\mathcal{J}_{\text{FD-SSM}}(\theta)$ is twice-differentiable with $\theta$;

(ii) The Assumption 4.1 in Bottou et al. : the gradient $\nabla_{\theta}\mathcal{J}_{\text{FD-SSM}}(\theta)$ is Lipschitz;

(iii) The Assumption 4.3 in Bottou et al. : the first and second moments of the stochastic gradients are bounded by the expected gradients;

(iv) The stochastic step size $\alpha_{k}$ satisfies the diminishing condition in Bottou et al. : $\sum_{k=1}^{\infty}\alpha_{k}=\infty$, $\sum_{k=1}^{\infty}\alpha_{k}^{2}<\infty$;

(v) The condition of Lemma 2 holds in each step $k$ of stochastic gradient update.

Note that the condition (v) only holds in the compact set $\overline{B}_{\epsilon_{0}}$, but we can choose it to be large enough to contain $(x,\theta_{k}),x\sim p(x)$, as well as containing the neighborhood of stationary points of $\nabla_{\theta}\mathcal{J}_{\text{SSM}}(\theta)$. These can be achieved by setting $\epsilon\rightarrow 0$. Thus we have

This means that stochastically optimizing the FD-SSM objective can make the parameters $\theta$ converge to the stationary point of the SSM objective when $\epsilon\rightarrow 0$. ∎

Appendix B Extended conclusions

In this section we provide extended and supplementary conclusions for the main text.

For the dependent operations like those in the gradient-based SM methods, it is possible to execute them on different devices via asynchronous parallelism . However, this asynchronous parallelization needs to perform across different data batches, requires complex design on the synchronization mechanism, and could introduce extra bias when updating the model parameters. These difficulties usually outweigh the gain from paralleling the operations in the gradient-based SM methods. In contrast, for our FD-based SM methods, the decomposed independent operations can be easily executed in a synchronous manner, which is further compatible with data or model parallelism.

B.2 Scaling the projection vector in training objectives

Below we explain why the scale of the random projection $v$ will not affect the training of SM objectives. For the original SM objective, we have

When we scale the basis vector $\bm{e}_{i}$ with a small value ${\epsilon^{\prime}}$, i.e., $\bm{e}_{i}\rightarrow{\epsilon^{\prime}}\bm{e}_{i}$, we have

Thus we can simply divide the objective by ${\epsilon^{\prime}}^{2}$ to recover the original SM objective $\mathcal{J}_{\text{SM}}(\theta)$. Similarly, for the DSM objective we have

When we scale the basis vector $\bm{e}_{i}$ with a small value ${\epsilon^{\prime}}$, i.e., $\bm{e}_{i}\rightarrow{\epsilon^{\prime}}\bm{e}_{i}$, we also have

Thus we can divide by ${\epsilon^{\prime}}^{2}$ to recover the DSM objective $\mathcal{J}_{\text{DSM}}(\theta)$. Finally as to SSM, we have

When we scale the random projection $v$ with a small value ${\epsilon^{\prime}}$, i.e., $v\rightarrow{\epsilon^{\prime}}v$, we should not that the adaptive factor $C_{v}$ will also be scaled to ${\epsilon^{\prime}}^{2}C_{v}$, then we can derive

This indicates that the SSM objective is already invariant to the scaling of $v$. It is trivial to also divide similar factors as $C_{v}$ in SM and DSM to result in similarly invariant objectives.

B.3 Mild regularity conditions for the FD-based SM methods

Remark. Note that the condition (iv) holds when we apply, e.g., average pooling layers and Softplus activation in the neural network models, while the condition (v) always holds as long as the support set of $p_{\textup{data}}(x)$ is bounded, e.g., for RGB-based image tasks there is $x\in^{d}$.

B.4 DSM under sliced Wasserstein distance

In this case, there is $\frac{v^{\top}(\widetilde{x}-x)}{\sigma^{2}}=\mathcal{O}(\epsilon)$ with high probability, thus we can approximate $v^{\top}\nabla_{\widetilde{x}}\log p_{\theta}(\widetilde{x})$ according to our FD decomposition.

B.5 Consistency between DSM and FD-DSM

Proof. Following the routines and notations in the proof of Lemma 2, we investigate the gradient $\nabla_{\theta}\mathcal{J}_{\text{FD-DSM}}(x,\widetilde{x},v;\theta)$, whose elements consist of $\frac{\partial}{\partial\omega}\mathcal{J}_{\text{FD-DSM}}(x,\widetilde{x},v;\theta)$ for $\omega\in\theta$. When $\log p_{\theta}(\widetilde{x})$ is three-times-differentiable in $\overline{B}_{\epsilon_{0}}$, we can obtain

Due to the continuity of $R_{\bm{\alpha}}^{\omega}(\widetilde{x}+v)$ and $R_{\bm{\alpha}}^{\omega}(\widetilde{x}-v)$ on the compact set $\overline{B}_{\epsilon_{0}}$, they have the uniform absolute upper bounds $U^{+}_{\omega}$ and $U^{-}_{\omega}$, respectively. Similarly, we have

where the first term equals to $\frac{\partial}{\partial\omega}\mathcal{J}_{\text{DSM}}(x,\widetilde{x},v;\theta)$. Due to the continuity of the norm functions $\|\nabla_{\widetilde{x}}\log p_{\theta}(\widetilde{x})\|_{2}$ and $\|\nabla_{\widetilde{x}}\frac{\partial}{\partial\omega}\log p_{\theta}(\widetilde{x})\|_{2}$ on the compact set $\overline{B}_{\epsilon_{0}}$, we denote their upper bound as $G$ and $G_{\omega}$, respectively. Then we have $|v^{\top}\nabla_{\widetilde{x}}\log p_{\theta}(\widetilde{x})|\leq\epsilon G$ and $\frac{\partial}{\partial\omega}v^{\top}\nabla_{\widetilde{x}}\log p_{\theta}(\widetilde{x})=v^{\top}\nabla_{\widetilde{x}}\frac{\partial}{\partial\omega}\log p_{\theta}(\widetilde{x})\leq\epsilon G_{\omega}$. Besides, since $\widetilde{x}$ and $x$ both come from bounded sets, we have an upper bound of $v^{\top}(\widetilde{x}-x)\leq\epsilon\sigma^{2}G_{x}$. Now we can derive the bound between the partial derivatives of FD-DSM and DSM as

where we denote $\Delta U=U^{+}+U^{-}$ and $\Delta U_{\omega}=U_{\omega}^{+}+U_{\omega}^{-}$. By setting $\epsilon_{0}<\frac{1}{d}$, we can omit the condition $\epsilon<\frac{1}{d}$ since $\epsilon<\epsilon_{0}=\min(\epsilon_{0},\frac{1}{d})$. Note that the condition $\epsilon<\frac{1}{d}$ can be generalize to, e.g., $\epsilon<1$ without changing our conclusions. Then it is easy to show that

where $\min_{(\widetilde{x},\theta)\in B,\|v\|_{2}<\epsilon_{0}}\left\|\nabla_{\theta}\mathcal{J}_{\text{DSM}}(x,\widetilde{x},v;\theta)\right\|_{2}$ must exist and larger than due to the continuity of $\nabla_{\theta}\mathcal{J}_{\text{DSM}}(x,\widetilde{x},v;\theta)$ and the condition that $\|\nabla_{\theta}\mathcal{J}_{\text{DSM}}(x,\widetilde{x},v;\theta)\|_{2}>0$ on the compact set. So we only need to choose $\xi$ as

When we choose $\|v\|_{2}=\epsilon<\min(\epsilon_{0},\xi)$, we can guarantee the angle between $\nabla_{\theta}\mathcal{J}_{\text{FD-DSM}}(x,\widetilde{x},v;\theta)$ and $\nabla_{\theta}\mathcal{J}_{\text{DSM}}(x,\widetilde{x},v;\theta)$ to be uniformly less than $\eta$ on $B$. ∎

B.6 Application on the latent variable models

For the latent variable models (LVMs), the log-likelihood is usually intractable. Unlike EBMs, this intractability cannot be easily eliminated by taking gradients. Recently, the proposed SUMO can provide an unbiased estimator for the intractable $\log p_{\theta}(x)$, which is defined as

Now we can derive an upper bound for our FD reformulated objectives exploiting SUMO. To see how to achieve this, we can first derive a tractable lower bound for the first-order squared term as

as well as a tractable unbiased estimator for the second-order term as

where we adjust the order of expectations to indicate the operation sequence in implementation. According to Eq. (52) and Eq. (53), we can construct upper bounds for our FD-SSM and FD-DSM objectives, and then train the LVMs via minimizing the induced upper bounds. In comparison, when we directly estimate the gradient-based terms $v^{\top}\nabla_{x}\log p_{\theta}(x)$ and $v^{\top}\nabla_{x}^{2}\log p_{\theta}(x)v$, we need to take derivatives on the SUMO estimator, which requires technical derivations .

B.7 Connection to MPF

We can provide a naive FD reformulation for the SSM objective as

Minimum probability flow (MPF) can fit probabilistic model parameters via establishing a deterministic dynamics. For a continues state space, the MPF objective is

where $V_{\epsilon}$ denotes the volume of $d$-dimensional hypersphere of radius $\epsilon$. Let $\Delta_{\theta}(x,v)=\log p_{\theta}(x+v)-\log p_{\theta}(x)$, then we can expand the exponential function around zero as

where the second equation holds because $\Delta_{\theta}(x,v)=\Theta(\epsilon)$. In this case, after removing the offset and scaling factor, the objective of MPF is directly equivalent to $\mathcal{R}(\theta)$ as to an $o(1)$ difference.

Appendix C Implementation details

In this section, we provide a pseudo code for the implementation of FD formulation for both SSM and DSM. Then we provide the specific details in our experiments.

C.2 Implementation details and definitions

Deep EBM directly defines the energy function with unnormalized models using a feed forward NN $f(\cdot)$ and the probability is defined as $p(x)=\frac{\exp(-f(x))}{\int\exp(-f(x))dx}$. The learning rate for DSM is $5\times 10^{-5}$ and the learning rate for SSM is $1\times 10^{-5}$ since the variance of SSM is larger than DSM. The optimizer is Adam with $\beta_{1}=0.9$ and $\beta_{2}=0.95$. The sampling method is annealed SGLD with a total of $2,700$ steps. The $\epsilon$ in the finite-difference formulation is set to $0.05$. When training with annealed DSM, the noise level is an arithmetic sequence from $0.05$ to $1.2$ with the same number of steps as the batch size. The default batch size is $128$ in all our experiments unless specified. The backbone we use is an 18-layer ResNet following Li et al. . No normalizing layer is used in the backbone and the output layer is of a generalized quadratic form. The activation function is ELU. All experiments adopt the ResNet with 128 filters. During testing, we randomly sample $1500$ test data to evaluate the exact score matching loss.

NICE is a flow-based model, which converts a simple distribution $p_{0}$ to the data space $p$ using a invertible mapping $f$. In this case, the probability is defined as $\log p(x)=\log p_{0}(z)+\log\det(\frac{\partial z}{\partial x})$, where $z=f^{-1}(x)$ and $\det(\cdot)$ denotes the determinant of a matrix. The NICE model has 4 blocks with 5 fully connected layers in each block. Each layer has $1,000$ units. The activation is Softplus. Models are trained using Adam with a learning rate of $1\times 10^{-4}$. The data is dequantized by adding a uniform noise in the range of $[-\frac{1}{512},\frac{1}{512}]$, which is a widely adopted dequantization method for training flow models. The $\epsilon$ in the finite-difference formulation is set to $0.1$.

NCSN models a probability density by estimating its score function, i.e., $\nabla_{x}\log p(x)$, which is modeled by a score net. We follow Song and Ermon and provide an excerpt on the description of the model architecture design in the original paper: "We use a 4-cascaded RefineNet and pre-activation residual blocks. We replace the batch normalizations with CondInstanceNorm++ , and replace the max-pooling layers in Refine blocks with average pooling. Besides, we also add CondInstanceNorm++ before each convolution and average pooling in the Refine blocks. All activation functions are chosen to be ELU. We use dilated convolutions to replace the subsampling layers in residual blocks, except the first one. Following the common practice, we increase the dilation by a factor of 2 when proceeding to the next cascade. For CelebA and CIFAR-10 experiments, the number of filters for layers corresponding to the first cascade is 128, while the number of filters for other cascades are doubled. For MNIST experiments, the number of filters is halved."

C.3 Details of the results on out-of-distribution detection

For out-of-distribution (OOD) detection, we apply the typicality as the detection metric. Specifically, we first use the training set $\mathcal{D}_{\textup{train}}$ to approximate the entropy of model distribution as

where $|\mathcal{D}_{\textup{train}}|=N$ indicates the number of elements in the training set. Then give a set of test data $\mathcal{D}_{\textup{test}}$, where we control $|\mathcal{D}_{\textup{test}}|=M$ as a hyperparameter, then we can calculate the typicality as

Note that the metric in Eq. (56) naturally adapt to unnormalized models like EBMs, since there is

where the intractable partition function $Z_{\theta}$ can be eliminated after subtraction in Eq. (56). Thus we can calculate the typicality for EBMs as

As shown in Nalisnick et al. , a higher value of $M$ usually lead to better detection performance due to more accurate statistic. Thus to have distinguishable quantitative results, we set $M=2$ in our experiments. As to training the deep EBMs for the OOD detection, the settings we used on SVHN and CIFAR-10 are identical to those that we introduced above. On the ImageNet dataset, the images are cropped into a size of 128$\times$128, and we change the number of filters to 64 limited by the GPU memory. On SVHN and CIFAR-10, the models are trained on two GPUs, while the model is trained on eight GPUs on ImageNet. For all datasets, we use $N=50,000$ to estimate the data entropy and randomly sample $1000M$ test samples to conduct OOD detection.

C.4 Results on the VAE / WAE with implicit encoders

VAE / WAE with implicit encoders enable more flexible inference models. The gradient of the intractable entropy term $H(q)$ in the ELBO can be estimated by a score net. We adopt the identical neural architectures as in Song et al. . The encoder, decoder, and score net are both 3-layer MLPs with 256 hidden units on MNIST and 4-layer CNNs on CelebA. For MNIST, the optimizer is RMSProp with the learning rate as $1\times 10^{-3}$ in all methods. The learning rate is $1\times 10^{-4}$ on CelebA. All methods are trained for 10K iterations. The $\epsilon$ in the finite-difference formulation is set to $0.1$.