Stable Global Weighting of Flow Mixtures using Simplex Exponential Moving Average

Resumen Técnico: Ponderación Global Estable de Mezclas de Flujos mediante el Promedio Móvil Exponencial en el Simpléx

Planteamiento del Problema
Los flujos normalizadores han expandido significativamente la flexibilidad de la inferencia variacional (VI) al transformar distribuciones base simples en densidades de probabilidad expresivas. Sin embargo, las arquitecturas de flujo individuales (por ejemplo, RealNVP, MAF, RBIG) a menudo fallan al generalizar a través de geometrías posteriores heterogéneas. Los diferentes flujos poseen distintos sesgos inductivos, dinámicas de optimización y características numéricas, lo que conduce a una situación en la que ninguna arquitectura única se desempeña de manera consistente a través de diversas distribuciones objetivo. Si bien los enfoques basados en mezclas ofrecen una estrategia natural para combinar flujos complementarios, los marcos existentes sufren limitaciones significativas:

1. Optimización Conjunta: El aprendizaje de los pesos de la mezcla y los parámetros de los expertos simultáneamente suele provocar una optimización inestable o "dominación del experto" debido al acoplamiento de las heterogéneas dinámicas de optimización.
2. Ponderación Global Fija: Estimar coeficientes fijos durante el entrenamiento y mantenerlos constantes durante la inferencia carece de la flexibilidad necesaria para explotar el comportamiento complementario de los expertos, lo que a menudo resulta en una asignación ineficiente de la capacidad.

En consecuencia, existe la necesidad de un marco de optimización estable y con principios que pueda combinar expertos heterogéneos sin requerir conocimiento previo de la arquitectura más apropiada o incurrir en la inestabilidad del entrenamiento conjunto.

Metodología: AMF-VI-sEMA
El artículo propone AMF-VI-sEMA (Mezcla Adaptativa de Flujos - Inferencia Variacional con Promedio Móvil Exponencial en el Simpléx), un marco de dos etapas diseñado para desacoplar el aprendizaje de los expertos de la optimización de la mezcla.

• Etapa 1: Especialización Independiente de Expertos
  Un conjunto heterogéneo de expertos (específicamente RealNVP, MAF y RBIG) se entrena de forma independiente para especializarse en distintos regímenes estructurales. Cada experto $q_k(z|\phi_k)$ es optimizado para maximizar su propio log-verosimilitud sobre los datos de entrenamiento. Una vez entrenado, los parámetros del experto $\phi_k$ se congelan. Esta etapa asegura que cada componente converja bajo sus propias dinámicas de optimización, evitando las patologías de coordinación de la optimización conjunta.

• Etapa 2: Ponderación Global sEMA (Promedio Móvil Exponencial en el Simpléx)
  Con los parámetros de los expertos congelados, el marco aprende un vector de peso de mezcla global $\pi$ en el simpléx de probabilidad. Esto se logra mediante un mecanismo impulsado por la verosimilitud y controlado por temperatura:

1. Cálculo de Responsabilidad: Para un lote de datos de validación mantenidos fuera del conjunto de entrenamiento, se computan las responsabilidades $r_k(z)$ utilizando un softmax de las log-verosimilitudes promedio de los expertos, escaladas por un parámetro de temperatura $\tau$.
2. Suavizado y Promedio: Las responsabilidades se promedian sobre el lote de validación y, opcionalmente, se suavizan hacia una distribución uniforme para amortiguar los transitorios.
3. Actualización de sEMA en el Simpléx: Los pesos globales $\pi$ se actualizan mediante un Promedio Móvil Exponencial (EMA) directamente en el simpléx: $\pi \leftarrow \alpha \pi + (1-\alpha)\bar{r}$. Se aplica un pequeño piso $\epsilon$ para evitar el cero, seguido de una renormalización.

Este diseño crea un mecanismo de compuerta tratable y agnóstico a los datos que reasigna la capacidad de forma adaptativa sin requerir compuertas por muestra o retropropagación a través de los pesos. El marco monitorea el número efectivo de expertos ($N_{eff} = \exp(H(\pi))$) para diagnosticar el colapso de los componentes.

Contribuciones Clave

1. Marco de Optimización Desacoplado: El artículo introduce una estrategia de dos etapas que separa el aprendizaje de los expertos de la adaptación de la mezcla. Al congelar los expertos antes de actualizar los pesos, evita la inestabilidad de la optimización ponderada por responsabilidad conjunta, permaneciendo agnóstico a la arquitectura.
2. Mecanismo Práctico de Ponderación Adaptativa: Los autores proponen una actualización de sEMA en el simpléx, computacionalmente eficiente y basada en la responsabilidad. Este método estima los pesos de la mezcla global sin requerir optimización conjunta o compuertas por muestra, proporcionando una cooperación estable de los expertos con un gasto computacional mínimo.
3. Evaluación Empírica Exhaustiva: El marco se evalúa en diez problemas de referencia, incluyendo seis familias sintéticas 2D canónicas (Banana, X-Shaped, Bimodal, Multimodal, Two-moons, Rings) y cuatro objetivos bayesianos reales/de baja dimensión (BLR, BPR, Weibull, Real-GMM2). La evaluación se extiende más allá de la log-verosimilitud estándar para incluir la distancia Wasserstein-2, la Discrepancia de Media Máxima (MMD), el análisis de calibración y la robustez a través de semillas aleatorias.
4. Análisis de la Dinámica de Optimización: El artículo proporciona un análisis de la utilización de los expertos, la evolución de los pesos de la mezcla y la sensibilidad de los hiperparámetros, ofreciendo ideas sobre por qué las mezclas de flujos heterogéneos tienen éxito y cómo el mecanismo sEMA estabiliza el entrenamiento.

Resultados
Las evaluaciones empíricas demuestran que AMF-VI-sEMA logra mejoras consistentes sobre su predecesor (AMF-VI) y evita los fallos catastróficos de transporte observados en los modelos base de un solo flujo (como ResFlow y EM-Mix).

• Desempeño: El marco logra el mejor o el segundo mejor NLL (Log-Verosimilitud Negativa) en todos los conjuntos de datos sintéticos y reales. Mantiene un comportamiento de transporte robusto (Wasserstein-2) y métricas de discrepancia (MMD), evitando los fallos geométricos vistos en arquitecturas únicas (por ejemplo, NICE y ResFlow muestran altos errores de transporte en geometrías curvas).
• Estabilidad: El método mantiene trayectorias de peso estables con un número efectivo de expertos $N_{eff} > 1.4$ en todos los conjuntos de datos, lo que indica que la mezcla no colapsa hacia un solo experto.
• Robustez: Las evaluaciones de cruce de semillas muestran una baja varianza en NLL, divergencia KL y distancia de Wasserstein, confirmando que el marco es reproducible y robusto a la inicialización aleatoria.
• Adaptabilidad: Los pesos globales aprendidos se adaptan a la geometría objetivo. Por ejemplo, en posteriores de estilo regresión (BLR, BPR), el mecanismo suprime a RealNVP en favor de MAF y RBIG, mientras que en estructuras de anillo (Rings), concentra fuertemente la masa en RealNVP.

Significancia y Reivindicaciones
El artículo afirma que AMF-VI-sEMA proporciona una estrategia de optimización robusta y fiable para la aproximación de la posterior a través de paisajes heterogéneos sin requerir conocimiento previo de qué arquitectura de flujo es la más apropiada. La significancia del trabajo radica en:

• Estabilidad: El mecanismo sEMA introduce un efecto de regularización temporal que estabiliza las dinámicas de compuerta y mitiga el cambio volátil entre expertos.
• Simplicidad y Eficiencia: El enfoque es agnóstico a la arquitectura, no requiere modificaciones en los modelos de flujo subyacentes e introduce un gasto computacional insignificante en comparación con el entrenamiento conjunto.
• Fiabilidad: Al evitar la fragilidad del entrenamiento conjunto ponderado por responsabilidad, el marco ofrece una solución práctica para modelar geometrías posteriores complejas, no convexas y multimodales.

Los autores señalan que, si bien las arquitecturas de un solo flujo especializado pueden seguir siendo preferibles para objetivos unimodales simples, AMF-VI-sEMA supera consistentemente a estas en robustez a través de diversas geometrías. Se sugiere como trabajo futuro explorar estrategias de compuerta condicional o amortizada y extensiones a entornos de mayor dimensión.