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

TMLR Paper5801 Authors

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 sincerely appreciate all the suggestions from the reviewers. They pointed out valid concerns, with the most critical ones stemming from the clarity of the paper. In this resubmission, 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.
- After the rebuttal, there were no notable critical concerns other than the initial criticism of readability. The AE acknowledged our efforts in revising the manuscript but also cited the AE-reviewer discussion, noting that the reviewers were not satisfied at that time (there were no outstanding criticisms regarding the work's correctness). Therefore, we made major revisions to improve readability for the general TMLR audience by rewriting the introduction and adding five more citations. We also included a new appendix (Appendix C) with technical details on the dynamical Schrödinger bridge problem.
- 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.

We hope that this resubmission will be evaluated based on the correctness of our claims and results, according to TMLR's criteria.

---

Here are detailed response on the reviews and AE comments from the previous version.

> Placing particular emphasis on strengthening the early framing, presenting the core problem, motivation, and contributions clearly in the introduction.

Following the suggestions, we carefully revised the manuscript, particularly Section 1, to coherently introduce the core problem and motivation. We enhanced the overall flow by repositioning major segments and included more relevant references to emphasize the contextual importance of our findings. Additionally, we overhauled the detailed introduction of the technical background and theoretical statements in Sections 2 & 3, which, despite being partially improved in the initial rebuttal, was not considered sufficient for acceptance.

> Experimental details on singular cell experiments and report with 10 trials.

We had already addressed this in the initial submission. The single-cell experiments follow the setup and data from Tong et al. (2023), and we report the average performance with a 95% confidence interval over 10 different trials, as is clearly stated in the manuscript. Although our approach outperforms the benchmark in this particular setup, the reviewers may not have felt certain about the statistical significance. However, we argue that the intention of our experiments was to demonstrate that our VMSB method can solve various SB problems with higher stability, which verifies the generality of our claim. Thus, the merit of our approach should be evaluated based on the breadth and consistency of VMSB. Furthermore, because the data consists of PCA projections of single cells from a donor, we note that the single-cell experiments should not be weighed significantly more than other synthetic datasets, as this is not a genuinely natural form of data. To address this, we have clarified the experimental goals and settings in the revision and added extra discussion on the conclusion section.

> Theoretical details on Theorem 1 and unclear statement of Lemma 14.

We corrected the ambiguity in Lemma 14 through careful revision and the addition of extra derivation steps.

Assigned Action Editor: ~Chris_J_Maddison1

Submission Number: 5801