Projection Optimal Transport on Tree-Ordered Lines

Authors: Hoang V. Tran, Huyen Trang Pham, Tho Tran Huu, Minh-Khoi Nguyen-Nhat, Thanh Chu, Tam Le, Tan Minh Nguyen
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: projection optimal transport, optimal transport

Abstract:
Many variants of Optimal Transport (OT) have been developed to address its heavy computation. Among them, notably, Sliced Wasserstein (SW) is widely used for application domains by projecting the OT problem onto one-dimensional lines, and leveraging the closed-form expression of the univariate OT to reduce the computational burden. However, projecting measures onto low-dimensional spaces can lead to a loss of topological information. To mitigate this issue, in this work, we propose to replace one-dimensional lines with a more intricate structure, called \emph{tree systems}. This structure is metrizable by a tree metric, which yields a closed-form expression for OT problems on tree systems. We provide an extensive theoretical analysis to formally define tree systems with their topological properties, introduce the concept of splitting maps, which operate as the projection mechanism onto these structures, then finally propose a novel variant of Radon transform for tree systems and verify its injectivity. This framework leads to an efficient metric between measures, termed Tree-Sliced Wasserstein distance on Systems of Lines (TSW-SL). By conducting a variety of experiments on gradient flows, image style transfer, and generative models, we illustrate that our proposed approach performs favorably compared to SW and its variants.

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

Reviews and discussion:
<thread>
creation date:
1730021617231

thread_type:
official_review

<thread_object>
creator: reviewer_hzfv

<content>
summary:
This paper introduces a new tree Sliced-Wasserstein distance. This new distance is obtained by projecting on trees instead of 1d lines. The authors provide a complete theoretical analysis of the tree structures introduced, of the Radon transform used to project on them, as well as a discussion on how to sample such trees. Finally, they compare their method with the original Sliced-Wasserstein distance and some of their variants on several tasks such as gradient flows, color transfer, generative modeling with SWGAN and for diffusion models.

strengths:
This paper provides a new tree Sliced-Wasserstein distance, which in contrast with previous works, project on trees, and is thus more flexible. 

The authors provide first a full theoretical analysis of the tree space introduced, as well as a way to sample them in practice. Defining the right Radon transform, this analysis allows defining naturally a Tree Sliced-Wasserstein distance, which generalizes the original Sliced-Wasserstein distance. They show that this is well a metric, and that it can be approximated efficiently in the same complexity as Sliced-Wasserstein (w.r.t samples). This construction is thus original, well motivated and with good properties.

Finally, they compare their distance with Sliced-Wasserstein and variants thereof, on different experiments. On these different tasks, they show consistently better performances than Sliced-Wasserstein. Moreover, their experiments are done on classical baselines such as gradient flows on 2D datasets and higher dimensional gaussians, or color transfer, but also on high dimensional data such as images with generative modeling.

weaknesses:
Some minor weaknesses are that there are no statistical analysis for the Tree Sliced-Wasserstein distance proposed, e.g. no sample complexity, nor topological analysis. One can wonder how it relates with other distances, it TSW-SL a lower-bound of the Sliced-Wasserstein or of the Wasserstein distance?

The theoretical section on trees is very interesting, but some parts are challenging to read (for instance Section 3.2). I note however that an effort has been made on being clear, notably through Figure 1 and 2, which help to understand the constructions.

questions:
When citing references about work which project on lower dimensional subspaces, some references are missing such as [1,2,3].

I think in Equation (2), there is an abuse of notation as the Radon transform of a density is a density, and not a measure.

The construction given in Algorithm 1 only samples a chain tree. Did you also test with sampling trees where nodes have different childs? 

For the experiments on SW generators (Section 6.3), the number of projections (50 and 500) seems very low compared to the original paper, where I think they use something like 100K projections.


Typos:
- Line 102: "(Helgason & Helgason, 2011)" -> (Helgason, 2011)
- Line 322: "If tree systems in $\mathbb{T}$ consists only one line"


[1] Lin, T., Zheng, Z., Chen, E., Cuturi, M., & Jordan, M. I. (2021, March). On projection robust optimal transport: Sample complexity and model misspecification. In International Conference on Artificial Intelligence and Statistics (pp. 262-270). PMLR.

[2] Huang, M., Ma, S., & Lai, L. (2021, July). A riemannian block coordinate descent method for computing the projection robust wasserstein distance. In International Conference on Machine Learning (pp. 4446-4455). PMLR.

[3] Muzellec, B., & Cuturi, M. (2019). Subspace detours: Building transport plans that are optimal on subspace projections. Advances in Neural Information Processing Systems, 32.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730289273747

thread_type:
official_review

<thread_object>
creator: reviewer_7kaa

<content>
summary:
The paper introduces a variant of the Sliced-Wasserstein distance that, rather than projecting distributions onto lines, projects them onto tree systems. This approach leverages the closed-form solution of optimal transport (OT) for trees to develop a fast algorithm based on the Radon transform applied to tree structures. It uses splitting maps to define how distributions are projected onto these structures. This tree-sliced Wasserstein distance on a System of Lines (TSW-SL) is shown to be a metric. Experiments compare the original SW with TSW-SL, demonstrating that TSW-SL provides slightly better performance than SW, with comparable computation times.

strengths:
Paper introduces a new variant of sliced-Wasserstein: relying on a tree metric is original and allows benefiting from closed-form solution that exists, allowing preserving computational efficiency. The paper is clear, well illustrated, experimental section validates the approach in different contexts.

weaknesses:
**Note:** I have reviewed an extension of TSW-SL for the same venue.

As noted in the conclusion, the paper *"introduces a straightforward alternative to SW by replacing one-dimensional lines with tree systems,"* aiming to provide a more geometrically meaningful space. This objective aligns with several SW variants (as mentioned in the related work section of the introduction) that have achieved performance improvements in learning tasks involving SW and/or offer alternative schemes with distinct properties (e.g., statistical characteristics, behavior in high-dimensional spaces, the ability to provide a transport plan). However, given that TSW-SL achieves similar (or slightly better) performance to SW while retaining the same limitations (e.g., difficulties in sampling meaningful lines in high dimensions, as noted in Table 2), it remains unclear why and under what circumstances TSW-SL should be preferred over SW or its many variants.

There is also no discussion regarding the impact of the number of lines, which seems to be a critical aspect of the method—especially as using only one line recovers the standard SW. Additionally, the influence of the splitting maps is not addressed, despite being central to the new Radon transform.

**Other Comments:**

- In the gradient flow experiment, differences in ground cost (e.g., $L_2$, tree metric) between methods make it challenging to compare results at a fixed number of iterations.
- Table 1: By *Wasserstein distance,* do you mean $W_2$? Additionally, could you clarify the timings reported in Table 2?
- Line 467: The loss should not be denoted as $\mathcal{L}$ since this letter represents a line.
- Section 6.2: It is difficult to visually confirm that *TSW-SL produces images that most closely resemble the target.*
- Appendix p.17, line 870: $\overline{L}$ should be $\overline{\mathcal{L}}$.
- Example A.14 is unclear: $n_i $ should have length 6 (i.e.,$i = 0$ to 5). What do the arrows signify? Additionally, an arrow seems to be missing on line 1035. The progression through step $i$ is confusing.
- Page 25: Please provide more detail about the transition from eq. (46) to eq. (47).

questions:
- What are the main arguments in favor of TSW-SL wrt other many variants of SW? Could you highlight the unique advantages of TSW-SL?

- Can you give some stronger argument to the claim that tree systems allow avoiding a loss of topological information? Is this claim related with the number of lines that should be sampled ? What is the impact of this parameter onto the perfomances?

- Is it possible to derive additional properties of TSW-SL? For instance, does it metrizes the weak convergence? What are its statistical properties? 

- It is stated in the conclusion that improved performances are expected by adapting recent advanced sampling schemes or techniques of SW to TSW-SL. It seems that the extension is not that straightforward: can you comment on that, and how improved performances can be expected?

- It is stated in the beginning of section 6 that "*the root is positioned near the mean of the target distribution*": what is the impact of this setting?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730560062679

thread_type:
official_review

<thread_object>
creator: reviewer_mttt

<content>
summary:
The authors advance the concept of Sliced Wasserstein (SW) distance, which involves projecting data onto one-dimensional lines and leveraging closed-form expressions for one dimensional optimal transport to reduce computational costs. The authors propose an alternative projection onto tree systems, asserting that this approach preserves topological information. Additionally, they rigorously demonstrate that their framework yields an efficient metric for comparing measures. They also present extensive experimental results showcasing the superior performance of their method compared to existing methods based on SW-type distances.

strengths:
The derivation is clear and well-structured, providing a thoroughly developed framework. The authors demonstrate improved performance relative to previous related methods while maintaining lightweight computation. The code is provided.

weaknesses:
Although the authors present a well-founded theoretical framework and demonstrate practical improvements on several setups, the proposed method essentially introduces (yet) another distance metric akin to the Sliced Wasserstein distance. The overall paper seems to be a rather straightforward and incremental extension of the prior related works in the field. That is, I am rather skeptical about the significance and impact of the construction proposed here on the field of ML.

It seems to me that the main purpose of the proposed TSW-SL seems to be to compare the probability distributions, a feature which is primary needed for GAN training nowadays (the other experiments in the paper with gradient flows, etc. are toy, so I do not consider them as particularly demonstrative and serious). Here I have a general concern regarding the usefulness of the proposed TSW-SL in GANs. The experiments with GANs indeed show some practical improvement in the task of image generation. However, according to the famous work [1] tuning GANs is more about tuning various hyperparameters rather than changing the particular loss functions. This is also confirmed by the fact that most more practically-related papers (from CVPR, ECCV, etc.) still rely on vanilla/NS/WGAN loss with additional regularizations rather than other more complex losses (like SW-based) proposed by the community later. This suggests that (in 2024) the contribution of the current paper (TSW-SL as a GAN loss) may be relatively minor, so I am more on the negative side about the paper.

Additionally, the experiments with GANs lack details and sufficient explanations. For example, it is not clear why the problem in lines (Appendix) 1501-1505 is a valid GAN optimization problem. In the DDGAN experiment, how exactly the DDGAN loss is used in TSW is also not explained in detail, while it should: the DD-GAN loss is a non-trivial model due to various conditioning (discriminator is conditioned on the point, on the time moment). How this conditioning fits to the the proposed TSW-SL framework remains explained (line 1618 is not enough).

[1] Lucic, Mario, Karol Kurach, Marcin Michalski, Sylvain Gelly, and Olivier Bousquet. "Are gans created equal? a large-scale study." Advances in neural information processing systems 31 (2018).

questions:
I would like to see more experimental results and discussion regarding the GAN training.
- Could you explain how does your training time (time per iteration, per epoch, overall time to convergence) compares with the one from DDGAN and the other basedlines considered in the experiments.
- Could you please provide the convergence plots (FID as a function of epoch, e.g., once per several epochs) for your model vs. DDGAN vs. some other SW baseline? I would like to understand how stable is the overall training of your model compared with the baselines.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721809865

thread_type:
official_review

<thread_object>
creator: reviewer_tdy2

<content>
summary:
The authors proposed a computation of the Sliced Wasserstein (SW) distance using a novel projections on systems of lines and achieved improvements compared to the original SW approach.

strengths:
Clarity:
The paper is well-written and logically structured, making it easy to follow.

Originality:
The approach of efficiently solving sliced optimal transport using tree-based priors presents a novel and intriguing perspective.

Significance:
This research represents a crucial advancement in the development of more efficient closed-form tools for computing Wasserstein distances, which is an important area of study.

weaknesses:
Quality:
The poor evaluation is the main weakness of the paper. The authors provided several experiments, but only Tables 1 and 2 offer an expressive justification of the proposed method.

questions:
What should I observe in Figure 5? The images appear identical to me. More expressive examples of the color transfer are necessary.
Why did you consider your distance as the regularization for the GAN model? In my understanding, gradient flow methods from section 6.1 in higher dimensions would provide better evidence.
The paper demonstrates how the proposed distance can help decrease the distance to the target datasets. However, there is no evaluation showing that the computed distance is actually closer to the real Wasserstein distance provided. I think it is crucial to demonstrate this across a range of datasets.
Table 4 does not demonstrate that the proposed method actually provides better results in FID.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732184261315

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (1/2)

comment:
Dear AC and reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us improve the paper significantly.

We sincerely thank the reviewers for their valuable feedback and constructive suggestions. We are encouraged by the positive endorsements regarding the following aspects of our work:

1. The derivation is clear, well-structured, and provides a thoroughly developed framework. (Reviewers tDY2, MtTT, 7kaA, HzfV)

2. The proposed metric is indeed a valid metric and can be approximated efficiently via a closed-form expression with the same computational complexity as Sliced-Wasserstein (with respect to samples). (Reviewers tDY2, 7kaA, HzfV)

3. The proposed metric demonstrates consistently better performance compared to Sliced-Wasserstein in tasks such as Gradient Flows, Color Transfer, and Generative Models. (Reviewers tDY2, 7kaA, HzfV)

---

Below, we address some common points raised in the reviews:

**Q1. Key contribution of Tree-sliced framework and promising directions for future work.**

The tree-sliced framework in our paper is built on two key insights:

- Local Perspective: Each line in a tree system is treated similarly to a line in the Sliced Wasserstein (SW) framework. Splitting maps determine how the mass at each point is distributed across the lines, and the projection of these mass portions onto the lines is then processed in the same way as in SW. A significant challenge with this approach lies in verifying whether the injectivity of the corresponding Radon Transform is preserved, as this determines whether the proposed metric is a true metric or merely a pseudo-metric. However, we addressed this concern by providing a proof in the paper.

- Global Perspective: Tree structures and splitting maps establish connections between the lines, creating a cohesive system. This introduces a novel aspect compared to SW, enabling interaction and integration among the lines in a tree system. The Wasserstein distance can now be computed on this space with a closed-form expression, analogous to how lines are treated in SW.

From these insights, several follow-up directions for TSW-SL can be pursued, focusing on each perspective:

- Local Perspective: Consider, for example, the Generalized Sliced Wasserstein (GSW) distance [1]. In GSW, the SW framework is retained, but the projection mechanism is altered. Specifically, GSW generalizes the integration level set in the Radon Transform, replacing the level set defined by the inner product (representing orthogonal hyperplanes) with one defined by an arbitrary function. Similarly, in TSW-SL, which currently relies on the inner product, a framework could be developed that generalizes this by using an arbitrary function, offering new flexibility and applications.

- Global Perspective: More detailed analysis can aid in designing improved tree structures and splitting maps. For instance, splitting maps could be enhanced to account for both positional information from points and lines, rather than relying solely on lines as in the current implementation. This idea is central to an extension of the TSW-SL paper, which is referenced by Reviewer 7kaA and has also been submitted to this venue. It can be accessed [here](https://openreview.net/forum?id=OiQttMHwce). 

It is important to note that while these ideas might seem straightforward, their development is non-trivial. For example, a critical property that must be preserved in TSW-SL variants is the injectivity of the corresponding Radon Transform. For these reasons, we believe the paper should be evaluated by the development of tree systems and the corresponding Radon Transform aspects (*i.e., compare apple-to-apple*), rather than focusing on numerical comparisons with state-of-the-art methods.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732186707298

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
Incorporating comments and suggestions from reviewers, as well as some further empirical studies we believe informative, we summarize here the main changes in the revised paper:

- We added an additional example of the color-transfer task in **Appendix E.2** to clearly demonstrate the qualitative advantages of our TSW-SL loss over other SW-variant losses.  
- We added a new column to **Table 4** (Denoising Diffusion Model task) to report the training time per iteration for each method.  
- We included the time per iteration of gradient flow task for high-dimensional dataset in **Table 2**. 
- We included the FID plot over epochs for DDGAN using TSW-SL and SW to showcase the empirical advantages of our method in terms of convergence rate.  
- We included the explanation of our method (TSW-SL) into Denoising Diffusion Model in **Appendix E.4, line 1632.**
- We corrected typos in **lines 467** and **870**, clarified **Example A.14** (as noted by Reviewer *7kaA*), corrected the typo in **line 322** (as suggested by Reviewer *HzfV*), and added references [1], [2], and [3] in **line 57** based on the recommendations of Reviewer *HzfV*.  

---

**Reference.**

[1] Lin, T., Zheng, Z., Chen, E., Cuturi, M., & Jordan, M. I. (2021, March). On projection robust optimal transport: Sample complexity and model misspecification. In International Conference on Artificial Intelligence and Statistics (pp. 262-270). PMLR.

[2] Huang, M., Ma, S., & Lai, L. (2021, July). A riemannian block coordinate descent method for computing the projection robust wasserstein distance. In International Conference on Machine Learning (pp. 4446-4455). PMLR.

[3] Muzellec, B., & Cuturi, M. (2019). Subspace detours: Building transport plans that are optimal on subspace projections. Advances in Neural Information Processing Systems, 32.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732390920188

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Any Questions about Our Rebuttal?

comment:
Dear Reviewers,

We would like to thank you very much for your feedback, and we hope that our response addresses your previous concerns. In case you have not responded to our rebuttal so far, please feel free to let us know if you have any further comments on our work as the discussion phase is expected to conclude in the next few days. We would be more than happy to address any additional concerns from you.

Thank you again for spending time on the paper. We really appreciate that!

Best regards,

Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733160688848

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion Phase Summary

comment:
Dear Area Chair and Reviewers,

As the discussion phase comes to a close, we would like to summarize the key contributions of our paper.

In this work, we introduce the Tree-Sliced Wasserstein (TSW) distance, an extension of the well-known Sliced Wasserstein (SW) distance. Instead of relying on traditional one-dimensional lines, the TSW distance utilizes a more intricate integration domain, referred to as *tree systems*. In essence, tree systems are structures where dimensional lines are interconnected, behaving locally as lines.

This innovative framework paves the way for the introduction of the *Tree-Sliced Wasserstein on Systems of Lines* (TSW-SL) distance, facilitating the application of Tree-Sliced distance to **dynamic-support problems** (e.g., GANs and Diffusion models). Addressing such applications has been a significant open challenge since the concept of Tree-Slicing was introduced in NeurIPS 2019 ([1]). To our knowledge, no successful attempts have been made to tackle this problem until now. Our work addresses this gap and opens new avenues for applying TSW to these complex tasks.

Our experimental results demonstrate the feasibility of the proposed method. While this is a foundational study, we show that TSW-SL outperforms traditional SW under equivalent computational costs. It is important to note that our primary goal is to establish a theoretical framework for TSW rather than competing with state-of-the-art SW-based methods. Furthermore, since tree systems are locally composed of lines, we anticipate that techniques developed for SW can be adapted to tree systems, potentially offering further improvements. This potential is discussed in more detail in our General Response.

Finally, the concepts introduced in this work serve as the foundation for extensions of the TSW-SL framework, which have also been submitted to this venue. Two examples of these extensions can be found [here](https://openreview.net/forum?id=OiQttMHwce) (referenced by Reviewer 7kaA) and also [here](https://openreview.net/forum?id=FPQzXME9NK). In these submissions, the proposed methods demonstrate either superior performance or results comparable to state-of-the-art SW approaches in their respective tasks.

We hope this summary helps the Area Chair and Reviewers form an accurate evaluation of the theoretical contributions presented in our paper.

Best regards,

Authors

---

**Reference.**

[1] Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. Tree-sliced variants of Wasserstein distances. NeurIPS, 2019.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces the Tree-Sliced Wasserstein distance on Systems of Lines (TSW-SL), a novel Optimal Transport (OT) variant that replaces one-dimensional projections in Sliced Wasserstein (SW) with tree systems to better preserve topological information. It provides a theoretical framework for tree systems, including their topological properties, splitting maps as projection mechanisms, and a Radon transform for injectivity. Experiments on tasks like gradient flows, image style transfer, and generative models are shown to perform similarly or favorably over SW and its variants.

Reviews for this paper are mixed. All the reviewers agree that the paper present an original novel formulation of SW, that is theoretically consistent. Nonetheless, after considering the different elements of discussions during the rebuttal, I concur with reviewer MtTT and 7kaA that the experiments are not fully convincing to show the superiority of TSW-SL compared to other variants of SW, nor it is demonstrated in a ML setting where its usage would be beneficial. For this reason, I am enclined to give a reject option. I encourage the authors to strengthen their papers by finding a ML setting where a clear advantage is shown for TSW-SL, not only with regards to variants of SW, but also with other distributions divergences or distances.