Resumen Técnico: Geodésicas de Contacto de Wasserstein para Puentes de Schrödinger No Conservativos

Enunciado del Problema

El problema del Puente de Schrödinger (SB) busca la trayectoria estocástica más probable que conecta dos distribuciones marginales, típicamente con respecto a un proceso de referencia como el movimiento browniano. Aunque los SB ofrecen un marco principista para modelar procesos estocásticos en campos que van desde la dinámica celular hasta la meteorología, los métodos existentes enfrentan dos limitaciones principales:

1. Supuesto de Conservación de Energía: Las formulaciones estándar de SB asumen un perfil de energía constante a lo largo de la trayectoria estocástica. Esto restringe el modelo a sistemas conservativos, fallando en capturar fenómenos del mundo real que involucran disipación de energía o interacciones externas (por ejemplo, tormentas que pierden intensidad o diferenciación de células madre).
2. Intratabilidad Computacional: Los solucionadores actuales a menudo dependen de simulaciones estocásticas iterativas hacia adelante y hacia atrás o del refinamiento progresivo de la dinámica de referencia. Estos enfoques operan en espacios de probabilidad de dimensión infinita, lo que conduce a altos costos computacionales, convergencia lenta y sensibilidad a la inicialización.

Metodología

1. Puente de Schrödinger Generalizado No Conservativo (NCGSB)

Los autores proponen una reformulación del problema SB basada en la mecánica hamiltoniana de contacto, una generalización de la mecánica hamiltoniana clásica que permite la variación de energía.

• Aumento del Estado: El espacio de estados del sistema se aumenta con una variable escalar $z_t$ que representa la acción lagrangiana. La dinámica de $z_t$ depende recursivamente del estado actual del sistema y de un factor de amortiguamiento $\gamma$.
• Variación de Energía: A diferencia del SB Generalizado (GSB), que minimiza un funcional de costo estático, el NCGSB minimiza un estado que varía con el tiempo. El funcional de costo se define como la integral de $\partial_t z_t$, donde la evolución de $z_t$ incluye un término de amortiguamiento $-\gamma z_t$. Esta estructura permite que la energía total del sistema varíe con el tiempo, modelando naturalmente comportamientos disipativos.
• Generación Guiada: El marco soporta la generación guiada incorporando un término de pérdida específico de la tarea $\|y - f(x_{t_s})\|^2$ en la acción lagrangiana. Esto dirige el proceso generativo hacia condiciones específicas en pasos de tiempo arbitrarios sin romper la variedad de datos subyacente.

2. Geodésicas de Contacto de Wasserstein (CWG)

Para hacer tratable el problema NCGSB, los autores elevan el problema desde la variedad de Wasserstein de dimensión infinita a un espacio de dimensión finita utilizando una perspectiva geométrica.

• Reformulación Geométrica: Se demuestra que las condiciones de optimalidad del NCGSB corresponden a un flujo geodésico en una variedad extendida $P_+(M) \times \mathbb{R}$, gobernado por una Métrica Jacobi de Contacto $\tilde{g}_J$. Esta métrica incorpora energía cinética, energía potencial y un término de entropía derivado del hamiltoniano de contacto.
• Parametrización ResNet: La geodésica continua se discretiza en una secuencia de transiciones de probabilidad. Los autores mapean esta discretización a una arquitectura de Red Residual (ResNet). Cada bloque de ResNet corresponde a un segmento de la geodésica, parametrizando el mapa de empuje (pushforward) entre pasos de tiempo consecutivos.
• Solucionador No Iterativo: Al enmarcar el problema como un cálculo de geodésica en un espacio de parámetros de dimensión finita $\Theta$, el método evita los costosos bucles de iteración externa típicos de los solucionadores de Ajuste Proporcional Iterativo (IPF) o de EDE hacia adelante y hacia atrás. El objetivo de entrenamiento minimiza una función de pérdida compuesta por:
  • Ajuste de marginales terminales e intermedias (distancia Wasserstein-2).
  • Minimización de energía a lo largo de la trayectoria, ponderada por el factor de escala $\Phi_{t_k}$ derivado de la métrica Jacobi de Contacto.

3. Complejidad y Escalabilidad

El solucionador propuesto logra una complejidad casi lineal con respecto tanto a la dimensionalidad de los datos $D$ como al tamaño del lote $N$. Esto representa una mejora significativa sobre los métodos existentes que escalan cuadrática o exponencialmente con la dimensión o requieren convergencia iterativa.

Contribuciones Clave

1. Formulación NCGSB: Una formulación no conservativa novedosa del Puente de Schrödinger que modela sistemas con energía variable extendiendo el marco hamiltoniano a la dinámica hamiltoniana de contacto.
2. Marco CWG: La introducción de la Geodésica de Contacto de Wasserstein (CWG), un solucionador geométrico general compatible con todas las variantes de SB. Transforma el problema en un cálculo de geodésica tratable en una variedad de dimensión finita, permitiendo un entrenamiento rápido y no iterativo mediante arquitecturas ResNet.
3. Generación Guiada: Un método para dirigir el puente aprendido utilizando pérdidas específicas de la tarea mediante la modulación de la métrica riemanniana, permitiendo la generación condicional sin comprometer la naturaleza estocástica del puente.

Resultados Experimentales

El marco se validó en diversas tareas, demostrando un rendimiento superior en precisión, factibilidad y eficiencia computacional en comparación con las líneas base (DSBM, SB-Flow, GSBM, SBIRR, DM-SB).

• Navegación en Variedad LiDAR: CWG calculó con éxito puentes en una variedad 3D compleja (Monte Rainier) respetando las restricciones de baja altitud. Logró costos de optimalidad más bajos y una mayor factibilidad (satisfacción de límites) que las líneas base, con tiempos de entrenamiento aproximadamente 5 veces más rápidos que GSBM.
• Secuenciación de Células Únicas: En la reconstrucción de la diferenciación de células madre (un proceso biológico no conservativo), CWG superó a las líneas base mmSB en el error de Wasserstein en conjuntos de validación. La formulación de energía variable capturó la dinámica rápida de las etapas tempranas de diferenciación con mayor precisión que los modelos conservadores de energía.
• Generación de Imágenes:
  • Temperatura de la Superficie del Mar: CWG predijo mapas de temperatura mensuales intermedios con puntuaciones FID significativamente más bajas que las líneas base y logró velocidades de entrenamiento órdenes de magnitud más rápidas (por ejemplo, 17s vs. 1227s para un año específico).
  • Manipulación Robótica: El modelo generó marcos intermedios realistas para tareas de agarre y colocación de robots, superando a las líneas base en calidad de imagen (FID) y tiempo de entrenamiento.
  • Transferencia de Imágenes No Pareadas (FFHQ): En el mapeo de rostros de adultos a niños, CWG satisfizo con éxito las restricciones de límites (Factibilidad), mientras que las líneas base no lograron coincidir con la distribución objetivo a pesar de tener costos de transporte más bajos.
• Generación Guiada: El método dirigió con éxito la generación en navegación LiDAR (apuntando a lados de montaña específicos) y tareas robóticas (colocando objetos en coordenadas específicas) con una degradación mínima en la calidad de la imagen.

Significado y Afirmaciones

El artículo afirma que, al pasar de formulaciones conservativas a no conservativas mediante geometría de contacto, el marco del Puente de Schrödinger puede modelar una clase más amplia de procesos estocásticos del mundo real que involucran disipación de energía. El significado principal radica en la eficiencia computacional del solucionador CWG; al aprovechar la estructura geométrica del problema y la parametrización ResNet, elimina la necesidad de bucles iterativos costosos. Esto permite la aplicación de SBs a tareas de alta dimensión y complejas (como la generación de imágenes y la dinámica molecular) que anteriormente eran computacionalmente prohibitivas. Los autores enfatizan que su enfoque proporciona un "marco principista" que es tanto más expresivo (manejando la variación de energía) como significativamente más rápido que los métodos existentes de última generación.