We thank the reviewer for their positive and constructive comments, and we address their comments and questions below. We first provide a brief rebuttal summary:

• Appendix: The appendix is located in the zip folder under ‘Supplementary Material’: Then go neurips_treeIMF_supplementary -> supplementary.pdf. Please let us know if you have issues accessing it.
• The Appendix contains many additional experiments and discussions relevant to the reviewer's questions, in particular comparisons with alternative continuous barycentre algorithms, and discussion of runtimes.
• We provide additional information and discussion regarding computational resources in this rebuttal.

We expand on these points in more detail below, and highlight any proposed changes for revisions in bold.

Comparison with other continuous barycentre methods

We do in fact provide comparisons with other continuous Wasserstein-2 barycentre solvers in the Appendix. Namely, we include comparisons with state-of-the-art methods WIN [1], W2CB [2] and NOTWB [3]. These methods are chosen to be representative of different approaches in the literature— WIN is an iterative fixed-point approach, W2CB is an Input Convex Neural Network-based method, and NOTWB is based on recent Neural OT methodology.

In particular, we include in Appendix E comparisons on a subset posterior aggregation experiment (used as a standard experiment in continuous Wasserstein-2 barycentre approaches, e.g. in [2],[4]), and also on high-dimensional Gaussian examples for $d=64,96,128$ (again a standard experiment in the literature, e.g. in [1],[2],[3]). In lines 1130-1145 we provide a discussion regarding the findings, highlighting the pros and cons of different approaches in the literature.

Importantly, if accepted we intend to move the data aggregation and high-dimensional Gaussian experiments to the main paper, using the extra page afforded to us.

Training time, inference time, GPU memory consumption

Training time:

• We discuss run-times of the various algorithms that we compare in the Appendix in lines 941-952, and in lines 1079-1095, and TreeDSBM appeared favourable compared to existing barycentre solvers.

Inference time:

• Inference time for TreeDSBM is slower than for methods that use a single function evaluation, as it involves solving an SDE (discussed in line 361 in main paper, and line 1139 in the Appendix). However, it can still be fast if one JIT-compiles the sampling procedure: transporting 1,000 samples along one of the maps in the data aggregation experiment takes 7.8 ms for the single-evaluation NOTWB method, and 15 ms for TreeDSBM when using 50 discretisation steps.

GPU memory consumption

• We comment here that the parallel nature of the algorithm during training means that a practitioner can have a trade-off between memory consumption and wall-clock time. Namely, one can train the edges simultaneously if compute allows (either on a single GPU if enough memory, or parallelised across GPUs). If this cannot be done, one can train sequentially instead (in which case memory requirements for training would be comparable with standard bridge-matching).
• We report GPU consumption for different experiments below (for sequential and joint training), for the hyperparameters used in the paper. These can be changed significantly by changing hyperparameters such as batch size.

• We note that in some of these experiments (e.g. Gaussian), peak GPU usage is due to the simulation and storage of the training samples for subsequent IMF steps, rather than during the network training. This can be reduced significantly by simulating in smaller batches, or by updating the cache during training rather than simulating all beforehand.
• We also note that TreeDSBM has two major benefits over TreeDSB regarding run-time and memory requirements. Firstly, TreeDSB cannot be parallelised in the same way, as the training of the networks in each IPF iteration rotates around the edges sequentially. Secondly, the memory required to store training data for TreeDSB is far higher, as training a time-reversal along an edge requires saving entire trajectories simulated along the reverse direction, whereas TreeDSBM only requires storing endpoints (this is one of the most significant benefits of IMF over IPF).

Application beyond barycentres

Our initial motivation for our study was to utilise the IMF procedure in combination with the highly-performant fixed-point approaches for barycentre approximation. We include a detailed exposition of this perspective in the response to Reviewer 3rwT.

We found that the tree-based framework provided the correct mathematical framework to combine these approaches in a simple and elegant way. Given our initial fixed-point motivation, the fact that our approach extends to more general tree settings was something of an added bonus that arose out of utilising the tree-based framework. Nevertheless, such problems have been utilised in works such as [5]. We conjecture that it could have applications in trajectory inference problems where populations might be known to 'split' according to a given structure (for example, maybe for single-cell dynamics). We unfortunately don't have the biological expertise to be confident in devising meaningful experiments for such settings, so instead provided a toy experiment in Appendix E.3 in which TreeDSBM appeared to work well, demonstrating fast and stable convergence. We have a brief paragraph discussing these wider problem settings in Appendix E.3, and propose extending this to include the above discussion, and referencing this paragraph at the end of the Related Work section. We hope future practitioners can build on our work for use in such settings.

We thank the reviewer for their already kind feedback, and we hope that the contents of the Appendix and this rebuttal have addressed their comments. If you have any further questions, please let us know!

References

[1] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation

[2] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization

[3] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport

[4] Fan et al, 2021, Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks

[5] Haasler et al, 2021, Multi-marginal Optimal Transport with a Tree-structured cost and the Schrödinger Bridge Problem

We thank the reviewer for their positive and constructive comments, and we address their comments and questions below. We first provide a brief rebuttal summary:

• 'Originality': We propose to add a paragraph to the introduction highlighting the fixed-point barycentre interpretation more prominently, to highlight that our approach occupies an important position in the literature.
• 'Figure 3': The noise in the samples in Figure 3 is expected, due to the entropy regularisation. There are ways to reduce/avoid this - see Figure 7 in the Appendix.

We expand on these points in more detail below, and highlight any proposed changes for revisions in bold.

"One could perhaps argue that this approach is only mixing already established techniques"

We appreciate the reviewer for raising this point (and for also noting this is not a strong limitation). Currently the narrative of the paper focuses on TreeDSBM being an IMF counterpart to TreeDSB. We discuss below a slightly alternative perspective (which in fact initially motivated our approach), which we hope highlights more clearly why we believe our method occupies an important and prominent position in the literature.

One of the most successful and elegant approaches for Wasserstein-2 barycentre computation is the iterative fixed-point approach, which involves iteratively updating a candidate barycentre $\nu_j$, by solving for the OT map $T_i$ from $\mu_i$ to $\nu_j$, and updating $\nu_{j+1}$ via the construction $\sum_i \lambda_i Y_i$ (where the $Y_i$ come from the coupling induced by the maps $T_i$). This approach was popularised by the seminal paper [1], and has formed the basis of many algorithmic developments, including in machine learning in [2] (one of the current SOTA approaches). Indeed, it has strong performance, and has been observed to converge quickly in only a few iterations [3]. However, a downside is that it requires solving $K \cdot N$ complete OT problems, where $K$ is the number of marginals and $N$ is the number of outer iterations, which is expensive as solving even a single OT problem can be challenging.

The IMF procedure [4] has recently gained popularity for solving the entropy-regularised OT problem, and proceeds by iteratively performing bridge matching on its induced couplings. It is therefore a prominent question whether IMF can be used in combination with a fixed-point approach for solving for Wasserstein-2 barycentres. A naive approach would be to insert IMF to solve for each OT subproblem, requiring $K \cdot N$ complete IMF problems, which would in turn be $K \cdot N \cdot M$ bridge matching problems (moreover, making this approach grounded in theory would be non-trivial, as the fixed-point method holds only for the unregularised OT problem).

The key aspect of our approach is that, rather than nesting the fixed-point iterations and the IMF iterations, they can in fact be elegantly combined into a single iterative procedure, along with a well-understood theoretical grounding. In particular, our approach shows that the expensive OT computations required for the fixed-point procedure (the main limitation of such approaches), can instead be switched out for inexpensive bridge-matching procedures. Each fixed-point iteration is therefore significantly cheaper, and empirically we found the TreeDSBM algorithm to retain the fast convergence of the fixed-point procedure in terms of the number of iterations required.

It is for this reason that we believe TreeDSBM is an extremely natural algorithm for Wasserstein-2 barycentre computation, and thus occupies a prominent position in the literature. We consider the fact it is a simple and elegant combination of the tree-structured framework with IMF to be one of its main strengths.

We believe that the paper would benefit from advertising this fixed-point perspective more prominently, as currently the emphasis focuses on being an IMF counterpart to TreeDSB. We thus propose to add a paragraph in the introduction highlighting this motivation and interpretation.

results of Figure 3

We currently provide a discussion regarding this point in Section D.2 of the Appendix (lines 978-985) and in Figure 7. We propose adding a reference to this discussion to this in the main paper. We recall the main points from this discussion below.

The results plotted in Figure 3 use a fairly large value of entropy-regularisation parameter $\varepsilon=0.02$ (which was selected to be more comparable with TreeDSB, which requires a large $\varepsilon$ to converge). This means that there should be noise present in an accurately computed solution, which is why the results look as they are. Indeed, the results shown in Figure 3 resemble those for other entropy-regularised barycentre approaches (see e.g. results for EGBARY in Figure 5 of the paper [5] linked by the reviewer).

Clearly this may not be desirable; fortunately this can be reduced/avoided in two ways (which we do in the Appendix):

• Using a smaller $\varepsilon$ value (e.g. $\varepsilon=0.001$)
• Generating samples as $\sum_i \lambda_i Y_i$, where $Y_i$ are samples from the coupling over the marginals induced by the learned transport maps.

Samples using both these approaches are included in Figure 7 in the Appendix (note minor typo, caption of (a) should read $\varepsilon=0.02$ rather than $0.2$). These samples are much more comparable with those in the work [5] referenced by the reviewer.

Other changes

Finally, we highlight here some additional changes that we propose, following feedback from other reviewers.

• Move the comparisons with other continuous barycentre solvers [2].[5],[6] to the main body (they are currently in Appendix E.1 and E.2 due to space constraints)
• Additional ablations on the role of epsilon, batch size (see response to reviewer QVwU).

We hope that these comments have addressed your questions and concerns. If you have any further questions, please let us know!

References

[1] Alvarez-Esteban et al, 2016, A fixed-point approach to barycenters in Wasserstein space

[2] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation

[3] von Lindheim, 2022, Simple Approximative Algorithms for Free-Support Wasserstein Barycenters

[4] Shi et al, 2023, Diffusion Schrödinger Bridge Matching

[5] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport

[6] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization

We thank the reviewer for their feedback, and we address their comments and questions below. We first provide a brief rebuttal summary:

• Experimental validation: There are several further experiments in the Appendix, many of which address concerns of the reviewer. In particular, there are comparisons with alternative barycentre algorithms, and discussion of runtimes. If accepted, we intend to use the additional space to move some of these to the main paper.
• Scalability: We address the reviewer's concerns in turn below. In particular, in the majority of use-cases, network sizes will be fairly small and thus not pose a computational bottleneck; in cases where it does, we outline ways it can be improved.
• CelebA: We have subsequently improved performance on this benchmark by using a shared architecture across the edges.

We expand on these points in more detail below, and highlight any proposed changes for revisions in bold.

Experimental validation

We note that there are several further experiments in the Appendix (Sections D and E), which could not be included in the main paper submission due to space constraints. We anticipate the reviewer may not have seen these, based on their comments. We hope these results allay the reviewer’s main concern regarding experimental validation.

In these experiments, we include comparisons with state-of-the-art methods WIN [1], W2CB [2] and NOTWB [3]. These methods are chosen to be representative of different approaches in the literature— WIN is an iterative fixed-point approach, W2CB is an Input Convex Neural Network-based method, and NOTWB is based on recent Neural OT methodology.

In particular, we include in Appendix E comparisons on a subset posterior aggregation experiment (a standard experiment in the continuous Wasserstein-2 barycentre literature, e.g. in [2],[4]), and also on high-dimensional Gaussian examples for $d=64,96,128$ (also a standard experiment, e.g. in [1],[2],[3]). In lines 1130-1145 we provide a discussion regarding the findings, highlighting the pros and cons of different approaches in the literature.

Importantly, we intend to move the subset posterior aggregation and high-dimensional Gaussian experiments to the main paper if accepted.

In Appendix D there are also more detailed examinations of the experiments in the main paper, including runtime analysis, barycentre progression during IMF iterations, and the effect of $\varepsilon$.

• Run-time analysis: We discuss run-times of the various algorithms in the Appendix in lines 941-952, and lines 1079-1095, and TreeDSBM appeared favourable compared to existing barycentre solvers.
• Ablations: We will add additional results concerning the role of batch size, and regularisation value $\varepsilon$ to the Appendix—please see response to reviewer QVwU.

Implementation details and scalability

Number of edges:

• The reviewer notes a concern that learning $2 |E|$ separate networks may cause scalability problems. We note that this is similar in many continuous Wasserstein barycentre algorithms, including TreeDSB and e.g. [1],[2],[3]. Fortunately, there are ways to improve this if necessary, and we also do not envisage this being a limiting factor for the vast majority of use cases.
• As discussed in Appendix C, one can use a single network along each edge and additionally condition on the direction. As the reversed drifts should be approximately the negative of each other, this can be efficiently parameterised as $(-1)^{dir} v_{\theta}(x,t,dir)$ (as done in [6], who show it has minimal impact on performance).
• More generally, one could use a single shared network by additionally conditioning on the edge $e$. This makes particular sense for problems with shared structure between edges.
• In our experiments, memory bottlenecks only arose in the CelebA experiment, which is exactly a scenario in which there is shared structure. We have subsequently used this single-network parameterisation and found improved performance (see below).
• We propose to extend the discussion in Appendix C to include conditioning on edges too.

Memory requirements, wall-clock time

• Our algorithm offers a trade-off between memory consumption and run time due to its parallel nature. Namely, one can train edges simultaneously if compute allows (on a single GPU if enough memory, or parallelised across GPUs). If this cannot be done, one can train sequentially instead (in which case memory requirements for training would be comparable with standard bridge-matching).
• In our experiments, we were (other than for CelebA) able to train simultaneously on a single GPU as the networks are not that large. Please see the response to reviewer ZjRa for further information regarding GPU usage. We anticipate that the majority of use-cases in practice would be for pointcloud data (e.g. the data aggregation experiment), which would use small MLP-like architectures with low memory requirements.
• Comparison to TreeDSB: We note that TreeDSBM has two major benefits over TreeDSB regarding run-time and memory requirements. Firstly, TreeDSB cannot be parallelised in the same way, as each IPF iteration rotates the training around the edges sequentially. Secondly, TreeDSB requires saving entire trajectories to train the time-reversals, while TreeDSBM only requires storing the endpoints (this is one of the most significant benefits of IMF over IPF).

CelebA

We currently provide an extended discussion regarding the CelebA experiment in lines 1006-1029 in the Appendix. In particular, we identify that the initial coupling (using bridge-matching models along each edge) does not always provided images that have structural similarity, so the resulting initial barycentre samples sometimes consist of distinct overlaid images. The subsequent IMF steps struggle to recover from this, resulting in the artifacts visible in some of the generations.

In this benchmark, the true maps are relatively simple (primarily just colour changes), but TreeDSBM struggles to recognise this simpler structure due to infrequent communication between edges (which occurs only after each IMF iteration). In contrast, methods which perform well on this benchmark optimise all the maps together and with much more interaction between them, providing more scope for identifying the shared structure present in this specific problem. We hope this helps clarify for practitioners when TreeDSBM is likely to perform well or poorly.

We hypothesised in the Appendix that different architectures could make TreeDSBM more competitive in this setting. We have since used a shared architecture as outlined in the previous section, which has improved the results (though they remain short of the best-performing methods for this problem). By learning features jointly across edges, the bridge-matching models behave more similarly along each edge, yielding coupling samples with greater structural similarity.

We thus propose to include in the Appendix a comparison of the results for the two parameterisations, to illustrate the role that architecture can play in such highly-structured settings. We emphasise that much of the behaviour discussed here is likely specific to this particular benchmark; we anticipate the majority of practical use-cases would be for pointcloud data, for which similar behaviour would not arise.

Convergence Rate

We are glad that you find the fast convergence of TreeDSBM compelling. As discussed in the response to reviewer 3rwT, we believe that the TreeDSBM algorithm inherits good convergence properties from the similar fixed-point barycentre algorithms, which are known to converge quickly (see e.g. [7], who show good results with just a single iteration).

We can include quantitative results regarding convergence speed in terms of IMF iterations in the Appendix. We provide such results for the 2d and data aggregation experiments below.

We can also provide similar results for TreeDSB. Note this would be expected to converge much slower, particularly for small epsilon, as fitting of the marginals only arises through iteratively learning time-reversals— if the forward processes rarely reach the other marginals (as is the case for smaller epsilon) this is very slow.

Theoretical assumptions

The key assumptions for the Brownian case are that $\int \lVert x \rVert^2 d \mu_i < \infty$ and $H(\mu_i)< \infty$ for each marginal $i$ (see proof of Theorem 3.1). We will mention these in the main paper. More generally, the required assumptions are stated in Proposition B.3 (for a general reference measure $\mathbb{Q}$), as well as standard technical assumptions inherited from the DSBM paper [8, Appendix C.1].

We hope that the presence of experiments in the Appendix that are relevant to many of your questions, along with our comments and results above, and our comments in other rebuttals, have addressed your concerns. If you have any further questions, please let us know!

References

[1] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation

[2] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization

[3] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport

[4] Fan et al, 2021, Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks

[5] Haasler et al, 2021, Multi-marginal Optimal Transport with a Tree-structured cost and the Schrödinger Bridge Problem

[6] De Bortoli et al, 2024, Schrödinger Bridge Flow for Unpaired Data Translation

[7] von Lindheim, 2022, Simple Approximative Algorithms for Free-Support Wasserstein Barycenters

[8] Shi et al, 2023, Diffusion Schrödinger Bridge Matching

We thank the reviewer for their positive and constructive comments, and we address their comments and questions below. We first provide a brief rebuttal summary:

• 'Overly technical presentation': We propose to:
  • add a paragraph to the introduction regarding the importance of the Wasserstein barycentre problem.
  • add a paragraph to the introduction emphasising the fixed-point interpretation, which we anticipate may be more intuitive for a reader less familiar with related works.
  • add an explicit description of the algorithm in the simplest $\frac{1}{K}$-barycentre case.
• 'Challenges in training/hyperparameters': We address each of these in turn below, with reference to relevant experiments in the Appendix, and with new ablations on the role of batch size and $\varepsilon$.

We expand on these points and the reviewer's questions in more detail below, and highlight any proposed changes for revisions in bold.

"Overly technical presentation"

We thank the reader for highlighting this, as (behind the mathematical notation) we believe our algorithm is in fact quite intuitive, so are keen to convey this to the reader.

"why it is even of interest to look at these problems"

We propose to add a paragraph in the introduction explaining the importance of the Wasserstein-2 barycentre problem. In particular, we will highlight its interpretation as the natural notion of ‘mean’ for probability distributions, and its applications in Bayesian learning, clustering, and representation learning to name a few.

"overly focuses on technical aspects and relevant methods", "Can authors provide a simple example"

Currently the narrative of the paper focuses on TreeDSBM being an IMF counterpart to TreeDSB, which we agree may be somewhat opaque for a reader unfamiliar with these works. Our approach was in fact initially motivated by trying to combine IMF with fixed-point approaches for barycentre computation (please see our response to Reviewer 3rwT). We believe that the fixed-point perspective — namely that each point in the 'average' measure, is the 'average' of its corresponding points in the marginals — may be more intuitive for a reader less familiar with related work.

While discussed in the paper, we believe a less familiar reader would benefit from having this interpretation highlighed more prominently. We thus propose to add a paragraph in the introduction highlighting this motivation and interpretation, and to add a fully explicit description of Algorithm 1 in the simplified case of barycentres with weights 1/K. This is the simplest setting (as the barycentre is constructed via $Y_0$ centred at $\sum_i \frac{1}{K} Y_i$, i.e. the ‘average’ of its corresponding points in the marginals), and should be intuitive for the reader.

"why is the tree setting of special interest beyond barycenters"

Given our initial fixed-point motivation described above, the fact our approach extends to the tree setting was something of an added bonus that arose out of using the tree-based mathematical framework. Nevertheless, such problems have been utilised in works such as [1]. We conjecture that it could have applications in trajectory inference problems where populations might be known to 'split' according to a given structure (for example, maybe for single-cell dynamics). We unfortunately don't have the biological expertise to be confident in devising meaningful experiments for such settings, so instead provided a toy experiment in Appendix E.3 in which TreeDSBM appeared to work well, demonstrating fast and stable convergence. We have a brief paragraph discussing these wider problem settings in Appendix E.3, and propose extending this to include the above discussion, and referencing this paragraph at the end of the Related Work section. We hope future practitioners can build on our work for use in such settings.

Challenges in training/hyperparameters

Scaling with number of datapoints/dim:

• As it uses bridge-matching (trained via gradient descent on minibatches), our approach inherits the same scalability with datapoints as other flow-matching style algorithms.
• In terms of dimension, we included experiments for high-dimensional Gaussians in Appendix E.2 for $d=64, 96, 128$, and observed no significant performance drop as $d$ increased.

Instability:

• We did not encounter instability issues, which we attribute to using stable bridge-matching losses, and an iterative procedure rather than adversarial losses.
• "demonstrate the stability claims": In the Appendix, we have also compared to a range of state-of-the-art solvers [2],[3],[4]. Regarding stability, we found the solvers [4] (adversarial loss) and [3] (single regularised loss) to fail on the 2d-experiment. We discuss such issues in lines 935-940 in the Appendix.
• We propose adding a sentence referencing this in the main body.

Error in marginals:

• We believe the main algorithmic limitation is that errors can accumulate in the marginals as IMF iterations proceed; this is a limitation inherited from standard IMF. We discuss this in lines 249-251 in the main paper, and in lines 851-858 in the Appendix. Rotating the starting marginal, as in [5], helps mitigate this.
• It is therefore important that the bridge-matching steps fit the marginals accurately, which is affected by the hyperparameters training steps, batch size, network width. The effect of these parameters on bridge-matching is well-understood and they act in the expected ways. To provide an indication of how the learned bridge-matching quality affects the overall solution, we provide an ablation for batch size below on the 2d and data aggregation experiments for the hyperparameters used in the paper (similar effects occur for increasing training steps).

• Network width: any sufficiently-expressive network will work similarly well. Fairly small MLP networks are sufficient for the pointcloud data experiments (i.e. 2$d$, subset aggregation, high-dim Gaussians).
• We propose to extend the discussion in Appendix lines 851-858, by incorporating the above discussion and ablations.

Role of $\varepsilon$

• This is a perennial question in entropic OT. Standard methods to choose $\varepsilon$ in commonly used OT libraries (e.g. POT, OTT-JAX—done in proportion to the costs) provide good guidance for choosing suitable values. We add a few additional comments below.
• Typically, one may want to use as small an epsilon as possible (for less entropic bias). One advantage of TreeDSBM over TreeDSB is that it allows for much smaller epsilon (TreeDSB does not converge for too-small epsilon as simulated trajectories do not reach the other marginals).
• There is a visualisation of the role of epsilon in Figure 5 in the Appendix, and also for two values of epsilon for the MNIST experiment in Figure 7. We propose to add quantitative results for varying epsilon to the Appendix. Results for the 2d and data aggregation experiments are shown below.

Neural Architectures and Assumptions

The theoretical analysis requires the SDE drifts to be locally Lipschitz, which is inherited from DSBM (for which one reason is to make use of known results regarding the well-posedness of the Doob’s $h$-transform). This is a weak regularity requirement, and is satisfied by standard neural architectures.

Beyond Quadratic Costs, non-Brownian Case

In its current form, the algorithm works only for the quadratic cost due to the Brownian reference measure (as with DSBM). We provide some thoughts regarding this below.

• While DSBM works only for the quadratic cost and $\sigma>0$, [6] shows that ReFlow (i.e. $\sigma=0$) converges to OT for a convex cost $c(x,y) = h(x-y)$ if one restricts the form of the drift, and uses a Bregman divergence loss. One could potentially replace the bridge-matching steps with this procedure. The barycentre update would have to be modified accordingly (likely from $\sum_i \lambda_i y_i$ to $argmin_z \sum_i c(z, y_i)$ [7], though this may not be available in closed form).
• Extensions of IMF to more general settings exist (e.g. state-dependent costs [8]). Incorporating such ideas into our framework is a promising direction.
• Non-Brownian reference measures: as with DSBM, our method supports any reference process for which conditional bridges can be sampled (e.g. linear drift processes, see Appendix B of [5]). We focus on the Brownian case in the main paper due to its connection to OT.

We hope that these comments have addressed your questions and concerns. If you have any further questions, please let us know!

References

[1] Haasler et al, 2021, Multi-marginal Optimal Transport with a Tree-structured cost and the Schrödinger Bridge Problem

[2] Korotin et al, 2022, Wasserstein Iterative Networks for Barycenter Estimation

[3] Korotin et al, 2021, Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization

[4] Kolesov et al, 2024, Estimating Barycenters of Distributions with Neural Optimal Transport

[5] Shi et al, 2023, Diffusion Schrödinger Bridge Matching

[6] Liu, 2022, Rectified Flow: A Marginal Preserving Approach to Optimal Transport

[7] Tanguy et al, 2025, Computing Barycentres of Measures for Generic Transport Costs

[8] Liu et al, 2023, Generalized Schrödinger Bridge Matching