Statistical data analysis in the Wasserstein space

TL;DR

This paper reviews recent advances in statistical inference for data modeled as probability measures, emphasizing Wasserstein distances and optimal transport tools like barycenters and geodesic PCA for analyzing geometric variations.

Contribution

It highlights the benefits of Wasserstein-based methods such as barycenter and geodesic PCA in understanding geometric data variations and discusses emerging research directions in statistical optimal transport.

Findings

Wasserstein distances effectively analyze probability measure data.

Barycenter and geodesic PCA reveal principal geometric variations.

The paper discusses future research perspectives in statistical optimal transport.

Abstract

This paper is concerned by statistical inference problems from a data set whose elements may be modeled as random probability measures such as multiple histograms or point clouds. We propose to review recent contributions in statistics on the use of Wasserstein distances and tools from optimal transport to analyse such data. In particular, we highlight the benefits of using the notions of barycenter and geodesic PCA in the Wasserstein space for the purpose of learning the principal modes of geometric variation in a dataset. In this setting, we discuss existing works and we present some research perspectives related to the emerging field of statistical optimal transport.

Figures

Equations

$$W_2(\mu ,\nu )=\underset{\pi \in \Gamma (\mu ,\nu )}{\inf }{\left({\iint }_{{\Omega }^2}|x-y{|}^{2}d\pi (x,y)\right)}^{1/2},\text{mbox}for\mu ,\nu \in {\mathcal{P}}_2(\Omega ),$$

$$W_2^2(\mu ,\nu )={\int }_0^1(F_{\mu }^{-}(u)-F_{\nu }^{-}(u){)}^{2}du,$$

$$W_2(r,q):=\underset{U\in U(r,q)}{\min }⟨C,U{⟩}^{1/2},$$

$$W_{2,\epsilon }^2(r,q):=\underset{U\in U(r,q)}{\min }⟨C,U⟩-\epsilon \text{h}(U).$$

$$\mu \mapsto \frac{1}{n}\sum _{i=1}^n{W}_2^2(\mu ,{\nu }_i),\text{mbox}over\mu \in {\mathcal{P}}_2(\Omega ).$$

$${\int }_{{\mathcal{P}}_2(\Omega )}{W}_2^2(\mu ,\nu )d\mathbb{P}(\nu )<+\infty ,\text{mbox}forsome\mu \in {\mathcal{P}}_2(\Omega ).$$

$$\mu \mapsto {\int }_{{\mathcal{P}}_2(\Omega )}{W}_2^2(\mu ,\nu )d\mathbb{P}(\nu ),\text{mbox}over\mu \in {\mathcal{P}}_2(\Omega ).$$

$${\mathbf{\nu }}_{p_i}=\frac{1}{p_i}\sum _{j=1}^{p_i}{\delta }_{{\mathbf{X}}_{i,j}},$$

$$\mu \mapsto \frac{1}{n}\sum _{i=1}^n{W}_2^2(\mu ,{\mathbf{\nu }}_{p_i}),\text{mbox}over\mu \in {\mathcal{P}}_2(\Omega ),$$

$${\hat{\mathbf{\mu }}}_{n,p}=\frac{1}{p}\sum _{j=1}^p{\delta }_{{\bar{\mathbf{X}}}_j^{\ast }},$$

$$F^{-}_{0}(\alpha)={\mathbb{E}}\left(\boldsymbol{F}^{-}(\alpha)\right)\mbox{ where $\boldsymbol{F}^{-}$ is the quantile function of the random measure $\boldsymbol{\nu}$ with distribution ${\mathbb{P}}$.}$$

$$\mathbb{E}\left(W_2^2({\hat{\mathbf{\mu }}}_{n,p},{\mu }_{\mathbb{P}})\right)\le \frac{1}{n}{\int }_0^1\mathrm{Var}\left({\mathbf{F}}^{-}(\alpha )\right)d\alpha +\mathbb{E}\left(W_2^2({\mathbf{\mu }}_p,{\mu }_{\mathbb{P}})\right).$$

$$\mathbb{E}\left(W_2^2({\mathbf{\mu }}_p,{\mu }_{\mathbb{P}})\right)\le \frac{2}{p+1}{J}_2({\mu }_{\mathbb{P}}),\text{mbox}where{J}_2({\mu }_{\mathbb{P}})={\int }_{\Omega }\frac{F_0(x)(1-F_0(x))}{f_0(x)}dx,$$

$$\underset{\mu \in {\mathcal{P}}_2(\Omega )}{\min }{\int }_{{\mathcal{P}}_2(\Omega )}{W}_2^2(\mu ,\nu )d\mathbb{P}(\nu )+\gamma E(\mu )$$

$$E(\mu)=\left\{\begin{array}[]{ll}\|f\|_{H^{k}(\Omega)}^{2},&\mbox{if}\ f=\frac{d\mu}{dx}\ \mbox{and}\ f\geq\alpha,\\ +\infty&\mbox{otherwise.}\end{array}\right.$$

$${\mu }_{\mathbb{P}}^{\gamma }$$

$${\hat{\mathbf{\mu }}}_{n,p}^{\gamma }$$

$$\mathbb{E}\left(\Vert {\hat{\mathbf{f}}}_{n,p}^{\gamma }-f_{\mathbb{P}}^{\gamma }{\Vert }_{{\mathbb{L}}_2(\Omega )}^2\right)\le c\left(\frac{1}{\gamma p^{1/4}}+\frac{1}{\gamma n^{1/2}}\right)$$

$${\tilde{\mathbf{q}}}_i^{p_i}=\frac{1}{p_i}\sum _{j=1}^{p_i}{\delta }_{{\tilde{X}}_{i,j}},\text{mbox}where{\tilde{X}}_{i,j}=\underset{x\in {\Omega }_N}{\mathrm{argmin}}|x-X_{i,j}|,$$

$${\hat{\mathbf{r}}}_{n,p}^{\epsilon }=\underset{r\in {\Sigma }_N}{\mathrm{argmin}}\frac{1}{n}\sum _{i=1}^n{W}_{2,\epsilon }^2(r,{\tilde{\mathbf{q}}}_i^{p_i}),$$

$${\Sigma }_N^{\rho }=\left\{r\in {\Sigma }_N:\underset{1\le \ell \le N}{\min }{r}_{\ell }\ge \rho \right\},$$

$$L_{\rho ,\epsilon }={\left(\sum _{1\le m\le N}{\left(2\epsilon \log (N)+\underset{1\le \ell ,k\le N}{\sup }|C_{m\ell }-C_{k\ell }|-2\epsilon \log (\rho )\right)}^2\right)}^{1/2}$$

$$\underset{\epsilon \to 0}{\lim }{L}_{\rho ,\epsilon }={\left(\sum _{1\le m\le N}{\left(\underset{1\le \ell ,k\le N}{\sup }|C_{m\ell }-C_{k\ell }|\right)}^2\right)}^{1/2}.$$

$$r^{\epsilon }$$

$${\hat{\mathbf{r}}}_{n,p}^{\epsilon }$$

$$\mathbb{E}(\Vert r^{\epsilon }-{\hat{r}}_{n,p}^{\epsilon }{\Vert }_{{\mathbb{R}}^N}^2)\le \frac{32L_{\rho ,\epsilon }^2}{{\epsilon }^{2}n}+\frac{2L_{\rho ,\epsilon }}{\epsilon }\left(\sqrt{\frac{N}{p}}+2\rho (N+\sqrt{N})\right).$$

$${\hat{\mathbf{\mu }}}_{n,p}^h=\underset{\mu \in {\mathcal{P}}_2(\Omega )}{\mathrm{argmin}}\frac{1}{n}\sum _{i=1}^n{W}_2^2({\hat{\mathbf{\nu }}}_i^{h_i},\mu ),$$

$$r^0\in \underset{r\in {\Sigma }_N^{\rho }}{\mathrm{argmin}}{\mathbb{E}}_{\mathbf{q}\sim \mathbb{P}}[W_{2,0}^2(r,\mathbf{q})]\text{mbox}with{W}_{2,0}^2(r,q):=\underset{U\in U(r,q)}{\min }⟨C,U⟩.$$

Full text

1 The emerging field of statistical optimal transport

In many fields of interest (e.g. in signal and image processing or bio-informatics), one records data in the form of high-dimensional vectors or matrices. For the purpose of learning information form such data sets, a standard statistical analysis consists in considering that the observations are realizations of random variables belonging to a linear space endowed with an Euclidean distance. However, it has now been widely recognized that a significant gain in statistical inference can be achieved by the use of non-Euclidean distances to better capture the geometry of the sets to which the data (in the form of histograms, curves, images or points clouds) truly belong. Such a choice of non-Euclidean distances comes from the observation that the data, in many applications, tend to vary along geodesics that are not straight lines. Therefore, the use of non-Euclidean distances is often meaningful to extract geometric information on non-linear sources of variability. Concepts such as nonlinearity and an appropriate modeling of the location of signals or images intensities are essential tools to solve high-dimensional regression problems for the purpose of estimation and classification. The use of Wasserstein distances based on the mathematics of optimal mass transport [63] allows to incorporate both spatial and intensity informations when comparing elements in a dataset that can be modeled as probability distributions. This approach leads to an intuitive geometric interpretation of the variability in such datasets.

Examples of such data are numerous. As an illustration, one may consider histograms that represent, for a list of first names, the distribution of children born with that name per year in France between years 1900 and 2013. In Figure 1, we display the histograms of four different names, as well as the whole dataset. Another example arises from the statistical analysis of data acquired from flow cytometry which is a high-throughput technique in biotechnology that can measure a large number of surface and intracellular markers of single cell in a biological sample. With this technique, one can assess individual characteristics (in the form of multivariate data) at a cellular level to determine the type of cell, their functionality and the way they differ. The development of this technology now leads to datasets made of multiple measurements (e.g. up to 18) of millions of individuals cells from different subjects. This leads to statistical inference problems from multiple points clouds such as those displayed in Figure 2 that can be modeled as discrete probability measures supported on ${\mathbb{R}}^{d}$.

Techniques based on optimal transport for data science have thus recently received an increasing interest in mathematical and computational statistics [8, 11, 13, 15, 12, 14, 10, 29, 49, 66, 62, 45, 44, 54, 58, 28, 67, 47, 51], machine learning [23, 55, 37, 38, 22, 6, 36, 40, 32, 57, 4], image processing and computer vision [41, 52, 7, 24, 30, 17, 31, 59, 60] or computational biology [56].

This paper, whose content is biased by the point of view of a statistician, is an introduction to some recent developments of optimal transport for statistical data analysis. It is focused on the notions of barycenter and geodesic principal component analysis (PCA) in the Wasserstein space and their connection to nonparametric statistics and functional data analysis. We also report the results of numerical experiments to highlight, in an intuitive manner, the potential benefits of such tools for statistical inference. A recent review on the statistical aspects of Wasserstein distances with a rich bibliography on the mathematics of optimal transport can also be found in [50]. More generally, computational optimal transport for data science is a rapidly growing research field. For a detailed introduction and examples of application beyond the prism of statistics we refer to [25].

In Section 2, we present the notions of Wasserstein space and regularized optimal transport. Section 3 is focused on Wasserstein barycenters and statistical inference problems related to this concept. Finally, geodesic principal component analysis (GPCA) in the Wasserstein space is briefly presented in Section 4. Note that the presentation of the paper is rater informal on the mathematical aspects, and we mainly give intuition of the benefits of statistical optimal transport through the presentation of various numerical examples. Throughout the paper, we also discuss some research perspectives and open problems related to the statistical aspects of barycenters and GPCA in the Wasserstein space.

2 Wasserstein space and regularized optimal transport

2.1 The Wasserstein metric

We denote by ${\mathcal{P}}_{2}(\Omega)$ the space of probability measures with finite second moment and supports included in a convex domain $\Omega\subset{\mathbb{R}}^{d}$. This set can be endowed with the Wasserstein metric $W_{2}$ defined as

[TABLE]

where $\Gamma(\mu,\nu)$ is the set of probability measures (also called transport plans or coupling measures) on the product space $\Omega\times\Omega$ with respective marginals $\mu$ and $\nu$, and $|\cdot|$ denotes the usual Euclidean norm on ${\mathbb{R}}^{d}$. The metric space $({\mathcal{P}}_{2}(\Omega),W_{2})$ is called the ($2$-)Wasserstein space.

Remark 2.1**.**

The minimisation problem (1) corresponds the so-called Kantorovich formulation of optimal transport. Other cost functions than the quadratic one $c(x,y)=|x-y|^{2}$ may be used, but we restrict the discussion to this setting to simplify the presentation. For further details on the formulation of optimal transport for probability measures support on general metric space we refer to [63].

For probability measures supported on the real (that is when $d=1$), the Wasserstein metric has a closed form expression

[TABLE]

for any $\mu,\nu\in{\mathcal{P}}_{2}({\mathbb{R}})$, where $F^{-}_{\mu}$ (resp. $F^{-}_{\nu}$) denotes the quantile function of the measure $\mu$ (resp. $\nu$). Hence, in the one-dimensional setting, the use of the Wasserstein metric amounts to perform data analysis in the space of quantile functions (supported on $[0,1]$) endowed with the usual Hilbert $\mathbb{L}_{2}$ metric. When $\mu$ and $\nu$ are absolutely continuous (a.c.) measures with square integrable probability density functions (pdf) $f_{\mu}$ and $f_{\nu}$, this remark allows to easily illustrate the difference between choosing either the Wasserstein metric $W_{2}(\mu,\nu)$ or the Hilbert metric $\|f_{\mu}-f_{\nu}\|_{2}:=\left(\int_{0}^{1}|f_{\mu}(x)-f_{\nu}(x)|^{2}dx\right)^{1/2}$ for statistical analysis. Indeed, this choice implies two different constructions of the modes of variation in a dataset as illustrated in Figure 3, where we display the geodesics between two Gaussian distributions in the Hilbert space of densities and in the Wasserstein space of probability measure. Recall that, in the one-dimensional setting, the value of the geodesic at time $t\in[0,1]$ between $\mu$ and $\nu$ (for the Wasserstein metric) is the measure $\mu_{t}$ whose quantile function is $F^{-}_{\mu_{t}}=(1-t)F^{-}_{\mu}+tF^{-}_{\nu}$. Therefore, the measure $\mu_{t}$ remains Gaussian along the geodesic path between two Gaussian measures $\mu$ and $\nu$. This illustrates the fact that performing statistical data analysis the Wasserstein space allows to simultaneously take into account variations in location and intensity (or mass) in a data set of densities (or histograms). A similar behavior holds in higher-dimensions ($d\geq 2$).

2.2 Regularized optimal transport

A serious limitation for the use of the Wasserstein metric for data analysis is its computational cost between probability measures supported in higher-dimensional spaces (that is beyond the case $d=1$). Indeed, the computational cost to evaluate a Wasserstein metric between two discrete probability distributions with supports of equal size $N$ is generally of order $\mathcal{O}(N^{3}\log N)$. In this discrete setting, computing a Wasserstein distance amounts to solve the linear program (1) whose solution is constrained to belong to the convex set $\Gamma(\mu,\nu)$. Hence, the cost of this convex minimization becomes prohibitive for moderate to large values of $N$. To reduce the complexity of a linear program, a classical approach in optimization [65] is to add a regularizing entropy term. This is the approach followed in [22, 23] by adding entropy regularization to the transport transport plan in (1), which yields the strictly convex problem (3) below.

To simplify its presentation, we consider discrete probability measures supported on the same finite space $\Omega_{N}=\{x_{1},\ldots,x_{N}\}\in({\mathbb{R}}^{d})^{N}$ of cardinal $N$. A discrete probability distribution $\mu$ (with fixed support included in $\Omega_{N}$) is identified by a vector of weights in the simplex $\Sigma_{N}=\{r=(r_{1},\ldots,r_{N})\in{\mathbb{R}}^{N}_{+}\ \mbox{with}\ \sum_{k=1}^{N}r_{k}=1\}$ such that $\mu=\sum_{k=1}^{N}r_{k}\delta_{x_{k}}$ where $\delta_{x}$ is the Dirac distribution at $x$. The Wasserstein distance between two discrete distributions $\mu=\sum_{k=1}^{N}r_{k}\delta_{x_{k}}$ and $\nu=\sum_{k=1}^{N}q_{k}\delta_{x_{k}}$ (identified to vectors $r$ and $q$ in $\Sigma_{N}$) then becomes

[TABLE]

where the set of coupling matrices is defined as $U(r,q):=\{U\in{\mathbb{R}}^{N\times N}_{+}\ \mbox{such that}\ U\mathds{1}_{N}=r,\ U^{T}\mathds{1}_{N}=q\}$ with $\mathds{1}_{N}$ the $N$ dimensional vector with all entries equal to $1$ and $C$ the cost matrix given by $C_{ml}=|x_{m}-x_{l}|^{2}$, for all $m,l\in\{1,\ldots,N\}$. Then, the notion of entropy regularized optimal transportation [22, 19] leads to the so-called Sinkhorn divergence $W_{2,\varepsilon}(r,q)$ between two discrete measures $r,q\in\Sigma_{N}$ which is defined as

[TABLE]

where $\varepsilon>0$ is a regularization parameter, and the discrete (negative) entropy for a given coupling matrix $U\in U(r,q)$ is given by $\textrm{h}(U):=-\sum_{m,\ell}U_{m\ell}\log U_{m\ell}$. The introduction of Sinkhorn divergence has been the starting point of the development of computational optimal transport [25] for machine learning as it makes feasible (from a computational point of view) the use of smoothed Wasserstein metrics for data analysis. Moreover, there now exists various toolbox and librairies to carry out data analysis using (possibly regularized) optimal transport in Matlab - https://optimaltransport.github.io/, Python - https://pot.readthedocs.io/en/stable/ and the statistical computing environment R - https://cran.r-project.org/package=Barycenter.

3 Wasserstein barycenters

The most basic statistical inference task from a data set is certainly to compute first order statistics that is to estimate an average location of the data. In this paper, we focus on the estimation of a population mean measure (or density function) from a dataset whose elements are probability measures. The notion of averaging depends on the metric that is chosen to compare elements in a given data set. In this paper, we consider the Wasserstein distance $W_{2}$ associated to the quadratic cost for the comparison of probability measures. As introduced in [2], an empirical Wasserstein barycenter $\bar{\mu}_{n}$ of set of $n$ probability measures $\nu_{1},\ldots,\nu_{n}$ (not necessarily random) in ${\mathcal{P}}_{2}(\Omega)$ is defined as a (possibly not unique) minimizer of

[TABLE]

The Wasserstein barycenter corresponds to the notion of empirical Fréchet mean [35] that is an extension of the usual Euclidean barycenter to nonlinear metric spaces. For the case $d=1$, the quantile formula (2) implies that $\bar{\mu}_{n}$ is the measure with quantile function $F^{-}_{\bar{\mu}_{n}}=\frac{1}{n}\sum_{i=1}^{n}F^{-}_{\nu_{i}}.$ For $n=2$, the Wasserstein barycenter $\bar{\mu}_{2}$ is the probability measure located at half distance (that is at “time $t=1/2$”) along the geodesic between two absolutely continuous (a.c.) measures $\nu_{1}$ and $\nu_{2}$. As illustrative examples, we display Wasserstein barycenters $\bar{\mu}_{2}$ between two probability measures in dimension $d=1$ in Figure 3 and $d=2$ in Figure 4.

Computing a Wasserstein barycenter between $n\geq 3$ probability measures is generally a delicate optimisation problem in dimension $d\geq 2$ that we do not discuss here. We refer to [25, 7, 24] for further details and references. In what follows, we focus on the mathematical statistical aspects of Wasserstein barycenters of random probability measures and the problem of their estimation from a nonparametric point of view.

3.1 Statistical models of random probability measures

Let us first introduce the notion of a population Wasserstein barycenter, as well as a two-sampling model of random probability measures that is appropriate to consider the problem of its nonparametric estimation. To this end, let ${\mathbb{P}}$ be a distribution on the space of probability measures $({\mathcal{P}}_{2}(\Omega),\mathcal{B}\left({\mathcal{P}}_{2}(\Omega)\right)$, where $\mathcal{B}\left({\mathcal{P}}_{2}(\Omega)\right)$ is the Borel $\sigma$-algebra generated by the topology induced by $W_{2}$. Throughout the paper, it is always assumed that ${\mathbb{P}}$ satisfies the square-integrability condition

[TABLE]

Definition 3.1**.**

A population Wasserstein barycenter $\mu_{{\mathbb{P}}}$ of a distribution ${\mathbb{P}}$ over ${\mathcal{P}}_{2}(\Omega)$ is defined as a (possibly not unique) minimizer of

[TABLE]

For a discrete distribution ${\mathbb{P}}_{n}=\frac{1}{n}\sum\delta_{\nu_{i}}$ on ${\mathcal{P}}_{2}(\Omega)$, one has that $\mu_{{\mathbb{P}}_{n}}$ corresponds to the empirical Wasserstein barycenter defined as a minimizer of (4). The existence and unicity of Wasserstein barycenters has been discussed in various works [2, 3, 47, 15, 43]. As argued in [3], a sufficient condition for the unicity of $\mu_{{\mathbb{P}}}$ is that the distribution ${\mathbb{P}}$ gives a strictly positive mass to the subset ${\mathcal{P}}_{2}^{ac}(\Omega)$ of a.c. measures in ${\mathcal{P}}_{2}(\Omega)$. Moreover, if ${\mathbb{P}}$ is supported on the set of measures ${\mathcal{P}}_{2}(\Omega)\cap\mathbb{L}_{q}(\Omega)$ for some $q\in(1,+\infty)$ (i.e. measures in ${\mathcal{P}}_{2}^{ac}(\Omega)$ with $\mathbb{L}_{q}(\Omega)$ densities), it follows that $\mu_{{\mathbb{P}}}$ admits a density in $\mathbb{L}_{q}(\Omega)$.

To define an appropriate sampling model, let us now assume that $\boldsymbol{\nu}_{1},\ldots,\boldsymbol{\nu}_{n}$ are independent and identically distributed a.c. random measures sampled from a distribution ${\mathbb{P}}$ satisfying ${\mathbb{P}}\left({\mathcal{P}}_{2}^{ac}(\Omega)\cap\mathbb{L}_{2}(\Omega)\right)=1$. Hence, under such assumptions, the population Wasserstein barycenter $\mu_{{\mathbb{P}}}$ is unique and absolutely continuous. As raw data in the form of densities are generally not directly available, we focus on the setting where one has access to a set of random vectors $(\boldsymbol{X}_{i,j})_{1\leq j\leq p_{i};\;1\leq i\leq n}$ in ${\mathbb{R}}^{d}$ organized in the form of $n$ experiment units, such that, conditionaly on $\boldsymbol{\nu}_{i}$, one has that $\boldsymbol{X}_{i,1},\ldots,\boldsymbol{X}_{i,p_{i}}$ are iid observations sampled from the measure $\boldsymbol{\nu}_{i}$ for each $1\leq i\leq n$. The observed data are thus in the form of multiple point clouds $\boldsymbol{X}=(\boldsymbol{X}_{i,j})_{1\leq i\leq n;\;1\leq j\leq p_{i}}.$ Elements in such a set can be modeled as $n$ discrete probability measures

[TABLE]

from which one may construct another empirical Wasserstein barycenter $\hat{\boldsymbol{\mu}}_{n,p}$ defined as a (not necessarily unique) minimizer of

[TABLE]

with the convention that $p=\min_{1\leq i\leq n}p_{i}$. We shall now discuss the issue of estimating the population Wasserstein barycenter $\mu_{{\mathbb{P}}}$ using either $\mu_{{\mathbb{P}}_{n}}$ (with ${\mathbb{P}}_{n}=\frac{1}{n}\sum\delta_{\boldsymbol{\nu}_{i}}$) or $\hat{\boldsymbol{\mu}}_{n,p}$ as $n\to+\infty$ and $p\to+\infty$.

3.2 Convergence of empirical Wasserstein barycenters

3.2.1 Law of large numbers

By the law of large numbers, we refer to the question of showing that $W_{2}(\mu_{{\mathbb{P}}_{n}},\mu_{{\mathbb{P}}})$ almost surely converges to zero as $n\to+\infty$. This problem has been addressed in [15] for a parametric model of random distributions ${\mathbb{P}}$ over ${\mathcal{P}}_{2}(\Omega)$ with $\Omega$ assumed to be compact. General results for arbitrary distributions can be found in [47]. Beyond the law of large numbers, it is natural to ask for a (functional) central limit theorem for the empirical Wasserstein barycenter. Nevertheless, this notion is more delicate to formulate as the Wasserstein space is not a Hilbert space. For recent contributions in this direction (but in specific parametric settings) we refer to [3, 45].

3.2.2 Rate of convergence in the one dimensional case

Let us now restrict to the one dimensional case of probability measure in ${\mathcal{P}}_{2}({\mathbb{R}})$ and assume that $p_{1}=\ldots=p_{n}=p$. Thanks to the quantile formula (2) for the expression of the Wasserstein metric when $d=1$, it follows that the empirical Wasserstein barycenter $\hat{\boldsymbol{\mu}}_{n,p}$ minimizing (7) is the unique discrete measure

[TABLE]

where $\boldsymbol{X}_{i,1}^{\ast}\leq\boldsymbol{X}_{i,2}^{\ast}\leq\ldots\leq\boldsymbol{X}_{i,p}^{\ast}$ are the order statistics the $i$-th sample, and $\bar{\boldsymbol{X}}_{j}^{\ast}=\frac{1}{n}\sum_{i=1}^{n}X_{i,j}^{\ast}$ for all $1\leq j\leq p$. Moreover, as shown in [13], the population Wasserstein barycenter $\mu_{{\mathbb{P}}}$ is the unique a.c. measure with quantile function

[TABLE]

In this setting, the following upper bound on the expected quadratic risk ${\mathbb{E}}\left(W_{2}^{2}\left(\hat{\boldsymbol{\mu}}_{n,p},\mu_{{\mathbb{P}}}\right)\right)$ has been obtained in [13].

Theorem 3.1**.**

Let $Y_{1},\ldots,Y_{p}$ be iid variables sampled from $\mu_{{\mathbb{P}}}$, and $\boldsymbol{\mu}_{p}=\frac{1}{p}\sum_{k=1}^{p}\delta_{Y_{k}}$ Then, the estimator $\hat{\boldsymbol{\mu}}_{n,p}$ satisfies

[TABLE]

Hence, the rate of convergence of $\hat{\boldsymbol{\mu}}_{n,p}$ is decomposed into the sum of two terms. The first one goes to zero at the parametric rate $\mathcal{O}(\frac{1}{n})$, while the second one depends on the rate of convergence of the empirical measure $\boldsymbol{\mu}_{p}$ towards it population counterpart in expected squared Wasserstein distance. For $d=1$, the question of deriving upper and lower bound for ${\mathbb{E}}\left(W^{2}_{2}(\boldsymbol{\mu}_{p},\mu_{{\mathbb{P}}})\right)$ has been thoroughly studied in [16]. In arbitrary dimension $d$, there is a rich literature in probability on studying the convergence rate of an empirical measure towards it population version in Wasserstein distance. We refer to [16, 34, 64, 27, 9] for extensive references and the most recent contributions on this topic. For example, using results from [16], one has that

[TABLE]

where $F_{0}$ is the cumulative distribution function of $\mu_{{\mathbb{P}}}$ and $f_{0}$ denotes its density. Hence, if $J_{2}(\mu_{{\mathbb{P}}})<+\infty$, it follows that $\hat{\boldsymbol{\mu}}_{n,p}$ converges to zero at the parametric rate $\mathcal{O}(\frac{1}{n}+\frac{1}{p})$. A detailed discussion on the derivation of upper bounds on the rate of converge of $\hat{\boldsymbol{\mu}}_{n,p}$ is proposed in [13] with some results on minimax optimality of such bounds.

3.2.3 Open problems in dimension $d\geq 2$

Beyond the one-dimensional case $d=1$, it appears that deriving the rate of convergence for empirical Wasserstein barycenters, e.g. for the expected quadratic risk ${\mathbb{E}}\left(W_{2}^{2}\left(\hat{\mu}_{{\mathbb{P}}_{n}},\mu_{{\mathbb{P}}}\right)\right)$, is a much more involved problem. A few works [45, 1] exist on the convergence rate of $\mu_{{\mathbb{P}}_{n}}$ to $\mu_{{\mathbb{P}}}$. However, the more general question of deriving the convergence rate of a Wasserstein barycenter $\hat{\boldsymbol{\mu}}_{n,p}$ (7) obtained from samples of random measure is open for $d\geq 2$ (to the best of our knowledge).

3.3 Regularization of Wasserstein barycenters

An empirical Wasserstein barycenter $\hat{\boldsymbol{\mu}}_{n,p}$ of the discrete measures $\boldsymbol{\nu}_{p_{1}},\ldots,\boldsymbol{\nu}_{p_{n}}$, defined by (6), is generally irregular (and even not unique). Irregularity has to be understood in the sense that $\hat{\boldsymbol{\mu}}_{n,p}$ is, in most situations, a discrete probability measure as shown by its expression (8) when $d=1$. Therefore, as a discrete measure, it poorly represents the smoothness of the population Wasserstein barycenter of a random a.c. measure $\boldsymbol{\nu}$ with distribution ${\mathbb{P}}$, as $\mu_{{\mathbb{P}}}$ is absolutely continuous in this setting as discussed previously. Hence, we now explain how regularizing the Wasserstein barycenter $\hat{\boldsymbol{\mu}}_{n,p}$ in order to construct a.c. and consistent estimators of an a.c. population barycenter in the asymptotic setting where both $n$ and $p=\min_{1\leq i\leq n}p_{i}$ tend to infinity. To this end, one may consider two ways of regularizing an empirical barycenter.

3.3.1 Penalized Wasserstein barycenters

A first possibility is to follow the approach proposed in [12]. Thus, we now introduce the penalized Wasserstein barycenter associated to the distribution ${\mathbb{P}}$ that is defined as a solution of the minimization problem

[TABLE]

where $\gamma>0$ is a penalization parameter and the penalty function writes

[TABLE]

where $\|\cdot\|_{H^{k}(\Omega)}$ denotes the Sobolev norm associated to the $\mathbb{L}^{2}(\Omega)$ space, $\alpha>0$ is arbitrarily small and $k>d-1$. The choice (10) for the penalty function $E$ is motivated by the discussion in Section 5 of [12]. Its main advantages are to impose that the penalized Wasserstein barycenter is absolutely continuous with a smooth density. Other examples of penalty functions are discussed in [12] including the class of relative G-functional described in Section 9.4 in [5].

Definition 3.2**.**

For $\gamma>0$, we consider the Wasserstein barycenters:

[TABLE]

called respectively penalized population barycenter (11) and penalized empirical barycenter (12).

Existence and uniqueness of these barycenters are discussed in [12] for a large class of penalty functions enforcing them to be a.c. One can thus let $\hat{\boldsymbol{f}}^{\gamma}_{n,p}$ (resp. $f^{\gamma}_{{\mathbb{P}}}$) be the densities associated to $\hat{\boldsymbol{\mu}}^{\gamma}_{n,p}$ (resp.  $\mu^{\gamma}_{{\mathbb{P}}}$). The choice (10) for the penalty function $E$ imposes that the densities of the penalized Wasserstein barycenters are bounded from below by a small constant $\alpha>0$ over their support that is equal to $\Omega$. In this setting, the following result (see Section 5 in [12]) gives the convergence rate in expected squared $\mathbb{L}_{2}(\Omega)$-distance between these densities for a fixed value of $\gamma>0$.

Theorem 3.2** (Section 5 in [12]).**

Assume that $\Omega\subset{\mathbb{R}}^{d}$ is compact. Let $\hat{\boldsymbol{f}}_{n,p}^{\gamma}$ and $f_{{\mathbb{P}}}^{\gamma}$ denote the density functions of $\hat{\boldsymbol{\mu}}_{n,p}^{\gamma}$ and $\mu_{{\mathbb{P}}}^{\gamma}$, induced by the choice (10) of the penalty function $E$. Then, provided that $d<4$, there exists a constant $c>0$ depending only on $\Omega$ such that

[TABLE]

where $p=\min_{1\leq i\leq n}p_{i}$.

Letting $f_{{\mathbb{P}}}^{0}$ be the density of the population Wasserstein barycenter $\mu_{{\mathbb{P}}}$, the convergence of the approximation error term $\|f_{{\mathbb{P}}}^{\gamma}-f_{{\mathbb{P}}}^{0}\|^{2}_{\mathbb{L}_{2}(\Omega)}$ is also discussed in [12], where it is shown that it converges to zero as $\gamma\to 0$. However, no rate of convergence (as a function of $\gamma$) has been derived in [12] for this approximation term, and deriving such a result remains an open problem for regularized Wasserstein barycenters. Therefore, only the consistency of $\hat{\boldsymbol{f}}_{n,p}^{\gamma}$ has been derived in [12] in the sense that the quadratic risk ${\mathbb{E}}\left(\|\hat{\boldsymbol{f}}_{n,p}^{\gamma}-f_{{\mathbb{P}}}^{0}\|^{2}_{\mathbb{L}_{2}(\Omega)}\right)\to 0$ under the assumption that $\gamma=\gamma_{n,p}$ decays to zero as $\min(n,p)\to+\infty$ at a rate which guarantees that the right hand size of inequality (13) converges to zero. In Theorem 3.2, the condition $d<4$ may be relaxed at the price of obtaining a slower rate of convergence as $p\to+\infty$ (see Section 5 in [12]).

Remark 3.1**.**

Consistency of regularized empirical Wasserstein barycenter is discussed in [12] for the Bregman divergence associated to the penalization function $E$. The specific choice (10) for $E$ finally yields to the use of the $\mathbb{L}_{2}(\Omega)$ norm to compare the densities of $\hat{\boldsymbol{f}}_{n,p}^{\gamma}$, $f_{{\mathbb{P}}}^{\gamma}$ and $f_{{\mathbb{P}}}^{0}$. The question of deriving consistency results for $\hat{\boldsymbol{\mu}}^{\gamma}_{n,p}$ using the Wasserstein metric (that is in the space of probability measures and not densities) has not been considered in [12]

3.3.2 Sinkhorn barycenters

Another way to regularize an empirical Wasserstein barycenter is to use the Sinkhorn divergence (3) to compare probability measures. As this divergence is defined for discrete measures supported on the same points, one first needs to choose a fixed grid $\Omega_{N}=\{x_{1},\ldots,x_{N}\}$ made of $N$ points $x_{k}\in\mathbb{R}^{d}$ (bin locations that are typically equally spaced). Then, one performs a binning of the data (6) on this grid, leading to a dataset of discrete measures (with supports included in $\Omega_{N}$) that we denote

[TABLE]

for $1\leq i\leq n$. Binning (i.e. choosing the grid $\Omega_{N}$) surely incorporates some sort of additional regularization, and a discussion on the influence of the grid size $N$ on the smoothness of the barycenter can be found in [11]. The size $N$ of the grid is also guided by numerical issues on the computational cost of the algorithms used to approximate a Sinkhorn divergence.

Then, another possibility to introduce regularization in the computation of Wasserstein barycenters [11, 23, 24] is to consider the following estimator

[TABLE]

that can be interpreted as a Fréchet mean with respect to a Sinkhorn divergence. We call $\hat{\boldsymbol{r}}^{\varepsilon}_{n,p}$ a Sinkhorn barycenter. From the point of view of statistics, it can be interpreted as a M-estimator defined over the compact set $\Sigma_{N}$. Two key properties of the Sinkhorn divergence have been used in [11] to study the convergence properties of such barycenters:

(i)

for any $q\in\Sigma_{N}$, the function $r\mapsto W_{2,\varepsilon}^{2}(r,q)$ is $\varepsilon$- strongly convex for the Euclidean 2-norm

(ii)

let $q\in\Sigma_{N}$ and $0<\rho<1/N$ an arbitrarily small constant. Then, one has that $r\mapsto W_{2,\varepsilon}^{2}(r,q)$ is $L_{\rho,\varepsilon}$-Lipschitz on

[TABLE]

with

[TABLE]

For a discussion on the behavior of the Lipschitz constant $L_{\rho,\varepsilon}$ as $\rho$ and $\varepsilon$ tend to zero, we refer to Section 3.1 in [11]. For example, if $\rho=\varepsilon^{\kappa}$ for some $\kappa>0$, then

[TABLE]

To guarantee the Lipschitz continuity of the mapping $r\mapsto W_{2,\varepsilon}^{2}(r,q)$, the analysis of the statistical properties of Sinkhorn barycenters in [11] is restricted to discrete measures belonging the convex set $\Sigma_{N}^{\rho}$ (thus having non-vanishing entries). Therefore, one has to slightly modify the data at hand so that the observations belong to $\Sigma_{N}^{\rho}$. This yields to the following definitions of sampling model, and empirical/population Sinkhorn barycenters in $\Sigma_{N}^{\rho}$.

Definition 3.3**.**

Let $0<\rho<1/N$, and ${\mathbb{P}}$ be a probability distribution on $\Sigma_{N}^{\rho}$. Let $\boldsymbol{q}_{1},\ldots,\boldsymbol{q}_{n}\in\Sigma_{N}^{\rho}$ be an iid sample drawn from the distribution ${\mathbb{P}}$. Assume that $(\tilde{X}_{i,j})_{1\leq j\leq p_{i}}$ are iid random variables sampled from $\boldsymbol{q}_{i}$, for each $1\leq i\leq n$ and consider the discrete measures $\tilde{\boldsymbol{q}}_{i}^{p_{i}}=\frac{1}{p_{i}}\sum_{j=1}^{p_{i}}\delta_{\tilde{X}_{i,j}}$ and $\hat{\boldsymbol{q}}_{i}^{p_{i}}=(1-\rho N)\tilde{\boldsymbol{q}}_{i}^{p_{i}}+\rho\hbox to0.0pt{1\hss}\kern 1.60004pt\hbox to0.0pt{1\hss}\kern 3.99994pt_{N},$ where $\hbox to0.0pt{1\hss}\kern 1.60004pt\hbox to0.0pt{1\hss}\kern 3.99994pt_{N}$ is the vector of ${\mathbb{R}}^{N}$ with all entries equal to one. Then, we define

[TABLE]

Theorem 3.3** (Section 3 in [11]).**

Let $\varepsilon>0$ and $p=\min_{1\leq i\leq n}p_{i}$. Then, one has that

[TABLE]

The upper bound (18) allows to shed some light on how the expected squared distance between the empirical Sinkhorn barycenter and its population version is influenced by the number $n$ of observed measures, the minimal number $p$ of samples per measure and the size $N$ of the grid. However, this upper bound does not converge to zero as $\min(n,p)\to+\infty$ if $\rho$ remains fixed. The main advantage of this upper bound is to allow the derivation of a data-driven strategy to select $\varepsilon$ that is presented in the following subsection below. Finally, one should remark that the behavior of the population Sinkhorn barycenter $r^{\varepsilon}$ as $\varepsilon\to 0$ has not been considered in [11].

Remark 3.2**.**

A third way to regularize an empirical Wasserstein barycenter (that might look more natural for statisticians) is to perform a preliminary smoothing step of the discrete measures $\boldsymbol{\nu}_{p_{1}},\ldots,\boldsymbol{\nu}_{p_{n}}$. For example, one may first smooth the data using standard kernel density estimation $\hat{\boldsymbol{f}}^{h_{i}}_{i}(x)=\frac{1}{p_{i}h_{i}^{d}}\sum_{j=1}^{p_{i}}K\left(\frac{x-\boldsymbol{X}_{i,j}}{h_{i}}\right)$ and then define a smoothed Wasserstein barycenter as

[TABLE]

where $d\hat{\boldsymbol{\nu}}_{i}^{h_{i}}(x)=\hat{\boldsymbol{f}}^{h_{i}}_{i}(x)dx$, and $h=\max_{1\leq i\leq}h_{i}$. The consistency of this approach (as $h\to 0$) has been studied in details by [49, 50] using multiple point processes for the sampling procedure, and it is also discussed in [13].

3.4 Data-driven regularization in computational optimal transport

In this paper, we report results on numerical experiments from [11] that are based on the approach proposed in [24] to compute Sinkhorn barycenters on a grid $\Omega_{N}$ of equi-spaced points. We do not report results on the computational aspects of penalized Wasserstein barycenters, and we refer to [11] for illustrations on their numerical performances.

As any nonparametric method, the use of regularized Wasserstein barycenters involves the delicate calibration of a regularization parameter. Choosing the value of $\varepsilon$ for the Sinkhorn barycenter $\hat{r}^{\varepsilon}_{n,p}$ might be guided by the classical bias and variance tradeoff principle as illustrated by the synthetic example displayed in Figure 5. For small values of $\varepsilon$ the Sinkhorn barycenter shows many oscillations while it is much smoother for larger values of $\varepsilon$. In some sense, the regularization parameters $\varepsilon$ can be interpreted as the usual bandwidth parameter in kernel density estimation and its choice greatly influences the shape of $\hat{r}^{\varepsilon}_{n,p}$.

Let us now briefly discuss the question of choosing $\varepsilon$ in an automatic way. In [11], the following (un-regularized) population Wasserstein is considered

[TABLE]

The term ${\mathbb{E}}(|r^{\varepsilon}-\hat{r}^{\varepsilon}_{n,p}|^{2})$ can then be interpreted as “variance term”, while $|r^{\varepsilon}-r^{0}|$ is referred to as a bias term. Ideally, one would like to select a value of $\varepsilon$ which gives a good tradeoff between these two terms whose values change in opposite directions as $\varepsilon$ varies. However, the decay of the bias term (as $\varepsilon\to 0$) is typically unknown and only an estimation of the variance term is available. In [11], it is thus proposed to use the upper bound (18) to automatically calibrate the parameter $\varepsilon>0$ by following Goldenshluger-Lepski (GL) principle [39] as formulated in [46] for standard kernel density estimation. In [11], an adaptation of the GL’s principle is used to to estimate a bias-variance trade-off functional that is minimized over a finite set of regularization parameters to obtain an automatic selection of $\varepsilon$ for Sinkhorn barycenters. The method consists in comparing estimators pairwise (for a given range of regularization parameters) with respect to a loss function, and we refer to [11] for further details. As an illustration of the usefulness of the GL’s principle in statistical optimal transport, we report in Figure 6 numerical experiments from [11] on the computation of a Sinkhorn barycenter (with a data-driven choice of $\varepsilon$) in dimension $d=2$ for the flow cytometry dataset displayed in Figure 2(b). This regularized barycenter clearly corrects some mis-alignment issues in these measurements. We also display in Figure 6(a) the Euclidean mean of this dataset (after a preliminary kernel smoothing step). The support of this Euclidean mean is more spread out due to the presence of a strong variance in spatial location in the data from one individual to another.

4 Geometry of the Wasserstein space and geodesic PCA

Principal Component Analysis (PCA) is certainly the most widely used approach to reduce the dimension of datasets whose elements are high-dimensional vectors or matrices. When such elements can be modeled as probability measures, we explain how using the Wasserstein metric to compute their modes of variation accordingly around their barycenter in the Wasserstein space. We briefly describe the setting where the data can be modeled as histograms supported on the real line, and we refer to extensions in the case of probability measures supported on ${\mathbb{R}}^{d}$ for $d\geq 2$.

4.1 Geometry of the Wasserstein space

The space ${\mathcal{P}}_{2}(\Omega)$ endowed with the 2-Wasserstein distance is not a Hilbert space. As a consequence, standard PCA in a linear space based on the diagonalisation of a covariance matrix cannot be applied directly to compute principal modes of variations with respect to the Wasserstein metric. Nevertheless, a meaningful notion of PCA [58, 14, 20] can still be defined by relying on the pseudo-Riemannian structure of the Wasserstein space which was extensively studied in [5].

This geometric structure allows to define a notion of geodesic PCA in the Wasserstein space which shares similarities with analogs of PCA for data belonging to a Riemannian manifolds. Principal Geodesic Analysis (PGA) [33, 61] is a geometric method for statistical analysis that extends the notion of classical PCA in Hilbert spaces [26, 53] to curved Riemannian manifolds by replacing Euclidean concepts of vector means, lines and orthogonality by the more general notions of Fréchet mean, geodesics in Riemannian manifolds and orthogonality in tangent spaces. For $d=1$, the geometric structure of the Wasserstein space is relatively easy to describe, and we refer to [14, 20] for a detailed presentation. The setting $d\geq 2$ is more involved. Some elements on the geometry of the Wasserstein space can be found in Section 4.4 in [50], and an extensive study can be found in [5].

4.2 Geodesic PCA in the one-dimensional case

Following these principles of statistical analysis in manifolds, a framework for geodesic PCA (GPCA) of probability measures supported on a interval $\Omega\subset{\mathbb{R}}$ has been introduced in [14]. GPCA is defined as the problem of estimating a principal geodesic subspace (of a given dimension) which maximizes the variance of the projection of the data to that subspace whose base point is the Wasserstein barycenter of the observations. In [14], a detailed study is proposed on the existence and consistency of GPCA. It is also shown that GPCA in the Wasserstein space is equivalent to map the data in the tangent space at the empirical Wasserstein barycenter (assumed to be a.c.) and then to perform a PCA that is constrained to lie in a convex and closed subset of functions. Mapping the data to this tangent space is not difficult in the one-dimensional case as it amounts to computing a set of optimal maps between the data and their Wasserstein barycenter, for which a closed form is available using their quantile functions. In particular, optimal maps belong to the set of non-decreasing functions. This convex constrained PCA can also be understood from the quantile formula (2) for the Wasserstein metric when $d=1$. Indeed, GPCA for probability measures supported on $\Omega\subset{\mathbb{R}}$ corresponds to a PCA of quantile functions that is constrained to lie in the convex subset of non-decreasing function in $\mathbb{L}_{2}([0,1])$. Nevertheless, such a convex PCA leads to the minimization of a non-convex and non-differentiable objective function. This is a delicate optimisation problem (even for $d=1$), and numerical methods to compute such a PCA have been proposed in [20].

To shed some lights on the benefits of GPCA for data analysis, we report numerical experiments on synthetic data in Figure 7 and for the dataset on children’s first name at birth in Figure 8. When the size or locations of significant bins in observed histograms significantly vary from one histogram to another, standard PCA with respect to the Euclidean metric (in the space of densities) fail in recovering geometric variations in location. To the contrary, GPCA in the Wasserstein space is able to capture meaningful sources of variations in a dataset by separating variability in location and scaling in the observed histograms along the first and second geodesic principal components.

4.3 Geodesic PCA in dimension $d\geq 2$

The study of Geodesic PCA of histograms in the Wasserstein space is much more involved in dimension $d\geq 2$, both on the theoretical and numerical aspects. Indeed, for $d\geq 2$, one could always define GPCA as the construction of a principal geodesic subspace maximizing the variance of the projection of the data to that subspace. However, a rigorous analysis of such a construction in dimension $d\geq 2$ has not been proposed so far. There also exist two main differences with the one-dimensional case. First, GPCA cannot be equivalent to map the data in the tangent space at their barycenter (assumed to be a.c.) and then to perform a standard PCA in that linear space as the Wasserstein space is a curved manifold for $d\geq 2$. Secondly, optimal maps sending a set of histograms to such a tangent space are more difficult to compute for $d\geq 2$. Moreover, it is well known [18, 21] that such maps are equal to the gradient of convex functions. This implies that an adaptation of existing algorithms for $d=1$ leads to optimisation problems with more difficult constraints to handle than a PCA constrained to lie in the convex set of non-decreasing functions.

Nevertheless, numerical approximation to the construction of GPCA in dimension $d=2$ have been proposed in [58, 20]. As an illustrative example, we report numerical experiments from [20] on the analysis of data from the MNIST database [48] which contains grayscale images of handwritten digits. All grayscale images in this dataset are of size $28\times 28$ pixels. We normalize them so that the sum of pixel grayscale values sum to one. In this manner, the images are histograms supported on a 2D grid of size $28\times 28$. The ground metric for the Wasserstein distance is then the squared Euclidean distance between the locations of bins in this 2D grid. In Figure 9, we display the first principal components computed from 1000 images of each digit using the algorithmic approach proposed in [20]. It can be seen that GPCA captures very well the main source of variability of each digit such as changes in the shape of the ’0’ or the presence or absence of the lower loop of the ’2’.