# Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers

arXiv: https://arxiv.org/abs/2404.09411 · HTML: https://mlsanity.com/paper/2404.09411/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 71 references that are in the map.
- First posted on arXiv 15 Apr 2024.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [ShapeNet: An Information-Rich 3D Model Repository](https://mlsanity.com/paper/1512.03012/) (2015)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Learning for 3D Point Clouds: A Survey](https://mlsanity.com/paper/1912.12033/) (2019)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Gromov-Wasserstein Learning for Graph Matching and Node Embedding](https://mlsanity.com/paper/1901.06003/) (2019)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Low-Rank Sinkhorn Factorization](https://mlsanity.com/paper/2103.04737/) (2021)
- [Manifold Interpolating Optimal-Transport Flows for Trajectory Inference](https://mlsanity.com/paper/2206.14928/) (2022)
- [Optimal Transport for Diffeomorphic Registration](https://mlsanity.com/paper/1706.05218/) (2017)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Learning Wasserstein Embeddings](https://mlsanity.com/paper/1710.07457/) (2017)
- [Wasserstein Task Embedding for Measuring Task Similarities](https://mlsanity.com/paper/2208.11726/) (2022)
- [Fast and Smooth Interpolation on Wasserstein Space](https://mlsanity.com/paper/2010.12101/) (2020)
- [Diffusion Earth Mover's Distance and Distribution Embeddings](https://mlsanity.com/paper/2102.12833/) (2021)
- [Learning Embeddings into Entropic Wasserstein Spaces](https://mlsanity.com/paper/1905.03329/) (2019)
- [Exploiting Local Geometry for Feature and Graph Construction for Better 3D Point Cloud Processing with Graph Neural Networks](https://mlsanity.com/paper/2103.15226/) (2021)
- [Improving Molecular Representation Learning with Metric Learning-enhanced Optimal Transport](https://mlsanity.com/paper/2202.06208/) (2022)
- [How can classical multidimensional scaling go wrong?](https://mlsanity.com/paper/2110.11430/) (2021)
- [Wasserstein t-SNE](https://mlsanity.com/paper/2205.07531/) (2022)
- [Permutation invariant networks to learn Wasserstein metrics](https://mlsanity.com/paper/2010.05820/) (2020)

## Built on by (most cited first, 14 of 14)

- [Wasserstein Flow Matching: Generative modeling over families of distributions](https://mlsanity.com/paper/2411.00698/) (2024)
- [Fourier Sliced-Wasserstein Embedding for Multisets and Measures](https://mlsanity.com/paper/2504.02544/) (2025)
- [Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances](https://mlsanity.com/paper/2509.20508/) (2025)
- [Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning](https://mlsanity.com/paper/2505.18150/) (2025)
- [Efficient Transferable Optimal Transport via Min-Sliced Transport Plans](https://mlsanity.com/paper/2511.19741/) (2025)
- [A Unified Foundation Model for All-in-One Multi-Modal Remote Sensing Image Restoration and Fusion with Language Prompting](https://mlsanity.com/paper/2604.05629/) (2026)
- [Revisiting Multi-Permutation Equivariance through the Lens of Irreducible Representations](https://mlsanity.com/paper/2410.06665/) (2024)
- [Constrained Sliced Wasserstein Embedding](https://mlsanity.com/paper/2506.02203/) (2025)
- [Distribution-Conditioned Transport](https://mlsanity.com/paper/2603.04736/) (2026)
- [Amortized Optimal Transport from Sliced Potentials](https://mlsanity.com/paper/2604.15114/) (2026)
- [Fourier Sliced-Wasserstein Embedding for Multisets and Measures](https://mlsanity.com/paper/2405.16519/) (2024)
- [Embedding Empirical Distributions for Computing Optimal Transport Maps](https://mlsanity.com/paper/2504.17740/) (2025)
- [Graph Alignment via Dual-Pass Spectral Encoding and Latent Space Communication](https://mlsanity.com/paper/2509.09597/) (2025)
- [ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection](https://mlsanity.com/paper/2605.12879/) (2026)

## Abstract (arXiv)

Optimal transport (OT) and the related Wasserstein metric (W) are powerful and ubiquitous tools for comparing distributions. However, computing pairwise Wasserstein distances rapidly becomes intractable as cohort size grows. An attractive alternative would be to find an embedding space in which pairwise Euclidean distances map to OT distances, akin to standard multidimensional scaling (MDS). We present Wasserstein Wormhole, a transformer-based autoencoder that embeds empirical distributions into a latent space wherein Euclidean distances approximate OT distances. Extending MDS theory, we show that our objective function implies a bound on the error incurred when embedding non-Euclidean distances. Empirically, distances between Wormhole embeddings closely match Wasserstein distances, enabling linear time computation of OT distances. Along with an encoder that maps distributions to embeddings, Wasserstein Wormhole includes a decoder that maps embeddings back to distributions, allowing for operations in the embedding space to generalize to OT spaces, such as Wasserstein barycenter estimation and OT interpolation. By lending scalability and interpretability to OT approaches, Wasserstein Wormhole unlocks new avenues for data analysis in the fields of computational geometry and single-cell biology.