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 the 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.

Applications

Barycenters in Image Processing

Barycenters have several applications in image processing. A basic example arises in handwriting recognition. One may wish to determine what an expected letter looks like, but simply taking the pixel-by-pixel Euclidean average won't capture the geometric information of the letter. Instead, a $k$-means algorithm can compute what a prototypical version of that letter looks like.

As another application, barycenters can be used to interpolate between two or more images. For example, suppose we are given two images of similar objects, one with a darker gray-scale histogram $\mu _{1}$ and the other with a lighter one $\mu _{2}$. One way to find the middle point of their histograms would be to take their Euclidean average $(\mu _{1}+\mu _{2})/2$, but this would result in a histogram with a distribution of some dark and some light shades as opposed to a gray-scaled version. The Wasserstein barycenter between these two histograms provides a gray image. Similar ideas can be used to enhance the contrast of an image.

Generalizations

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