A Fast Proximal Point Method for Computing Exact Wasserstein Distance

TL;DR

This paper introduces a fast, stable proximal point algorithm for computing the exact Wasserstein distance, overcoming limitations of regularized methods in machine learning and image processing.

Contribution

It develops an Inexact Proximal Point method (IPOT) that guarantees convergence to the exact Wasserstein distance with improved numerical stability and comparable complexity to existing methods.

Findings

Converges to exact Wasserstein distance with theoretical guarantees

Alleviates numerical stability issues in optimal transport computations

Avoids shrinking problem in generative model applications

Abstract

Wasserstein distance plays increasingly important roles in machine learning, stochastic programming and image processing. Major efforts have been under way to address its high computational complexity, some leading to approximate or regularized variations such as Sinkhorn distance. However, as we will demonstrate, regularized variations with large regularization parameter will degradate the performance in several important machine learning applications, and small regularization parameter will fail due to numerical stability issues with existing algorithms. We address this challenge by developing an Inexact Proximal point method for exact Optimal Transport problem (IPOT) with the proximal operator approximately evaluated at each iteration using projections to the probability simplex. The algorithm (a) converges to exact Wasserstein distance with theoretical guarantee and robust regularization parameter selection, (b) alleviates numerical stability issue, (c) has similar computational complexity to Sinkhorn, and (d) avoids the shrinking problem when apply to generative models. Furthermore, a new algorithm is proposed based on IPOT to obtain sharper Wasserstein barycenter.