Research map / Learning theory and algorithms / Optimal transport
Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformer
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains a transformer autoencoder (Wasserstein Wormhole) so that Euclidean distances between per-point-cloud embeddings approximate entropic/Sinkhorn OT distances, with a decoder reconstructing distributions for barycenters and interpolation.
Problem
Computing pairwise Wasserstein distances across large cohorts of empirical distributions (point clouds) is quadratic in cohort size and becomes intractable; existing OT-preserving embeddings do not scale to high dimensions or arbitrary supports.
Approach
Encode each point cloud with a permutation-invariant transformer into a single vector whose squared Euclidean distances are trained (stress loss on mini-batch Sinkhorn divergences) to match OT distances, plus a transformer decoder trained to reconstruct point clouds by OT, and extend MDS theory with lower/upper bounds on the stress when embedding non-Euclidean distance matrices.
Key results
- Wormhole achieves OT correlation of 0.97-0.99 across datasets and outperforms DiffusionEMD, DWE and Fréchet on the 2D and 3D benchmarks (e.g., MNIST 0.98 vs DWE 0.92).
- On MNIST, Wormhole needs about 2000-fold fewer OT computations than naive pairwise computation, training in about 10 minutes.
- Wormhole is the only evaluated embedding method that scales to the 256,058-cell, 254-dimensional MERFISH cohort; competing methods run out of memory there.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · MNIST | OT correlation: 0.92 -> 0.98 | Helped |
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · Fashion-MNIST | MSE: 1.88\cdot10^{-4} -> 1.41\cdot10^{-5} | Helped |
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · ModelNet40 | label accuracy: 0.78 -> 0.86 | Helped |
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · ShapeNet | MSE: 6.66\cdot10^{-4} -> 1.14\cdot10^{-4} | Helped |
| Use GW Wormhole instead of Fréchet | Rotated ShapeNet embeddings · Rotated ShapeNet (GW) | OT correlation: 0.68 -> 0.98 | Helped |
| Use GW Wormhole instead of Fréchet | Rotated ShapeNet embeddings · Rotated ShapeNet (GW) | label accuracy: 0.16 -> 0.82 | Helped |
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · ModelNet40 | MSE: 3.52\cdot10^{-4} -> 8.89\cdot10^{-5} | Helped |
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · MNIST | MSE: 1.63\cdot10^{-4} -> 2.57\cdot10^{-5} | Helped |
| Use Wasserstein Wormhole instead of DWE | OT-preserving embeddings · ShapeNet | OT correlation: 0.97 -> 0.99 | Helped |
Surprising findings
- Classical MDS on non-Euclidean Wasserstein distance matrices yields stress outside the derived bounds and gets worse as dimension increases.
- Wormhole embeddings generalize to a held-out class ('Bag') in Fashion-MNIST with Pearson correlation 0.995 and unchanged label accuracy.
Limitations
- Only handles uniform-weight point clouds; arbitrary weights are future work.
- Relies on entropic Sinkhorn computations during training, requiring mini-batch OT calculations.
- Evaluation of pairwise OT correlation is on random subsamples of 128 point clouds rather than full cohorts.
Experimental setup
Datasets: MNIST Fashion-MNIST ModelNet40 ShapeNet MERFISH cell niches SeqFISH cell niches scRNA-seq COVID atlas (MetaCells)
Benchmarks: MNIST Fashion-MNIST ModelNet40 ShapeNet Rotated ShapeNet (GW) MERFISH Cell Niches SeqFISH Cell Niches scRNA-seq Atlas
Metrics: Pearson correlation of OT distances MSE between embedding and OT distances label classification accuracy stress
Baselines: DiffusionEMD DWE Fréchet distance Low-Rank Sinkhorn MetaOT classical MDS PGD
Builds on: Sinkhorn divergence multidimensional scaling Deep Wasserstein Embedding DiffusionEMD Set Transformer / attention
Closest papers by meaning
- Geometry-Aware Optimal Transport: Fast Intrinsic Dimension and Wasserstein Distance Estimation (NeurIPS 2026 · similarity 0.88)
- Tree-Sliced Wasserstein Distance with Nonlinear Projection (ICML 2025 · similarity 0.88)
- Wasserstein Flow Matching: Generative Modeling Over Families of Distributions (ICML 2025 · similarity 0.87)
- Tree-Sliced Wasserstein Distance: A Geometric Perspective (ICML 2025 · similarity 0.87)
- Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy (ICLR 2025 · similarity 0.87)
- Distance-Based Tree-Sliced Wasserstein Distance (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.