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

Resumen Técnico: ¡Detén el Muestreador! Clasificador de Parada Adaptativa para Núcleos de Muestreo

Planteamiento del Problema
El muestreo de densidades de probabilidad no normalizadas complejas, definidas como $\pi(x) = r(x)/Z$, es un desafío fundamental en la inferencia bayesiana y el modelado probabilístico. Aunque los métodos de Monte Carlo por Cadenas de Markov (MCMC) proporcionan garantías asintóticas, frecuentemente sufren de una mezcla lenta y altos costos computacionales, particularmente en espacios de alta dimensión. Una limitación crítica del MCMC estándar es la dependencia de longitudes de trayectoria fijas o ajustadas manualmente. Los enfoques existentes para determinar los tiempos de parada suelen depender de diagnósticos en línea (por ejemplo, Gelman–Rubin, Geweke) computados a partir de muestras generadas, los cuales no siempre proporcionan criterios de parada afirmativos y pueden conducir a una terminación prematura o estimaciones sesgadas. Por el contrario, los enfoques amortizados (por ejemplo, flujos normalizantes, muestreadores de difusión, GFlowNets) trasladan la carga computacional al entrenamiento, pero a menudo omiten las garantías de muestreo exacto de MCMC.

Metodología
Los autores proponen un marco novedoso que trata la terminación de la trayectoria como un componente aprendible de la dinámica de muestreo, uniendo MCMC con el muestreo amortizado. La metodología central involucra:

1. Formulación de GFlowNet No Aclíclica: Los autores enmarcan el MCMC con parada adaptativa dentro de la teoría de las Redes de Flujo Generativo (GFlowNets) no aclíclicas. Definen un espacio de estados estructurado $\bar{S} = \{s_0\} \cup S \cup \{s_f\}$, donde $S = \mathbb{R}^d$ representa el espacio de estados continuos, $s_0$ es la fuente y $s_f$ es un estado de sumidero absorbente. A diferencia de las GFlowNets acyclicas estándar, este entorno permite ciclos dentro de $S$.
2. Parada Basada en Clasificadores: En lugar de ejecutar una cadena de Markov durante un número fijo de pasos, el método emplea clasificadores neuronales dependientes del estado, $d_F(s)$ y $d_B(s)$, para decidir cuándo debe terminar una trayectoria.
   • El clasificador hacia adelante $d_F(s)$ determina la probabilidad de transicionar directamente al estado sumidero $s_f$ (parada).
   • El clasificador hacia atrás $d_B(s)$ controla la probabilidad de moverse al estado fuente $s_0$.
   • La densidad de transición hacia adelante se define como $p_F(s'|s) = (1 - d_F(s))q_F(s'|s)$ para transiciones continuas, y $p_F(s_f|s) = d_F(s)$ para la terminación.
3. Conexiones Teóricas: El artículo establece una conexión rigurosa entre las políticas de parada aprendidas y la densidad objetivo $\pi$ mediante condiciones de equilibrio detallado. Teóricamente, se demuestra que el clasificador hacia adelante óptimo satisface $d_F(s) = r(s)/f(s)$, donde $f(s)$ es la medida de flujo. Esto relaciona la decisión de parada con la discrepancia acumulada entre la masa inyectada y la distribución objetivo.
4. Algoritmo de Entrenamiento: Los autores utilizan un objetivo de equilibrio de trayectoria de prefijo (prefix trajectory balance - TB), el cual es computacionalmente eficiente para entornos no aclíclicos. El proceso de entrenamiento involucra:
   • Muestrear trayectorias parciales de longitud máxima fija $N_{max}$ con los clasificadores desactivados para fomentar la exploración.
   • Extender los prefijos con transiciones terminales para formar trayectorias completas.
   • Minimizar una función de pérdida que combina la pérdida de equilibrio de trayectoria de prefijo con un término de regularización de flujo. La regularización fomenta trayectorias más cortas penalizando el flujo total, minimizando efectivamente la longitud esperada de la trayectoria.
5. Generalización Multinivel: Para manejar geometrías complejas, el marco se extiende a un esquema multinivel. El espacio de estados se aumenta con variables de nivel $\ell \in \{1, \dots, L\}$, permitiendo al modelo cambiar entre regímenes. Se definen recompensas intermedias $r(s, \ell)$ para interpolar entre la distribución inicial y la objetivo, estabilizando el entrenamiento y la exploración de manera similar al muestreo de importancia con temple (annealed importance sampling).

Contribuciones Clave

1. Marco Teórico: Los autores demuestran que la dinámica de muestreo en $\mathbb{R}^d$ con la posibilidad de detenerse en cualquier punto puede verse como un entorno de GFlowNet no aclíclico. Derivan la relación entre los clasificadores óptimos y la densidad objetivo utilizando condiciones de flujo y equilibrio detallado.
2. Algoritmos Prácticos: Introducen algoritmos prácticos que combinan la parada aprendida con núcleos hacia adelante y hacia atrás adaptativos. Esto incluye una extensión multinivel diseñada para mejorar la estabilidad del entrenamiento y la exploración en geometrías complejas.
3. Evaluación Empírica: El método se evalúa en varios conjuntos de densidades de referencia (GMM9, Funnel, ManyWell y tareas de alta dimensión como Digits) Los resultados demuestran que el enfoque reduce significativamente las longitudes de trayectoria promedio en comparación con los baselines de MCMC estándar (como ULA) mientras mejora la cobertura de modos y la mezcla.

Resultados Experimentales

• Longitud de Trayectoria: El enfoque basado en clasificadores reduce sustancialmente la longitud media de la trayectoria ($E[n_\tau]$) en comparación con ULA de pasos fijos y ULA con diagnósticos de convergencia estándar (Gelman–Rubin, Geweke). Por ejemplo, en la tarea GMM9 ($d=2$), el clasificador aprendido redujo la longitud de la trayectoria de ~104 pasos (ULA) a ~15 pasos.
• Calidad de la Muestra: Mientras que el clasificador por sí solo reduce la longitud de la trayectoria, el modelo completo (incluyendo correcciones del núcleo hacia adelante aprendidas) mejora significativamente la calidad de la aproximación. Esto se evidencia mediante distancias de Sinkhorn (SD) más bajas, menor Discrepancia de Media Máxima (MMD) y límites inferior/superior de evidencia (ELBO/EUBO) más ajustados.
• Cobertura de Modos: Las visualizaciones y métricas cuantitativas muestran que el método propuesto transiciona eficazmente entre regiones de alta densidad distantes y modos, mientras que el ULA estándar a menudo permanece atrapado en regiones locales.
• Desempeño Multinivel: La extensión multinivel ($L=5$) supera consistentemente a los modelos de un solo nivel y se vuelve competitiva con métodos de muestreo de difusión establecidos (PIS, DDS) en tareas de alta dimensión como Digits ($d=196$).

Significancia y Reivindicaciones
El artículo afirma proporcionar un marco unificado donde las trayectorias de muestreo se acortan dinámicamente en regiones de alta densidad mientras permanecen exploratorias en otros lugares. Al tratar la parada como un componente aprendible, el método aborda la ineficiencia de las trayectorias de MCMC de longitud fija sin sacrificar las garantías teóricas asociadas con el muestreo basado en flujo. Los autores señalan que, si bien el enfoque introduce una complejidad de optimización adicional (balanceando exploración y regularización), ofrece una forma fundamentada de eliminar el sesgo de los núcleos fijos (como ULA) y adaptarse a la geometría de las distribuciones objetivo complejas. El trabajo sugiere que las direcciones futuras podrían incluir la mejora de los objetivos de entrenamiento para GFlowNets continuos no aclíclicos y la extensión del marco a dinámicas más expresivas como el Monte Carlo Hamiltoniano.