Multi-Marginal Schrödinger Bridge Matching

We thank the reviewer for both the recognition of our study’s strengths and the valuable, comprehensive feedback. Our point-by-point responses are laid out below.

---

1.Multi-Marginal Schrödinger Bridge.

• The global network parameterization is widely used, but in our work it serves only as a practical means to acheive our goal of solving the multi-marginal Schrödinger Bridge problem (mSBP) efficiently.

• Prior works [1, 2] that use a single shared network across intervals focus on pairwise SB or on classic optimal transport, then they glue together several two-point solutions. As a result, the final composite evolution can satisfy each local pairwise constraint.

• These approaches may overlook the global requirement of a single trajectory that satisfies the multi-marginal constraint over the entire time horizon, so they are not guaranteed to produce the KL-optimal—or most probable evolution. Furthermore, ensuring OT-optimality presumes the local or global coupling is known or can be cheaply estimated, a requirement that quickly becomes computationally prohibitive as the sample size grows.

• Our approach is different. We lift the standard IMF algorithm itself to the multi-marginal case, and prove that the alternating our proposed multi-marginal reciprocal and Markov projections still converge to the true mSBP minimiser (SB optimality).

• This optimality gives two benefits: (1) the recovered dynamics is the single most probable evolution with minimum control energy; (2) any existing IMF-based SB solver can now handle an arbitrary (countable) number of snapshots by substituting our multi-marginal projections without losing convergence or optimality guarantees.

• On the practical side, based on these theories, we propose an MSBM algorithm using a single shared network because it is the simplest way to enforce boundary-matching drifts (Corollary 5). Consequently, our method requires no pre-computed couplings, reaches the joint KL minimum for all snapshots, and does so in markedly less runtime than prior multi-marginal SB approaches.

---

2.Baseline and Benchmark Results.

• Thank you for pointing us to additional multi-snapshot benchmarks. Our original experiments focused on previous multi-marginal SB methods (DMSB, SBIRR) because our primary goal is extending IMF to the multi-marginal constraints for an efficient mSBP solver rather than a new state-of-the-art.

• We fully agree, however, that flow-matching variants have shown strong results and merit inclusion. In response, we have run MSBM, DMSB on CITE-seq and MULTI-seq at both 5- and 100-dimensional PCA resolutions. Baseline figures marked with an asterisk are taken directly from [4]. We will integrate the full tables and additional baselines in the revised manuscript.

• We omitted SBIRR on these settings because their public implementation did not scale beyond the small-scale regimes reported in our paper (e.g., small datasets such as hESC). On the larger dataset such as CITE-seq or MULti-seq, it ran out of memory.

---

3. Computational Efficiency.

• Moreover, we measured the full training time on the CITE‑seq 100-dim benchmark using the same hardware for every method (a single NVIDIA RTX 3090 GPU; the main paper used an NVIDIA  A6000 GPU). Although MSBM is not the state-of-the-art for every configuration, it still clearly outperforms all baselines in training speed. Because IMF-type algorithms do not require pre-computed OT-coupling and eliminates the need to store intermediate states.

• OT‑CFM was excluded because its implementation details for this specific task are unavailable, preventing a fair comparison.

---

[1] Tong et al., Improving and Generalizing Flow-based Generative Models with Minibatch Optimal Transport
[2] Tong et al., Simulation-Free Schrödinger Bridges via Score and Flow Matching
[3] Rohbeck et al., Modeling Complex System Dynamics with Flow Matching Across Time and Conditions
[4] Kapusniak et al., Metric flow matching for smooth interpolations on the data manifold.

Thank you for the rebuttal and for investing the time to add the CITE-seq and MULTI-seq benchmarks, as well as the detailed runtime comparison. The extra results and clarifications are much appreciated.

Your extension of Iterative Markovian Fitting to the multi-marginal Schrödinger Bridge setting is now clearer, and the convergence proof addresses the global-consistency issue. At the same time, I still struggle to see a practical advantage (and still practical novelty) over recent flow-matching that already share a network across intervals, particularly as the new tables show MSBM's performance is lower than that of OT-MFM or WLF-SB. I suspect the W1 metric may be insensitive to the theoretical advantages of your method outlined above (pathwise continuity). Could a more suitable metric (e.g., RNA-velocity cosine similarity) better capture these benefits?

The rebuttal states that MSBM “requires no pre-computed couplings” and therefore scales better. However, Tong et al. (2023, “Simulation-Free SB via Score & Flow Matching”) report successful training with a geometric coupling in a 1000 dimensions, suggesting that coupling estimation may not be as prohibitive a bottleneck as the authors state.

While it is outside of the scope of current submission, the provided comparison outline that state-of-the-art baselines increasingly integrate geometric information—for instance, the coupling prior in Tong et al. (2023), the manifold interpolation in Kapusniak et al. (2024). Particularly, Liu et al. (2023, “Generalized SB Matching”) incorporate this in SB problem. Do you foresee a straightforward way to incorporate Multi marginality into GSBM?

I realise that adding new experiments during the rebuttal window may not be feasible, but any clarification on the above points would make the paper stronger. Depending on your answers, I am open to revisiting my overall score.

Please let me know if I have misunderstood or misstated any of your claims—I am happy to correct them.

[1] Tong et al., Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport
[2] Zhang et al., Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport
[3] Baradat et al., Minimizing relative entropy of path measures under marginal constraints
[4] Liu et al., Deep Generalized Schrödinger Bridge

We hope the clarification we posted has fully resolved the reviewer's earlier concerns. With one day remaining in the discussion period, we would like to provide this additional rebuttal.

---

3.Novelty.

• The "global-network + glue" idea is not itself the novelty we emphasize. In our work, this parameterization is simply the most practical means to realize our primary contribution: an extension of the SB Matching framework to the mSBP with full theoretical guarantees of optimality.
• The novelty lies in lifting the IMF algorithm to the multi-marginal setting and proving that alternating our multi-marginal reciprocal and Markov projections converges to the true mSBP minimizer. And we further propose an efficient learning algorithm that can aid the error accumulation of the direct extension of SBM into multi-marginal SBM which we called Naïve extension in the manuscript.
• This is, to the best of our knowledge, the first SB-based algorithm that achieves simulation-free training in the multi-marginal case without requiring any pre-computed OT couplings, while retaining optimality guarantees for multiple constraints rather than just two.
• While recent multi-marginal flow-matching approaches also use a shared network across intervals, they typically solve a sequence of pairwise problems and "glue" the results, which does not ensure KL-optimality over the entire horizon. In contrast, our formulation directly targets the global mSBP objective.
• Even if current W1-based metrics show comparable empirical results, achieving optimality has intrinsic advantages, such as principled scalability to larger datasets and the ability to integrate richer prior structures without losing convergence guarantees. Moreover, our algorithm achieves significantly shorter training times (x6~27 times faster compared to baselines), likely due to the rapid and stable convergence properties of IMF-style optimization.

---

We hope this additional clarification removes any remaining misunderstanding about the source of novelty in our work and makes clear the theoretical and practical significance of our contribution. In light of this, we respectfully invite the reviewer to reconsider their rating so that it reflects both the theoretical advance and its practical implications.

We sincerely appreciate the reviewer’s acknowledgement of our contributions and the depth of their analysis. Detailed responses to each point are presented below.

---

1. Clarification on Naive extension, Algorithm 1, Corollary 5.

• We sincerely thank the reviewer for raising this important question, which helped us clarify the conceptual distinction between the naive extension and MSBM.
• Even though the naive extension based on two multi-marginal projection operators (Section 3) yields a valid mSBP solution in theory (based on Proposition 3), naively following original SBM algorithm (Equation (7)) with our developed multi-marginal projection operators suffers from bias in practice (shown in Figure 1). This is because intermediate marginals are not directly anchored to training data but are instead inferred sequentially from earlier steps, leading to compounding errors.
• What we want to emphasize in Section 4 is that Corollary 5 directly addresses and mitigates these limitations. Although our proposed MSBM (based on Corollary 5 and Algorithm 1) differs algorithmically from the naive extension, both are theoretically designed to recover the same optimal mSBP solution! Corollary 5 guarantees that if pairwise SBs over adjacent intervals are continuously joined—meaning their drifts are continuous at boundary points—then the piecewise formulation in equation (18c) induces a globally continuous process that satisfies the optimality condition in Proposition 3. In this sense, MSBM is not just a heuristic alternative but a practical and principled implementation of the same theoretical solution, designed to avoid bias by directly supervising all intermediate marginals during training.
• We will reflect these clarifications in the revised manuscript.

---

2. Global parameterization.

• The condition on $v_i$ and $u_i$​ refers to whether they are continuous functions with respect to time $t \in [0, T]$. In our implementation, we enforce this by using a global neural network parameterization for both $v$ and $u$, and since standard neural networks are continuous with respect to their input (assuming typical architectures and activation functions), this guarantees continuity across adjacent intervals.
• In this regard, we agree with the reviewer that smoothness is a stronger condition and generally not required in our formulation. Accordingly, we will revise Line 164 to replace "smooth between adjacent timestamps" with "continuous between adjacent timestamps" to better reflect our intended meaning.

---

3. Baselines and Experimental Details.

• We thank the reviewer for the valuable suggestions to enhance our work. We conducted the experiment with various baselines with additional benchmarks. Please refer to the response to the Reviewer chJE section for a detailed explanation. *In table 3, “Full” column denoting every snapshots $t_0, t_1, t_2, t_3, t_4$ is used during training, and the reconstruction error is measured on the same set of snapshots. In the three settings $t_1$, $t_2$ and $t_3$, we remove one snapshot that is omitted during training and inferred at test time.
• Comparing the errors in these three columns with the “Full” column shows how much each method degrades when it must interpolate to a time point it never saw during training.
• This evaluation protocol follows exactly in DMSB, so it differs from the settings employed in Table 1 and in the additional experiments.
• Table 2 uses the same “Full” protocol as Table 3. The model is trained with data from every snapshot time. Within each snapshot we randomly split the dataset into training and testing subsets.

---

4. Reference Dynamics.

• The reference process $\mathbb{Q}$ is defined by the drift $f_t$. In principle, it does not have to be a pure Brownian motion. Indeed, any SDE with continuous drift can serve as a reference as long as its bridge kernel $\mathbb{Q}_{s|t}$ is tractable. Classic examples include pure Brownian motion, Ornstein-Uhlenbeck, and, more generally and linear time-inhomogeneous SDE with constant diffusion [1, Appendix B] list the main tractable cases.
• We agree that when domain knowledge provides a reasonable prior on the underlying dynamics, it is beneficial to embed that information in $f_t$ such as in [2]. In many data-driven settings such as scRNA-seq, however, we lack a reliable physical model, so the reasonable option is to set $f_t = 0$ and use a Brownian reference.

---

5. Minor Comments.

• We appreciate the feedback regarding readability. We will conduct a thorough proofreading pass to address any syntactic errors, enhance sentence flow, and improve overall clarity.

• Equations (14) and (17) : Thank you for catching the typo; the correct definition is $\gamma_{\mathcal{T}}(t) = \max_{u}[u < t |u \in \mathcal{T}]$, not $\gamma_{\mathcal{T}}(t) = \max_{u}[u < t |t \in \mathcal{T}]$.

• MMD kernel: We adopt a multi-scale Gaussian kernel. Specifically, we set the bandwidth to the batch’s average squared distance, generate four more by repeatedly doubling or halving that scale, and sum the resulting five Gaussian kernels to compute MMD.

• Line 86-87 and 108-110, Figure 1 and Layout on Page 9: Because of the strict page limit, some explanations were omitted and the formatting became somewhat cramped. The extra page allowed for the camera-ready version will let us restore the omitted explanations and reorganize the formatting for clarity.

---

[1] Shi et al., Diffusion Schrödinger Bridge Matching.
[2] Shen et al., Multi-marginal Schrödinger Bridges with Iterative Reference Refinement.

Thank you for recognizing the contributions of our work and for the careful, thoughtful review. We address every comment in detail in the responses provided below.

---

1. Inconsistent Baseline Comparisons.

• We appreciate the reviewer’s suggestions. In response, we conducted additional experiments across various RNA-seq datasets and summarized the results in our reply to reviewer chJE in below. We excluded SBIRR because, in our attempt, its public implementation exhausted memory on larger datasets such as CITE-seq and MULTI-seq. However, we believe that our evaluation remains thorough, as it already includes well-established flow matching and Schrödinger bridge baselines.

---

2. Limited Diversity and Feature Dimensionality

• We agree with the reviewer that showing wider generalizability is important. In this study we concentrated on making IMF practical for the multi-marginal case, so we evaluated our method on benchmark tasks widely used in recent work on multi-marginal SB and trajectory inference, namely single-cell datasets reduced to 5–100-dimensional PCA space. Exploring additional domains such as high dimensional image data, longitudinal time series, and other generative OT tasks forms a clear direction for future work.

---

3. Generative Capabilities, Robustness and Sparse Time Settings.

• The central objective of MSBM and trajectory inference is precisely to learn a joint path measure from which new samples can be drawn at any time. In other words, the method is generative by construction. We already evaluate this property in the leave-one-out experiment (Table 1,2) that removes one intermediate snapshot during training. After fitting on the remaining snapshots, the trained model synthesises data at the held-out time, and we compare the generated samples with real observations. The close match demonstrates that MSBM does more than interpolate observed marginals. Specifically, in Table 1, MSBM keeps the computational advantage over SBIRR. These findings confirm that our method does not trade robustness for speed.

• Concretely, It produces new data that follow the true distribution, which is the essence of generative modelling in dynamic OT. We will clarify in the main text that the task itself is framed from a generative modelling perspective, since the learned stochastic process defines a sampler capable of generating unseen temporal snapshots that align with the underlying data manifold.

• Moreover, repeating the same protocol with OT-CFM and WLF-SB shows accuracy differences that lie within the statistical margin, confirming that MSBM is at least as accurate as these well-known baselines. These experiments demonstrate that MSBM delivers robustness in addition to its computational advantage.

---

4. SDE Formulation and Error Accumulation.

• For simulation in Algorithm 2, we use the dynamics in Eq (10). Specifically, after learning, global parameterization lets the dynamics in Eq (18a) recover the dynamics in Eq(10) while assuring that it is an mSBP solution.
• If the model fails to capture the local dynamics in any interval, the resulting path becomes suboptimal and can no longer be guaranteed to solve the mSBP. In practice, an error introduced at an early stage could snowball through subsequent intervals, with bias accumulating as the trajectory unfolds. In our experiments we did not observe such propagation. We attribute this to the rapid and stable convergence of IMF-style optimisation, in which joint updates across all intervals quickly correct local mismatches before they have a chance to cascade.

---

5. Random dropout.

• When many snapshots are available, the shared drift $v_{\theta}$ can infer the missing segments from neighbouring data, so the procedure should curb local overfitting and improve generalisation to unseen times. In our setting the number of snapshots is modest, so additional dropout could remove critical information and push the model toward underfitting. For datasets with dense temporal sampling the idea looks useful as a regularisation technique and is an interesting direction for future work.

---

6. Appendix Citation.

• We thank the reviewer for pointing this out and will insert the pointer appropriately in the revised manuscript.

We gratefully thank the reviewers for their valuable feedback and recognition the contribution of our work. Here, we address the concerns raised by the reviewer.

---

1. The continuity condition,

• Ensuring path-wise continuity for each local bridge is difficult if one treats the intervals independently. We eliminate the difficulty by embedding continuity in the parametrisation itself such as a single neural network, shared across all sub-intervals. Because common network architectures are continuous in their inputs, the map $t \mapsto v_{\theta}(t, \cdot)$ is continuous on $[0,T]$. Consequently the left limit $v_i(t_i^{-})$ and the right limit $v_i(t_i^{+})$ coincide at every internal time $t_i$​.
• Equation 20a reinforces this property: its aggregated loss evaluates samples that lie just before and just after every boundary, so each gradient step simultaneously penalises any mismatch and pulls the two sides together. The continuity requirement is therefore satisfied by construction; it is not an extra constraint we must enforce post-hoc.

---

2. Flow Matching-Based Baselines.

• We appreciate the reviewer’s helpful suggestions. In response to requests from the reviewers, we have run additional experiments that include flow-matching based algorithms. Upon computational efficiency, our proposed MSBM clearly outperforms base line methods including the reviewer's suggestion, while achieving comparable performance on the single-cell RNA sequencing benchmarks. Please We thank the reviewer for their valuable feedback. Following the request, we conducted additional experiments that incorporate flow‑matching baselines.
• Notably, In terms of computational efficiency, our MSBM consistently surpasses all baseline methods (including the reviewer suggested methods) while delivering comparable performance on single‑cell RNA‑sequencing benchmarks. Detailed results can be found in the response the reviewer chJE in below.

---

3. Minor Comments.

• Thank you for pointing out the readability issues. We will incorporate the suggested edits to improve clarity in the revised manuscript.