Estimating Complex Densities using Two-Stage Normalizing Flows

Resumen Técnico: Estimación de Densidades Complejas mediante Flujos Normalizantes de Dos Etapas

1. Planteamiento del Problema

En muchas aplicaciones científicas (estadística aplicada, astrofísica, modelos jerárquicos bayesianos), la distribución de probabilidad objetivo (densidad objetivo) no puede evaluarse en forma cerrada ni muestrearse directamente. A menudo, la información sobre esta distribución es heterogénea y parcial:

• Algunos componentes son accesibles solo a través de muestras generadas por simuladores o conjuntos de datos externos (cajas negras).
• Otros componentes admiten expresiones matemáticas analíticas tratables o se especifican mediante suposiciones estructurales.
• En modelos jerárquicos, la verosimilitud puede ser inaccesible, pero se dispone de muestras de una distribución que la contiene junto con una prior específica del software.

Los métodos de inferencia estándar (MCMC, Inferencia Variacional, ABC) fallan en estos escenarios porque requieren:

1. Acceso completo a la densidad no normalizada.
2. Muestras conjuntas de todo el vector objetivo.
3. O bien, solo pueden realizar muestreo o evaluación de densidad, pero no ambas tareas simultáneamente de manera eficiente.

El desafío central es reconstruir la distribución objetivo completa integrando coherentemente estas fuentes de información dispares (muestras parciales + expresiones analíticas) sin imponer suposiciones restrictivas que simplifiquen excesivamente el problema.

2. Metodología: Flujos Normalizantes de Dos Etapas (TSNF)

Los autores proponen un marco de trabajo basado en Flujos Normalizantes (Normalizing Flows), que modelan distribuciones complejas mediante transformaciones invertibles aplicadas a una distribución base simple (ej. normal estándar). El método se divide en dos etapas secuenciales:

Etapa 1: Aprendizaje de Componentes Basados en Muestras (Forward-KL)

• Objetivo: Aproximar las densidades de los componentes de la distribución objetivo para los cuales solo se dispone de muestras (ej. verosimilitudes de cajas negras o subvectores de datos).
• Procedimiento: Se entrenan $m$ modelos de Flujos Normalizantes independientes, uno para cada subvector de muestras disponible.
• Función de Pérdida: Se minimiza la divergencia KL hacia adelante ($D_{KL}(f_{real} \parallel f_{modelo})$). Esto es ideal cuando se tienen muestras de la distribución real pero no se conoce su densidad.
• Resultado: Se obtienen estimaciones de densidad $\hat{g}_i(y_i)$ que son evaluables en cualquier punto del espacio.

Etapa 2: Reconstrucción de la Distribución Objetivo (Reverse-KL)

• Objetivo: Combinar las densidades estimadas en la Etapa 1 con los términos analíticos conocidos ($h(y_0)$) para formar la distribución objetivo completa.
• Construcción: Se define una densidad aproximada no normalizada:
  $$\tilde{f}_X(x) \propto h(y_0) \prod_{i=1}^{m} \hat{g}_i(y_i)$$
• Procedimiento: Se entrena un único modelo de Flujo Normalizante sobre esta densidad compuesta $\tilde{f}_X$.
• Función de Pérdida: Se minimiza la divergencia KL inversa ($D_{KL}(f_{modelo} \parallel \tilde{f}_X)$). Esto es posible porque ahora se puede evaluar la densidad objetivo (compuesta) en cualquier punto.
• Salida Final: Un modelo capaz de:
  1. Muestrear eficientemente de la distribución objetivo (inyectando ruido base a través de la transformación inversa).
  2. Evaluar la densidad en puntos arbitrarios (calculando el determinante jacobiano de la transformación).

Consideraciones de Implementación:

• Se prioriza la estabilidad numérica en la Etapa 1 para asegurar que los determinantes jacobianos sean finitos y bien comportados, evitando inestabilidades en el entrenamiento de la Etapa 2.
• El marco es flexible y puede generalizarse a funciones $F$ que no sean simplemente productos (ej. sumas o composiciones deterministas).

3. Contribuciones Clave

1. Integración de Información Heterogénea: El método permite combinar datos de simuladores (cajas negras) con conocimiento analítico sin necesidad de acceder a la densidad conjunta completa ni a muestras conjuntas.
2. Dualidad Muestreo-Evaluación: A diferencia de métodos como ABC o MCMC estándar, el modelo resultante permite tanto la generación rápida de muestras representativas como la evaluación puntual de la densidad, facilitando la cuantificación de incertidumbre y la comparación de modelos.
3. Robustez ante Priors de Software: En modelos jerárquicos con verosimilitud inaccesible, el método corrige automáticamente el efecto de la "prior específica del software" utilizada en la generación de muestras, algo que los métodos existentes (como el algoritmo TSFB) no logran hacer eficazmente.
4. Escalabilidad: El enfoque es aplicable a problemas de alta dimensión donde los métodos no paramétricos tradicionales (como estimadores de densidad kernel) fallan debido a la maldición de la dimensionalidad.

4. Resultados de los Estudios de Simulación y Caso Real

A. Inferencia de Densidad Conjunta (Simulación):

• Se reconstruyó una densidad tridimensional altamente no lineal utilizando muestras de subvectores $(X, Y)$ y $(X, Z)$ y una expresión analítica para $X$.
• Resultado: El modelo TSNF capturó con precisión las estructuras de dependencia no lineal y las formas complejas, superando a los métodos tradicionales en la recuperación de la forma de la distribución.

B. Modelos Jerárquicos con Verosimilitud Inaccesible (Simulación):

• Se comparó TSNF con el algoritmo Two-Stage Fully Bayesian (TSFB) existente.
• Hallazgo Crítico: El método TSFB es altamente sensible a la prior del software utilizado para generar las muestras iniciales. Si la prior desplaza las muestras lejos de las regiones de alta densidad de la verosimilitud real, TSFB produce estimaciones truncadas y sesgadas.
• Ventaja TSNF: Al incorporar explícitamente la prior del software en la función $h(y_0)$ de la Etapa 2, el modelo TSNF cancela este efecto, recuperando la forma correcta de la posterior y la estructura de covarianza, independientemente de la prior del software.

C. Caso de Estudio: Astrofísica (Estrellas Enanas Blancas):

• Se aplicó el método a un problema real de inferencia jerárquica sobre 30 estrellas, utilizando el software BASE-9 (caja negra) para generar muestras de verosimilitud.
• Comparación:
  • TSFB: Mostró estimaciones marginales con falta de suavidad (discontinuidades) y superposición excesiva entre estrellas, debido a la acumulación de masa de probabilidad en puntos discretos de las propuestas de MCMH.
  • TSNF: Proporcionó estimaciones suaves, separando claramente las distribuciones marginales específicas de cada estrella y capturando la variabilidad poblacional real.
• Dado que la densidad posterior real no es evaluable directamente, la superioridad se demostró cualitativamente mediante la suavidad, la estabilidad y la capacidad de explorar el espacio de parámetros.

5. Significado y Conclusión

El marco Two-Stage Normalizing Flows representa un avance significativo para la inferencia estadística en escenarios modernos donde los modelos son complejos, parciales o computacionalmente intratables.

• Flexibilidad: Permite a los científicos utilizar toda la información disponible (simuladores, teoría física, datos parciales) sin sacrificar la fidelidad del modelo.
• Eficiencia: Una vez entrenado, el modelo permite muestreo y evaluación de densidad con un costo computacional comparable al muestreo de una distribución normal, eliminando la necesidad de cadenas MCMC largas y costosas para cada nueva consulta.
• Robustez: Mitiga los sesgos introducidos por las herramientas de simulación (priors de software), ofreciendo una aproximación más fiel a la distribución objetivo verdadera.

El trabajo sugiere que este enfoque es particularmente valioso en campos como la astrofísica, la biología y la física, donde los modelos de simulación son "cajas negras" pero se requiere una inferencia bayesiana rigurosa y completa.