Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers (2404.09411v4)

Abstract: 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.

• The paper presents a transformer-based autoencoder that maps empirical distributions to embedded spaces for linear-time optimal transport distance approximation.
• The method outperforms existing OT acceleration techniques in accuracy and scalability on datasets such as MNIST and spatial transcriptomics.
• The study extends MDS theory for non-Euclidean metrics by providing error bounds and a convergent projected gradient descent algorithm.

Analysis of "Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers"

The paper introduces a novel technique, Wasserstein Wormhole, aimed at efficiently computing optimal transport (OT) distances using a transformer-based autoencoder for embedding empirical distributions. The authors address the notoriously high computational complexity associated with calculating pairwise Wasserstein distances in large cohorts of distributions, presenting a method that approximates OT distances through an embedded space where Euclidean distances are used. This approach is both innovative and practical, extending the domain of scalable OT applications significantly.

Summary of Contributions

The contribution of this paper lies in proposing the Wasserstein Wormhole model, which leverages transformers for scalable computation of Wasserstein distances. The primary achievements and findings in the paper include:

1. Transformer-Based Embedding: The paper details the design of a transformer-based autoencoder that maps empirical distributions into an embedded space. This space allows for efficient Euclidean distance calculations, approximating OT distances in linear time.
2. Comparison to Other Methods: Through comparisons with existing OT accelerations like DiffusionEMD and DWE, the authors demonstrate superior performance in terms of accuracy and scalability on diverse datasets, ranging from MNIST to high-dimensional spatial transcriptomics data.
3. Innovative Theoretical Insights: Extending MDS theory to non-Euclidean metrics, the authors devise upper and lower bounds on the error incurred during non-Euclidean embeddings. They introduce a projected gradient descent algorithm with guaranteed convergence to the global optimum for any distance matrix.
4. Practicability and Versatility: The paper showcases that Wormhole not only computes OT distances effectively but also maintains versatility across multiple domains including computational geometry and single-cell biology. The model adapts to various dataset structures, including high-dimensional niches in spatial transcriptomics.

Implications

Practically, Wasserstein Wormhole offers significant computational advantages, enabling OT-based analyses on datasets of thousands of distributions without the overhead of conventional OT calculations. This potentially positions the method as a standard in fields requiring frequent distribution comparisons, such as image processing, biology, and beyond.

Theoretically, the paper advances the understanding of embedding non-Euclidean metrics, providing a framework for examining such embeddings. The derivation of bounds for the embedding error represents a valuable contribution to computational geometry and distance learning theories.

Future Directions

The research paves the way for further investigation into embedding algorithms for non-Euclidean metrics. Future work could explore expanding the model to handle other OT-based distance metrics, such as the Gromov-Wasserstein distance, more efficiently. Additionally, applying the framework to other high-dimensional data scenarios beyond those examined could extend its utility.

In conclusion, the introduction of Wasserstein Wormhole marks a substantial progression in scalable OT computations. Its transformer-based embedding mechanism, combined with the theoretical guarantees provided, promises to significantly enhance the efficiency and applicability of OT analysis in various scientific and engineering disciplines.