Estimating High Order Gradients of the Data Distribution by Denoising

1 Introduction . 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 . 2 Background . 2.1 Scores of a distribution . 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 . 2.2 Denoising score matching . 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 ] . 2.3 Tweedie ’ s formula . 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 . 3 Estimating Higher Order Scores by Denoising . 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 . 3.1 DSM in the view of Tweedie ’ s formula . 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 . 3.2 Second order denoising score matching . 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 .

Summary: 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.