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）及其相关的瓦瑟斯坦度量（W）是比较分布时功能强大且普适的工具。然而，随着样本规模的增大，计算成对瓦瑟斯坦距离的复杂度会迅速变得难以处理。一种有吸引力的替代方案是寻找一个嵌入空间，使其中成对欧几里得距离能够映射为OT距离，类似于经典的多维标度（MDS）。我们提出了瓦瑟斯坦虫洞（Wasserstein Wormhole），一种基于Transformer的自编码器，它能将经验分布嵌入到潜在空间中，使得其中的欧几里得距离近似于OT距离。通过扩展MDS理论，我们证明了目标函数隐含了非欧几里得距离嵌入时的误差界。实验表明，虫洞嵌入之间的距离与瓦瑟斯坦距离高度吻合，从而实现了OT距离的线性时间计算。除了将分布映射为嵌入的编码器，瓦瑟斯坦虫洞还包含一个将嵌入解码回分布的解码器，使得嵌入空间中的操作能够推广到OT空间，例如瓦瑟斯坦重心估计和OT插值。通过赋予OT方法可扩展性和可解释性，瓦瑟斯坦虫洞为计算几何和单细胞生物学领域的数据分析开辟了新途径。