Variational Online Mirror Descent for Robust Learning in Schrödinger Bridge
Authors: Dong-Sig Han, Jaein Kim, HEE BIN YOO, Byoung-Tak Zhang
TMLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | 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. ... 6 Experimental Results Experiment goals. We aimed to test our online learning hypothesis and verify that the VMSB effectively induces OMD updates. |
| Researcher Affiliation | Academia | Dong-Sig Han,1 Jaein Kim,2 Hee Bin Yoo,2 and Byoung-Tak Zhang2 1Department of Computing, Imperial College London 2Artificial Intelligence Institute, Seoul National University |
| Pseudocode | Yes | Algorithm 1 Variational Mirrored SB (VMSB). Input: SB models ( πθ, πφ) parameterized by Gaussian mixtures, step sizes (η1, ηT), ny, B N. 1: for t 1 to T do 2: Acquire φt with an external data-driven SB solver. 3: θt θ, ηt 1/ η91 1 + (η91 T η91 1 )(t 1/T 1) 4: for n 1 to N do 5: {xi}B i=1 sample mini batch data from µ. B PB i=1ηt WFRgrad(θ; φt, xi, ny) + (1 ηt)WFRgrad(θ; θt, xi, ny) 7: Update θ with the gradient L θ . 8: end for 9: end for Output: Trained SB model πθ. |
| Open Source Code | Yes | Reproducibility statement. Comprehensive justification and theoretical background are presented in Appendices A and B. Since the primary contributions of this paper pertain to the learning methodology, we ensured that all architectures and hyperparameters remained consistent across the Light SB variants. All datasets utilized in this study are available for download alongside the training scripts. Please refer to Appendix D for more information on the experimental setups. |
| Open Datasets | Yes | Section 6.2 Quantitative Evaluation EOT benchmark. Next, we considered the EOT benchmark proposed by Gushchin et al. (2024b), which contains 12 entropic OT problems with different volatility and dimensionality settings. ... SB on single cell dynamics. We evaluated VMSB on unpaired single-cell data problems in the high-dimensional single cell dynamics experiment (Tong et al., 2024a). ... MNIST-EMNIST. We applied VMSB to unpaired image translation tasks for MNIST and EMNIST datasets. ... FFHQ. Following the latent SB setting of Korotin et al. (2024), we assessed our method by utilizing a pretrained ALAE model for generating 1024 1024 images of the FFHQ dataset (Karras et al., 2019). |
| Dataset Splits | Yes | We applied an angle-based rotating filter, making the marginal as a data stream where only 12.5% (or 45-degree angle) of the total data is accessible for each step t. ... The first setting spanned from Day 2 to Day 4 (evaluated at Day 3), while the second setting considers duration from Day 3 to Day 7 (evaluated at Day 4). |
| Hardware Specification | No | In Table 7, we report the wall-clock time for a 100-dimensional single-cell data problem Vargas et al. (2021); Korotin et al. (2024), where the performance is reported in Table 3. Additionally, training time in the MNIST-EMNIST translation is reported in Table 11 in the ablation study. This property also holds for generation, allowing practitioners to deploy the model much faster on GPUs. In Table 8, we also report that generating 100 MNIST samples from 4096 Gaussian particles, equipped with competitive performance, can be done 1,854 times faster under the same hardware. |
| Software Dependencies | No | For fast computation, we utilized the JAX automatic differentiation library (Bradbury et al., 2018) for computing gradients and Hessians in Proposition 3. |
| Experiment Setup | Yes | D.2 Hyperparameters. The hyperparameters are displayed in Table 5. For step size scheduling, we followed the theoretical result in Theorem 1 and Proposition 1, and chose η1 = 1 and ηT {0.05, 0.1} with harmonic sequences, as illustrated in Fig. 5. For high dimensional tasks in MSCI (1000d), MNIST-EMNIST (784d), and latent FFHQ Image-to-Image transfer tasks (512d), the initial warm up steps for 10% of the total learning helped starting a training sequence from a reasonable starting point as this set ηt = 1 for a certain period of the early stage. Table 5: Hyperparameters. ... Modality K {8, 20, 50} [5, 100] 50 {256, 1024, 4096} {256, 1024} 10 Volatility ε 0.1 {0.1, 1, 10} 0.1 10 4 10 3 {0.1, 0.5, 1.0, 10.0} Total steps (τ) 20,000 30,000 10,000 100,000 30,000 20,000 OMD steps (t) 400 600 200 1000 375 400 |