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.

翻译：最优输运（OT）及其相关的Wasserstein度量（W）是用于比较分布的强大且普遍的工具。然而，随着样本群规模的增加，计算成对Wasserstein距离会迅速变得难以处理。一种有吸引力的替代方案是寻找一个嵌入空间，使得其中的欧几里得距离映射到OT距离，类似于标准多维尺度分析（MDS）。我们提出了Wasserstein虫洞，这是一种基于Transformer的自编码器，它将经验分布嵌入到一个潜在空间中，在该空间中欧几里得距离近似于OT距离。通过扩展MDS理论，我们证明了我们的目标函数意味着在嵌入非欧几里得距离时所产生误差的一个上界。实证表明，虫洞嵌入之间的距离与Wasserstein距离高度吻合，从而实现了OT距离的线性时间计算。除了将分布映射到嵌入的编码器外，Wasserstein虫洞还包含一个将嵌入映射回分布的解码器，这使得嵌入空间中的操作能够推广到OT空间，例如Wasserstein重心估计和OT插值。通过为OT方法提供可扩展性和可解释性，Wasserstein虫洞为计算几何和单细胞生物学领域的数据分析开辟了新途径。