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

Dong-Sig Han, Jaein Kim, HEE BIN YOO, Byoung-Tak Zhang

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.

Submission Length: Long submission (more than 12 pages of main content)

Previous TMLR Submission Url: https://openreview.net/forum?id=xGsg2L3ppf

Changes Since Last Submission:

This is a resubmission of submission #4592.

- We have put significant effort into improving general readability and enhancing the accessibility of the background material for TMLR readers.
- The technical comments from the reviews have been addressed in this resubmission. The notations and derivations were enhanced to improve the presentation of our claims.
- The AE gave us solid and specific direction on strengthening the early framing—presenting the core problem, motivation, and contributions clearly in the introduction in a way that would be accessible to the broader machine learning community. Following this guidance, we focused the majority of our revision efforts on addressing these points, as well as particular concerns about the theory and experiments.

Code: https://github.com/dshan4585/vmsb

Assigned Action Editor: ~Chris_J_Maddison1

Submission Number: 5801