The Polynomial Stein Discrepancy for Assessing Moment Convergence

1. Formulação do Problema

Na inferência bayesiana, avaliar a qualidade das amostras geradas por algoritmos de amostragem aproximada (particularmente os enviesados, como a Dinâmica de Langevin com Gradiente Estocástico, SGLD) é um "grande desafio" crítico.

• Limitações dos Métodos Existentes:
  • Tamanho Efetivo da Amostra (ESS): Não é adequado para algoritmos enviesados assintoticamente.
  • Discrepância de Stein com Kernel (KSD): O padrão-ouro atual. No entanto, sofre de complexidade computacional quadrática ($O(n^2)$) em relação ao número de amostras, tornando-a inviável para MCMC em grande escala.
  • Aproximações de KSD em Tempo Linear (FSSD, RFSD): Embora mais rápidas, estas frequentemente sofrem com a maldição da dimensionalidade, exigem extenso ajuste de hiperparâmetros (por exemplo, locais de teste, mapas de características) e, crucialmente, falham em detectar a convergência de momentos.
• A Lacuna Específica: A KSD padrão (usando kernels IMQ) controla a convergência fraca (métrica de Lipschitz limitada), mas não garante a convergência em momentos. Como os momentos (média, variância) são frequentemente as quantidades primárias de interesse em aplicações de grandes dados (onde a posteriori é frequentemente semelhante a uma Gaussiana via o teorema de Bernstein-von Mises), a incapacidade de detectar vieses nos momentos é uma desvantagem significativa.

2. Metodologia: Discrepância de Stein Polinomial (PSD)

Os autores propõem a Discrepância de Stein Polinomial (PSD), uma alternativa em tempo linear à KSD projetada especificamente para detectar discrepâncias nos primeiros $r$ momentos de uma distribuição.

Formulação Central

• Operador de Stein: Utiliza o operador de Stein-Langevin de segunda ordem $A$ definido para uma densidade alvo $p(x)$:
  $$A g(x) = \Delta_x g(x) + \nabla_x g(x) \cdot \nabla_x \log p(x)$$
• Classe de Funções: Em vez de um Espaço de Hilbert de Reprodutor de Kernel (RKHS), a PSD restringe a classe de funções $\mathcal{G}$ ao span de polinômios de ordem $r$.
  $$\mathcal{G} = \text{span}\left\{ \prod_{i=1}^d x_i^{\alpha_i} : \sum \alpha_i \leq r \right\}$$
• Definição da Discrepância: A PSD é definida como o supremo da expectativa do operador de Stein sobre esta classe polinomial, restrita à bola unitária:
  $$\text{PSD} = \sup_{g \in \mathcal{G}, \|g\| \leq 1} \left| \mathbb{E}_{Q}[Ag(X)] \right|$$
• Solução de Forma Fechada: Ao parametrizar $g$ como uma combinação linear de monômios com coeficientes $\beta$, a otimização produz uma solução de forma fechada:
  $$\text{PSD} = \sqrt{\sum_{k=1}^J \bar{z}_k^2}$$
  onde $\bar{z}_k = \mathbb{E}_Q[A P_k(X)]$ e $P_k$ são os monômios de base.
• Complexidade Computacional:
  • A PSD ao quadrado pode ser computada como uma V-estatística ou U-estatística.
  • A complexidade é $O(nJ)$, onde $J = \binom{d+r}{d} - 1$.
  • Isso é linear no número de amostras $n$, uma melhoria significativa em relação aos $O(n^2)$ da KSD.
  • Para dimensões muito altas, uma versão aproximada que exclui termos de interação pode reduzir a complexidade para $O(nd^r)$.

Teste de Ajuste (Goodness-of-Fit)

• O artigo propõe um teste baseado em bootstrap para a hipótese nula $H_0: Q = P$.
• Diferentemente dos testes assintóticos para alternativas de KSD em tempo linear (que exigem estimar uma matriz de covariância sob $P$), o teste de bootstrap da PSD reamostra pesos dos dados observados $Q$.
• Esta abordagem evita a necessidade de amostrar a partir do alvo intratável $P$ e demonstra empiricamente maior poder estatístico do que aproximações assintóticas.

3. Principais Contribuições Teóricas

• Garantia de Detecção de Momentos (Proposição 1):
  • Se o alvo $P$ for Gaussiano, $\text{PSD} = 0$ se e somente se os primeiros $r$ momentos da distribuição amostral $Q$ coincidirem com os primeiros $r$ momentos de $P$.
  • Isso fornece uma justificativa teórica rigorosa para o uso da PSD especificamente para convergência de momentos, uma propriedade que a KSD padrão carece.
• Poder Assintótico (Corolário 3.0.2):
  • No limite de Bernstein-von Mises (regime de grandes dados onde a posteriori se aproxima de uma Gaussiana), os testes baseados em PSD têm poder assintótico aproximando-se de 1 para detectar discrepâncias nos primeiros $r$ momentos.
• Invariância: O método é invariante sob transformações lineares invertíveis (por exemplo, branqueamento) quando aplicado a alvos Gaussianos.

4. Resultados Experimentais

Os autores avaliaram a PSD contra concorrentes: KSD IMQ (tempo quadrático), KSD Gaussiana, FSSD (tempo linear) e RFSD (tempo linear).

• Poder de Ajuste (Goodness-of-Fit Power):
  • A PSD superou consistentemente os concorrentes em tempo linear (FSSD, RFSD) na detecção de discrepâncias de momentos (variância, curtose).
  • Por exemplo, na detecção de perturbações de variância em altas dimensões ($d=20$), a PSD alcançou poder estatístico até 4 vezes maior do que os concorrentes em tempo linear e foi competitiva com a KSD de tempo quadrático.
  • A PSD com $r=4$ foi o único método a alcançar consistentemente poder $\approx 1$ na detecção de discrepâncias de curtose.
• Ajuste de Hiperparâmetros (SGLD):
  • Em uma tarefa de seleção de tamanho de passo para Dinâmica de Langevin com Gradiente Estocástico (SGLD), a PSD ($r=2,3,4$) identificou corretamente o tamanho de passo ótimo ($\epsilon=0.005$) que minimizou o viés na variância da posteriori, igualando o desempenho da muito mais lenta KSD IMQ.
• Eficiência Computacional:
  • A PSD é ordens de magnitude mais rápida que a KSD.
  • No exemplo SGLD, a PSD foi ~70 vezes mais rápida que a KSD e ~7 vezes mais rápida que a RFSD.
• Robustez:
  • A PSD performou bem mesmo em cenários não Gaussianos (por exemplo, Regressão Logística com priors de esparsidade, Máquinas de Boltzmann Restritas), embora sua garantia teórica seja mais forte para alvos Gaussianos.

5. Significado e Impacto

• Escalabilidade: A PSD preenche a lacuna entre o rigor estatístico da KSD e a eficiência computacional exigida para inferência bayesiana moderna em grande escala.
• Utilidade Direcionada: Ela aborda especificamente o "ponto cego" das discrepâncias de Stein existentes: a incapacidade de detectar vieses em momentos. Isso é crucial para algoritmos de MCMC enviesados onde a convergência de momentos é o objetivo primário.
• Praticidade: O método requer ajuste mínimo (apenas a ordem do polinômio $r$), possui uma interpretação clara (rastreamento de momentos específicos) e é computacionalmente viável para conjuntos de dados de alta dimensão.
• Recomendação: Os autores recomendam o uso da PSD com $r=2$ para praticantes que utilizam amostradores enviesados (como SG-MCMC) para avaliar eficientemente a qualidade da amostra e ajustar hiperparâmetros.

Conclusão

A Discrepância de Stein Polinomial oferece um novo framework em tempo linear para avaliar a qualidade da amostra na inferência bayesiana. Ao alavancar funções polinomiais dentro do framework de Stein, ela fornece uma ferramenta teoricamente fundamentada, computacionalmente eficiente e estatisticamente poderosa para detectar a convergência de momentos, superando as limitações de escalabilidade e ajuste dos métodos atuais de Discrepância de Stein com Kernel.