Local Curvature Smoothing with Stein's Identity for Efficient Score Matching

Authors: genki  osada
takashi  nishide
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 05:09:10

Abstract: The training of score-based diffusion models (SDMs) is based on score matching. The challenge of score matching is that it includes a computationally expensive Jacobian trace. While several methods have been proposed to avoid this computation, each has drawbacks, such as instability during training and approximating the learning as learning a denoising vector field rather than a true score.
We propose a novel score matching variant, local curvature smoothing with Stein's identity (LCSS). The LCSS bypasses the Jacobian trace by applying Stein's identity, enabling regularization effectiveness and efficient computation. We show that LCSS surpasses existing methods in sample generation performance and matches the performance of denoising score matching, widely adopted by most SDMs, in evaluations such as FID, Inception score, and bits per dimension. Furthermore, we show that LCSS enables realistic image generation even at a high resolution of $1024 \times 1024$.

Main paper: 
 introduction: 
 Score-based diffusion models (SDMs) [35, 30, 7, 31, 12] have emerged as powerful generative models that have achieved remarkable results in various fields [27, 26, 24] . While likelihood-based models learn the density of observed (i.e., training) data as points [25, 5, 28, 38, 37, 14] , SDMs learn the gradient of logarithm density called the score -a vector field pointing toward increasing data density. The sample generation process of SDMs has two steps: 1) learning the score for a given dataset and 2) generating samples by guiding a random noise vector toward high-density regions based on the learned score using stochastic differential equation (SDE).Score matching used for learning the score includes a computationally expensive Jacobian trace, making it challenging to apply to high-dimensional data. While some methods have been proposed to avoid computing the Jacobian trace, each has its drawbacks. Denoising score matching (DSM) [39] , ubiquitously employed in SDMs, learns not the ground truth score but its approximation and imposes constraints on the design of SDE. On the other hand, sliced score matching (SSM) [33] and its variant, finite-difference SSM (FD-SSM) [23] , suffer from high variance due to using random projection.In this paper, we propose a novel score matching variant, local curvature smoothing with Stein's identity (LCSS). The key idea of LCSS is to use Stein's identity to bypass the expensive computation of Jacobian trace. To apply Stein's identity, we take the expectation over a Gaussian distribution centered on input data points, which is indeed equivalent to the regularization with local curvature 38th Conference on Neural Information Processing Systems (NeurIPS 2024).smoothing. Exploiting this equivalence, we propose a score matching method that offers both regularization benefits and faster computation.We first establish a method as an independent score matching technique, then propose a timeconditional version for its application to SDMs. We present the experimental results using synthetic data and several popular datasets. Our proposed method is highly efficient compared to existing score matching methods and enables the generation of high-resolution images with a size of 1024. We show that LCSS outperforms SSM, FD-SSM, and DSM in the quality of generated images and is comparable to DSM in the qualitative evaluation of the FID score, Inception score, and the negative likelihood measured in bits per dimension. While DSM requires the drift and diffusion coefficients of an SDE to be affine, our LCSS has no such constraint, allowing for a more flexible SDE design (Sec. 2.4). Hence, this paper contributes to opening up new directions in SDMs' research based on more flexible SDEs.Related works. Liu et al. [20] proposed a method for directly estimating scores using Stein's identity without using score matching. Shi et al. [29] further enhanced that method by applying spectral decomposition to the function in Stein's identity. However, Song et al. [33] reported that these methods underperform compared to SSM. In our approach, we use Stein's identity specifically to avoid computing the Jacobian trace in score matching. The regularization effect attained by adding noise to data has been recognized for a long time [3] , which our method utilizes. The relationship between noise-adding regularization and curvature smoothing in the least square function is elucidated in Bishop [2] . The previous studies of score matching variants are described in the next section. In efforts to remove the affine constraint of SDE in SDMs, Kim et al. [13] proposed running SDE in the latent space of normalizing flows. This constraint stems from using DSM for score matching, and we propose a score matching method free from such constraint. 
preliminary 2.1 score-based diffusion models: 
 Score-based diffusion models (SDMs) [35, 31] define an stochastic differential equation (SDE) for x t ∈ R d in continuous time t ∈ [0, T ] asEQUATIONwhere f (•, t) : R d → R d is the drift coefficient, g(t) ∈ R is the diffusion coefficient, and w t denotes a standard Wiener process. Eq. ( 1), known as the forward process, has a corresponding reverse process from time T to 0 [1] :dx t = f (x t , t) -g(t) 2 ∇ x log p t (x t ) dt + g(t)d wt (2)where wt is a standard Wiener process in reverse-time and p t (x t ) denotes the ground truth marginal density of x t following the forward process. Samples from a dataset are represented as x 0 ∼ p 0 , while initial vectors for sample generation with Eq. ( 2) are x T ∼ p T . In Eq. ( 2), the only unknown term is ∇ x log p t (x t ), referred to as the score of density p t (x t ). To estimate ∇ x log p t (x t ), SDMs train a score network s θ parametrized by θ by score matching.A score network s θ that estimates the score of the ground truth density is trained through score matching [9] . Score matching, a technique independent of SDMs and SDE, has no concept of time. So, as long as our discussion is focused on score matching, we use the notation of x and p, without the subscript of t, and treat a score network without conditioning on t, i.e., denote it as s θ (x) instead of s θ (x t , t). Score matching is defined as the minimization ofEQUATION2 . Calculating J (θ) is generally impractical since it requires knowing the ground truth ∇ x log p(x), but Hyvärinen [9] has shown that J (θ) is equivalent to the following J SM (θ) up to constant:EQUATIONwhere J s SM (θ, x) is the version of J SM (θ) for a single data point x, defined asEQUATIONScore matching in SDMs and its problem. SDMs train a time-conditional score function s θ (x t , t) using score matching. The loss function of SDMs is defined as the integral of J SM (θ) over time t ∈ [0, T ] asEQUATIONThe weight function λ(t) is determined by the form of the SDE, and λ(t) used for typical SDEs can be found in Table 1 in [34] . The p t is obtained from the SDE in Eq. ( 1), with its mean dependent on x 0 , and its specific form in typical SDEs is given as Eq. ( 29) in Song et al. [35] . The problem is that since s θ (x t , t) has the same dimension as input x t , computing its Jacobian trace, Tr(∇ x s θ (x t , t)), is costly. It renders training with score matching impractical in high-dimensional data.To avoid the computation of the Jacobian trace, the following scalable variants of J SM have been developed. While any score matching method can be used to train s θ , SDMs predominantly employ DSM due to its empirical performance [31, 35, 12] .Sliced score matching (SSM) and finite-difference sliced score matching (FD-SSM). SSM [33] approximates Tr(∇ x s θ (x)) with Hutchinson's trick [8] and minimizes the following:EQUATIONwhere is a small scaler value and v ∼ p v is a d-dimensional random vector such thatE v∼pv [vv T ] = 2 I d .To enhance the efficiency of SSM further, FD-SSM [23] adopts finite difference to Eq. ( 6). The objective function isEQUATIONThe drawback of these two methods is the high variance induced by random projection with v. In particular, the error between the true trace of matrix A, Tr(A), and the estimate by Hutchinson's trick, TA , is | Tr(A) -TA | ≤ 1 √ M A F where • F is the Frobenius norm and M is the sampling times from p v [22] . Typically, M = 1 setting is employed in these methods, potentially making the error magnitude non-negligible and causing instability in training process, as we see in Sec 4.2.3.Denoising score matching (DSM). DSM [39] circumvents the computation of Tr(∇ x s θ (x)) by perturbing x with a Gaussian noise distribution q σ (x|x) with noise scale σ and then estimating the score of the perturbed distribution q σ (x) := q σ (x|x)p(x)dx. The DSM minimizesEQUATIONIn SDMs, the following time-conditional version is used:EQUATIONwhere σ t is designed to increase as t progresses from 0 to T . Almost all SDMs use DSM for score matching because it performs faster and is more stable than SSM and FD-SSM. However, DSM has three drawbacks. 1) Approximation: in DSM, s θ learns ∇ x log q σt (x t |x 0 ) rather than the ground true score, ∇ x log p t (x t ). 2) Constraining the design of SDE: DSM constrains SDE coefficients to be affine. We will describe this in Sec. 2.4.3) The dilemma regarding σ t : Only when σ t → 0 does ∇ x log q σt (x t |x 0 ) match ∇ x log p t (x t ). However, as σ → 0, both the numerator and denominator of x0-xtσ 2 tapproach 0, leading to potential numerical instability [19] .The design of SDEs directly influences the performance of SDMs, as demonstrated in previous studies [12] . The benefits of non-linear SDE, particularly highlighted in [13] , enable more accurate alignment of scores with the ground-truth data distributions than affine SDE and thus enhance the quality of generated samples. (Fig. 2 in [13] illustrates this.) However, unless specific modifications are made as proposed in these studies, the general SDEs [35] used in almost all existing SDMs must be affine. This constraint comes from the fact that the SDMs, consciously or unconsciously, select DSM for their score matching methods. The loss function of DSM requires ∇ x log q σt (x t |x 0 ) as Eq. (9) . Thus, to compute Eq. ( 9) at every training iteration, ∇ x log q σt (x t |x 0 ) needs to be in closed form. DSM models q σt (x t |x 0 ) as a Gaussian distribution, for which this requirement is satisfied as∇ x log q σt (x t |x 0 ) = x0-xt σ 2 t. However, this Gaussian modeling comes at the cost of imposing a constraint on the SDE design: the drift and diffusion terms of SDE, i.e., f (x t , t) and g(t) in Eq. ( 1), need to be affine. The existing SDMs are DSM-based, so the SDEs used in these SDMs, including the VE SDE and subVP SDE we use in our experiments, are designed to adhere to this constraint. The details of the same discussion and the specific form of the Gaussian distribution q σt (x t |x 0 ) for the typical SDEs can be found in Sec. 3.3 in Song et al. [35] . Unlike DSM, SSM and FD-SSM do not have this limitation, allowing for more flexible SDE design and thus removing the requirement to limit the forward process's convergence destination to Gaussian distributions. Unfortunately, as we will see later, SSM and FD-SSM cannot handle high-dimensional data due to the high-variance they cause. Our proposed method uniquely satisfies both the flexible design of SDEs and compatibility with high-dimensional data.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: Section 5 contains a paragraph explaining the limitations. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
our method: 
 We propose a novel score matching variant that avoids the expensive computation of the Jacobian trace. The crux of our method is using Stein's identity to bypass Jacobian computation. Our approach comprises three steps: 1) introducing local curvature smoothing regularization into score matching (Definition 1), 2) treating the regularization of 1) as taking an expectation over a Gaussian distribution (Lemma 1), and 3) applying Stein's identity (Corollary 2). While introducing regularization may appear to cause extra computational costs, it enables faster computation by the use of Stein's identity trick. We begin by discussing our method separately from SDMs, without involving the time variable t, and then explain its incorporation into SDMs at the end of this section.We first introduce some lemmas and corollaries that constitute our method. Definition 1 (Score matching with local curvature smoothing [15] ). Regularizing the score matching objective J s SM at a data point x ∈ R d with local curvature smoothing (LCS) is defined as:EQUATIONGiven ∇ x s θ (x) approximating the Hessian of log p(x), minimizing the regularization termEQUATIONF acts as a local curvature smoothing where the square of the curvature of the surface of the log-density at x are penalized. Curvature smoothing is one of the commonly employed regularizations in machine learning [3] . Lemma 1 (Kingma and LeCun [15] ). Score matching with local curvature smoothing (Definition 1) is equivalent to the expectation of J s SM (θ, x) over a Gaussian distribution centered at x, i.e., x ∼ N (x, σ 2 I d ):EQUATIONwhere:= x -x 2 .Lemma 1 states that taking the expectation of score matching objective with respect to a Gaussian distribution centered around x yields an effect equivalent to a curvature smoothing regularization. Definition 2 (Stein class [36] ). Assume thatEQUATIONThe condition for Eq. ( 12) to hold isEQUATION)Lemma 2 (Stein's identity, Liu et al. [20] , Gorham and Mackey [6] ). Let h : Z → R d be a smooth (i.e., continuous and differentiable) vector valued functionh(z) = [h 1 (z), h 2 (z), . . . h d (z)] T . Then , if h i (z)∀i = 1, .. . , d is the Stein class of a smooth density Q(z), the following identity holds:EQUATIONIn Eq. ( 14),∇ z log Q(z) T is a 1 × d matrix, ∇ z h(z) is a d × d matrix, and 0 d ,d is a d × d zero matrix.Corollary 1 (Li and Turner [18] ). When Q(z) = N (z; µ, σ 2 I d ), we haveEQUATIONEq. ( 15) holds for the i-th element of the vector h. The condition Eq. ( 13) holds for Gaussian distribution Q, since Q(z) → 0 as z → ∞. Then, h i (z)∀i = 1, . . . , d are the Stein class of Q, and thus Lemma 2 is valid for a Gaussian distribution Q. As we also know∇ z log Q(z) = -1 σ 2 (z -µ), by substituting it into Lemma 2, we obtain Corollary 1. Corollary 2 (Bypassing Jacobian trace computation). Let x ∈ R d×1 , s θ (x) ∈ R d×1 , and Q(x ) = N (x ; x, σ 2 I d ).With Corollary 1 and a few assumptions, we have the following:EQUATION)The s θ , which represents a score network in our context, corresponds to h in Lemma 2 and Corollary 1.The derivation of Eq. ( 16) is presented in Appendix A, in which we assume the interchangeability between the expectation and summation regarding s θ (x ).Objective function of LCSS. We propose a variant of score matching method, local curvature smoothing with Stein's identity (LCSS). The development of the objective function of LCSS, J s LCSS , begins with the curvature smoothing regularization of Eq. ( 10), followed by the application of Lemma 1 and Corollary 2. Since J s LCS (θ, x, σ) in Eq. ( 10) involves computationally expensive ∇ x s θ (x), alongside the original challenge of Tr(∇ x s θ (x)) in J s SM , training with J s LCS is impractical for high-dimensional data. However, by inserting the transformation of Lemma 1, it enables the application of Corollary 2 to J s LCS . By substituting Eq. ( 4) into Eq. ( 11) and ignoring O( 2), we haveEQUATIONand by applying Eq. ( 16) to the first term, we obtain J s LCSS as:EQUATIONInJ s LCSS , Tr(∇ x s θ (x)) is replaced with the inner product, s θ (x ) T • x -x σ 2, which is computed efficiently, thereby bypassing the issue of high computational cost.Comparing LCSS with existing score matching methods. Unlike SSM and FD-SSM, LCSS does not use random projection, eliminating the high variance issue. While DSM learns the approximation of ground truth score ∇ x log q σ (x|x), LCSS learns the ground truth score ∇ x log p(x). Furthermore, unlike DSM, J s LCSS does not require ∇ x log q σ (x|x), thus eliminating the need for affine restrictions on the SDE coefficients. The original score matching, i.e., minimizing J s SM , involves the following two: (1) Increasing the first term Tr(∇ x s θ (x)) ≈ ∇ x • ∇ x log p(x), the divergence of the score, in the negative direction promotes s θ to learn the vector field flowing into points where x exists.(2) Minimizing the second term 1 2 s θ (x) 2 2 promotes s θ to learn that its length approaches 0 at points where x exists. The LCSS also performs both ( 1) and ( 2), but instead of at a single point x, it considers a Gaussian cloud centered around x. By applying Stein's identity, LCSS bypasses the challenge of (1), thereby making score matching feasible even for high-dimensional data.We define time-conditional version of LCSS for training SDMs as:EQUATIONand formulate the loss function of SDMs based on LCSS as:EQUATIONWe replace σ in Eq. ( 18) with a time-varying σ t . By making σ t take on a wide range of values depending on t, we aim to facilitate robust learning of score vectors even in low-density regions in p 0 , mirroring the original motivation of NCSN [31] . With Eq. ( 19), s θ learns a vector in the direction of -(x t -x 0 ) to minimize the inner product of the first term, weighted by 1σ 2 t, while minimizing its L 2 norm, s θ (x t , t) 2 . Sampling x t in the expectation in Eq. ( 19) only once yields satisfactory performance, as evidenced by our experimentation. SDEs for LCSS-based SDMs can be designed flexibly without restricting the drift and diffusion coefficients to be affine. However, devising a new SDE is beyond the scope of this paper and is left for future work, and our experiments use existing SDEs designed for use with DSM: the Variance Exploding (VE) SDE and the sub Variance Preserving (subVP) SDE [35] . Taking advantage of the fact that p t is a Gaussian distribution in both SDEs, we employ the standard deviation of p t as the value ot σ t in each SDE in our experiments. For example, for VE SDE, σ t = g(t). For both SDEs, σ t increases as t goes from 0 to T , but the way it increases is different for each SDE.Following Song and Ermon [31] , we set λ(t) = g(t) 2 . With this setting, λ(t) becomes λ(t) = g(t) 2 = σ 2 t , effectively cancelling out σ 2 t in the denominator of Eq. ( 19) and avoiding unstable situations where the denominator could become zero. For other SDE types (VP and sub VP), λ(t) is more elaborate but similarly cancels out σ 2 t in the denominator. For fairness, we note that, similarly, in training SDMs with DSM, applying the coefficient λ(t) = g(t) 2 allows for the cancellation of σ 2 t in the denominator, thus circumventing the weakness of DSM. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: The details of the experiments of observing SDM likelihood behavior are in the main text. Almost all settings are the same as Song et al. [35] or [31] . Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiments: 
 In this section, we demonstrate that our LCSS enables fast model training and high-quality image generation on several commonly used image datasets.We [35] . The same SDEs are applied to all the score matching methods we evaluate, including our LCSS. We use the official codes from the original papers, and the hyperparameters are kept as in the official code, unless stated otherwise. 4 For LCSS, we perform only one sampling iteration to calculate the expectation in J s LCSS (Eq. ( 19)). We set σ t = g(t) in each SDEs. All experiments are performed on a server with 128 GB RAM, 32 Intel Xeon, Silver 4316 CPUs, and eight NVIDIA A100 SXM GPUs.We evaluate the proposed LCSS against existing score matching methods, SSM, FD-SSM, and DSM, in density estimation, training efficiency, and qualitative and quantitative sample generation evaluation.We first compare LCSS to SSM, FD-SSM, and DSM in score matching performance. We visualize estimated densities on Checkerboard dataset, whose density is multi-modal. The details of the experiments, including the training loss curve, are presented in Appendix B. Table 1 depicts the density distribution learned by the model. Compared to SSM and FD-SSM, LCSS demonstrates higher accuracy in density estimation with faster convergence and stability in loss reduction. DSM exhibits similar accuracy in density estimation and stability in loss reduction to LCSS. However, LCSS shows slightly better consistency in estimating high-density regions (bright-colored areas) and maintains stable loss.We compare LCSS with the existing score matching methods for training efficiency. We measure the time taken for model training on Checkerboard and FFHQ dataset [11] resized to 256 × 256. Table 2 shows the average elapsed time over 100 epochs for Checkerboard and 1000 iterations for FFHQ, respectively. It shows that LCSS is the most efficient.We show generated samples on CIFAR-10 using NCSN++ deep and DDPM++ deep trained with LCSS in Appendix C.1. In this subsection, we qualitatively compare the sample generation capability of LCSS with existing methods. Comparison with SSM and FD-SSM. We first focus on comparing LCSS with SSM and FDSSM. 5We generate samples using NCSNv2 on CIFAR-10 We compare generated samples on FFHQ, AFHQ, and FFHQ + AFHQ. The size of images in the three datasets is (256×256), and we train NCSNv2 for 600k with batch size 16 on each of them. On FFHQ, LCSS can generate more realistic images than DSM, as shown in Fig. 4 . We note that during the training with DSM, around 210k training steps, a sharp decline in the quality of generated images was observed. On AFHQ [4] , Fig. 5 shows that LCSS generates realistic samples, but DSM does not. We also create and examine with a dataset FFHQ + AFHQ, a fusion of FFHQ and AFHQ, designed to increase learning difficulty by diversifying data modalities. On FFHQ + AFHQ, Fig. 6 shows LCSS's superior capability in generating realistic images over DSM. Table 3 shows the qualitative results on CIFAR-10. Regardless of SDMs, LCSS tends to surpass DSM in IS but underperform in FID. Compared to the values in Song et al. [35] , our experimental results generally exhibit higher (better) IS values and higher (worse) FID values. 6 In BPD, LCSS surpasses DSM in DDPM++ variants but underperforms in NCSN++ variants. Overall, qualitative evaluation on CIFAR-10 suggests no decisive superiority between LCSS and DSM, hinting at distinct characteristics.(Summary. Table 4 illustrates a highly simplified comparison of the relative performance between LCSS and DSM. Model training is more complex in Case #2 than in Case #1. It was observed that in challenging conditions like Case #2, DSM suffered from performance degradation. We regularly monitored the quality of generated images during model training. In the experiments of Case #2 with DSM, as noted above, although the quality of generated images was improving up to a certain stage (around 210k iterations, for example), it suddenly deteriorated. Also, frequent spikes in loss values were observed during training with DSM, which appeared to be a trigger for the deterioration. Unlike DSM, LCSS retained superior performance without suffering from such instability.We demonstrate that learning with LCSS enables models to generate high-resolution images. We train NCSN++ and DDPM++ on CelebA HQ (1024 × 1024) [10] , using hyperparameters consistent with those used to train DSM-based models in Song et al. [35] . In Fig. 1 , we show generated images: NCSN++ images are from the model trained for around 1.3M iterations, and DDPM++ ones are trained for around 0.3M iterations, with batch size 16 for both. The figures show that LCSS is promising as a score matching method. More generated samples are presented in Appendix C.2.The loss function of LCSS, similar to the original score matching, consists of two terms. To study the roles of each term, using the modified version of J s LCSS in Eq. ( 19) with a balancing coefficient γ asEQUATIONwe train NCSNv2 on FFHQ. Images generated from the models trained with different γ are shown in Fig. 7 . When γ = 0.5, only noisy images akin to those at time t = T , x T , are produced. With γ < 1, the force to minimize the second term, s θ (x t , t)2 2, is more emphasized than when using the original J s LCSS , leading to shorter score vector lengths. The score vector length is directly linked to the spatial movement distance of x t during the reverse process for sample generation. Since the score vector is forced to be short, the noise x T generated at the start of the sample generation process cannot reach the regions corresponding to realistic images with high density as it traces back from time T to 0. On the other hand, when γ > 1, particularly for γ = 10, it is observed that as the number of training iterations increased, images with emphasized contours but lost textures are generated. It suggests that the involvement of the first term of J s LCSS in contour formation. The object contours in an image are characterized by rapid changes in pixel values, which can be associated with high curvature or changes in second-order derivatives. Sinces θ (x t , t) T • x t -xt σ 2in Eq. ( 21) or ( 19) corresponds to the Hessian trace of log-density, this observation can be interpreted as natural. It is implied that the first and second terms in the loss function of LCSS are dedicated to the formation of contours and texture, respectively. 
conclusion: 
 Limitation. While LCSS, unlike DSM, can design SDEs flexibly without restricting them to affine forms, we used existing affine SDEs designed for use together with DSM, i.e., VE SDE and subVP SDE, in this work. Proposing more flexible SDEs leveraging LCSS is left for future work.We proposed a local curvature smoothing with Stein's identity (LCSS), a regularized score matching method expressed in a simple form, enabling fast computation. We demonstrated LCSS's effectiveness in training on high-dimensional data and showed that LCSS-based SDMs enable high-resolution image generation. Currently, SDMs primarily rely on DSM, constraining the design of SDE. LCSS offers an alternative to DSM, opening avenues for SDM research based on more flexible SDEs. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: For the qualitative evaluation using FID, IS, and BPD, we calculated the mean over all samples in the CIFAR-10 test set, and in line with common practice, error bars are not shown, as in Song et al. [35] , for example. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We presented it in the second paragraph of Section 4.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: Our proposal enables the development of score-based generative models based on nonlinear SDEs, potentially significantly improving generation quality over existing methods. As noted in the guidelines below, this could lead to the misuse in creating more sophisticated deepfakes.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: -Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We provided proper citations for the code, trained models, and datasets used in our experiments.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA]Justification: -Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: -Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: -Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: To derive Corollary 2, we first introduce the following two assumptions. Assumption 1. Let x be a real-valued vector, x = [x 1 , . . . , x d ] T . Considering the derivative of a function s θ : R d → R d , we assume that the expectation and summation are interchangeable as follows:EQUATION)Assumption 2. For the same function s θ in Assumption 1, we assume the following interchangeability:EQUATION)The interchangeability holds when a (score) function s θ is integrable and differentiable. As s θ is implemented by neural networks in our case, we assume these conditions are met. We noted that if there are correlations between the dimensions of x over which expectations are taken, interchangeability does not always hold. However, because x is a sample from a diagonal Gaussian, x ∼ N (x, σ 2 I d ), there are no correlations between dimensions, thus fulfilling this condition.Derivation of Corollary 2. Using these assumptions, Eq. ( 16) is derived as follows:EQUATION)Assumption 1 is applied in transitioning from Eq. ( 24) to Eq. ( 25). The transition from Eq. ( 25) to Eq. ( 26) is achieved by applying the chain rule, with x i = x i + σ where ∼ N (0, 1), due to x i ∼ N (x i , σ 2 ), and thus we have 27). We use Corollary 1 in the transition from Eq. ( 27) to Eq. ( 28). We finally apply Assumption 2 in the transition from Eq. ( 28) to Eq. ( 29).∂x i ∂xi = 1 in Eq. (We describe the setup for the experiments on Checkerboard dataset. We use the publicly available code 7 . The generation of the Checkerboard dataset can be found, for example, in Appendix D.1 in Lai et al. [17] . In the data space X ∈ R 2 , we train a function f θ : X → R parameterized by θ to estimate density directly. The score, calculated as ∇ x f θ , corresponds to s θ in the main text. We apply the score matching methods we examine to estimate ∇ x f θ , resulting in the trained density estimation function f θ . Sampling is performed using Langevin dynamics based on ∇ x f θ as:x t-1 = x t + 2 ∇ x f θ (x t ) + z (31)where z ∼ N (0, I 2 ) and = 0.1. The function f θ is implemented as a simple multilayer perceptron (MLP) with two hidden layers, each with 300 units, which is the same architecture as the one used in Song and Ermon [31] . The number of Langevin dynamics steps is 1000, with the initial vector, x 1000 , being randomly sampled from a uniform distribution. We use stochastic gradient descent with a batch size of 10k, a learning rate of 1e -3, and train for 200 epochs. Fig. 1 shows the sampling results of 250M points. The loss curve during training on Checkerboard is shown in Fig. 8 . Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Justification: We read and confirmed that this study is in compliance with it.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).github.com/taufikxu/FD-ScoreMatching/github.com/Ending2015a/toy_gradlogpgithub.com/yang-song/score_sde_pytorchAlthough FD-DSM has also been proposed, it was excluded from comparative evaluation due to reported performance below DSM in Pang et al.[23] and failure to generate images appropriately in our experiments.Although we use the official code from Song et al.[35] in our experiments, the difference between JAX version in Song et al.[35] and PyTorch version in our experiments is considered to be the cause.github.com/Ending2015a/toy_gradlogp 


Decision: This paper introduces a novel score matching variant called Local Curvature Smoothing with Stein’s Identity (LCSS), designed to address the computational challenges associated with calculating the Jacobian trace in score matching, particularly in high-dimensional settings. By leveraging Stein’s identity, LCSS bypasses the need for direct Jacobian evaluations, offering both computational efficiency and regularization benefits. Reviewers generally found the idea to be innovative and the method effective, as demonstrated through experiments on both synthetic and real datasets. However, concerns were raised about the paper's clarity, particularly regarding the advantages of LCSS over existing methods like Denoising Score Matching (DSM). Additionally, the paper could benefit from more detailed explanations of certain assumptions and a deeper comparison with methods that directly compute the Jacobian. Despite these issues, the contribution is deemed valuable, with promising potential for future exploration. I recommend accepting the paper.

Reviews: review: paper_summary: paper_summary: This manuscript proposes a new score matching method that bypasses the Jacobian trace by applying Stein’s identity, enabling effective regularization and efficient computation.
strength_weakness: strength_weakness: 1. The advantage of the proposed method compared to denoising score matching (DSM) is unclear. The manuscript mentions that it restricts the SDE to be affine, but it does not clarify the benefit of using a non-affine SDE. Furthermore, the influence of the SDE on the generative model needs to be elaborated.
2. The experimental results do not show significant improvements over DSM. The proposed method achieves comparable sample quality, as shown in Table 3.
questions: questions: Please refer to weaknesses.
limitations: limitations: N/A
paper_summary: paper_summary: The paper proposes to use Stein's lemma to obtain a computationally efficient way in implementing a local-curvature regularized variant of the score matching objective. The main idea is to rewrite the Jacobian-trace term in a way that requires no Jacobian evaluations. In numerical experiments, the effectiveness of this approach is clearly demonstrated.
strength_weakness: strength_weakness: - Approximating Jacobian trace through Stein's identity potentially leads to an estimator with large variance -- I found the claims that it solves Hutchinson's high variance problem to be a bit misleading.
questions: questions: Can there be a formal argument that the proposed estimator has lower variance than random projections? Essentially, the gradient is estimated through random (zero-order) sampling, which is not exactly low-variance?
limitations: limitations: All limitations are addressed.
paper_summary: paper_summary: This paper provides a new way for score matching with the purpose of resolving some of the limitations of the existing methods such as high variance of sliced score matching and Gaussian constraints of denoising score matching (DSM). The new method is based on the local curvature smoothing proposed in [15]. A new score matching objective function is proposed by combining the Stein's Identity with the local curvature smoothing. The authors empirically show that the new method is more efficient in training than DSM and also has comparable performance to DSM.
strength_weakness: strength_weakness: I think some parts of the paper are not stated clearly and further clarification is needed. See questions for more details.
questions: questions: - In section 2.4, the authors criticize the DSM method for having a Gaussian constraint. However, later there is no clarification showing how the new method is different from DSM in this regard. Can you please clarify this?
- In line 108, the authors criticize the DSM for having 0 numerator and denominator. However, in the final LCSS (equation (16)), the denominator can also be 0 and be problematic. Can the authors provide more discussion on why the new method is better in this regard?
- In Corollary 2, there is an assumption that an integral must be 0. How restrictive is this assumption? It seems to me that later on when designing LCSS objective, formula (14) is directly used without any further discussion on this assumption. Can the authors explain why this assumption can be dropped?
limitations: limitations: N/A
paper_summary: paper_summary: The paper proposes a novel score matching variant called Local Curvature Smoothing with Stein’s Identity (LCSS). This method addresses the computational challenges associated with the Jacobian trace in score matching, particularly for high-dimensional data, by leveraging Stein’s identity. LCSS aims to bypass the expensive computation of the Jacobian trace, offering both regularization benefits and efficient computation. The method is validated through experiments on synthetic and real datasets.
strength_weakness: strength_weakness: 1. In lines 161-162, interchangeability is assumed. However, in the analysis, interchangeability requires some properties of the interested function. The reason why the assumption holds is missing.
2. This paper does not approximate the Jacobian but instead circumvents the Jacobian. The empirical and theoretical differences against the method using Jacobian should be discussed, such as the difference in the estimated error bound. 
3. In Tab. 3, the improvement seems to be marginal, while in figures, such as Fig. 4, the selected picture is much better under LCSS. The discrepancy should be discussed.
questions: questions: see weakness
limitations: limitations: Yes

Comments: title: Response to Reviewer j6aU (2)
Response to Reviewer j6aU
comment: 1. Please make sure to include an explanation somewhere in Section 3. Right now, the whole non-affine thing is missing in this section.
2. That makes sense.
3. This is good. 

Thanks the authors for the clarification and I have raised my score.
Thanks for the interesting reply. These experiments are inspiring.
### Response to 1
Thanks to this question, we have understood the insight behind the original 3rd question.
The original score matching (SM) proposed in [7] does not approximate the Jacobian.

In [8], the loss between SM and SSM-VR (= SSM in our paper) for low-dimensional table datasets (dimensionality is 11 ~ 22) is compared, demonstrating that the SSM approximation is nearly equivalent to the original SM. 
The experimental results in our paper, along with the error-bound inequality above in response to Q2, demonstrate that LCSS approximates better than SSM, indirectly indicating a minor discrepancy between LCSS and SM.

After receiving this question, we conducted an experiment on the Checkerboard dataset akin to those in the paper.
The necessity to drastically reduce batch-size to avoid out-of-memory errors during Jacobian computation probably led to the unsuccessful density estimation, hence no comparison of performance could be made.
Instead, we compare the computational efficiencies.
The training time (sec.) measured over 20k iterations with a batch size of 10,000 are presented below.

| SSM  |   DSM   |  LCSS   | SM (no approx.)  | 
| :--------: |:--------------:|:-----------:  |:--------: |
|  26.13    |    22.59      |  21.75      |  645.61    |

It shows that even on a mere 2-dimension Checkerboard dataset, the computational cost is about 25 times greater than other methods. 
This underscores that, in score matching training for high-dimensional data, such as images with dimensions up to several hundreds of thousands, Jacobian approximation for acceleration is indispensable in today’s typical computational environments.

- [7] Hyvärinen, A., & Dayan, P. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4).
- [8] Song, Y., Garg, S., Shi, J., & Ermon, S. (2020, August). Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence (pp. 574-584). PMLR.
---
### Response to 2
We agree with the reviewer's suggestion. Through this rebuttal, we have recognized that presenting the differences with existing methods in a tabular comparison will highlight our contributions.

---
Lastly, we deeply appreciate the reviewer for dedicating time to the discussion and for the inquiries to clarify the ambiguities.
Dear reviewers,

Could you please respond to the rebuttal, discuss with the authors and finalize your score?
Thanks for the reply:
1 is interesting to read. 
2 is okay.

1. Perhaps I have missed something. Is there any work involving the accurate Jacobian? If so, how much of your methods approximate the accurate one empirically? This is the key concern in my previous 3rd question.

2. I think it's good for readability to involve a table to present the difference between your method and the existing one. This can highlight your contribution. Also, the summary table you present is good for presenting the paper.

Anyway, Although I am not an expert in this field, I think this paper is worth reading. I have raised my score to 6.
We appreciate the reviewer's questions to ensure accurate understanding and hope the responses below will contribute to the reviewer's clarity.

---
## Regarding Q.3
The table below illustrates a highly simplified comparison of the relative performance between the performance of LCSS (ours) and DSM.

| Case  | Model  |   Resolution   |  Results   |  LCSS  |  DSM   |
| :-----: | :--------: |:--------------:|:-----------:  |:--------: |:--------:|
| #1     |  Large    |    32 x 32      |  Table 3      |  good    | good    |
| #2     |  Small    |   256 x 256   |  Figs. 4-6   |  good    | poor     |

Under the stricter conditions of Case #2 compared to Case #1, LCSS performance remained stable, whereas DSM performance degraded.
While LCSS, SSM, and FD-SSM approximate the Jacobian in their respective ways, DSM does not approximate it but circumvents it by replacing the true score, $\nabla _ {\bf x} \log p({\bf x})$, with the score of the perturbed distribution 
corrupted by Gaussian noise, $\nabla _ {\bf x} \log q _ {\sigma} (\tilde{{\bf x}}|{\bf x})$, as the learning target.
The drawbacks of DSM noted in Lines 104-108 are all caused by this replacement.

We regularly monitored the quality of generated images during model training.
In the experiments of Case #2, as noted in lines 256-258, 
although the quality of generated images was improving up to a certain stage (around 210k iterations, for example), it suddenly deteriorated.
Frequent spikes in loss values were observed during training, which appeared to be a trigger for the deterioration. 
Although the exact cause was not precisely identified,
we attribute this phenomenon to the replacement by $\nabla _ {\bf x} \log q _ {\sigma} (\tilde{{\bf x}}|{\bf x})$, as it was not observed in other score matching methods.
We argue that the instability of DSM becomes apparent under stringent training conditions, such as those in Case #2. 

This is our response, but to ensure clarity for the reviewers, please request clarification if any uncertainties remain.

---
## Regarding non-affine SDE
Designing non-affine SDE demands an in-depth understanding of SDE, and we are in the process of developing that.
As such, we are not currently ready for conducting empirical comparisons, and let us leave proposing non-linear SDE leveraging LCSS for future work.
I thank the authors for further clarifications.

1-2 is good for me.

3 is kind of mysterious to me. Perhaps I have missed something. In some sense, I think giving up Jacobian should be considered as an approximation. If so, when model capacity is high, it should reach comparable performance, while when capacity is small, approximation causes more deviation. This is different from your findings, why?

I noticed that Reviewer AvFd also mentioned non-affine SDE. Is it possible to do an empirical comparison?