Spherical Tree-Sliced Wasserstein Distance

TL;DR

This work introduces the Spherical Tree-Sliced Wasserstein (STSW) distance for comparing probability measures on the sphere by extending tree-sliced OT to spherical geometry. It develops spherical trees and a Spherical Radon Transform on these trees, establishes injectivity with orthogonal-invariance, and defines STSW as an integral over tree-structured projections with a closed-form, Monte Carlo computable approximation. The approach yields a metric that respects sphere symmetries and improves computational efficiency, demonstrated through gradient flows, self-supervised learning, Earth-density estimation, and SWAE experiments where STSW often outperforms recent spherical Wasserstein variants. The framework opens avenues for richer topological mass-transport representations on manifolds and scalable, parallelizable OT computations on spherical domains.

Abstract

Sliced Optimal Transport (OT) simplifies the OT problem in high-dimensional spaces by projecting supports of input measures onto one-dimensional lines and 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 problems for measures supported on a sphere. As a 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 and 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.

Figures (15)

• Figure 1: Illustrations of stereographic projection, rays in $\mathbb{R}^3$, and spherical rays in $\mathbb{S}^2$.
• Figure 2: Spherical Tree
• Figure 3: Radon Transform
• Figure 5: Runtime Comparison, averaged over 15 runs.
• Figure 6: Evolution between vMF($\mu, \kappa$) and vMF($\cdot, 0)$) w.r.t $\kappa$ on $\mathbb{S}^{d - 1}$ across various methods. We use $\kappa \in \{1, 5, 10, 20, 30, 40, 50, 75, 100, 150, 200, 250\}$

Theorems & Definitions (16)

• Remark
• Definition 3.1: Spherical rays in $\mathbb{R}^{d+1}$
• Definition 3.2: Spherical Trees in $\mathbb{S}^d$
• Theorem 3.3: Spherical trees are metric spaces with tree metric
• Definition 4.1: Spherical Radon Transform on Spherical Trees
• Definition 4.2
• Theorem 4.3
• Remark
• Definition 5.1: Spherical Tree-Sliced Wasserstein Distance
• Remark