PaperlayerView exact revision on arXiv

Source-linked AI summary

Distributionally Robust Schr\"odinger Bridge

Jinhwan Sul, Panagiotis Theodoropoulos, Vincent Pacelli, Jaemoo Choi, Evangelos Theodorou

arXiv:2610.02043v1cs.LGcs.AI

TL;DR

Standard Schrödinger Bridge dynamics may fail to preserve the target distribution when the initial law shifts at test time. DRSB learns one controller robust to an ambiguity set of initial distributions; experiments report improved robustness over standard SB across transport and image-translation tasks, with a tradeoff in nominal performance.

  • Problem

    A Schrödinger Bridge controller trained for a prescribed initial distribution may produce a different terminal distribution after a test-time shift, violating the target marginal constraint.

    Source: §1 · p. 1

  • Method

    DRSB minimizes the worst-case control-energy and terminal-mismatch objective over an ambiguity set, using a variational formulation and alternating adversarial-distribution, terminal-cost, and controller updates.

    Source: §1 · pp. 1–2

  • Results

    DRSB improves robustness to noisy and shifted initial distributions over standard SB on 2D transport and image-to-image translation, while Sinkhorn DRSB outperforms Wasserstein DRSB under most tested 2D perturbations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The robust controller need not preserve nominal SB transport or be optimal for each initial distribution separately, so nominal performance may be traded for lower worst-case cost.

    Source: §3.1 · p. 4

    Minimax adversarial training can be unstable in high-dimensional settings and requires careful hyperparameter tuning and initialization.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

1 INTRODUCTION

DRSB addresses the failure of standard Schrödinger Bridge dynamics to preserve target matching under initial-distribution shifts by learning one controller for uncertain initial distributions. Its formulation and update methods support improved robustness in the reported transport and image-translation experiments.

1 INTRODUCTION

  • Standard SB dynamics are trained for a prescribed initial distribution, so a test-time shift can change the terminal distribution and violate the target marginal constraint.This motivates accounting for uncertainty in the initial distribution.Source: §1 · p. 1
  • DRSB seeks one controller that transports every initial distribution in an ambiguity set close to the target, using a terminal cost for mismatch between induced and target distributions.The terminal cost is based on their log-density ratio and is approximated with a neural network to provide a controller training signal.Source: §1 · pp. 1–2
  • DRSB improves robustness to noisy and shifted initial distributions over standard SB, while Sinkhorn DRSB outperforms Wasserstein DRSB under most tested 2D perturbations.Evaluations cover 2D Gaussian-to-Gaussian and GMM-to-GMM transport and image-to-image translation.Source: §1 · p. 2
  • The paper derives an exact variational formulation of DRSB and connects its fixed-terminal-cost subproblem to stochastic optimal control and distributionally robust optimization.Source: §1 · p. 2
  • The proposed block-update algorithms for Wasserstein and Sinkhorn variants use stochastic optimal-control optimality conditions.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define the Schrödinger Bridge setting and explain why a controller trained for a nominal initial law may not preserve terminal matching after a shift. They then frame distributionally robust optimization as minimizing against a worst-case distribution in an ambiguity set.

2.1 SCHRÖDINGER BRIDGE UNDER INITIAL DISTRIBUTION UNCERTAINTY

  • SB seeks a controlled stochastic process that transports an initial law to a target law while minimizing its KL divergence from a reference process.The path measure is induced by the controlled dynamics; the zero-control process supplies the reference path measure.Source: §2.1 · p. 2
  • A controller trained for nominal initial law µ need not preserve terminal matching when deployment inputs instead follow a shifted, unknown law.The deployment setting assumes one controller is used without test-time retraining.Source: §2.1 · p. 3
  • An ambiguity set represents plausible deployment initial distributions, and the goal is to minimize the worst-case sum of control energy and terminal mismatch using one controller.The target distribution remains unchanged under the initial-distribution shift.Source: §2.1 · p. 3, Figure 1 · p. 3

2.2 DISTRIBUTIONALLY ROBUST OPTIMIZATION

  • DRO minimizes over a decision variable while maximizing over distributions within a prescribed distance of the nominal distribution.Source: §2.2 · p. 3

3 DISTRIBUTIONALLY ROBUST SCHRÖDINGER BRIDGE

DRSB selects one controller to minimize worst-case transport cost over initial distributions near the nominal distribution, while retaining a terminal KL penalty. Its variational formulation connects fixed-terminal-cost optimization to stochastic optimal control and distributionally robust optimization, motivating alternating updates and Wasserstein and Sinkhorn variants.

3.1 DISTRIBUTIONAL ROBUSTNESS IN SB AND CONNECTION TO SOC

  • DRSB minimizes worst-case control energy plus a weighted terminal KL penalty over an ambiguity set, trading nominal performance for lower worst-case cost.The controller is chosen before the test-time input distribution is known; it need not preserve nominal SB transport or be optimal for each initial distribution separately.Source: §3.1 · p. 4
  • The exact variational formulation retains optimization over the terminal cost, whose optimum depends on both the controller and initial distribution.For fixed terminal cost, the subproblem has a stochastic-control cost and a worst-case distributional objective.Source: §3.1 · p. 5

3.2 ITERATIVE OPTIMIZATION SCHEME FOR DRSB

  • Training alternates adversarial initial-distribution, terminal-cost, and controller updates, but finite-step block updates do not establish convergence of the nested problem.The procedure initializes the controller from a pretrained standard SB model.Source: §3.2 · pp. 5–6

3.3 VARIANTS OF DRSB

  • The Wasserstein variant represents the adversary with one perturbed initial sample per nominal sample and optimizes worst-case samples by gradient descent.It reuses the previous iteration's dual variable; substituting the controller's initial-time action for the value-function gradient is exact at the fixed-cost SOC optimum and a surrogate after finite optimization.Source: §3.3 · p. 6
  • The Sinkhorn variant linearizes the fixed-cost value function to obtain Gaussian worst-case sampling and a closed-form dual-variable update, avoiding iterative inner optimization.The Sinkhorn distance uses entropy regularization, and its worst-case distribution has support over all of Rd even when the nominal distribution is discrete.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

Across tested perturbations, Sinkhorn DRSB usually outperformed Wasserstein DRSB, while image translation favored DRSB over DSBM under noise but favored DSBM on nominal inputs.

4.1 ROBUSTNESS AGAINST NOISY AND SHIFTED DISTRIBUTIONS ON 2D TRANSPORT TASKS

  • The 2D transport evaluation measures terminal-to-target SWD under Gaussian noise, directional shifts, and worst-case test inputs.The standard SB baseline uses DSBM; Table 1 reports mean ± standard deviation over 10 seeds.Source: §4.1 · p. 7, Table 1 · p. 7

4.2 COMPARISON BETWEEN WASSERSTEIN AND SINKHORN DRSB

  • DRSB-S achieved lower mean SWD in eight of ten conditions; its GMM worst-case samples were more dispersed than DRSB-W’s.The authors suggest the broader spread may expose the controller to a wider range of inputs during training.Source: §4.2 · p. 8, Table 2 · p. 8, Figure 2 · p. 8

4.3 NOISE AUGMENTATION AND GENERALIZATION TO UNSEEN NOISE LEVELS

  • The noise-augmentation comparison tests GMM transport at δ = 1, 2, 3, and 4, with δ = 3 and 4 unseen by both augmentation baselines.Augmentation uses training noise δtrain = 1 or 2 and perturbs 20% of clean samples.Source: §4.3 · p. 8, Table 3 · p. 8

4.4 IMAGE-TO-IMAGE TRANSLATION UNDER NOISE

  • Under all tested L2-norm noise levels, DRSB outperformed DSBM on image translation, while DSBM achieved lower FID on nominal inputs.The table reports FFHQ FID under Gaussian noise.Source: §4.4 · p. 9, Table 4 · p. 9

4.5 ADDITIONAL EXPERIMENTAL ANALYSIS

  • On one RTX 4090, the additional DRSB stage raised total training cost to approximately 1.54× DSBM on Gaussian transport and 1.16× on FFHQ.Worst-case optimization accounted for less than 1.5% of measured runtime.Source: §4.5 · p. 9
  • On Gaussian transport, the variational objective stabilized near 312, but these observations do not establish convergence to a saddle point.The final objective was 311.37; the authors also state that stationarity of all three blocks is not established.Source: §4.5 · p. 9
  • Across five evaluated epochs, cost-gradient directions and magnitudes approximately satisfied the SOC relation used in adversarial updates.At epoch 50, the direction cosine was 0.9998 and the gradient-to-control magnitude ratio was 1.0089; this is an empirical consistency check, not proof of exact optimality.Source: §4.5 · p. 9

5 RELATED WORK

Related work addresses distributional robustness through ambiguity sets and transport formulations, while DRSB targets shifts in the test-time initial distribution.

5 RELATED WORK

  • DRO moved from moment-based to distance-based ambiguity sets; Wasserstein offers a geometric interpretation, while KL offers convexity and tractable inner maximization but has stated limitations.The passage also identifies Sinkhorn DRO as a recent approach intended to improve computational tractability while modeling realistic uncertainty.Source: §5 · p. 10
  • Robust optimal-transport methods address outliers or noisy training samples, whereas DRSB focuses on distribution shifts at test time.The cited approaches include relaxed or unbalanced transport constraints and soft-constrained Schrödinger bridges.Source: §5 · p. 10

6 CONCLUSION AND LIMITATIONS

DRSB transports noisy or shifted input distributions toward the target, with Wasserstein and Sinkhorn variants. The conclusion notes limits in distance choices, high-dimensional training stability, and convergence evidence.

6 CONCLUSION AND LIMITATIONS

  • DRSB targets transport from noisy or shifted input distributions to the target distribution.The conclusion frames this as the problem the method formulates and solves.Source: §6 · p. 10
  • The Wasserstein inner problem is globally solved to bound the fixed-terminal-cost objective, while the linearized Sinkhorn inner problem permits Gaussian sampling.Source: §6 · p. 10
  • The study leaves other distributional distances for future work; high-dimensional minimax training can be unstable, and empirical diagnostics do not establish full-algorithm convergence.The conclusion also notes careful hyperparameter tuning and initialization are required.Source: §6 · p. 10

A.1 PROOF OF PROPOSITION 3.3

The proof passages establish the variational log-density-ratio result and derive a fixed-terminal-cost formulation without treating one optimizer as universal across controller and distribution choices.

A.1 PROOF OF PROPOSITION 3.3

  • The binary cross-entropy minimizer recovers the log-density ratio, and substituting it into the variational form achieves the KL divergence.Source: §A.1 · p. 14

A.2 PROOF OF THEOREM 3.4

  • For each fixed controller–distribution pair, the normalized maximizer depends on both, so it cannot be held fixed while optimizing over both variables.With a separately fixed g, the logarithmic term is constant in those variables, yielding the fixed-g subproblem.Source: §A.2 · p. 14

B DETAILS ON THE DERIVATION OF DRSB VARIANTS

B.1 DETAILS ON WASSERSTEIN DRSB

  • The fixed-g Wasserstein inner maximization is expressed in dual form using the interchangeability principle.Source: §B.1 · pp. 14–15

B.1.1 DERIVATION OF UPDATE RULES (EQUATION (18) AND EQUATION (19))

The update derivations use a fixed-terminal-cost stochastic-control condition to approximate adversarial gradients, while the Sinkhorn derivation simplifies its inner problem through cost approximation.

B.1.1 DERIVATION OF UPDATE RULES (EQUATION (18) AND EQUATION (19))

  • Replacing the fixed-g value-function gradient with the trained drift gives an approximate Wasserstein adversarial update unless the fixed-g SOC optimality condition holds.The relation is evaluated against discrete-adjoint gradients on a Gaussian task in the cited derivation.Source: §B.1.1 · p. 15

B.2 DETAILS ON SINKHORN DRSB

  • For positive dual parameters and a finite normalizing integral, the entropy-regularized Sinkhorn inner maximization has a closed-form dual and optimal conditional distribution.The stated boundary case at λ = 0 is understood by limits.Source: §B.2 · p. 17

B.2.1 SECOND-ORDER APPROXIMATION

  • A second-order expansion of the value function yields a Gaussian conditional distribution when its resulting matrix A is positive definite.The passage defines g as a local gradient vector and A from the local Hessian and regularization parameters.Source: §B.2.1 · p. 17

B.2.2 SIMPLIFICATION OF THE DUAL FORM

  • The second-order dual form complicates solving for λ because A^-1 depends on λ.Source: §B.2.2 · p. 18

B.2.3 FIRST-ORDER APPROXIMATION

  • Because computing the Hessian is expensive in high dimensions, the linearized-cost approximation omits it and gives exact expressions for that approximated inner problem.The conditional distribution then simplifies.Source: §B.2.3 · p. 18

C STOCHASTIC OPTIMAL CONTROL THEORY

  • The fixed-terminal-cost SOC derivation obtains the HJB equation and an optimal-control expression, with the identity applying to an optimal control and value function.The derivation assumes a sufficiently smooth value function and admissible controls.Source: C STOCHASTIC OPTIMAL CONTROL THEORY · pp. 18–19

D.1 ADJOINT MATCHING

  • Adjoint matching is presented for a fixed initial distribution and terminal cost, under differentiability and integrability assumptions for pathwise differentiation.Source: §D.1 · p. 19

D.1.1 DERIVATION OF THE ADJOINT SYSTEM

The appendix derives adjoint-based controller training, relates its stationarity to stochastic optimal control, and outlines the alternating DRSB loop and experiment implementation.

D.1.1 DERIVATION OF THE ADJOINT SYSTEM

  • A Lagrange multiplier for the dynamics yields an adjoint system through first variation, which is then used to obtain the initial-state cost-to-go gradient.Source: §D.1.1 · p. 20

D.1.2 DERIVATION OF THE BASIC ADJOINT MATCHING OBJECTIVE

  • The basic adjoint-matching objective can be optimized by gradient descent, with trajectory dependence on controller parameters canceled by the adjoint equations.Trajectories are detached during simulation, avoiding retention of the full simulation graph but not trajectory-generation or network-update costs.Source: §D.1.2 · pp. 21–22

D.1.3 CONNECTION TO OPTIMAL CONTROL, THE LEAN ADJOINT STATE, AND THE ADJOINT MATCHING OBJECTIVE

  • A verified SOC optimum satisfies the functional stationarity condition, but this does not establish global optimality for arbitrary neural-network stationary points.The lean adjoint is constant in time in the stated zero-drift, zero-state-cost setting and is not equivalent to the cost-to-go gradient.Source: §D.1.3 · pp. 22–23

E FULL DRSB ALGORITHM

  • Each DRSB iteration updates the worst-case distribution, terminal-cost network, and controller; Wasserstein and Sinkhorn differ in the worst-case subroutine.The Wasserstein routine iteratively updates samples and λ, whereas Sinkhorn uses a closed-form λ and Gaussian worst-case sampling.Source: E FULL DRSB ALGORITHM · p. 24

F IMPLEMENTATION DETAILS

  • Experiments use a single NVIDIA RTX 4090 GPU and multilayer perceptrons for both the drift and terminal-cost networks.For FFHQ translation, the model operates in the 512-dimensional latent space of the ALAE autoencoder.Source: F IMPLEMENTATION DETAILS · p. 25

F.1 PRACTICAL TECHNIQUES FOR ADVERSARIAL TRAINING

The appendix describes adversarial-training stabilizers, examines Sinkhorn sensitivity to κ, and reports noise-robustness results across two transport tasks. DRSB-S leads on GMM at all tested noise levels except the lowest, while both DRSB variants outperform DSBM at higher noise on the Gaussian task.

F.1 PRACTICAL TECHNIQUES FOR ADVERSARIAL TRAINING

  • Three stabilizers ease adversarial training: ramp λ for DRSB-W or ε for DRSB-S, pretrain terminal-cost gϕ, and regularize uθ toward the pretrained SB drift.For DRSB-S, the ε ramp limits extreme shifts from the start; the pretrained drift also guides terminal-cost initialization and controller regularization.Source: §F.1 · p. 25

F.2 SENSITIVITY TO κ IN DRSB-S

  • For DRSB-S, larger κ widens worst-case support but excessive κ destabilizes training; on FFHQ with d = 512, κ = 0.005 and 0.32 are stable, while κ = 0.25 or ε > 30 diverges.κ near 0 or 1/π reduces the sensitivity to dimension d.Source: §F.2 · p. 26

G ADDITIONAL EXPERIMENTS / G.1 NOISE ROBUSTNESS

  • On eight noise levels, DSBM leads at low noise on Gaussian transport, while both DRSB variants win at higher noise; DRSB-S leads GMM except δ = 0.5.The additional experiments report SWD for both tasks across L2-norm noise levels from δ = 0.5 to 4.0.Source: Table 6 · p. 27, §G.1 · p. 27

G.2 SHIFT ROBUSTNESS

DRSB-S achieves the lowest SWD under most tested shift directions in both tasks. The cost analysis finds the computed adversarial distribution has the largest measured fixed-terminal-cost value among tested distributions, but does not certify global optimality or the full objective.

G.2 SHIFT ROBUSTNESS

  • DRSB-S has the lowest SWD under most shift directions in both tasks; DRSB-W also beats DSBM on the original GMM distribution.The reported comparisons cover the original distribution, directional shifts, and each variant's worst-case inputs.Source: Table 7 · p. 27, §G.2 · p. 27

H.1 WORST-CASE COST INTERPRETATION

  • A globally maximizing initial distribution upper-bounds fixed-controller, fixed-terminal-cost values across the ambiguity set, but this is not a certificate for the full KL objective.The bound requires global inner optimality; approximate adversarial updates need not attain it, and the ambiguity constraint need not be active.Source: §H.1 · p. 28

H.2 EMPIRICAL COST COMPARISON

  • Among the computed adversarial, four shifted, and nominal distributions, the adversarial distribution has the largest measured cost; this finite comparison does not establish global inner optimality.The control-cost component saturates as ρ increases while the terminal-cost term receives greater weight.Source: Figure 5 · p. 28, §H.2 · p. 28

I ADDITIONAL EXPERIMENTAL ANALYSIS

The appendix reports measured training costs, stabilization diagnostics, and an empirical check of the stochastic-optimal-control gradient relation. The diagnostics support approximate stabilization and relation-consistency on the evaluated 2D Gaussian experiment, not general convergence or exact optimality.

I.1 MEASURED COMPUTATION COST

  • Including DRSB's additional stage, total training costs are about 1.54× DSBM on Gaussian transport and 1.16× on FFHQ; worst-case optimization uses less than 1.5% of measured runtime.Gradient updates account for about 80% of Gaussian training time and over 90% on FFHQ; these workloads do not establish runtime scaling with dimension.Source: Table 8 · p. 29, §I.1 · p. 29

I.2 EMPIRICAL CONVERGENCE

  • On the 2D Gaussian task, Rk < 0.005 from epoch 25 onward and λ50 = 3.483, close to the active-constraint SOC prediction of 3.496.The terminal critic output and control magnitude change less at later epochs; these summaries measure stabilization, not stationarity of all three blocks.Source: Table 9 · p. 29, §I.2 · p. 29
  • The variational objective stabilizes near 312 after a peak of 313.44 at k = 32; a separate constant-learning-rate run also settles near 312.This provides an additional check beyond the decreasing learning-rate schedule, without proving convergence of the block algorithm.Source: Table 10 · p. 29, §I.2 · p. 29

I.3 EMPIRICAL CHECK OF THE SOC OPTIMALITY CONDITION

  • On evaluated Gaussian samples, the SOC gradient and control directions have cosine 0.9996–0.9998 and magnitude ratios 0.9310–1.0761, supporting approximate—not exact or general—relation satisfaction.At epoch 50, the cosine is 0.9998 and the ratio is 1.0089.Source: Table 11 · p. 30, §I.3 · p. 30

Related papers

Loading 2610.02043v1…