Distributionally Robust Schr\"odinger Bridge
Schrödinger bridge (SB) transport assumes a fixed initial distribution, which breaks at test time when inputs shift. This paper introduces Distributionally Robust Schrödinger Bridge (DRSB), which learns a single controller by minimizing the worst-case sum of control energy $\int_0^1 \frac{1}{2}\|u_t^\theta(X_t)\|^2 dt$ and a terminal KL penalty $\rho \mathrm{KL}(p_1^{u,\mu}\|\nu)$ over an ambiguity set $\mathfrak{M}_\epsilon(\hat{\mu})$ of initial distributions. The authors derive a variational reformulation, connect it to stochastic optimal control and distributionally robust optimization, and develop Wasserstein and Sinkhorn variants with alternating block updates. Experiments on 2D transport and FFHQ image-to-image translation show improved robustness to input perturbations, with a tradeoff in nominal performance.
The paper presents a conceptually natural but largely incremental extension of SB to distributional robustness, and the Sinkhorn variant’s closed-form Gaussian adversary is an algorithmic highlight. Overall, however, the block-coordinate training scheme lacks convergence guarantees, the empirical gains are modest and narrowly benchmarked, and the reliance on an approximate SOC optimality condition weakens the theoretical grounding. The result is a competent, well-motivated contribution that does not significantly advance the state of the art in either robust optimal transport or generative modeling.
The strongest aspects are the rigorous variational derivation and the practical algorithmic reduction of the inner adversarial problem. Proposition 3.3 gives a sound neural estimator for the log-density ratio that drives the terminal KL penalty, and Corollary 3.7 yields a closed-form Gaussian sampling rule for the Sinkhorn adversary that avoids iterative inner optimization.
The block-coordinate algorithm is a heuristic for a nonconvex minimax problem. The authors explicitly note that "These updates neither exchange the minimax operators nor establish convergence of the full nested problem," leaving open whether the procedure reaches a saddle point or cycles. Moreover, the Wasserstein and Sinkhorn adversarial updates rely on the SOC relation $u_0^*(x)=-\nabla V(u^*,x)$, which the paper admits is "exact at the fixed-g SOC optimum and is a surrogate for a controller after finite optimization"; no empirical check is provided for the high-dimensional image model.
Empirically, the image-to-image translation results (Table 3) report FID differences of less than one point with overlapping standard deviations, suggesting the improvements may not be statistically significant. The evaluation is also narrow: noise is injected in a fixed ALAE latent space, and comparisons with stronger robustness baselines—such as multi-level augmentation or unbalanced OT—are absent. The claim of better generalization to unseen noise rests on only two training levels and four test levels on a single 2D task.
The evidence supports the narrow claim that DRSB improves robustness relative to a plain SB baseline. Table 1 shows consistent reductions in sliced Wasserstein distance under Gaussian, shifted, and worst-case perturbations, and the Sinkhorn variant convincingly dominates the Wasserstein variant in most conditions. However, comparisons to related work are only conceptual: the authors do not empirically benchmark against robust OT methods such as unbalanced OT or soft-constrained SB, and the noise-augmentation comparison in Table 2 is limited to two fixed training levels. The nominal-versus-robust tradeoff is demonstrated but not theoretically characterized.
Algorithm 1 and the reported training times give a coarse recipe, but crucial implementation details—network architectures, full hyperparameter schedules, and random seeds—are relegated to Appendix I, which is not fully reproduced in the accessible text. The authors state that "Appendix I.1 reports the training configurations and stage-wise costs," yet without open-source code or checkpoints, independent reproduction would require substantial guesswork. The Wasserstein variant uses up to 90 inner gradient steps with manually tuned step sizes, while the Sinkhorn variant depends on a first-order approximation whose accuracy is not publicly validated.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.