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

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$.

Paper

Similar papers

Peer review

Reviewer j6aU6/10 · confidence 3/52024-06-18

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.

Strengths

1. the idea of LCSS is novel 2. Jacobian is not computed directly, but implicitly respected. 3. Experiments on high and low resolution are performed.

Weaknesses

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

see weakness

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Reviewer j6aU2024-08-09

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.

Authorsrebuttal2024-08-11

Response to Reviewer j6aU (2)

### 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.

Reviewer j6aU2024-08-11

Thanks for the interesting reply. These experiments are inspiring.

Reviewer qUJ66/10 · confidence 3/52024-07-11

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.

Strengths

Although DSM is the default method used nowadays for score matching, the authors provide a nice novel alternative which may have some advantages over DSM. I'm interested to see more theoretical study in the future of this new method.

Weaknesses

I think some parts of the paper are not stated clearly and further clarification is needed. See questions for more details.

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Reviewer tfBV6/10 · confidence 2/52024-07-12

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.

Strengths

- The paper is well-written and the main idea is clear and easy to understand. - Other works which the paper builds upon are referenced and fairly attributed. - Experiments on small-scale data clearly demonstrate the effectiveness of the approach. - Also on larger datasets, the method appears to give strong empirical results.

Weaknesses

- 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

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?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

All limitations are addressed.

Reviewer AvFd6/10 · confidence 3/52024-07-25

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.

Strengths

1. The method is computationally efficient compared to other SSM variants. 2. Experimental results demonstrate the effectiveness of the proposed method.

Weaknesses

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

Please refer to weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

N/A

Reviewer j6aU2024-08-07

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?

Authorsrebuttal2024-08-09

Response to Reviewer j6aU

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.

Area Chair yk5F2024-08-10

Dear reviewers, Could you please respond to the rebuttal, discuss with the authors and finalize your score?

Reviewer qUJ62024-08-11

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.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC