$\lambda$-PSD: Scalable Approximate SNR-Optimised Polynomial Stein Discrepancies

Technical Summary: $\lambda$-PSD: Scalable Approximate SNR-Optimised Polynomial Stein Discrepancies

Problem Statement
Polynomial Stein Discrepancies (PSD) have emerged as a scalable, linear-time alternative to Kernel Stein Discrepancies (KSD) for measuring sample quality and conducting goodness-of-fit (GoF) tests when the target distribution $p$ is known only up to a normalizing constant. While PSDs offer computational advantages by replacing kernel functions with polynomial features of order $Q$, their statistical properties remain poorly understood. Specifically, empirical evidence suggests that test power is highly sensitive to the choice of polynomial degree $Q$, often exhibiting non-monotonic behavior where intermediate values of $Q$ yield poor performance. The paper identifies a fundamental theoretical limitation: increasing the polynomial degree primarily amplifies the signal without adequately controlling variance, leading to a failure mode where the Signal-to-Noise Ratio squared ($SNR^2$) decays exponentially with $Q$ under Gaussian settings. This decay undermines the reliability of PSD-based tests as the model complexity increases.

Methodology
The authors reformulate the construction of Stein discrepancies as an explicit $SNR^2$ maximization problem.

1. Theoretical Analysis of PSD: The paper first analyzes the standard $L_1$-based Polynomial Stein Discrepancy ($PSD_{L_1}$). By deriving the asymptotic scaling of $SNR^2$ for the corresponding U-statistic and split-sample z-statistic under Gaussian assumptions ($p = \mathcal{N}(0, I_d)$ and $q = \mathcal{N}(\delta, \sigma^2 I_d)$), they prove that $SNR^2$ decays exponentially as $Q \to \infty$. This is attributed to the uncontrolled inflation of variance in higher-order polynomial terms relative to the signal.
2. SNR-Optimised Framework: To address this, the authors propose maximizing the ratio of the squared expected value to the variance of the Stein operator output:
   $$\sup_h \frac{(\mathbb{E}_q[L_1 h(x)])^2}{\text{Var}_q[L_1 h(x)]}$$
   This formulation leads to a generalized Rayleigh quotient over Stein features. Theoretically, the optimal solution involves inverting the full covariance matrix of the feature vector, which is computationally prohibitive ($O(Q^3 d^6)$) and statistically ill-posed in high dimensions due to rank deficiency.
3. $\lambda$-PSD (Approximate Solution): To achieve scalability, the authors introduce $\lambda$-PSD. Instead of optimizing over the full high-dimensional feature space, they restrict the optimization to a low-dimensional subspace spanned by the unweighted $PSD_{L_1}$ solution.
   • They define a vector of Stein features $\Psi(X)$ based on the theoretically optimal coefficients of the standard PSD.
   • They introduce a scalar weight vector $\lambda \in \mathbb{R}^{Q+1}$ to reweight these features.
   • The optimization problem reduces to finding $\lambda$ that maximizes the Rayleigh quotient over this reduced subspace: $\sup_{\lambda} \frac{(\lambda^\top \mu_\Psi)^2}{\lambda^\top \Sigma_\Psi \lambda}$.
   • To solve this efficiently and avoid numerical instability in the empirical covariance matrix $\hat{\Sigma}_\Psi$, the authors employ a stochastic optimization strategy based on the Cross-Entropy (CE) method. This approach adaptively learns a sampling distribution for $\lambda$, effectively performing implicit regularization by pruning components that contribute little to the signal.

Key Contributions

• Theoretical Characterization of Failure Modes: The paper provides the first theoretical characterization showing that the $SNR^2$ of standard PSDs can decay exponentially with polynomial degree under Gaussian settings, explaining the observed sensitivity to $Q$.
• SNR-Optimised Framework: The authors propose a new framework that treats Stein discrepancy construction as an explicit $SNR^2$ maximization problem, shifting the objective from merely maximizing signal to optimizing the signal-to-noise ratio.
• $\lambda$-PSD Estimator: They introduce $\lambda$-PSD, a scalable, approximate estimator that achieves a stable $SNR^2$ under Gaussian settings. It retains linear-time complexity with respect to the number of samples while substantially improving robustness to the choice of polynomial order $Q$.
• Practical Implementation: The paper details a split-sample testing framework using a one-sample z-test to bypass the computational overhead of bootstrapping high-dimensional U-statistics, making the method applicable to large-scale settings.

Results

• Theoretical Stability: Under Gaussian settings, the authors prove that $\lambda$-PSD avoids the exponential $SNR^2$ collapse inherent in standard PSDs, achieving a stable $SNR^2$ bounded by the chi-squared divergence between the distributions.
• Empirical Performance: Simulation studies demonstrate that $\lambda$-PSD substantially improves test power compared to standard PSDs. Crucially, it achieves competitive goodness-of-fit testing performance while retaining linear-time complexity in the number of samples. The method shows reduced sensitivity to the specific choice of polynomial degree $Q$, addressing the non-monotonic performance issues observed in prior work.

Significance
The paper highlights the critical importance of $SNR$-aware design in the development of scalable Stein discrepancies. By identifying that existing PSD frameworks optimize the wrong objective (signal alone rather than signal-to-noise ratio), the authors provide a principled path toward more robust and reliable statistical tests. The proposed $\lambda$-PSD offers a practical solution that bridges the gap between theoretical optimality and computational feasibility, enabling effective goodness-of-fit testing in large-scale Bayesian computation settings where target distributions are known only up to a normalizing constant.