# A Riemannian Block Coordinate Descent Method for Computing the Projection Robust Wasserstein Distance

arXiv: https://arxiv.org/abs/2012.05199 · HTML: https://mlsanity.com/paper/2012.05199/

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Dec 2020.

## Built on by (most cited first, 27 of 27)

- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Statistical, Robustness, and Computational Guarantees for Sliced Wasserstein Distances](https://mlsanity.com/paper/2210.09160/) (2022)
- [Sinkhorn Distributionally Robust Optimization](https://mlsanity.com/paper/2109.11926/) (2021)
- [Spherical Sliced-Wasserstein](https://mlsanity.com/paper/2206.08780/) (2022)
- [Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution](https://mlsanity.com/paper/2204.01188/) (2022)
- [Efficiently Escaping Saddle Points in Bilevel Optimization](https://mlsanity.com/paper/2202.03684/) (2022)
- [Amortized Projection Optimization for Sliced Wasserstein Generative Models](https://mlsanity.com/paper/2203.13417/) (2022)
- [Monge, Bregman and Occam: Interpretable Optimal Transport in High-Dimensions with Feature-Sparse Maps](https://mlsanity.com/paper/2302.04065/) (2023)
- [Accelerating Inexact HyperGradient Descent for Bilevel Optimization](https://mlsanity.com/paper/2307.00126/) (2023)
- [Riemannian Hamiltonian methods for min-max optimization on manifolds](https://mlsanity.com/paper/2204.11418/) (2022)
- [Markovian Sliced Wasserstein Distances: Beyond Independent Projections](https://mlsanity.com/paper/2301.03749/) (2023)
- [Re-evaluating Word Mover's Distance](https://mlsanity.com/paper/2105.14403/) (2021)
- [Max-Sliced Mutual Information](https://mlsanity.com/paper/2309.16200/) (2023)
- [Projection Robust Wasserstein Barycenters](https://mlsanity.com/paper/2102.03390/) (2021)
- [Tree-Sliced Wasserstein Distance: A Geometric Perspective](https://mlsanity.com/paper/2406.13725/) (2024)
- [Learning Elastic Costs to Shape Monge Displacements](https://mlsanity.com/paper/2306.11895/) (2023)
- [Provably Convergent Policy Optimization via Metric-aware Trust Region Methods](https://mlsanity.com/paper/2306.14133/) (2023)
- [Distributional Counterfactual Explanations With Optimal Transport](https://mlsanity.com/paper/2401.13112/) (2024)
- [Statistical and Computational Guarantees of Kernel Max-Sliced Wasserstein Distances](https://mlsanity.com/paper/2405.15441/) (2024)
- [Riemannian coordinate descent algorithms on matrix manifolds](https://mlsanity.com/paper/2406.02225/) (2024)
- [On the Convergence of Projected Alternating Maximization for Equitable and Optimal Transport](https://mlsanity.com/paper/2109.15030/) (2021)
- [Efficient Optimization with Orthogonality Constraint: a Randomized Riemannian Submanifold Method](https://mlsanity.com/paper/2505.12378/) (2025)
- [Low-complexity subspace-descent over symmetric positive definite manifold](https://mlsanity.com/paper/2305.02041/) (2023)
- [Last-Iterate Convergence of Adaptive Riemannian Gradient Descent for Equilibrium Computation](https://mlsanity.com/paper/2306.16617/) (2023)
- [A Copula Graphical Model for Multi-Attribute Data using Optimal Transport](https://mlsanity.com/paper/2404.06735/) (2024)
- [Accelerated Fully First-Order Methods for Bilevel and Minimax Optimization](https://mlsanity.com/paper/2405.00914/) (2024)
- [Collision-based Dynamics for Multi-Marginal Optimal Transport](https://mlsanity.com/paper/2412.16385/) (2024)

## Abstract (arXiv)

The Wasserstein distance has become increasingly important in machine learning and deep learning. Despite its popularity, the Wasserstein distance is hard to approximate because of the curse of dimensionality. A recently proposed approach to alleviate the curse of dimensionality is to project the sampled data from the high dimensional probability distribution onto a lower-dimensional subspace, and then compute the Wasserstein distance between the projected data. However, this approach requires to solve a max-min problem over the Stiefel manifold, which is very challenging in practice. The only existing work that solves this problem directly is the RGAS (Riemannian Gradient Ascent with Sinkhorn Iteration) algorithm, which requires to solve an entropy-regularized optimal transport problem in each iteration, and thus can be costly for large-scale problems. In this paper, we propose a Riemannian block coordinate descent (RBCD) method to solve this problem, which is based on a novel reformulation of the regularized max-min problem over the Stiefel manifold. We show that the complexity of arithmetic operations for RBCD to obtain an ε-stationary point is O(ε-3). This significantly improves the corresponding complexity of RGAS, which is O(ε-12). Moreover, our RBCD has very low per-iteration complexity, and hence is suitable for large-scale problems. Numerical results on both synthetic and real datasets demonstrate that our method is more efficient than existing methods, especially when the number of sampled data is very large.