Generalized Schrödinger Bridge on Graphs

Résumé Technique : Pont de Schrödinger Généralisé sur Graphes (GSBoG)

Énoncé du Problème
Le transport sur les graphes est un défi fondamental dans des domaines allant de la chaîne d'approvisionnement et des infrastructures de trafic à la cinétique moléculaire. Bien que le transport optimal (OT) sur les graphes fournisse des méthodes rigoureuses pour l'appariement de distributions de probabilité sources et cibles, les approches existantes produisent souvent des couplages statiques qui ne parviennent pas à décrire le mouvement de masse dépendant du temps ou à fournir des politiques de contrôle exécutables. De plus, les formulations dynamiques pour le transport sur graphes, telles que celles résolvant des flux hamiltoniens sur le simplexe de probabilité, souffrent d'une faible scalabilité sur des topologies larges et creuses, et manquent souvent d'expressivité pour optimiser les trajectoires intermédiaires sous des coûts dépendants de l'état (par exemple, la congestion ou les barrières énergétiques). Le défi central consiste à construire une évolution respectant la topologie et réalisable, capable de contrôler les distributions intermédiaires tout en minimisant les coûts du réseau et en tenant compte des contraintes opérationnelles comme les limites de capacité.

Méthodologie
Les auteurs introduisent le Pont de Schrödinger Généralisé sur Graphes (GSBoG), un cadre fondé sur les données pour apprendre des politiques de chaînes de Markov à temps continu (CTMC) exécutables et contrôlées sur des graphes arbitraires. La méthode étend les principes du Pont de Schrödinger Généralisé (GSB) des espaces d'états continus vers des domaines discrets structurés en graphes.

1. Formulation Mathématique :

   • Le problème est formulé comme la minimisation d'une fonctionnelle de coût sur des mesures de chemin d'une CTMC. L'objectif combine un coût courant $f_t(x, p_t)$ (dépendant de l'état et de la distribution marginale) avec un terme de divergence de Kullback-Leibler (KL) qui régularise la dynamique contrôlée vers un processus CTMC de référence.
   • Les auteurs dérivent les conditions d'optimalité discrètes parallèles à la formulation en état continu. Ils établissent que les taux de transition contrôlés optimaux $u^*_t(y, x)$ satisfont une transformation de Hopf-Cole : $u^*_t(y, x) = r_t(y, x) \exp(-V_t(y) + V_t(x))$, où $V_t$ est une fonction de potentiel dépendante du temps et $r_t$ est le taux de référence.
   • Cela conduit à un système d'équations couplées pour les potentiels de Schrödinger $(\phi_t, \hat{\phi}_t)$, analogues à l'équation de Fokker-Planck dans les contextes continus, mais définis via l'équation de continuité du graphe.

2. Approche Algorithmique (gIPF et TD) :

   • Pour éviter des solveurs globaux coûteux en calcul sur l'ensemble du graphe et de l'horizon temporel, le GSBoG adopte une approche basée sur les particules et les données utilisant l'ajustement proportionnel itératif généralisé (gIPF).
   • La méthode utilise la formule de Dynkin pour dériver des analogues discrets des équations différentielles stochastiques antérieures-postérieures (FBSDE) utilisées dans le GSB continu.
   • L'objectif d'entraînement se compose de deux composantes alternées :
     • Perte IPF : Minimise l'écart entre la mesure de chemin apprise et le processus de référence pour satisfaire les contraintes marginales des extrémités ($p_0 = \mu, p_1 = \nu$).
     • Perte de Différence Temporelle (TD) : Impose explicitement la cohérence avec les coûts courants dépendants de l'état et de la distribution en pénalisant les écarts entre les incréments prédits par le générateur et les incréments réels de la trajectoire.
   • L'algorithme paramètre les log-potentiels à l'aide de réseaux de neurones (ou d'embeddings) et met à jour ces derniers via une descente de gradient stochastique, en évaluant les taux uniquement sur les voisinages locaux du graphe.

Contributions Clés

• Premier GSB fondé sur les données sur les graphes : L'article présente la première formulation du Pont de Schrödinger Généralisé sur des graphes arbitraires de manière pilotée par les données, se distinguant des récents travaux de diffusion sur graphes qui se concentrent sur la génération de graphes.
• Apprentissage de Politiques Scalable : En apprenant des politiques au niveau de la trajectoire via l'échantillonnage par particules plutôt qu'en résolvant des flux temporels globaux, le GSBoG atteint une scalabilité sur de grands graphes creux là où les méthodes traditionnelles d'OT dynamique échouent en raison des contraintes de mémoire.
• Dynamique Sensible aux Coûts : Le cadre optimise explicitement le comportement intermédiaire sous des coûts courants généraux (ex: pénalités de congestion), garantissant que les politiques résultantes sont non seulement topologiquement réalisables mais aussi opérationnellement efficaces.
• Dérivation Théorique : Les auteurs fournissent une dérivation rigoureuse des conditions d'optimalité discrètes et un objectif d'apprentissage traitable basé sur le générateur de la CTMC, comblant le fossé entre la théorie du contrôle stochastique et le transport sur graphes discrets.

Résultats Expérimentaux
Les auteurs évaluent le GSBoG sur trois topologies de graphes réels exigeantes :

1. Réseau de Chaîne d'Approvisionnement (9 559 nœuds) : Le GSBoG a réussi à apprendre des politiques de routage qui correspondent aux distributions d'offre/demande terminales tout en réduisant considérablement la congestion intermédiaire et l'occupation de pointe des nœuds par rapport aux modèles de référence comme l'Attraction Flow et le Dynamic W1 flow. Notamment, le GSBoG était la seule méthode capable de maintenir les flux de bordure induits dans les limites de capacité sans contraintes strictes explicites, alors que les solveurs globaux (GrSB) ont échoué en raison de l'épuisement de la mémoire.
2. Tâche d'Assignation Équilibrée : Sur des graphes bipartis de tailles variables ($n=6$ à $20$), le GSBoG a obtenu un chevauchement quasi parfait avec la solution OT statique optimale (100 % de précision pour $n \le 10$) tout en apprenant un couplage souple qui minimise les coûts d'assignation.
3. Dynamique Moléculaire Discrétisée (Repliement de la Chignolin) : Appliqué à un modèle d'état de Markov (MSM) du repliement des protéines, le GSBoG a dirigé les transitions d'événements rares d'un état déplié vers un état replié. Il a atteint un taux de repliement de 99,7 % avec une barrière énergétique significativement réduite (1,80 $k_B T$) par rapport aux dynamiques non contrôlées ou aux biais guidés par le committor, et s'est encore amélioré à 1,39 $k_B T$ en incorporant des coûts courants d'énergie libre.

Analyse de Scalabilité
Les expériences démontrent que le GSBoG passe à l'échelle linéairement avec le facteur de branchement local plutôt qu'avec le nombre total de nœuds. Sur des graphes comprenant jusqu'à $10^6$ nœuds, le GSBoG est resté réalisable avec une faible utilisation de la mémoire (1,4 Go) et un temps d'exécution (1 heure), tandis que les méthodes d'OT dynamique devenaient infaisables (N/A) ou nécessitaient beaucoup plus de ressources (ex: 22 Go de mémoire, 4,6 heures).

Signification et Revendications
L'article affirme que le GSBoG ouvre de nouvelles voies pour le transport dynamique sensible aux coûts sur des graphes généraux. En allant au-delà des couplages statiques et des solveurs globaux, le cadre permet l'apprentissage de politiques exécutables et respectant la topologie qui optimisent à la fois l'appariement des extrémités et les coûts de trajectoire intermédiaires. Les auteurs positionnent ce travail comme une étape vers l'intégration de la dynamique moderne basée sur l'apprentissage avec les principes de transport optimal pour les systèmes réseau complexes, particulièrement dans les scénarios où les contraintes opérationnelles (comme la congestion ou les barrières énergétiques) sont critiques. Le travail est modeste quant à ses limites, notant qu'il suppose actuellement une topologie de graphe fixe et entièrement connue, et identifie l'extension aux graphes variant dans le temps ou partiellement observés comme un travail futur.