The Polynomial Stein Discrepancy for Assessing Moment Convergence

1. Problem Statement

In Bayesian inference, assessing the quality of samples generated by approximate sampling algorithms (particularly biased ones like Stochastic Gradient Langevin Dynamics, SGLD) is a critical "grand challenge."

• Limitations of Existing Methods:
  • Effective Sample Size (ESS): Not suitable for asymptotically biased algorithms.
  • Kernel Stein Discrepancy (KSD): The current gold standard. However, it suffers from quadratic computational complexity ($O(n^2)$) regarding the number of samples, making it infeasible for large-scale MCMC.
  • Linear-time KSD Approximations (FSSD, RFSD): While faster, these often suffer from the curse of dimensionality, require extensive hyperparameter tuning (e.g., test locations, feature maps), and crucially, fail to detect moment convergence.
• The Specific Gap: Standard KSD (using IMQ kernels) controls weak convergence (bounded-Lipschitz metric) but does not guarantee convergence in moments. Since moments (mean, variance) are often the primary quantities of interest in big-data applications (where the posterior is often Gaussian-like via the Bernstein-von Mises theorem), the inability to detect biases in moments is a significant drawback.

2. Methodology: Polynomial Stein Discrepancy (PSD)

The authors propose the Polynomial Stein Discrepancy (PSD), a linear-time alternative to KSD designed specifically to detect discrepancies in the first $r$ moments of a distribution.

Core Formulation

• Stein Operator: Uses the second-order Langevin-Stein operator $A$ defined for a target density $p(x)$:
  $$A g(x) = \Delta_x g(x) + \nabla_x g(x) \cdot \nabla_x \log p(x)$$
• Function Class: Instead of a Reproducing Kernel Hilbert Space (RKHS), PSD restricts the function class $\mathcal{G}$ to the span of polynomials of order $r$.
  $$\mathcal{G} = \text{span}\left\{ \prod_{i=1}^d x_i^{\alpha_i} : \sum \alpha_i \leq r \right\}$$
• Discrepancy Definition: The PSD is defined as the supremum of the expectation of the Stein operator over this polynomial class, constrained to the unit ball:
  $$\text{PSD} = \sup_{g \in \mathcal{G}, \|g\| \leq 1} \left| \mathbb{E}_{Q}[Ag(X)] \right|$$
• Closed-Form Solution: By parameterizing $g$ as a linear combination of monomials with coefficients $\beta$, the optimization yields a closed-form solution:
  $$\text{PSD} = \sqrt{\sum_{k=1}^J \bar{z}_k^2}$$
  where $\bar{z}_k = \mathbb{E}_Q[A P_k(X)]$ and $P_k$ are the basis monomials.
• Computational Complexity:
  • The squared PSD can be computed as a V-statistic or U-statistic.
  • Complexity is $O(nJ)$, where $J = \binom{d+r}{d} - 1$.
  • This is linear in the number of samples $n$, a significant improvement over KSD's $O(n^2)$.
  • For very high dimensions, an approximate version excluding interaction terms can reduce complexity to $O(nd^r)$.

Goodness-of-Fit Testing

• The paper proposes a bootstrap-based test for the null hypothesis $H_0: Q = P$.
• Unlike asymptotic tests for linear-time KSD alternatives (which require estimating a covariance matrix under $P$), the PSD bootstrap test resamples weights from the observed data $Q$.
• This approach avoids the need to sample from the intractable target $P$ and empirically demonstrates higher statistical power than asymptotic approximations.

3. Key Theoretical Contributions

• Moment Detection Guarantee (Proposition 1):
  • If the target $P$ is Gaussian, $\text{PSD} = 0$ if and only if the first $r$ moments of the sample distribution $Q$ match the first $r$ moments of $P$.
  • This provides a rigorous theoretical justification for using PSD specifically for moment convergence, a property standard KSD lacks.
• Asymptotic Power (Corollary 3.0.2):
  • In the Bernstein-von Mises limit (large data regime where the posterior approaches a Gaussian), the PSD-based tests have asymptotic power approaching 1 for detecting discrepancies in the first $r$ moments.
• Invariance: The method is invariant under invertible linear transformations (e.g., whitening) when applied to Gaussian targets.

4. Experimental Results

The authors evaluated PSD against competitors: IMQ KSD (quadratic time), Gaussian KSD, FSSD (linear time), and RFSD (linear time).

• Goodness-of-Fit Power:
  • PSD consistently outperformed linear-time competitors (FSSD, RFSD) in detecting moment discrepancies (variance, kurtosis).
  • For example, in detecting variance perturbations in high dimensions ($d=20$), PSD achieved statistical power up to 4 times higher than linear-time competitors and competitive with quadratic-time KSD.
  • PSD with $r=4$ was the only method to consistently achieve power $\approx 1$ in detecting kurtosis discrepancies.
• Hyperparameter Tuning (SGLD):
  • In a Stochastic Gradient Langevin Dynamics (SGLD) step-size selection task, PSD ($r=2,3,4$) correctly identified the optimal step size ($\epsilon=0.005$) that minimized bias in the posterior variance, matching the performance of the much slower IMQ KSD.
• Computational Efficiency:
  • PSD is orders of magnitude faster than KSD.
  • In the SGLD example, PSD was ~70 times faster than KSD and ~7 times faster than RFSD.
• Robustness:
  • PSD performed well even in non-Gaussian settings (e.g., Logistic Regression with sparsity priors, Restricted Boltzmann Machines), though its theoretical guarantee is strongest for Gaussian targets.

5. Significance and Impact

• Scalability: PSD bridges the gap between the statistical rigor of KSD and the computational efficiency required for modern large-scale Bayesian inference.
• Targeted Utility: It specifically addresses the "blind spot" of existing Stein discrepancies: the inability to detect biases in moments. This is crucial for biased MCMC algorithms where moment convergence is the primary goal.
• Practicality: The method requires minimal tuning (only the polynomial order $r$), has a clear interpretation (tracking specific moments), and is computationally feasible for high-dimensional datasets.
• Recommendation: The authors recommend using PSD with $r=2$ for practitioners using biased samplers (like SG-MCMC) to efficiently assess sample quality and tune hyperparameters.

Conclusion

The Polynomial Stein Discrepancy offers a novel, linear-time framework for assessing sample quality in Bayesian inference. By leveraging polynomial functions within the Stein framework, it provides a theoretically grounded, computationally efficient, and statistically powerful tool for detecting moment convergence, overcoming the scalability and tuning limitations of current Kernel Stein Discrepancy methods.