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

We appreciate the reviewer for thoroughly reading our paper and asking important questions, which we believe will clarify the contributions of our work.

Response to Q.1

Interchangeability holds when the score function $S_{\theta}$ is both integrable and differentiable. $S_{\theta}$ is implemented by a neural network, and 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' \sim \mathcal{N}(x, \sigma^2 \mathbb{I} _ {d})$, in our case, there are no correlations between dimensions, thus fulfilling this condition. We plan to include an extra sentence in the camera-ready version for better reader comprehension.

Response to Q.2

The loss curve in Fig. 8 (Appendix B) empirically demonstrates that our LCSS has a lower variance than SSM (and even DSM).  The theoretical explanation is provided below, (essentially identical to the response given to Reviewer tfBV's question).

We replaced the Jacobian computation with Stain's identity, approximating the expectation computation with a single sampling. Compared to Hutchinson's trick, which approximates the Jacobian by random projection in the existing methods (SSM and FD-SSM), we show that our approximation error is smaller with high probability when the variance of Gaussian, $\sigma^{2}$, is small.

2-1. Error in single sampling approximation

Let $x$ be a $d$-dimension vector, $S(x)$ be a $L$-Lipschitz (score) function, $S: \mathbb{R}^{d} \rightarrow \mathbb{R}$. We approximate $M := \mathbb{E} _ { x' \sim N(x, \sigma^{2} I_d) } \left[ \mathcal{J} _ \text{SM}^{s} (\theta, x') \right]$ by its single sampling, $M' := \mathcal{J} _ \text{SM}^{s} (\theta, x')$. Then, Chernoff bound for Gaussian variable tells us that $\Pr[ |M - M'| \geq \delta] \leq 2 \exp \left(- \frac{\delta^{2}}{2 L^{2} \sigma^{2}}\right), \forall \delta \geq 0.$ Letting $p = 2 \exp \left(- \frac{\delta^{2}}{2 L^{2} \sigma^{2}} \right)$, we obtain $|M - M'| \leq \delta = \sqrt{2 L^{2} \sigma^{2} \log \left(\frac{2}{p}\right)}$ with probability at least $1 - p$.

(See Thm. 2.4 in [5], for example.)

• [5] Wainwright, M. (2015). "Mathematical Statistics: Chapter 2 Basic tail and concentration bounds."

2-2. Upper bound of Hutchinson's trick error

We denote the Jacobian matrix of $S(x)$, $\nabla _ {x} S(x)$, by $A$. The error between the true trace of $A$, $\text{Tr}(A)$, and the estimate by Hutchinson's trick, $\tilde{T}$, is bounded as $|\text{Tr}(A) - \tilde{T}| \leq |A| _ {F}$ where $|\cdot|_{F}$ is the Frobenius norm. The upper bound of Hutchinson's trick error is thus $|A| _ {F}$.

(See line 94-96 of our paper.)

2-3. Comparison

We compare the upper bound of $|M - M'|$ to $|A _ {F}|$. By substituting $|A _ {F}|$ into $\delta$ in the above, we know that $|M - M'| \leq |A _ {F}|$ holds with probability at least $q := 1 - p = 1 - 2 \exp \left( - \frac{|A| _ {F}^{2}}{2 L^{2} \sigma^{2}} \right)$. In the case of small $\sigma$ (a region particularly crucial for SDM training), $q$ is nearly 1, indicating that a single sampling approximation of the expected value on a Gaussian distribution almost consistently yields smaller errors compared to Hutchinson's trick.

Response to Q.3

Evaluation in Table 3 was conducted using large models like DDPM++ at a low resolution of $32 \times 32$, where LCSS and DSM showed comparable performance. In contrast, the experiments in Figs. 4-6 were conducted with the smaller NCSNv2 model at $256 \times 256$ resolution, presented more challenging conditions. The qualitative evaluation of Figs. 4-6 indicates that LCSS significantly outperforms DSM under such severe conditions. The results demonstrate that LCSS operates stably even when the model capacity is small.

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.

The detailed questions from the reviewer reflect her/his careful reading of our paper. We are grateful for the constructive questions and hope our responses address their concerns.

Response to Q.1

The loss function of DSM includes $\nabla _ {x} \log q _ {\sigma _ {t}} ( \tilde{x} _ {t} | x _ {0} )$. To compute it for any $0 \leq t \leq T$, $\nabla _ {x} \log q _ {\sigma _ {t}} ( \tilde{x} _ {t} | x _ {0} )$ must be Gaussian, necessitating a Gaussian prior and constraining the SDE to be linear. In constrast, our LCSS does not include $\nabla _ {x} \log q _ {\sigma _ {t}} ( \tilde{x} _ {t} | x _ {0} )$ in the loss function, allowing for a non-Gaussian prior and the design of nonlinear SDEs.

We note the following two:

• The LCSS loss function involves sampling from a Gaussian, $x' \sim \mathcal{N}(x, \sigma^2 \mathbb{I} _ {d})$, but this stems from curvature smoothing regularization and does not relate to or constrain the prior distribution form or SDE design.
• Like our LCSS, SSM and FD-SSM do not impose Gaussian constraints, yet, as shown in the paper, SSM and FD-SSM fail to learn effectively at resolutions larger than $32 \times 32$.

Response to Q.2

Actual model training of SDMs with our method minimizes Eq.(18), which incorporates Eq.(16). In Eq.(18), the coefficient $\lambda(t)=g^2(t)$ and for the VE SDE, $\lambda(t)=g^2(t)=\sigma_{t}^2$, effectively canceling out $\sigma_{t}^2$ in the denominator of Eq. (18) and avoiding unstable situations where the denominator could become zero. For other SDE types (VP and sub VP), $\lambda(t)$ is more elaborate but similarly cancels out $\sigma_{t}^2$ in the denominator.

We note that, similarly, in training SDMs with DSM, applying the coefficient $\lambda(t)=g^2(t)$ allows for the cancellation of $\sigma_{t}^2$ in the denominator, thus circumventing the weakness of DSM the authors mentioned in line 108. For fairness, we plan to add this point to the camera-ready version.

Response to Q.3

The assumption in Corollary 2 holds almost surely. The assumption can be confirmed if $\lim _ {\lVert x \rVert \rightarrow \infty} s_{\theta_{i}}(x) Q(x) = 0$, as noted in [6]. Because $Q$ is defined as a probability density function of Gaussian in Eq. (14), $Q(x)$ tending towards zero at $x=\pm \infty$, this condition is satisfied. For clarity, let us also include this explanation in the camera-ready version.

• [6] Liu, Q., Lee, J., & Jordan, M. A kernelized Stein discrepancy for goodness-of-fit tests. (ICML 2016)

We appreciate the reviewer's meaningful question; addressing it will clarify our paper's contributions. For reader comprehension, we plan to include the following argument in the camera-ready version (maybe in the appendix).

Question: a formal argument that the proposed estimator has lower variance than random projections?

The expectation in Eq.(16) is approximated by a single sampling, leading to an error, as the reviewer asks. However, the key lies in taking the expectation over Gaussian samples, resulting in sufficiently small errors. Below, we show the comparison between the error in a single sampling approximation of Stain's identiy with the error caused in random projections (Hutchinson's trick error).

1. Error in single sampling approximation

Let $x$ be a $d$-dimension vector, $S(x)$ be a $L$-Lipschitz (score) function, $S: \mathbb{R}^{d} \rightarrow \mathbb{R}$. In our implementation, we approximate $M := \mathbb{E} _ { x' \sim N(x, \sigma^{2} I_d) } \left[ \mathcal{J} _ \text{SM}^{s} (\theta, x') \right]$ by its single sampling, $M' := \mathcal{J} _ \text{SM}^{s} (\theta, x')$. Then, Chernoff bound for Gaussian variable tells us that $\Pr[ |M - M'| \geq \delta] \leq 2 \exp \left(- \frac{\delta^{2}}{2 L^{2} \sigma^{2}}\right), \forall \delta \geq 0.$ Letting $p = 2 \exp \left(- \frac{\delta^{2}}{2 L^{2} \sigma^{2}} \right)$, we obtain $|M - M'| \leq \delta = \sqrt{2 L^{2} \sigma^{2} \log \left(\frac{2}{p}\right)}$ with probability at least $1 - p$.

(See Thm. 2.4 in [5], for example.)

• [5] Wainwright, M. (2015). "Mathematical Statistics: Chapter 2 Basic tail and concentration bounds."

2. Upper bound of Hutchinson's trick error

We denote the Jacobian matrix of $S(x)$, $\nabla _ {x} S(x)$, by $A$. The error between the true trace of $A$, $\text{Tr}(A)$, and the estimate by Hutchinson's trick, $\tilde{T}$, is bounded as $|\text{Tr}(A) - \tilde{T}| \leq |A| _ {F}$ where $|\cdot|_{F}$ is the Frobenius norm. The upper bound of Hutchinson's trick error is thus $|A| _ {F}$.

(See line 94-96 of our paper.)

3. Comparison

We compare the upper bound of $|M - M'|$ to $|A _ {F}|$. By substituting $|A _ {F}|$ into $\delta$ in the above, we know that $|M - M'| \leq |A _ {F}|$ holds with probability at least $q := 1 - p = 1 - 2 \exp \left( - \frac{|A| _ {F}^{2}}{2 L^{2} \sigma^{2}} \right)$. In the case of small $\sigma$ (a region particularly crucial for SDM training), $q$ is nearly 1, indicating that a single sampling approximation of the expected value on a Gaussian distribution almost consistently yields smaller errors compared to Hutchinson's trick.

We note that empirically, the stability of the loss curve in Fig. 8 in Appendix B indicates that the proposed LCSS has lower variance than SSM (random projection) and even DSM.

We appreciate the reviewer taking the time to read our paper thoroughly. We hope the following responses clarify our contribution. Additionally, our response to Reviewer tfBV below provides an argument about the advantage over SSM and FD-SSM, which we would like the reviewer to examine.

Response to Q.1

The design of SDE directly influences the performance of score-based diffusion models, as demonstrated in [1][2]. The benefits of non-linear SDE, particularly highlighted in [3], 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 [3] illustrates this.)

Our LCSS allows for the engineering of non-linear SDEs with DSM-equivalent performance. This paper does not cover the creation of new non-linear SDE based on LCSS, which is noted as a limitation, but we focus on describing LCSS itself as a novel score-matching method in this work.

• [1] Dockhorn, T., Vahdat, A., & Kreis, K. Score-based generative modeling with critically-damped langevin diffusion. (ICLR 2022)
• [2] Karras, T., Aittala, M., Aila, T., & Laine, S. Elucidating the design space of diffusion-based generative models. (NeurIPS 2022)
• [3] Kim, D., Na, B., Kwon, S. J., Lee, D., Kang, W., & Moon, I. C. Maximum likelihood training of implicit nonlinear diffusion model. (NeurIPS 2022)

Response to Q.2

The evaluations in Table 3 utilize large models, such as DDPM++, at a low resolution of $32 \times 32$, where the results show that LCSS performs comparably to DSM, as the reviewer stated. However, qualitative evaluations presented in Figs. 4-6 conducted with the smaller model NCSNv2 at $256 \times 256$ resolution demonstrate noticeable differences between LCSS and DSM; DSM exhibits poor image generation performance under these conditions, as suggested in [4], whereas LCSS can still produce high-quality images. The stable performance of LCSS with limited model capacity is its advantage over DSM.

• [4] Song, Y., & Ermon, S. (2020). Improved techniques for training score-based generative models. (NeurIPS 2020)