Generalized Schrödinger Bridge on Graphs

Resumen Técnico: Puente de Schrödinger Generalizado en Grafos (GSBoG)

Planteamiento del Problema
El transporte en grafos es un desafío fundamental en dominios que van desde las cadenas de suministro y la infraestructura de tráfico hasta la cinética molecular. Si bien el transporte óptimo (OT) en grafos proporciona métodos fundamentados para el emparejamiento de distribuciones de probabilidad de origen y destino, los enfoques existentes suelen producir acoplamientos estáticos que no logran describir el movimiento de masa dependiente del tiempo ni proporcionar políticas de control ejecutables. Además, las formulaciones dinámicas para el transporte en grafos, como aquellas que resuelven flujos hamiltonianos en el simplex de probabilidad, sufren de una baja escalabilidad en topologías grandes y dispersas, y a menudo carecen de la expresividad para optimizar trayectorias intermedias bajo costos específicos del estado (por ejemplo, congestión o barreras de energía). El desafío central es construir una evolución que respete la topología y sea factible, que controle las distribuciones intermedias mientras minimiza los costos de la red y acomoda restricciones operativas como los límites de capacidad.

Metodología
Los autores presentan el Puente de Schrödinger Generalizado en Grafos (GSBoG), un marco de trabajo basado en datos para aprender políticas de cadenas de Markov de tiempo continuo (CTMC) controladas y ejecutables en grafos arbitrarios. El método extiende los principios del Puente de Schrödinger Generalizado (GSB) de espacios de estados continuos a dominios discretos estructurados en grafos.

1. Formulación Matemática:

   • El problema se formula como la minimización de un funcional de costo sobre medidas de trayectoria de una CTMC. El objetivo combina un costo de ejecución $f_t(x, p_t)$ (dependiente del estado y la distribución marginal) con un término de divergencia de Kullback-Leibler (KL) que regulariza la dinámica controlada hacia un proceso CTMC de referencia.
   • Los autores derivan condiciones de optimalidad discretas paralelas a la formulación de estado continuo. Establecen que las tasas de transición controladas óptimas $u^*_t(y, x)$ satisfacen una transformación de Hopf-Cole: $u^*_t(y, x) = r_t(y, x) \exp(-V_t(y) + V_t(x))$, donde $V_t$ es una función de potencial dependiente del tiempo y $r_t$ es la tasa de referencia.
   • Esto conduce a un sistema acoplado de ecuaciones para los potenciales de Schrödinger $(\phi_t, \hat{\phi}_t)$, análogo a la ecuación de Fokker-Planck en entornos continuos, pero definido a través de la ecuación de continuidad del grafo.

2. Enfoque Algorítmico (gIPF y TD):

   • Para evitar la resolución de problemas computacionalmente prohibitivos mediante solvers globales sobre todo el grafo y el horizonte temporal, el GSBoG adopta un enfoque basado en partículas y orientado a datos utilizando el Ajuste Proporcional Iterativo Generalizado (gIPF).
   • El método utiliza la fórmula de Dynkin para derivar análogos discretos de las Ecuaciones Diferenciales Estocásticas de Retroalimentación hacia Adelante (FBSDEs) utilizadas en el GSB continuo.
   • El objetivo de entrenamiento consiste en dos componentes alternos:
     • Pérdida IPF: Minimiza la discrepancia entre la medida de trayectoria aprendida y el proceso de referencia para satisfacer las restricciones de marginales de los extremos ($p_0 = \mu, p_1 = \nu$).
     • Pérdida de Diferencia Temporal (TD): Fuerza explícitamente la consistencia con los costos de ejecución dependientes del estado y la distribución, penalizando las desviaciones entre los incrementos predichos por el generador y los incrementos reales de la trayectoria.
   • El algoritmo parametriza los log-potenciales utilizando redes neuronales (o embeddings) y los actualiza mediante descenso de gradiente estocástico, evaluando las tasas solo en vecindades locales del grafo.

Contribuciones Clave

• Primer GSB basado en datos en grafos: El artículo presenta la primera formulación del Puente de Schrödinger Generalizado en grafos arbitrarios de manera basada en datos, distinta de los trabajos recientes de difusión discreta que se centran en la generación de grafos.
• Aprendizaje de Políticas Escalable: Al aprender políticas a nivel de trayectoria mediante el muestreo basado en partículas en lugar de resolver flujos de tiempo expandido globales, el GSBoG logra escalabilidad en grafos grandes y dispersos donde los métodos tradicionales de OT dinámico fallan debido a restricciones de memoria.
• Dinámica Consciente del Costo: El marco optimiza explícitamente el comportamiento intermedio bajo costos de ejecución generales (por ejemplo, penalizaciones por congestión), asegurando que las políticas resultantes no solo sean topológicamente factibles, sino también operativamente eficientes.
• Derivación Teórica: Los autores proporcionan una derivación rigurosa de las condiciones de optimalidad discretas y un objetivo de aprendizaje tratable basado en el generador de la CTMC, cerrando la brecha entre la teoría de control estocástico y el transporte en grafos discretos.

Resultados Experimentales
Los autores evalúan el GSBoG en tres topologías de grafos del mundo real desafiantes:

1. Red de Cadena de Suministro (9,559 nodos): El GSBoG aprendió con éxito políticas de enrutamiento que coincidían con las distribuciones de oferta/demanda de los extremos, reduciendo significativamente la congestión intermedia y la ocupación máxima de nodos en comparación con líneas base como Attraction Flow y Dynamic W1 flow. Notablemente, el GSBoG fue el único método que mantuvo los flujos de aristas inducidos dentro de los límites de capacidad sin restricciones estrictas explícitas, mientras que los solvers globales (GrSB) no pudieron completarse debido al agotamiento de la memoria.
2. Tarea de Asignación Balanceada: En grafos bipartitos con tamaños variables ($n=6$ a $20$), el GSBoG logró un solapamiento casi perfecto con la solución de OT estático óptimo (100% de precisión para $n \le 10$) mientras aprendía un acoplamiento suave que minimizaba los costos de asignación.
3. Dinámica Molecular Discretizada (Plegamiento de Chignolin): Aplicado a un Modelo de Estados de Markov (MSM) de plegamiento de proteínas, el GSBoG dirigió las transiciones de eventos raros desde un estado desplegado hacia un estado plegado. Logró una tasa de plegamiento del 99.7% con una barrera de energía significativamente reducida (1.80 $k_B T$) en comparación con la dinámica no controlada o los sesgos guiados por comitentes, y mejoró aún más a 1.39 $k_B T$ al incorporar costos de ejecución de energía libre.

Análisis de Escalabilidad
Los experimentos demuestran que el GSBoG escala linealmente con el factor de ramificación local en lugar del número total de nodos. En grafos con hasta $10^6$ nodos, el GSBoG se mantuvo factible con un bajo uso de memoria (1.4 GB) y tiempo de ejecución (1 hora), mientras que los métodos de OT dinámico se volvieron inviables (N/A) o requirieron significativamente más recursos (por ejemplo, 22 GB de memoria, 4.6 horas).

Significancia y Reivindicaciones
El artículo afirma que el GSBoG abre nuevos caminos para el transporte dinámico consciente del costo en grafos generales. Al ir más allá de los acoplamientos estáticos y los solvers globales, el marco permite el aprendizaje de políticas ejecutables y que respetan la topología que optimizan tanto el emparejamiento de los extremos como los costos de las trayectorias intermedias. Los autores posicionan este trabajo como un paso hacia la integración de la dinámica moderna basada en el aprendizaje con los principios del transporte óptimo para sistemas de red complejos, particularmente en escenarios donde las restricciones operativas (como la congestión o las barreras de energía) son críticas. El trabajo es modesto en sus limitaciones, señalando que actualmente asume una topología de grafo fija y totalmente conocida, e identifica las extensiones a grafos variables en el tiempo o parcialmente observados como trabajo futuro.