$λ$-PSD: Scalable Approximate SNR-Optimised Polynomial Stein Discrepancies

Abstract

Polynomial Stein discrepancies (PSD) provide a scalable alternative to kernel Stein methods for measuring sample quality and goodness-of-fit testing, but their statistical properties remain poorly understood. We show that increasing polynomial degree primarily amplifies signal without adequately controlling variance, rather than directly optimising the signal-to-noise ratio (SNR). Under suitable assumptions, this might lead to a failure mode in which the $\text{SNR}^2$ can provably decay exponentially with polynomial degree. Motivated by this observation, we reformulate Stein discrepancy construction as an explicit $\text{SNR}^2$ maximisation problem, yielding a Rayleigh quotient over Stein features. This perspective motivates $λ$-PSD, an approximate scalable covariance-aware reweighting scheme defined in a low-dimensional subspace. Under Gaussian settings, we show that $λ$-PSD avoids the exponential $\text{SNR}^2$ collapse and achieves a stable $\text{SNR}^2$. Empirically, $λ$-PSD substantially improves test power while retaining linear-time complexity in the number of samples, highlighting the importance of SNR-aware design for scalable Stein discrepancies.