Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

Resumen Técnico: Estimación de la Media Armónica Aprendida con Flow Matching para Posteriores Multimodales

Planteamiento del Problema
La verosimilitud marginal (o evidencia bayesiana), denotada como $z = p(d|M)$, es una cantidad fundamental para la comparación de modelos bayesianos, lo que permite una evaluación cuantitativa mediante el factor de Bayes. Sin embargo, calcular esta cantidad requiere evaluar una integral multidimensional sobre todo el espacio de parámetros, lo cual es computacionalmente prohibitivo para modelos complejos y de alta dimensión. Mientras que los métodos de Monte Carlo por Cadenas de Markov (MCMC) y el Monte Carlo Hamiltoniano (HMC) sobresalen en la generación de muestras posteriores en altas dimensiones, no proporcionan inherentemente la verosimilitud marginal. Por el contrario, el muestreo anidado (nested sampling) proporciona la evidencia, pero restringe la estrategia de muestreo a una jerarquía anidada específica.

El "estimador de la media armónica aprendida" fue introducido previamente para estimar la verosimilitud marginal utilizando únicamente muestras posteriores, independientemente del método de muestreo utilizado para obtenerlas. Este enfoque se basa en el aprendizaje de una densidad objetivo interna $\phi(\theta)$ que está contenida dentro de la posterior pero que posee colas más delgadas. Sin embargo, las implementaciones anteriores utilizando flujos normalizantes discretos o estimadores de densidad personalizados presentaban dificultades ante posteriores altamente multimodales. Estas arquitecturas a menudo fallaban al capturar topologías complejas sin modificaciones heurísticas (por ejemplo, la construcción manual de mezclas gaussianas para la distribución base) o sufrían de la asignación de una masa de probabilidad no despreciable a los puentes de densidad entre los modos, lo que conducía a una varianza infinita en el estimador.

Metodología
Este trabajo aborda las limitaciones de los estimadores de densidad previos integrando el Flow Matching en el marco de la media armónica aprendida. La metodología procede de la siguiente manera:

1. Estimador de la Media Armónica Aprendida: El estimador aproxima el recíproco de la verosimilitud marginal $\rho = z^{-1}$ utilizando muestras posteriores $\theta_i \sim p(\theta|d)$ y una densidad objetivo aprendida $\phi(\theta)$:
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   Para que el estimador tenga una varianza finita, $\phi(\theta)$ debe estar concentrada dentro de las regiones de alta probabilidad de la posterior y decaer más rápido en las colas.

2. Arquitectura de Flow Matching: En lugar de flujos normalizantes discretos, los autores emplean Flujos Normalizantes Continuos (CNFs) entrenados mediante flow matching.

   • Definición: Un CNF parametriza la transformación entre una distribución base simple (Gaussiana estándar) y la posterior objetivo mediante un campo de velocidad dependiente del tiempo $v_t(\theta)$.
   • Objetivo de Entrenamiento: La red se entrena utilizando un objetivo de regresión libre de simulación. Se define una trayectoria de probabilidad condicional como una línea recta entre una muestra de ruido $\theta_0$ y una muestra posterior $\theta_1$: $\theta_t = (1-t)\theta_0 + t\theta_1$. La red minimiza el error cuadrático medio entre el campo vectorial predicho y el verdadero campo vectorial condicional $(\theta_1 - \theta_0)$.
   • Ventajas: Este enfoque permite utilizar redes neuronales de alta capacidad y arquitectónicamente sin restricciones (Perceptrones Multicapa con activación Swish) y evita las costosas simulaciones de ODE requeridas en los métodos de CNF anteriores durante el entrenamiento.

3. Escalamiento de Temperatura: Para asegurar que la densidad objetivo $\phi(\theta)$ tenga colas más delgadas que la posterior, se aplica un parámetro de temperatura $T \in (0, 1)$ para escalar la varianza de la distribución base del flujo. Esto concentra la distribución del flujo, satisfaciendo el requisito de contención para el estimador de la media armónica.

Contribuciones Clave

• Arquitectura Novedosa: La introducción de flujos normalizantes continuos entrenados por flow matching como el estimador de densidad interna para el método de la media armónica aprendida.
• Gestión de la Multimodalidad: Demostrar que esta arquitectura puede capturar con precisión estructuras posteriores altamente multimodales y complejas sin la necesidad de modificaciones heurísticas a la distribución base o ajustes finos.
• Escalabilidad: Proporcionar un método que es agnóstico a la estrategia de muestreo (compatible con MCMC, HMC, NUTS, SBI, etc.) y capaz de operar en espacios de dimensión moderada a alta.

Resultados Experimentales
Los autores validaron el enfoque en dos entornos de prueba desafiantes:

1. Función de Rastrigin (2 Dimensiones):

   • Un benchmark estándar con muchos modos locales.
   • El modelo de flow matching (6 capas, 256 unidades ocultas) capturó exitosamente todos los modos con una temperatura $T=0.98$.
   • Los resultados de 100 ejecuciones experimentales mostraron que los valores estimados de la verosimilitud marginal fueron precisos y coincidieron estrechamente con la verdad de terreno obtenida mediante integración numérica. Las estimaciones de error fueron fiables.

2. Mezcla de Gaussianas en 20 Dimensiones:

   • Una mezcla de cinco componentes gaussianas en 20 dimensiones con modos no superpuestos y correlaciones introducidas.
   • El modelo (10 capas, 256 unidades ocultas) fue entrenado con 40,000 muestras.
   • Se demostró que la distribución objetivo aprendida estaba adecuadamente concentrada dentro de la posterior, capturando la topología desafiante.
   • Las verosimilitudes marginales estimadas a través de 100 ejecuciones se alinearon con la verdad de terreno analítica, confirmando la precisión del método y la robustez de la estimación de error incluso en este escenario complejo de dimensión moderada.

Significancia y Reivindicaciones
El artículo sostiene que, al aprovechar la expresividad de los CNF entrenados por flow matching, el estimador de la media armónica aprendida puede ahora estimar de forma robusta la evidencia para posteriores altamente complejas y multimodales. La principal significancia radica en eliminar la dependencia de ajustes heurísticos (como las mezclas gaussianas seleccionadas manualmente) que anteriormente eran necesarios para manejar la multimodalidad. Los autores afirman que este enfoque permite el cálculo preciso de la evidencia utilizando muestras de cualquier método de muestreo acelerado o escalable, facilitando la comparación de modelos bayesianos en campos ricos en datos como la astrofísica, sin los cuellos de botella computacionales del muestreo anidado tradicional o la inestabilidad del estimador original de la media armónica.

El trabajo futuro, como señalan los autores, incluye la aplicación de esta metodología a problemas del mundo real, como la astrofísica de ondas gravitacionales, y la exploración de nuevas refinaciones arquitectónicas para tareas de muy alta dimensión, como la imagenología y la inferencia cosmológica a nivel de campo.