Spherical Sliced-Wasserstein
Abstract: Many variants of the Wasserstein distance have been introduced to reduce its original computational burden. In particular the Sliced-Wasserstein distance (SW), which leverages one-dimensional projections for which a closed-form solution of the Wasserstein distance is available, has received a lot of interest. Yet, it is restricted to data living in Euclidean spaces, while the Wasserstein distance has been studied and used recently on manifolds. We focus more specifically on the sphere, for which we define a novel SW discrepancy, which we call spherical Sliced-Wasserstein, making a first step towards defining SW discrepancies on manifolds. Our construction is notably based on closed-form solutions of the Wasserstein distance on the circle, together with a new spherical Radon transform. Along with efficient algorithms and the corresponding implementations, we illustrate its properties in several machine learning use cases where spherical representations of data are at stake: sampling on the sphere, density estimation on real earth data or hyperspherical auto-encoders.
Summary
- The paper introduces SSW as an extension of the sliced-Wasserstein distance by adapting Wasserstein computations to spherical manifolds via geodesic projections.
- It presents a novel spherical Radon transform that integrates over half-circles, enabling efficient computation on hyperspherical data.
- Experiments demonstrate SSW's scalability and improved performance in tasks like density estimation on geophysical data and generative modeling.
Spherical Sliced-Wasserstein
Introduction
In the paper "Spherical Sliced-Wasserstein", the authors aim to extend the concept of the Sliced-Wasserstein (SW) distance, traditionally defined over Euclidean spaces, to spherical manifolds. This extension leverages the Wasserstein distance on the circle and introduces the spherical Radon transform, potentially broadening the applicability of SW distances to datasets that naturally reside on spherical domains, such as geophysical data and hyperspherical representations in machine learning models.
Background
Optimal transport (OT) has become a staple in machine learning, with the Wasserstein distance offering a robust metric for comparing probability measures. However, its computational expense has led to approximations like the SW distance, which reduces dimensionality through one-dimensional projections. The paper's innovation is to bring this reduction to spherical domains.
The SW distance on Euclidean spaces is formulated using one-dimensional projections wherein the Wasserstein distance can be efficiently computed. The authors propose using geodesic projections onto great circles for spherical manifolds to create the analogous Spherical Sliced-Wasserstein (SSW) distance. Notably, the computation relies on closed-form solutions for the Wasserstein distance on circles, aiding computational efficiency.
Spherical Sliced-Wasserstein Distance
The SSW distance is a measure defined over the hypersphere . It uses geodesic projections to map data onto great circles, then computes the Wasserstein distance over these projections. Formally, SSW is given by:
where denotes the geodesic projection onto a circle determined by , within the set of all such projections , and denotes the Wasserstein distance on these circles. This integration is performed over the Stiefel manifold with uniform measure .
Spherical Radon Transform
The authors introduce a novel spherical Radon transform which plays a critical role in defining SSW. This transform is key to understanding how the slicing approach is adapted to manifolds. The transform integrates functions over half-circles, a deviation from traditional hyperplane integration in Euclidean space, which reflects the spherical geometry.
Figure 1: Set of integration of the spherical Radon transform, illustrating the integration over spherical caps instead of planes.
Implementation and Complexity
The implementation of SSW on manifolds involves the QR decomposition for uniform sampling from the Stiefel manifold and efficient roll-outs to exploit the closed-form of the Wasserstein distance on circles. The computational complexity of SSW is advantageous compared to direct OT methods, particularly for large-scale spherical datasets.
The complexity estimate for using SSW is , where is the number of projections, highlighting scalability in terms of data size and dimensionality.
Experiments
The effectiveness of SSW is illustrated through applications in gradient flows for distribution approximation, density estimation on earth data, and generative modeling via autoencoders. These experiments demonstrate SSW's ability to capture complex distributions on spherical domains efficiently and more effectively than traditional SW.
Density Estimation on Earth Data
For example, a comparison of density estimation techniques using a mix of von Mises-Fisher (vMF) distributions and real-world geophysical data highlights SSW's computational efficiency and accuracy in modeling spherical distributions.
Figure 2: Density estimation of models trained on earth data, showing improved estimation accuracy using SSW methods.
Conclusion
The introduction of SSW as an extension of sliced-Wasserstein distances to the hypersphere opens new avenues for applying OT in machine learning. It maintains computational efficiency while addressing manifold constraints that Euclidean-based methods overlook. Future work could extend SSW to other non-Euclidean geometries, such as hyperbolic spaces and further explore its statistical properties, especially concerning uniformity and injectivity of the Radon transform on spheres.
Paper Prompts
Sign up for free to create and run prompts on this paper using GPT-5.
Top Community Prompts
Collections
Sign up for free to add this paper to one or more collections.