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

TL;DR

This paper introduces a new stochastic fixed-point algorithm for computing Wasserstein barycenters of continuous measures, with proven convergence and practical efficiency demonstrated through experiments.

Contribution

It provides the first rigorous convergence analysis for estimator-based stochastic fixed-point methods in Wasserstein barycenter computation and proposes a scalable, provably convergent algorithm.

Findings

Established almost sure convergence of the proposed method.

Achieved geometric rates of convergence under certain conditions.

Demonstrated strong computational efficiency and accuracy in experiments.

Abstract

We develop an estimator-based stochastic fixed-point framework for approximately computing the 2-Wasserstein barycenter of continuous, non-parametric probability measures. Notably, we provide the first rigorous convergence analysis for implementable estimator-based stochastic extensions of the fixed-point iterative scheme proposed by \'Alvarez-Esteban, del Barrio, Cuesta-Albertos, and Matr\'an (2016). In particular, we establish almost sure convergence, and identify sufficient conditions for geometric rates of convergence under controlled errors in optimal transport (OT) map estimation. We subsequently propose a concrete, provably convergent, and computationally tractable stochastic algorithm that accommodates input measures satisfying Caffarelli-type regularity conditions, which form a dense subset of the Wasserstein space. This algorithm leverages a modified entropic OT map estimator to enable efficient and scalable implementation. To facilitate quantitative evaluation, we further propose a novel and efficient procedure for synthetically generating benchmark instances, in which the input measures exhibit non-trivial features and the corresponding barycenters are approximately known. Numerical experiments on both synthetic and real-world datasets demonstrate the strong computational efficiency, estimation accuracy, and sampling flexibility of our approach.