Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances
This work addresses the computational bottleneck of estimating Wasserstein distances for researchers and practitioners in fields like machine learning and data analysis, offering an incremental improvement over existing methods.
The paper tackles the problem of efficiently computing Wasserstein distances for multiple distribution pairs by proposing a fast estimation method based on regression using sliced Wasserstein distances as predictors. The result is a method that provides a better approximation than the state-of-the-art Wasserstein Wormhole, especially in low-data regimes, and can accelerate its training.
We address the problem of efficiently computing Wasserstein distances for multiple pairs of distributions drawn from a meta-distribution. To this end, we propose a fast estimation method based on regressing Wasserstein distance on sliced Wasserstein (SW) distances. Specifically, we leverage both standard SW distances, which provide lower bounds, and lifted SW distances, which provide upper bounds, as predictors of the true Wasserstein distance. To ensure parsimony, we introduce two linear models: an unconstrained model with a closed-form least-squares solution, and a constrained model that uses only half as many parameters. We show that accurate models can be learned from a small number of distribution pairs. Once estimated, the model can predict the Wasserstein distance for any pair of distributions via a linear combination of SW distances, making it highly efficient. Empirically, we validate our approach on diverse tasks, including Gaussian mixtures, point-cloud classification, and Wasserstein-space visualizations for 3D point clouds. Across various datasets such as MNIST point clouds, ShapeNetV2, MERFISH Cell Niches, and scRNA-seq, our method consistently provides a better approximation of Wasserstein distance than the state-of-the-art Wasserstein embedding model, Wasserstein Wormhole, particularly in low-data regimes. Finally, we demonstrate that our estimator can also accelerate Wormhole training, yielding \textit{RG-Wormhole}.