MLOT: Extending the Bipartite Structure towards Multi-Layered Structure for Optimal Transport

Authors: Liangliang Shi, Rui Tao, Junchi Yan
Primary area: optimization
Keywords: Optimal Transport, Multi-Layered Structure, Entropic Regularization, Sinkhorn Algorithm

Abstract:
Despite its remarkable success and widespread adoption in various domains, optimal transport (OT) has a rather simple structure, relying on bipartite graphs with only two layers of nodes for transportation. In this paper, we propose a multi-layered OT approach that extends the original two-layer structure to handle transportation problems across multiple hierarchical levels. Within this framework, the source distribution flows through intermediate layers, before reaching the target distribution. Unlike previous variants of OT that involve multiple distributions, our multi-layered OT typically involves uncertain intermediate distributions, which need to be computed based on the relationships between the preceding and succeeding distributions. Under entropic regularization, MLOT-Sinkhorn algorithm is further proposed for multi-layered OT, which can be accelerated using GPUs and significantly outperforms general solvers such as Gurobi. The theoretical results of our entropic MLOT are also given in this paper. In the experiments, we validate its speed advantage and convergence performance. We further validate its feasibility through Text-Image retrieval and intermediate image computing task, which demonstrates reformulating the problems as MLOT can achieve better results. Source code will be made available.

Decision: reject
Reviewer scores: 3, 3, 5, 6

Reviews and discussion:
<thread>
creation date:
1730087602611

thread_type:
official_review

<thread_object>
creator: reviewer_b2pr

<content>
summary:
This paper introduces a Multi-Layered Optimal Transport (MLOT) problem and proposes a Sinkhorn algorithm for solving this problem. The MLOT problem generalizes the standard OT problem by modeling transport from the source distribution to the target distribution through the multiple intermediate layers. The proposed Sinkhorn algorithm offers superior efficiency on complex problems while maintaining comparable accuracy to the LP solver, such as Gurobi. The proposed method is also evaluated on real-world tasks, such as Text-Image retrieval and intermediate image computation.

strengths:
- This paper introduces an interesting variant of the standard OT problem and demonstrates its practical utility by applying the MLOT framework to real-world tasks.
- This paper presents theoretical results on the convergence of the proposed algorithm.
- This paper is easy to follow.

weaknesses:
- Please refer to the questions section.

questions:
- Could you provide the wall-clock computation time for the Image-Text Retrieval task in Table 2? The proposed method is likely to be slower compared to the standard softmax operation.
- In the Image-Text Retrieval task (Fig 6), the two couplings (the coupling between Query image and Caption, and the coupling between Caption and Augmented image) can indicate different image-caption pairings. How coherent are these two couplings?
- In the Image-Text Retrieval task, the standard OT problem can be directly applied between the Query image and Caption using the same cost as in the MLOT, i.e., the negative cosine similarity between the normalized embeddings. Is there a meaningful difference in performance between the MLOT and OT frameworks (Table 2)?
- In the intermediate image experiments (Fig 7), how do the MLOT results differ from the standard OT barycenter?
- In the synthetic data experiments (Sec 4.1), the data points in the intermediate layers are also sampled, e.g., $(n\_{k})\_{k} = \\{3, 10, 20, 20, 10, 3\\}$ in Line 373. Could you provide the reason for sampling these intermediate samples? Is it just for setting the support of the intermediate distributions $a\_{k}$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730482490673

thread_type:
official_review

<thread_object>
creator: reviewer_prsw

<content>
summary:
The paper introduces multi-layered traditional optimal transport, which finds an optimal transport map between the input and output distribution through multiple layers as an extension of traditional optimal transport. Compared with other optimal transport methods involving multiple distributions, the authors introduce uncertain intermediate distributions computed based on the previous and the next distributions. Additionally, the authors further propose MLOT-Sinkorn, which is designed specifically for MLOT to accelerate computation. The experimental results show that MLOT achieved better results than previous optimization methods.

strengths:
1. This paper introduces a variant of multiple distributions optimal transport and a corresponding variant of the Sinkhorn algorithm to solve this problem.
2. The results of the method significantly outperform existing methods for multiple distributions.
3. The experimental section is thorough and clearly written, detailing the settings and configurations for each task.

weaknesses:
1. The motivation for the introduction of MLOT compared with other multiple distributions for optimal transport is not written. Please refer to Question 1 in the Questions section for more details.
2. The related works section doesn't mention any previous works on multiple distributions in optimal transport although in the abstract, the authors state that "Unlike previous variants of OT that involve multiple distributions,..." (L19)
3. Although the experimental results demonstrate significant improvements with MLOT over other optimization methods for multiple distributions, a comparison with other optimal transport (OT) methods of similar scope is missing. Specifically, the authors evaluated MLOT-Sinkhorn against non-OT methods, yet did not compare it with other OT-based approaches for finding optimal transport distances between multiple distributions.
4. The overall writing is vague and the flow of the paper is difficult to follow. The related works are non-relevant to the methods proposed. For example, in Section 2, the authors mention Optimal Transport on a Graph and state that "However, the above works rely on the shortest distances on the graph and do not directly compute the transport couplings in the graph. In this paper, we attempt to directly compute the transportation between nodes in a multi-layer structure.". However, the paper lacks further comparisons or references to this aspect in later sections.
5. Typo: In line 259, "As shown in Fig. ,", the figure is not specified.

questions:
1. My biggest question lies in the motivation for finding the OT distance between multiple distributions, which involves uncertain intermediate distributions, and only the source and the target distribution are known. As far as I understand, the motivation of the paper is from real-world problems, where finding the OT distance between two distributions is not enough, the case in which we assume that the source and the target distribution are known.
2. How to compute the intermediate distributions as mentioned in the paper?
3. Can you provide more baselines or related works of optimal transports in literature and explain clearly what is the difference between your methods and these methods?
4. As far as I understand, the authors of (1) propose Sobolev Transport, which aims to find the optimal transport distance between 2 distributions, in which supports of these two distributions are nodes of the graph, which is different from the main purpose of this paper to find optimal transport map between multiple distributions. What does it mean by "we attempt to directly compute the transportation between nodes in a multi-layer structure." (L135-136) since in (1), the nodes of one graph are support of two distributions?

**Remark**: I gave a score of 5, but I am open to raising it once the above concerns are addressed and I have a clearer understanding of the paper. I will also reconsider my score following the rebuttal to ensure my evaluation is fully justified.

(1) Tam Le, Truyen Nguyen, Dinh Phung, and Viet Anh Nguyen. Sobolev transport: A scalable metric for probability measures with graph metrics. In International Conference on Artificial Intelligence and Statistics, pp.9844–9868. PMLR, 2022.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714760451

thread_type:
official_review

<thread_object>
creator: reviewer_xovs

<content>
summary:
The work proposed multi-step optimal transport when finding the OT mapping from one distribution to other distribution, given the supports of "transitional" distributions are known (cost matrices are known). The method is presented in Section 3 with entropic regularization that leads to adapted version of Sinkhorn algorithm to solve the problem. Proposition 3 prove the convergence property of the method. Real applications are text-image retrieval task and generating intermediate images.

strengths:
Empirical results shown in text-image retrieval task, although the setting is quite simple, by creating layers based on flipping the image. 
I wonder if there is another work that use augmented data to do this task.

weaknesses:
The theory part is just straight adaptation from known work of Sinkhorn algorithm. In proposition 4, convergence rate depends on $\gamma$, which is not easy to quantify. 

For task of generating intermediate images, there is no baselines, or quantity to compare/assess the current method with others. This contribution is limited. 

It is not clear (theoretically, intuitively) that in text-image retrieval task, why do we need to use data augmentation to create another layer for OT?

questions:
Cost metric in the task text-image retrieval  could be negative, since it is defined as the cosine similarity, does this affect the definition and algorithm of and using OT?

In both real applications, is there any empirical rule to choose parameter?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731306150202

thread_type:
official_review

<thread_object>
creator: reviewer_mxuv

<content>
summary:
The authors propose to discretize the path of intermediate distributions between source and target distributions in the OT problem. They added up the total cost between each two adjacent transportation problems on the path and adopted the Sinkhorm iterations to alternatively solve for the intermediates and the couplings which ensemble the transportation plan. Convergence was provided and experiments showed performance gains in solving the text-image retrieval problem.

strengths:
The authors wrapped their MLOT formation with the Sinkhorn algorithm and thus inherited all the computational benefits of it. Theoretical results seem solid to me. Convergence analysis was provided and verified with synthetic data.

weaknesses:
The main problem I find in the paper is that the authors didn't answer the very question that they raised. In the second paragraph (line 039), they referenced transportation as the "real-world" example of the benefit of MLOT, by optimizing the freight transportation to reduce the cost and environmental impact, as suggested by their reference Bektas et al. 2019. However, throughout the paper, they didn't answer the question. Can MLOT help find the optimal routes and minimize the total cost? 

In real-world problems, the intermediate distributions have certain constraints like the routes (cost matrix) and max capacity of the stops (intermediate distributions) in freight transportation example. The paper doesn't mention such constraints either in methodology or experiments. 

In 4.2, MLOT was compared against softmax so it's not clear whether the gains were from OT or from the multi-layer formulation or both. 

Figure 7 in 4.3 also demonstrated that intermediate distributions may not be physically meaningful and smooth and thus not applicable to solving "real-world problems".

questions:
Line 323: Therefore, out MLOT can offer new perspectives and approximate computations for Dynamic OT and the Schrodinger bridge. What are the new perspectives? That MLOT can "approximate computations for Dynamic OT and the Schrodinger bridge" is a serious claim. I'd like the authors to provide evidence for that or tune down the sentence to something like whether we can approximate dynamic OT via MLOT is a future direction. And it's not clear in what aspect do we approximate. Is it the total cost or the transportation map or something else.

Comparing MLOT with barycenters should be expanded. In line 182, "when $S=2$ in Eq.6 and $K=3$ in Eq.5, the optimization of our MLOT is equivalent to solving the Wasserstein barycenter." It's not, because for that to be true you have to assume $\lambda = \frac{1}{S} = \frac{1}{2}$, don't you? As we change $\lambda$ in barycenters, we could also form a path of intermediate distributions between the target and the source. How is that process compared to MLOT?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733071378930

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response and Summary of Updates to Manuscript

comment:
We sincerely thank the reviewers and AC for investing their valuable time and efforts in this paper.

We thank the reviewers for noting that we introduce a variant of multiple distributions optimal transport and a corresponding variant of the Sinkhorn algorithm to solve this problem (MXuv, XoVS, Prsw, b2Pr), with solid (MXuv) theoretical analysis on the convergence property (XoVS, b2Pr), as well as experiments that demonstrates superior efficiency on complex problems while maintaining accuracy (Prsw, b2Pr). The experimental section is thorough and clearly written (Prsw), and the paper is easy to follow (b2Pr). 

First, we provide a high-level summary of the changes that we've made to the draft to address your feedback, and conclude with an overview of our key contributions, and how they differ from previous work.

**Summary of changes made to the draft**
* Revamped **introduction and related work section** to clarify how our MLOT formulation was motivated from real-world scenario, and how MLOT differs from other multi-distribution OT. (**MXuv, Prsw**)
* Added detailed discussion in **Appendix. G** about why MLOT can offer new perspectives and approximate computations for Dynamic OT and Schrödinger Bridge (**MXuv**).
* Added detailed experiment on Retrieval task in **Appendix. F Table 3**, recording the performance using traditional OT. Without formulating the problem as MLOT, it cannot outperform using softmax. (**MXuv, b2Pr**)
* Added additional experiment on Constraint-MLOT problem in **Appendix. F Table 6**, i.e. exists flow constraint on each intermediate distributions. The results demonstrate our MLOT-Sinkhorn algorithm's strong adaptability to such constrainted problem. (**MXuv, Prsw**)
* Added additional experiment on sythetic data with different shape of $(\mathbf{n}_k)_k$ in **Appendix. F Table 5**, to show that our MLOT-Sinkhorn can work on arbitary intermediate shape. (**b2Pr**)
* Added additional experiment result of comparison of two methods computing intermediate images: via Barycenter and via MLOT, in **Appendix. F Figure 8**. And further record their computing time, which shows our simplify in computing procedure and improvement in efficiency. (**Prsw, b2Pr**)

**Novelty and technical insight**

We would like to clarify that our proposed MLOT has different problem setting with previous multiple distribution OT, and ours were motivated by a more complicated real-world scenario compared to traditional OT. Details are summarized in the following table.
||Structure|Known Information|Transport Result|Handle Constraints while Maintaining Original Speed|Requirement for Shortest Path|
|-|-|-|-|-|-|
|**Traditional OT**|Bipartite distributions|Ground metric, source & target distributions|One $\mathbf{P}$, couples source and target|Yes|-|
|**Multi-Marginal OT**|Multiple layered graph, with multiple marginal distributions|Ground metric,all marginals' distributions|One coupling, simultaneously couples all distributions|Yes|No need|
|**Graph OT**|Arbitrary graph|Ground metric, original & final distributions|Flow $s$, transfers original distribution to final|No|Need|
|**Multi-Layered OT**|Multiple layered graph, with multiple marginal distributions|Ground metric, source & target distributions, **not include intermediates'**|A sequnce of couplings, couples each pair of adjacent distributions|Yes|No need|

We now summarize our key contributions:

1. **Algorithm solving MLOT with high efficiency and accuracy.** We propose MLOT-Sinkhorn, **prove its global convergence, verify its convergence, accuracy and high efficiency by a wide range of sythetic experiment**. Compared to previous OT-based method for solving MLOT, which requires firstly computing shortest path on graph based on ground metric. In contrast, our algorithm **compute optimal flow directly by ground metric**. What's more, our algorithm can **naturally adapt the constrainted-MLOT**, which Graph OT cannot handle.
2. **A Training-Free improvment for Zero-shot Image-Text Retrieval task.** Following OT-CLIP^[ICML24], which introduce graph structure and perform graph matching for classification/matching, we **reformulate the image-text retrieval task as MLOT problem**, by introducing data augmentation. MLOT can leverage two coupling's information, outperform methods using softmax or traditional OT. What's more important is that, differing from other works using data augmentation technique, our method is a **training-free improvement.**
3. **A new perspective to computing intermediate images.** When source and target in MLOT are each an image, the intermediate distributions will exactly be intermediate images. Traditionally, we compute barycenter for this task. The drawback of barycenter is that when multiple intermediate images are need, we have to compute several barycenter respectively. In contrast, **our method has significant time efficiency, which only requires a single computation process and returns all intermediate images.**
</content>
</thread_object>

</thread>


Metareview:
metareview:
In the paper, the authors proposed multi-layered optimal transport (OT) framework that generalizes the traditional two-layer OT structure to handle transportation challenges across multiple hierarchical levels. For the practical computation, the authors also proposed entropic regularization of the multi-layered OT in which the Sinkhorn can be used and accelerated using GPUs. Finally, convergence results are also provided and are verified under the synthetic settings.

After the rebuttal, there are still several weaknesses of the current paper: (1) The theoretical results are not novel, especially those with the Sinkhorn algorithm, as they are adapted directly from the known results. (2) The experiments lack baselines and are not extensive. It is not clear if the proposed method can outperform methods with similar scope. (3) The current writing of the paper is generally hard to follow. Several parts are not closely connected. (4) Finally, the motivation of the current paper is not convincing.

Given the above weaknesses of the paper, I recommend rejecting the paper at the current stage. The authors are encouraged to incorporate the suggestions and feedback of the reviewers into the revision of their manuscript.