An Optimal Transport-Based Generative Model for Bayesian Posterior Sampling

1. Planteamiento del Problema

El problema central en la inferencia bayesiana moderna es realizar inferencias válidas sobre los parámetros de un modelo de manera computacionalmente eficiente. La distribución posterior $\pi_n(\theta)$ se define mediante la regla de Bayes, pero su implementación práctica se ve obstaculizada por la constante de normalización (la densidad marginal de los datos), que implica una integral multidimensional a menudo intratable matemáticamente.

Las soluciones actuales tienen limitaciones significativas:

• MCMC (Cadenas de Markov Monte Carlo): Aunque exactas, sufren de mezcla lenta (slow mixing) en espacios de alta dimensión o estructuras complejas, generando muestras altamente correlacionadas que requieren un tiempo de computación excesivo.
• Inferencia Variacional (VI): Escala bien a grandes conjuntos de datos pero no garantiza el muestreo exacto de la densidad objetivo, lo que puede resultar en estimaciones sesgadas y cuantificación de incertidumbre poco fiable.
• Modelos de Transporte de Medida (anteriores): Métodos previos basados en mapas de transporte a menudo permiten mapas no únicos, lo que conduce a inestabilidad computacional y falta de interpretabilidad.

El objetivo es desarrollar un método que aprenda un mapa determinista desde una distribución de referencia simple hacia la posterior objetivo, permitiendo la generación eficiente de muestras independientes sin necesidad de simular cadenas de Markov ni calcular constantes de normalización.

2. Metodología Propuesta

La propuesta central es un modelo generativo basado en Transporte Óptimo (OT). En lugar de buscar cualquier mapa que empuje la distribución de referencia $\mu$ hacia la posterior $\pi_n$, el método restringe la búsqueda al mapa de transporte óptimo (en la formulación de Monge), que minimiza el costo cuadrático de transporte.

A. Formulación Variacional

El problema se formula como la minimización de la divergencia de Kullback-Leibler (KL) entre la distribución empujada por el mapa y la posterior objetivo:
$$\min_{T \in \mathcal{T}} d_{KL}(T_\#\mu \parallel \pi_n)$$
Gracias a una propiedad de invariancia de la divergencia KL bajo transformaciones invertibles (Lema 1), esto se reescribe como un problema de optimización que no requiere conocer la constante de normalización de $\pi_n$:
$$\max_{T \in \mathcal{T}} \mathbb{E}_{\mu} \left[ \log(\pi_n \circ T) + \log |\det(J_T)| \right]$$
donde $J_T$ es la matriz jacobiana del mapa.

B. Restricciones Estructurales (Clase de Mapas $\mathcal{T}$)

La contribución clave es la definición de una clase de mapas $\mathcal{T}$ basada en la teoría del transporte óptimo para garantizar unicidad y estabilidad:

1. Parámetros Continuos: Se aprovecha el teorema de factorización polar de Brenier, que establece que el mapa óptimo es el gradiente de una función convexa ($T^* = \nabla u^*$).

   • Se construye una clase aproximadora $\mathcal{T}_{conv}$ utilizando funciones convexas definidas como el máximo de varias "unidades convexas" (combinaciones de funciones no lineales crecientes y términos afines).
   • La estructura de "máximo" permite capturar distribuciones multimodales de manera natural, donde cada componente local corresponde a un modo de la distribución objetivo.

2. Parámetros Mixtos (Discretos y Continuos): Para modelos con variables latentes discretas (común en modelos bayesianos), el mapa óptimo puede no ser invertible.

   • Se propone una caracterización teórica (Teorema 1) donde el potencial óptimo tiene la forma $u^*(x^{(1)}, x^{(2)}) = \kappa \max_k \{ \langle x^{(1)}, b_k \rangle + \phi_k(x^{(2)}) \}$.
   • El componente discreto se selecciona mediante un mecanismo tipo "Probit multiclase" (arg max), y el componente continuo se genera mediante un mapa de transporte local $\nabla \phi_k$ condicionado a la clase discreta.

3. Adaptaciones Específicas:

   • Posteriores Aproximadamente Gaussianas: Se utiliza una clase de mapas afines, lo que reduce drásticamente la complejidad computacional manteniendo la precisión estadística (Lema 6).
   • Modelos de Variables Latentes: Se incorpora una aproximación de "campo medio" (mean-field) en la estructura del potencial para reducir la cardinalidad exponencial del espacio latente, haciendo escalable el método a grandes conjuntos de datos.

3. Contribuciones Clave

1. Unicidad e Interpretabilidad: A diferencia de los flujos normalizadores (Normalizing Flows) o mapas triangulares generales, la restricción de transporte óptimo garantiza que el mapa aprendido sea único (salvo transformaciones de medida). Esto proporciona una interpretación geométrica clara: el mapa representa el camino geodésico más corto entre la prior y la posterior en el espacio de Wasserstein.
2. Manejo de Multimodalidad y Espacios Mixtos: El método aborda explícitamente distribuciones multimodales (usando la estructura de máximo de potenciales) y espacios de parámetros mixtos (discretos/continuos), un desafío para muchos métodos de flujos normalizadores que requieren invertibilidad estricta.
3. Herramientas Inferenciales Nuevas: Al ser el mapa el gradiente de una función convexa, preserva el orden de los cuantiles. Esto permite construir:
   • Cuantiles Multivariados Centro-Exteriores: Contornos de cuantiles que capturan la geometría no convexa de la posterior.
   • Rangos y P-valores Bayesianos: Una nueva forma de evaluar la plausibilidad de valores de parámetros y realizar análisis exploratorio visual.
4. Eficiencia Computacional: Una vez entrenado el mapa, la generación de muestras es instantánea y produce muestras independientes, evitando el costo de la simulación de cadenas de Markov.

4. Resultados Experimentales

Los autores evaluaron el método en varios escenarios simulados y con datos reales, comparándolo con Transport Monte Carlo (TMC), Flujos Normalizadores (Planar, MAF, NSF) y Mapas Triangulares.

• Mezclas de Gaussianas (Multimodalidad):
  • En distribuciones multimodales de alta dimensión, los métodos de flujos normalizadores (especialmente NSF y MAF) y TMC fallaron en capturar todos los modos o perdieron la estructura de correlación.
  • El método propuesto logró la menor distancia de Wasserstein ($W_2$) en la mayoría de los casos, capturando con precisión la separación entre modos y la geometría de los contornos de cuantiles.
• Regresión Logística Bayesiana:
  • El método superó consistentemente a TMC y fue comparable o superior a los flujos normalizadores en la precisión de los intervalos de credibilidad y la selección de variables.
  • La aproximación lineal (para posteriones gaussianas) demostró ser altamente efectiva y rápida.
• Modelos de Mezclas Gaussianas (Variables Latentes):
  • El enfoque de "campo medio" permitió estimar simultáneamente las variables latentes discretas y los parámetros continuos, algo que los flujos normalizadores estándar no pueden hacer directamente.
  • Aunque la distancia $W_2$ fue ligeramente mayor que en métodos puramente continuos cuando los clusters se superponían, el método proporcionó estimaciones para las variables latentes, ofreciendo una interpretación completa del modelo.
• Análisis de Datos Reales (Set de Levadura):
  • En una tarea de clasificación binaria con genes, el método seleccionó las mismas variables significativas que el MCMC (el estándar de oro) y logró una precisión de prueba comparable (76.8%), superando a otros métodos de transporte.
  • La cuantificación de la incertidumbre (longitud de intervalos de credibilidad) fue la más cercana a la obtenida por MCMC.

5. Significado e Impacto

Este trabajo representa un avance significativo en la inferencia bayesiana al cerrar la brecha entre la precisión de los métodos basados en muestreo (MCMC) y la velocidad de los métodos variacionales.

• Nuevas Herramientas de Visualización: La capacidad de generar cuantiles multivariados y rangos basados en transporte óptimo ofrece a los investigadores herramientas poderosas para el análisis exploratorio de datos, permitiendo visualizar dependencias complejas y estructuras no convexas que los intervalos de credibilidad marginales tradicionales ocultan.
• Escalabilidad: La capacidad de manejar espacios mixtos y grandes espacios latentes mediante aproximaciones estructuradas (campo medio) abre la puerta a la aplicación de estos métodos en modelos jerárquicos complejos y grandes conjuntos de datos.
• Fundamento Teórico Sólido: Al basarse en la teoría del transporte óptimo, el método ofrece garantías teóricas sobre la unicidad de la solución y la estabilidad, resolviendo problemas de identificabilidad y entrenamiento que afectan a otros modelos generativos.

En resumen, el artículo propone un marco unificado que no solo acelera el muestreo bayesiano, sino que también enriquece el proceso inferencial con nuevas capacidades de diagnóstico y visualización geométrica.