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 Wormhole，一种基于Transformer的自编码器，它将经验分布嵌入到一个潜空间中，在该空间中欧氏距离近似于OT距离。通过扩展MDS理论，我们证明目标函数隐含了对嵌入非欧氏距离时产生误差的界限。实验表明，Wormhole嵌入之间的距离与Wasserstein距离高度吻合，从而实现了OT距离的线性时间计算。除了一种将分布映射为嵌入的编码器外，Wasserstein Wormhole还包含一个将嵌入映射回分布的解码器，使得嵌入空间中的操作能够泛化到OT空间，例如Wasserstein重心估计和OT插值。通过为OT方法赋予可扩展性和可解释性，Wasserstein Wormhole为计算几何和单细胞生物学领域的数据分析开辟了新的途径。