Rigid-Invariant Sliced Wasserstein via Independent Embeddings

TL;DR

Rigid-Invariant Sliced Wasserstein via Independent Embeddings (RISWIE) is introduced, a scalable pseudometric that combines the invariance of NP-hard approaches with the efficiency of projection-based OT to achieve rigid invariance with near-linear complexity in the sample size.

Abstract

Comparing probability measures when their supports are related by an unknown rigid transformation is an important challenge in geometric data analysis, arising in shape matching and machine learning. Classical optimal transport (OT) distances, including Wasserstein and sliced Wasserstein, are sensitive to rotations and reflections, while Gromov-Wasserstein (GW) is invariant to isometries but computationally prohibitive for large datasets. We introduce \emph{Rigid-Invariant Sliced Wasserstein via Independent Embeddings} (RISWIE), a scalable pseudometric that combines the invariance of NP-hard approaches with the efficiency of projection-based OT. RISWIE utilizes data-adaptive bases and matches optimal signed permutations along axes according to distributional similarity to achieve rigid invariance with near-linear complexity in the sample size. We prove bounds relating RISWIE to GW in special cases and empirically demonstrate dimension-independent statistical stability. Our experiments on cellular imaging and 3D human meshes demonstrate that RISWIE outperforms GW in clustering tasks and discriminative capability while significantly reducing runtime.

Figures (15)

• Figure 1: RISWIE-PCA vs. OT: bias (left) and variance (right). RISWIE bias and variance do not become worse in higher dimensions. Ground-truth population distances are calculated with the Gaussian closed form that exists for both distances, and the empirical distances are calculated repeatedly and averaged across sampled distributions. The exponent $\alpha$ corresponds to the empirical decay rate in the log--log plot: we fit a power law of the form $A n^{-\alpha}$ to each curve (separately for bias and variance), which estimates the convergence rate.
• Figure 2: Runtime scaling with sample size $n$ for different distance metrics in $d=3$ (solid) and $d=64$ (dashed).
• Figure 3: 3D Example of RISWIE alignment. We illustrate how RISWIE aligns two point clouds by matching their marginal distributions along embedded axes. This method naturally extends to higher dimensions. For each axis of the anchor shape, we evaluate all possible pairings with axes of the target, including sign flips (reflections) to minimize the 1D Wasserstein cost. The second row shows the optimal axis matching determined by this process, and we show the poses overlaid with this alignment procedure.
• Figure 4: Matrix-build time versus number of points per mesh $n$ (log-scale). Markers show means across repeats; shaded ribbons are 95% CIs. Euclidean is fastest; RISWIE grows gently with $n$ and stays well below Sliced/OT, while Gromov--Wasserstein is the slowest by far.
• Figure 5: RISWIE Distance matrix for the HuBMAP tissue slices. Each block along the diagonal corresponds to slices from the same tissue stack. Within a block, RISWIE distances are consistently near zero, indicating strong invariance to small perturbations and local alignment of slices from the same sample. Across blocks, RISWIE captures larger geometric variation between tissues from different regions, producing higher inter-block distances.

Theorems & Definitions (19)

• Definition 1: Signed Permutation Group
• Definition 2: RISWIE Distance
• Theorem 1: Rigid-Invariance
• Theorem 2: Pseudometric
• Theorem 3: RISWIE Distance for Gaussians under PCA Embeddings
• Theorem 4: RISWIE–GW Comparison for Gaussians
• Definition 3: Soft RISWIE (SRISWIE) Distance
• proof : Proof of Theorem \ref{['thm:rigid_invariance']}
• proof : Proof of Theorem \ref{['thm:pseudometric']}
• proof : Proof of Theorem \ref{['thm:RISWIE_PCA']}