Variational Online Mirror Descent for Robust Learning in Schrödinger Bridge

TL;DR

This work addresses the stability and reliability of Schrödinger Bridge learning under uncertain, online learning signals. It introduces Variational Mirrored Schrödinger Bridge (VMSB), a simulation-free SB solver built on Variational Online Mirror Descent (VOMD) and Wasserstein-Fisher-Rao geometry, with a Gaussian mixture parameterization enabling tractable updates. The authors establish convergence and regret bounds for the online MD framework and derive closed-form, gradient-flow-based updates that maintain stability across online data streams, transport benchmarks, and high-dimensional image translation tasks. Empirical results show that VMSB consistently outperforms existing simulation-free SB solvers, highlighting robustness and generality of the proposed approach.

Abstract

The Schrödinger bridge (SB) has evolved into a universal class of probabilistic generative models. In practice, however, estimated learning signals are innately uncertain, and the reliability promised by existing methods is often based on speculative optimal case scenarios. Recent studies regarding the Sinkhorn algorithm through mirror descent (MD) have gained attention, revealing geometric insights into solution acquisition of the SB problems. In this paper, we propose a variational online MD (OMD) framework for the SB problems, which provides further stability to SB solvers. We formally prove convergence and a regret bound for the novel OMD formulation of SB acquisition. As a result, we propose a simulation-free SB algorithm called Variational Mirrored Schrödinger Bridge (VMSB) by utilizing the Wasserstein-Fisher-Rao geometry of the Gaussian mixture parameterization for Schrödinger potentials. Based on the Wasserstein gradient flow theory, the algorithm offers tractable learning dynamics that precisely approximate each OMD step. In experiments, we validate the performance of the proposed VMSB algorithm across an extensive suite of benchmarks. VMSB consistently outperforms contemporary SB solvers on a wide range of SB problems, demonstrating the robustness as well as generality predicted by our OMD theory.

Figures (17)

• Figure 1: A schematic illustration. The primal and dual spaces $(\mathcal{C}, \mathcal{D})$ retain bidirectional maps $(\delta_{\mspace{-1mu}\mathcal{C}} \Omega, \delta_{\mspace{-1mu}\mathcal{D}} \Omega^\ast\mspace{-1.5mu})$. $\Pi^\perp_\nu$ and $\Pi^\perp_\mu$ indicate projection spaces of $\gamma_1\pi = \mu$ and $\gamma_2\pi = \nu$, respectively. The current $\pi_t$ performs an online learning update following a "unreliable" leader $\pi^\circ_{\raisebox{.9pt}{$t$}}$ in a region shaded in gray.
• Figure 2: A technical overview. Combining chracteristics of existing methods, VMSB offers a simulation-free SB solver that produces iterative OMD solutions. Our VMSB additionally provides a strong theoretical guarantee of convergence based on a regert analysis.
• Figure 3: The SB problem.
• Figure 4: Learning for an SB model $\{\pi_t\}_{t=1}^\infty$ in the primal space $\mathcal{C}$ (see \ref{['fig:schema']} for the details). Left: Sinkhorn (\ref{['lem:sink']}). Middle: Wasserstein gradient descent in the distributional space $\mathcal{C}$ for fixed $F$ (\ref{['lem:wd']}). Right: Variational online mirror descent with sequence of convex costs using uncertain estimates $\{\pi^\circ_{\raisebox{.9pt}{$t$}}\}_{t=1}^\infty$.
• Figure 5: Loss landscapes and gradient dynamics in a 2D problem. Left: In an early stage, parameters of three modalities $\{m_k\}_{k=1}^3$ (mean estimations) for both LightSB (top) and VMSB (bottom) methods approach the optimality with different costs. Right: When magnified the landscapes in the late stages (10 times), while LightSB is vibrant, whereas our method emits strictly convex landscape and stable dynamics.

Theorems & Definitions (51)

• Lemma 1: Duality; introeot
• Definition 1: Directional derivative
• Definition 2: First variation
• Definition 3: Generalized Bregman divergence
• Remark 1: LSI for the standard Gaussian
• Remark 2: conforti2024weak
• Lemma 2: Sinkhorn
• Lemma 3: MD in the Wasserstein space
• Theorem 1: Step size considerations
• Proposition 1: Convergence