Closing the Gap: Efficient Algorithms for Discrete Wasserstein Barycenters

Abstract

Given a collection of probability measures, the Wasserstein barycenter problem seeks a probability measure that minimizes the weighted sum of their squared type-$2$ Wasserstein distances. We study the discrete setting, in which every input measure has finite support; this setting arises frequently in machine learning and operations research. Since computing a discrete Wasserstein barycenter is NP-hard, we develop approximation algorithms with provable performance guarantees. The best previously known polynomial-time guarantee is a factor of $2$. We improve this bound by proposing a polynomial-time approximation scheme (PTAS) that generalizes the existing $2$-approximation: for any prescribed $\alpha\in(0,1]$, the proposed algorithms return a $(1+\alpha)$-approximate barycenter in time polynomial in $(nk)^{1/\alpha}$ and $d$. We provide both randomized and deterministic constructions and derive a sharper guarantee when the input measures are equally weighted. Numerical experiments on synthetic and real datasets show that the algorithms are computationally practical and produce near-optimal barycenters.