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

Resumen Técnico: $\lambda$-PSD: Discrepancias de Stein Polinomiales Aproximadas y Optimizadas para la SNR Escalable

Planteamiento del Problema
Las Discrepancias de Stein Polinomiales (PSD) han surgido como una alternativa escalable y de tiempo lineal a las Discrepancias de Stein de Kernel (KSD) para medir la calidad de las muestras y realizar pruebas de bondad de ajuste (GoF) cuando la distribución objetivo $p$ se conoce solo hasta una constante de normalización. Si bien las PSD ofrecen ventajas computacionales al reemplazar las funciones kernel por características polinomiales de orden $Q$, sus propiedades estadísticas siguen estando insuficientemente comprendidas. Específicamente, la evidencia empírica sugiere que la potencia de la prueba es altamente sensible a la elección del grado polinomial $Q$, exhibiendo a menudo un comportamiento no monotónico donde valores intermedios de $Q$ producen un rendimiento deficiente. El artículo identifica una limitación teórica fundamental: el aumento del grado polinomial amplifica principalmente la señal sin controlar adecuadamente la varianza, lo que conduce a un modo de fallo donde la relación señal-ruido al cuadrado ($SNR^2$) decae exponencialmente con $Q$ bajo configuraciones gaussianas. Este decaimiento socava la fiabilidad de las pruebas basadas en PSD a medida que aumenta la complejidad del modelo.

Metodología
Los autores reformulan la construcción de las discrepancias de Stein como un problema explícito de maximización de $SNR^2$.

1. Análisis Teórico de PSD: El artículo analiza primero la Discrepancia de Stein Polinomial estándar basada en $L_1$ ($PSD_{L_1}$). Al derivar la escala asintótica de $SNR^2$ para el correspondiente estadístico U y el estadístico z de muestra dividida bajo supuestos gaussianos ($p = \mathcal{N}(0, I_d)$ y $q = \mathcal{N}(\delta, \sigma^2 I_d)$), demuestran que $SNR^2$ decae exponencialmente cuando $Q \to \infty$. Esto se atribuye al inflado incontrolado de la varianza en los términos polinomiales de orden superior en relación con la señal.
2. Marco Optimizado para SNR: Para abordar esto, los autores proponen maximizar la razón entre el valor esperado al cuadrado y la varianza de la salida del operador de Stein:
   $$\sup_h \frac{(\mathbb{E}_q[L_1 h(x)])^2}{\text{Var}_q[L_1 h(x)]}$$
   Esta formulación conduce a un cociente de Rayleigh generalizado sobre las características de Stein. Teóricamente, la solución óptima implica invertir la matriz de covarianza completa del vector de características, lo cual es computacionalmente prohibitivo ($O(Q^3 d^6)$) y estadísticamente mal planteado en altas dimensiones debido a la deficiencia de rango.
3. $\lambda$-PSD (Solución Aproximada): Para lograr la escalabilidad, los autores introducen $\lambda$-PSD. En lugar de optimizar sobre el espacio de características completo y de alta dimensión, restringen la optimización a un subespacio de baja dimensión generado por la solución $PSD_{L_1}$ no ponderada.
   • Definen un vector de características de Stein $\Psi(X)$ basado en los coeficientes teóricamente óptimos de la PSD estándar.
   • Introducen un vector de pesos escalar $\lambda \in \mathbb{R}^{Q+1}$ para reponderar estas características.
   • El problema de optimización se reduce a encontrar $\lambda$ que maximice el cociente de Rayleigh sobre este subespacio reducido: $\sup_{\lambda} \frac{(\lambda^\top \mu_\Psi)^2}{\lambda^\top \Sigma_\Psi \lambda}$.
   • Para resolver esto eficientemente y evitar la inestabilidad numérica en la matriz de covarianza empírica $\hat{\Sigma}_\Psi$, los autores emplean una estrategia de optimización estocástica basada en el método de Entropía Cruzada (CE). Este enfoque aprende adaptativamente una distribución de muestreo para $\lambda$, realizando efectivamente una regularización implícita mediante la poda de componentes que contribuyen poco a la señal.

Contribuciones Clave

• Caracterización Teórica de Modos de Fallo: El artículo proporciona la primera caracterización teórica que muestra cómo la $SNR^2$ de las PSD estándar puede decaer exponencialmente con el grado polinomial bajo configuraciones gaussianas, explicando la sensibilidad observada a $Q$.
• Marco Optimizado para SNR: Los autores proponen un nuevo marco que trata la construcción de la discrepancia de Stein como un problema explícito de maximización de la $SNR^2$, cambiando el objetivo de simplemente maximizar la señal a optimizar la relación señal-ruido.
• Estimador $\lambda$-PSD: Introducen $\lambda$-PSD, un estimador aproximado y escalable que logra una $SNR^2$ estable bajo configuraciones gaussianas. Mantiene una complejidad de tiempo lineal respecto al número de muestras mientras mejora sustancialmente la robustez ante la elección del orden polinomial $Q$.
• Implementación Práctica: El artículo detalla un marco de prueba de muestra dividida utilizando una prueba z de una muestra para evitar la carga computacional del bootstrapping de estadísticos U de alta dimensión, haciendo que el método sea aplicable a entornos de gran escala.

Resultados

• Estabilidad Teórica: Bajo configuraciones gaussianas, los autores demuestran que $\lambda$-PSD evita el colapso exponencial de $SNR^2 inherente a las PSD estándar, logrando una$SNR^2$ estable acotada por la divergencia chi-cuadrado entre las distribuciones.
• Rendimiento Empírico: Los estudios de simulación demuestran que $\lambda$-PSD mejora sustancialmente la potencia de la prueba en comparación con las PSD estándar. Crucialmente, logra un rendimiento competitivo en las pruebas de bondad de ajuste manteniendo la complejidad de tiempo lineal en el número de muestras. El método muestra una reducción en la sensibilidad a la elección específica del grado polinomial $Q$, abordando los problemas de rendimiento no monotónico observados en trabajos previos.

Significancia
El artículo destaca la importancia crítica del diseño consciente de la $SNR$ en el desarrollo de discrepancias de Stein escalables. Al identificar que los marcos de PSD existentes optimizan el objetivo incorrecto (la señal sola en lugar de la relación señal-ruido), los autores proporcionan un camino fundamentado hacia pruebas estadísticas más robustas y fiables. El $\lambda$-PSD propuesto ofrece una solución práctica que cierra la brecha entre la optimalidad teórica y la viabilidad computacional, permitiendo pruebas de bondad de ajuste efectivas en entornos de computación bayesiana de gran escala donde las distribuciones objetivo se conocen solo hasta una constante de normalización.