A Unified Approach for Computing Wasserstein Barycenters of Discrete and Continuous Measures

TL;DR

This paper introduces a primal mirror descent algorithm that efficiently computes Wasserstein barycenters for both discrete and continuous measures within a unified framework, with proven convergence guarantees.

Contribution

It presents a novel, unified algorithm capable of handling discrete and continuous measures simultaneously, improving flexibility and convergence in Wasserstein barycenter computation.

Findings

Algorithm converges to the exact barycenter in both discrete and continuous cases.

Demonstrates promising accuracy and computational efficiency on synthetic and real data.

Produces absolutely continuous iterates that approximate discrete barycenters.

Abstract

Computing the unregularized Wasserstein barycenter for measure-valued data is a challenging optimization task. Recent algorithms have been tailored to either discrete measures as point clouds or continuous measures discretized on regular grids. In this work, we propose a primal mirror descent algorithm for computing the exact Wasserstein barycenter in the Fisher-Rao geometry. Our algorithm is a unified approach that is flexible enough to simultaneously cover discrete and absolutely continuous input measures, with convergence guarantees in both settings. In particular, when all input measures are discrete, our algorithm, initialized from any probability density, solves a sequence of semi-discrete optimal transport subproblems and produces absolutely continuous iterates that converge to the discrete barycenter. We use synthetic and real data examples to demonstrate the promising result in terms of accuracy and computational cost.