Abstract

Comparing spherical probability distributions is of great interest in various fields, including geology, medical domains, computer vision, and deep representation learning. The utility of optimal transport-based distances, such as the Wasserstein distance, for comparing probability measures has spurred active research in developing computationally efficient variations of these distances for spherical probability measures. This paper introduces a high-speed and highly parallelizable distance for comparing spherical measures using the stereographic projection and the generalized Radon transform, which we refer to as the Stereographic Spherical Sliced Wasserstein (S3W) distance. We carefully address the distance distortion caused by the stereographic projection and provide an extensive theoretical analysis of our proposed metric and its rotationally invariant variation. Finally, we evaluate the performance of the proposed metrics and compare them with recent baselines in terms of both speed and accuracy through a wide range of numerical studies, including gradient flows and self-supervised learning.

Depiction of stereographic projection $\phi: \mathbb{S}^2\backslash \{s_n\} \to \mathbb{R}^2$ (a), the stereographic Radon transform integration surfaces on the sphere, i.e., the level sets of $\langle \phi(x),\theta\rangle$ for a fixed $\theta\in \mathbb{R}^d$ (b), and the generalized stereographic Radon transform integration surfaces on the sphere, i.e. the level sets of $\langle h\circ\phi(x),\theta\rangle$ for a fixed $\theta\in \mathbb{R}^{d'}$.

Contributions

Introducing a computationally efficient transport distance, Stereographic Spherical Sliced Wasserstein distance (S3W), for spherical probability measures.
Providing a rotationally invariant variation of the proposed distance, Rotationally Invariant Stereographic Spherical Sliced Wasserstein distance (RI-S3W).
Offering theoretical analysis of the proposed distances.
Demonstrating the performance, both in terms of speed and accuracy, of the proposed distances in diverse applications, including gradient flows on the sphere, representation learning with Sliced-Wasserstein Auto-Encoders (SWAEs), spherical density estimation via normalizing flows, sliced-Wasserstein variational inference on the sphere, and self-supervised learning.

Runtime Comparison

Runtime comparison for $ARI\text{-}S3W_2$, $RI\text{-}S3W_2$, $S3W$, $SW_2$, $SSW_1$ (using level median formula), $SSW_2$ with binary search (BS) and antipodal closed form (for uniform distribution), Wasserstein, and Sinkhorn with geodesic distance as cost function.

Experimental Highlights

Gradient Flows on Sphere

Learning a mixture of $12$ von Mises–Fisher distributions. $ARI\text{-}S3W$ (30) has $30$ rotations, pool size of $1000$. $S3W$ variants use $\text{LR}=0.01$. $SSW$ has an additional $\text{LR}=0.05$ for better comparison. The plots show convergence of different distances w.r.t. iterations and runtime.

Self-Supervised Contrastive Learning

SSL latent space visualization on CIFAR-10 using $S3W$, $RI$-$S3W$, and $ARI$-$S3W$.

Sliced-Wasserstein Autoencoders

$S3W$-SWAE latent space visualization and reconstruction.

$RI$-$S3W$-SWAE latent space visualization and reconstruction.

$ARI$-$S3W$-SWAE latent space visualization and reconstruction.

Earth Density Estimation

Earth density estimation using normalizing flow with $S3W$ on Earthquake, Flood, and Fire datasets.

Earth density estimation using normalizing flow with $RI$-$S3W$ on Earthquake, Flood, and Fire datasets.

Earth density estimation using normalizing flow with $ARI$-$S3W$ on Earthquake, Flood, and Fire datasets.

BibTeX

Huy Tran*, Yikun Bai*, Abihith Kothapalli*, Ashkan Shahbazi, Xinran Liu, Rocio Diaz Martin, Soheil Kolouri
Stereographic Spherical Sliced Wasserstein Distances
International Conference on Machine Learning, 2024.

If you find our work useful, please consider citing us using the following BibTeX:

[BibTeX]

Acknowledgements

This template was originally made by Phillip Isola and Richard Zhang for a colorful ECCV project; the code can be found here.