Variational Online Mirror Descent for Robust Learning in Schrödinger Bridge
Journal:
arXiv
Published Date:
Apr 3, 2025
Abstract
Sch\"odinger bridge (SB) has evolved into a universal class of probabilistic
generative models. In practice, however, estimated learning signals are often
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\"odinger Bridge (VMSB) by utilizing the Wasserstein-Fisher-Rao geometry of
the Gaussian mixture parameterization for Schr\"odinger 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 range of
SB problems, demonstrating the robustness predicted by our theory.