Resumen Técnico: Dinámicas Estocásticas de Datos-a-Energía

Planteamiento del Problema

El problema del puente de Schrödinger (SB, por sus siglas en inglés) busca encontrar un sistema dinámico estocástico que conecte dos distribuciones marginales, $p_0$ y $p_1$, minimizando un costo de transporte relativo a un proceso de referencia (típicamente una ecuación diferencial estocástica de Itô). Este problema generaliza el transporte óptimo al entorno estocástico y fundamenta técnicas modernas de modelado generativo como los modelos de difusión y el flow matching.

Los algoritmos existentes para resolver el problema SB, particularmente aquellos basados en el procedimiento de Ajuste Proporcional Iterativo (IPF), generalmente requieren acceso a muestras de ambas distribuciones marginales. Esta limitación impide la aplicación de los métodos SB en escenarios donde una o ambas distribuciones están definidas únicamente por densidades no normalizadas (funciones de energía) $p(x) = e^{-E(x)}/Z$, donde la función de partición $Z$ es desconocida y no hay muestras disponibles directamente. Esta brecha es crítica para aplicaciones como la inferencia bayesiana de la distribución posterior en espacios latentes de modelos generativos, donde la posterior suele conocerse solo hasta una constante de normalización mediante una recompensa o función de verosimilitud.

Metodología

El artículo propone el IPF de Datos-a-Energía (y de Energía-a-Energía), un método general para modelar puentes de Schrödinger cuando las distribuciones se presentan como densidades no normalizadas. El enfoque extiende el algoritmo IPF estándar integrando técnicas de aprendizaje por refuerzo off-policy y muestreo de difusión.

1. Marco de Trabajo IPF Generalizado

El algoritmo IPF estándar actualiza iterativamente los procesos hacia adelante y hacia atrás ($\overrightarrow{P}_t$ y $\overleftarrow{P}_t$) resolviendo problemas de medio puente (half-bridge).

• Paso hacia atrás (Datos-a-Datos): Cuando se dispone de muestras de $p_0$, el proceso hacia atrás se entrena mediante máxima verosimilitud (paso estándar de IPF).
• Paso hacia adelante (Datos-a-Energía): Cuando $p_1$ está dada por una densidad no normalizada $e^{-E_1(x)}$, el proceso hacia adelante no puede entrenarse mediante verosimilitud estándar. En su lugar, los autores introducen una pérdida de varianza de logaritmo (gradiente de varianza condicional o VarGrad). Esta pérdida minimiza la varianza del logaritmo de la razón entre la probabilidad de la trayectoria hacia adelante y la densidad objetivo, ponderada por la probabilidad de la trayectoria hacia atrás:
  $$\text{LLV}(x_0, \theta) = \text{Var} \left( \log \frac{\overrightarrow{p}_\theta(\tau | x_0)}{p_1(x_1) \overleftarrow{p}_\phi(\tau | x_1)} \right)$$
  Este objetivo impone la proporcionalidad requerida por la solución del SB sin necesidad de muestras de $p_1$.

2. Estrategias de Aprendizaje por Refuerzo Off-Policy

Para abordar el desafío de muestrear desde densidades no normalizadas complejas y de alta dimensión donde el entrenamiento on-policy falla en explorar todos los modos, los autores adaptan técnicas off-policy:

• Buffer de Replay (Replay Buffer): Un buffer almacena muestras finales $x_1$ del proceso hacia adelante. Durante el entrenamiento, $x_0$ se genera muestreando una $x_1$ del buffer y ejecutando una trayectoria hacia atrás. Esto enfoca la exploración en regiones de alta densidad del objetivo.
• Actualizaciones de Langevin: El buffer se actualiza periódicamente utilizando la dinámica de Langevin no ajustada sobre la densidad objetivo $p_1$ para corregir ajustes imperfectos.
• Reutilización de Trayectorias: El algoritmo reutiliza trayectorias hacia atrás para formar lotes de trayectorias que comparten un punto inicial $x_0$, mejorando la eficiencia de la estimación de la pérdida basada en la varianza.
• Mezcla (Mixing): La política de entrenamiento mezcla muestras on-policy con muestras off-policy del buffer para equilibrar la exploración y la explotación.

3. Coeficientes de Difusión Aprendibles

Una contribución metodológica secundaria implica generalizar el algoritmo IPF para aprender no solo las funciones de deriva (drift), sino también los coeficientes de difusión ($\sigma_t$) de la dinámica. Al reemplazar las varianzas fijas en los pasos de Euler-Maruyama discretizados por funciones aprendibles, el método busca corregir los errores de discretización temporal, un concepto inspirado en trabajos recientes sobre muestreadores de difusión.

Contribuciones Clave

1. Primer Método de SB de Datos-a-Energía General: El artículo introduce el primer algoritmo capaz de inferir puentes de Schrödinger cuando una o ambas distribuciones se proporcionan como densidades no normalizadas sin acceso a muestras. Logra aprender puentes estocásticos con bajo costo de transporte entre conjuntos de datos sintéticos y densidades, con un rendimiento similar al IPF de Datos-a-Datos que utiliza muestras reales.
2. Coeficientes de Difusión Aprendibles: Los autores demuestran que los algoritmos actuales de IPF de Datos-a-Datos pueden mejorar sustancialmente aprendiendo el coeficiente de difusión además de la deriva. Esto generaliza los hallazgos previos para los muestreadores de difusión al contexto más amplio de SB, mostrando que la varianza aprendible permite un modelado más preciso, especialmente con discretizaciones temporales gruesas.
3. Muestreo de la Posterior Outsourced: El método se aplica al problema de muestrear distribuciones posteriores en los espacios latentes de modelos generativos (ej. GANs, VAEs). Al modelar un puente entre la distribución latente previa y una posterior ponderada por recompensa, los autores crean un método de traducción imagen-a-imagen libre de datos. Este enfoque preserva el contenido semántico (estilo, fondo) que no está restringido por la condición objetivo, a diferencia de los muestreadores de difusión estándar.

Resultados Experimentales

La eficacia del método propuesto se demuestra a través de varios entornos de prueba:

• Benchmarks Sintéticos 2D: En tareas como de Gaussiano a Mezcla de Gaussianos (GMM) y distribuciones "Two Moons", el IPF de Datos-a-Energía logra un rendimiento comparable al de IPF de Datos-a-Datos. La inclusión de varianza aprendible reduce significamente la distancia Wasserstein-2 y el Path KL, especialmente cuando el número de pasos de discretización es pequeño.
• Puente de Energía-a-Energía: Los resultados preliminares muestran que el algoritmo puede conectar con éxito dos distribuciones definidas únicamente por funciones de energía, sin acceso a muestras de ninguna de ellas.
• Traducción de Imágenes (MNIST y CIFAR-10): En experimentos de espacio latente utilizando SN-GAN y StyleGAN, el método traduce muestras previas a posteriores de clase objetivo (ej. generar "perros" o "cincos").
  • Métricas: El método logra puntuaciones FID más bajas en comparación con el muestreo por rechazo y la dinámica de Langevin en varias clases.
  • Cualitativo: La inspección visual confirma que el puente preserva las características de estilo (ej. grosor del dígito, orientación del fondo) mientras transforma el contenido a la clase objetivo, lo que sugiere una preservación efectiva de la información semántica no restringida.
• Estudios de Ablación: Los experimentos confirman que los buffers de replay, las actualizaciones de Langevin y la reutilización de trayectorias mejoran significativamente las métricas de Path KL. La reutilización de trayectorias hacia atrás mejora el Path KL pero puede impactar negativamente la media de la recompensa del logaritmo, sugiriendo un compromiso (trade-off) entre la cobertura de modos y el costo de transporte.

Significado y Reivindicaciones

El artículo afirma proporcionar un marco general para puentes de Schrödinger de datos-a-energía, llenando un vacío en la literatura donde los métodos SB estaban previamente restringidos a entornos con acceso a muestras para ambas marginales.

Los autores posicionan su trabajo como un paso significativo hacia el modelado generativo escalable y libre de datos. Al aprovechar las técnicas de aprendizaje por refuerzo off-policy, el método permite el aprendizaje de la dinámica de transporte óptimo en espacios latentes de alta dimensión donde solo se dispone de funciones de energía. La capacidad de aprender coeficientes de difusión mejora aún más la robustez de estas dinámicas frente a los errores de discretización.

Los autores señalan modestamente que, aunque el método escala a tareas de generación de imágenes de alta dimensión, se requiere trabajo futuro para abordar el colapso de modo (mode collapse) y mejorar la cobertura de modos. También identifican posibles direcciones futuras, incluyendo el ajuste fino de recompensas de imagen condicionadas por texto y puentes de Schrödinger discretos, pero no afirman éxito inmediato en estas áreas inexploradas.