Translation-invariant optimal transport distance

Many modern datasets represent the same object in different orientations, making direct comparison challenging. For example, two adjacent slices of tissue from the same region can appear rotated or reflected relative to each other after preparation. Standard tools for comparing data distributions, like the Wasserstein distance from optimal transport, are popular because they account for geometric structure. However, they incorrectly treat rigid motions like rotations, reflections, and translations as meaningful differences. Other methods, such as Gromov-Wasserstein, are invariant to such motions but scale poorly with sample size and dimension. Our project developed Rigid-Invariant Sliced Wasserstein via Independent Embeddings (RISWIE), a new pseudometric that efficiently compares high-dimensional probability measures while ignoring rigid misalignments.

 

RISWIE builds on the idea of Sliced Wasserstein, which projects data onto 1D directions and averages Wasserstein distances between the marginals. Although efficient, Sliced Wasserstein relies on random projections, which can miss structure in high dimensions, are sensitive to rigid misalignments, and lose theoretical guarantees when using a finite number of axes. We extend Sliced Wasserstein by (1) choosing data-dependent axes instead of random ones, (2) computing Wasserstein distances between projections onto different axes rather than the same, and (3) pairing these axes by solving an optimal matching problem that minimizes dissimilarity, including reflections. RISWIE constructs orthonormal bases $\{\varphi_j\}$ and $\{\psi_j\}$ for each dataset using geometry-adaptive embeddings such as PCA or diffusion maps. Each measure is pushed forward onto its respective basis, producing one-dimensional marginals $\{\alpha_j\}$ and $\{\beta_j\}$. RISWIE then finds the signed permutation $R$ that minimizes the average Wasserstein-2 cost between matched marginals $\{\alpha_j\}$ and $\{\beta_{R(j)}\}$. We proved that RISWIE is a pseudometric, invariant under rigid motions when the embedding spectra are simple, and derived a closed-form formula for Gaussian measures along with bounds relating it to Gromov-Wasserstein. We also established statistical stability, showing its bias and variance decay quickly and independently of dimension. In experiments, RISWIE outperformed competing metrics in clustering 3D human poses from the FAUST dataset and biological tissue slices grouped into stacks. On FAUST, RISWIE achieved separability ratios close to Gromov-Wasserstein while reducing computation time by over five orders of magnitude. On tissue slices, RISWIE correctly grouped 46 of 48 slices into their original stacks, outperforming Wasserstein, Sliced Wasserstein and Gromov-Wasserstein.