PaperlayerView exact revision on arXiv

Source-linked AI summary

Tangent Schr\"odinger Bridge Matching: Learning Stochastic Transport with Mechanistic Sensitivities

Jowaria Khan, Elizabeth Bondi-Kelly

arXiv:2610.02906v1cs.LG

TL;DR

Stochastic surrogate models can reduce repeated simulation costs, but matching observed outcomes does not determine responses to interventions. Tangent-SBM supervises trajectory sensitivities alongside outcome distributions, using independent rollouts for mean-response targets. Across four systems, comparisons support improved intervention responses alongside comparable stochastic distributional fit, and the Navier–Stokes decision study reports useful intervention choices.

  • Problem

    Stochastic intervention comparisons require costly repeated simulations, while agreement with observed outcomes at finite parameter values leaves local intervention responses undetermined.

    Source: §1 · pp. 1–2

  • Method

    Tangent-SBM learns stochastic transports by propagating parameter sensitivities alongside trajectories and supervising terminal responses, with separate objectives for individual and average targets.

    Source: §1 · p. 2

  • Results

    Across four systems, comparisons support the value of explicit response supervision, with gains coexisting with distributional fit comparable to Conditional DSBM on stochastic benchmarks.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The Navier–Stokes decision study shows that learned responses can guide useful intervention choices.

    Source: §6 · p. 9

    Benefits depend on sensitivity quality and coverage, and tangent supervision adds computation.

    Source: §6 · p. 10

Abstract

from arXiv · show

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\"odinger 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.

1 INTRODUCTION

Repeated simulation makes stochastic intervention comparison costly, and fitting outcomes at observed parameters does not guarantee accurate intervention responses. Tangent-SBM adds mechanistic sensitivity supervision to stochastic transport, with objectives tailored to individual or average responses.

1 INTRODUCTION

  • Repeated simulator solves make stochastic intervention comparisons costly, motivating reusable models that represent outcome distributions rather than only their means.The cost arises when comparing interventions across candidate settings and initial conditions.Source: §1 · p. 1
  • Endpoint agreement at finitely many parameter values leaves local responses undetermined, so a model can fit observations yet recommend an ineffective intervention.Mechanistic sensitivities describe responses to parameter changes and can be supplied by physical models or simulators.Source: §1 · pp. 1–2
  • Tangent-SBM propagates parameter sensitivities alongside learned trajectories and supervises terminal responses, using distinct objectives for individual and average response targets.For mean-only targets, two independent rollouts match the average without the additional penalty on response variability.Source: §1 · p. 2
  • The paper establishes conditions linking sensitivity accuracy to accurate finite-change intervention predictions and better decisions.Source: §1 · p. 2
  • Evaluation spans Gaussian, double-well, reaction–diffusion, and stochastic Navier–Stokes systems, with controls for target correctness and simulation budgets and a calibrated intervention-selection study.The experiments assess response accuracy, endpoint fidelity, distributional fit, and downstream usefulness.Source: §1 · p. 2

2 RELATED WORK

Prior work develops stochastic bridges and scientific surrogates, including derivative-informed learning. Tangent-SBM adds mechanistic sensitivity supervision to stochastic transport and distinguishes individual from mean responses.

Schrödinger bridges and stochastic generative models

  • Schrödinger bridges connect endpoint distributions through processes close to reference diffusions, with neural, conditional, generalized, and twisted formulations among the developed variants.The cited literature also links bridges to entropy-regularized optimal transport and describes alternative distribution-learning approaches.Source: §2 · p. 2
  • Tangent-SBM adds mechanistic sensitivity supervision to constrain intervention responses beyond endpoint fitting.Source: §2 · p. 3

Scientific surrogates and derivative-informed learning

  • Scientific learning includes physics-informed models, neural operators, and derivative-informed methods such as Sobolev training and Jacobian matching.Derivative-informed neural operators have been applied to scientific surrogates and optimization under uncertainty.Source: §2 · p. 3

Stochastic sensitivities and response objectives

  • For stochastic outputs, regressing each random response against a mean target with squared error penalizes response variance as well as mean error; independent rollouts isolate the mean-error term.The paper applies the double-sampling principle to mechanistic response supervision for stochastic bridges.Source: §2 · p. 3

3 TANGENT SCHR ¨ODINGER BRIDGE MATCHING

Tangent-SBM fits conditional stochastic transports and tracks intervention sensitivities along trajectories. Its response objectives distinguish individual-trajectory targets from mean targets, and training combines the selected response loss with bridge drift matching.

3 TANGENT SCHR ¨ODINGER BRIDGE MATCHING

  • Tangent-SBM trains stochastic transports to match observed outcomes and intervention responses, with different objectives for individual and average response targets.The framework is summarized in Figure 1.Source: §3 · p. 3

3.1 CONDITIONAL STOCHASTIC TRANSPORT

  • The conditional transport generates an outcome from context and intervention using a learned drift, time-dependent additive noise, and an intervention-independent initial state.Bridge time parameterizes the transport and need not equal physical simulation time.Source: §3.1 · p. 3

3.2 TRACKING TRAJECTORY RESPONSES

  • Holding the Brownian path fixed isolates intervention effects, and differentiating the rollout tracks trajectory responses whose conditional expectation equals the directional response of the mean prediction.The implementation differentiates discretized rollouts using Jacobian–vector products.Source: §3.2 · pp. 3–4

3.3 SUPERVISING INDIVIDUAL AND AVERAGE RESPONSES

  • Individual-response supervision requires a specified, parameter-independent pairing of model and simulator randomness, whereas mean-only targets do not require trajectory pairing.A single-rollout squared loss against a mean target penalizes response variability; two conditionally independent rollouts isolate mean-response error without that additional penalty.Source: §3.3 · p. 4

3.4 TRAINING WITH RESPONSE SUPERVISION

  • Training augments the bridge drift-matching loss with a response loss chosen according to whether targets describe individual paths or the mean.Response updates sample context, intervention, and direction, then propagate model states and tangents using free rollouts that measure the sampler used for prediction.Source: §3.4 · pp. 4–5
  • Tangent supervision adds tangent arrays and derivative computation, while mean supervision also needs independent replicas; equal simulator-call budgets therefore need not imply equal training costs.The added loss balances endpoint and response fit and does not generally preserve the original DSBM solution.Source: §3.4 · p. 5

4 FROM LOCAL SENSITIVITIES TO FINITE INTERVENTIONS

The section shows that endpoint agreement alone cannot determine intervention responses, then relates sensitivity accuracy to finite-change prediction and intervention selection. These guarantees depend on response accuracy along the intervention segment and, for decisions, across the finite action set.

What endpoint observations leave undetermined

  • Smooth conditional distributions can agree at finitely many interventions yet have arbitrarily different mean Jacobians there, so endpoint agreement alone does not identify responses.Sensitivity supervision supplies information about this otherwise unconstrained parameter dependence.Source: §4 · p. 5

From sensitivities to finite changes

  • For continuously differentiable outcomes, finite-change prediction error is bounded using sensitivity errors along the intervention segment, including both model-fit and target errors.The bound requires accurate sensitivities along the segment, not only at its starting point.Source: §4 · p. 5

Consequences for intervention selection

  • For a finite action set, uniformly bounded errors in predicted mean changes bound the tracking error of the selected action, even if it differs from the oracle’s choice.For linear predictions, the bound must also include linearization errors.Source: §4 · p. 5

5 EXPERIMENTS

Experiments test response learning across four systems, assess which supervision choices matter, and evaluate whether learned responses support useful decisions. Tangent-SBM improves intervention-response accuracy and calibrated viscosity selection in the evaluated settings.

5.1 SYSTEMS AND EVALUATION

  • The evaluation covers Gaussian, double-well, reaction–diffusion, and stochastic Navier–Stokes systems, with Seen, ID, Near-OOD, and Far-OOD intervention-value settings.Baselines include four bridge methods without sensitivity targets and a single-rollout Sobolev-DSBM control.Source: §5.1 · pp. 6–7

5.2 ENDPOINT FIT AND INTERVENTION RESPONSE

  • Across Gaussian, double-well, PDEBench, and SNS experiments, Tangent-SBM improves response accuracy over conditional-bridge baselines while maintaining comparable endpoint and distributional accuracy.The reported results include lower Gaussian and double-well ID response errors, 37–44% lower PDEBench response error, and lower SNS ID sensitivity error.Source: §5.2 · pp. 7–8, Table 1 · p. 7

5.3 WHAT MAKES RESPONSE SUPERVISION EFFECTIVE?

  • On double well, two-rollout mean supervision lowers ID mean-response RMSE from 0.104 ± 0.025 to 0.066 ± 0.014 versus single-rollout supervision.Single-rollout squared error against a mean target also penalizes response variability.Source: §5.3 · p. 8
  • On PDEBench ID, wrong-example and reversed-sign sensitivity targets raise EJ from 0.508 to 1.010 and 1.647; SNS errors fall from 2.075 to 1.002 as response coverage expands.The SNS comparison retains response labels at endpoint-training intervention values.Source: §5.3 · p. 8
  • Endpoint-only budget controls yield PDEBench ID EJ of approximately 0.85–0.88 versus 0.508 for Tangent-SBM, and SNS errors of 2.929 and 3.649 versus 1.002.These tested endpoint allocations do not recover the gains from explicit response supervision.Source: §5.3 · p. 9
  • Response supervision lowers ID sensitivity error from 0.865 to 0.017 for a Gaussian terminal predictor and from 0.416 to 0.204 for a PDEBench terminal predictor.Direct response-supervised predictors are alternatives when only mean outcomes are needed; Tangent-SBM also models outcome distributions.Source: §5.3 · p. 9

5.4 INTERVENTION SELECTION

  • Tangent-SBM reduces calibrated SNS tracking error by 54% versus Conditional DSBM and 65% versus TSBM.It improves over no intervention on all 96 seed–task pairs and selects the simulator’s best candidate in 79.2% of decisions.Source: §5.4 · p. 9, Table 2 · p. 9

6 DISCUSSION AND LIMITATIONS

Tangent-SBM combines endpoint observations with mechanistic sensitivities to improve intervention responses, with benefits depending on sensitivity quality and coverage and added computation. The appendix details the conditions and objectives behind its sensitivity and decision guarantees.

6 DISCUSSION AND LIMITATIONS

  • Across stochastic transport, Tangent-SBM improves intervention-response accuracy while retaining comparable distributional fit, but benefits depend on sensitivity quality and coverage and supervision adds computation.The discussion also notes that better endpoint predictions alone do not necessarily improve intervention decisions.Source: §6 · pp. 9–10

A.1 TRAJECTORY SENSITIVITIES AND MEAN RESPONSES

  • Pathwise differentiation requires a sufficiently regular drift, intervention-independent additive diffusion and initial state, and numerical convergence to identify the discrete tangent with the continuous-SDE derivative.Under integrability conditions, the expected response derivative equals the expected pathwise derivative.Source: §A.1 · p. 15

A.2 ENDPOINT AMBIGUITY

  • Finite endpoint observations can leave conditional outcome distributions unchanged while making mean-Jacobian discrepancies arbitrarily large; finite endpoint and sensitivity observations still do not identify the entire response function.Additional common derivative bounds can constrain this ambiguity.Source: §A.2 · p. 16

A.3 RESPONSE OBJECTIVES AND INDEPENDENT ROLLOUTS

  • For average-response targets, independent model rollouts isolate squared mean error, whereas single-response squared loss adds a response-variance penalty.A realization-level pairing instead affects the full objective, which can penalize response variance as well as mean error.Source: §A.3 · pp. 16–17

A.4 BIASED AND NOISY MECHANISTIC TARGETS

  • With reused noisy targets, target covariance does not change the population gradient, but target bias shifts the fitted response; independent target estimates remove covariance, not bias.Finite-difference truncation and simulation error can also contribute to target-approximation error.Source: §A.4 · p. 17

A.5 FINITE-INTERVENTION ERROR; A.6 DIRECTION AND COLLOCATION COVERAGE

  • Finite-change prediction error is bounded by sensitivity-fit and target errors along the intervention segment, while extending coverage from sampled directions or collocation points requires additional assumptions.The collocation-point bound assumes bounded response error there and a Lipschitz bound between points.Source: §A.5 · p. 18, §A.6 · p. 18

A.7 INTERVENTION SELECTION AND CALIBRATED PREDICTIONS

  • Intervention-selection guarantees bound the selected action’s tracking error from prediction error, without requiring it to match the oracle action.For linearized forecasts, the bound must also include linearization error.Source: §A.7 · pp. 18–19

B.1 SHARED SETUP

  • Endpoint examples provide physical outcomes, while response supervision adds information about how outcomes change with intervention; conditional methods receive the intervention as input.Finite-change evaluation uses the same model noise at original and perturbed parameter values to isolate intervention effects.Source: §B.1 · pp. 19–20

B.2 DATASET CONSTRUCTION

The appendix documents the benchmark construction, model adaptations, evaluation metrics, and calibrated intervention-selection protocol. It specifies how sensitivity targets and test actions are obtained for the simulated systems.

B.2 DATASET CONSTRUCTION

  • The passages describe stochastic Navier–Stokes data with intervention-controlled viscosity and final vorticity fields, and PDEBench reaction–diffusion data evolving two fields on a 128 × 128 grid.For Navier–Stokes, the response target uses paired simulations at u + 0.03 and u − 0.03 with the same random forcing, divides their difference by 0.06, and averages over independent forcing realizations.Source: §B.2 · pp. 20–21

B.3 MODELS AND OPTIMIZATION

  • Tangent-SBM continues from conditional-bridge checkpoints with response supervision, while GSBM and TSBM adaptations receive intervention conditioning and trajectory costs but no sensitivity labels.The field-system trajectory cost is a spatial smoothness prior, not a governing-equation residual.Source: §B.3 · pp. 21–22

B.4 METRICS AND EVALUATION

  • Evaluation measures endpoint, sensitivity, and finite-intervention errors, with finite changes compared against physical mean changes and dataset-specific error metrics.The reported sampling budgets vary by system; the main SNS comparison uses 32/1,024/16 model samples for endpoint, sensitivity, and finite-change estimates.Source: §B.4 · pp. 22–23

B.5 INTERVENTION-SELECTION PROTOCOL

  • For each method and training seed, the protocol estimates a local response from 16 model rollouts, freezes α before testing, and evaluates three seeds on 96 decisions.The candidate set is {0, −0.125, 0.125, −0.25, 0.25}. Tracking error measures distance from the goal; regret measures additional squared loss relative to the simulator-estimated best candidate action, defining the finite-grid oracle.Source: §B.5 · p. 23

C RESULTS ACROSS INTERVENTION REGIONS

Table 3 extends the ID comparison to Seen, Near-OOD, and Far-OOD regions defined relative to endpoint-training parameter values.

C RESULTS ACROSS INTERVENTION REGIONS

  • Table 3 compares endpoint, sensitivity, and finite-change errors across intervention regions, reporting three-seed means and sample standard deviations; lower is better.The regions are defined relative to parameter values used for endpoint training.Source: C RESULTS ACROSS INTERVENTION REGIONS · p. 24

D DISTRIBUTIONAL EVALUATION

Tangent-SBM retains distributional accuracy comparable to Conditional DSBM while improving responses. The section also tests mean-response objectives and where response supervision is measured.

D DISTRIBUTIONAL EVALUATION

  • For Gaussian outputs, the evaluation compares predicted and reference covariance; for double well, it measures error in the probability of ending in the right well.Covariance is estimated from residuals after subtracting conditional means.Source: D DISTRIBUTIONAL EVALUATION · pp. 24–25, Table 4 · p. 25
  • SNS distributional diagnostics include energy distance per field dimension, sliced Wasserstein distance, enstrophy Wasserstein distance, and spread error.Spread error compares model and simulator standard-deviation fields at each grid location.Source: D DISTRIBUTIONAL EVALUATION · p. 25
  • Across stochastic systems, Tangent-SBM’s response improvements accompany distributional errors similar to Conditional DSBM.Tables 4 and 5 cover Gaussian, double well, and stochastic Navier–Stokes; Tangent-SBM has similar errors on the two low-dimensional systems and comparable accuracy across the three stochastic systems.Source: D DISTRIBUTIONAL EVALUATION · pp. 24–25, Table 4 · p. 25, Table 5 · p. 25

E RESPONSE OBJECTIVE AND ROLLOUT CONSTRUCTION

  • For conditional-mean targets, the two-rollout objective lowers mean-response RMSE and increases trajectory-response variation relative to the single-rollout objective.The variance ratio is diagnostic only and does not establish recovery of physical trajectory-level responses because model and simulator trajectories need not be paired.Source: E RESPONSE OBJECTIVE AND ROLLOUT… · pp. 25–26, Table 6 · p. 26
  • On SNS, free-rollout supervision with the two-rollout mean loss gives the lowest mean response error, while controls have slightly lower mean endpoint errors.For double well, free and endpoint-constrained supervision give similar results; all models are evaluated using free rollouts.Source: E RESPONSE OBJECTIVE AND ROLLOUT… · p. 26, Table 7 · p. 26

F TARGET CORRECTNESS AND RESPONSE COVERAGE

Correct mechanistic targets, rather than merely adding a derivative loss, underpin the response gains; the tested coverage requirement varies by system.

F TARGET CORRECTNESS AND RESPONSE COVERAGE

  • Correct response targets give the lowest mean sensitivity error in every system; shuffled or sign-reversed targets do not reproduce the same gains.The controls retain response supervision but assign labels to wrong conditions or reverse their signs.Source: F TARGET CORRECTNESS AND RESPONSE… · p. 26, Table 8 · p. 27
  • Gaussian, double well, and PDEBench show little response-error change across tested coverage levels, while SNS sensitivity error decreases with broader coverage.SNS retains response labels at endpoint-training settings in every variant; only additional response conditions are subsampled, and 100% denotes the full available pool.Source: F TARGET CORRECTNESS AND RESPONSE… · p. 27, Table 9 · p. 27
  • For SNS, correct supervision improves response alignment and reduces magnitude overestimation; sign-reversed targets yield responses opposed to the reference.The sign-reversed model predicts a response of similar size to the correctly supervised model, so magnitude alone does not explain the difference.Source: F TARGET CORRECTNESS AND RESPONSE… · p. 27, Table 10 · p. 28

G SIMULATOR-BUDGET CONTROLS

Under the tested simulator-run budgets, explicit response supervision outperforms endpoint-only allocations on sensitivity and finite-change errors.

G SIMULATOR-BUDGET CONTROLS

  • LocalBudget uses additional perturbation outputs as endpoint examples, while SpreadBudget distributes the same simulation budget over a broader intervention range.Neither control uses sensitivity labels; Table 11 compares their prediction and response errors with Tangent-SBM.Source: G SIMULATOR-BUDGET CONTROLS · p. 28, Table 11 · p. 29
  • The PDEBench response-label set includes 108 anchor conditions and 216 additional conditions, with each sensitivity label requiring two simulator runs.The added runs are beyond the original 270 endpoint-training examples; SpreadBudget uses 648 runs and tests 800 versus 2,720 continuation updates.Source: G SIMULATOR-BUDGET CONTROLS · p. 28
  • Tangent-SBM achieves lower mean sensitivity and finite-change errors than tested endpoint-only allocations on both systems, and extra PDEBench SpreadBudget updates do not close the gap.The controls match simulator-run counts, not total training time; the comparison supports spending the additional simulations on response supervision under these allocations.Source: G SIMULATOR-BUDGET CONTROLS · pp. 28–29, Table 11 · p. 29

H RESPONSE SUPERVISION BEYOND BRIDGES

Response supervision improves errors across tested deterministic predictors, while direct mean predictors and stochastic transport serve distinct evaluation needs. In the Gaussian study, response accuracy in one transport direction does not automatically transfer to the other.

H RESPONSE SUPERVISION BEYOND BRIDGES

  • Response supervision lowers mean endpoint, sensitivity, and finite-change errors in every tested deterministic predictor pair.Sensitivity error falls from 0.865 to 0.017 for the Gaussian terminal predictor and from 0.416 to 0.204 for the PDEBench terminal predictor.Source: H RESPONSE SUPERVISION BEYOND BRIDGES · p. 29
  • Direct response-supervised predictors are strong alternatives when only mean outcomes are needed, especially on deterministic PDEBench.That benchmark tests high-dimensional response learning rather than the need for stochastic outcome modeling.Source: H RESPONSE SUPERVISION BEYOND BRIDGES · p. 30
  • On SNS, Tangent-SBM has lower reported sensitivity and finite-change errors than tested direct operators and conditional flow matching despite their lower endpoint errors.Tangent-SBM retains a sampler of outcome distributions, distinguishing response accuracy from mean prediction and variability modeling.Source: H RESPONSE SUPERVISION BEYOND BRIDGES · pp. 29–30

I ADDITIONAL SUPERVISION STUDIES

  • In the Gaussian study, forward-only supervision lowers forward error but leaves backward error high, and backward-only supervision shows the opposite pattern.Supervising both directions gives low errors in both and the lowest mean endpoint RMSE in this study; forward-only retains the lowest mean forward response error.Source: I ADDITIONAL SUPERVISION STUDIES · p. 30

Related papers

Loading 2610.02906v1…