Provably convergent stochastic fixed-point algorithm for free-support Wasserstein barycenter of continuous non-parametric measures
This work addresses the need for scalable and theoretically grounded algorithms for Wasserstein barycenters, which are important in machine learning and statistics for averaging probability distributions.
The paper develops a provably convergent stochastic fixed-point algorithm for computing the 2-Wasserstein barycenter of continuous non-parametric measures, providing the first rigorous convergence analysis with geometric rates under controlled errors. Numerical experiments show strong computational efficiency and estimation accuracy.
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 Álvarez-Esteban, del Barrio, Cuesta-Albertos, and Matrán (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.