A Riemannian Exponential Augmented Lagrangian Method for Computing the Projection Robust Wasserstein Distance

Projecting the distance measures onto a low-dimensional space is an efficient way of mitigating the curse of dimensionality in the classical Wasserstein distance using optimal transport. The obtained maximized distance is referred to as projection robust Wasserstein (PRW) distance. In this paper, we equivalently reformulate the computation of the PRW distance as an optimization problem over the Cartesian product of the Stiefel manifold and the Euclidean space with additional nonlinear inequality constraints. We propose a Riemannian exponential augmented Lagrangian method (ReALM) with a global convergence guarantee to solve this problem. Compared with the existing approaches, ReALM can potentially avoid too small penalty parameters. Moreover, we propose a framework of inexact Riemannian gradient descent methods to solve the subproblems in ReALM efficiently. In particular, by using the special structure of the subproblem, we give a practical algorithm named as the inexact Riemannian Barzilai-Borwein method with Sinkhorn iteration (iRBBS). The remarkable features of iRBBS lie in that it performs a flexible number of Sinkhorn iterations to compute an inexact gradient with respect to the projection matrix of the problem and adopts the Barzilai-Borwein stepsize based on the inexact gradient information to improve the performance. We show that iRBBS can return an $\epsilon$-stationary point of the original PRW distance problem within $\mathcal{O}(\epsilon^{-3})$ iterations. Extensive numerical results on synthetic and real datasets demonstrate that our proposed ReALM as well as iRBBS outperform the state-of-the-art solvers for computing the PRW distance.