Paper deep dive
Distributionally Robust Schrödinger Bridge
Jinhwan Sul, Panagiotis Theodoropoulos, Vincent Pacelli, Jaemoo Choi, Evangelos Theodorou
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 10/4/2026, 2:20:11 AM
Summary
The paper introduces Distributionally Robust Schrödinger Bridge (DRSB), a method to learn a single stochastic transport controller that is robust to shifts in the initial distribution. DRSB minimizes the worst-case sum of control energy and terminal distribution mismatch (KL penalty) over an ambiguity set of initial distributions. The authors derive a variational formulation connecting DRSB to stochastic optimal control and distributionally robust optimization, proposing an alternating algorithm with Wasserstein and Sinkhorn variants for adversarial updates. Experiments demonstrate improved robustness to input perturbations in 2D transport and image-to-image translation tasks compared to standard Schrödinger Bridge.
Entities (9)
Relation Signals (8)
Distributionally Robust Schrödinger Bridge → extends → Schrödinger Bridge
confidence 95% · The DRSB formulation is Problem 3.1... DRSB addresses the choice of a single transport controller... When additionally ρ→∞, the KL penalty enforces p1u,μ=ν and we recover the standard SB
Distributionally Robust Schrödinger Bridge → uses → Ambiguity Set
confidence 92% · 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.
Wasserstein DRSB → isvariantof → Distributionally Robust Schrödinger Bridge
confidence 90% · We develop Wasserstein and Sinkhorn variants using stochastic control optimality conditions
Sinkhorn DRSB → isvariantof → Distributionally Robust Schrödinger Bridge
confidence 90% · We develop Wasserstein and Sinkhorn variants using stochastic control optimality conditions
Distributionally Robust Schrödinger Bridge → appliedto → Gaussian Mixture Transport
confidence 88% · On Gaussian mixture transport, Sinkhorn DRSB also achieves lower mean sliced Wasserstein distance
Distributionally Robust Schrödinger Bridge → appliedto → Image-to-Image Translation
confidence 88% · Experiments on two-dimensional transport tasks and image-to-image translation show improved robustness
Distributionally Robust Schrödinger Bridge → connectsto → Stochastic Optimal Control
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.02043v1
- Canonical: https://arxiv.org/abs/2610.02043v1
Trouble viewing inline? Open PDF directly →
Full Text
106,931 characters extracted from source content.
Expand or collapse full text
Distributionally Robust Schrödinger Bridge Jinhwan Sul Panagiotis Theodoropoulos Vincent Pacelli Jaemoo Choi Evangelos Theodorou Department of Aerospace Engineering Georgia Institute of Technology, Atlanta, GA, USA jsul7,ptheodor3,vpacelli3,jchoi843,evangelos.theodorou@gatech.edu 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. 1 Introduction Generative modeling is fundamentally a task of learning a transport map from a source distribution to a target distribution. While early approaches, such as GANs (Goodfellow et al., 2014), focused on learning static maps, recent approaches such as diffusion models (Ho et al., 2020; Song et al., 2021) learn dynamic transport maps via stochastic differential equations (SDEs). The Schrödinger Bridge (SB) problem generalizes standard diffusion models and likewise seeks an optimal transport mapping described by an SDE. By definition, SB seeks a stochastic process that connects two arbitrary distributions and minimizes the KL divergence with respect to a reference stochastic process (Schrödinger, 1931; Léonard, 2014). SB has been used in applications such as sampling (Liu et al., 2025), image-to-image translation (Theodoropoulos et al., 2024), and image restoration (Liu et al., 2023). However, the SB dynamics are trained for a prescribed initial distribution. When this distribution shifts at test time, the same dynamics produce a different terminal distribution, violating the original target marginal constraint. This motivates a distributionally robust SB formulation that accounts for uncertainty in the initial distribution. In this paper, we formulate the Distributionally Robust Schrödinger Bridge (DRSB) problem to achieve robust generation under initial distribution shifts. The goal of DRSB is to find a single controller that transports every initial distribution in a prescribed set close to the target distribution. This set, called the ambiguity set, consists of distributions within a prescribed distance of the nominal initial distribution. We connect DRSB to stochastic optimal control through a terminal cost that penalizes mismatch between the induced terminal distribution and the target distribution. This cost takes the form of the log-density ratio between these distributions, which we approximate with a neural network to provide a training signal for the controller. Since the induced terminal distribution depends on both the controller and the initial distribution, the terminal cost must be re-estimated as these change. For a fixed terminal cost, the resulting subproblem has the standard distributionally robust optimization (DRO) form (Mohajerin Esfahani & Kuhn, 2018; Gao & Kleywegt, 2023), where the control variable serves as the DRO decision variable. Strong duality results from the DRO literature provide a dual formulation of the inner worst-case maximization over the ambiguity set. We update the adversarial initial distribution using this dual formulation and train the controller via adjoint matching (Domingo-Enrich et al., 2024) to minimize the expected transport cost under the current adversarial distribution. We alternate these updates with re-estimation of the terminal cost, training its network as a binary classifier that distinguishes the resulting terminal samples from target samples. We develop Wasserstein and Sinkhorn variants of DRSB, with adversarial input updates derived from their DRO dual formulations. Using an optimality condition from stochastic optimal control theory (Kappen, 2005), we approximate cost gradients with controller evaluations. The Wasserstein variant iteratively updates adversarial inputs. For the Sinkhorn variant, a first-order approximation of the cost yields a Gaussian sampling rule for each nominal input and a closed-form update of the dual variable. This approximation allows adversarial inputs to be sampled without iterative inner optimization. We evaluate DRSB on 2D Gaussian-to-Gaussian and GMM-to-GMM transport tasks and on image-to-image translation. The results show that DRSB improves robustness to noisy and shifted initial distributions compared with standard SB. On the 2D transport tasks, Sinkhorn DRSB outperforms Wasserstein DRSB under most tested perturbation conditions. The main contributions of this paper are summarized as follows. • We formulate Distributionally Robust Schrödinger Bridge problems to address initial distribution shift. • We derive an exact variational formulation of DRSB and connect its fixed-terminal-cost subproblem to SOC and DRO. • We develop block update algorithms for the Wasserstein and Sinkhorn variants using SOC optimality conditions. 2 Preliminaries 2.1 Schrödinger Bridge under Initial Distribution Uncertainty The standard Schrödinger bridge (SB) problem (Schrödinger, 1931; Léonard, 2014) is Problem 2.1 (Schrödinger Bridge). minutσ2KL(pu,μ∥pbase)=minutX∼pu,μ[∫0112∥utθ(Xt)∥2dt], \,\, _u_t\,σ^2KL(p^u,μ\|p^base)= _u_tE_X p^u,μ [ _0^1 12\|u^θ_t(X_t)\|^2dt ], (1a) s.t. dXt=utθ(Xt)dt+σdWt,X0∼μ(X0),X1∼ν(X1). dX_t=u^θ_t(X_t)dt+σ dW_t, X_0 μ(X_0), X_1 ν(X_1). (1b) Here, pu,μp^u,μ denotes the path measure induced by Equation 1b with the initial law μ, and pbase:=pu:=0p^base:=p^u:=0 denotes the base path measure when ut:=0u_t:=0. The base path measure has the same initial law μ. The equality assumes the conditions for Girsanov’s theorem. Xt∈ℝdX_t ^d is the d-dimensional state vector at time t. The control uθ:[0,1]×ℝd→ℝdu^θ:[0,1]×R^d ^d is parameterized by θ. WtW_t is a standard d-dimensional Brownian motion, and σ is a constant positive diffusion coefficient. p0u,μp^u,μ_0 and p1u,μp^u,μ_1 are the marginal distributions of pu,μp^u,μ at times t=0t=0 and t=1t=1, respectively. Then, by definition, p0u,μ=μp_0^u,μ=μ always holds. The Schrödinger bridge (SB) problem seeks the optimal control by minimizing the KL divergence between pu,μp^u,μ and pbasep^base while transporting the initial data distribution μ to the target data distribution ν. Deployment with an unknown initial distribution. Let u⋆u solve Problem 2.1 for the nominal initial distribution μ. For this fixed controller, the Fokker–Planck equation defines the map ΦSB _SB from initial to terminal distributions, ΦSB:p0u⋆,μ↦p1u⋆,μ,ΦSB(μ)=ν. _SB:p^u ,μ_0 p^u ,μ_1, _SB(μ)=ν. (2) At deployment, inputs may instead follow an unknown distribution μtest≠μ _test≠μ. Applying the same controller need not preserve terminal matching, i.e. ΦSB(μtest)≠νin general. _SB( _test)≠ν general. (3) We consider a deployment scenario where the test-time initial law is not known during training and a single learned controller must be used without test-time retraining. Figure 1 illustrates this deployment test setting using point clouds. In Figure 1(a), blue circles represent samples from the nominal initial distribution μ, and purple squares represent samples from a perturbed distribution μtest _test. The same state-dependent dynamics act on both input clouds, but a change in the distribution of starting points can change the terminal cloud. The nominal terminal distribution matches ν, whereas the perturbed terminal distribution may deviate from this unchanged target. We represent uncertainty in the test-time initial distribution by a prescribed set of plausible initial distributions, called an ambiguity set. Figure 1(b) illustrates the goal of learning one controller for this set. Without a known probability law over deployment shifts, we aim to minimize the worst-case sum of control energy and terminal mismatch over the ambiguity set, using the same learned controller at deployment. (a) Standard SB (b) DRSB (ours) Figure 1: Deployment under initial distribution uncertainty. Blue circles denote nominal samples from μ, and purple squares denote shifted test samples from μtest _test. The desired target ν is unchanged. (a) A standard SB controller matches ν for μ, 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. 2.2 Distributionally Robust Optimization In stochastic optimization, the underlying data distribution is often unknown, and only a few data samples are available. DRO hedges against this uncertainty by optimizing under the worst-case distribution within a prescribed set of distributions called the ambiguity set (Wiesemann et al., 2014). The DRO problem takes the minimax form minu∈maxμ∈ϵ(μ^)X0∼μ[V(u,X0)], _u _μ∈ M_ε( μ)E_X_0 μ[V(u,X_0)], (4) where the outer minimization selects the decision variable u, and the inner maximization identifies the worst-case distribution μ within the ambiguity set ϵ(μ^):=μ:D(μ,μ^)≤ϵ M_ε( μ):=\μ:D(μ, μ)≤ε\, defined as the set of distributions within distance ϵε of the nominal distribution μ μ under distributional distance D. 3 Distributionally Robust Schrödinger Bridge 3.1 Distributional Robustness in SB and Connection to SOC We formulate and solve the Distributionally Robust Schrödinger Bridge (DRSB) problem, which seeks a control that can robustly steer shifted test-time distributions μtest _test toward the target distribution. All proofs are provided in Appendix A. The DRSB formulation is Problem 3.1 (Distributionally Robust Schrödinger Bridge). minumaxμ∈ϵ(μ^) _u _μ∈ M_ε( μ) X∼pu,μ[∫0112∥utθ(Xt)∥2dt]+ρKL(p1u,μ∥ν), _X p^u,μ [ _0^1 12\|u^θ_t(X_t)\|^2dt ]+ (p^u,μ_1\|ν), (5a) s.t. dXt=utθ(Xt)dt+σdWt,X0∼μ, dX_t=u^θ_t(X_t)dt+σ dW_t, X_0 μ, (5b) ϵ(μ^)=μ:D(μ,μ^)≤ϵ. M_ε( μ)=\μ:D(μ, μ)≤ε\. (5c) Here, the initial distribution μ is no longer assumed to be known exactly. Instead, the inner maximization in Equation 5a finds the worst-case initial distribution from the ambiguity set ϵ(μ^) M_ε( μ). μ μ is the empirical distribution μ^=1N∑i=1NδX^0i μ= 1N _i=1^N _ X_0^i, where δX^0i _ X_0^i is a Dirac delta function at each sample point X^0i X_0^i. The ambiguity set is centered at μ μ with radius ϵε under the distributional distance D. DRSB addresses the choice of a single transport controller before the test-time input distribution is known. We consider input perturbations for which the desired target distribution ν remains unchanged. The control-energy term in Equation 5a retains the transport-cost criterion of SB, while the soft penalty ρKL(p1u,μ∥ν)ρ\,KL(p^u,μ_1\|ν) relaxes the hard terminal constraint p1u,μ=νp^u,μ_1=ν in Equation 1b. Optimality in Problem 3.1 is defined by minimizing the worst-case sum of control energy and the weighted terminal KL divergence over the ambiguity set. The resulting controller need not preserve the nominal SB transport or be optimal for each initial distribution separately. Nominal performance may therefore be traded for lower worst-case cost. This formulation provides a criterion for selecting one controller when adaptation to each deployment distribution is not possible. Remark 3.2. For a nonnegative distance with D(μ,μ^)=0D(μ, μ)=0 if and only if μ=μ^μ= μ, when ϵ=0ε=0, the ambiguity set reduces to 0(μ^)=μ M_0( μ)=\ μ\, eliminating the inner maximization. When additionally ρ→∞ρ→∞, the KL penalty enforces p1u,μ=νp^u,μ_1=ν and we recover the standard SB in Problem 2.1. To optimize the DRSB objective through stochastic optimal control, we require a terminal cost that represents the mismatch between p1u,μp_1^u,μ and ν. We introduce a terminal cost neural network gϕ:ℝd→ℝg^φ:R^d parameterized by ϕφ that approximates the log-density ratio between p1u,μp_1^u,μ and ν. Proposition 3.3 (Terminal cost approximation). For a fixed admissible pair (u,μ)(u,μ), suppose p1u,μp_1^u,μ and ν have positive densities and finite KL divergence. Suppose the function class contains their log-density ratio and a terminal cost function gϕ:ℝd→ℝg^φ:R^d is an optimal solution of the optimization problem gϕ=argmingX1∼p1u,μ[log(1+e−g(X1))]+X1∼ν[log(1+eg(X1))].g^φ= _gE_X_1 p^u,μ_1 [ (1+e^-g(X_1) ) ]+E_X_1 ν [ (1+e^g(X_1) ) ]. (6) Then, gϕ(x)=log(p1u,μ(x)ν(x)) g^φ(x)= ( p_1^u,μ(x)ν(x) ) and KL(p1u,μ∥ν)=X1∼p1u,μ[gϕ(X1)]−logX1∼ν[egϕ(X1)].KL(p_1^u,μ\|ν)=E_X_1 p_1^u,μ[g^φ(X_1)]- _X_1 ν[e^g^φ(X_1)]. (7) Proposition (3.3) shows that the KL penalty in Equation 5a can be expressed in terms of the terminal cost gϕg^φ using Equation 7. The terminal cost is obtained by solving (6). The discriminator gϕg^φ distinguishes generated samples (X1∼p1u,μX_1 p_1^u,μ) from target samples (X1∼νX_1 ν). As a terminal cost, gϕg^φ guides the terminal distribution p1u,μp_1^u,μ toward the target distribution ν through this representation of the KL penalty. Theorem 3.4 (Variational formulation of DRSB). For each admissible pair (u,μ)(u,μ), apply the variational representation of the terminal KL divergence. Then Problem 3.1 is equivalent to minumaxsupgμ∈ϵ(μ^)X0∼μ[V(u,X0,g)]−ρlogX1∼ν[eg(X1)], _u _μ∈ M_ε( μ) _g \E_X_0 μ[V(u,X_0;g)]-ρ _X_1 ν[e^g(X_1)] \, (8) subject to Equations 5b and 5c, where V(u,x,g):=X⋅|0∼p⋅|0u[∫0112‖utθ(Xt)‖2t+ρg(X1)|X0=x].V(u,x;g):=E_X_·|0 p^u_·|0 [ _0^1 12\|u_t^θ(X_t)\|^2dt+ρ g(X_1)\, |\,X_0=x ]. (9) Here, p⋅|0up^u_·|0 is the controlled path distribution conditioned on the initial state, and the supremum is over measurable g for which the expectations are well defined and 0<ν[eg]<∞0<E_ν[e^g]<∞. The optimizing terminal cost depends on (u,μ)(u,μ) and cannot generally be fixed across the outer optimization. For fixed g, the normalization term is constant in (u,μ)(u,μ), giving the subproblem minumaxμ∈ϵ(μ^)X0∼μ[V(u,X0,g)],s.t. Equations 5b and 5c. _u _μ∈ M_ε( μ)E_X_0 μ[V(u,X_0;g)], .t. @crefcreftypeplural~refnumeq:drsb_dyn and @nobreakspace @crefrefnumeq:drsb_ambiguity. (10) We write V(u,x)=V(u,x,g)V(u,x)=V(u,x;g) when g is held fixed. From the SOC perspective, V(u,x,g)V(u,x;g) in Equation 9 is the cost of a stochastic control problem with terminal cost ρgρ g. For fixed (μ,g)(μ,g), minimizing it gives a control adapted to that initial distribution and terminal cost (Chen et al., 2021; Liu et al., 2022). A shift in the initial distribution can change the terminal marginal. DRSB instead minimizes the worst-case objective over the ambiguity set, retaining the optimization over g in Equation 8. For fixed g, introducing a Lagrange multiplier λ≥0λ≥ 0 gives the dual of Equation 10, with equality when strong duality holds: minuminλ≥0λϵ+maxμX0∼μ[V(u,X0)]−λD(μ,μ^),s.t. (5b). _u _λ≥ 0λε+ _μ \E_X_0 μ [V(u,X_0) ]-λ D(μ, μ) \, .t. ( eq:drsb_dyn). (11) 3.2 Iterative Optimization Scheme for DRSB We use block updates for Equation 8 in the order μ(k)→g(k)→u(k)μ^(k)→ g^(k)→ u^(k), holding the other two variables fixed at each step. First, with (u(k−1),g(k−1))(u^(k-1),g^(k-1)) fixed, the adversarial initial distribution is updated through the DRO dual: μ(k),λ(k)←minλ≥0λϵ+maxμX0∼μ[V(u(k−1),X0,g(k−1))]−λD(μ,μ^). μ^(k),λ^(k)← _λ≥ 0\;λε+ _μ \E_X_0 μ[V(u^(k-1),X_0;g^(k-1))]-λ D(μ, μ) \. (12) We simulate the dynamics with u(k−1)u^(k-1) from μ(k)μ^(k) and denote the resulting terminal marginal by p1(k)=p1u(k−1),μ(k)p_1^(k)=p_1^u^(k-1),μ^(k). The terminal cost update is g(k)=argmingX1∼p1(k)[log(1+e−g(X1))]+X1∼ν[log(1+eg(X1))].g^(k)= _gE_X_1 p_1^(k) [ (1+e^-g(X_1) ) ]+E_X_1 ν [ (1+e^g(X_1) ) ]. (13) At the population optimum, this loss recovers the same log-density ratio as the KL variational objective. With (μ(k),g(k))(μ^(k),g^(k)) fixed, the control subproblem is u(k)←argminuX0∼μ(k)[V(u,X0,g(k))],s.t.(5b).u^(k)← _uE_X_0 μ^(k)[V(u,X_0;g^(k))], .t. ( eq:drsb_dyn). (14) We use the adjoint matching objective (Domingo-Enrich et al., 2024): u(k)=argminuX0∼μ(k)t∼[0,1]Xt,X1∼p⋅|0u¯[∥uθ(t,Xt)+ρ∇g(k)(X1)∥2],u^(k)= _uE_X_0 μ^(k)E_t [0,1]E_X_t,X_1 p u_·|0 [\|u^θ(t,X_t)+ρ∇ g^(k)(X_1)\|^2 ], (15) where p⋅|0u¯p u_·|0 is the path measure induced by u¯=stopgrad(uθ) u= stopgrad(u^θ), conditioned on the initial state. This objective regresses the control toward the negative terminal cost gradient. In practice, each block uses finitely many optimization steps. These updates neither exchange the minimax operators nor establish convergence of the full nested problem. The DRSB training procedure is summarized in Algorithm 1. Following the successive approximation method from robust control theory (Beard & McLain, 1998), we initialize uθu^θ with a pretrained standard SB model. Steps 3–5 remain the same across all DRSB variants. Only Step 2, the adversarial initial distribution update, depends on the choice of ambiguity set distance D. The cost of Step 2 depends on this choice, while the total training cost also includes simulation and network updates. Algorithm 1 Distributionally Robust SB Require: Drift uθu^θ, terminal cost gϕg^φ, nominal distribution μ μ, target ν. Initialize: uθu^θ with a pretrained SB model, gϕg^φ, λ 1: for k=1,2,…,Kk=1,2,…,K do 2: μ(k),λ(k)←μ^(k),λ^(k)← Update via (12) 3: Simulate (5b) from μ(k),u(k−1)μ^(k),u^(k-1) 4: Update g(k)g^(k) via (13) 5: Update u(k)u^(k) via (15) 6: end for 3.3 Variants of DRSB Different choices of the distributional distance D can lead to different DRSB variants. We consider two variants based on the Wasserstein and Sinkhorn distances. We use a fixed-terminal-cost SOC optimality condition to simplify the adversarial updates. See Appendix B for details. DRSB with Wasserstein Distance. The following proposition shows the dual formulation of DRSB with Wasserstein distance. Proposition 3.5 (Dual form of Wasserstein DRSB). The fixed-g subproblem (10) with the Wasserstein-2 ambiguity set ϵW(μ^)=μ:W22(μ,μ^)≤ϵ M^W_ε( μ)=\μ:W^2_2(μ, μ)≤ε\ admits strong duality (Zhang et al., 2025; Gao & Kleywegt, 2023), and the corresponding dual problem (11) specializes to minuminλ≥0λε−1N∑i=1NminX0λ‖X0−X^0i‖2−V(u,X0),s.t. (5b). _u _λ≥ 0 \λ - 1N _i=1^N _X_0 \λ\|X_0- X^i_0\|^2-V(u,X_0) \ \, .t. 10.22217pt( eq:drsb_dyn). (16) Since the nominal distribution is a discrete empirical distribution, the inner minimization must be solved for every sample. Our implementation represents the adversary by one perturbed sample per nominal sample (Gao & Kleywegt, 2023). Equation 16 is a nested optimization problem over λ and the per-sample initial state X0X_0. In practice, we use the value of λ from the previous iteration when optimizing the worst-case samples, just as we use the controller from the previous iteration in Equation 12. We optimize X0X_0 using gradient descent. For a smooth value function and an optimal controller u∗u^* of the fixed-g SOC subproblem, the optimality condition is (Kappen, 2005): u0∗(x)=−∇V(u∗,x).u^*_0(x)=-∇ V(u^*,x). (17) Replacing −∇V(u,x)-∇ V(u,x) by u0(x)u_0(x) uses one network evaluation. This is exact at the fixed-g SOC optimum and is a surrogate for a controller after finite optimization. For outer iteration k, hold u(k−1)u^(k-1) fixed and let j index the inner updates: X0i,(j+1)←X0i,(j)−η[u0(k−1)(X0i,(j))+2λ(j)(X0i,(j)−X^0i)],η : step size.X^i,(j+1)_0← X^i,(j)_0-η [u_0^(k-1)(X^i,(j)_0)+2λ^(j)(X^i,(j)_0- X_0^i) ], η : step size. (18) The dual update is λ(j+1)←λ(j)−ζ(ε−1N∑i=1N∥X0i,(j+1)−X^0i∥2),ζ : step size,λ^(j+1)←λ^(j)-ζ ( - 1N _i=1^N\|X^i,(j+1)_0- X^i_0\|^2 ), ζ : step size, (19) which is valid as a dual update only while λ(j+1)≥0λ^(j+1)≥ 0. DRSB with Sinkhorn Distance. The Sinkhorn distance (Cuturi, 2013) between μ and μ μ is κ(μ^,μ):=minγ∈Γ(μ^,μ)(X^0,X0)∼γ[‖X^0−X0‖2+κlog(dγ(X^0,X0)dμ^(X^0)dX0)],W_κ( μ,μ):= _γ∈ ( μ,μ) \E_( X_0,X_0) γ [\| X_0-X_0\|^2+κ ( dγ( X_0,X_0)d μ( X_0)dX_0 ) ] \, (20) where Γ(μ^,μ) ( μ,μ) is the set of joint distributions between μ μ and μ, and κ>0κ>0 is the entropy regularization parameter. Proposition 3.6 (Dual form of Sinkhorn DRSB). The fixed-g subproblem (10) with the Sinkhorn ambiguity set ϵS(μ^)=μ:κ(μ^,μ)≤ϵ M^S_ε( μ)=\μ:W_κ( μ,μ)≤ε\ admits strong duality (Wang et al., 2025), and the corresponding dual problem (11) specializes to mininfλ>0uλε+1N∑i=1Nλκlog∫e(V(u,X0)−λ‖X0−X^0i‖2)/λκdX0,s.t. (5b), _u _λ>0λ + 1N _i=1^Nλκ e (V(u,X_0)-λ\|X_0- X^i_0\|^2 )/λκdX_0, .t. 10.22217pt( eq:drsb_dyn), (21) The λ=0λ=0 boundary is interpreted through a limit. For λ>0λ>0 and a finite normalizing integral, the worst-case conditional distribution γ∗(X0|X^0i)γ^*(X_0| X^i_0) is obtained as γ∗(X0|X^0i)∝e(V(u,X0)−λ‖X0−X^0i‖2)/λκ.γ^*(X_0| X^i_0) e (V(u,X_0)-λ\|X_0- X^i_0\|^2 )/λκ. (22) In contrast to the Wasserstein case, the worst-case distribution in Sinkhorn DRSB is supported on all of ℝdR^d, even when the nominal distribution μ μ is discrete. We approximate this distribution by linearizing the fixed-g cost. Corollary 3.7 (Gaussian worst-case distribution). Fix (u,g)(u,g) and assume the SOC optimality condition in Equation 17 holds. Suppose ε+κd2log(πκ)>0 + κ d2 (πκ)>0. Under the first-order approximation V(X0)≈V(X^0i)+∇V(X^0i)⊤(X0−X^0i)V(X_0)≈ V( X_0^i)\!+\!∇ V( X_0^i)\! \!(X_0\!-\! X_0^i), where V(x)=V(u,x,g)V(x)\!=\!V(u,x;g), the linearized inner dual in Equation 21 is infλ>0λ(ε+κd2logπκ)+1N∑i=1N[V(X^0i)+‖u0θ(X^0i)‖24λ], _λ>0\;λ ( + κ d2 πκ )+ 1N _i=1^N [V( X_0^i)+ \|u^θ_0( X_0^i)\|^24λ ], (23) where the worst-case conditional distribution is γ∗(X0|X^0i)=(X^0i−u0θ(X^0i)2λ,κ2I),γ^*(X_0| X^i_0)=N ( X^i_0- u^θ_0( X_0^i)2λ\,, κ2I ), (24) and the optimal dual variable admits the closed form: λ∗=(ε+κd2logπκ)−11N∑i=1N[‖u0θ(X^0i)‖24].λ^*= ( + κ d2 πκ )^-1 1N _i=1^N [ \|u^θ_0( X_0^i)\|^24 ]. (25) Corollary 3.7 gives closed-form sampling for the linearized inner problem. 4 Experiments 4.1 Robustness against noisy and shifted distributions on 2D transport tasks We demonstrate the robustness of our DRSB algorithm on the 2D Gaussian-to-Gaussian and GMM-to-GMM transport tasks. We use the sliced Wasserstein distance (SWD) between the terminal and target distributions as our metric. We use a standard SB model trained with the Diffusion Schrödinger Bridge Matching (DSBM) algorithm (Shi et al., 2023) as the baseline. We test each trained model against three perturbation types, Gaussian noise (Noisy), directional shift (Shifted), and the DRSB worst-case distribution (Worst), to the initial distribution at test time. Table 1: SWD (↓ ) under input perturbations (mean ± std, 10 seeds). Org denotes the original input; the other conditions are defined in the text. Gaussian-to-Gaussian GMM-to-GMM Input DSBM DRSB DSBM DRSB Org 0.038± 0.007 0.123± 0.012 0.902± 0.242 1.165± 0.138 Noisy 0.684± 0.013 0.638± 0.016 2.407± 0.243 2.212± 0.413 Shifted 0.666± 0.011 0.501± 0.014 5.917± 0.127 5.840± 0.095 Worst 0.951± 0.013 0.710± 0.013 1.988± 0.100 1.478± 0.100 The results are summarized in Table 1. On the original distribution, DSBM achieves lower SWD in both tasks. This is because DRSB optimizes for worst-case performance over the ambiguity set, which allocates transport capacity across a wider region of the input space at the cost of nominal accuracy. Under all three perturbation types, DRSB outperforms DSBM on both Gaussian and GMM tasks. DSBM degrades under perturbations because the learned transport map is fitted to the training distribution, and inputs that deviate from the nominal support produce suboptimal trajectories. The gap between DRSB and DSBM is largest on the worst-case distribution trained on DRSB. Figure 2: SWD (↓ ), mean ± std over 10 seeds. W-wst and S-wst are the worst-case inputs of DRSB-W and DRSB-S. Gaussian-to-Gaussian GMM-to-GMM Input DRSB-W DRSB-S DRSB-W DRSB-S Org 0.144± 0.007 0.123± 0.012 0.885± 0.226 1.165± 0.138 Noisy 0.678± 0.013 0.638± 0.016 2.315± 0.236 2.212± 0.413 Shifted 0.517± 0.014 0.501± 0.014 5.967± 0.153 5.840± 0.095 W-wst 0.972± 0.007 0.881± 0.018 1.503± 0.052 1.791± 0.102 S-wst 0.812± 0.012 0.710± 0.013 1.946± 0.100 1.478± 0.100 Figure 3: GMM adversarial samples from DRSB-S (left) and DRSB-W (right). Black crosses mark nominal centers. 4.2 Comparison between Wasserstein and Sinkhorn DRSB We compare the two DRSB variants across perturbation types (Figures 3 and 3). DRSB-S achieves lower mean SWD in eight of the ten reported conditions. These results indicate greater robustness of DRSB-S under most of the tested perturbation conditions. Figure 3 shows the computed worst case samples for the GMM task. DRSB-S samples are more dispersed due to the entropy regularization, whereas DRSB-W samples are more concentrated. One possible interpretation is that the broader spread exposes the DRSB-S controller to a wider range of inputs during training, which may help explain its performance under the tested perturbations. 4.3 Noise Augmentation and Generalization to Unseen Noise Levels We compare DRSB-S with a noise-augmented DSBM baseline on GMM-to-GMM transport to determine whether DRSB provides benefits beyond noise augmentation. Each augmentation baseline uses Gaussian noise at δtrain=1 _train=1 or 22, with the number of perturbed samples equal to 20% of the clean dataset size. This makes the number of training steps involving perturbed samples comparable to that in the additional DRSB stage. We evaluate each model at test-time noise levels δ∈1,2,3,4δ∈\1,2,3,4\. Table 2: GMM-to-GMM SWD (↓ ) with noise augmentation (mean ± std, 10 seeds). Test levels δ=3,4δ=3,4 are unseen by both augmentation baselines. DSBM DSBM + Aug DRSB-S δ δtrain=1 _train=1 δtrain=2 _train=2 Org 0.902± 0.242 0.850± 0.178 0.907± 0.153 1.165± 0.138 1.01.0 1.287± 0.112 1.042± 0.246 1.119± 0.196 1.150± 0.206 2.02.0 1.842± 0.086 1.881± 0.145 1.800± 0.153 1.636± 0.157 3.03.0 2.175± 0.118 2.298± 0.146 2.220± 0.165 1.953± 0.095 4.04.0 2.407± 0.243 2.425± 0.129 2.308± 0.138 2.212± 0.413 Table 2 shows that augmentation is effective at the noise level used in training, but its benefit does not transfer consistently to unseen levels. Augmentation at δtrain=1 _train=1 gives the lowest mean SWD at δ=1δ=1 (1.042), yet performs worse than unaugmented DSBM at every tested δ≥2δ\!≥\!2. Augmentation at δtrain=2 _train=2 improves performance at the matched level and at some other levels, but is worse than DSBM at the unseen level δ=3δ=3. DRSB-S achieves lower mean SWD than both augmentation baselines at every tested δ≥2δ\!≥\!2. DRSB trains against an adaptive adversarial distribution within an ambiguity set, without specifying a test noise level. The results support broader generalization across the tested noise magnitudes than fixed-level augmentation. This robustness comes at the cost of higher Org input SWD (1.165, compared with 0.850 and 0.907 for augmentation). 4.4 Image-to-Image Translation under Noise Table 3: FFHQ FID (↓ ) under Gaussian noise δ (mean ± std, 3 seeds). Woman → Man Man → Woman Noise level DSBM DRSB DSBM DRSB Org 25.55± 0.50 26.52± 0.50 29.32± 0.44 29.57± 0.48 δ=0.25δ=0.25 26.59± 0.48 25.82± 0.49 30.02± 0.49 29.69± 0.42 δ=0.5δ=0.5 26.15± 0.49 25.98± 0.49 30.46± 0.43 29.70± 0.51 δ=1.0δ=1.0 26.81± 0.52 25.94± 0.47 29.63± 0.50 29.48± 0.51 Figure 4: Woman-to-Man translation. Left to right: perturbed input, DSBM, and DRSB. We apply DRSB to image-to-image translation on the FFHQ dataset (Karras et al., 2019), which contains images at a resolution of 1024×1024 (C BY-NC-SA 4.0). Following the experimental settings in Theodoropoulos et al. (2024) and Gushchin et al. (2024), we transport images in the 512-dimensional latent space of the ALAE autoencoder (Pidhorskyi et al., 2020). We use FID to evaluate generation quality after decoding the latent vector. The results are shown in Table 3. DRSB outperforms DSBM across all L2L_2-norm noise levels. DSBM achieves a lower FID score than DRSB on nominal inputs. Figure 4 shows a qualitative example. Even when the effect of the added latent-space noise is visually imperceptible in the decoded input image, DSBM generates visible artifacts, whereas DRSB produces a clean translation. 4.5 Additional Experimental Analysis Computation cost. On a single NVIDIA RTX 4090, the additional DRSB stage takes 404 s on Gaussian transport and 1.4 h on FFHQ, after DSBM pretraining of 754 s and 8.5 h. Total training costs are therefore approximately 1.54×1.54× and 1.16×1.16× DSBM. Network gradient updates dominate runtime, accounting for about 80% on Gaussian transport and over 90% on FFHQ. Worst-case optimization accounts for less than 1.5% in the measured runs. DRSB-W uses up to 90 inner gradient steps, whereas DRSB-S samples a Gaussian adversary using the closed-form dual variable of the linearized problem. Appendix I.1 reports the training configurations and stage-wise costs. Empirical convergence. We track the variational objective and update diagnostics over 50 epochs on 2D Gaussian transport with ρ=10ρ=10 and ϵ=12ε=12. The objective stabilizes near 312, with a final value of 311.37. At all reported epochs k≥25k≥ 25, the normalized ambiguity-constraint residual is below 0.005. The final dual variable is 3.483, close to the value 3.496 predicted by the active-constraint SOC relation. The terminal critic mean and mean squared control magnitude also stabilize, ending at 1.550 and 24.30. A separate 120-iteration run at constant learning rate 10−710^-7 also settles near 312, supporting stabilization beyond the decaying learning-rate schedule. These observations do not establish stationarity of all three blocks or convergence to a saddle point (Appendix I.2). SOC optimality condition. The adversarial updates use the relation ∇X0V(u,X0,g)=−u0(X0) _X_0V(u,X_0;g)=-u_0(X_0) from Equation 17. To assess its accuracy for the trained controller, we compute the cost gradient by discrete adjoint recursion and compare it with −uθ(0,X0)-u^θ(0,X_0) on the same Gaussian task. Across the five evaluated epochs, their direction cosines are 0.9996–0.9998 and gradient-to-control magnitude ratios are 0.9310–1.0761. At epoch 50, these values are 0.9998 and 1.0089. Thus, both direction and magnitude approximately satisfy the SOC relation on the evaluated samples, supporting its use in the adversarial update for this experiment. This is an empirical consistency check, not a proof of exact SOC optimality (Appendix I.3). 5 Related Work Distributionally Robust Optimization. Beginning with the work of Scarf et al. (1957), early DRO focused on moment-based ambiguity sets (Ghaoui et al., 2003; Delage & Ye, 2010). Because moments have limited expressive power, the DRO literature shifted toward distance-based ambiguity sets. Wasserstein DRO (Mohajerin Esfahani & Kuhn, 2018; Zhao & Guan, 2018; Blanchet & Murthy, 2019; Gao & Kleywegt, 2023; Zhang et al., 2025) has attracted attention because of the geometric interpretation of the Wasserstein distance. KL divergence offers convexity and a tractable inner maximization (Ben-Tal et al., 2013; Hu & Hong, 2013). However, its absolute continuity requirement and abstract interpretation are limitations. Recently, Sinkhorn DRO has been proposed to improve computational tractability while modeling realistic uncertainty (Wang et al., 2025). Robust Optimal Transport. The OT literature studies how to reduce the sensitivity of transport maps to outliers. One approach relaxes the marginal constraints in OT (Mukherjee et al., 2021). In particular, unbalanced OT (UOT) formulations are used for this purpose (Balaji et al., 2020; Fatras et al., 2021; Choi et al., 2023). More closely related to our approach, the soft-constrained Schrödinger Bridge (Garg et al., 2024) relaxes terminal boundary constraints to prevent overfitting to noisy training samples. Wang et al. (2024) use a UOT distance within a DRO framework to handle outliers in training data. These methods and DRSB address different forms of robustness. DRSB focuses on distribution shifts at test time. 6 Conclusion and Limitations We formulate and solve the DRSB problem, which aims to transport noisy or shifted input distributions to the target distribution. We focus on the Wasserstein and Sinkhorn variants of DRSB. A globally solved Wasserstein inner problem bounds the fixed-terminal-cost objective, while the linearized Sinkhorn inner problem permits Gaussian sampling. Exploring other distances for DRSB is a direction for future work. Minimax adversarial training can be unstable in high-dimensional settings and requires careful hyperparameter tuning and initialization. The additional experiments quantify noise robustness and training overhead. The convergence diagnostics are empirical and do not establish convergence of the full algorithm. References Balaji et al. (2020) Yogesh Balaji, Rama Chellappa, and Soheil Feizi. Robust optimal transport with applications in generative modeling and domain adaptation. Advances in Neural Information Processing Systems, 33:12934–12944, 2020. Beard & McLain (1998) Randal W Beard and Timothy McLain. Successive Galerkin approximation algorithms for nonlinear optimal and robust control. International Journal of Control, 71(5):717–743, 1998. Bellman (1954) Richard Bellman. The theory of dynamic programming. Bulletin of the American Mathematical Society, 60(6):503–515, 1954. Ben-Tal et al. (2013) Aharon Ben-Tal, Dick Den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2):341–357, 2013. Blanchet & Murthy (2019) Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565–600, 2019. Chen et al. (2021) Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard Sinkhorn meets Gaspard Monge on a Schrödinger bridge. SIAM Review, 63(2):249–313, 2021. Choi et al. (2023) Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. Advances in Neural Information Processing Systems, 36:42433–42455, 2023. Cuturi (2013) Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013. Delage & Ye (2010) Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations Research, 58(3):595–612, 2010. Domingo-Enrich et al. (2024) Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.08861, 2024. Fatras et al. (2021) Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In International Conference on Machine Learning, p. 3186–3197. PMLR, 2021. Gao & Kleywegt (2023) Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with Wasserstein distance. Mathematics of Operations Research, 48(2):603–655, 2023. Garg et al. (2024) Jhanvi Garg, Xianyang Zhang, and Quan Zhou. Soft-constrained Schrödinger bridge: a stochastic control approach. In International Conference on Artificial Intelligence and Statistics, p. 4429–4437. PMLR, 2024. Ghaoui et al. (2003) Laurent El Ghaoui, Maksim Oks, and Francois Oustry. Worst-case value-at-risk and robust portfolio optimization: A conic programming approach. Operations Research, 51(4):543–556, 2003. Goodfellow et al. (2014) Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in Neural Information Processing Systems, 27, 2014. Gushchin et al. (2024) Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal Schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024. Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020. Hu & Hong (2013) Zhaolin Hu and L Jeff Hong. Kullback-Leibler divergence constrained distributionally robust optimization. Available at Optimization Online, 1(2):9, 2013. Kappen (2005) Hilbert J Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of Statistical Mechanics: Theory and Experiment, 2005(11):P11011, 2005. Karras et al. (2019) Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 4401–4410, 2019. Léonard (2014) Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A, 34(4):1533–1574, 2014. Liu et al. (2022) Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized Schrödinger bridge. Advances in Neural Information Processing Systems, 35:9374–9388, 2022. Liu et al. (2023) Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. I2SB: Image-to-Image Schrödinger bridge. In International Conference on Machine Learning, 2023. Liu et al. (2025) Guan-Horng Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky TQ Chen. Adjoint Schrödinger bridge sampler. arXiv preprint arXiv:2506.22565, 2025. Mohajerin Esfahani & Kuhn (2018) Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming, 171(1):115–166, 2018. Mukherjee et al. (2021) Debarghya Mukherjee, Aritra Guha, Justin M Solomon, Yuekai Sun, and Mikhail Yurochkin. Outlier-robust optimal transport. In International Conference on Machine Learning, p. 7850–7860. PMLR, 2021. Pidhorskyi et al. (2020) Stanislav Pidhorskyi, Donald A Adjeroh, and Gianfranco Doretto. Adversarial latent autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 14104–14113, 2020. Scarf et al. (1957) Herbert E Scarf, KJ Arrow, and S Karlin. A min-max solution of an inventory problem. Technical report, Rand Corporation Santa Monica, 1957. Schrödinger (1931) Erwin Schrödinger. Über die Umkehrung der Naturgesetze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u. Company, 1931. Shi et al. (2023) Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. Advances in Neural Information Processing Systems, 36:62183–62223, 2023. Song et al. (2021) Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. Theodoropoulos et al. (2024) Panagiotis Theodoropoulos, Nikolaos Komianos, Vincent Pacelli, Guan-Horng Liu, and Evangelos A Theodorou. Feedback Schrödinger bridge matching. arXiv preprint arXiv:2410.14055, 2024. Wang et al. (2025) Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization. Operations Research, 2025. Wang et al. (2024) Zifan Wang, Yi Shen, Michael M Zavlanos, and Karl H Johansson. Outlier-robust distributionally robust optimization via unbalanced optimal transport. Advances in Neural Information Processing Systems, 37:52189–52214, 2024. Wiesemann et al. (2014) Wolfram Wiesemann, Daniel Kuhn, and Melvyn Sim. Distributionally robust convex optimization. Operations Research, 62(6):1358–1376, 2014. Zhang et al. (2025) Luhao Zhang, Jincheng Yang, and Rui Gao. A short and general duality proof for Wasserstein distributionally robust optimization. Operations Research, 73(4):2146–2155, 2025. Zhao & Guan (2018) Chaoyue Zhao and Yongpei Guan. Data-driven risk-averse stochastic optimization with Wasserstein metric. Operations Research Letters, 46(2):262–267, 2018. Appendix A Proofs of Results in Section 3 A.1 Proof of Proposition 3.3 Fix an admissible pair (u,μ)(u,μ) and abbreviate p1u,μp^u,μ_1 as p1up^u_1 in this subsection. Assume that p1up^u_1 and ν have positive densities on a common support and finite KL divergence. The function class is unrestricted, and expectations are population expectations. Lemma A.1 (Variational form of KL). Let Φ(g):=p1u[g(X)]−logν[eg(X)] (g):=E_p^u_1[g(X)]- _ν[e^g(X)] be defined over measurable functions g:→ℝg:X with p1u[|g|]<∞E_p^u_1[|g|]<∞ and 0<ν[eg]<∞0<E_ν[e^g]<∞. Then, for every such g, Φ(g)≤KL(p1u∥ν) (g) (p^u_1\|ν) (26) and equality holds when g∗(x):=log(p1u(x)ν(x)).g^*(x):= ( p^u_1(x)ν(x) ). (27) Proof. Introduce a probability measure QgQ_g such that dQgdν=egν[eg] dQ_gdν= e^gE_ν[e^g] (28) Then g can be written as g(X)=logdQgdν+logν[eg(X)]g(X)= dQ_gdν+ _ν[e^g(X)] (29) and Φ(g) (g) becomes Φ(g) (g) =∫g(X)dp1u(X)−logν[eg(X)] = g(X)\,dp^u_1(X)- _ν[e^g(X)] (30) =∫logdQgdν+logν[eg(X)]dp1u(X)−logν[eg(X)] = \ dQ_gdν+ _ν[e^g(X)] \dp^u_1(X)- _ν[e^g(X)] (31) =∫logdQgdνdp1u = dQ_gdν\,dp^u_1 (32) =∫(logdQgdp1u+logdp1udν)dp1u = ( dQ_gdp^u_1+ dp^u_1dν )\,dp^u_1 (33) =KL(p1u∥ν)−KL(p1u∥Qg)≤KL(p1u∥ν) =KL(p^u_1\|ν)-KL(p^u_1\|Q_g) (p^u_1\|ν) (34) Setting g∗(X):=log(dp1u/dν)g^*(X):= (dp^u_1/dν) gives logν[eg∗]=log∫(dp1u/dν)dν=0 _ν[e^g^*]= (dp^u_1/dν)\,dν=0, hence Φ(g∗)=∫log(dp1udν)dp1u=KL(p1u∥ν). (g^*)= ( dp^u_1dν )dp^u_1=KL(p^u_1\|ν). (35) ∎ Therefore, the optimal g recovers the log-density ratio log(p1u/ν) (p^u_1/ν). Lemma A.2 (Binary Cross-Entropy Minimizer). The minimizer of the binary cross-entropy loss ℒ(g)=x∼p1u[log(1+e−g(x))]+x∼ν[log(1+eg(x))]L(g)=E_x p^u_1 [ (1+e^-g(x) ) ]+E_x ν [ (1+e^g(x) ) ] (36) is g∗(x)=log(p1u(x)/ν(x))g^*(x)= (p^u_1(x)/ν(x) ). Proof. For each x, minimize the population-loss integrand over the scalar z=g(x)z=g(x): ℓx(z) _x(z) =p1u(x)log(1+e−z)+ν(x)log(1+ez), =p^u_1(x) (1+e^-z )+ν(x) (1+e^z ), (37) ℓx′(z) _x (z) =−p1u(x)+ezν(x)1+ez, = -p^u_1(x)+e^zν(x)1+e^z, (38) ℓx′(z) _x (z) =(p1u(x)+ν(x))ez(1+ez)2>0. = (p^u_1(x)+ν(x) )e^z (1+e^z )^2>0. (39) Thus the pointwise minimum satisfies eg∗(x)ν(x)=p1u(x)e^g^*(x)ν(x)=p^u_1(x), giving g∗(x)=logp1u(x)ν(x).g^*(x)= p^u_1(x)ν(x). (40) Integrating the pointwise inequality proves the result over unrestricted measurable functions. A stationary point in neural-network parameters need not satisfy this pointwise condition. ∎ Proposition (3.3) follows from Lemma (A.1) and Lemma (A.2). By Lemma (A.2), the minimizer of the binary cross-entropy loss recovers the log-density ratio g∗(x)=log(p1u(x)/ν(x))g^*(x)= (p^u_1(x)/ν(x)). Substituting g∗g^* into the variational form in Lemma (A.1) achieves Φ(g∗)=KL(p1u∥ν) (g^*)=KL(p^u_1\|ν), which gives Equation 7. A.2 Proof of Theorem 3.4 Proof. For every admissible pair (u,μ)(u,μ), Lemma A.1 gives KL(p1u,μ∥ν)=supgp1u,μ[g(X1)]−logν[eg(X1)].KL (p^u,μ_1\|ν )= _g \E_p^u,μ_1[g(X_1)]- _ν[e^g(X_1)] \. (41) Substitution into Equation 5a, followed by conditioning on X0X_0, yields minumaxsupgμ∈ϵ(μ^)pu,μ[∫0112‖utθ(Xt)‖2t+ρg(X1)]−ρlogν[eg(X1)] _u _μ∈ M_ε( μ) _g \E_p^u,μ [ _0^1 12\|u_t^θ(X_t)\|^2dt+ρ g(X_1) ]-ρ _ν[e^g(X_1)] \ (42) =minumaxsupgμ∈ϵ(μ^)X0∼μ[V(u,X0,g)]−ρlogν[eg(X1)], = _u _μ∈ M_ε( μ) _g \E_X_0 μ[V(u,X_0;g)]-ρ _ν[e^g(X_1)] \, (43) which is Equation 8. This substitution is pointwise in (u,μ)(u,μ) and does not interchange the minimization over u with an inner maximization. For a fixed pair, the normalized maximizer is g∗(x,u,μ)=logp1u,μ(x)ν(x),logν[eg∗(X1,u,μ)]=0.g^*(x;u,μ)= p^u,μ_1(x)ν(x), _ν[e^g^*(X_1;u,μ)]=0. (44) Because this maximizer depends on both u and μ, it cannot be fixed while optimizing them. For a separately fixed g, the logarithmic term is constant with respect to u and μ and may be omitted from their optimization. This gives the fixed-g subproblem in Equation 10, rather than an equivalent formulation with a single fixed g∗g^* for all pairs. ∎ Appendix B Details on the Derivation of DRSB Variants Throughout this section, g is fixed and V(u,x)V(u,x) abbreviates V(u,x,g)V(u,x;g); in an inner distribution update, u is also fixed. The dual identities below concern this fixed-cost subproblem, under the finiteness and regularity assumptions of the cited DRO duality results. B.1 Details on Wasserstein DRSB We start from the fixed-g Wasserstein DRSB subproblem minumaxμ∈ϵ(μ^) _u _μ∈ M_ε( μ) pu[∫0112‖utθ(Xt)‖2t+ρg(X1)] _p^u [ _0^1 12\|u^θ_t(X_t)\|^2dt+ρ g(X_1) ] (45a) s.t. dXt=utθ(Xt)dt+σdWt,X0∼μ dX_t=u^θ_t(X_t)dt+σ dW_t,\,X_0 μ (45b) ϵ(μ^)=μ:W22(μ,μ^)≤ϵ M_ε( μ)=\μ:W_2^2(μ, μ)≤ε\ (45c) where the squared Wasserstein-2 distance is defined as W22(μ,μ^):=min∫γ∈Γ(μ,μ^)‖X0−X^0‖2γ(dX0,dX^0).W_2^2(μ, μ):= _γ∈ (μ, μ) \|X_0- X_0\|^2γ(dX_0,d X_0). (46) Consider the inner maximization maxμ∈ϵ(μ^) _μ∈ M_ε( μ) pu[∫0112‖utθ(Xt)‖2t+ρg(X1)] _p^u [ _0^1 12\|u^θ_t(X_t)\|^2dt+ρ g(X_1) ] (47a) s.t. dXt=utθ(Xt)dt+σdWt,X0∼μ dX_t=u^θ_t(X_t)dt+σ dW_t,\,X_0 μ (47b) ϵ(μ^)=μ:W22(μ,μ^)≤ϵ M_ε( μ)=\μ:W_2^2(μ, μ)≤ε\ (47c) Define the Lagrangian as ℒ(μ,λ) (μ,λ) :=∼pu[∫0112‖utθ(Xt)‖2t+ρg(X1)]−λ(∫‖X0−X^0‖2π(dX0,dX^0)−ε) :=E_X p^u [ _0^1 12\|u^θ_t(X_t)\|^2dt+ρ g(X_1) ]-λ ( \|X_0- X_0\|^2π(dX_0,d X_0)- ) (48) Conditioning on X0X_0, ℒ(μ,λ) (μ,λ) =X0∼μ[V(u,X0)]−λ(∫‖X0−X^0‖2π(dX0,dX^0)−ε) =E_X_0 μ [V(u,X_0) ]-λ ( \|X_0- X_0\|^2π(dX_0,d X_0)- ) (49) where V(u,X0):=⋅|0∼p⋅|0u[∫0112‖utθ(Xt)‖2t+ρg(X1)|X0]V(u,X_0):=E_X_·|0 p^u_·|0 [ _0^1 12\|u^θ_t(X_t)\|^2dt+ρ g(X_1) |X_0 ] is the cost function at t=0t=0. Then we have ℒ(μ,λ) (μ,λ) =λε+∫V(u,X0)π(X0,X^0)dX0dX^0−λ∫‖X0−X^0‖2π(X0,X^0)dX0dX^0 =λ + V(u,X_0)π(X_0, X_0)\,dX_0\,d X_0-λ \|X_0- X_0\|^2π(X_0, X_0)\,dX_0\,d X_0 (50) The dual problem is minλ≥0λε+maxμ(X0,X^0)∼π[V(X0)−λ‖X0−X^0‖2] _λ≥ 0\,λ + _μ\,E_(X_0, X_0) π [V(X_0)-λ\|X_0- X_0\|^2 ] (51) By the interchangeability principle (Zhang et al., 2025), the dual problem can be written as minλ≥0λε+∫maxX0V(X0)−λ‖X0−X^0‖2μ^(X^0)dX^0 _λ≥ 0\,λ + _X_0 \V(X_0)-λ\|X_0- X_0\|^2 \ μ( X_0)\,d X_0 (52) =minλ≥0λε−∫minX0λ‖X0−X^0‖2−V(X0)μ^(X^0)dX^0 = _λ≥ 0\,λ - _X_0 \λ\|X_0- X_0\|^2-V(X_0) \ μ( X_0)\,d X_0 (53) For the empirical distribution μ^=1N∑i=1NδX^0i μ= 1N _i=1^N _ X_0^i, the expectation reduces to a finite sum, giving minλ≥0λε−1N∑i=1NminX0λ‖X0−X^0i‖2−V(X0). _λ≥ 0\,λ - 1N _i=1^N _X_0 \λ\|X_0- X^i_0\|^2-V(X_0) \. (54) B.1.1 Derivation of Update Rules (Equation 18 and Equation 19) Equivalently, the samplewise maximization objective is maxX0ℒ1(X0):=V(X0)−λ‖X0−X^0i‖2. _X_0L_1(X_0):=V(X_0)-λ\|X_0- X_0^i\|^2. (55) Taking the gradient of ℒ1(X0)L_1(X_0) with respect to X0X_0 gives ∇xℒ1(X0)=∇xV(X0)−2λ(X0−X^0i) _xL_1(X_0)= _xV(X_0)-2λ(X_0- X_0^i) (56) If u=u∗u=u^* solves the SOC problem for this fixed g and its value function is differentiable, Equation 17 gives ∇xℒ1(X0)=−u0∗(X0)−2λ(X0−X^0i) _xL_1(X_0)=-u^*_0(X_0)-2λ(X_0- X_0^i) (57) Replacing u∗u^* by the trained drift gives the implemented ascent update. This substitution is approximate unless the fixed-g SOC optimality condition is satisfied: X0←X0+η[−u0θ(X0)−2λ(X0−X^0i)]X_0← X_0+η [-u^θ_0(X_0)-2λ(X_0- X_0^i) ] where η is the gradient ascent step size. Each sample update in Equation 18 requires one evaluation of u0θu^θ_0; total overhead also depends on the number of sample and dual updates. Similarly, we can derive the update rule in Equation 19 from the optimality condition for the worst-case optimization. For Equation 55, the first-order optimality condition is ∇xℒ1(X0(i)∗)=∇xV(X0(i)∗)−2λ(X0(i)∗−X^0i)=0 _xL_1(X^(i)*_0)= _xV(X^(i)*_0)-2λ(X^(i)*_0- X_0^i)=0 (58) Therefore, the worst-case sample satisfies ∇xV(X0(i)∗)=2λ(X0(i)∗−X^0i) _xV(X^(i)*_0)=2λ(X^(i)*_0- X_0^i) (59) Define the objective for the dual variable as infλ≥0f(λ):=λε−1N∑i=1N[λ∥X0(i)∗−X^0i∥2−V(X0(i)∗)] _λ≥ 0f(λ):=λ - 1N _i=1^N [λ\|X^(i)*_0- X^i_0\|^2-V(X^(i)*_0) ] (60) By taking the gradient with respect to λ, we have df(λ)dλ df(λ)dλ =ε+1N∑i=1N[∇xV(X0(i)∗)dX0(i)∗dλ−∥X0(i)∗−X^0i∥2 = + 1N _i=1^N [ _xV(X^(i)*_0) dX_0^(i)*dλ-\|X^(i)*_0- X^i_0\|^2 . −2λ(X0(i)∗−X^0i)dX0(i)∗dλ] .-2λ(X_0^(i)*- X_0^i) dX_0^(i)*dλ ] (61) From Equation 59, we can cancel out the dX0(i)∗dλ dX_0^(i)*dλ terms: df(λ)dλ=ε−1N∑i=1N∥X(i)∗0−X^i0∥2 df(λ)dλ= - 1N _i=1^N\ \|X^(i)*_0- X^i_0\|^2 (62) Consequently, we can perform gradient descent on λ: λ←λ−ζ(ε−1N∑i=1N∥X0(i)∗−X^0i∥2)λ←λ-ζ ( - 1N _i=1^N\|X^(i)*_0- X^i_0\|^2 ) where ζ is the gradient descent step size. This update is valid only while λ remains nonnegative; otherwise feasibility requires projection onto [0,∞)[0,∞). B.2 Details on Sinkhorn DRSB We start from the fixed-g Sinkhorn DRSB subproblem minumaxμ∈ϵS(μ^) _u _μ∈ M^S_ε( μ) pu[∫0112‖utθ(Xt)‖2t+ρg(X1)] _p^u [ _0^1 12\|u^θ_t(X_t)\|^2dt+ρ g(X_1) ] (63a) subject to dXt=utθ(Xt)dt+σdWt,X0∼μ dX_t=u^θ_t(X_t)dt+σ dW_t, X_0 μ (63b) ϵS(μ^)=μ:κ(μ^,μ)≤ϵ M^S_ε( μ)=\μ:W_κ( μ,μ)≤ε\ (63c) where κ(μ^,μ)W_κ( μ,μ) is the Sinkhorn distance with entropy regularization parameter κ: κ(μ^,μ)=minγ∈Γ(μ^,μ)(X^0,X0)∼γ[‖X^0−X0‖2+κlog(dγ(X^0,X0)dμ^(X^0)dX0)].W_κ( μ,μ)= _γ∈ ( μ,μ) \E_( X_0,X_0) γ [\| X_0-X_0\|^2+κ ( dγ( X_0,X_0)d μ( X_0)dX_0 ) ] \. (64) Denote the value of the inner maximization by V. :=maxμ∈ϵS(μ^)X0∼μ[V(X0)] := _μ∈ M^S_ε( μ)E_X_0 μ [V(X_0) ] (65) =maxγminλ≥0λε+X^0∼μ^X0∼γ(X0|X^0) = _γ _λ≥ 0λ +E_ X_0 μE_X_0 γ(X_0| X_0) [V(X0)−λ‖X^0−X0‖2−λκlog(dγ(X0|X^0)dX0)] [V(X_0)-λ\| X_0-X_0\|^2-λκ ( dγ(X_0| X_0)dX_0 ) ] (66) Under the assumptions for strong duality in Sinkhorn DRO (Wang et al., 2025), we switch the order: =minλ≥0λε+X^0∼μ^[maxγ(X0|X^0)X0∼γ(X0|X^0) = _λ≥ 0λ +E_ X_0 μ [ _γ(X_0| X_0)E_X_0 γ(X_0| X_0) [V(X0)−λ∥X^0−X0∥2−λκlog(dγ(X0|X^0)dX0)]] [V(X_0)-λ\| X_0-X_0\|^2-λκ ( dγ(X_0| X_0)dX_0 ) ] ] (67) For λ>0λ>0, κ>0κ>0, and a finite normalizing integral, the entropy-regularized inner maximization admits the following closed-form solution. Boundary cases at λ=0λ=0 are understood by limits. The dual form is =infλ>0λε+X^0∼μ^[λκlog∫e(V(X0)−λ‖X0−X^0‖2)/λκdX0]V= _λ>0λ +E_ X_0 μ [λκ e (V(X_0)-λ\|X_0- X_0\|^2 )/λκdX_0 ] (68) and the optimal conditional distribution is γ∗(X0|X^0)∝e(V(X0)−λ‖X0−X^0‖2)/λκγ^*(X_0| X_0) e (V(X_0)-λ\|X_0- X_0\|^2 )/λκ (69) B.2.1 Second-order approximation We approximate V(X0)V(X_0) up to second order: V(X0)≈V(X^0)+∇V(X^0)⊤δX0+12δX0⊤∇xxV(X^0)δX0V(X_0)≈ V( X_0)+∇ V( X_0) δ X_0+ 12δ X_0 _xV( X_0)\,δ X_0 (70) where δX0=X0−X^0δ X_0=X_0- X_0. Define g:=∇V(X^0)λκg:= ∇ V( X_0)λκ and A:=2κId−∇xxV(X^0)λκA:= 2κI_d- _xV( X_0)λκ. Here g denotes a local gradient vector; the terminal-cost function remains fixed. For the quadratic approximation and A≻0A 0, the corresponding conditional distribution is Gaussian. The equalities below apply to the approximated cost: γ∗(X0|X^0) γ^*(X_0| X_0) ∝exp(g⊤δX0−12δX0⊤AδX0) (g δ X_0- 12δ X_0 A\,δ X_0 ) (71) ∝exp(−12(δX0−A−1g)⊤A(δX0−A−1g)) (- 12(δ X_0-A^-1g) A(δ X_0-A^-1g) ) (72) As a result, γ∗(X0|X^0) γ^*(X_0| X_0) =(X^0+A−1g,A−1) =N( X_0+A^-1g,\,A^-1) (73) with the full expression γ∗(X0|X^0)=(X^0+(2κI−∇xxVλκ)−1∇Vλκ,(2κI−∇xxVλκ)−1) γ^*(X_0| X_0)=N ( X_0+ ( 2κI- _xVλκ )^-1 ∇ Vλκ,\, ( 2κI- _xVλκ )^-1 ) (74) where V and its derivatives are evaluated at X^0 X_0. B.2.2 Simplification of the dual form Denote I:=∫e(V(X0)−λ‖X0−X^0‖2)/λκdX0I:= e (V(X_0)-λ\|X_0- X_0\|^2 )/λκdX_0. For the same quadratic approximation, completing the square in terms of g and A gives the following approximate integral and dual objective: I I =eV(X^0)λκ∫expg⊤δX0−12δX0⊤AδX0dX0 =e V( X_0)λκ \g δ X_0- 12δ X_0 A\,δ X_0 \dX_0 (75) =eV(X^0)λκ+12g⊤A−1g∫exp−12(δX0−A−1g)⊤A(δX0−A−1g)dX0 =e V( X_0)λκ+ 12g A^-1g \- 12(δ X_0-A^-1g) A(δ X_0-A^-1g) \dX_0 (76) Evaluating the Gaussian integral under the condition A≻0A 0 gives I=(2π)d/2det(A−1)1/2exp(V(X^0)λκ+12g⊤A−1g)I=(2π)^d/2 (A^-1)^1/2 ( V( X_0)λκ+ 12g A^-1g ) (77) Therefore, λκlogI=V(X^0)+λκ2g⊤A−1g+λκlog((2π)d/2det(A−1)1/2) λκ I=V( X_0)+ λκ2g A^-1g+λκ ((2π)^d/2 (A^-1)^1/2 ) (78) and the dual objective becomes =infλ>0λ(ε+κd2log2π)+X^0∼μ^[V(X^0)+λκ2g⊤A−1g+λκ2logdet(A−1)] = _λ>0λ ( + κ d2 2π )+E_ X_0 μ [V( X_0)+ λκ2g A^-1g+ λκ2 (A^-1) ] (79) Solving for λ in this form is complicated due to the dependence of A−1A^-1 on λ. B.2.3 First-order approximation In high-dimensional cases, computing the Hessian ∇xxV(X^0) _xV( X_0) is expensive. For the linearized cost, the Hessian term is omitted, giving A=2κIdA= 2κI_d and A−1=κ2IdA^-1= κ2I_d. The following expressions are exact for this linearized inner problem. The conditional distribution simplifies to γ∗(X0|X^0)=(X^0+12λ∇V(X^0),κ2I)γ^*(X_0| X_0)=N ( X_0+ 12λ∇ V( X_0)\,, κ2I ) (80) The λ optimization reduces to infλ>0λ(ε+κd2logπκ)+X^0∼μ^[V(X^0)+‖∇V(X^0)‖24λ] _λ>0λ ( + κ d2 πκ )+E_ X_0 μ [V( X_0)+ \|∇ V( X_0)\|^24λ ] (81) For ε+κd2log(πκ)>0 + κ d2 (πκ)>0 and μ^[‖∇V‖2]>0E_ μ[\|∇ V\|^2]>0, this yields the positive closed-form solution λ∗=(ε+κd2logπκ)−1X^0∼μ^[‖∇V(X^0)‖24]λ^*= ( + κ d2 πκ )^-1E_ X_0 μ [ \|∇ V( X_0)\|^24 ] (82) Appendix C Stochastic Optimal Control Theory Here we derive Equation 17 using stochastic optimal control theory. We refer the reader to Kappen (2005) for more details. The terminal cost is fixed, and the HJB derivation assumes a sufficiently smooth value function and admissible controls for the verification argument. To apply this section to DRSB, replace the terminal cost g below by ρgρ g. The identity holds for an optimal control and its optimal value, not for an arbitrary intermediate neural drift. Consider the stochastic optimal control (SOC) problem minu _u [∫0112‖ut(Xt)‖2t+g(X1)] [ _0^1 12\|u_t(X_t)\|^2dt+g(X_1) ] (83a) subject to dXt=ut(Xt)dt+σdWt dX_t=u_t(X_t)dt+σ dW_t (83b) X0∼μ, X_0 μ,\,\, (83c) Define the cost function as ct(Xt,ut):=[∫t112‖us(Xs)‖2s+g(X1)]c_t(X_t,u_t):=E [ _t^1 12\|u_s(X_s)\|^2ds+g(X_1) ] (84) The value function at time t is the optimal cost from time t to 11. Vt(x):=minu(s)t≤s≤1[∫t112‖us(Xs)‖2s+g(X1)|Xt=x]V_t(x):= _ subarraycu(s)\\ t≤ s≤ 1 subarrayE [ _t^1 12\|u_s(X_s)\|^2ds+g(X_1)\,\, |\,\,X_t=x ] (85) The cost function depends on the control utu_t, whereas the value function is obtained by minimizing over the control. We split the integral into the intervals from t to t+dt+dt and from t+dt+dt to 1. Substituting t=1t=1 gives V1(x)=g(x)V_1(x)=g(x). Vt(Xt)=minu(s)t≤s≤1[∫t+dt12‖us(Xs)‖2s+∫t+dt112‖us(Xs)‖2s+g(X1)|Xt]V_t(X_t)= _ subarraycu(s)\\ t≤ s≤ 1 subarrayE [ _t^t+dt 12\|u_s(X_s)\|^2ds+ _t+dt^1 12\|u_s(X_s)\|^2ds+g(X_1)\,\, |\,\,X_t ] (86) Using Bellman’s principle of optimality (Bellman, 1954), we have Vt(Xt) V_t(X_t) =minu(s)t≤s≤1[∫t+dt12∥us(Xs)∥2ds = _ subarraycu(s)\\ t≤ s≤ 1 subarrayE [ _t^t+dt 12\|u_s(X_s)\|^2ds +minu(s)t+dt≤s≤1[∫t+dt112∥us(Xs)∥2ds+g(X1)|Xt+dt]|Xt] + _ subarraycu(s)\\ t+dt≤ s≤ 1 subarrayE [ _t+dt^1 12\|u_s(X_s)\|^2ds+g(X_1)\, |\,X_t+dt ]\, |\,X_t ] (87) =minu(s)t≤s≤1[∫t+dt12‖us(Xs)‖2s+Vt+dt(Xt+dt)|Xt] = _ subarraycu(s)\\ t≤ s≤ 1 subarrayE [ _t^t+dt 12\|u_s(X_s)\|^2ds+V_t+dt(X_t+dt)\, |\,X_t ] (88) Applying a Taylor expansion to Vt+dt(Xt+dt)V_t+dt(X_t+dt) at (Xt,t)(X_t,t) gives [Vt+dt(Xt+dt)]≈[Vt(Xt)+∂Vt(Xt)∂tdt+(∇Vt)⊤dXt+12(dXt)⊤∇2V(dXt)]E[V_t+dt(X_t+dt)] [V_t(X_t)+ ∂ V_t(X_t)∂ tdt+(∇ V_t) dX_t+ 12(dX_t) ∇^2V(dX_t) ] (89) Applying stochastic calculus to Equation 83b gives [dXt] [dX_t] =ut(Xt)dt =u_t(X_t)dt (90) [(dXt)(dXt)⊤] [(dX_t)(dX_t) ] =σ2dt =σ^2Idt (91) Therefore, [Vt+dt(Xt+dt)]=[Vt(Xt)+∂Vt(Xt)∂tdt+(∇Vt)⊤ut(Xt)dt+σ22ΔVtdt]E[V_t+dt(X_t+dt)]=E [V_t(X_t)+ ∂ V_t(X_t)∂ tdt+(∇ V_t) u_t(X_t)dt+ σ^22 V_tdt ] (92) Taking the limit as dt→0dt→ 0 and substituting Equation 92 into Equation 88 gives −∂Vt(Xt)∂t=minu12‖ut(Xt)‖2+(∇Vt)⊤ut(Xt)+σ22ΔVt- ∂ V_t(X_t)∂ t= _u \ 12\|u_t(X_t)\|^2+(∇ V_t) u_t(X_t)+ σ^22 V_t \ (93) which is the Hamilton–Jacobi–Bellman (HJB) equation. By differentiating the right-hand side, we obtain the optimal control u∗u^*: ut∗(Xt)=−∇Vt(Xt)u_t^*(X_t)=-∇ V_t(X_t) (94) At time t=0t=0, we have Equation 17. Appendix D Solving the SOC with Adjoint Matching D.1 Adjoint Matching This section summarizes results from the original adjoint matching paper. We refer the reader to Domingo-Enrich et al. (2024), Sections 5.2 and E, for more details. The initial distribution and terminal cost are fixed throughout this section. We assume sufficient differentiability and integrability for pathwise differentiation with additive noise and interchange with expectation. Gradients are column vectors and ∇xu _xu denotes the Jacobian. Adjoint matching solves the optimal control problem minuθpu[∫0112‖utθ(Xtu)‖2t+g(X1u)] _u^θ\,\,E_p^u [ _0^1 12\|u_t^θ(X^u_t)\|^2dt+g(X_1^u) ] (95) s.t.dXtu=utθ(Xtu)dt+σtdWt .t.\,\,dX^u_t=u_t^θ(X^u_t)dt+ _tdW_t (96) X0∼μ \,\, X_0 μ (97) Here, we explicitly include the superscript u in the state notation XuX^u, unlike in previous sections. This notation is useful in the derivation of the adjoint matching objective. D.1.1 Derivation of the Adjoint System We introduce a Lagrange multiplier process a:Ω×[0,1]→ℝna: ×[0,1] ^n for the dynamics. The Lagrangian of the optimal control problem is ℒ(utθ,at,tu,t) (u_t^θ,a_t; X^u_t,t) =∫t112‖utθ(Xtu)‖2t+g(X1u) = _t^1 12\|u_t^θ(X^u_t)\|^2dt+g(X_1^u) −∫t1at⊤(dXtu−utθ(Xtu)dt−σtdWt) - _t^1a_t (dX^u_t-u_t^θ(X_t^u)dt- _tdW_t ) (98) =∫t112‖utθ(Xtu)‖2t+g(X1u)−∫t1at⊤dXtu = _t^1 12\|u_t^θ(X^u_t)\|^2dt+g(X_1^u)- _t^1a_t dX^u_t +∫t1at⊤utθ(Xtu)dt+∫t1σtat⊤dWt + _t^1a_t u_t^θ(X_t^u)dt+ _t^1 _ta_t dW_t (99) Applying integration by parts to the second integral gives ℒ(utθ,at,tu) (u_t^θ,a_t, X^u_t) =∫t112‖utθ(Xtu)‖2t+g(X1u)−a1⊤X1+at⊤Xt = _t^1 12\|u_t^θ(X^u_t)\|^2dt+g(X_1^u)-a_1 X_1+a_t X_t +∫t1(Xtu)⊤(datdt)dt+∫t1at⊤utθ(Xtu)dt+∫t1σtat⊤dWt + _t^1(X_t^u) ( da_tdt )dt+ _t^1a_t u_t^θ(X_t^u)dt+ _t^1 _ta_t dW_t (100) Taking the first variation with respect to the trajectory t X_t yields δtuℒ δ X^u_tL =(∇xg(X1u)−a1)⊤h1 =( _xg(X_1^u)-a_1) h_1 +∫t1(datdt+(∇xutθ(Xtu))⊤at+∇x(12∥utθ(Xtu)∥2))⊤htdt=0, + _t^1 ( da_tdt+ ( _xu_t^θ(X_t^u) ) a_t+ _x ( 12\|u_t^θ(X^u_t)\|^2 ) ) h_t\,dt=0, (101) which must hold for all hth_t. Thus, the coefficients of hth_t must be zero, and the adjoint system datdt da_tdt =−(∇xutθ(Xtu))⊤at−∇x(12‖utθ(Xtu)‖2) =- ( _xu_t^θ(X_t^u) ) a_t- _x ( 12\|u_t^θ(X_t^u)\|^2 ) (102a) a1 a_1 =∇xg(X1u) = _xg(X_1^u) (102b) is derived. Next, we add the dynamics constraint to the cost-to-go function (u,xt)J(u,x_t), using the adjoint state as its multiplier. (u,xt) (u,x_t) :=∫t112‖utθ(Xtu)‖2t+g(X1u) := _t^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1) (103) =∫t112‖utθ(Xtu)‖2t+g(X1u)−∫t1at⊤(dXtu−utθ(Xtu)t−σtdWt) = _t^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1)- _t^1a_t (dX^u_t-u_t^θ(X_t^u)dt- _tdW_t ) (104) =∫t112‖utθ(Xtu)‖2t+g(X1u)−a1⊤X1u+at⊤Xtu = _t^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1)-a_1 X^u_1+a_t X^u_t +∫t1(Xtu)⊤(datdt)dt+∫t1at⊤utθ(Xtu)dt+∫t1σtat⊤dWt + _t^1(X_t^u) ( da_tdt )dt+ _t^1a_t u_t^θ(X_t^u)dt+ _t^1 _ta_t dW_t (105) Taking the gradient with respect to the initial state Xt=xtX_t=x_t yields ∇xt(u,xt) _x_tJ(u,x_t) =∫t1(∇xtXsu)⊤(dasds+(∇xusθ(Xsu))⊤as+∇x(12‖usθ(Xsu)‖2))s = _t^1 ( _x_tX_s^u ) ( da_sds+ ( _xu_s^θ(X_s^u) ) a_s+ _x ( 12\|u_s^θ(X_s^u)\|^2 ) )ds +(∇xtX1u)⊤(∇xg(X1u)−a1)+at + ( _x_tX_1^u ) ( _xg(X_1^u)-a_1 )+a_t (106) =at(by Equation 102). =a_t (by @crefcreftype~refnumeq:adjoint-system ). (107) As a result, the solution of the adjoint system in Equation 102 is at(tu,uθ)=∇xt(∫t112‖usθ(Xsu)‖2s+g(X1u))a_t( X^u_t,u^θ)= _x_t ( _t^1 12\|u_s^θ(X^u_s)\|^2ds+g(X^u_1) ) (108) D.1.2 Derivation of the Basic Adjoint Matching Objective Now we derive the basic adjoint matching objective ℒbasic-AM(uθ,u¯)=12∫01‖utθ(Xtu¯)+at(u¯,u¯)‖2t,u¯∼pu,u¯=stopgrad(uθ)L_basic-AM(u^θ; X u)= 12 _0^1\|u_t^θ(X u_t)+a_t( X u, u)\|^2dt, X u p^u,\,\, u= stopgrad(u^θ) (109) The trajectories are generated by the neural network with fixed parameters (u¯ u), detached from gradient computation during simulation. This avoids retaining the full simulation computation graph. It does not remove the costs of trajectory generation and network updates. To derive this objective, we start from the optimal control objective ℒ(u)=∫0112‖utθ(Xtu)‖2t+g(X1u)L(u)= _0^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1) (110) We follow a similar procedure by adding the dynamics constraint with the adjoint state at(u¯,u¯)a_t( X u, u) as its multiplier. Here, at(u¯,u¯)a_t( X u, u) is independent of θ. ℒ(u) (u) =∫0112‖utθ(Xtu)‖2t+g(X1u)−∫01at⊤(dXtu−utθ(Xtu)t−σtdWt) = _0^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1)- _0^1a_t (dX^u_t-u_t^θ(X_t^u)dt- _tdW_t ) (111) =∫0112‖utθ(Xtu)‖2t+g(X1u)−a1⊤X1u+a0⊤X0u = _0^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1)-a_1 X^u_1+a_0 X^u_0 +∫01(Xtu)⊤(datdt)dt+∫01at⊤utθ(Xtu)dt+∫01σtat⊤dWt + _0^1(X_t^u) ( da_tdt )dt+ _0^1a_t u_t^θ(X_t^u)dt+ _0^1 _ta_t dW_t (112) Differentiating with respect to the neural network parameter θ gives dℒ(u)dθ dL(u)dθ =dθ(∫0112‖utθ(Xtu)‖2t+g(X1u)−a1⊤X1u+a0⊤X0uCLOSE = ddθ ( _0^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1)-a_1 X^u_1+a_0 X^u_0 +∫01(Xtu)⊤(datdt)dt+∫01at⊤utθ(Xtu)dt+∫01σtat⊤dWt) + _0^1(X_t^u) ( da_tdt )dt+ _0^1a_t u_t^θ(X_t^u)dt+ _0^1 _ta_t dW_t ) (113) =dθ(∫0112‖utθ(Xtu)‖2t+g(X1u)−a1⊤X1uCLOSE = ddθ ( _0^1 12\|u_t^θ(X^u_t)\|^2dt+g(X^u_1)-a_1 X^u_1 +∫01(Xtu)⊤(datdt)dt+∫01at⊤utθ(Xtu)dt) + _0^1(X_t^u) ( da_tdt )dt+ _0^1a_t u_t^θ(X_t^u)dt ) (114) Here, X0X_0 is fixed and therefore independent of θ. Differentiating the first term gives dθ(∫0112‖utθ(Xtu)‖2t) ddθ ( _0^1 12\|u_t^θ(X_t^u)\|^2dt ) =∫01(∂utθ∂θ(Xtu¯)+∇xutθ(Xtu¯)∂Xtu∂θ)⊤utθ(Xtu¯)t = _0^1 ( ∂ u_t^θ∂θ(X_t u)+ _xu_t^θ(X_t u) ∂ X_t^u∂θ ) u_t^θ(X_t u)dt (115) =∫01(∂utθ∂θ(Xtu¯))⊤utθ(Xtu¯)+(∂Xtu∂θ)⊤∇x(12‖utθ(Xtu¯)‖2)t. = _0^1 ( ∂ u_t^θ∂θ(X_t u) ) u_t^θ(X_t u)+ ( ∂ X_t^u∂θ ) _x ( 12\|u_t^θ(X_t u)\|^2 )dt. (116) The X1X_1 terms are dθ(g(X1u)−a1⊤X1u)=(∂X1u∂θ)⊤(∇xg(X1u¯)−a1)=0(by Equation 102b). ddθ (g(X_1^u)-a_1 X_1^u )= ( ∂ X_1^u∂θ ) ( _xg(X_1 u)-a_1 )=0 (by @crefcreftype~refnumeq:adjoint-terminal ). (117) The last two integrals are dθ(∫01(Xtu)⊤datdtt+∫01at⊤utθ(Xtu)t) ddθ ( _0^1(X_t^u) da_tdtdt+ _0^1a_t u_t^θ(X_t^u)dt ) =∫01(∂utθ∂θ(Xtu¯))⊤att+∫01(∂Xtu∂θ)⊤(datdt+(∇xutθ(Xtu¯))⊤at)t. = _0^1 ( ∂ u_t^θ∂θ(X_t u) ) a_tdt+ _0^1 ( ∂ X_t^u∂θ ) ( da_tdt+ ( _xu_t^θ(X_t u) ) a_t )dt. (118) Gathering all the terms, we have dℒ(u)dθ dL(u)dθ =∫01(∂utθ∂θ(Xtu¯))⊤(utθ(Xtu¯)+at)t = _0^1 ( ∂ u_t^θ∂θ(X_t u) ) (u_t^θ(X_t u)+a_t )dt +∫01(∂Xtu∂θ)⊤(datdt+(∇xutθ(Xtu¯))⊤at+∇x(12∥utθ(Xtu¯)∥2))dt + _0^1 ( ∂ X_t^u∂θ ) ( da_tdt+ ( _xu_t^θ(X_t u) ) a_t+ _x ( 12\|u_t^θ(X_t u)\|^2 ) )dt (119) =∫01(∂utθ∂θ(Xtu¯))⊤(utθ(Xtu¯)+at)dt(by Equation 102a). = _0^1 ( ∂ u_t^θ∂θ(X_t u) ) (u_t^θ(X_t u)+a_t )dt (by @crefcreftype~refnumeq:adjoint-ode ). (120) Throughout this derivation, the ∂Xtu∂θ ∂ X_t^u∂θ terms account for the dependence of trajectories on θ. These terms cancel by the adjoint equations. The remaining terms are evaluated using trajectories generated by u¯=stopgrad(u) u= stopgrad(u). We can rewrite Equation 120 as dℒ(u)dθ dL(u)dθ =∫01∂θ(12‖utθ(Xtu¯)‖2)+∂θ(at⊤utθ(Xtu¯))t = _0^1 ∂θ ( 12\|u^θ_t(X_t u)\|^2 )+ ∂θ (a_t u_t^θ(X_t u) )dt (121) =dθ(∫0112‖utθ(Xtu¯)‖2+at⊤utθ(Xtu¯)t) = ddθ ( _0^1 12\|u^θ_t(X_t u)\|^2+a_t u_t^θ(X_t u)dt ) (122) =dθ(∫0112∥utθ(Xtu¯)∥2+at⊤utθ(Xtu¯)+12∥at∥2dt)(∵at is independent of θ) = ddθ ( _0^1 12\|u^θ_t(X_t u)\|^2+a_t u_t^θ(X_t u)+ 12\|a_t\|^2dt ) ( a_t is independent of θ) (123) =dθ(∫0112‖utθ(Xtu¯)+at(u¯,u¯)‖2t) = ddθ ( _0^1 12\|u^θ_t(X_t u)+a_t( X u, u)\|^2dt ) (124) Thus, the basic AM objective in Equation 109 can be used to optimize utθu^θ_t using gradient descent. D.1.3 Connection to Optimal Control, the Lean Adjoint State, and the Adjoint Matching Objective The pathwise adjoint system applies to any sufficiently regular admissible control. We show the consistency of the functional stationarity condition with a verified SOC optimum. During the following variation, trajectories and adjoints are held fixed, and the perturbation is a state-feedback function ht=ht(Xt)h_t=h_t(X_t). The critical-point condition for adjoint matching follows from the first variation of pu[ℒbasic-AM]E_p^u[L_basic-AM]. δpu[ℒbasic-AM]δu _p^u[L_basic-AM]δ u =dϵ(pu[12∫01‖(utθ+ϵht)(Xtu¯)+at(u¯)‖2t])|ϵ=0 = . ddε (E_p^u [ 12 _0^1\|(u_t^θ+ε h_t)(X u_t)+a_t( X u)\|^2dt ] ) |_ε=0 (125) =pu[∫01ht⊤(utθ(Xtu¯)+at(u¯))t] =E_p^u [ _0^1h_t (u_t^θ(X u_t)+a_t( X u))dt ] (126) =pu[∫01ht⊤(utθ(Xtu¯)+∼pu[at(u¯)|Xtu¯])t] =E_p^u [ _0^1h_t (u_t^θ(X u_t)+E_ X p^u[a_t( X u)|X_t u])dt ] (127) The last equality follows from the tower property, X[X]=Y[X[X|Y]]E_X[X]=E_Y[E_X[X|Y]]. Thus, if u∗u^* satisfies the following condition, then u∗u^* is a critical point of the adjoint matching objective. ut∗(xt)=−∗∼pu∗[at(u∗)|Xtu∗=xt]u_t^*(x_t)=-E_ X^* p^u^*[a_t( X^u^*)|X_t^u^*=x_t] (128) For a verified SOC optimum u∗u^*, the conditional expected cost-to-go equals VtV_t. Using Equation 108 and interchanging differentiation and expectation gives ut∗(xt) u_t^*(x_t) =−pu∗[at(u∗)∣Xtu∗=xt] =-E_p^u^* [a_t( X^u^*) X_t^u^*=x_t ] (129) =−pu∗[∇xt(u∗,xt)∣Xtu∗=xt] =-E_p^u^* [ _x_tJ(u^*,x_t) X_t^u^*=x_t ] (130) =−∇xtpu∗[(u∗,xt)∣Xtu∗=xt] =- _x_tE_p^u^* [J(u^*,x_t) X_t^u^*=x_t ] (131) =−∇xtVt(xt). =- _x_tV_t(x_t). (132) Thus the SOC optimum satisfies the functional stationarity condition. This calculation does not establish global optimality of an arbitrary stationary point in neural-network parameters. We can further simplify the adjoint system to the lean adjoint system. Multiplying Equation 128 by (∇xut∗(xt))⊤ ( _xu_t^*(x_t) ) , we have (∇xut∗(xt))⊤ut∗(xt)=−pu∗[(∇xut∗(xt))⊤at(u∗)∣Xtu∗=xt] ( _xu_t^*(x_t) ) u_t^*(x_t)=-E_p^u^* [ ( _xu_t^*(x_t) ) a_t( X^u^*) X_t^u^*=x_t ] (133) ⇔ pu∗[−∇x(12‖ut∗(xt)‖2)−(∇xut∗(xt))⊤at(u∗)∣Xtu∗=xt]=0. _p^u^* [- _x ( 12\|u_t^*(x_t)\|^2 )- ( _xu_t^*(x_t) ) a_t( X^u^*) X_t^u^*=x_t ]=0. (134) The adjoint ODE has zero conditional mean drift at the SOC optimum. This motivates the lean adjoint below; it does not imply that each pathwise adjoint is constant. The functional critical-point result for lean Adjoint Matching is provided by Domingo-Enrich et al. (2024), Proposition 7, under its stated assumptions. da~tdt d a_tdt =0 =0 (135a) a~1 a_1 =∇xg(X1u) = _xg(X_1^u) (135b) where a~ a is called the lean adjoint state. The lean adjoint state is no longer equivalent to the gradient of the cost-to-go function. Since the base SDE has no drift and the state cost is zero, the lean adjoint state is constant in time, with a~t=∇xg(X1u) a_t= _xg(X_1^u). The resulting adjoint matching objective is ℒAM(uθ,u¯)=12∫01‖utθ(Xtu¯)+a~t(u¯)‖2t=12∫01‖utθ(Xtu¯)+∇xg(X1u¯)‖2twhereu¯∼pu,u¯=stopgrad(uθ) splitL_AM(u^θ; X u)&= 12 _0^1\|u_t^θ(X u_t)+ a_t( X u)\|^2dt= 12 _0^1\|u_t^θ(X u_t)+ _xg(X_1 u)\|^2dt\\ &where X u p^u, u= stopgrad(u^θ) split (136) Replacing the terminal cost by ρgρ g and taking expectation over trajectories gives Equation 15, up to the constant factor 1/21/2. The final adjoint matching algorithm is outlined in Algorithm 2. Algorithm 2 Adjoint Matching (Domingo-Enrich et al., 2024) Require: Parameterized drift uθu^θ, differentiable terminal cost g(x)g(x), nominal distribution μ μ, target distribution ν, replay buffer ℬB, number of epochs E. Initialize: uθu^θ 1: for epoch in 1,2,…,E1,2,…,E do 2: Sample X^0ii=1M∼μ^\ X^i_0\_i=1^M μ 3: Sample tii=1M∼[0,1]\t^i\_i=1^M [0,1]. 4: Simulate (5b) to get (X~0i,ti,Xti,X1i)i=1M\( X^i_0,t^i,X_t^i,X_1^i)\_i=1^M. 5: Update replay buffer ℬ←ℬ∪(ti,Xti,X1i)i=1MB ∪\(t^i,X^i_t,\,X_1^i)\_i=1^M 6: Take gradient steps ∇θℒAM _θL_AM to minimize ℒAM(θ)=(t,Xt,X1)∼ℬ[‖uθ(t,Xt)+∇g(X1)‖2]L_AM(θ)=E_(t,X_t,X_1) [\|u^θ(t,X_t)+∇ g(X_1)\|^2 ]. 7: end for Appendix E Full DRSB Algorithm The main DRSB training loop is given in Algorithm 3. Each iteration consists of three steps: worst-case distribution update, terminal cost network update, and control network update. DRSB-W and DRSB-S differ only in the worst-case subroutine. Algorithm 4 details the Wasserstein worst-case optimization, which iteratively updates each sample via gradient ascent (Equation 18) and adjusts λ via gradient descent (Equation 19). Algorithm 5 details the Sinkhorn worst-case optimization, which computes λ∗λ^* in closed form (Equation 25) and samples from the Gaussian worst-case distribution (Equation 24). Algorithm 3 Distributionally Robust Schrödinger Bridge Require: Parameterized drift uθu^θ, parameterized terminal cost gϕg^φ, nominal distribution μ μ, target distribution ν, replay buffer ℬB, number of epochs E. Initialize: uθu^θ with a pretrained SB model from μ μ to ν, gϕg^φ, λ 1: for epoch in 1,2,…,E1,2,…,E do 2: X~0ii=1M,λ←\ X^i_0\_i=1^M,λ← WassersteinWorstCase(OPENuθ,λ)u^θ,λ) (Algorithm 4) or X~0ii=1M←\ X^i_0\_i=1^M← SinkhornWorstCase(OPENuθ)u^θ) (Algorithm 5) 3: Sample tii=1M∼[0,1]\t^i\_i=1^M [0,1]. 4: Simulate (5b) to get (X~0i,ti,Xti,X1i)i=1M\( X^i_0,t^i,X_t^i,X_1^i)\_i=1^M. 5: Update replay buffer ℬ←ℬ∪(ti,Xti,X1i)i=1MB ∪\(t^i,X^i_t,\,X_1^i)\_i=1^M 6: Take gradient steps ∇ϕℒD _φL_D to minimize ℒD(ϕ)=X1∼ℬ[log(1+e−gϕ(X1))]+X1∼ν[log(1+egϕ(X1))]L_D(φ)=E_X_1 [ (1+e^-g^φ(X_1)) ]+E_X_1 ν [ (1+e^g^φ(X_1)) ] 7: Take gradient steps ∇θℒAM _θL_AM to minimize ℒAM(θ)=(t,Xt,X1)∼ℬ[∥uθ(t,Xt)+ρ∇gϕ(X1)∥2]L_AM(θ)=E_(t,X_t,X_1) [\|u^θ(t,X_t)+ρ∇ g^φ(X_1)\|^2 ]. 8: end for Algorithm 4 Wasserstein worst-case optimization Require: Minibatch size β, number of λ updates L, worst-case update threshold τ. Function WassersteinWorstCase(uθ,λu^θ,λ) 1: Prepare result set ℛ=R=\\ 2: for λ update iteration 1,2,…,L1,2,…,L do 3: Sample X^0jj=1β∼μ^\ X^j_0\_j=1^β μ 4: X~0jj=1β←X^0jj=1β\ X^j_0\_j=1^β←\ X^j_0\_j=1^β 5: repeat 6: X0jj=1β←X~0jj=1β\X^j_0\_j=1^β←\ X^j_0\_j=1^β 7: X~0j←X0j+η[−uθ(0,X0j)−2λ(X0j−X^0j)]j=1β \ X^j_0← X^j_0+η [-u^θ(0,X^j_0)-2λ(X^j_0- X_0^j) ] \_j=1^β 8: until 1β∑j=1β‖X~0j−X0j‖2<τ 1β _j=1^β\ \| X^j_0-X^j_0\|^2<τ. 9: Update λ←λ−ζ[ε−1β∑j=1β‖X~0j−X^0j‖2]λ←λ-ζ [ - 1β _j=1^β\ \| X^j_0- X^j_0\|^2 ] 10: Save ℛ←ℛ∪X~0jj=1βR ∪\ X^j_0\_j=1^β 11: end for Return: ℛ=X~0ii=1LβR=\ X^i_0\_i=1^Lβ, λ Algorithm 5 Sinkhorn worst-case optimization Require: Batch size B, number of worst-case samples per clean sample w. Function SinkhornWorstCase(uθu^θ) 1: Sample X^0jj=1B∼μ^\ X^j_0\_j=1^B μ 2: Compute λ←(ε+κd2logπκ)−11B∑j=1B‖uθ(0,X^0j)‖24λ← ( + κ d2 πκ )^-1 1B _j=1^B \|u^θ(0, X_0^j)\|^24 3: Sample X~0ll=1w∼(X^0j−u0θ(X^0j)2λ,κ2I)j=1B \\ X^l_0\_l=1^w ( X^j_0- u^θ_0( X^j_0)2λ, κ2I ) \_j=1^B Return: X~0ii=1wB\ X^i_0\_i=1^wB Appendix F Implementation Details All experiments are conducted on a single NVIDIA RTX 4090 GPU. The drift network uθu^θ and the terminal cost network gϕg^φ are both parameterized as multi-layer perceptrons. For the 2D transport tasks (Gaussian and GMM), the baseline SB model is pretrained with the DSBM algorithm (Shi et al., 2023), and the DRSB training is initialized from this pretrained model following the successive approximation scheme in Algorithm 1. For the FFHQ image translation task, the model operates in the 512-dimensional latent space of the ALAE autoencoder (Pidhorskyi et al., 2020). F.1 Practical Techniques for Adversarial Training The minimax structure of DRSB can make training unstable, particularly in high-dimensional settings. The following techniques stabilize the adversarial optimization. Gradual adversarial training. The adversary in DRSB is the worst-case optimization, and its penalty is controlled by the dual variable λ. Gradually decreasing λ from a large initial value increases the adversarial shift slowly. In DRSB-W, this occurs naturally through the λ gradient descent steps (Figure 5(a)). In DRSB-S, the closed-form λ∗λ^* can produce extreme worst-case shifts from the start. To mitigate this, we linearly increase the ambiguity set radius ε over the first few iterations (Figure 5(b)). (a) DRSB-W λ evolution. (b) DRSB-S ε schedule. Figure 5: Training dynamics. (a) Evolution of λ in DRSB-W. (b) ε -schedule for DRSB-S and corresponding λ evolution. Pretraining the terminal cost network gϕg^φ. Consistent with findings in robust control (Beard & McLain, 1998), proper initialization is critical for convergence in adversarial training. We pretrain gϕg^φ using the pretrained SB drift u¯ u by minimizing the joint objective ℒD(ϕ)+ξℒAM(ϕ)L_D(φ)+ _AM(φ), where ξ is a weight hyperparameter. Generated samples for ℒD(ϕ)L_D(φ) are obtained from u¯ u, and ℒAM(ϕ)L_AM(φ) regresses ρ∇gϕ(X1)ρ∇ g^φ(X_1) to the negative drift −u¯(t,Xt)- u(t,X_t). Regularization with a pretrained SB model. Worst-case samples can push the controller away from the correct data manifold, leading to semantic degradation. Since the pretrained SB control u¯ u captures the data manifold direction, we regularize uθu^θ to remain close to u¯ u via the loss ℒAM(θ)+α[‖uθ−u¯‖2]L_AM(θ)+ [\|u^θ- u\|^2], where α is a weight hyperparameter. F.2 Sensitivity to κ in DRSB-S The entropy regularization parameter κ in the Sinkhorn distance controls the support of the worst-case distribution. As κ→0κ→ 0, the Sinkhorn distance approaches the Wasserstein distance, and the ambiguity set preserves geometric interpretability. As κ increases, higher-entropy couplings γ are favored, and the worst-case distribution has wider support. Training stability is sensitive to κ. From Equation 24, larger κ produces a more diffuse worst-case distribution. If κ is too large, training becomes unstable due to extreme deviation from the nominal samples. In high-dimensional settings, κ interacts with the dimension d through the optimal dual variable in Equation 25. The term h(κ,ε):=ε+κd2log(πκ)h(κ, ):= + κ d2 (πκ) must satisfy h(κ,ε)>0h(κ, )>0 for λ∗λ^* to be well-defined. If h(κ,ε)h(κ, ) is too large, λ∗λ^* becomes small and the worst-case distribution deviates excessively from the nominal in Equation 24. Setting κ near 00 or 1/π1/π eliminates the log(πκ) (πκ) contribution and reduces the sensitivity to d. Table 4 reports FID scores on the FFHQ Woman→ task for two values of κ and several values of the ambiguity set radius ε . Both κ=0.005κ=0.005 and κ=0.32κ=0.32 yield stable training and robustness to input noise. In the FFHQ task with d=512d=512, values of κ far from 00 or 1/π1/π (e.g., κ=0.25κ=0.25) or ε>30 >30 cause training to diverge. Table 4: FID (↓ ) sensitivity to κ and ε on FFHQ Woman→ under noise δ. Model Org δ=0.25δ=0.25 0.500.50 1.001.00 DSBM 25.55± 0.50 26.59± 0.48 26.15± 0.49 26.81± 0.52 DRSB-S κ=0.005, ε =9 26.76± 0.49 26.28± 0.44 26.60± 0.52 27.15± 0.49 DRSB-S κ=0.005, ε =12 26.36± 0.48 26.48± 0.48 26.10± 0.50 26.66± 0.52 DRSB-S κ=0.005, ε =15 26.08± 0.47 26.57± 0.50 26.35± 0.53 26.40± 0.53 DRSB-S κ=0.32, ε =7 26.70± 0.48 25.93± 0.48 26.05± 0.47 26.05± 0.50 DRSB-S κ=0.32, ε =12 26.52± 0.50 25.82± 0.49 25.98± 0.49 25.94± 0.47 DRSB-S κ=0.32, ε =15 26.26± 0.50 25.88± 0.44 26.51± 0.52 25.89± 0.49 Appendix G Additional Experiments The main text reports selected perturbation conditions in Table 1 and Figure 3. This section provides the full results across all noise levels and shift directions. G.1 Noise Robustness Table 5 reports SWD under L2L_2-norm noise at eight levels (δ=0.5δ=0.5 to 4.04.0) for both tasks. DSBM performs best at low noise in the Gaussian task, while DRSB-W and DRSB-S outperform DSBM at higher noise levels. In the GMM task, DRSB-S outperforms both DSBM and DRSB-W at all noise levels except δ=0.5δ=0.5. Table 5: SWD (↓ ) on 2D transport tasks under varying L2L_2-norm noise levels δ (mean ± std over 10 seeds). Bold indicates best performance; gray indicates outperforming the baseline DSBM. Model δ=0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 (a) Gaussian-to-Gaussian DSBM 0.052± 0.006 0.072± 0.009 0.128± 0.012 0.214± 0.010 0.313± 0.015 0.431± 0.012 0.555± 0.013 0.684± 0.013 DRSB-W 0.080± 0.010 0.093± 0.012 0.138± 0.013 0.216± 0.013 0.315± 0.012 0.425± 0.013 0.548± 0.013 0.678± 0.013 DRSB-S 0.131± 0.011 0.134± 0.014 0.165± 0.015 0.227± 0.016 0.314± 0.015 0.414± 0.015 0.522± 0.015 0.638± 0.016 (b) GMM-to-GMM DSBM 1.164± 0.137 1.287± 0.112 1.597± 0.077 1.842± 0.086 2.051± 0.116 2.175± 0.118 2.254± 0.108 2.407± 0.243 DRSB-W 1.096± 0.216 1.220± 0.172 1.547± 0.114 1.821± 0.123 1.985± 0.111 2.139± 0.128 2.222± 0.166 2.315± 0.236 DRSB-S 1.191± 0.195 1.150± 0.206 1.293± 0.145 1.636± 0.157 1.836± 0.127 1.953± 0.095 2.013± 0.118 2.212± 0.413 G.2 Shift Robustness Table 6 reports SWD under directional distribution shifts and worst-case inputs. DRSB-S achieves the lowest SWD under most shift directions in both tasks. DRSB-W outperforms DSBM on the original GMM distribution, suggesting that the worst-case update can act as a regularizer on complex manifolds. Table 6: SWD (↓ ) on 2D transport tasks under distribution shifts (mean ± 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. Model Org Up Down Left Right W-wst S-wst (a) Gaussian-to-Gaussian DSBM 0.038± 0.007 0.665± 0.012 0.673± 0.020 0.666± 0.011 0.651± 0.009 1.081± 0.021 0.951± 0.013 DRSB-W 0.144± 0.007 0.691± 0.020 0.698± 0.020 0.517± 0.014 0.792± 0.014 0.972± 0.007 0.812± 0.012 DRSB-S 0.123± 0.012 0.647± 0.018 0.640± 0.014 0.501± 0.014 0.769± 0.016 0.881± 0.018 0.710± 0.013 (b) GMM-to-GMM DSBM 0.902± 0.242 6.167± 0.233 5.798± 0.245 5.917± 0.127 5.917± 0.132 2.086± 0.166 1.988± 0.100 DRSB-W 0.885± 0.226 5.949± 0.262 5.865± 0.282 5.967± 0.153 5.785± 0.120 1.503± 0.052 1.946± 0.100 DRSB-S 1.165± 0.138 5.887± 0.247 5.698± 0.203 5.840± 0.095 5.675± 0.185 1.791± 0.102 1.478± 0.100 Appendix H Worst-Case Cost Analysis H.1 Worst-Case Cost Interpretation Remark H.1 (Fixed-terminal-cost bound). Fix a controller u and terminal cost g. If μ~ μ globally maximizes X0∼μ[V(u,X0,g)]E_X_0 μ[V(u,X_0;g)] over ϵ(μ^) M_ε( μ), then ∀μ∈ϵ(μ^),X0∼μ[V(u,X0,g)]≤X0∼μ~[V(u,X0,g)].∀μ∈ M_ε( μ), _X_0 μ[V(u,X_0;g)] _X_0 μ[V(u,X_0;g)]. (137) This is a direct consequence of global inner optimality. It does not require an optimal controller or imply that the ambiguity constraint is active. Approximate adversarial updates need not attain this bound. Since g is fixed, the inequality is not a certificate for the full objective with terminal KL divergence. Applying this bound to a test distribution requires verifying its membership in the ambiguity set. For empirical training and test distributions, the Wasserstein distance is well defined and can be computed from their samples. With the squared Euclidean transport cost, the Wasserstein ambiguity set constrains the average squared displacement from the nominal distribution. H.2 Empirical Cost Comparison We compare the measured costs of the computed adversarial distribution, four shifted distributions (Up, Down, Left, Right), and the nominal distribution within the Wasserstein ambiguity set. For fixed uθu^θ and gϕg^φ, we estimate [∑t12‖ut‖2Δt+ρgϕ(X1)]E[ _t 12\|u_t\|^2 t+ρ g^φ(X_1)] over 1000 uniform time steps. In Figure 6, the computed adversarial distribution has the largest measured cost among these distributions. This finite comparison does not verify global inner optimality or certify the full objective with terminal KL divergence. The control cost in Figure 6 saturates as ρ increases, while the terminal-cost term receives greater weight. Figure 6: Cost comparison within the Wasserstein ambiguity set. Left: Evaluated distributions with the ambiguity-set boundary (dashed). Right: Measured cost versus ρ. Appendix I Additional Experimental Analysis I.1 Measured Computation Cost DRSB starts from a pretrained DSBM drift. Table 7 separates the two stages. The DRSB stage costs approximately 0.54×0.54× the Gaussian DSBM time and 0.16×0.16× the FFHQ DSBM time. Including pretraining gives approximately 1.54×1.54× and 1.16×1.16×, respectively. Gradient updates account for about 80% of Gaussian training time and over 90% on FFHQ, while worst-case optimization accounts for less than 1.5%. DRSB-W uses up to 90 worst-case gradient steps in these runs; DRSB-S uses Gaussian sampling with a closed-form λ update under the linear approximation. FFHQ DRSB draws four Monte Carlo samples per clean sample. These measurements describe the tested workloads and do not establish how runtime scales with dimension. Table 7: Measured training costs. DRSB columns report the additional stage after DSBM pretraining. Trajectory counts follow the reported training configurations. 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×1× 0.54×0.54× 1×1× 0.16×0.16× I.2 Empirical Convergence We examine the 2D Gaussian task with ρ=10ρ=10, ϵ=12ε=12, and 50 epochs, using a cosine learning-rate schedule from 3×10−73× 10^-7 to 10−710^-7. In Table 8, Rk=|1−1Nϵ∑i∥X0i∗−X^0i∥2|R_k= |1- 1Nε _i\|X_0^i*- X_0^i\|^2 | is the constraint residual, Gk=p1u(k),μ(k)[gϕ(X1)]G_k=E_p_1^u^(k),μ^(k)[g^φ(X_1)] is the mean terminal critic output, and UkU_k is the mean squared control magnitude on a fixed clean evaluation set. At the reported epochs k≥25k≥ 25, Rk<0.005R_k<0.005. The final dual variable is λ50=3.483 _50=3.483, close to the value 3.496 predicted by the active-constraint SOC relation. The terminal critic output and control magnitude change less at later epochs. These summaries measure stabilization, not gradients or stationarity residuals for all three blocks. Table 8: Reported update diagnostics on the 2D Gaussian task. k 5 10 15 20 25 30 35 40 45 50 RkR_k 0.260 0.319 0.240 0.111 0.002 0.004 0.002 0.000 0.002 0.001 λk _k 6.601 5.696 4.535 3.627 3.427 3.457 3.477 3.488 3.494 3.483 GkG_k 0.752 1.232 1.735 2.142 2.068 1.757 1.672 1.631 1.524 1.550 UkU_k 22.94 22.92 23.19 23.74 24.10 24.22 24.29 24.36 24.32 24.30 We also evaluate the variational objective ℒ(u,μ,g)=X0∼μ[V(u,X0,g)]−ρlogX1∼ν[eg(X1)],L(u,μ,g)=E_X_0 μ[V(u,X_0;g)]-ρ _X_1 ν[e^g(X_1)], (138) where V(u,x,g)V(u,x;g) includes the control effort and terminal cost ρgρ g. The control effort is estimated with a Riemann sum. The sampled values in Table 9 stabilize near 312 after an observed peak of 313.44 at k=32k=32. A separate run with constant learning rate 10−710^-7 for 120 iterations also settles near 312. This provides an additional check beyond the decreasing learning-rate schedule, without proving convergence of the block algorithm. Table 9: Variational objective at the reported epochs. k 5 10 20 30 40 50 ℒL 270.99 276.16 304.99 311.99 311.44 311.37 I.3 Empirical Check of the SOC Optimality Condition The simplified adversarial update replaces ∇X0V(u,X0,g) _X_0V(u,X_0;g) by −uθ(0,X0)-u^θ(0,X_0) using Equation 17. The identity is exact for an optimal controller of the fixed-terminal-cost SOC problem, but finite training need not satisfy it. We assess this approximation on the same 2D Gaussian experiment by computing ∇X0V _X_0V with a discrete adjoint recursion requiring one vector-Jacobian product per time step. The controller and terminal cost are held fixed for this gradient evaluation. Table 10 reports both direction cosines and gradient-to-control magnitude ratios. The cosine tests whether the surrogate points in the cost-gradient direction, while the ratio tests its scale. Across the evaluated epochs, cosines are 0.9996–0.9998 and ratios are 0.9310–1.0761. At epoch 50, the cosine is 0.9998 and the ratio is 1.0089, corresponding to a gradient magnitude approximately 0.9% larger than the control magnitude. These measurements support approximate satisfaction of the SOC relation on the evaluated samples. They do not establish exact SOC optimality or its validity across other tasks and input distributions. Table 10: SOC consistency diagnostics on the 2D Gaussian task. k 10 20 30 40 50 cos(∇X0V,−uθ(0,X0)) ( _X_0V,-u^θ(0,X_0) ) 0.9997 0.9996 0.9996 0.9997 0.9998 ‖∇X0V‖/‖uθ(0,X0)‖\| _X_0V\|/\|u^θ(0,X_0)\| 0.9310 1.0761 1.0249 1.0278 1.0089