A Riemannian exponential augmented Lagrangian method for computing the projection robust Wasserstein distance

TL;DR

This paper introduces a novel Riemannian exponential augmented Lagrangian method (ReALM) and an inexact Riemannian gradient descent framework (iRBBS) for efficiently computing the projection robust Wasserstein distance, outperforming existing methods.

Contribution

It develops a new optimization algorithm ReALM with global convergence guarantees and a practical inexact Riemannian Barzilai-Borwein method (iRBBS) for solving PRW distance computation problems.

Findings

ReALM effectively avoids small penalty parameters.

iRBBS computes an $ ext{epsilon}$-stationary point within $ ext{O}( ext{epsilon}^{-3})$ iterations.

Numerical results show superior performance over state-of-the-art solvers.

Abstract

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.