Spherical Tree-Sliced Wasserstein Distance

Authors: Hoang V. Tran, Thanh Chu, Minh-Khoi Nguyen-Nhat, Huyen Trang Pham, Tam Le, Tan Minh Nguyen
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: tree-sliced wasserstein distance, spherical optimal transport, equivariance

Abstract:
Sliced Optimal Transport (OT) simplifies the OT problem in high-dimensional spaces by projecting supports of input measures onto one-dimensional lines, then exploiting the closed-form expression of the univariate OT to reduce the computational burden of OT. Recently, the Tree-Sliced method has been introduced to replace these lines with more intricate structures, known as tree systems. This approach enhances the ability to capture topological information of integration domains in Sliced OT while maintaining low computational cost. Inspired by this approach, in this paper, we present an adaptation of tree systems on OT problem for measures supported on a sphere. As counterpart to the Radon transform variant on tree systems, we propose a novel spherical Radon transform, with a new integration domain called spherical trees. By leveraging this transform and exploiting the spherical tree structures, we derive closed-form expressions for OT problems on the sphere. Consequently, we obtain an efficient metric for measures on the sphere, named Spherical Tree-Sliced Wasserstein (STSW) distance. We provide an extensive theoretical analysis to demonstrate the topology of spherical trees, the well-definedness and injectivity of our Radon transform variant, which leads to an orthogonally invariant distance between spherical measures. Finally, we conduct a wide range of numerical experiments, including gradient flows and self-supervised learning, to assess the performance of our proposed metric, comparing it to recent benchmarks.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1729925546791

thread_type:
official_review

<thread_object>
creator: reviewer_gf54

<content>
summary:
The paper introduces a novel way to measure distances between probability distributions on hyperspheres, coined Spherical Tree-Sliced Wasserstein (STSW). The core technical contribution lies in adapting tree-based structures from [1] to work on spheres (spherical trees) and defining a new type of Radon transform for these structures. The authors prove this approach leads to closed-form solutions for optimal transport problems and show that STSW is a valid distance metric on $\mathcal{P}(\mathbb{S}^d)$. Through various experiments including gradient flows, density estimation, and self-supervised learning, they show that STSW can perform competitively with/better than the baselines while having faster runtime.

---

[1] Tran, Viet-Hoang, et al. "Tree-Sliced Wasserstein Distance on a System of Lines." arXiv preprint arXiv:2406.13725 (2024).

strengths:
S1. This method addresses an important problem of comparing distributions on the sphere, which has applications in many fields.

S2. Extending the tree-sliced concept to spherical domains is novel and non-trivial.

S3. The writing in the main paper is rigorous, coherent, and easy to follow.

S4. STSW is proved to be a proper metric (compared to SSW [2] which is only known to be pseudometric) with a novel splitting maps that preserve orthogonal invariance.

S5. STSW appears to outperforms various baselines in terms of runtime and other quantitative metrics.

S6. The algorithm is straightforward to implement.

---

[2] Bonet, Clément, et al. "Spherical sliced-wasserstein." arXiv preprint arXiv:2206.08780 (2022).

weaknesses:
W1. Sampling: In the algorithm, the authors propose sampling uniformly from R^{d+1} and then normalizing to get points on S^d, which does not produce a uniform distribution on the sphere. Would this induce a bias (and implications)? 

W2. Ablation: The paper would be strengthened if there are more insights provided via ablations on different design choices (i.e., rays, trees). How does the current tree structure help capture the data better than existing methods? Are there limitations, theoretical issues, or numerical instability associated with different components of the method (i.e., S3W [3] has the north pole issue). Can the splitting maps be learned? etc.

W3. Runtime and Complexity: It would be nice to have an explicit discussion of the computational/memory complexity (this is aside from the information provided in Appendix B).

W4. Experiments: This may be a minor point, but setup and hyperparameters could be better documented for all methods. In addition, there is no comparison with Vertical SW in appropriate setups. For generative experiments, there are no samples, of quantitative measures of the quality of images (i.e. the FID score).

---

[3] Tran, Huy, et al. "Stereographic Spherical Sliced Wasserstein Distances." Forty-first International Conference on Machine Learning. 2024.

questions:
- What makes STSW run faster than S3W [3], and in some cases, SW?

- The tree structure is supposed to capture 'richer' topological information per the authors' claim. How does that translate to practical results? Have the authors explored different hyperparameters to confirm that the superior performance in these setups is due to the tree component of the method? Traditional trees in euclidean spaces often have a hierarchical structure; here, it appears that our design choice is not hierarchical? If so, then what are the concrete benefits?

- Are there relationships to the OT distance on the spheres?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730393735107

thread_type:
official_review

<thread_object>
creator: reviewer_mqxd

<content>
summary:
The paper proposes a variant of the Wasserstein distance for distributions defined on the hypersphere. The proposed metric, called the Spherical Tree-Sliced Wasserstein (STSW) distance, adapts the Tree-Sliced Wasserstein distance to be applicable to the hypersphere by defining a novel spherical Radon transform. The proposed metric is invariant to orthogonal transformations, and can be computed efficiently.

strengths:
* The paper is well-written. 
* The paper proposes an efficient extension of the Sliced Wasserstein distances applicable to hyperspheres and the construction of the Radon transform is well-motivated.
* The authors also show desirable properties of their proposal, such as invariance to orthogonal transforms.
* On the practical side, the authors propose an efficient way to compute such a metric and show its effectiveness and superior performance in various experiments.

weaknesses:
* The proposed metric is limited to hyperspheres. The paper's contributions appear to be incremental, primarily combining previously established concepts (tree systems and Sliced Wasserstein distances) and extending them to the sphere setting.
* One of the claimed contributions is a bit misleading. Specifically, the authors claim to derive a closed-form expression; however, its computation still relies on approximations: first, due to the need for sampling trees, and second, by considering only discrete distributions in the explanation of how it is computed. This should be clarified, as the current claim gives the impression that the distance can be computed exactly due to the closed-form expression.
* Line 65 states that the use of tree systems enhances the capture of topological information. However, it is not so clear to me why is this the case. An experiment demonstrating this advantage would be useful.

questions:
- **Impact uniformity loss**: How should we understand $STSW(z^A,\nu)$ in eq. 20?  $z^A$ is a single point, so I assume we are considering a Dirac distribution at $z^A$. Theoretically, the distance from a Dirac distribution to a uniform distribution in the sphere is constant regardless of where the Dirac distribution is placed, due to invariance to rotations, right?  Why does this term favour uniformity if it is a constant term? Or am I misunderstanding something?
- **Radon Transform measure preservation**: Why does the proposed Radon transformation in eq. 8 transform a probability distribution $\mu$ defined on $\mathbb{S}^d$ into a probability distribution defined on $\mathcal{T}$? This is mentioned in lines 332-333, but in line 268 it says that $||\mathcal{R}_{\mathcal{T}}^{\alpha}f||\leq||f||_1$.  So it does not immediately follow that the Radon transform preserves the measure.
- **STSW Computation on continuous measures**: In section 5 you explain how to compute STSW in practice, but it is assumed that the probability distributions are discrete.Is it possible to get a closed form analogous to that in eq. 19 for non-discrete distributions? 
- **Injectivity of the radon transform**: In Theorem 4.3 it is proved that if the splitting map is $\mathcal{O}(d+1)$-invariant, then the spherical Radon transform is invariant. What would be the consequences of using a non-injective spherical Radon transform? What structure might be missed?

**Minor & Typos:**  
- l. 170 "be **a** positive"
- l. 372 $\nu(x)=\sum_{j=1}^n$ 
- In l. 322 change notation of $\delta$ to other letter, in order to avoid possible confusion with the Dirac delta function.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730670690439

thread_type:
official_review

<thread_object>
creator: reviewer_k7wt

<content>
summary:
This paper introduces the Spherical Tree-Sliced Wasserstein (STSW) distance, a novel metric designed for optimal transport on spherical domains. The key innovation lies in the integration over spherical trees as the domain, rather than traditional one-dimensional lines or great semicircles used in existing spherical Sliced Wasserstein approaches. This change allows STSW to better capture the underlying topology of spherical data, offering closed-form solutions that enhance both performance and computational efficiency.

The authors introduce a variant of the spherical Radon Transform tailored for spherical trees and prove its injectivity. Defining the STSW in terms of this transform is essential for establishing the metric properties of the distance, including its invariance under the action of the orthogonal group.

strengths:
The paper is well-structured, presenting clear objectives and a comprehensive review of related work. 
The efficiency of the new metric is well presented in the experiments. 

Leveraging on the ideas presented in Tran et al. (2024b), as said before, this reviewer things that the key innovation of this article (over articles as Bonet et al. (2022) and Tran et al. (2024a)) lies in the integration over spherical trees for defining the new metric between spherical probability measures.

weaknesses:
The approach builds on previous work by Bonet et al. (2022) and Tran et al. (2024a), in the sense that uses the same hight-level ideas. 
However, while the research incrementally follows the line of previous studies by Bonet et al. (2022), Tran et al. (2024a), and Tran et al. (2024b), it offers meaningful advancements by developing a metric specifically adapted for spherical data analysis.
Also, the experiments closely follows experiments previously presnted in papers as Bonet et al. (2022).

questions:
Besides the experimental comparisons of the new STSW with SW, SSW, and S3W variants, are there any analytic comparisons among them? Which are the differents in the topologies defined by these different approach?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718029129

thread_type:
official_review

<thread_object>
creator: reviewer_8auo

<content>
summary:
This paper introduces the Spherical Tree-Sliced Wasserstein (STSW) distance, an efficient optimal transport (OT) metric for measures on a spheres. By leveraging a novel spherical Radon transform that integrates over spherical tree structures, it provides closed-form OT solutions and maintains computational efficiency. Theoretical analysis and experiments, including gradient flows and self-supervised learning, confirm STSW’s effectiveness and its performance against recent benchmarks.

strengths:
* Paper is very well-written.
* Through extensive experiments, the effectiveness of the STSW has been investigated.
* Paper introduces a novel Radon transform for the measures on the spherical trees.

weaknesses:
* While STSW aims to be efficient, the spherical Radon transform and tree-slicing require considerable computation, especially as the number of edges or the dimension of the hypersphere increases. This could limit scalability for very high-dimensional or densely-sampled spherical data, impacting runtime in large-scale applications. Although I understand that the authors have provided extensive runtime comparisons, I would like to see the scalability of the method on higher-dimensional tasks beyond CIFAR, MNIST, and similar datasets.

* The effectiveness of STSW relies on the quality of sampled spherical trees, which introduces variability in metric accuracy. If the sampling process fails to capture diverse spherical structures adequately, STSW’s results might be inconsistent, especially in complex distributions where more refined tree structures are necessary. I would like to see which strategies (e.g., Markov Chains, Random Paths, etc.) could be applied here to sample more informative slices.

questions:
* I was wondering if incorporating a Markov chain over the distributions of the slices, instead of using a uniform distribution, could help in generating more informative tree-slices.

* I am interested in understanding the topology of the tree corresponding to the most informative slice in spherical trees and compare its effectiveness to the most informative slice in SSW, and essentially comparing them to the most informative slice in SWD. (By the most informative slice I mean Max-slice). This can be done on a chosen benchmark.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731403462520

thread_type:
official_review

<thread_object>
creator: reviewer_ocq3

<content>
summary:
The paper is a natural extension of sliced spherical OT to incorporate tree systems. The authors propose a topological space on spheres called spherical trees by connecting spherical rays with a common root. They then adapt Radon transform onto the spheres and slice the trees which is equivalent to slicing the spheres with trees. After showing that spherical tree are metric spaces, the authors followed classic approach of adapting the OT computation in the tree-sliced spheres. They provided comprehensive theoretical and empirical results to support their theories.

strengths:
The motivation is justified and the narrative follows standard ones when substituting simpler structures with more sophisticated variants for solving specific OT problems. Claims have been proved by theoretical results and verified by empirical results.

weaknesses:
While I didn't find major weaknesses in the paper, the authors didn't provide sufficient theoretical explanation in several critical places. The proposed STSW outperforms baselines in all most all metrics. First, why did the variances in Table 1 so small, smaller than one tenth of other methods in most lines, under Monte Carlo sampling? Why did STSW outperforms other baselines in terms of runtime? The theoretical complexity the authors provided cannot explain it. Why didn't the better runtime translate to a similar margin in reducing the training time in Table 2? Another point is that the authors attributes the better performances to "the ability to capture topological information of integration
domain" of STSW but they didn't show a direct connection in the paper. The results from the CIFAR dataset are not that impressive and the visualization didn't help either in explaining them.

questions:
419: "STSW outperforms the baselines in all metrics and achieves faster convergence." Why is that? Is this result theoretically predictable? 

466: "We also conduct experiments with $d=2$ to visualize learned representations." Why don't we directly project the learned features in the original image space to a sphere, rather than redoing the experiments on the sphere?

443: The variances from STSW are quite small. What's the explanation of that? Why doing tree-slicing on spheres is more robust (against different sampling?)?

We have seen spherical trees in the following article, although it's for a different application and the construction of the trees is different as well. Is there any connection between this work and that? 

Meng, Yu, Yunyi Zhang, Jiaxin Huang, Yu Zhang, Chao Zhang, and Jiawei Han. "Hierarchical topic mining via joint spherical tree and text embedding." In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 1908-1917. 2020.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732188442614

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

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 paper is well-written, well-structured, and the writing is rigorous, coherent, and easy to follow. (Reviewer Gf54, mQXD, k7Wt, 8AUo, oCQ3)

2.  The idea of proposing an efficient extension of Sliced Wasserstein distances to hyperspheres, with a well-motivated Radon transform is both novel and non-trivial, addressing the important problem of comparing distributions on spheres with applications across various fields; Furthermore, the proposed distance is rigorously shown to be a proper metric with orthogonal invariance. (Reviewer Gf54, mQXD, k7Wt, 8AUo, oCQ3)

3. The proposed distance is shown to be effective and efficient, outperforming baselines in runtime and quantitative metrics through extensive experiments. (Reviewer Gf54, mQXD, k7Wt, 8AUo, oCQ3)

4. The algorithm is straightforward to implement. (Reviewer Gf54)

---

Below, we address a common question raised in the reviews:

**Q1. Tree structure enhances the capture of topological information.**

**Answer.** The motivation for our paper comes from a simple yet intriguing idea: In the framework of Sliced Wasserstein (SW), a probability distribution on $\mathbb{R}^d$ is pushed forward onto a line. This raises the question: what does the resulting distribution reveal about the original one? It is evident that distinct distributions, when projected onto the same line, can become indistinguishable. 

The situation is similar in the spherical setting. Given, for example, vertical SW, where each slice corresponds to a great semicircle, after rotating a spherical distribution around the diameter of a slice, the projected distribution of on that slice remains unchanged. This means that distinct distributions can become indistinguishable when projected onto the same slice. However, with spherical trees that include more than one great semicircle, the splitting map comes into play. It allows for differentiating distributions that are otherwise indistinguishable under rotation. As a result, two distributions that vertical SW cannot distinguish due to rotational symmetry can now be separated using the spherical tree structure in STSW.

In summary, with the same number of edges (as vertical SW), and thus the same computational cost, spherical trees in STSW provide a significantly deeper understanding of probability distributions compared to individual edges as in vertical SW.

A natural question arises: if a better representation space is desired, why not replace one-dimensional manifold with higher-dimensional manifolds? The answer lies in computational feasibility. Optimal Transport in $\mathbb{R}^d$ for $d>1$ is computationally prohibitive due to the lack of a closed-form solution. In contrast, both vertical SW and STSW offer efficient closed-form expressions, making them more practical.

We believe this explanation sufficiently answers the question raised.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732203094982

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 clarified in **lines 83, 90, 327, 378, 531, 769, 1160** and **1162** that we derive a closed-form **approximation** for STSW (as noted by Reviewer *mQXD*).
- We changed $\delta$ to $\zeta$ to avoid confusion with the Dirac delta function (as suggested by Reviewer *mQXD*) in **lines 322, 1237, 1329, 1371** and **1495**.
- We corrected typos in **lines 170** and **372** (as suggested by Reviewer *mQXD*).
- We corrected the typo $x^{(l)}, y^{(l)}_1, \ldots, y^{(l)}_k \sim \mathcal{U}(\mathbb{R}^{d+1})$ to $x^{(l)}, y^{(l)}\_1, \ldots, y^{(l)}_k \sim \mathcal{N}(0, Id\_{d+1})$ in  **line 392** (as suggested by Reviewer *Gf54*).
- We added an additional experiment for the most informative sliced methods (as suggested by Reviewer *8AUo*), including MAX-STSW, MAX-SSW and MAX-SW in **Appendix B.3**. We report in **Table 5** and **Figure 11** the negative log-likelihood (NLL) and converged log 2-Wasserstein curves.
- We added an additional experiment on generative task using MNIST dataset (based on recommendations of Reviewer *Gf54*) in **Appendix B.6**. We included the FID scores in **Table 8** and generated images produced by the trained models in **Figure 14**.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces the Spherical Tree-Sliced Wasserstein (STSW) distance, a novel metric for optimal transport on spherical domains, which leverages a unique spherical Radon Transform and tree structures to achieve computational efficiency and enhanced topological representation. A key strength of the paper lies in its ability to adapt tree-sliced methods to hyperspherical distributions, offering both theoretical rigor and practical relevance. The reviewers identified several areas for improvement, including the need for deeper theoretical explanations for runtime and convergence performance, clearer connections between tree structures and topological information capture, and additional experiments on scalability to high-dimensional data. However, the authors addressed these concerns during the rebuttal period by clarifying key concepts, refining experimental sections, and demonstrating responsiveness to reviewer feedback. These updates, coupled with the novelty and practicality of the proposed method, make this paper a strong candidate for acceptance, with the potential to advance research in optimal transport on spherical domains.