Variational Mirror Descent for Robust Learning in Schrödinger Bridge

We appreciate your helpful feedback and concerns. We address your comments below.

What's newly proposed in § 4?

Thank you for the question. Karimi et al. [1] introduced "a continuous-time analogue of the Sinkhorn algorithm" by leveraging the perspective of considering Sinkhorn as MD [2, 3]. They proposed γ-IPF formulation of MD, which particularly involves one of the marginals in SB (Definition 1 of [1]). Overall, this can be understood as performing an algorithm akin to Lemma 1, but taking the limit $\eta\to0$.

In contrast to [4], our work takes a distinct learning-theoretical approach. In § 4, we constructed an online learning objective (Lemma 2) in a Wasserstein sense with a formal proof, allowing us to analyze convergence properties (Theorem 1 & Proposition 1) and a regret bound (Theorem 2). We would like to stress that conducting a precise analysis is important in learning theory. Moreover, we found that the variational principle can yield the OMD iterates in Proposition 2. To the best of our knowledge, the particular fact that the Wasserstein gradient flow of OMD in SB can be decomposed into a linear combination of KL functions is a novel result in the context of SB.

We would like to emphasize that § 4 focuses on the essential mechanisms of SB learning and shows that many achievements in convex analysis can be translated into SB learning. Furthermore, § 4 establishes the fundamentals for VMSB, including the online cost function and learning via Wasserstein gradient flow. Consequently, § 4 is of critical importance in our paper and presents various novel ideas on the learning of SB models.

To me, the main contribution lies in § 5 ... If the proposed method still requires NNs, it's unclear why we should consider mixture of Gaussian models.

Thank you for your constructive feedback. As you mentioned, § 5 also makes significant contributions, particularly in implementing the Wasserstein gradient flow as described in Proposition 2. These contributions lie in enabling computation for performing OMD, thanks to the LightSB parameterization by Korotin et al. [5]. Our paper does not present strong arguments on the advantages or disadvantages of parametrization in the main text (we delineate these points in Appendix D), because the pros and cons of parametrization are extensively discussed in [5]. Our main arguments were to show consistent performance improvement from VMSB under the identical parameterization setting.

In some of the experiments related to high-dimensional images, our algorithmic settings rely on discriminator and latent diffusion models for effective training. We would like to point out that the scalability limitations of GMMs in our cases are mostly problem-specific. Certain NN operations (e.g., convolution, U-Net) specialized for specific problems can help improve the SB solver’s performance. Once fully trained through interaction with NNs, we can efficiently sample from GMM-based SB models with competitive performance in generating MNIST data, achieving sampling speeds 1,854 times faster under the same hardware conditions.

One of the distinctive advantages of the non-NN SB parameterization is that we can obtain SB samples without simulation inside an SDE, which reduces the time of sampling and avoids discretization errors from SDE solving. Therefore, we would like to emphasize that non-NN SB solvers can solve complex SB problems under certain conditions, and we believe GMM-based models still hold an irreplaceable role in theoretical and applied subfields of machine learning due to their simplicity and theoretical properties.

Dear Reviewer hbEJ,

We would like to sincerely thank you for your thoughtful feedback and earnest reassessment of our work. We are pleased to hear that our previous rebuttal has addressed your initial concerns regarding clarity and theory. We address your concern about the use of a discriminative network in training as follows.

Overall, the experiments are designed to validate the generality of our claims across various OT-related ML problems. In particular, we observed that the current limitations of LightSB in image translation tasks at a certain level primarily can be learning-related issues rather than parameterization issues. Our VMSB demonstrated that the claimed theoretical properties of variational OMD can mitigate the mode collapsing issue considerably, which has also been considered one of the prominent learning-related issues in this field. We believe that proposing a practical adversarial algorithm that is “tolerant of unreliable objective estimation” is fundamentally connected to our central claim on learning theory, and is within our experimental scope.

Thank you once again for your dedication and commitment. Please let us know if you have any further questions during the discussion phase.

Best regards,

Submission #5991 authors

We sincerely thank you for your overall support of our submission. We address your question below.

Question. ... This is due to the convexity of SB cost used for VOMD, compared to the non-convex nature of the other methods. I wonder if it's possible to show directly that the proposed method is better than other methodsunder some situations?

As you mentioned, the distinct properties of the SB cost have significantly advanced our understanding of the MD-Sinkhorn relationship [1-3]. The properties are also highly important for our theoretical analysis of OMD improvement. In this paper, we claim that VOMD can be efficiently performed since infinitesimal variations of the Bregman divergence fundamentally resemble those of the KL divergences.

While convex objectives are common in ML, non-convex learning is also prevalent in probabilistic generative models. In non-convex settings, general convergence guarantees are typically unavailable, and deriving regret bounds becomes considerably more challenging. Alternative notions, such as dynamic regret [4] and local regret [5], have been proposed to address these issues. However, adequate assumptions for regret analysis in the context of probabilistic generative models have not yet been rigorously established.

To effectively show that our proposed method outperforms others under certain situations, we consider Assumption 1 in our paper. This main assumption allows our algorithm to converge to an asymptotic dual mean even under nonstationary objectives. Although this approach involves an indirect comparison based on a theoretical postulate, the practicality and generality of this assumption are supported by the performance improvements observed in our experiments. Through a wide range of experiments, we demonstrated that the method can outperform existing methods and that the proposed assumptions hold for various situations.

Thank you for your question. Please feel free to reach out if you have any further questions.

Sincerely,

Submission #5991 authors

[1] Flavien Léger. A gradient descent perspective on Sinkhorn. Applied Mathematics & Optimization, 2021

[2] Pierre-Cyril Aubin-Frankowski, Anna Korba, and Flavien Léger. Mirror descent with relative smoothness in measure spaces, with application to Sinkhorn and EM. Advances in NeurIPS, 2022.

[3] Mohammad Reza Karimi, Ya-Ping Hsieh, and Andreas Krause. Sinkhorn flow as mirror flow: a continuous-time framework for generalizing the sinkhorn algorithm. In AISTATS, 2024.

[4] Omar Besbes, Yonatan Gur, and Assaf Zeevi. Non-stationary Stochastic Optimization. Operations research, 2015

[5] Elad Hazan, Karan Singh, and Cyril Zhang. Efficient Regret Minimization in Non-Convex Games. In ICML, 2017

Adversarial training

Adversarial training has provided a powerful baseline in the MNIST-EMNIST translation tasks, which other LightSB algorithms cannot match. However, like GANs, the mode collapse phenomenon still presents a practical issue. From our theory, we considered that one cause of this problem is the unstable learning signals due to limited discriminative capability. In the adversarial training experiment, the proposed VMSB variant “VMSB-adv” showed a clear advantage since constrained OMD updates govern the progress of the adversarial training. This point is strongly related to our theoretical claims of convergence and is predicted by the assumptions.

Q1

Thank you for the question. We respectfully highlight that we have substantially reduced the amount of required computation for OMD by using the LightSB parameterization and variational inference, and we have verified its usefulness through experimental results. Since the algorithm can be performed more efficiently than expected, the proposed VMSB obtained strong numerical results, especially in low- to mid-range generation problems with relatively low computational time.

For high-dimensional problems, the computation of WGF may become more parallel, and we essentially favor a GPU for training. Once it is fully trained, we can deploy and utilize the model on CPUs and GPUs, which can be applied efficiently to various applications, considering its complexity. Here, we also report that collecting SB samples from 4096 Gaussian particles, which are equipped with competitive performance, can be done 1,854 times faster under the same hardware, and we can also generate samples with a CPU.

Note that this generation process does not suffer from discretization errors of SDE. Although the LightSB parameterization might suffer from very high-dimensional problems, we believe GMM-based models still hold an irreplaceable role in practice due to their simplicity and theoretical properties. As our primary intention was to validate our theoretical claims on learning, we hope that our experiments could help the training of non-NN models that has potential to be used in theoretical and applied ML problems

Q2

Following your suggestion, we reported experimental results with EMA above. We would like to summarize our method from both perspectives as below.

• Optimization for “high-level” SB probability distributions.
  • Current SB models can be seen as solving a convex optimization problem. We claim that this does not fully guarantee stability in SB learning.
  • We propose VMSB, which performs online MD updates utilizing the variational approach for computationally efficient and theoretically explainable learning processes.
• Optimization for “low-level” parameters with EMA.
  • We do not specifically argue for or against training using optimizers and other parametric techniques.
  • VMSB. VMSB optimizes SB along with the WFR gradients using an optimizer such as Adam, but the proximal local target of learning is the MD updates at the high-level probabilistic perspective, which makes the policy learning more robust.
  • EMA. Even if the SB updates are smoothed, EMA might have practical issues in high-level distributional spaces, and this approach can be unreliable depending on the problems, as there are no theoretical guarantees in such spaces.

Q3

As mentioned earlier, the mode collapsing phenomenon was one of the considerations while contemplating Assumption 1. Our theory finds the best SB model $\pi_\theta$ that is close to $\delta_D \Omega^\ast(\lim_{t\to\infty}\mathbb{E}_t[\delta_C \Omega(\pi{\phi_t})])$, where the first variations of Bregman divergence are used for mappings between primal and dual spaces. In our theoretical framework, we claim that the OMD framework yields favorable outcomes by covering all modalities through the statistical divergence induced by $\Omega$.

Please let us know if there are any remaining questions or comments.

Sincerely,

Submission #5991 authors

[1] Maria L Rizzo and Gábor J. Székely. Energy distance. Wiley interdisciplinary reviews: Computational statistics, 2016

[2] Marc Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. In NeurIPS, 2022.

[3] Mohammad Reza Karimi, Ya-Ping Hsieh, and Andreas Krause. Sinkhorn flow as mirror flow: a continuous-time framework for generalizing the sinkhorn algorithm. In AISTATS, 2024.

[4] Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light Schrödinger bridge. In ICLR, 2024.

[5] Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, Alexander Korotin. Light and Optimal Schrödinger Bridge Matching, In ICML.

We appreciate your constructive feedback and the time spent reviewing our work. Please check the revised manuscript reflecting your important suggestions as follows.

• (Notation & Clarity) We agree that the enumeration style of $a_t$ is too brief. Following your suggestion, we enhanced the clarity of expressions in L209. Following fruitful suggestions, we also included additional theoretical details for clarity.
• (Maximum Mean Discrepancy) We agree that the energy distance is a more appropriate term for our evaluation metric. We have unified the terminology to energy distance [1] and ED.
• (Typos) Thank you for your detailed comments. We corrected spelling and typographical errors in the revised manuscript.

We address your comments below. Please refresh the webpage if the latex expressions are not properly rendered.

EMA baselines.

Thank you for the insightful comment. We respectfully emphasize that MD works in more meaningful probabilistic spaces.

• The SB problem is convex for distributions. We assumed that the expression $\lim_{t\to\infty} \mathbb{E}_t[\delta_C\Omega(\pi^\circ_t)]$ is stationary as a mapping of Gaussian mixtures (by consistently changing its mode) using Bregman potentials.
• The problem is non-convex to its parameters. The exponential moving averaging (EMA) is the technique favors the Euclidean average of historical parameters $\{\phi_t\}^\infty_t$. It stabilizes the learning process, but analyzing its meaning in the probabilistic space could be difficult.
• Our OMD algorithm finds the best model that is closest to $\delta_D(\lim_{t\to\infty}\mathbb{E}_t[\delta_C\Omega(\pi^\circ_t)]$) where the Bregman potential is used to measure local distance. Therefore, we claimed that OMD brings theoretically pleasing outcomes regarding the statistical divergence induced by KL.

To further validate our point, we added a row of LightSB-EMA in Tables 6 & 7. The table below presents the ε=1 results of cBW-UVP. This shows that VSMB consistently achieves better results than LightSB-EMA, and our theoretical argument rigorously supports these gains.

The theory is a combination of prior works

We agree that our contributions are built upon the prior works on SB. We respectfully point out that our main contribution lies in the learning theory, and the significance our theoretical contribution can be to be justified by a fine level of convergence analysis. In Theorems 1 & 2, and also Proposition 1, we have provided novel convergence proofs of OMD in SB problems. We also established the fundamental connection between the variational principle and OMD update rules in the Wasserstein gradient flow, enabling our technical contribution to utilize the important theoretical framework from Lambert et al. [2].

• Inspired by Karimi et al. [3], we took a distinct theoretical learning approach. We constructed an online learning objective $F_t(\pi)$ and answered some important questions related to SB learning, such as conditions for ideal and general convergence (Theorem 1 & Proposition 1) and regret bounds (Theorem 2). Although the suggestions of Karimi et al. [3] for various possibilities of adaptations for large-scale problems in their appendix are promising, our work implemented the VMSB algorithm with real experiments, validating our theoretical analysis.
• Korotin et al. [4] We would like to stress that we found a novel theoretical value of the LightSB parameterization by proposing its application in variational algorithms. Our analysis not only strengthened the mathematical rigor of the SB problem but also predicted a practical learning algorithm, which complements the contributions of [4].
• Lambert et al. [2]. We formally stated in Proposition 2 that the OMD iterates can be attained by the variational principle. We discovered that the Wasserstein gradient flow of OMD can be decomposed into a linear combination of KL functions in SB problems. As a minor contribution, we applied a correction to the ODE system equation in § 6 of [2], allowing mixture weights ($\alpha_{k,\tau}$; $r^{(i)}_t$ in the notation of [2]) to correctly change their weights according to the geometry.

Practical gain is not convincing.

We respectfully remind you that our experiments were designed in a highly controlled setting with identical model configurations. Our primary goal was to verify our theory and method, which we achieved by surpassing LightSB [4] and LightSB-M [5] under identical parameterization conditions and numbers of parameters. With all other aspects controlled and outcomes determined solely by learning methods, the consistent performance gains of our work were an well-anticipated result from our theoretical analysis.

Dear Reviewer 7fDv,

We appreciate for your thoughtful feedback and the positive remark on our response to Question 2. We further address Questions 1 & 3 below. As requested, we provide an anonymous link containing a single-page PDF file that addresses one of your comments: [anonymous link].

Question 1

In addition to the exceptionally fast generation time of simulation-free samples, we respectfully emphasize the importance of VMSB-adv experiments due to the various practical implications in the image domain.

• Table 9 in the appendix presents the FID scores on both the training and testing datasets. The performance difference demonstrate that VMSB-adv suffer less from overfitting compared to its counterpart LightSB-adv, across all parameterization settings. This supports our claims regarding the stability of SB solution acquisition in pixel spaces.
• We acknowledge the limitation on GMM architecture as noted in L1664-L1672. Due to the architectural limitation, training GMM-based SB models tends to focus on instance-level associations within high-dimensional problems, rather than the sub-instance or feature-level associations intrinsic to deep neural network architectures. Consequently, the generated letters may appear synthetic, yet the proposed approach produces statistically sound pixel representations within the given constraints. We hope that the proposed OMD theory will inspire future studies across various domains, including novel neural architecture research.
• We have included a demonstration of generated samples with their neighboring images from the training set based on MSD via the anonymous [link]. The figure shows that VMSB-adv demonstrates robust EOT images which are coherent to input and top-5 similar images. We would like to stress that the proposed method does not favor strict memorization of exact pixel positions and values of the training set, since such deterministic behavior does not meaningfully reduce adversarial loss throughout the time.
• Additionally, we present quantitative results comparing MSD metrics with 500 samples, respectively.

• We identified two key trends: (1) Samples generated by VMSB-adv consistently produce more representative images within local groups compared to LightSB-adv, as observed by improved Top-1 and Top-5 similarity metrics, and (2) VMSB-adv does not show signs of overfitting, and maintains a certain level of Top-1 similarity (VMSB-1024 -> VMSB-4096). Also, a substantial decrease in Top-5 similarity is observed when scaling from VMSB-1024 to VMSB-4096. These trends suggest that VMSB-adv enhances the stability of the learning process by aligning the modalities of distribution within local groups of multiple images, and the discriminator plays a central role in determining the convergence point. Consequently, our approach is less likely to reproduce identical instances from the data source, addressing the concern you raised.

• Learning vs. overfitting. Similar to other ML models, overfitting might happens when K is very high due to (1) the discrete nature of the dataset in the pixel domain and (2) the models’ sufficient capacity to represent all the data, potentially leading to overly complex model structures. Since the learning of SB models is closely related to the instance-level discovery of optimal couplings, some degree of memorization occurs when data pairs have clear associations. Informally speaking, full memorization would require setting $K^\ast =(\textrm{Number of images in} \ \mu) \times (\textrm{Number of images in} \ \nu)$ to ensure allocation of such associations. Therefore, we strongly expect generalization in SB learning for the GMM-based models, as $K\in$ {256,1024,4096} is relatively small compared to $K^\ast$.

We greatly appreciate your insightful feedback and thank you for the encouraging comments. We address your concerns below.

Heavy theoretical descriptions

Thank you for your constructive feedback. In response, we have enhanced the clarity of the theoretical descriptions in the manuscript by providing additional information and detailed explanations. While our approach is theoretical, we believe that a rigorous method offers a unique contribution toward better understanding for the target problem and fundamentally closing the gap between theory and practice. Additionally, we respectfully highlight that we have made considerable efforts to provide helpful resources and intuitive figures, and we have conducted a wide range of experiments in the SB domain to strengthen our claims.

WFR gradient

We appreciate your insightful comment. We would like to emphasize our contributions to the computational aspects of OMD using first variations. We intentionally designed Algorithm 1 to be a simple first-order method in parameterized SB models in order to verify our theoretical claims through experiments in a straightforward manner. Nevertheless, we have demonstrated that our algorithm can model complex distributions and solve various online learning problems. As you suggested, one can extend the efficiency of OMD with higher-order methods. However, technical limitations might arise since we did not assume high-order moments to be finite in $\mathcal{C}$ throughout the analysis. Therefore, a different theoretical approach would be required for a higher-order extension of OMD.

Q1 [efficiency]

Driven by parallel nature of Gaussian particles, we observed that the computation of Proposition 2 favors vectorized instructions, and the expected speed enhancement from using GPUs is much more evident in the NN case. In the following table, we report the wall-clock time for a 100-dimensional single-cell data problem [1,2,3], where the performance is reported in Table 3.

Additionally, training time in the MNIST-EMNIST translation is reported in Table 8 in the appendix**.** We believe this property also holds for sample generation, allowing us to deploy the model much faster on GPUs. Here, 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.

Note that the GMM generation process does not suffer from discretization errors of SDE. As we stated in Appendix D, we believe GMM-based models still hold an irreplaceable role in practice due to their simplicity and parallel nature.

Q2 [model details]

We used the LightSB parameterization [1] with diagonal GMMs. In response, we provide each modality $K$ for benchmarks in the following table.

The rest of the step size scheduling and trading details are stated in Appendix C and can be found in setting files of the provided code.

Please let us know if there are any further questions or comments.

Sincerely,

Submission #5991 authors

[1] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schr¨ odinger bridges via maximum likelihood. Entropy, 23(9):1134, 2021.

[2] Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light Schrödinger bridge. In ICLR, 2024.

[3] Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, Alexander Korotin. Light and Optimal Schrödinger Bridge Matching, In ICML.

Q7 [§ 4.2]

The high-level idea of Assumption 1 is from the observation that a computational SB algorithm should be implemented through approximation. We devised a mild assumption in the sense that the dual points could be stationary only with their mean when $t\to\infty$. It is apparent that if the algorithm is optimal and finds probabilistic convergence to π*, VMSB also converges according to Theorem 1.

In Fig. 4, we viewed the reference estimation fitted using LightSB as $\pi^\circ_t$. This setting can be changed to LightSB-M and adversarial learning, and we assumed that such an approximation of SB learning would create a suboptimal region that does not deviate drastically after a certain amount of time. The validity of assumption was considerably supported by the experiments.

Q8 [§ 4.3]

Following your comment, our revised manuscript has fixed the notation mismatch, which caused confusion about Theorem 3. In this paper, we mainly used $\rho_\tau \in \mathcal{C}$ for $\tau\in[0,\infty)$ in § 4.3 to denote the one-parameter semigroup for modeling OMD update from conditional distribution $\vec{\pi}$ₜ to $\vec{\pi}_{t+1}$. Therefore, $t$ is the time step for MD, which outputs static SB model $\pi$, and $\tau$ is the sub-time step. To avoid confusion, we have fixed the notation in Theorem 3.

Q9 [Eq. (14)]

Yes, Eq. (14) utilizes the parameterization from LightSB to perform variational OMD in the submanifold of Wasserstein-2 space, termed the Wasserstein Fisher Rao geometry.

Q10

In theory, if a push-forward measure from a non-trivial marginal distribution can be analytically drawn without critical errors, we can easily extend our results to dynamical models. For instance, one could apply multiple reparameterization tricks across the Wiener measure and draw a dynamic version of Proposition 2. Currently, flow-based NNs might be advancing in this direction. However, there can be multiple technical hurdles to overcome. Therefore, we mostly considered the theoretical side of learning and substituted the computational aspect with WFR geometry. As we believe GMMs are powerful nonlinear models for SB, we think the performance benefits under identical architecture constraints are sufficient within the scope of this paper to validate our theoretical claims.

Q11

Thank you for the insightful comment. We respectfully emphasize that MD works more meaningfully in the SB learning circumstance.

• The SB problem is convex for distributions. We assumed that the expression $\lim_{t\to\infty} \mathbb{E}_t[\delta_C\Omega(\pi^\circ_t)]$ is stationary as a mapping of Gaussian mixtures (by consistently changing its mode) using Bregman potentials.
• The problem is non-convex to its parameters. The exponential moving averaging (EMA) is the technique favors the "average" of historical parameters $\{\phi_t\}^\infty_t$. It stabilizes the learning process, but analyzing its meaning in the probabilistic space is difficult.
• Our OMD algorithm finds the best model that is closest to $\delta_D(\lim_{t\to\infty}\mathbb{E}_t[\delta_C\Omega(\pi^\circ_t)]$) where the Bregman potential is used to measure distance. Therefore, we claim that MD brings pleasing outcomes regarding statistical divergence induced by KL.

To further validate our point, we added a row of LightSB-EMA in Tables 6 & 7. The table below is the ε=1 result of cBW-UVP. This shows that VSMB consistently shows better results than LightSB-EMA, and our theoretical argument rigorously supports these gains.

Please feel free to reach out to us if you have further questions.

Sincerely,

Submission #5991 authors

[1] Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Communications in Partial Differential Equations, 2001.

[2] Mohammad Reza Karimi, Ya-Ping Hsieh, and Andreas Krause. Sinkhorn flow as mirror flow: a continuous-time framework for generalizing the sinkhorn algorithm. In AISTATS, 2024.

[3] Shun-ichi Amari. Information geometry and its applications. Springer, 2016.

[4] Frank Nielsen. An elementary introduction to information geometry. Entropy, 2020

[5] Christian Léonard. A survey of Schrödinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.0215, 2013.

[6] Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet. Schrödinger bridge flow for unpaired data translation. In NeurIPS, 2024.

[7] Stefano Peluchetti. BM$^2$: coupled Schrödinger bridge matching. arXiv preprint arXiv:2409.09376, 2024.

[8] Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, Alexander Korotin. Light and Optimal Schrödinger Bridge Matching, In ICML, 2024.

Q1 [Fig. 1]

The space $\mathcal{C}$ represents the coupling space of bidirectional transportation given by $d\pi = e^{\varphi\oplus\psi - c_{\varepsilon}} d (\mu \otimes \nu)$, subject to the probabilistic constraints. This is also equipped with extra conditions, as stated in the definition of $\mathcal{C}$ (L200), to ensure a nice geometric structure. Fig. 1 is a simplified representation of the solution space, where each SB solution $\pi_t$ is represented as a point.

One could attain a better understanding by considering the 3D illustration in Fig. 3. The figure shows that the projection spaces $\Pi^\perp_\mu$ and $\Pi^\perp_\nu$ represent distinct surfaces. Thus $\mathcal{C}$ and $(\Pi^\perp_\mu,\Pi^\perp_\nu)$ meet along curves, and this concept is simplified as Fig. 1. In Fig. 1, both types of learning involve $\mathcal{C}$, yet IPF iterations have a slightly more sophisticated geometric interpretation of alternating Bregman projections. For ideal algorithms, both MD iterations will iteratively find solutions $\pi^\ast$. In this paper, we are dealing with problems where we can not expect such ideal conditions.

Q2 [Lemma 2]

There was a missing superscript mistake in Lemma 2; please check the revised version addressing this point. The original intention was to link KL functionals to the original formulation Wasserstein gradient operator $\nabla_{\mathbb{W}}$ by Otto [1]. Of course, both expressions $\nabla_\mathbb{W} F(\pi_t)$ and $\nabla_\mathbb{W} f(\vec{\pi}^x_t)$ are equivalent in SB under the disintegration theorem. The paper deliberately chose the latter form for its actual usage in our VMSB algorithm in § 5.

Q3 [Eq. (6)]

We respectfully point out the influence of the divergence in Eq (6) indicates the proximity from the current point and is controlled by $\frac{1}{\eta_t}$. We find that gradually lowering $eta_t$ guarantees general convergence of OMD in SB. As you noticed, this is a fundamental equation. All of first-order algorithms discussed in this paper fall into (6) with specialization, where Lemmas 1 & 2 are dedicated to this standpoint. Arbitrary $F_t(\pi)$ that makes convergence to $\pi^\ast$ can be used to train SB can be chosen, and our choice of online cost $F_t$ is drawn from the static case of Wasserstein gradient descent in Lemma 2.

Q4 [$\alpha_t$]

We agree that certain enumerating styles $a_t$ and $\alpha_t$ were too short or did not match the overall formal presentation style. Following your insightful comment, we enhanced the clarity of expressions in L209 and L211, which became more formal and appropriate for both expressions.

Q5 [Lemma 2]

Yes, it is the apparent, yet vital premise that algorithms do not access $\pi^\ast$ ($\mu$ & $\nu$ might be accessed in assumptions of some OT papers). Historically, this is also the reason why $\tilde{F}_t$ of Sinkhorn is defined as it is. Therefore, implementation of Wasserstein gradient flows and acquiring plausible cost $F_t$ without access of $\pi^\ast$ were our core algorithmic considerations.

Q6 [α-IMF & BM$^2$]

We appreciate your insightful comments and great papers [6] and [7]. We would like to clarify that we were unaware of these recent articles that were released this September. We generally agree that both achieved great progress in the generalization of SB learning, yet the directions of progress can also be understood from this paper’s perspective.

• α-IMF [6]
  • Similar to Sinkhorn, IMF deals with two distinct surfaces, named Markovian projection space and reciprocal projection space (Fig. 1 of [6]). To the best of our knowledge, there is no definite answer to which choice is better.
  • The paper proposes α-DSBM, where we can recover DSBM-IMF (Shi et al., 2023) by α=1. Considering $\alpha \to 0$ enables performing non-iterative, online versions of IMF, and α becomes a step size for approximation.
  • We conducted several comparison studies with bridge matching by considering LightSB-M [8] and DSBM. See that Theorem 3.3 of Gushchin et al. [8] matches Eq. (9) of [6]. Thus, we respectfully point out that the conducted experiments can cover the α-DSBM case by considering the performance of LightSB-M and DSBM.
• BM$^2$ [7]
  • BM$^2$, or coupled bridge matching, demonstrated an intriguing idea aiming for a simple loss. We think BM$^2$ offers streamlined procedures and balanced samples.
  • In Algorithm 1 of [7], the loss is computed by bridge matching, and the reference dynamics of the Brownian bridge are efficiently collected. We think this approach is promising, although it would be challenging to analyze what the theoretical benefits would be without some additional assumptions.
  • By comparing with the experimental results of [7], we believe our VMSB method shows better performance in the EOT benchmark.

We greatly appreciate the time and effort you have invested in reviewing our submission. We are currently incorporating all of the invaluable feedback in the manuscript, and we are pleased to report that many of your suggestions are used in the current revision. We address your comments below.

Suggestions

Following your fruitful suggestions, we have carefully revised the manuscript. The clarity of the manuscript has improved significantly from the submission, thanks to your insightful feedback. We summarize the improvements below:

• (§ 1) We have enhanced the clarity of sentences around Eq. (1), improving the focus on optimization and stability.
• (§ 4.2) We overhauled L288 and L296-L302 to explain theoretical results in a less formal tone; some important formal parts of the current sentences have been moved moved to the end of Appendix A.4.
• (§ 5.1) We added an opening sentence to describe the LightSB formulation. For more technical details for LightSB, we also added Appendix C.1, which provides theoretical and technical rationales for our choice of the GMM parameterization.
• (§ 5.2) We enhanced the clarity of writing by including an explicit sentence regarding the algorithmic implication of the proposition.

Dense presentation style

We appreciate your insightful feedback. Starting from Eq. (6), we devised a variational online mirror descent (OMD) algorithm (Algorithm 1) for the derived update rule Eq. (10), using the tractable form of Eq. (10). We agree that the sheer volume to reach our algorithmic decision could be overwhelming to readers who are not familiar with the subjects. We would like to highlight our contributions and provide the context below.

Contributions & Key aspects of the proposal

The contributions of this paper stated in L89-L94 can be explicitly delineated as follows.

• A new OMD method with strong theoretical guarantees. This paper provides a fine level of convergence analysis of an online learning problem in SB with a rigorous regret bound.
• Simulation-free variational SB solver. Compared to contemporary SB algorithms, the proposed VSMB is well-grounded and can work under various real-world circumstances.
• A computational SB method validated in a wide range of benchmarks. Our work successfully verified the MD framework in SB and corresponding theoretical claims with experiments.

In a less formal tone, we respectfully emphasize the following points.

• Solving SBP with learning theory. Inspired by prior studies, our study casts a fresh geometric perspective to minimize online costs. We effectively show that each learning step of SB can be precisely taken analytically, with a gradient-based algorithm.
• New gradient-based SB algorithm. We show a new learning method for SB, which is drawn from variational inference. Applying the proposed gradient dynamics to the LightSB parameterization, this technique effectively outperforms prior methods in our experiments.
• As a minor contribution, we believe that readers can better understand the EOT/SB problem with a broader perspective from our study.

Fundamental & required concepts for methodology

The following points are a list of important concepts for drawing VMSB.

• Information geometry. Bregman divergences and MD are closely related to information geometry [3,4]. Our paper deals with the geometric problem with infinite-dimensional analysis, and we use Wasserstein geometry of distributions as an infinite-dimensional analogue of the Riemannian geometry.
• Learning theory. SB algorithms fundamentally struggle with imperfect learning signals (Figs. 4 & 5). Inspired by learning theory, we devised an online learning algorithm for a stable SB algorithm through the OMD theory.
• Measure theory. One of our theoretical contribution lies in the application of log Sobolev inequalities (Assumption 3) into our analysis. Understanding of the measure theory is required to understand Theorem 2 on its technical side, and the choice of GMM parameterization.