Distributionally Robust Schrödinger Bridge

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

cs.LGOct 1, 2026

Distributionally Robust Schrödinger Bridge

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

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

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Deployment under initial distribution uncertainty. Blue circles denote nominal samples from $\mu$ , and purple squares denote shifted test samples from $\mu_{\text{test}}$ . The desired target $\nu$ is unchanged. (a) A standard SB controller matches $\nu$ for $\mu$ , but this matching need not persist under an initial distribution shift. (b) DRSB trains one controller to reduce the worst-case sum of control energy and terminal mismatch over an ambiguity set of initial distributions.

	Gaussian-to-Gaussian		GMM-to-GMM
Input	DSBM	DRSB	DSBM	DRSB
Org	0.038 $\pm$ 0.007	0.123 $\pm$ 0.012	0.902 $\pm$ 0.242	1.165 $\pm$ 0.138
Noisy	0.684 $\pm$ 0.013	0.638 $\pm$ 0.016	2.407 $\pm$ 0.243	2.212 $\pm$ 0.413
Shifted	0.666 $\pm$ 0.011	0.501 $\pm$ 0.014	5.917 $\pm$ 0.127	5.840 $\pm$ 0.095
Worst	0.951 $\pm$ 0.013	0.710 $\pm$ 0.013	1.988 $\pm$ 0.100	1.478 $\pm$ 0.100

Table 1: SWD (

\downarrow

) under input perturbations (mean

\pm

std, 10 seeds). Org denotes the original input; the other conditions are defined in the text.

	DSBM	DSBM + Aug		DRSB-S
$\delta$		$\delta_{\mathrm{train}}=1$	$\delta_{\mathrm{train}}=2$
Org	0.902 $\pm$ 0.242	0.850 $\pm$ 0.178	0.907 $\pm$ 0.153	1.165 $\pm$ 0.138
$1.0$	1.287 $\pm$ 0.112	1.042 $\pm$ 0.246	1.119 $\pm$ 0.196	1.150 $\pm$ 0.206
$2.0$	1.842 $\pm$ 0.086	1.881 $\pm$ 0.145	1.800 $\pm$ 0.153	1.636 $\pm$ 0.157
$3.0$	2.175 $\pm$ 0.118	2.298 $\pm$ 0.146	2.220 $\pm$ 0.165	1.953 $\pm$ 0.095
$4.0$	2.407 $\pm$ 0.243	2.425 $\pm$ 0.129	2.308 $\pm$ 0.138	2.212 $\pm$ 0.413

Table 2: GMM-to-GMM SWD (

\downarrow

) with noise augmentation (mean

\pm

std, 10 seeds). Test levels

\delta=3,4

are unseen by both augmentation baselines.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Training dynamics. (a) Evolution of $\lambda$ in DRSB-W. (b) $\varepsilon$ -schedule for DRSB-S and corresponding $\lambda$ evolution.

Model	Org	$\delta=0.25$	$0.50$	$1.00$
DSBM	25.55 $\pm$ 0.50	26.59 $\pm$ 0.48	26.15 $\pm$ 0.49	26.81 $\pm$ 0.52
DRSB-S $\kappa$ =0.005, $\varepsilon$ =9	26.76 $\pm$ 0.49	26.28 $\pm$ 0.44	26.60 $\pm$ 0.52	27.15 $\pm$ 0.49
DRSB-S $\kappa$ =0.005, $\varepsilon$ =12	26.36 $\pm$ 0.48	26.48 $\pm$ 0.48	26.10 $\pm$ 0.50	26.66 $\pm$ 0.52
DRSB-S $\kappa$ =0.005, $\varepsilon$ =15	26.08 $\pm$ 0.47	26.57 $\pm$ 0.50	26.35 $\pm$ 0.53	26.40 $\pm$ 0.53
DRSB-S $\kappa$ =0.32, $\varepsilon$ =7	26.70 $\pm$ 0.48	25.93 $\pm$ 0.48	26.05 $\pm$ 0.47	26.05 $\pm$ 0.50
DRSB-S $\kappa$ =0.32, $\varepsilon$ =12	26.52 $\pm$ 0.50	25.82 $\pm$ 0.49	25.98 $\pm$ 0.49	25.94 $\pm$ 0.47

Appendix

Table 4: FID (

\downarrow

) sensitivity to

\kappa

and

\varepsilon

on FFHQ Woman

\to

Man under noise

\delta

Model	$\delta$ =0.5	1.0	1.5	2.0	2.5	3.0	3.5	4.0
(a) Gaussian-to-Gaussian
DSBM	0.052 $\pm$ 0.006	0.072 $\pm$ 0.009	0.128 $\pm$ 0.012	0.214 $\pm$ 0.010	0.313 $\pm$ 0.015	0.431 $\pm$ 0.012	0.555 $\pm$ 0.013	0.684 $\pm$ 0.013
DRSB-W	0.080 $\pm$ 0.010	0.093 $\pm$ 0.012	0.138 $\pm$ 0.013	0.216 $\pm$ 0.013	0.315 $\pm$ 0.012	0.425 $\pm$ 0.013	0.548 $\pm$ 0.013	0.678 $\pm$ 0.013
DRSB-S	0.131 $\pm$ 0.011	0.134 $\pm$ 0.014	0.165 $\pm$ 0.015	0.227 $\pm$ 0.016	0.314 $\pm$ 0.015	0.414 $\pm$ 0.015	0.522 $\pm$ 0.015	0.638 $\pm$ 0.016
(b) GMM-to-GMM
DSBM	1.164 $\pm$ 0.137	1.287 $\pm$ 0.112	1.597 $\pm$ 0.077	1.842 $\pm$ 0.086	2.051 $\pm$ 0.116	2.175 $\pm$ 0.118	2.254 $\pm$ 0.108	2.407 $\pm$ 0.243

Appendix

Table 5: SWD (

\downarrow

) on 2D transport tasks under varying

L_{2}

-norm noise levels

\delta

(mean

\pm

std over 10 seeds). Bold indicates best performance; gray indicates outperforming the baseline DSBM.

(a) Gaussian-to-Gaussian
Model	Org	Up	Down	Left	Right	W-wst	S-wst
DSBM	0.038 $\pm$ 0.007	0.665 $\pm$ 0.012	0.673 $\pm$ 0.020	0.666 $\pm$ 0.011	0.651 $\pm$ 0.009	1.081 $\pm$ 0.021	0.951 $\pm$ 0.013
DRSB-W	0.144 $\pm$ 0.007	0.691 $\pm$ 0.020	0.698 $\pm$ 0.020	0.517 $\pm$ 0.014	0.792 $\pm$ 0.014	0.972 $\pm$ 0.007	0.812 $\pm$ 0.012
DRSB-S	0.123 $\pm$ 0.012	0.647 $\pm$ 0.018	0.640 $\pm$ 0.014	0.501 $\pm$ 0.014	0.769 $\pm$ 0.016	0.881 $\pm$ 0.018	0.710 $\pm$ 0.013
(b) GMM-to-GMM
DSBM	0.902 $\pm$ 0.242	6.167 $\pm$ 0.233	5.798 $\pm$ 0.245	5.917 $\pm$ 0.127	5.917 $\pm$ 0.132	2.086 $\pm$ 0.166	1.988 $\pm$ 0.100

Appendix

Table 6: SWD (

\downarrow

) on 2D transport tasks under distribution shifts (mean

\pm

std over 10 seeds). Org: original distribution; Up/Down/Left/Right: directional shift; W-wst: DRSB-W worst case; S-wst: DRSB-S worst case. Bold indicates best performance; gray indicates outperforming the baseline DSBM.

Panel of Figure 6 — Figure 6: Cost comparison within the Wasserstein ambiguity set. Left : Evaluated distributions with the ambiguity-set boundary (dashed). Right : Measured cost versus $\rho$ .

	Gaussian DSBM	Gaussian DRSB	FFHQ DSBM	FFHQ DRSB
Iterations	140	50	70	50
Gradient steps (total)	224k	101k	3.36M	320k
Trajectories	4,096	1,024	20,480	40,960
Wall-clock time	754 s	404 s	8.5 h	1.4 h
Stage time / DSBM time	$1\times$	$0.54\times$	$1\times$	$0.16\times$

Appendix

Table 7: Measured training costs. DRSB columns report the additional stage after DSBM pretraining. Trajectory counts follow the reported training configurations.

$k$	5	10	15	20	25	30	35	40	45	50
$R_{k}$	0.260	0.319	0.240	0.111	0.002	0.004	0.002	0.000	0.002	0.001
$\lambda_{k}$	6.601	5.696	4.535	3.627	3.427	3.457	3.477	3.488	3.494	3.483
$G_{k}$	0.752	1.232	1.735	2.142	2.068	1.757	1.672	1.631	1.524	1.550
$U_{k}$	22.94	22.92	23.19	23.74	24.10	24.22	24.29	24.36	24.32	24.30

Appendix

Table 8: Reported update diagnostics on the 2D Gaussian task.

$k$	5	10	20	30	40	50
$\mathcal{L}$	270.99	276.16	304.99	311.99	311.44	311.37

Appendix

Table 9: Variational objective at the reported epochs.

$k$	10	20	30	40	50
$\cos\left(\nabla_{X_{0}}V,-u^{\theta}(0,X_{0})\right)$	0.9997	0.9996	0.9996	0.9997	0.9998
$\\|\nabla_{X_{0}}V\\|/\\|u^{\theta}(0,X_{0})\\|$	0.9310	1.0761	1.0249	1.0278	1.0089

Appendix

Table 10: SOC consistency diagnostics on the 2D Gaussian task.

Explore similar work

Data-to-Energy Stochastic Dynamics
Sep 30, 2025Kirill Tamogashev, Esmeralda S. WhitammerSchrödinger Bridges Diffusion Models
University of Edinburgh
QDSB: Quantized Diffusion Schrödinger Bridges
May 12, 2026Tobias Fuchs, Florian Kalinke, Nadja KleinSchrödinger Bridges Bridge
Karlsruhe Institute of Technology, Karlsruhe, Germany
Twisted Schrödinger Bridge Matching
Jul 18, 2026Maxence Noble, Marie Scheid, Yazid Janati +2Schrödinger Bridges Differentiable Optimal Transport
CMAP, Ecole polytechnique · Institute of Foundation Models · MBZUAI +1

Sep 30, 2025, cs.LGJ/K move · Enter open · S save

Data-to-Energy Stochastic Dynamics

Kirill Tamogashev, Esmeralda S. Whitammer

University of Edinburgh

Schrödinger Bridges Diffusion Models

Open paper AskPDF arXiv

The Schrödinger bridge problem is concerned with finding a stochastic dynamical system bridging two marginal distributions that minimises a certain transportation cost. This problem, which represents a generalisation of optimal transport to the stochastic case, has received attention due to its connections to diffusion models and flow matching, as well as its applications in the natural sciences. However, all existing algorithms allow to infer such dynamics only for cases where samples from both distributions are available. In this paper, we propose the first general method for modelling Schrödinger bridges when one (or both) distributions are given by their unnormalised densities, with no access to data samples. Our algorithm relies on a generalisation of the iterative proportional fitting (IPF) procedure to the data-free case, inspired by recent developments in off-policy reinforcement learning for training of diffusion samplers. We demonstrate the efficacy of the proposed data-to-energy IPF on synthetic problems, finding that it can successfully learn transports between multimodal distributions. As a secondary consequence of our reinforcement learning formulation, which assumes a fixed time discretisation scheme for the dynamics, we find that existing data-to-data Schrödinger bridge algorithms can be substantially improved by learning the diffusion coefficient of the dynamics. Finally, we apply the newly developed algorithm to the problem of sampling posterior distributions in latent spaces of generative models, thus creating a data-free image-to-image translation method. Code: https://github.com/mmacosha/d2e-stochastic-dynamics

Abstract

Explore similar work

Data-to-Energy Stochastic Dynamics

QDSB: Quantized Diffusion Schrödinger Bridges

Twisted Schrödinger Bridge Matching