Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances
Khai Nguyen, Hai Nguyen, Nhat Ho
Abstract
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 than the state-of-the-art method, Wasserstein Wormhole, and classical methods, particularly in low-data regimes. To illustrate its robustness, we also experiment the method with intra- and inter-class settings. Finally, we demonstrate that RG can accelerate Wasserstein Wormhole training, yielding RG-Wormhole.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 69deef73-572b-482c-97f2-3cfa7ae4fa4cBuilds on14
- Geometric Dataset Distances via Optimal TransportDavid Alvarez-Melis, Nicolò FusiNeurIPS 2020 · 267 citations
- Distributional Sliced-Wasserstein and Applications to Generative ModelingKhai Nguyen, Nhat Ho, Tung Pham, Hung BuiICLR 2021 · 111 citations
- Wasserstein Embedding for Graph LearningSoheil Kolouri, Navid NaderiAlizadeh, Gustavo K. Rohde, Heiko HoffmannICLR 2021 · 99 citations
- Projection Robust Wasserstein Distance and Riemannian OptimizationTianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi et al.NeurIPS 2020 · 84 citations
- Low-Rank Sinkhorn FactorizationMeyer Scetbon, Marco Cuturi, Gabriel PeyréICML 2021 · 76 citations
Related papers
- Point-set Distances for Learning Representations of 3D Point CloudsTrung Nguyen, Quang-Hieu Pham, Tam Le, Tung Pham et al.ICCV 2021 · 89 citations
- Sliced Wasserstein Estimation with Control VariatesKhai Nguyen, Nhat HoICLR 2024 · 16 citations
- Wasserstein Wormhole: Scalable Optimal Transport Distance with TransformerDoron Haviv, Russell Zhang Kunes, Thomas Dougherty, Cassandra Burdziak et al.ICML 2024 · 15 citations
- Augmented Sliced Wasserstein DistancesXiongjie Chen, Yongxin Yang, Yunpeng LiICLR 2022 · 23 citations
- A Riemannian Block Coordinate Descent Method for Computing the Projection Robust Wasserstein DistanceMinhui Huang, Shiqian Ma, Lifeng LaiICML 2021 · 45 citations
