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

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 SNR2 can provably decay exponentially with polynomial degree. Motivated by this observation, we reformulate Stein discrepancy construction as an explicit SNR2 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 SNR2 collapse and achieves a stable SNR2. 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.