This repository contains code for TMLR 2026 accepted paper titled Variational Online Mirror Descent for Robust Learning in Schrödinger Bridge.
This code is the official implementation of Variational Mirrored Schrödinger Bridge (VMSB), which is based on Variational Online Mirror Descent (VOMD) technique for the Wasserstein-Fisher-Rao Geometry.
Some of the code contains the exact reimplemntation of LightSB PyTorch official repository.
To run this code, every scripts for runnning experiments are placed in run/. Note that we did not ship the dataset for the code; but, we leave downloading snippets with --cfg.data_preload option. For example, you can
python vmsb_eot.py --cfg.data_preload
for donwloading each data and then run
python vmsb_eot.py
for remaining runs. The algorithms typically decised with the options --cfg.alg_name lsb, --cfg.alg_name lsbm, --cfg.alg_name md0, and --cfg.alg_name md0_lsbm, which means LightSB, LightSB-M, VMSB, and VMSB-M in the submission. Zeros in md0 indicatess we are only using x=0 for VOMD (B=1), and you can set --cfg.alg_name md where B=cfg.batch_size. Other default configuration settings are stored in the cfg/ directory.
The demo/ repository contains code for the toy 2D GMM experiment in the theoretical analysis (Section 4.2) with the complete plotting code.
@article{
han2025variational,
title={Variational Online Mirror Descent for Robust Learning in Schr\"odinger Bridge},
author={Dong-Sig Han and Jaein Kim and Hee Bin Yoo and Byoung-Tak Zhang},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2025},
url={https://openreview.net/forum?id=g3SsM9FLpm},
}