The geometry of Stein's method of moments: A canonical decomposition via score matching

1. Problem Statement

The paper addresses the challenge of parameter estimation for non-normalized statistical models (also known as energy-based models). These models have a probability density of the form:
$$q_\theta(x) = \frac{1}{Z(\theta)} \tilde{q}_\theta(x)$$
where $Z(\theta)$ is a computationally intractable normalizing constant.

• The Limitation of MLE: The Maximum Likelihood Estimator (MLE) requires calculating $Z(\theta)$, making it infeasible for many complex models.
• Existing Solutions:
  • Score Matching (SM): Introduced by Hyvärinen (2005), this method minimizes the distance between the model's score function ($\nabla_x \log q_\theta$) and the true score, avoiding $Z(\theta)$. However, its statistical efficiency (asymptotic variance) relative to the MLE is not fully characterized, and it is not always asymptotically efficient.
  • Stein's Method of Moments (SMoM): A recent framework (Ebner et al., 2025) using Stein operators to construct estimators that do not depend on $Z(\theta)$. While flexible, the geometric relationship between SMoM, Score Matching, and the MLE remains unclear.

Core Question: Can we geometrically characterize the relationship between Score Matching and SMoM, and construct an SMoM estimator that improves upon the asymptotic variance of the Score Matching estimator?

2. Methodology and Framework

The authors develop a geometric framework based on Stein operators and orthogonality concepts to analyze and improve estimators.

A. Stein's Method of Moments (SMoM)

The SMoM estimator $\hat{\theta}_{SMoM}$ is defined by solving estimating equations using a Stein operator $A_\theta$:
$$\frac{1}{n} \sum_{i=1}^n A_\theta f_{\theta, j}(X_i) = 0, \quad j=1,\dots,d$$
where $f_{\theta, j}$ are test functions. The authors utilize the divergence-based Stein operator:
$$A_\theta f_\theta := \frac{\nabla_x \cdot (q_\theta f_\theta)}{q_\theta}$$
Crucially, Lemma 1 establishes that the Score Matching estimator is a specific instance of an SMoM estimator where the test functions are chosen as $f_{\theta, j} = \nabla_x \partial_{\theta_j} \log q_\theta$.

B. Canonical Decomposition (Theorem 1)

The authors derive an asymptotic linear representation for any SMoM estimator by centering it around the Score Matching estimator. They introduce two types of orthogonality:

1. W-orthogonality: Based on the standard inner product in $L^2(q_\theta)$, i.e., $E[\langle f, g \rangle] = 0$. This relates to the Wasserstein geometry.
2. $A_\theta$-orthogonality: Based on the Stein operator, i.e., $E[(A_\theta f)(A_\theta g)] = 0$.

Key Insight: The authors prove that any SMoM estimator can be decomposed as:
$$\hat{\theta}_{SMoM} - \theta^* = (\hat{\theta}_{SM} - \theta^*) - G^{-1} \frac{1}{n} \sum_{i=1}^n \begin{pmatrix} A_{\theta^*} u^*_1(X_i) \\ \vdots \\ A_{\theta^*} u^*_d(X_i) \end{pmatrix} + o_p(n^{-1/2})$$
where $u^*_j$ are W-orthogonal to the Score Matching test functions.

• Significance: The Score Matching estimator is the "base" SMoM estimator with no W-orthogonal correction terms. The deviation of any other SMoM estimator from Score Matching is entirely determined by these W-orthogonal terms.

C. Improving Asymptotic Variance (Theorem 2)

Since W-orthogonality does not imply $A_\theta$-orthogonality, the authors construct a new estimator $\hat{\theta}[\theta^*]$ by adding specific W-orthogonal terms to the Score Matching test functions.

• They define a correction term using auxiliary functions $\tilde{v}_\alpha$ (e.g., neural networks) orthogonalized against the Score Matching gradients.
• Result: The asymptotic variance of the new estimator is strictly smaller than or equal to that of the Score Matching estimator:
  $$AVar[\hat{\theta}_{SM}] - AVar[\hat{\theta}[\theta^*]] = G^{-1} S_{\theta^*} T_{\theta^*}^{-1} S_{\theta^*}^\top G^{-1} \succeq 0$$
  Equality holds only if the correction term is zero.

3. Key Contributions

A. Geometric Interpretation via Wasserstein Geometry

The paper establishes a profound connection between SMoM and Wasserstein geometry.

• The Wasserstein score function $\Phi_{\theta, j}$ is defined as the solution to $A_\theta (\nabla_x \Phi_{\theta, j}) = -\partial_{\theta_j} \log q_\theta$.
• The MLE corresponds to the SMoM estimator using test functions $f_{\theta, j} = -\nabla_x \Phi_{\theta, j}$.
• The gap in asymptotic variance between Score Matching and MLE is characterized by the W-orthogonal component of the Wasserstein score function relative to the Fisher score function.

B. Characterization of Asymptotic Efficiency (Theorem 3)

The authors provide a necessary and sufficient condition for the Score Matching estimator to be asymptotically efficient (i.e., achieve the Cramér-Rao bound of the MLE):

The Score Matching estimator is asymptotically efficient IF AND ONLY IF the space spanned by the Fisher score functions ($\nabla_x \partial_{\theta_j} \log q_\theta$) is identical to the space spanned by the Wasserstein score functions ($\nabla_x \Phi_{\theta, j}$).

This explains why Score Matching is efficient for Gaussian distributions (where the spaces coincide) but inefficient for others (like Generalized Normal distributions with high shape parameters).

C. Practical Construction of Improved Estimators

The paper proposes a practical algorithm to improve Score Matching:

1. Choose a set of flexible functions (e.g., neural networks) $\tilde{v}_\alpha$.
2. Orthogonalize them against the Score Matching gradients (W-orthogonalization).
3. Construct the SMoM estimator using these corrected test functions.
4. As the number of correction terms $K$ increases, the estimator approaches the efficiency of the MLE.

4. Results and Numerical Experiments

The authors validate their theory through numerical experiments on two models:

1. Generalized Normal Distribution:

   • For $\beta=1$ (Gaussian), Score Matching is efficient.
   • For $\beta > 1$, Score Matching is inefficient.
   • Result: The proposed SMoM estimator ($\hat{\theta}[\hat{\theta}_{SM}]$) significantly reduces the Mean Squared Error (MSE) compared to standard Score Matching, approaching the performance of the MLE as the number of correction terms $K$ increases.

2. Polynomially Tilted Pairwise Interaction (PPI) Model:

   • A model on the unit sphere (truncated domain).
   • Result: The improved SMoM estimator outperforms the weighted Score Matching estimator for moderate $K$ (e.g., $K=12$). However, for very large $K$ (e.g., $K=24$), numerical instability arises, highlighting a trade-off between variance reduction and estimation stability.

3. Matrix Bingham Distribution:

   • A model on the Stiefel manifold.
   • Result: The Score Matching estimator was already nearly efficient (MSE ratio $\approx 1$), and the proposed improvement yielded negligible gains, confirming the theoretical condition that efficiency is achieved when the score spaces align.

5. Significance and Conclusion

• Theoretical Unification: The paper unifies Score Matching, SMoM, and Wasserstein geometry, revealing that Score Matching is a specific projection in a larger geometric space.
• Efficiency Gap Explanation: It provides a clear geometric reason for why Score Matching fails to be efficient in certain models (the misalignment between Fisher and Wasserstein score spaces).
• Practical Algorithm: It offers a constructive method to build estimators that are computationally tractable (like Score Matching) but statistically more efficient (approaching MLE), bridging the gap between intractable MLE and standard Score Matching.
• Generalizability: The framework is extended to Riemannian manifolds and truncated domains, making it applicable to a wide range of modern statistical models in machine learning and physics.

In summary, Nagai and Yano demonstrate that by leveraging the geometry of the Wasserstein space and Stein operators, one can systematically improve the statistical efficiency of Score Matching without incurring the computational cost of calculating the normalizing constant.