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为计算几何和单细胞生物学领域的数据分析开辟了新的途径。