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