Estimating High Order Gradients of the Data Distribution by Denoising
Year: 2022

Abstract:
The first order derivative of a data density can be estimated efficiently by denoising score matching, and has become an important component in many applications, such as image generation and audio synthesis. Higher order derivatives provide additional local information about the data distribution and enable new applications. Although they can be estimated via automatic differentiation of a learned density model, this can amplify estimation errors and is expensive in high dimensional settings. To overcome these limitations, we propose a method to directly estimate high order derivatives (scores) of a data density from samples. We first show that denoising score matching can be interpreted as a particular case of Tweedie’s formula. By leveraging Tweedie’s formula on higher order moments, we generalize denoising score matching to estimate higher order derivatives. We demonstrate empirically that models trained with the proposed method can approximate second order derivatives more efficiently and accurately than via automatic differentiation. We show that our models can be used to quantify uncertainty in denoising and to improve the mixing speed of Langevin dynamics via Ozaki discretization for sampling synthetic data and natural images.

heading: 1 Introduction
text: The first order derivative of the log data density function, also known as score, has found many applications including image generation [23, 24, 6], image denoising [20, 19] and audio synthesis [9]. Denoising score matching (DSM) [29] provides an efficient way to estimate the score of the data density from samples and has been widely used for training score-based generative models [23, 24] and denoising [20, 19]. High order derivatives of the data density, which we refer to as high order scores, provide a more accurate local approximation of the data density (e.g., its curvature) and enable new applications. For instance, high order scores can improve the mixing speed for certain sampling methods [2, 18, 12], similar to how high order derivatives accelerate gradient descent in optimization [11]. In denoising problems, given a noisy datapoint, high order scores can be used to compute high order moments of the underlying noise-free datapoint, thus providing a way to quantify the uncertainty in denoising.
Existing methods for score estimation [8, 29, 25, 32], such as denoising score matching [29], focus on estimating the first order score (i.e., the Jacobian of the log density). In principle, high order scores can be estimated from a learned first order score model (or even a density model) via automatic differentiation. However, this approach is computationally expensive for high dimensional data and score models parameterized by deep neural networks. For example, given a D dimensional
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
distribution, computing the pn ` 1q-th order score value from an existing n-th order score model by automatic differentiation is on the order of D times more expensive than evaluating the latter [25]. Moreover, computing higher-order scores by automatic differentiation might suffer from large estimation error, since a small training loss for the first order score does not always lead to a small estimation error for high order scores.
To overcome these limitations, we propose a new approach which directly models and estimates high order scores of a data density from samples. We draw inspiration from Tweedie’s formula [4, 16], which connects the score function to a denoising problem, and show that denoising score matching (DSM) with Gaussian noise perturbation can be derived from Tweedie’s formula with the knowledge of least squares regression. We then provide a generalized version of Tweedie’s formula which allows us to further extend denoising score matching to estimate high order scores. In addition, we provide variance reduction techniques to improve the optimization of these newly introduced high order score estimation objectives. With our approach, we can directly parameterize high order scores and learn them efficiently, sidestepping expensive automatic differentiation.
While our theory and estimation method is applicable to scores of any order, we focus on the second order score (i.e., the Hessian of the log density) for empirical evaluation. Our experiments show that models learned with the proposed objective can approximate second order scores more accurately than applying automatic differentiation to lower order score models. Our approach is also more computationally efficient for high dimensional data, achieving up to 500ˆ speedups for second order score estimation on MNIST. In denoising problems, there could be multiple clean datapoints consistent with a noisy observation, and it is often desirable to measure the uncertainty of denoising results. As second order scores are closely related to the covaraince matrix of the noise-free data conditioned on the noisy observation, we show that our estimated second order scores can provide extra insights into the solution of denoising problems by capturing and quantifying the uncertainty of denoising. We further show that our model can be used to improve the mixing speed of Langevin dynamics for sampling synthetic data and natural images. Our empirical results on second order scores, a special case of the general approach, demonstrate the potential and applications of our method for estimating high order scores.

heading: 2 Background

heading: 2.1 Scores of a distribution
text: Definition 1. Given a probability density ppxq over RD, we define the k-th order score skpxq : RD Ñ bkRD, where bk denotes k-fold tensor multiplications, to be a tensor with the pi1, i2, . . . , ikqth index given by rskpxqsi1i2...ik fi B k
Bxi1Bxi2 ¨¨¨Bxik log ppxq, where pi1, i2, . . . , ikq P t1, ¨ ¨ ¨ , Duk.
As an example, when k “ 1, the first order score is the gradient of log ppxq w.r.t. to x, defined as s1pxq fi rx log ppxq. Intuitively, this is a vector field of the steepest ascent directions for the log-density. Note that the definition of first order score matches the definition of (Stein) score [8]. When k “ 2, the second order score is the Hessian of log ppxq w.r.t. to x. It gives the curvature of a density function, and with s1pxq it can provide a better local approximation to log ppxq.

heading: 2.2 Denoising score matching
text: Given a data distribution pdatapxq and a model distribution ppx;✓q, the score functions of pdatapxq and ppx;✓q are defined as s1pxq fi rx log pdatapxq and s1px;✓q fi rx log ppx;✓q respectively. Denoising score matching (DSM) [29] perturbs a data sample x „ pdatapxq with a pre-specified noise distribution q px̃ | xq and then estimates the score of the perturbed data distribution q px̃q “ ≥ q px̃ | xqpdatapxqdx which we denote s̃1px̃q fi rx̃ log q px̃q. DSM uses the following objective 1
2 EpdatapxqEq px̃|xqr}s̃1px̃;✓q ´ rx̃ log q px̃ | xq}22s. (1)
It is shown that under certain regularity conditions, minimizing Eq. (1) is equivalent to minimizing the score matching [8] loss between s̃1px̃;✓q and s̃1px̃q [29] defined as
1 2 EpdatapxqEq px̃|xqr}s̃1px̃;✓q ´ s̃1px̃q}22s. (2)
When q px̃ | xq “ N px̃|x, 2Iq , the objective becomes
LDSMp✓q “ 1
2 EpdatapxqEq px̃|xq
„›››s̃1px̃;✓q ` 1
2 px̃ ´ xq
››› 2
2
⇢ . (3)
Optimizing Eq. (3) can, intuitively, be understood as predicting x̃´x 2 , the added “noise" up to a constant, given the noisy input x̃, and is thus related to denoising. Estimating the score of the noise perturbed distribution q px̃q instead of the original (clean) data distribution pdatapxq allows DSM to approximate scores more efficiently than other methods [8, 25]. When is close to zero, q px̃q « pdatapxq so the score of q px̃q estimated by DSM will be close to that of pdatapxq. When is large, the estimated score for q px̃q plays a crucial role in denoising [20] and learning score-based generative models [23, 24].

heading: 2.3 Tweedie’s formula
text: Given a prior density pdatapxq, a noise distribution q px̃|xq “ N px̃|x, 2Iq, and the noisy density q px̃q “ ≥ pdatapxqq px̃|xqdx, Tweedie’s formula [16, 4] provides a close-form expression for the posterior expectation (the first moment) of x conditioned on x̃:
Erx | x̃s “ x̃ ` 2s̃1px̃q, (4) where s̃1px̃q fi rx̃ log q px̃q. Equation 4 implies that given a “noisy” observation x̃ „ q px̃q, one can compute the expectation of the “clean” datapoint x that may have produced x̃. As a result, Equation 4 has become an important tool for denoising [19, 20]. We provide the proof in Appendix B.
A less widely known fact is that Tweedies’ formula can be generalized to provide higher order moments of x given x̃, which we will leverage to derive the objective for learning higher order scores.

heading: 3 Estimating Higher Order Scores by Denoising
text: Below we demonstrate that DSM can be derived from Tweedie’s formula [4, 16]. By leveraging the generalized Tweedie’s formula on high order moments of the posterior, we extend DSM to estimate higher order score functions.

heading: 3.1 DSM in the view of Tweedie’s formula
text: The optimal solution to the least squares regression problem
min ✓
EpdatapxqEq px̃|xqr}hpx̃;✓q ´ x}22s (5)
is well-known to be the conditional expectation hpx̃;✓˚q “ Erx | x̃s. If we parameterize hpx̃;✓q “ x̃ ` 2s̃1px̃;✓q where s̃1px̃;✓q is a first order score model with parameter ✓, the least squares problem in Eq. (5) becomes equivalent to the DSM objective:
min ✓ EpdatapxqEq px̃|xqr} 2s̃1px̃;✓q ` x̃ ´ x}22s “ min✓ 2 4 ¨ LDSMp✓q. (6)
From Tweedie’s formula, we know the optimal ✓˚ satisfies hpx̃;✓˚q “ x̃ ` 2s̃1px̃;✓˚q “ Erx | x̃s “ x̃` 2s̃1px̃q, from which we can conclude that s̃1px̃;✓˚q “ s̃1px̃q. This proves that minimizing the DSM objective in Eq. (6) recovers the first order score.
There are other ways to derive DSM. For example, [15] provides a proof based on Bayesian least squares without relying on Tweedie’s formula. Stein’s Unbiased Risk Estimator (SURE) [27] can also provide an alternative proof based on integration by parts. Compared to these methods, our derivation can be easily extended to learn high order scores, leveraging a more general version of Tweedie’s formula.

heading: 3.2 Second order denoising score matching
text: As a warm-up, we first consider the second order score, and later generalize to any desired order. Leveraging Tweedie’s formula on ErxxT | x̃s and Erx | x̃s, we obtain the following theorem.
Theorem 1. Given a D-dimensional distribution ppxq and q px̃q fi ≥ ppxqq px̃|xqdx, we have
ErxxT | x̃s “ fpx̃, s̃1, s̃2q (7) ErxxT ´ xx̃T ´ x̃xT | x̃s “ hpx̃, s̃1, s̃2q, (8)
where fpx̃, s̃1, s̃2q and hpx̃, s̃1, s̃2q are polynomials of x̃, s̃1px̃q, s̃2px̃q defined as
fpx̃, s̃1, s̃2q “ x̃x̃T ` 2x̃s̃1px̃qT ` 2s̃1px̃qx̃T ` 4s̃2px̃q ` 4s̃1px̃qs̃1px̃qT ` 2I, (9) hpx̃, s̃1, s̃2q “ ´x̃x̃T ` 4s̃2px̃q ` 4s̃1px̃qs̃1px̃qT ` 2I. (10)
Here s̃1px̃q and s̃2px̃q denote the first and second order scores of q px̃q.
In Theorem 1, Eq. (9) is directly given by Tweedie’s formula on ErxxT | x̃s, and Eq. (10) is derived from Tweedie’s formula on both Erx | x̃s and ErxxT | x̃s. Given a noisy sample x̃, Theorem 1 relates the second order moment of x to the first order score s̃1px̃q and second order score s̃2px̃q of q px̃q. A detailed proof of Theorem 1 is given in Appendix B. In the same way as how we derive DSM from Tweedie’s formula in Section 3.1, we can obtain higher order score matching objectives with Eq. (9) and Eq. (10) as a least squares problem. Theorem 2. Suppose the first order score s̃1px̃q is given, we can learn a second order score model s̃2px̃;✓q by optimizing the following objectives
✓ ˚ “ argmin
✓ EpdatapxqEq px̃|xq
„›››xxT ´ fpx̃, s̃1px̃q, s̃2px̃;✓qq ››› 2
2
⇢ , (11)
✓ ˚ “ argmin
✓ EpdatapxqEq px̃|xq
„›››xxT ´ xx̃T ´ x̃xT ´ hpx̃, s̃1px̃q, s̃2px̃;✓qq ››› 2
2
⇢ (12)
where fp¨q and hp¨q are polynomials defined in Eq. (9) and Eq. (10). Assuming the model has an infinite capacity, then the optimal parameter ✓ ˚ satisfies s̃2px̃;✓˚q “ s̃2px̃q for almost any x̃.
Here Eq. (11) and Eq. (12) correspond to the least squares objective of Eq. (7) and Eq. (8) respectively, and have the same set of solutions assuming sufficient model capacity. In practice, we find that Eq. (12) has a much simpler form than Eq. (11), and will therefore use Eq. (12) in our experiments.

heading: 3.3 High order denoising score matching
text: Below we generalize our approach to even higher order scores by (i) leveraging Tweedie’s formula to connect higher order moments of x conditioned on x̃ to higher order scores of q px̃q; and (ii) finding the corresponding least squares objective.
Theorem 3. Erbnx|x̃s “ fnpx̃, s̃1, ..., s̃nq, where bnx P RD n
denotes n-fold tensor multiplications, fnpx̃, s̃1, ..., s̃nq is a polynomial of tx̃, s̃1px̃q, ..., s̃npx̃qu and s̃kpx̃q represents the k-th order score of q px̃q “ ≥ pdatapxqq px̃|xqdx.
Theorem 3 shows that there exists an equality between (high order) moments of the posterior distribution of x given x̃ and (high order) scores with respect to x̃. To get some intuition, for n “ 2 the polynomial f2px̃, s̃1, s̃2q is simply the function f in Eq. (9). In Appendix B, we provide a recursive formula for obtaining the coefficients of fn in closed form. Leveraging Theorem 3 and the least squares estimation of Erbkx|x̃s, we can construct objectives for approximating the k-th order scores s̃kpx̃q as in the following theorem. Theorem 4. Given score functions s̃1px̃q, ..., s̃k´1px̃q, a k-th order score model s̃kpx̃;✓q, and
✓ ˚ “ argmin ✓ EpdatapxqEq px̃|xqr} bk x ´ fkpx̃, s̃1px̃q, ..., s̃k´1px̃q, s̃kpx̃;✓qq}2s.
We have s̃kpx̃;✓˚q “ s̃kpx̃q for almost all x̃.
As previously discussed, when approaches 0 such that q px̃q « pdatapxq, s̃kpx̃;✓˚q wellapproximates the k-th order score of pdatapxq.

heading: 4 Learning Second Order Score Models
text: Although our theory can be applied to scores of any order, we focus on second order scores for empirical analysis. In this section, we discuss the parameterization and empirical performance of the learned second order score models.

heading: 4.1 Instantiating objectives for second order score models
text: In practice, we find that Eq. (12) has a much simpler expression than Eq. (11). Therefore, we propose to parameterize s̃2px̃q with a model s̃2px̃;✓q, and optimize s̃2px̃;✓q with Eq. (12), which can be simplified to the following after combining Eq. (10) and Eq. (12):
LD2SMp✓q fi EpdatapxqEq px̃|xq „›››s̃2px̃;✓q ` s̃1px̃;✓qs̃1px̃;✓qT ` I ´ zzT 2 ››› 2
2
⇢ , (13)
where z fi x̃´x . Note that Eq. (13) requires knowing the first order score s̃1px̃q in order to train the second order score model s̃2px̃;✓q. We therefore use the following hybrid objective to simultaneously train both s̃1px̃;✓q and s̃2px̃;✓q:
Ljointp✓q “ LD2SMp✓q ` ¨ LDSMp✓q, (14) where LDSMp✓q is defined in Eq. (3) and P R°0 is a tunable coefficient. The expectation for LD2SMp✓q and LDSMp✓qin Eq. (14) can be estimated with samples, and we optimize the following unbiased estimator
L̂jointp✓q “ 1
N
Nÿ
i“1
„›››s̃2px̃i;✓q ` s̃1px̃i;✓qs̃1px̃i;✓qT ` I ´ zizTi
2
››› 2
2 ` 2
›››s̃1px̃i;✓q ` zi ››› 2
2
⇢ , (15)
where we define zi fi x̃i´xi , and tx̃iuNi“1 are samples from q px̃q “ ≥ pdatapxqq px̃|xqdx which can be obtained by adding noise to samples from pdatapxq. Similarly to DSM, when Ñ 0, the optimal model s̃2px̃;✓˚q that minimizes Eq. (15) will be close to the second order score of pdatapxq because q px̃q « pdatapxq. When is large, the learned s̃2px̃;✓q can be applied to tasks such as uncertainty quantification for denoising, which will be discussed in Section 5.
For downstream tasks that require only the diagonal of s̃2, we can instead optimize a simpler objective Ljoint-diagp✓q fi LD2SM-diagp✓q ` ¨ LDSMp✓q, where (16)
LD2SM-diagp✓q fi EpdatapxqEq px̃|xq „›››diagps̃2px̃;✓qq ` s̃1px̃;✓q d s̃1px̃;✓q ` 1 ´ z d z 2 ››› 2
2
⇢ . (17)
Here diagp¨q denotes the diagonal of a matrix and d denotes element-wise multiplication. Optimizing Eq. (16) only requires parameterizing diagps̃2px̃;✓qq, which can significantly reduce the memory and computational cost for training and running the second order score model. Similar to L̂jointp✓q, we estimate the expectation in Eq. (17) with empirical means.

heading: 4.2 Parameterizing second order score models
text: In practice, the performance of learning second order scores is affected by model parameterization. As many real world data distributions (e.g., images) tend to lie on low dimensional manifolds [13, 3, 21], we propose to parametrize s̃2px̃;✓q with low rank matrices defined as below
s̃2px̃;✓q “ ↵px̃;✓q ` px̃;✓q px̃;✓qT, where ↵p¨;✓q : RD Ñ RDˆD is a diagonal matrix, p¨;✓q : RD Ñ RDˆr is a matrix with shape D ˆ r, and r § D is a positive integer.

heading: 4.3 Antithetic sampling for variance reduction
text: As the standard deviation of the perturbed noise approximates zero, training score models with denoising methods could suffer from a high variance. Inspired by a variance reduction method for DSM [30, 26], we propose a variance reduction method for D2SM
LD2SM-VR “ Ex„pdatapxqEz„N p0,Iq „ px̃`q2 ` px̃´q2 ` 2 I ´ zzT 2 d p px̃`q ` px̃´q ´ 2 pxqq ⇢ ,
where x̃` “ x` z, x̃´ “ x´ z and “ s̃2 ` s̃1s̃T1 . Instead of using independent noise samples, we apply antithetic sampling and use two correlated (opposite) noise vectors centered at x. Similar to Eq. (14), we define Ljoint-VR “ LD2SM-VR ` ¨ LDSM-VR, where LDSM-VR is proposed in [30]. We empirically study the role of variance reduction (VR) in training models with DSM and D2SM. We observe that VR is crucial for both DSM and D2SM when is approximately zero, but is optional when is large enough. To see this, we consider a 2-d Gaussian distribution N p0, Iq and train s̃1px̃;✓q and s̃2px̃;✓q with DSM and D2SM respectively. We plot the learning curves in Figs. 1a and 1b, and visualize the first dimension of the estimated scores for multiple noise scales in Figs. 1c and 1d. We observe that when “ 0.001, both DSM and D2SM have trouble converging after a long period of training, while the VR counterparts converge quickly (see Fig. 1). When gets larger, DSM and D2SM without VR can both converge quickly and provide reasonable score estimations (Figs. 1c and 1d). We provide extra details in Appendix C.

heading: 4.4 The accuracy and efficiency of learning second order scores
text: We show that the proposed method can estimate second order scores more efficiently and accurately than those obtained by automatic differentiation of a first order score model trained with DSM. We observe in our experiments that s̃1px̃;✓q jointly optimized via L̂joint or L̂joint-diag has a comparable empirical performance as trained directly by DSM, so we optimize s̃1px̃;✓q and s̃2px̃;✓q jointly in later experiments. We provide additional experimental details in Appendix C.
Learning accuracy We consider three synthetic datasets whose ground truth scores are available—a 100-dimensional correlated multivariate normal distribution and two high dimensional mixture of logistics distributions in Table 1. We study the performance of estimating s̃2 and the diagonal of s̃2. For the baseline, we estimate second order scores by taking automatic differentiation of s̃1px̃;✓q trained jointly with s̃2px̃;✓q using Eq. (15) or Eq. (17). As mentioned previously, s̃1px̃;✓q trained with the joint method has the same empirical performance as trained directly with DSM. For our method, we directly evaluate s̃2px̃;✓q. We compute the mean squared error between estimated second order scores and the ground truth score of the clean data since we use small and q px̃q « pdatapxq (see Table 1). We observe that s̃2px̃;✓q achieves better performance than the gradients of s̃1px̃;✓q.
Computational efficiency Computing the gradients of s̃1px̃;✓q via automatic differentiation can be expensive for high dimensional data and deep neural networks. To see this, we consider two models—a 3-layer MLP and a U-Net [17], which is used for image experiments in the subsequent sections. We consider a 100-d data distribution for the MLP model and a 784-d data distribution for the U-Net. We parameterize s̃1 and s̃2 with the same model architecture and use a batch size of 10 for both settings. We report the wall-clock time averaged in 7 runs used for estimating second order scores during test time on a TITAN Xp GPU in Table 2. We observe that s̃2px̃;✓q is 500ˆ faster than using automatic differentiation for s̃1px̃;✓q on the MNIST dataset.

heading: 5 Uncertainty Quantification with Second Order Score Models
text: Our second order score model s̃2px̃;✓q can capture and quantify the uncertainty of denoising on synthetic and real world image datasets, based on the following result by combining Eqs. (4) and (9)
Covrx | x̃s fi ErxxT | x̃s ´ Erx | x̃sErx | x̃sT “ 4s̃2px̃q ` 2I (18) By estimating Covrx | x̃s via s̃2px̃q, we gain insights into how pixels are correlated with each other under denoising settings, and which part of the pixels has large uncertainty. To examine the uncertainty given by our s̃2px̃;✓q, we perform the following experiments (details in Appendix D). Synthetic experiments We first consider 2-d synthetic datasets shown in Fig. 2, where we train s̃1px̃;✓q and s̃2px̃;✓q jointly with Ljoint. Given the trained score models, we estimate Erx | x̃s and Covrx | x̃s using Eq. (4) and Eq. (18). We approximate the posterior distribution ppx|x̃q with a conditional normal distribution N px | Erx | x̃s,Covrx | x̃sq. We compare our result with that of Eq. (4), which only utilizes s̃1 (see Fig. 2). We observe that unlike Eq. (4), which is a point estimator, the incorporation of covariance matrices (estimated by s̃2px̃;✓q) captures uncertainty in denoising. Covariance diagonal visualizations We visualize the diagonal of the estimated Covrx | x̃s for MNIST and CIFAR-10 [10] in Fig. 3. We find that the diagonal values are in general larger for pixels near the edges where there are multiple possibilities corresponding to the same noisy pixel. The diagonal values are smaller for the background pixels where there is less uncertainty. We also observe that covariance matrices corresponding to smaller noise scales tend to have smaller values on the diagonals, implying that the more noise an image has, the more uncertain the denoised results are.
Full convariance visualizations We visualize the eigenvectors (sorted by eigenvalues) of Covrx | x̃s estimated by s̃2px̃;✓q in Fig. 4. We observe that they can correspond to different digit identities, indicating uncertainty in the identity of the denoised image. This suggests Covrx | x̃s can capture additional information for uncertainty beyond its diagonal.

heading: 6 Sampling with Second Order Score Models
text: Here we show that our second order score model s̃2px̃;✓q can be used to improve the mixing speed of Langevin dynamics sampling.
Table 2: Speed analysis: direct modeling vs. autodiff.
Dataset 1 Dataset 2
Method Metric ESS Ò ESS Ò Langevin 21.81 26.33
Ozaki 28.89 46.57

heading: 6.1 Background on the sampling methods
text: Langevin dynamics Langevin dynamics [1, 31] samples from pdatapxq using the first order score function s1pxq. Given a prior distribution ⇡pxq, a fixed step size ✏ ° 0 and an initial value x̃0 „ ⇡pxq, Langevin dynamics update the samples iteratively as follows
x̃t “ x̃t´1 ` ✏ 2
s1px̃t´1q ` ? ✏zt, (19)
where zt „ N p0, Iq. As ✏ Ñ 0 and t Ñ 8, x̃t is a sample from pdatapxq under suitable conditions. Ozaki sampling Langevin dynamics with Ozaki discretization [28] leverages second order information in s2pxq to pre-condition Langevin dynamics:
x̃t “ x̃t´1 ` Mt´1s1px̃t´1q ` ⌃1{2t´1zt, zt „ N p0, Iq (20) where Mt´1 “ pe✏s2px̃t´1q ´ Iqs2px̃t´1q´1 and ⌃t´1 “ pe2✏s2px̃t´1q ´ Iqs2px̃t´1q´1. It is shown that under certain conditions, this variation can improve the convergence rate of Langevin dynamics [2]. In general, Eq. (20) is expensive to compute due to inversion, exponentiation and taking square root of matrices, so we simplify Eq. (20) by approximating s2px̃t´1q with its diagonal in practice. In our experiments, we only consider Ozaki sampling with s2 replaced by its diagonal in Eq. (20). As we use small , s̃1 « s1 and s̃2 « s2. We observe that diagps̃2px̃;✓qq in Ozaki sampling can be computed in parallel with s̃1px̃;✓q on modern GPUs, making the wall-clock time per iteration of Ozaki sampling comparable to that of Langevin dynamics. Since we only use the diagonal of s̃2px̃;✓q in sampling, we can directly learn the diagonal of s̃2px̃q with Eq. (16).

heading: 6.2 Synthetic datasets
text: We first consider 2-d synthetic datasets in Fig. 5 to compare the mixing speed of Ozaki sampling with Langevin dynamics. We search the optimal step size for each method and observe that Ozaki sampling can use a larger step size and converge faster than Langevin dynamics (see Fig. 5). We use the optimal step size for both methods and report the smallest effective sample size (ESS) of all the dimensions [22, 5] in Table 3. We observe that Ozaki sampling has better ESS values than Langevin dynamics, implying faster mixing speed. Even when using the same step size, Ozaki sampling still converges faster than Langevin dynamics on the two-model Gaussian dataset we consider (see Fig. 6). In all the experiments, we use “ 0.1 and we provide more experimental details in Appendix E.

heading: 6.3 Image datasets
text: Ozaki discretization with learned s̃2px̃;✓q produces more diverse samples and improve the mixing speed of Langevin dynamics on image datasets (see Fig. 7) To see this, we select ten different digits from MNIST test set and initialize 1000 different sampling chains for each image. We update the
Ozaki
Langevin
5 iterations 10 iterations 20 iterations 50 iterations 100 iterations 200 iterations
(a) Percentage of changes in class label w.r.t. iterations.
(b) Different chains initialized with the same left panel image after 1000 iterations of update with step size ✏ “ 0.03.
Figure 8: Sample diversity analysis. The number in the parenthesis in Fig. 8a denotes the step size. We initialize the chain with MNIST test images and report the percentage of images that have changed class labels from the initialization w.r.t. sampling iterations. We observe that Ozaki sampling has more diverse samples.
chains with Ozaki sampling and report the percentage of images that have class label changes after a fixed number of sampling iterations in Fig. 8a. We compare the results with Langevin dynamics with the same setting and observe that Ozaki sampling has more diverse samples within the same chain in a fixed amount of iterations. We provide more details in Appendix E.

heading: 7 Conclusion
text: We propose a method to directly estimate high order scores of a data density from samples. We first study the connection between Tweedie’s formula and denoising score matching (DSM) through the lens of least squares regression. We then leverage Tweedie’s formula on higher order moments, which allows us to generalize denoising score matching to estimate scores of any desired order. We demonstrate empirically that models trained with the proposed method can approximate second order scores more efficiently and accurately than applying automatic differentiation to a learned first order score model. In addition, we show that our models can be used to quantify uncertainty in denoising and to improve the mixing speed of Langevin dynamics via Ozaki discretization for sampling synthetic data and natural images. Besides the applications studied in this paper, it would be interesting to study the application of high order scores for out of distribution detection. Due to limited
computational resources, we only consider low resolution image datasets in this work. However, as a direct next step, we can apply our method to higher-resolution image datasets and explore its application to improve the sampling speed of score-based models [23, 24, 6] with Ozaki sampling. In general, when approximating the high-order scores with a diagonal or a low rank matrix, our training cost is comparable to standard denoising score matching, which is scalable to higher dimensional data. A larger rank typically requires more computation but could give better approximations to second-order scores. While we focused on images, this approach is likely applicable to other data modalities such as speech.

heading: Acknowledgements
text: The authors would like to thank Jiaming Song and Lantao Yu for constructive feedback. This research was supported by NSF (#1651565, #1522054, #1733686), ONR (N000141912145), AFOSR (FA95501910024), ARO (W911NF-21-1-0125) and Sloan Fellowship.

Review summaries:
- The author show that denoising score matching with Gaussian noise can be derived from Tweedie's formula through the lens of least square regressions. - This provides a new interpretation of the first score (gradient of log density). - They use the generalization of Tweedie's formula to higher orders to derive a similar estimation of higher order scores. - This is accompanied by a thorough experimental validation for the second order score. 

This paper utilizes Tweedie's formula to generalize denoising score matching to higher order derivatives. The proposed method can approximate second order derivatives more efficiently and accurately than differentiating first order derivatives. Experiments on synthetic data and small image datasets also demonstrates its applications to uncertainty estimation, and Langevin dynamics sampling.

Using Tweedie's formula, the paper proposes a $n$-th order denoising score matching: a parametric model learns $n$-th order gradient of the log density of perturbed data. More specifically, the paper considers a data $x \sim p(x)$ and a perturbed random variable $\tilde{x} \sim p(\tilde{x})$, which is a marginal distribution of $\tilde{x}$ for a given a conditional distribution $p(\tilde{x}|x)$ and $p(x)$.  Tweedie's formula states that when $p(\tilde{x}|x)$ is an exponential family distribution, the expected value of the $n$-th order moment of the posterior distribution of $x$ given $\tilde{x}$ can be represented as a polynomial function of $\tilde{x}$ and from-first-to-$n$th order gradients of the log marginal $\log p(\tilde{x})$. Nothing that the optimal solution of the least-squares regression to $x$'s moments for a given $\tilde{x}$ is the expected value of the posterior moments, the paper proposes to learn $n$-th order score models by performing the regression. The authors emphasize that for guaranteeing optimality of $n$-th order denoising score matching, ground-truth scores should be accessible up to $n$-1-th order.  Nevertheless, the paper proposes to jointly train the first and second-order score models for empirical studies. In particular, the authors suggest a memory-efficient parameterization of the second-order score and a variance reduction method for the training. For the application of the joint training, the authors first demonstrate that the learned second-order scores provide information about the uncertainty of denoising. Second, the paper proposes to plug in the learned second-order gradients into the Ozaki discretization of Langevin dynamics, which uses the second-order gradients to improve the mixing speed of the Langevin dynamics. The experiments on synthetic data and natural images demonstrate that the second-order score model improves sampling quality for the same discretization steps compared to the Euler-Maruyama discretization (the first order).

The paper generalizes denoising score matching, which estimates the score, i.e. the gradient of the log density, of perturbed data. In particular, it shows how higher-order derivates of perturbed data densities can be estimated and learnt. Much of this is derived via interesting connections of denoising score matching to Tweedie's formula, a technique for denoising noisy data. While the paper theoretically shows how to learn derivatives of the log data density of arbitrary order, its experiments focus on second-order derivatives, corresponding to the Hessian of the log density. The learnt second-order score is leveraged for improved image (and toy data) denoising, where the second-order score allows to parametrize a covariance matrix, that captures the uncertainty in the denoising operation. Furthermore, it uses the second-order score for sampling applications where it is leveraged to design more efficient samplers via preconditioning. It is also shown that directly learning the second-order score as proposed is more accurate and faster than computing it via auto-differentiation from first-order scores at test time.