Cost-augmented Schrödinger bridges on graphs are exactly solvable: a Feynman-Kac tilt replaces learned control

Akshay Balsubramani

cs.LGOct 1, 2026

Cost-augmented Schrödinger bridges on graphs are exactly solvable: a Feynman-Kac tilt replaces learned control

Authors: Akshay Balsubramani

Abstract

The generalized Schrödinger bridge on a graph moves mass between two distributions while charging a cost for the states visited. It has been approached by learning the rates of a controlled continuous-time Markov chain, with a temporal-difference penalty that restores the cost. A state cost folds into the reference process as a Feynman-Kac tilt. The cost-augmented bridge is then a plain bridge against the tilted reference, and the penalty is unnecessary. The bridge is computed exactly by alternating two endpoint rescalings, each one sparse matrix-exponential application; nothing is discretized in time or learned. The alternation converges at a rate set by the endpoint coupling alone. For a quadratic congestion cost on time-averaged occupancies, damped best response around the exact bridge is gradient descent on a strongly convex function, and its residual bounds its error. On a protein-folding model, a free-energy cost lowers the expected barrier of the folding paths. On the learned approach's road network, roll-outs of the exact bridge match the target within sampling error, and on networks with millions of intersections its memory grows linearly.

Figures & tables

Figure 1: The number of fitting iterations follows the second singular value of the endpoint coupling, while memory and time per iteration grow with the size of the problem. Road networks of the ROAD-FLOW benchmark from Washington DC (

n=9{,}559

) to Texas (

n=2{,}073{,}870

); 40 sources and 40 targets within 50 hops of one random center, running cost uniform on

[0,0.5]

, horizon 100. (a) The solver’s peak allocation grows in proportion to the node count, 346 to 412 bytes per node; dashed, proportional to

n

. (b) Wall-clock time per fitting iteration against the arc count times

\tau\max_{x}|A_{xx}|

, for the seven networks timed in one container (filled) and for Washington DC with the endpoints of [ 40 ] at horizons 25 to 400, timed together in a second container (open). The dashed lines, one proportional fit per container, nearly coincide. Washington DC, the one network timed in both containers, costs

2.0\times 10^{-8}

seconds per arc and unit of

\tau\max_{x}|A_{xx}|

in the first and

4.1\times 10^{-8}

in the second. The networks lie within a factor 2.4 of one constant (Washington DC, timed again at the end, within 1.1% of its first timing); the horizons lie within 3% (horizon 25, timed again at the end, within 1.3% of its first timing). (c) Fitting iterations to a change of

10^{-12}

in the logarithm of the scalings, against

\log(10^{12})/\log(1/\sigma_{2}^{2})

: the seven networks (filled), four further endpoint sets on New Hampshire (open circles) and the horizons of (b) (open squares); dashed, equality; dotted, a factor 2 either way. At the end of every solve the error contracts per iteration by

\sigma_{2}^{2}

to within 0.03%. The counts are 0.81 to 1.76 times the prediction and in its order on the networks (rank correlation 1.00). Every solve in (c) meets both endpoint marginals to a total variation below

3\times 10^{-13}

Figure 2: Inside the concentration band of the main table of [ 40 ] , the exact bridge recovers the minimum-cost assignment on every instance. Departures from the optimum appear only at higher entropy. (a) Mass the soft assignment places on the optimal assignment and (b) fraction of instances whose hard read-out is cost-optimal, against the row entropy of the coupling, for

K=6,8,10,20

. Each point averages 20 random integer cost matrices at one cost scale; along each curve the scales are 25, 50, 100, 200, 400 and 800, from right to left. The shaded band is the row-entropy range 0.07–0.17 reported for the learned solver in the main table of [ 40 ] . Ringed points are the scale matched to it at each size. There the hard read-out agrees with the optimum on every row of every instance. The largest endpoint-marginal violation over the 20 instances is

1.8\times 10^{-6}

9.4\times 10^{-7}

1.9\times 10^{-6}

and

1.0\times 10^{-6}

for

K=6,8,10,20

, set on this layered construction by the budget of the two-step fitting iteration. These violations change no read-out. That paper’s appendix repeats the table with a row entropy of 0.22 at

K=20

in place of 0.17; at the scale matching 0.22, 19 of the 20 instances at that size are cost-optimal. Mass on the optimum falls monotonically as the entropy rises. The departures in (b) scatter at a sample size of 20 instances. Bars are 95% intervals, bootstrap in both coordinates in (a) and Wilson in (b).

Figure 3: On a Markov state model of chignolin folding, a free-energy running cost lowers the expected barrier monotonically. (a) Expected barrier, the maximum free energy along a path, measured from the lowest basin state, against the strength

\lambda

of the running cost

\lambda F

(

\lambda=0

is the plain bridge; the axis is linear below 0.1 and logarithmic above). Filled markers are exact, from the tilted generator restricted to each sublevel set of

F

. Open markers average over the roll-outs that fold; they lie within 0.02 of the exact values. The dashed line is the uncontrolled process conditioned on folding, exact; its roll-out estimate from its 42 folding trajectories out of 4,000 is 5.60. (b) The exact distribution of the barrier,

\text{Pr}(\max_{t}F(X_{t})\leq h)

Figure 4: On the compared paper’s road network the bridge meets both endpoint marginals, exactly in its read-outs and within sampling error in its roll-outs. Washington DC roads (

n=9{,}559

intersections, 29,682 arcs) with the six supply and six demand nodes of [ 40 ] , horizon 100 mean jump times; four processes, one panel per metric, with exact values at the bar ends. (a) The larger total variation of the two exact endpoint read-outs; for the uncontrolled chain, which has no endpoint constraint, the terminal total variation. (b) The terminal total variation of 5,000 roll-outs of the controlled rate ( 4 ). The gray band spans exact draws of 5,000 particles from the target up to their 95th percentile (0.019), and the dotted line marks their mean (0.012). (c) to (e) The compared paper’s crowding metrics in expectation, per 5,000 particles: the mean time-averaged occupancy of the 100 most occupied nodes, the largest occupancy of one node at one unit time (both excluding the endpoints), and the largest flow along one arc within one unit time over its capacity. The congestion cost of Section 4 (bottom two bars) spreads the load. For the uncontrolled chain only the terminal total variation and the peak were computed.

Figure 5: The tested dampings converge below the critical damping

2/(1+L)

and fail to converge above it. The road network of Figure 4 . (a) The residual

\|p-\bar{p}(\beta p)\|_{2}

, which bounds the distance to the equilibrium (Proposition 4.1 ), against the round, at two congestion strengths. At damping 0.35 the curves for the two strengths start a decade apart and agree to within a factor 1.5 from round 10 on. (b) The ratio of successive residuals over the last three rounds of each run, against the damping in units of its critical value. Lines are the local rate

\max\{1-\alpha(1+m),\ \alpha(1+L)-1\}

of Proposition 4.1 . Triangles mark runs that did not converge: the ratio lies above one where the residual grows and at one where it cycles. At

\beta\tau=250

the curvature is

L=0.155

(bound

\beta\tau\max_{x}p^{\circ}(x)=1.64

), so the critical damping is 1.732; dampings 0.35, 1 and 1.56 converge and 1.91 does not. At

\beta\tau=2500

the curvature is

L=1.506

(bound

\beta\tau\max_{x}p^{\circ}(x)=12.29

), so the critical damping is 0.798; dampings 0.35 and 0.718 converge and 0.878 and 1 do not. On the branch set by

L

the contraction over the last three rounds is within 1.0% of

\alpha(1+L)-1

at both strengths. The smallest eigenvalue

m

\nabla^{2}h(p^{\circ})-I

is zero at both strengths, because the equilibrium occupancy is zero at the 37 congestible nodes that lie on no path from a supply node to a demand node. On the branch set by

m

, the local rate is therefore

1-\alpha

. The runs at damping 0.35 contract faster than that over their last three rounds, by 2.3% at

\beta\tau=250

and 1.0% at

\beta\tau=2500

. Above the threshold at

\beta\tau=250

(damping 1.91) the residual grows by 1.16 per round, against

\alpha(1+L)-1=1.20

Explore similar work

Jun 2, 2026math.OC

Nonlocal Mean Field Schrödinger Bridge with Learned Interactions

The Schrödinger Bridge Problem connects an initial distribution to a terminal one along a minimum-energy stochastic process. Its mean-field extension, the Mean-Field Schrödinger Bridge, governs interacting populations whose dynamics and costs depend on the collective distribution. When these interactions are nonlocal, their direct evaluation scales quadratically with the population size, making large ensembles intractable within FBSDE-based solvers. We replace these terms with neural surrogates in state and time, trained on empirical interaction values along sampled trajectories and embedded in a four-stage alternating scheme that updates the forward and backward potentials and the surrogates in turn, while preserving forward--backward consistency and the prescribed endpoint marginals. We derive Grönwall-type stability bounds quantifying how surrogate errors propagate to the generated trajectories under a small-gain condition. On crowd-navigation and high-dimensional opinion-dynamics benchmarks, the surrogates reproduce the trajectories obtained with exact evaluation at reduced training cost. The advantage is most significant when the interaction is a nonlinear functional of the measure, such as the normalized bounded-confidence drift, for which random-batch subsampling is biased and unstable whereas the learned surrogate remains accurate.

Daisuke Inoue, Dante Kalise, Mathieu Laurière

Department of Mathematics, Imperial College London, London SW7 2AZ, United Kingdom. · Shanghai Frontiers Science Center of Artificial Intelligence and Deep Learning, NYU-ECNU Institute of Mathematical Sciences, NYU Shanghai, Shanghai 200126, People’s Republic of China.

Schrödinger Bridges Mean-Field Limit Stochastic Processes Stochastic Interpolants Bridge

Ask PDF

Jul 18, 2026stat.ML

Twisted Schrödinger Bridge Matching

Over the past few years, diffusion-based Schrödinger bridge models have been proposed to approximate optimal transport dynamics between two prescribed boundary distributions, with successful applications to generative modeling. More precisely, these methods aim to estimate a path measure whose initial and terminal marginals match the two boundary distributions, while minimizing the Kullback-Leibler divergence with respect to a reference Markov process. In this work, we consider the generalized Schrödinger bridge problem, in which the reference process is a twisted Brownian motion, that is, a Feynman-Kac transform of a Brownian motion induced by a time-dependent differentiable potential. Building on the Iterative Markovian Fitting (IMF) paradigm, and in particular on its special case Diffusion Schrödinger Bridge Matching (DSBM), which corresponds to the zero potential case, we introduce Twisted Schrödinger Bridge Matching (TSBM), a diffusion-based method designed to handle both continuous- and discrete-time potentials. Unlike previous approaches, TSBM provides a rigorous extension of the IMF scheme to the generalized Schrödinger bridge problem. This derivation leads to a new bridge-matching loss that depends explicitly on the gradient of the potential and recovers the DSBM objective when the potential vanishes, yielding improved performance. We further introduce trajectory-based variance-reduction techniques that substantially stabilize optimization and may be useful beyond the present setting. Finally, we empirically demonstrate the benefits of TSBM for trajectory inference across increasingly high-dimensional settings, including crowd navigation and single-cell data. Code available at https://github.com/maxencenoble/twisted-sb-matching.

Maxence Noble, Marie Scheid, Yazid Janati +2

CMAP, Ecole polytechnique · Institute of Foundation Models · MBZUAI +1

Schrödinger Bridges Differentiable Optimal Transport Stochastic Differential Equations

Ask PDF

May 12, 2026cs.LG

QDSB: Quantized Diffusion Schrödinger Bridges

Learning generative models in settings where the source and target distributions are only specified through unpaired samples is gaining in importance. Here, one frequently-used model are Schrödinger bridges (SB), which represent the most likely evolution between both endpoint distributions. To accelerate training, simulation-free SBs avoid the path simulation of the original SB models. However, learning simulation-free SBs requires paired data; a coupling of the source and target samples is obtained as the solution of the entropic optimal transport (OT) problem. As obtaining the optimal global coupling is infeasible in many practical cases, the entropic OT problem is iteratively solved on minibatches instead. Still, the repeated cost remains substantial and the locality can distort the global transport geometry. We propose quantized diffusion Schrödinger bridges (QDSB), which compute the endpoint coupling on anchor-quantized endpoint distributions and lift the resulting plan back to original data points through cell-wise sampling. We show that the regularized optimal coupling is stable w.r.t. anchor quantization, with an error controlled by the quality of the anchor approximation. In real-world experiments, QDSB matches the sample quality of existing baselines, requiring substantially less time. Code and data are available at github.com/mathefuchs/qdsb.

Tobias Fuchs, Florian Kalinke, Nadja Klein

Karlsruhe Institute of Technology, Karlsruhe, Germany

Schrödinger Bridges Bridge Generative Models

Ask PDF