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

Jul 18, 2026

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 Trajectory Inference Optimal Transport Diffusion Models

Ask PDF

Oct 1, 2026

Distributionally Robust Schrödinger Bridge

Schrödinger bridge (SB) learns stochastic transport between prescribed initial and target distributions. When the initial distribution shifts at test time, the learned dynamics can fail to recover the target distribution. We introduce the Distributionally Robust Schrödinger Bridge (DRSB), which learns a single controller that accounts for uncertainty in the initial distribution. The DRSB objective consists of control energy and a KL penalty between the resulting terminal distribution and the target distribution. DRSB seeks a single controller that minimizes the worst-case value of this objective as the initial distribution varies within an ambiguity set around the nominal distribution. We derive an exact variational formulation of this objective and connect its fixed-terminal-cost subproblem to stochastic optimal control and distributionally robust optimization. This formulation motivates an alternating algorithm that updates the adversarial initial distribution, estimates the terminal log-density ratio, and trains the controller. We develop Wasserstein and Sinkhorn variants using stochastic control optimality conditions to approximate the gradients required for adversarial updates. Experiments on two-dimensional transport tasks and image-to-image translation show improved robustness to input perturbations relative to standard SB, with a tradeoff in nominal performance. On Gaussian mixture transport, Sinkhorn DRSB also achieves lower mean sliced Wasserstein distance than fixed-level noise augmentation at both tested unseen noise levels.

Jinhwan Sul, Panagiotis Theodoropoulos, Vincent Pacelli +2

Department of Aerospace Engineering Georgia Institute of Technology, Atlanta, GA, USA

Distribution Shift Robustness Schrödinger Bridges Distributionally Robust Optimization Optimal Transport

Ask PDF

Sep 28, 2026

Tilted Schrödinger Bridge Matching

Schrödinger bridges provide an entropy-regularized framework and a principled solution for unpaired domain translation. In practice, a pretrained bridge may need to be adapted to human preferences or physical constraints through a reward a problem closely related to reward tilting in diffusion models but underexplored for Schrödinger bridges. We introduce Tilted Schrödinger Bridge Matching (TSBM), a post-training method for fine-tuning a learned bridge

P

between source

p_0

and target

p_1

toward a reward-tilted target

p_1^r\propto p_1e^r

, while preserving source

p_0

. We formulate this adaptation as alternating optimization initialized from

P

, provide theoretical justification, and derive a practical algorithm based on Adjoint Matching. We evaluate TSBM on unpaired image-to-image translation targeting digit properties in MNIST and facial attributes in CelebA.

Sergei Kholkin, Evgeny Burnaev, Alexander Korotin

Applied AI Institute

Schrödinger Bridges Unpaired Image-to-Image Translation

Ask PDF