Resumen Técnico: Inferencia Variacional mediante el Descenso de Transporte Entrópico

Planteamiento del Problema

El muestreo aproximado de distribuciones objetivo intratables $\pi(x) \propto \exp(-V(x))$ es un desafío fundamental en la inferencia bayesiana, el modelado generativo y la simulación científica. Los métodos de Inferencia Variacional basada en Partículas (ParVI) abordan esto evolucionando un conjunto de partículas que interactúan para aproximar $\pi$. Sin embargo, los enfoques dominantes existentes, como el Gradiente de Descenso de Stein (SVGD), dependen de la repulsión basada en kernels. El artículo identifica que estos métodos sufren de colapso de varianza en altas dimensiones y colapso de modo en objetivos multimodales. Estas patologías se atribuyen a la ausencia de una estructura de transporte global en las interacciones mediadas por kernels, la cual no logra coordinar eficazmente a las partículas a través de todo el espacio de estados.

Metodología: Descenso de Transporte Entrópico (ETD)

Los autores introducen el Descenso de Transporte Entrópico (ETD), una familia de algoritmos ParVI que enmarca las actualizaciones de las partículas como problemas de transporte óptimo con regularización entrópica (EOT). La metodología se deriva a través de los siguientes pasos teóricos:

1. Levantamiento del Esquema Próximo JKO: El enfoque se basa en el esquema de Jordan-Kinderlehrer-Otto (JKO), que ve el muestreo como un flujo de gradiente de la divergencia de Kullback-Leibler (KL) con respecto a la métrica de Wasserstein. En lugar de optimizar directamente sobre distribuciones, los autores "levantan" la optimización al espacio de los acoplamientos (distribuciones conjuntas entre las partículas actuales y las propuestas).
2. Relajación mediante la Regla de la Cadena de KL: Para hacer el problema tratable, la divergencia KL se descompone utilizando la regla de la cadena. Esto produce un límite superior para el objetivo original, transformando el problema en un problema de EOT semi-relajado.
3. La Familia $\tau$: El objetivo resultante está parametrizado por un escalar $\tau \ge 0$, que controla la fidelidad del acoplamiento respecto a la marginal de la medida objetivo:
   • Semi-relajado ($\tau=0$): El objetivo entra solo como una medida de referencia en el regularizador de KL.
   • Balanceado ($\tau \to \infty$): La marginal del objetivo se impone exactamente.
   • Desbalanceado ($0 < \tau < \infty$): Una interpolación entre ambos extremos.
4. Implementación Algorítmica: Cada iteración implica:
   • Proponer: Generar posiciones candidatas (propuestas) $\{y_j\}$, ya sea mediante un paseo aleatorio o un paso guiado por el score (Euler-Maruyama).
   • Acoplar: Resolver el problema de EOT entre las partículas actuales y las propuestas utilizando el algoritmo de Sinkhorn. Esto computa un plan de transporte (matriz de acoplamiento) que minimiza el costo de transporte más la regularización entrópica.
   • Actualizar: Remuestrear las partículas basándose en la distribución condicional definida por el plan de transporte.

Crucialmente, el ETD puede operar sin necesidad de score (score-free), requiriendo solo evaluaciones puntuales de la densidad no normalizada del objetivo, aunque también puede incorporar información del score.

Contribuciones Clave

El artículo realiza tres contribuciones teóricas y prácticas:

1. Marco Algorítmico: Introduce el ETD, una familia de ParVI derivada del esquema JKO que reduce cada iteración a un cómputo de Sinkhorn. Esto proporciona un mecanismo de coordinación global para las partículas, a diferencia de la repulsión local por kernel.
2. Caracterización Teórica de la Estacionariedad:
   • Teorema 2: Los autores caracterizan la distribución estacionaria del ETD balanceado. Demuestran que el sesgo es independiente del costo de transporte y de la regularización entrópica, dependiendo únicamente del ancho de banda de la propuesta.
   • Teorema 3: Proponen pesos de objetivo corregidos por importancia ($b_j \propto \pi(y_j)/q_\mu(y_j)$) que eliminan este sesgo por completo, haciendo que la distribución objetivo $\pi$ sea exactamente estacionaria sin requerir un paso de corrección de Metropolis-Hastings.
3. Desempeño Empírico: Los experimentos demuestran que el ETD iguala o supera a SVGD, AGF-SVGD y SGLD, con mejoras significativas en entornos de alta dimensión y multimodales.

Resultados Experimentales

El artículo evalúa el ETD en cuatro categorías de pruebas de referencia (benchmarks):

• Diagnósticos de Colapso de Varianza: En objetivos gaussianos isotrópicos en dimensiones $d \in \{10, \dots, 200\}$, SVGD exhibe un rápido colapso de varianza (DAMV $\approx 0.11$ en $d=200$), mientras que las variantes de ETD logran una recuperación de varianza casi perfecta (DAMV $\approx 1.0$), igualando o superando a SGLD.
• Funciones de Energía Multimodales: En objetivos 2D con estructuras complejas (por ejemplo, anillos, modos sinusoidales), SVGD y AGF-SVGD sufren de colapso de modo. Las variantes de ETD, particularmente aquellas que utilizan la corrección de muestreo por importancia, capturan con éxito la estructura multimodal completa y logran distancias de energía más bajas.
• Regresión Logística Bayesiana: En el conjunto de datos Covertype ($d=56$), mientras que los puntajes de NLL son comparables entre métodos, las variantes de ETD muestran una cobertura marginal significativamente superior (0.97–0.99) en comparación con SVGD y SGLD (0.27–0.32), lo que indica una mejor cuantificación de la incertidencia posterior.
• Redes Neuronales Bayesianas y Muestreo Molecular: El ETD supera a los modelos de base en tareas de regresión UCI y en distribuciones de Boltzmann molecular (Lennard-Jones 13). Notablemente, en el objetivo LJ-13, los métodos de base (SVGD, SGLD, AGF-SVGD) produjeron energías numéricamente divergentes, mientras que el ETD produjo muestras físicamente significativas con una baja distancia de variación total.

Significado y Reivindicaciones

El artículo afirma que el ETD aborda las limitaciones fundamentales de los métodos ParVI basados en kernels al introducir una estructura de transporte global. Al mediar las interacciones a través del transporte óptimo entrópico, el ETD preserva naturalmente la estructura multimodal y evita el colapso de varianza en altas dimensiones.

Los autores enfatizan que el ETD ofrece un marco flexible que puede operar sin información de gradiente (score-free), lo que lo hace aplicable a problemas donde los gradientes no están disponibles o son costosos. Además, el resultado teórico de que el sesgo estacionario puede eliminarse mediante una simple corrección de pesos (Teorema 3) proporciona un camino principista hacia el muestreo exacto sin la carga computacional de los pasos de Metropolis-Hastings.

El artículo reconoce las limitaciones, señalando que el costo por iteración es $O(NML)$ (donde $L$ es el número de iteraciones de Sinkhorn), el cual es mayor que el de SVGD ($O(N^2)$) o SGLD ($O(N)$). Adicionalmente, al igual que todos los métodos que dependen de distancias euclidianas, el ETD enfrenta la maldición de la dimensionalidad, ya que el costo de transporte concentra el kernel de Gibbs en altas dimensiones, aunque los autores sugieren que el trabajo futuro podría abordar esto mediante costos aprendidos o métricas geométricas.