Distributionally Robust Schrödinger Bridge

2026-10-01

AI Summary

This study addresses the failure of standard Schrödinger bridges to recover target distributions under initial distribution shifts by proposing a robust Schrödinger bridge framework. By integrating stochastic optimal control with distributionally robust optimization, we establish an exact variational formulation and design an alternating algorithm to learn a single controller that minimizes the worst-case objective under initial distributional uncertainty. Furthermore, Wasserstein and Sinkhorn gradient approximations are derived to enable efficient computation. Experimental results demonstrate that the proposed method significantly enhances robustness to input perturbations in tasks such as image translation, effectively reducing the sliced Wasserstein distance even under unseen noise levels.

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.

Problem

• Schrödinger Bridge
• Distributional Robustness
• Distribution Shift
• Stochastic Transport

Innovation

• Distributionally Robust Schrödinger Bridge
• Stochastic Optimal Control
• Variational Formulation
• Alternating Algorithm
• Sinkhorn Divergence