Tree-Based Diffusion Schrödinger Bridge with Applications to Wasserstein Barycenters

Summary

The authors develop a novel method for approximating Wasserstein barycenters, which is based on diffusion models. This method combines ideas from the recently introduced Diffusion Schrodinger bridge and a framework for multi-marginal optimal transport with tree-structured costs. Thus, although the paper focuses on the barycenter problem (corresponding to a star shaped tree), the proposed approach can be used for any multi-marginal optimal transport problem, where the cost is composed of pairwise terms that are structured as a tree. A diffusion Schrodinger bridge type algorithm is presented that approximates the multimarginal iterative proportional fitting (mIPF) algorithm. The authors present novel convergence results of mIPF, which in contrast to previous works neither assume that the space is compact nor that the cost is bounded. Finally, the application to Wasserstein barycenters is evaluated for several experiments.

Strengths

I see the strengths of this paper in its originality and significance. The paper extends results for Diffusion Schrodinger bridges to a class of multi-marginal problems, including Wasserstein barycenters. Moreover, the paper extends results for tree structured multi-marginal optimal transport problems by developing a diffusion based method for these, which makes the problem scalable to high-dimensional settings.

Weaknesses

The clarity of presentation should be improved in many parts of the paper. Please see comments below.

Questions

My main question concerns the different update schemes mentioned in Section 3. The (mIPF) algorithm requires sampling from the whole tree. The authors argue that each (mIPF) update can be done on a path of the tree. In lines 160-173 it is described that one first samples along all edges of the path, and then learns the reverse dynamics for the path. In contrast, Algorithm 1 samples and learns the reverse dynamics for each edge on the path individually, and this is repeated iteratively along the path. That is, three different update schemes are mentioned: a) sample along whole tree and learn reverse dynamics; b) sample along whole path and learn reverse dynamics; c) iterate over the edges on the path, and for each edge sample along each and learn reverse dynamics. Please clarify why these are equivalent, especially the equivalence with scheme c) is not obvious to me. The following list is ordered mostly chronological, and not according to relevance. 1. In Section 2 the authors should make sure that the background material and setting are introduced in a clear and concise manner. This includes my following comments, but I suggest to carefully work through the whole section. - Are all the details on dynamic and static versions of Schrodinger bridges really needed? E.g., it is not clear what $\Phi$ and $a$ are, and I believe they are not used anywhere in the paper anyways, so this comment could be skipped. - The link between (EmOT) and Schrodinger bridges is not well explained. A tree is defined in the second sentence, but it is only explained how it is related to (TreeSB) in the last sentence. Also you probably need a reference to equation (1) somewhere? - The final paragraph "our framework" could be used to give a much more helpful summary of the main idea of this work. I suggest to state more clearly the motivation, contribution, and main result, and tone down the technical details (which are hard to understand at this point). 2. Proposition 3 should be presented in a self-sufficient way. It is not completely obvious here what $\mathbb{P}^n$ and $\pi^n$ denote. 3. What is the convergence criterion in Algorithm 1, line 5? 4. I feel that Section 4 does not really follow the flow of the paper, and I wonder why the authors don't stress the convergence result more as a main contribution of the paper. Also, I'm wondering if the result relies on any additional assumptions compared to previous results? Moreover, I'm a bit concerned about the fact that the one page of Section 4 relies on 10 pages of proofs, which are deferred to the appendix (and which I was not able to check). I find it problematic to have to hide 95% of the derivation for a strong theoretical result in the supplementary material. It might be preferable to not include this result in the paper due to space limitations, and publish it separately. This would also guarantee a thorough peer review of the result. In my opinion the paper makes a strong enough contribution without this Section. 5. For the synthetic two dimensional dataset it would be interesting to also see the true (possibly regularized) Wasserstein barycenter for reference. 6. For the MNIST application, I suggest to clarify in the main paper that you consider the problem in 784 dimensional space. This may not be clear, since images could be seen as distributions of 2-dimensional data, similar to the first numerical experiment. 7. What do you mean by one IPF cycle, and how do you chose the number of IPF cycles in the different experiments? 8. Since the paper presents methods for optimal transport problems with general tree structure, it would be very interesting to also include some experiments that are not a Wasserstein barycenter problem. 9. Typos: 43: hyperparameter and; 112: Section 5. We; 220 applied to trees

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

yes

Summary

Authors develop a new method for computing multi-marginal optimal transport in the case where the cost function is tree structured. Their algorithm cleverly exploits the dynamic formulation of the problem. The correctness of the algorithm is thoroughly demonstrated, and benefits are shown over a number of toy examples as well as competitive baselines involving the computation of the Wasserstein Barycenter.

Strengths

It is a solid paper (theoretically and experimentally). Although the mathematical part is mostly technical and algorithm-related (but correct me if I am wrong), it is still quite sophisticated and praise deserving. Results section is compelling.

Weaknesses

There are some weaknesses: first, the paper is slightly hard to read and while there are some helpful guides (as the first figure) at some points the paper gets helplessly hard to follow. Second, the paper should better explain what the contribution over the already existing methods (i.e Hasslear et al). It is not clear from the reading what this adds on top of that. Third, the paper is not compelling on its motivation. Why if the object of interest is a tree-based cost, all experiments focus on the barycenter problem? it gives the impression that the paper addresses a problem that doesn't exist in practice. Fourth, while the experimental results are super interesting, there is a lack of commentary as to why we should expect this.

Questions

1) What is the difference between your method and the one of Haessler et al? 2)Could authors spell out applications where multimarginal OT inference in tree structured costs is important? Why is the focus only on the Barycenter? 3)Authors are improving upon the competitive baselines such as Cuturi and Doucet (table 1). Why the method was expected to bring this benefit? is it related to the tree-ness or the fact that it implements a Schrodinger Bridge?.  4)Likewise, why better results are obtained in data fussion? (Table 2). It seems to be implied that this is due to better performance in high dimensions. Why this should be the case?  5)Line 175 seems to be an error in the Ext() statement. 6)Type in line 305

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

n.a.

Summary

This paper generalizes the diffusion Schrödinger bridge framework to multidimensional problems with tree-structured quadratic cost function. The method builds on representing each egde in the graph as two SDEs (one for each direction) parameterized by neural networks. These are used to propagate the information (from the known marginals) in the graph. Under a set of assumptions the method is shown to converge in TV.

Strengths

It is a nice extension of the Diffusion Schrödinger Bridge framework, which allows for addressing a relatively large class of multimarginal optimal transport problems where each marginal is a probability measure. Also it is nice that convergence of the method is provided.

Weaknesses

I think it would be nice with a little bit more information on the neural nets, about their structure and how they are used and updated. See the question below.

Questions

I do not see how the assumptions A1-A5 uses the neural nets. The neural nets typically have finite prespecified parameters. Can you explain how one can represent any measure (satisfying the assumptions), which are infinite dimensional objects.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations are adequately addressed.

Summary

The authors consider the multi-marginal optimal transport (mOT) with entropic regularization and a tree-structured quadratic cost. The authors propose to address it by developing the Tree-based Diffusion Schrodinger Bridge (TreeDSB), i.e., dynamic and continuous state-space of the multimarginal Sinkhorn algorithm. The authors apply the proposed approach to compute Wasserstein barycenter, i.e., equivalent to mOT with star-shaped tree.

Strengths

+ The authors propose Tree-based Diffusion Schrodinger Bridge, a dynamic and continuous state-space counterpart of the multimarginal Sinkhorn algorithm. + The authors illustrate some advantages of the proposed approach for Wasserstein barycenter.

Weaknesses

+ The motivation of developing the dynamic and continuous state-space counterpart of the multimarginal Sinkhorn algorithm and experiments seem quite weak. For Wasserstein barycenters and problems illustrated in the experiments, it suffices to address the multimarginal OT (i.e., the static problem). + Although there is some discussion in the appendix, it is essential that the submission is very closely related to (Hassler et al. 2021). It is an important part and should elaborate with details in the main manuscript to emphasize the contribution. + The notations seem heavy and hard to follow. It is better to simplify the notations or give more explanation for the notations when they are introduced. At the present, it is hard to follow the meaning of the notation system.

Questions

+ It is not clear why one needs to address the dynamic and continuous state-space of the multimarginal Sinkhorn. It suffices to consider the multimarginal OT directly for Wasserstein barycenters and problems in the experiments. Could the authors comment on this? + Although the authors give some discussion about the relation with Hassler et al. (2021) in the appendix, it is essential to clarify their difference in the main manuscript (e.g., in related work section) to emphasize the contribution. Could the authors summary their main differences? (since they are very closely similar). + In line 110, $S$ is the set of probability measures on the “leaves” of $T$. However, in line 118, why $r \in S$ (it seems that $r$ is the root of tree T as in line 113)? So, what does the tree $T$ look like? + Please explain the Figure 2, its description is too short and hard to follow. + As in line 157-158, there are $2|E|$ neural networks, how are these networks related? + Kitagawa and Pass (2015) showed the equivalent between multimarginal OT and the Wasserstein barycenter. Could the authors comment the result in Proposition 8 and the claim in line 271-272 with those in Kitagawa and Pass (2015)? Ref: Kitagawa, J. and Pass, B., 2015, September. The multi-marginal optimal partial transport problem. In Forum of Mathematics, Sigma (Vol. 3, p. e17). Cambridge University Press. + In line 284, how does the choice of $\epsilon$ affect the result, and how to choose it in applications? + For results in Tables 1, 2, Figures 6, 7, it is better to also report the corresponding time consumption? + It seems unclear how to choose $\epsilon$ (hyperparameters) for the barycenter problem? Does it affect the result? Some hyperparamters seem missing, e.g., how many free supports are used for fsWB? + For results in Figures 6 and 7, how to interpret these results? (especially Figure 7) + in line 330, “a scalable scheme”? it is unclear whether the proposed approach is scalable from the proposed algorithm and experiments? Please elaborate the scalability of the proposed method. Some minor points: + In line 2, what is $C([O, T], R^d)$? There is no explanation for the new notation $C( , )$ + in line 78, it is better to explain the Stochastic Differential Equation form, e.g., there is no explanation about several new notations, $X_t, a, B_t$ + Please explain the Figures. The description for each Figure is very short, it is unclear the meaning of each Figure, e.g., how to understand the Figures 2, 3? + In line 125, what is the role of $K$? + In Proposition 1, what is the relation between $T_{K-1}$ and $T$ (although there is some explanation in line 71-72, it is hard to follow it). It is better to consider to simplify the notations or explain the notations with more details. --- Thank you for the rebuttal.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have discussed the limitation of the work. However, there is no discussion about the potential negative societal impact of their work.

Summary

Schröninger bridges (SB) can be understood as dynamical counterpart of the Entropic Optimal Transport problem. This paper introduces the Tree Diffusion Schröninger Bridge (TDSB) model, which can be seen as a multi-marginal generalization of SB when the cost $c$ has a tree structure: given $\ell$ marginals $\pi_1,\dots,\pi_\ell$, and a tree $T = (V,E)$ with $|V| = \ell$ nodes, one has $c(x_1,\dots,x_\ell) = \sum_{(v,v') \in E} \| x_v - x_{v'}\|^2$. Let $S$ denote the leaves of $T$. In that setting, the goal of the paper is to solve the TreeSB problem: $$ \min_\pi \mathrm{KL}(\pi | \pi^0),\quad \pi_i = \mu_i,\ \forall i \in S. $$ where the $\mu_i$ are prescribed marginales, and $\pi^0 = \pi_r^0 \bigotimes_{(v,v')} \pi^0_{v,v'}$ is a reference measure that factorize along $T$, and $r \in V$ is identified as a root of $T$, and the factors $\pi^0_{v,v'}$ are chosen as normal distributions. If $T$ has a star-structure, this problem is equivalent to the (doubly-regularized) Wasserstein barycenter problem (in sense that, at optimality, the marginal of $\pi$ in the center of the star is the Wasserstein barycenter of the marginals on the leaves). The contribution of the work is to solve this problem using an iterative algorithm along _paths_ of $T$ which can be summarized as a succession of $\mathrm{KL}$ projections where only one of the marginals are constrained; if the tree is only made of two nodes (and one edge) we retrieve the Sinkhorn algorithm. In practice, the transition along a given edge (during one step of the algorithm) are performed by SDE of the form $\dot X_t = f_\theta(X_t) \mathrm{d}t + \mathrm{d} B_t$, where $B_t$ is the Brownian motion, and $f_\theta$ is a drift that is learned (and parametrized by a neural network). This leads to an iterative scheme where one repetitively select path between two leaves and learns the drift to solve the KL projection along the path. From a theoretical standpoint, authors prove that (assuming the intermediate step converge, e.g. the drifts $f_\theta$ end up being optimal) the proposed algorithm solves (asymptotically) the TDSB problem. From a numerical perspective, they showcase the practicability of their approach on synthetic and simple benchmark datasets.

Strengths

- I find the developed approach quite elegant. It mixes different notions in optimal transport literature (multi-maginal OT, entropic OT/Schrödinger bridge, doubly-regularized Wasserstein barycenters, parametrization of "transport maps" (the drifts) by neural networks, etc.). It is nice to see all these approaches implemented in a _practical_ setting (which can be implemented). - Comparison with previous work is well presented. - The convergence results of the proposed approach are strong.

Weaknesses

## Numerical experiments My main remark would be about the numerical experiments. They are quite limited in their comparison with state-of-the-art methods and rather acts like proof-of-concept. In details: - They do not showcase how the proposed approach solves new problems. Though reasonable, none of the results showcase feel "groundbreaking". This may just be a matter of exposition, but right now, loosely speaking, I am not "impressed" by the experimental results. For instance, the choice of synthetic datasets (Figure 4) is questionable in that the actual barycenter is not very intuitive and thus it is hard to visually tell whether the displayed output (Figure 5) is good or not; in addition, the "consistency" of the proposed approach is not striking (I admit that there is a decent similarity between the three outputs, but we cannot faithfully consider that the model did converge). - The comparison with SotA methods is incomplete; e.g. there is no running time involved. Given the (maybe only superficial) complexity of the proposed approach---such as the need for a possibly huge amount of NN (which are not very small according to Section G, as they are MLPs), one may be feared by a possibly huge running time (compared to simpler approaches like, say, the one of Cuturi and Doucet). If there is a computational tradeoff (better result for longer training time), it must be mentioned and discussed in details. - I would expect a comparison with the _Iterative Bregman Projection_ of (Benamou et al., 2015), given the "similarity" (in that they are iterative KL projections) with the proposed approach (and I would consider IPB as being "more SotA" than Cuturi&Doucet, for instance; for instance the IPB paper claims in Section 1.3 that "IPB converges (...) orders of magnitude faster than Cuturi&Doucet"). - A bit minor, but the description of the _subset posterior aggregation_ experiment is very minimalist. It is hard for the reader that is not already familiar with such experiment to have an idea of what is done there.

Questions

- My main question, which can be developed (see below) would be _why should one use this approach rather than existing ones?_ (e.g. IPB, fsWb, etc.). Of course, the paper gives some insight, but: 1. As said above, there is no insight on the computational burden of the proposed approach. 2. While I find very interesting that the approach works on any tree-structured coast, the work only focuses on Wasserstein barycenter, somewhat limiting its impact (that is, one strength of this work is that it solves a more general problem than previous Wasserstein-barycenter-dedicated methods). Giving an example + experiment where a (non-stared) tree cost is involved would be insightful (it would give a situation where the proposed approach as no challenger, as far as I know!).

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

authors have addressed the limitations of their approach. I do not identify ethical concern specific to this work.

Thanks

Thank you for taking time to answer my comments, along with the global answer with additional experiments. I am increasing my grade.