Linear Spherical Sliced Optimal Transport:
A Fast Metric for Comparing Spherical Data


ICLR 2025 (Spotlight)
MY ALT TEXT

The proposed Linear Spherical Sliced Optimal Transport framework (LSSOT) projects the spherical distributions onto great circles $\mathbb{S}^1$, and then embeds circular distributions into $L^2$ spaces. The LSSOT distance is excellent in preserving the spherical geometry.

Abstract

Efficient comparison of spherical probability distributions becomes important in fields such as computer vision, geosciences, and medicine. Sliced optimal transport distances, such as spherical and stereographic spherical sliced Wasserstein distances, have recently been developed to address this need. These methods reduce the computational burden of optimal transport by slicing hyperspheres into one-dimensional projections, i.e., lines or circles. Concurrently, linear optimal transport has been proposed to embed distributions into $L^2$ spaces, where the $L^2$ distance approximates the optimal transport distance, thereby simplifying comparisons across multiple distributions. In this work, we introduce the Linear Spherical Sliced Optimal Transport (LSSOT) framework, which utilizes slicing to embed spherical distributions into $L^2$ spaces while preserving their intrinsic geometry, offering a computationally efficient metric for spherical probability measures. We establish the metricity of LSSOT and demonstrate its superior computational efficiency in applications such as cortical surface registration, 3D point cloud interpolation via gradient flow, and shape embedding. Our results demonstrate the significant computational benefits and high accuracy of LSSOT in these applications.

Contributions

  • We propose Linear Spherical Sliced Optimal Transport (LSSOT) to embed spherical distributions into $L^2$ space while preserving their intrinsic spherical geometry.
  • We prove that LSSOT defines a metric (LSSOT distance), and demonstrate the superior computation efficiency over other baseline metrics.
  • We conduct a comprehensive set of experiments to show the effectiveness and efficiency of LSSOT in diverse applications, from point cloud analysis to cortical surface registration.
  • Computation Efficiency

    MY ALT TEXT

    Figure: pairwise distances runtime (log scale) comparison w.r.t the number of distributions. $N$ denotes sample sizes in each distribution. The number of slices is 500 for all slice-based methods.

    Experiments

    Cortical Surface Registration

    We verify the validity and efficiency of the proposed LSSOT method in comparing spherical cerebral cortex data for registration tasks. Cortical surface registration seeks to establish meaningful anatomical correspondences across subjects or time points.
    We leverage the Superfast Spherical Surface Registration (S3Reg) [1] neural network to perform atlas-based registration, which focuses on registering all surfaces to one atlas surface. We replace the original mean squared error (MSE) similarity loss with our LSSOT distance and other baselines, and evaluate the registration performance achieved by each similarity measure. Baselines include Spherical Sliced-Wasserstein (SSW) [2] and Stereographic Spherical Sliced Wasserstein Distances (S3W) [3] and the Sliced Wasserstein distance (SWD) [4] under Euclidean geometry.

    MY ALT TEXT

    Figure: qualitative registration results (middle columns) from a moving surface (left column) to the fixed surface (right column).

    MY ALT TEXT

    Table: evaluation metrics on the test datasets (mean $\pm$ standard deviation) for sulcal depth registration of NKI dataset (left hemispheres). The best two methods are bolded. The average training time for each epoch is also included in the bottom row of each scenario. We perform significance testing for the difference on Dice, Edge Dist. and Area Dist. against LSSOT. $^{**}$ denotes $p<0.001$ and $^{*}$ denotes $p<0.01$.

    Point Cloud Interpolation

    LSSOT can also be utilized in point cloud analysis, once each point cloud is endowed with a spherical representation. In this experiment, we explore the interpolations between point cloud pairs from the ModelNet dataset. Specifically, we train an autoencoder to project the original point clouds to a spherical latent space, to represent each of them as a spherical distribution. Then we apply gradient flow between the pairs of spherical distributions using the LSSOT metric along with SSW and Spherical OT. Finally, this transformation is reconstructed in the original space by the trained decoder, resulting in an interpolation between the original pairs.

    MY ALT TEXT

    Figure: Top panel: the interpolation process between a pair of point clouds using gradient flow in the latent space $\mathbb{S}^2$. Bottom panel: gradient flow interpolations from a range hood to a bottle (left), and from a stool to a chair (right) using three metrics LSSOT, SSW and spherical OT.

    Poster

    References


    [1] Zhao, Fenqiang, Zhengwang Wu, Fan Wang, Weili Lin, Shunren Xia, Dinggang Shen, Li Wang, and Gang Li. "S3Reg: superfast spherical surface registration based on deep learning." IEEE transactions on medical imaging 40, no. 8 (2021): 1964-1976.

    [2] Bonet, Clément, Paul Berg, Nicolas Courty, François Septier, Lucas Drumetz, and Minh Tan Pham. "Spherical Sliced-Wasserstein." In The Eleventh International Conference on Learning Representations. (2024)

    [3] Tran, Huy, Yikun Bai, Abihith Kothapalli, Ashkan Shahbazi, Xinran Liu, Rocio P. Diaz Martin, and Soheil Kolouri. "Stereographic Spherical Sliced Wasserstein Distances." In International Conference on Machine Learning, pp. 48494-48564. PMLR, 2024.

    [4] Bonneel, Nicolas, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. "Sliced and radon wasserstein barycenters of measures." Journal of Mathematical Imaging and Vision 51 (2015): 22-45.