Computing Wasserstein Barycenter via operator splitting: the method of averaged marginals

TL;DR

This paper introduces a scalable, parallelizable algorithm for computing Wasserstein barycenters, including unbalanced cases, using operator splitting and averaging marginals, suitable for large datasets.

Contribution

It proposes a novel convex nonsmooth optimization formulation and a decomposition scheme based on Douglas-Rachford splitting for efficient computation of Wasserstein barycenters.

Findings

The method performs competitively against state-of-the-art algorithms.

It effectively handles large-scale and unbalanced measures.

The algorithm allows parallel and randomized projections for efficiency.

Abstract

The Wasserstein barycenter (WB) is an important tool for summarizing sets of probability measures. It finds applications in applied probability, clustering, image processing, etc. When the measures' supports are finite, computing a (balanced) WB can be done by solving a linear optimization problem whose dimensions generally exceed standard solvers' capabilities. In the more general setting where measures have different total masses, we propose a convex nonsmooth optimization formulation for the so-called unbalanced WB problem. Due to their colossal dimensions, we introduce a decomposition scheme based on the Douglas-Rachford splitting method that can be applied to both balanced and unbalanced WB problem variants.Our algorithm, which has the interesting interpretation of being built upon averaging marginals, operates a series of simple (and exact) projections that can be parallelized and even randomized, making it suitable for large-scale datasets. Numerical comparisons against state-of-the-art methods on several data sets from the literature illustrate the method's performance.