Generalized Schrödinger Bridge on Graphs

Technical Summary: Generalized Schrödinger Bridge on Graphs (GSBoG)

Problem Statement
Transportation on graphs is a fundamental challenge in domains ranging from supply chains and traffic infrastructure to molecular kinetics. While optimal transport (OT) on graphs provides principled methods for matching source and target probability distributions, existing approaches often yield static couplings that fail to describe time-dependent mass movement or provide executable control policies. Furthermore, dynamic formulations for graph transport, such as those solving Hamiltonian flows on the probability simplex, suffer from poor scalability on large, sparse topologies and often lack the expressivity to optimize intermediate trajectories under application-specific state-dependent costs (e.g., congestion or energy barriers). The core challenge is to construct a feasible, topology-respecting evolution that controls intermediate distributions while minimizing network costs and accommodating operational constraints like capacity limits.

Methodology
The authors introduce the Generalized Schrödinger Bridge on Graphs (GSBoG), a data-driven framework for learning executable, controlled continuous-time Markov chain (CTMC) policies on arbitrary graphs. The method extends the principles of the Generalized Schrödinger Bridge (GSB) from continuous state spaces to discrete, graph-structured domains.

1. Mathematical Formulation:

   • The problem is formulated as minimizing a cost functional over path measures of a CTMC. The objective combines a running cost $f_t(x, p_t)$ (dependent on state and marginal distribution) with a Kullback-Leibler (KL) divergence term that regularizes the controlled dynamics toward a reference CTMC process.
   • The authors derive discrete optimality conditions parallel to the continuous-state formulation. They establish that the optimal controlled transition rates $u^*_t(y, x)$ satisfy a Hopf-Cole transformation: $u^*_t(y, x) = r_t(y, x) \exp(-V_t(y) + V_t(x))$, where $V_t$ is a time-dependent potential function and $r_t$ is the reference rate.
   • This leads to a coupled system of equations for the Schrödinger potentials $(\phi_t, \hat{\phi}_t)$, analogous to the Fokker-Planck equation in continuous settings but defined via the graph's continuity equation.

2. Algorithmic Approach (gIPF and TD):

   • To avoid computationally prohibitive global solvers over the entire graph and time horizon, GSBoG adopts a particle-based, data-driven approach using Generalized Iterative Proportional Fitting (gIPF).
   • The method utilizes Dynkin's formula to derive discrete analogues of the Forward-Backward Stochastic Differential Equations (FBSDEs) used in continuous GSB.
   • The training objective consists of two alternating components:
     • IPF Loss: Minimizes the discrepancy between the learned path measure and the reference process to satisfy endpoint marginal constraints ($p_0 = \mu, p_1 = \nu$).
     • Temporal-Difference (TD) Loss: Explicitly enforces consistency with the state- and distribution-dependent running costs by penalizing deviations between the generator-predicted increments and actual trajectory increments.
   • The algorithm parameterizes the log-potentials using neural networks (or embeddings) and updates them via stochastic gradient descent, evaluating rates only on local graph neighborhoods.

Key Contributions

• First Data-Driven GSB on Graphs: The paper presents the first formulation of the Generalized Schrödinger Bridge on arbitrary graphs in a data-driven manner, distinct from recent discrete diffusion works that focus on graph generation.
• Scalable Policy Learning: By learning trajectory-level policies via particle-based sampling rather than solving global time-expanded flows, GSBoG achieves scalability on large, sparse graphs where traditional dynamic OT methods fail due to memory constraints.
• Cost-Aware Dynamics: The framework explicitly optimizes intermediate behavior under general running costs (e.g., congestion penalties), ensuring that the resulting policies are not only topologically feasible but also operationally efficient.
• Theoretical Derivation: The authors provide a rigorous derivation of discrete optimality conditions and a tractable learning objective based on the generator of the CTMC, bridging the gap between stochastic control theory and discrete graph transport.

Experimental Results
The authors evaluate GSBoG on three challenging real-world graph topologies:

1. Supply Chain Network (9,559 nodes): GSBoG successfully learned routing policies that matched endpoint supply/demand distributions while significantly reducing intermediate congestion and peak node occupancy compared to baselines like Attraction Flow and Dynamic W1 flow. Notably, GSBoG was the only method that kept induced edge flows within capacity limits without explicit hard constraints, whereas global solvers (GrSB) failed to complete due to memory exhaustion.
2. Balanced Assignment Task: On bipartite graphs with varying sizes ($n=6$ to $20$), GSBoG achieved near-perfect overlap with the optimal static OT solution (100% accuracy for $n \le 10$) while learning a soft coupling that minimized assignment costs.
3. Discretized Molecular Dynamics (Chignolin Folding): Applied to a Markov State Model (MSM) of protein folding, GSBoG steered rare-event transitions from an unfolded to a folded state. It achieved a 99.7% folding rate with a significantly reduced energy barrier (1.80 $k_B T$) compared to uncontrolled dynamics or committor-guided biases, and further improved to 1.39 $k_B T$ when incorporating free energy running costs.

Scalability Analysis
Experiments demonstrate that GSBoG scales linearly with the local branching factor rather than the total number of nodes. On graphs with up to $10^6$ nodes, GSBoG remained feasible with low memory usage (1.4 GB) and runtime (1 hour), whereas Dynamic OT methods became infeasible (N/A) or required significantly more resources (e.g., 22 GB memory, 4.6 hours).

Significance and Claims
The paper claims that GSBoG paves new avenues for cost-aware dynamical transport on general graphs. By moving beyond static couplings and global solvers, the framework enables the learning of executable, topology-respecting policies that optimize both endpoint matching and intermediate trajectory costs. The authors position this work as a step toward integrating modern learning-based dynamics with optimal transport principles for complex networked systems, particularly in scenarios where operational constraints (like congestion or energy barriers) are critical. The work is modest in its limitations, noting that it currently assumes a fixed, fully known graph topology, and identifies extensions to time-varying or partially observed graphs as future work.