Wasserstein barycenters and applications in image processing

In optimal transport, a Wasserstein barycenter is a probability measure that acts as a center of mass between two or more probability measures. It generalizes the notions of physical barycenters and geometric centroids.

Introduction

Motivation

Barycenters in physics and geometry are points that represent a notion of a mean of a number of objects. In astronomy, the barycenter is the center of mass of two or more objects that orbit each other, and in geometry, a centroid is the arithmetic mean of all the points in an object. Given countably many points $\{x_{i}\}_{i\in I}$ in $\mathbb {R} ^{n}$ with nonnegative weights $\{\lambda _{i}\}_{i\in I}$, the weighted $L^{2}$ barycenter of the points is the unique point $y$ minimizing $\sum _{i\in I}\lambda _{i}||y-x_{i}||_{2}^{2}$. Wasserstein barycenters attempt to capture this concept for probability measures by replacing the Euclidean distance with the Wasserstein distance of two probability measures, $W_{2}(\rho ,\mu )$.

Definition

Let $\Omega$ be a domain and ${\mathcal {P}}(\Omega )$ be the set of probability measures on $\Omega$. Given a collection of probability measures $\{\mu _{i}\}_{i\in I}$ and nonnegative weights $\{\lambda _{i}\}_{i\in I}$, we define a weighted barycenter of $\{\mu _{i}\}_{i\in I}$ as any probability measure $\mu$ that minimizes  $\sum _{i\in I}\lambda _{i}W_{2}(\mu _{i},\mu )^{2}$ over the space $\mu \in {\mathcal {P}}(\Omega )$. Here $W_{2}$ denotes the 2-Wasserstein distance, which may be replaced with the $p$-Wasserstein distance, $W_{p}$, though this is not always as convenient.

This minimization problem was originally introduced by Agueh and Carlier , who also proposed an alternative formulation of the problem above when there are finitely many measures $\{\mu _{i}\}_{i=1}^{N}$. Instead of considering the minimum over all probability measures $\mu \in {\mathcal {P}}(\Omega )$, one can equivalently consider the optimization problem over all multi-marginal transport plans $\gamma \in {\mathcal {P}}(\Omega ^{N+1})$ whose push forwards satisfy $(\pi _{i})_{\#}\gamma =\mu _{i}$ for $i\in \{1,\ldots ,N\}$ and $(\pi _{0})_{\#}\gamma =\mu$ for some unspecified probability measure $\mu$. The problem then tries to minimize

$\int \left(\sum _{i=1}^{N}\lambda _{i}\|x_{i}-x_{0}\|^{2}\right)d\gamma$

Existence and Uniqueness

Agueh and Carlier demonstrated that when at least one of the measures $\mu _{i}$ is absolutely continuous, the Wasserstein barycenter exists and is unique.

Examples

In the case where all of the measures $\mu _{i}$ are finitely supported measures, computing the Wasserstein barycenter reduces to a problem in linear programming.

When our family of measures $\{\mu _{i}\}_{i\in I}$ consists of only one finitely supported measure $\mu _{1}$ and we restrict the input probability measure $\mu$ to finitely supported probability measures whose supports have at most $k$ points, the problem of finding the Wasserstein barycenter is equivalent to the $k$-means clustering problem.

Generalizations

Wasserstein barycenters are examples of Karcher and Fréchet means where the distance function used is the Wasserstein distance.