Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers

Area of Science:

• Computational geometry
• Machine learning
• Single-cell biology

Background:

• Optimal transport (OT) and Wasserstein distances are crucial for comparing distributions.
• Calculating pairwise Wasserstein distances is computationally intensive for large datasets.
• Existing methods lack scalability and efficient computation for high-dimensional data.

Purpose of the Study:

• To develop a scalable method for approximating Wasserstein distances.
• To create an embedding space where Euclidean distances approximate OT distances.
• To enable efficient computation and generalization of OT-based analyses.

Main Methods:

• Developed Wasserstein Wormhole, a transformer-based autoencoder.
• The autoencoder embeds empirical distributions into a latent space.
• Utilized an objective function extending multidimensional scaling (MDS) theory to bound embedding errors.

Main Results:

• Wasserstein Wormhole embeddings closely approximate Wasserstein distances.
• Achieved linear time computation of OT distances, significantly improving scalability.
• Demonstrated the decoder's ability to generalize embedding space operations (barycenter estimation, interpolation) to OT spaces.

Conclusions:

• Wasserstein Wormhole provides a scalable and interpretable approach to optimal transport.
• Enables efficient analysis of distributions in computational geometry and single-cell biology.
• Facilitates new data analysis avenues by bridging Euclidean and Wasserstein spaces.