Variational Online Mirror Descent for Robust Learning in Schr""odinger Bridge

2025-04-03

AI Summary

Schrödinger bridge (SB) generative modeling suffers from insufficient robustness due to learning signal uncertainty. Method: We propose a variational online mirror descent (OMD) framework, yielding the first simulation-free variational Monte Carlo SB (VMSB) algorithm. By introducing online optimization into SB solving, we establish— for the first time—theoretical guarantees on convergence and regret bounds. Furthermore, we design a Gaussian mixture potential parameterization grounded in Wasserstein–Fisher–Rao (WFR) geometry, enabling stable, differentiable, and simulation-free SB computation. Contribution/Results: Our method unifies variational inference, Wasserstein gradient flows, and Fisher–Rao metrics. Across multiple benchmark tasks, it consistently outperforms existing SB solvers, demonstrating significantly improved robustness and stability—especially under suboptimal conditions.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Quantum Machine Learning
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

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.

Problem

• Addresses uncertainty in learning signals for Schru00f6dinger Bridge models
• Enhances stability of SB solvers via variational online mirror descent
• Proposes simulation-free VMSB algorithm for robust SB problem solutions

Innovation

• Variational online mirror descent framework
• Wasserstein-Fisher-Rao geometry utilization
• Simulation-free VMSB algorithm