Provably convergent algorithm for free-support Wasserstein barycenter of continuous non-parametric measures

Abstract

We propose a provably convergent algorithm for approximating the Wasserstein barycenter of continuous and non-parametric probability measures. Our algorithm is inspired by the fixed-point scheme of Alvarez-Esteban, Del Barrio, Cuesta-Albertos, and Matran (2016), which begins with an initial reference measure and iteratively performs updates via its pushforward by a weighted sum of optimal transport (OT) maps to the input measures. However, the convergence of their scheme relies on obtaining exact OT maps which is computationally intractable for general non-parametric input measures. Hence, we develop a tailored class of smooth and strongly convex consistent estimators of OT maps, of which a concrete example named kernel-smoothed shape-constrained least square (KS-SCLS) estimator is provided. Specifically, in order to obtain consistent OT map estimators, we develop tailored set truncation techniques to preserve the regularity properties of both the true OT maps and the estimated OT maps throughout the iterations. Replacing the OT maps in the fixed-point scheme with their estimated counterparts then gives rise to a stochastic fixed-point algorithm which is provably convergent to the true Wasserstein barycenter. In particular, this algorithm does not restrict the support of the barycenter to be fixed and can be implemented in a distributed computing environment, which make it suitable for large-scale information aggregation problems. In our numerical experiments, we showcase the efficacy of our proposed algorithm by comparing it with benchmark exact algorithms in special cases, and illustrate its application in probabilistic forecasts aggregation via a toy example.