Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels

Resumo Técnico: Pare o Amostrador! Classificador para Parada Adaptativa em Núcleos de Amostragem

Definição do Problema
A amostragem de densidades de probabilidade não normalizadas complexas, definidas como $\pi(x) = r(x)/Z$, é um desafio fundamental na inferência Bayesiana e na modelagem probabilística. Embora os métodos de Monte Carlo via Cadeias de Markov (MCMC) forneçam garantias assintóticas, eles frequentemente sofrem com mistura lenta e altos custos computacionais, particularmente em espaços de alta dimensão. Uma limitação crítica do MCMC padrão é a dependência de comprimentos de trajetória fixos ou ajustados manualmente. As abordagens existentes para determinar tempos de parada baseiam-se frequentemente em diagnósticos online (ex: Gelman–Rubin, Geweke) computados a partir de amostras geradas, que nem sempre fornecem critérios de parada afirmativos e podem levar a uma terminação prematura ou estimativas enviesadas. Por outro lado, abordagens amortizadas (ex: fluxos normalizadores, amostradores de difusão, GFlowNets) deslocam o fardo computacional para o treinamento, mas muitas vezes ignoram as garantias exatas de amostragem do MCMC.

Metodologia
Os autores propõem um novo framework que trata a terminação da trajetória como um componente aprendível da dinâmica de amostragem, unindo o MCMC à amostragem amortizada. A metodologia central envolve:

1. Formulação de GFlowNet Não-acíclica: Os autores enquadram o MCMC com parada adaptativa dentro da teoria das Redes de Fluxo Generativo (GFlowNets) não-acíclicas. Eles definem um espaço de estados estruturado $\bar{S} = \{s_0\} \cup S \cup \{s_f\}$, onde $S = \mathbb{R}^d$ representa o espaço de estados contínuos, $s_0$ é a fonte e $s_f$ é um estado de sumidouro absorvente. Diferente das GFlowNets acíclicas padrão, este cenário permite ciclos dentro de $S$.
2. Classificador Baseado em Classificador: Em vez de executar uma cadeia de Markov por um número fixo de passos, o método emprega classificadores neurais dependentes do estado, $d_F(s)$ e $d_B(s)$, para decidir quando uma trajetória deve terminar.
   • O classificador de avanço $d_F(s)$ determina a probabilidade de transicionar diretamente para o estado de sumidouro $s_f$ (parada).
   • O classificador de retrocesso $d_B(s)$ controla a probabilidade de mover-se para o estado de origem $s_0$.
   • A densidade de transição de avanço é definida como $p_F(s'|s) = (1 - d_F(s))q_F(s'|s)$ para transições contínuas, e $p_F(s_f|s) = d_F(s)$ para terminação.
3. Conexões Teóricas: O artigo estabelece uma conexão rigorosa entre as políticas de parada aprendidas e a densidade alvo $\pi$ via condições de equilíbrio detalhado. Teoricamente, mostra-se que o classificador de avanço ótimo satisfaz $d_F(s) = r(s)/f(s)$, onde $f(s)$ é a medida de fluxo. Isso relaciona a decisão de parada à discrepância acumulada entre a massa injetada e a distribuição alvo.
4. Algoritmo de Treinamento: Os autores utilizam um objetivo de equilíbrio de trajetória de prefixo (TB - prefix trajectory balance), que é computacionalmente eficiente para ambientes não-acíclicos. O processo de treinamento envolve:
   • Amostrar trajetórias parciais de comprimento máximo fixo $N_{max}$ com os classificadores desativados para incentivar a exploração.
   • Estender os prefixos com transições terminais para formar trajetórias completas.
   • Minimizar uma função de perda que combina a perda de equilíbrio de trajetória de prefixo com um termo de regularização de fluxo. A regularização incentiva trajetórias mais curtas ao penalizar o fluxo total, efetivamente minimizando o comprimento esperado da trajetória.
5. Generalização Multinível: Para lidar com geometrias complexas, o framework é estendido para um esquema multinível. O espaço de estados é aumentado com variáveis de nível $\ell \in \{1, \dots, L\}$, permitindo que o modelo alterne entre regimes. Recompensas intermediárias $r(s, \ell)$ são definidas para interpolar entre a distribuição inicial e o alvo, estabilizando o treinamento e a exploração de forma semelhante à amostragem de importância por recozimento (annealed importance sampling).

Principais Contribuições

1. Framework Teórico: Os autores demonstram que a dinâmica de amostragem em $\mathbb{R}^d$ com a possibilidade de parada em qualquer ponto pode ser vista como um ambiente de GFlowNet não-acíclico. Eles derivam a relação entre classificadores ótimos e a densidade alvo usando condições de fluxo e equilíbrio detalhado.
2. Algoritmos Práticos: Eles introduzem algoritmos práticos que combinam parada aprendida com núcleos de avanço e retrocesso adaptativos. Isso inclui uma extensão multinível projetada para melhorar a estabilidade de treinamento e a exploração em geometrias complexas.
3. Avaliação Empírica: O método é avaliado em várias densidades de referência (GMM9, Funnel, ManyWell e tarefas de Digits de alta dimensão). Os resultados demonstram que a abordagem reduz significativamente os comprimentos médios de trajetória em comparação com baselines padrão de MCMC (como ULA) enquanto melhora a cobertura de modos e a mistura.

Resultados Experimentais

• Comprimento de Trajetória: A abordagem baseada em classificador reduz substancialmente o comprimento médio da trajetória ($E[n_\tau]$) comparado ao ULA de passo fixo e ao ULA com diagnósticos de convergência padrão (Gelman–Rubin, Geweke). Por exemplo, na tarefa GMM9 ($d=2$), o classificador aprendido reduziu o comprimento da trajetória de ~104 passos (ULA) para ~15 passos.
• Qualidade da Amostra: Embora o classificador sozinho reduza o comprimento da trajetória, o modelo completo (incluindo correções de núcleo de avanço aprendidas) melhora significativamente a qualidade da aproximação. Isso é evidenciado por distâncias de Sinkhorn (SD) menores, menor MMD (Maximum Mean Discrepancy) e limites inferior/superior de evidência (ELBO/EUBO) mais estreitos.
• Cobertura de Modos: Visualizações e métricas quantitativas mostram que o método proposto transita efetivamente entre regiões de alta densidade distantes e modos, enquanto o ULA padrão frequentemente permanece preso em regiões locais.
• Desempenho Multinível: A extensão multinível ($L=5$) supera consistentemente modelos de nível único e torna-se competitiva com métodos de amostragem de difusão estabelecidos (PIS, DDS) em tarefas de alta dimensão como Digits ($d=196$).

Significância e Alegações
O artigo afirma fornecer um framework unificado onde as trajetórias de amostragem são dinamicamente encurtadas em regiões de alta densidade enquanto permanecem exploratórias em outros lugares. Ao tratar a parada como um componente aprendível, o método aborda a ineficiência de trajetórias de MCMC de comprimento fixo sem sacrificar as garantias teóricas associadas à amostragem baseada em fluxo. Os autores observam que, embora a abordagem introduza complexidade de otimização adicional (equilibrando exploração e regularização), ela oferece uma maneira fundamentada de desenviesar núcleos fixos (como o ULA) e adaptar-se à geometria de distribuições alvo complexas. O trabalho sugere que direções futuras poderiam envolver a melhoria dos objetivos de treinamento para GFlowNets contínuas não-acíclicas e a extensão do framework para dinâmicas mais expressivas como o Monte Carlo de Hamiltonian (HMC).