Tangent Schrödinger Bridge Matching: Learning Stochastic Transport with Mechanistic Sensitivities

Jowaria Khan; Elizabeth Bondi-Kelly

cs.LGOct 2, 2026

Tangent Schrödinger Bridge Matching: Learning Stochastic Transport with Mechanistic Sensitivities

Authors: Jowaria Khan, Elizabeth Bondi-Kelly

Organizations: University of Michigan

Abstract

Predicting how stochastic systems respond to changes in viscosity, reaction rates, or external forces requires costly simulations, motivating reusable learned models. Yet matching observed outcome distributions does not ensure accurate intervention responses. We introduce Tangent Schrödinger Bridge Matching (Tangent-SBM), which learns stochastic transports from endpoint observations and mechanistic sensitivities. It propagates parameter derivatives alongside trajectories and supervises them against supplied targets. For average-response targets, single-rollout squared error also penalizes response variability; our objective uses two independent rollouts to match the mean without this additional penalty. We establish conditions under which sensitivity accuracy bounds finite-change prediction error and decision regret. Across Gaussian, stochastic double-well, PDEBench reaction--diffusion, and stochastic Navier--Stokes systems, Tangent-SBM improves sensitivity and finite-change prediction over matched conditional-bridge baselines while maintaining comparable endpoint and distributional accuracy. Controls examine target correctness, response objectives, and simulator-budget allocation. To test decision usefulness, we evaluate calibrated viscosity selection in Navier--Stokes: Tangent-SBM reduces tracking error relative to taking no action on every evaluated task.

Figures & tables

Figure 1: Tangent-SBM learns outcome distributions and intervention responses by supervising trajectory sensitivities. Independent rollout pairs match mean-response targets without penalizing response variability.

Dataset	Method	Endpoint	$E_{J}$	Finite change
Gaussian	DSBM	$0.714\pm 0.002$	$1.000\pm 0.000$	$0.321\pm 0.000$
	Conditional DSBM	$0.433\pm 0.014$	$0.866\pm 0.003$	$0.275\pm 0.001$
	GSBM	$0.440\pm 0.002$	$0.865\pm 0.001$	$0.276\pm 0.000$
	TSBM	$0.440\pm 0.000$	$0.865\pm 0.001$	$0.275\pm 0.000$
	Tangent-SBM	$\mathbf{0.236\pm 0.006}$	$\mathbf{0.050\pm 0.005}$	$\mathbf{0.016\pm 0.001}$
Double well	DSBM	$0.336\pm 0.005$	$1.000\pm 0.000$	$0.110\pm 0.000$

Table 1: ID endpoint, sensitivity, and finite-change errors across four systems. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds. Metric definitions and method-specific evaluation budgets appear in Appendix B .

Figure 2: SNS predictions for a fixed Seen example. Rows show the mean endpoint field, its change under a viscosity intervention, and the error in that change. Color scales are shared within rows; the bottom-left panel repeats the physical change from the second row on a narrower color scale.

Method	Tracking error	Regret	Agreement	Improvement
DSBM	$1.000\pm 0.000$	$0.965\pm 0.000$	0.0%	0.0%
Conditional DSBM	$0.553\pm 0.390$	$0.625\pm 0.909$	53.1%	84.4%
Tangent-SBM	$0.256\pm 0.040$	$0.048\pm 0.028$	79.2%	100.0%
TSBM	$0.724\pm 0.313$	$0.742\pm 0.460$	35.4%	47.9%
Finite-grid oracle	$0.186$	$0$	100%	100%

Table 2: Calibrated SNS intervention selection on 32 tasks, shared across three training seeds. Tracking error and normalized regret are lower-is-better. Agreement and improvement percentages average seed-level task rates.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

		Seen			Near-OOD			Far-OOD
Dataset	Method	Endpoint	$E_{J}$	Finite change	Endpoint	$E_{J}$	Finite change	Endpoint	$E_{J}$	Finite change
Gaussian	DSBM	$.559\pm.001$	$1.000\pm.000$	$.345\pm.000$	$.511\pm.004$	$1.000\pm.000$	$.237\pm.000$	$1.397\pm.006$	$1.000\pm.000$	$.512\pm.000$
	Conditional DSBM	$.245\pm.008$	$1.048\pm.007$	$.357\pm.003$	$.481\pm.036$	$1.247\pm.034$	$.249\pm.004$	$.397\pm.028$	$.664\pm.008$	$.329\pm.004$
	GSBM	$.264\pm.003$	$1.051\pm.001$	$.358\pm.001$	$.505\pm.003$	$1.274\pm.008$	$.252\pm.002$	$.462\pm.009$	$.681\pm.008$	$.339\pm.005$
	TSBM	$.263\pm.002$	$1.051\pm.002$	$.358\pm.001$	$.504\pm.005$	$1.266\pm.006$	$.250\pm.001$	$.437\pm.022$	$.672\pm.011$	$.333\pm.006$
	Tangent-SBM	$\mathbf{.244\pm.008}$	$\mathbf{.035\pm.004}$	$\mathbf{.017\pm.005}$	$\mathbf{.245\pm.007}$	$\mathbf{.129\pm.025}$	$\mathbf{.021\pm.004}$	$\mathbf{.260\pm.003}$	$\mathbf{.137\pm.026}$	$\mathbf{.085\pm.006}$

Appendix

Table 3: Endpoint, sensitivity, and finite-change errors across intervention regions. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds.

Method	Gaussian covariance error	Double-well probability RMSE
DSBM	$0.485\pm 0.009$	$0.170\pm 0.001$
Conditional DSBM	$0.330\pm 0.009$	$0.108\pm 0.003$
GSBM	$0.346\pm 0.004$	$0.107\pm 0.001$
TSBM	$0.345\pm 0.003$	$0.106\pm 0.003$
Tangent-SBM	$0.329\pm 0.011$	$0.105\pm 0.001$

Appendix

Table 4: ID distributional errors for Gaussian and double well. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds.

Method	Spread error	ED	SWD	Enstrophy $W_{1}$
DSBM	$.207\pm.017$	$.0207\pm.0028$	$.0345\pm.0019$	$2.222\pm.439$
Conditional DSBM	$.188\pm.008$	$.0170\pm.0048$	$.0319\pm.0045$	$1.654\pm.123$
GSBM	$.189\pm.002$	$.0130\pm.0014$	$.0288\pm.0014$	$2.613\pm.806$
TSBM	$.181\pm.001$	$.0092\pm.0001$	$.0239\pm.0005$	$1.322\pm.113$
Tangent-SBM	$.183\pm.003$	$.0166\pm.0027$	$.0317\pm.0025$	$1.673\pm.490$

Appendix

Table 5: SNS ID distributional errors. ED is energy distance divided by

\sqrt{d}

; SWD is sliced Wasserstein distance. Enstrophy

W_{1}

is displayed in units of

10^{-3}

. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds.

Objective	Mean-response RMSE	Variance ratio
Single-rollout squared loss	$0.104\pm 0.025$	$0.312\pm 0.055$
Two-rollout mean loss	$0.066\pm 0.014$	$0.665\pm 0.102$

Appendix

Table 6: Double-well ID mean-target ablation. Entries are mean

\pm

sample standard deviation over three training seeds. The variance ratio compares conditional model response variance with conditional simulator response variance; it is a diagnostic rather than a training target.

System	Training variant	Endpoint error	$E_{J}$
Double well	Free, two-rollout mean loss	$0.136\pm 0.001$	$0.197\pm 0.031$
	Free, single-rollout squared loss	$0.145\pm 0.006$	$0.272\pm 0.051$
	Endpoint-constrained	$0.177\pm 0.002$	$0.252\pm 0.020$
SNS	Free, two-rollout mean loss	$0.265\pm 0.028$	$1.002\pm 0.075$
	Free, single-rollout squared loss	$0.253\pm 0.034$	$2.624\pm 1.103$
	Endpoint-constrained	$0.252\pm 0.043$	$1.914\pm 0.646$

Appendix

Table 7: ID response-objective and rollout controls. The variants describe training choices; all models are evaluated using free rollouts. Entries are mean

\pm

sample standard deviation over three training seeds. Lower errors are better.

System	Correct	Shuffled	Sign-reversed
Gaussian	$.050\pm.005$	$.858\pm.007$	$1.817\pm.008$
Double well	$.197\pm.031$	$.357\pm.026$	$1.027\pm.043$
PDEBench	$.508\pm.045$	$1.010\pm.008$	$1.647\pm.060$
SNS	$1.002\pm 0.075$	$1.925\pm.822$	$2.077\pm.077$

Appendix

Table 8: ID sensitivity error

E_{J}

with correct, shuffled, and sign-reversed targets. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds.

System	25%	50%	100%
Gaussian	$.050\pm.007$	$.050\pm.006$	$.050\pm.005$
Double well	$.197\pm.042$	$.198\pm.022$	$.197\pm.031$
PDEBench	$.516\pm.031$	$.509\pm.015$	$.508\pm.045$
SNS	$2.075\pm.828$	$1.726\pm.748$	$1.002\pm 0.075$

Appendix

Table 9: ID sensitivity error

E_{J}

when retaining different fractions of the additional response-label pool. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds.

Variant	Cosine alignment	Magnitude ratio
Conditional DSBM	$.121\pm.082$	$2.547\pm.310$
TSBM	$.058\pm.102$	$2.253\pm.166$
Tangent-SBM, correct targets	$.592\pm.040$	$1.140\pm.142$
Tangent-SBM, shuffled targets	$.229\pm.125$	$1.794\pm.817$
Tangent-SBM, sign-reversed targets	$-.336\pm.103$	$1.518\pm.152$

Appendix

Table 10: SNS ID response alignment and magnitude. Higher cosine indicates better alignment; a magnitude ratio closer to one indicates a better match in response size. Entries are mean

\pm

sample standard deviation over three training seeds.

System	Variant	Endpoint	$E_{J}$	Finite change
PDEBench	LocalBudget	$.636\pm.071$	$.849\pm.094$	$.846\pm.094$
	SpreadBudget-800	$.646\pm.039$	$.875\pm.032$	$.872\pm.036$
	SpreadBudget-2720	$.661\pm.063$	$.855\pm.109$	$.855\pm.111$
	Tangent-SBM	$.598\pm.069$	$.508\pm.045$	$.510\pm.046$
SNS	LocalBudget	$.272\pm.005$	$2.929\pm.403$	$2.517\pm.307$
	SpreadBudget	$.283\pm.011$	$3.649\pm.332$	$3.145\pm.281$

Appendix

Table 11: ID simulator-budget comparisons. LocalBudget and SpreadBudget use additional simulations as endpoint examples; Tangent-SBM uses them for sensitivity supervision. Lower errors are better. Entries are mean

\pm

sample standard deviation over three training seeds.

System	Predictor	Endpoint	$E_{J}$	Finite change
Gaussian	Terminal, endpoint-only	$.361\pm.001$	$.865\pm.002$	$.276\pm.001$
	Terminal, response	$.023\pm.002$	$.017\pm.001$	$.005\pm.000$
Double well	Terminal, endpoint-only	$.125\pm.003$	$.268\pm.008$	$.031\pm.001$
	Terminal, response	$.120\pm.001$	$.170\pm.005$	$.018\pm.001$
PDEBench	Terminal, endpoint-only	$.159\pm.003$	$.416\pm.032$	$.406\pm.017$
	Terminal, response	$.130\pm.003$	$.204\pm.011$	$.208\pm.010$

Appendix

Table 12: ID prediction errors beyond bridge models. “Response” denotes training with endpoint examples and mechanistic response targets. Lower is better. Entries are mean

\pm

sample standard deviation over three training seeds. Metrics follow Appendix B .

Supervision	Endpoint RMSE	Forward $E_{J}$	Backward $E_{J}$
Forward only	$.238\pm.008$	$.050\pm.005$	$.905\pm.004$
Backward only	$.427\pm.010$	$.869\pm.001$	$.101\pm.015$
Both directions	$.219\pm.009$	$.066\pm.013$	$.108\pm.017$

Appendix

Table 13: Gaussian ID forward/backward supervision study. Lower errors are better. Entries are mean

\pm

sample standard deviation over three training seeds.

Explore similar work

Twisted Schrödinger Bridge Matching
Jul 18, 2026Maxence Noble, Marie Scheid, Yazid Janati +2Schrödinger Bridges Trajectory Inference
CMAP, Ecole polytechnique · Institute of Foundation Models · MBZUAI +1
Distributionally Robust Schrödinger Bridge
Oct 1, 2026Jinhwan Sul, Panagiotis Theodoropoulos, Vincent Pacelli +2Distribution Shift Robustness Schrödinger Bridges
Department of Aerospace Engineering Georgia Institute of Technology, Atlanta, GA, USA
Tilted Schrödinger Bridge Matching
Sep 28, 2026Sergei Kholkin, Evgeny Burnaev, Alexander KorotinSchrödinger Bridges Unpaired Image-to-Image Translation
Applied AI Institute

Jul 18, 2026J/K move · Enter open · S save

Twisted Schrödinger Bridge Matching

Maxence Noble, Marie Scheid, Yazid Janati, Eric Moulines+1

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

Schrödinger Bridges Trajectory Inference Optimal Transport Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Twisted Schrödinger Bridge Matching

Distributionally Robust Schrödinger Bridge

Tilted Schrödinger Bridge Matching