# Measure Estimation in the Barycentric Coding Model

Matthew Werenski  Ruijie Jiang  Abiy Tasissa  Shuchin Aeron  James M. Murphy

###### Abstract

This paper considers the problem of measure estimation under the _barycentric
coding model (BCM)_ , in which an unknown measure is assumed to belong to the
set of Wasserstein-2 barycenters of a finite set of known measures. Estimating
a measure under this model is equivalent to estimating the unknown barycentric
coordinates. We provide novel geometrical, statistical, and computational
insights for measure estimation under the BCM, consisting of three main
results. Our first main result leverages the Riemannian geometry of
Wasserstein-2 space to provide a _procedure for recovering the barycentric
coordinates as the solution to a quadratic optimization problem_ assuming
access to the true reference measures. The essential geometric insight is that
the parameters of this quadratic problem are determined by inner products
between the optimal displacement maps from the given measure to the reference
measures defining the BCM. Our second main result then establishes an
_algorithm for solving for the coordinates in the BCM when all the measures
are observed empirically via i.i.d. samples_. We prove precise _rates of
convergence_ for this algorithm—determined by the smoothness of the underlying
measures and their dimensionality—thereby guaranteeing its statistical
consistency. Finally, we demonstrate the utility of the BCM and associated
estimation procedures in three application areas: (i) _covariance estimation_
for Gaussian measures; (ii) _image processing_ ; and (iii) _natural language
processing_.

…

##  1 Introduction

A number of recent machine learning applications including computer vision
(Schmitz et al., 2018; Bonneel et al., 2016), domain adaptation (Montesuma &
Mboula, 2021; Redko et al., 2019), natural language processing (NLP) (Singh et
al., 2020; Colombo et al., 2021; Xu et al., 2018), and unsupervised
segmentation of multivariate time series data (Cheng et al., 2021) have shown
the utility of representing and modeling high-dimensional data as probability
distributions. The essential insight in these applications is to utilize the
Riemannian geometry of the space of probability distributions induced by the
Wasserstein-2 metric and the Wasserstein-2 barycenters as alternatives to the
Euclidean distance and linear combinations, respectively.

In this context, the methods that use Wasserstein barycenters for data
modeling and inference involve solving two core problems which we term the
_synthesis problem_ and _analysis problem_. To be precise, let
$\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ denote the space of
absolutely continuous distributions on $\mathbb{R}^{d}$ with finite second
moment. For $\mu,\nu\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$,

$W_{2}^{2}(\mu,\nu)\triangleq\min_{T\\#\mu=\nu}\int_{\mathbb{R}^{d}}||T(x)-x||_{2}^{2}d\mu(x)$
(1)

defines the _Wasserstein-2 metric_ where the minimum is over all measurable
maps $T:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ and the _pushforward_
$T\\#\mu=\nu$ is such that for all Borel sets $B$ we have
$\nu[B]=\mu[T^{-1}(B)]$ (Villani, 2003; Santambrogio, 2015). Under the
assumption $\mu,\nu\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$,
unique minimizing $T$ are known to exist and the resulting $W_{2}$ distance is
finite (Villani, 2003). We refer to the space
$\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ equipped with the metric
$W_{2}$ as the _Wasserstein-2 space_ , denoted $\mathcal{W}_{2}$.

Let
$\Delta^{p}=\left\\{\lambda=(\lambda_{1},\dots,\lambda_{p})\in\mathbb{R}^{p}:\lambda_{i}\geq
0,\sum_{i=1}^{p}\lambda_{i}=1\right\\}$ and let $\\{\mu_{i}\\}_{i=1}^{p}$ be
known _reference measures_. The _$W_{2}$ barycenter_ for coordinates
$\lambda\in\Delta^{p}$ with respect to $\\{\mu_{i}\\}_{i=1}^{p}$ is defined as

$\nu_{\lambda}\triangleq\operatorname*{arg\,min}_{\nu\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}}\frac{1}{2}\sum_{i=1}^{p}\lambda_{i}W_{2}^{2}(\nu,\mu_{i}).$
(2)

The measure $\nu_{\lambda}$ is an advection of $\\{\mu_{i}\\}_{i=1}^{p}$ and
blends their geometric features; it may be thought of as a “displacement
interpolation" (McCann, 1997; Villani, 2003). In contrast, the linear mixture
$\sum_{i=1}^{p}\lambda_{i}\mu_{i}$ simply sets $\lambda_{i}$ to be the rate we
draw samples from $\mu_{i}$. Figure 1 demonstrates the important difference
between linear mixture models and Wasserstein barycenters on Gaussian
distributions: the linear mixture is a standard Gaussian mixture model
(Murphy, 2012), while the Wasserstein barycenter is itself Gaussian. The
important takeaway is that Wasserstein barycenters offer a more geometrically
meaningful interpolation between measures compared to mixtures, and as
indicated above have been successfully leveraged in a number of applications.

Now, let
$\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})\triangleq\left\\{\nu_{\lambda}:\lambda\in\Delta^{p}\right\\}$
denote the set of all of possible barycenters of $\\{\mu_{i}\\}_{i=1}^{p}$.
The _synthesis_ problem (2) seeks to combine the reference measures
$\\{\mu_{i}\\}_{i=1}^{p}$ in a notion of weighted average, with relative
contributions determined by $\lambda$. This problem is well-studied (McCann,
1997; Agueh & Carlier, 2011; Kroshnin et al., 2019) and known to have a unique
solution under mild assumptions on the $\\{\mu_{i}\\}_{i=1}^{p}$. Moreover,
there exist fast algorithms for computing or approximating $\nu_{\lambda}$
when samples from the $\\{\mu_{i}\\}_{i=1}^{p}$ are available (Rabin et al.,
2011; Cuturi & Doucet, 2014; Bonneel et al., 2015; Claici et al., 2018; Yang
et al., 2021), which have spurred the application of Wasserstein barycenters
to machine learning problems.

In this paper, our focus is on the less well-studied _analysis_ problem: given
$\mu_{0}$ and reference measures $\\{\mu_{i}\\}_{i=1}^{p}$, solve

$\operatorname*{arg\,min}_{\lambda\,\in\Delta^{p}}W_{2}^{2}\left(\mu_{0},\nu_{\lambda}\right).$
(3)

Instead of _combining_ the measures $\\{\mu_{i}\\}_{i=1}^{p}$ according to the
weights $\lambda\in\Delta^{p}$ as in (2), the analysis problem (3)
_decomposes_ a given measure $\mu_{0}$ into its optimal representation in the
set $\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})$ as parameterized by the
learned barycentric coordinates $\lambda\in\Delta^{p}$. If the measure
$\mu_{0}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})$, then
$\min_{\lambda\,\in\Delta^{p}}W_{2}^{2}(\mu_{0},\nu_{\lambda})=0$ and solving
(3) is equivalent to learning the barycentric coordinates
$\lambda\in\Delta^{p}$ such that $\nu_{\lambda}=\mu_{0}$. We refer to this
general model of parameterizing measures
$\mu_{0}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})$ through the
associated coordinates $\lambda\in\Delta^{p}$ as the _barycentric coding
model_ (BCM).

Figure 1: _Top four_ : reference Gaussians. _Bottom left_ : linear mixture
with $\lambda=(\frac{1}{4},\frac{1}{4},\frac{1}{4},\frac{1}{4})$. _Bottom
right_ : Wasserstein barycenter with same $\lambda$. (Intensities adjusted to
show detail.)

Main Contributions: This paper addresses the analysis problem (3) as a measure
estimation problem under the BCM. For this general problem, we make the
following geometrical, statistical, and computational contributions:

  1. (1)

When all the measures are directly observed, Theorem 1 provides a _geometric_
characterization of the solution to (3) as corresponding to a solution to a
quadratic program. Our analysis leverages the Riemannian geometry of
$\mathcal{W}_{2}$ space and establishes that the quadratic program is
determined by the angles between the optimal displacement maps between the
measures. When specialized to the case of Gaussian measures, Corollary 1
provides a closed-form characterization of the solution to the analysis
problem.

  2. (2)

When the measures are indirectly observed via i.i.d. samples, we provide an
algorithm for estimating the barycentric coordinates $\lambda$ of a given
measure with respect to the reference measures. Theorem 3 gives precise rates
of convergence of the angles between the optimal displacement map by a
computationally tractable estimator that exploits entropic regularization,
thereby establishing the statistical consistency of our approach in Corollary
2.

  3. (3)

We showcase the utility of the BCM and the proposed algorithm on several data
analysis problems: covariance estimation for Gaussian measures; image
inpainting and denoising; and document classification from limited labeled
data. We show that the BCM provides a simple yet effective approach that
outperforms competing methods in the regimes considered.111Code to reproduce
results is available at https://github.com/MattWerenski/BCM

###  1.1 Related Work

The synthesis problem (2) is well-studied. (Agueh & Carlier, 2011) propose
$W_{2}$ barycenters for a finite set of reference measures
$\\{\mu_{i}\\}_{i=1}^{p}$ and prove their existence and uniqueness in a very
general setting. (Le Gouic & Loubes, 2017) extend this to barycenters
constructed from infinitely many reference measures. Statistical rates of
estimation in the setting of i.i.d. samples from the measures
$\\{\mu_{i}\\}_{i=1}^{p}$ are established in (Kroshnin et al., 2019). Building
on fundamental advances in estimating transport maps between measures using
entropic regularization (Cuturi, 2013), fast estimation methods for solving
the synthesis problem have been developed for entroptically regularized
barycenters (Cuturi & Doucet, 2014). Approaches to computing a barycenter
based on alternating direction of multipliers on the dual problem (Yang et
al., 2021) and parallelization schemes (Claici et al., 2018) have also been
considered.

The literature on the analysis problem (3) is much sparser. To the best of our
knowledge, our geometrical characterization of solutions to the analysis
problem (Theorem 1) and sample complexity results for estimating the optimal
coordinates $\lambda$ (Theorem 3, Corollary 2) are the first of their kind.
Note that our approach extends (Bonneel et al., 2016; Schmitz et al.,
2018)—which consider measure estimation under the BCM but only for measures
with finite support—in three fundamental ways. First, we enable the
application of the BCM to absolutely continuous measures. Second, we prove
precise theoretical characterizations of the solutions to (3). Third, we avoid
the expensive procedure of differentiating through an iterative algorithm as
required in (Bonneel et al., 2016) and instead provide the direct Algorithm 1
that enjoys statistical consistency (Corollary 2). At the level of technical
analysis, our Theorem 1 relies on an alternative characterization of $W_{2}$
barycenters that leverages known conditions under which Karcher means are
guaranteed to be $W_{2}$ barycenters (Panaretos & Zemel, 2020).

For the special case when all the measures are zero-mean Gaussians, the
_synthesis_ problem is again well-studied, with a fixed-point algorithm for
computing the barycenter given in (Álvarez-Esteban et al., 2016) that is shown
to have a linear convergence rate in (Chewi et al., 2020; Altschuler et al.,
2021). The corresponding _analysis_ problem for the specific case of Gaussian
measures has only been recently studied in (Musolas et al., 2021) but limited
to the case when $p=2$. When specialized to this case, our results extend the
model considered in (Musolas et al., 2021) to the $p\geq 3$ case and
characterize solutions to barycenter parameterized covariance estimation.

In Section 4, we use the proposed BCM framework to model and address several
problems in image processing and natural language processing. In this context
our work is related to a long lineage of methods in signal processing and
machine learning that use coefficients in a basis or a dictionary for data
modeling and processing (Mallat, 1999; Donoho, 2006; Tošić & Frossard, 2011).
Our approach of representing a general measure in terms of its barycentric
coordinates with respect to a fixed set of reference measures is a novel
approach in this tradition, and our results in Section 4 suggest its utility
for data modeled as probability measures.

##  2 Convex Optimization for the Analysis Problem

To solve the analysis problem (3), we propose to solve a convex optimization
problem based on an alternative characterization of $W_{2}$-barycenters as
minimizers of a certain functional. We leverage this alternative
characterization to show in Theorem 1 that, under mild assumptions on
$\\{\mu_{i}\\}_{i=1}^{p}$, one can use the optimal value of a convex program
to both check if $\mu_{0}\in\operatorname*{Bary}(\\{\mu_{i}\\})_{i=1}^{p}$ and
if so, recover its barycentric coordinate. Before proceeding, we must
introduce some background on the Riemannian geometry of the $\mathcal{W}_{2}$
space. The geodesic curve from $\nu_{0}$ to $\nu_{1}$ in $\mathcal{W}_{2}$ is
given by $[\nu_{t}]_{t=0}^{1}$ where
$\nu_{t}=(tT+(1-t)\operatorname*{Id})\\#\nu_{0}$ where $T$ is the optimal
transport map from $\nu_{0}$ to $\nu_{1}$ in (1) (Santambrogio, 2015). These
facts together with Brenier’s Theorem (see Supplementary Material Section 1)
motivate the definition of the tangent space at $\nu$:

$\displaystyle T_{\nu}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$
$\displaystyle\triangleq$
$\displaystyle\overline{\\{\beta(\nabla\varphi-\operatorname*{Id}):\beta>0,\varphi\in
C_{c}^{\infty}(\mathbb{R}^{d}),\varphi\text{ convex}\\}},$

where the closure is in $L^{2}(\nu)$ (Ambrosio et al., 2005). For $u,v\in
T_{\nu}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$, the inner product
is defined by $\langle u,v\rangle_{\nu}\triangleq\int_{\mathbb{R}^{d}}\langle
u(x),v(x)\rangle d\nu(x),$ which also gives the standard norm of $u\in
T_{\nu}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$,
$||u||_{\nu}\triangleq\sqrt{\langle u,u\rangle_{\nu}}$.

###  2.1 Alternative Characterization of Barycenters

Given the reference measures $\\{\mu_{i}\\}_{i=1}^{p}$ and barycentric
coordinates $\lambda\in\Delta^{p}$, the _variance functional_ is
$G_{\lambda}:\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}\rightarrow\mathbb{R}$
defined by

$\displaystyle
G_{\lambda}(\nu)\triangleq\sum_{i=1}^{p}\frac{\lambda_{i}}{2}W_{2}^{2}(\nu,\mu_{i}),$

so that the barycenter for $\\{\mu_{i}\\}_{i=1}^{p}$ and
$\lambda\in\Delta^{p}$ is

$\nu_{\lambda}=\displaystyle\operatorname*{arg\,min}_{\nu\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}}G_{\lambda}(\nu).$

In many unconstrained optimization problems, the most direct way of checking
if a point is optimal is to use first-order optimality conditions: compute the
gradient at the point and check if it is zero. Without further assumptions
this is a necessary, but not sufficient, condition for optimality. Applying
this to a general variance functional, one arrives at the definition of a
_Karcher mean_ (Karcher, 1977).

######  Definition 1.

Let $(X,\rho)$ be a metric space furnished with a tangent space at every $x\in
X$. For a collection of points $\\{x_{i}\\}_{i=1}^{p}\subset X$ and
coordinates $\lambda\in\Delta^{p}$, a point $y\in X$ is said to be a _Karcher
mean_ of the set for $\lambda$ if
$G_{\lambda}=\sum_{i=1}^{p}\lambda_{i}\rho^{2}(\cdot,x_{i})$ is differentiable
at $y$ and $||\nabla G_{\lambda}(y)||_{y}=0$ where $\nabla G_{\lambda}(y)$ is
understood as an element of the tangent space at $y$.

Figure 2: Visualization of the tangent space
$T_{\nu}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ and $\nabla
G_{\lambda}$. The red arc is the path along which $G_{\lambda}$ is maximized,
starting at $\nu$ and the green arrow points in that direction in the tangent
space at $\nu$.

Any barycenter $y$ is a global minimizer of $G_{\lambda}$. Using first order
optimality conditions, this implies $||\nabla G_{\lambda}(y)||_{y}=0$, and
therefore $y$ must be a Karcher mean. Special care must be made to ensure that
being a Karcher mean is sufficient to be a barycenter. When it is sufficient,
it immediately gives an alternative test for $y$ to be a barycenter for
$\lambda$: simply check if $||\nabla G_{\lambda}(y)||_{y}=0$.

An example of a Karcher mean which is not a barycenter is given in the
Supplementary Material Section 7. To ensure that being a Karcher mean is
sufficient to be a barycenter we require the following assumptions:

A1: The measures $\\{\mu_{i}\\}_{i=0}^{p}$ are absolutely continuous and
supported on either all of $\mathbb{R}^{d}$ or a bounded open convex subset.
Call this shared support set $\Omega$.

A2: The measures $\\{\mu_{i}\\}_{i=0}^{p}$ have respective densities
$\\{g_{i}\\}_{i=0}^{p}$ which are bounded above and $g_{1},...,g_{p}$ are
strictly positive on $\Omega$.

A3: If $\Omega=\mathbb{R}^{d}$ then $\\{g_{i}\\}_{i=0}^{p}$ are locally Hölder
continuous. Otherwise $\\{g_{i}\\}_{i=0}^{p}$ are bounded away from zero on
$\Omega$.

Under these assumptions if $\mu_{0}$ is a Karcher mean with coordinates
$\lambda$, then $\mu_{0}=\nu_{\lambda}$ (Panaretos & Zemel, 2020); see the
Supplementary Material Section 1 for a precise statement.

This characterization of barycenters can be much easier to work with and the
assumptions required are mild in practice. To accompany the notion of a
Karcher mean, the _Fréchet derivative_ (Fréchet, 1948) of the variance
functional $G_{\lambda}$ is given by

$\nabla
G_{\lambda}(\nu)=-\sum_{i=1}^{p}\lambda_{i}\left(T_{i}-\operatorname*{Id}\right),$
(4)

where $T_{i}$ is the optimal transport map from $\nu$ to $\mu_{i}$ as in (1)
(Ambrosio et al., 2005; Panaretos & Zemel, 2020).

We now have all the tools to begin building up to Theorem 1. The first step is
to derive a formula for $||\nabla G_{\lambda}(\mu_{0})||_{\mu_{0}}^{2}$ in
terms of $\lambda$ and show its convexity.

######  Proposition 1.

Let
$\\{\mu_{i}\\}_{i=0}^{p}\subset\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$.
Then $||\nabla G_{\lambda}(\mu_{0})||_{\mu_{0}}^{2}=\lambda^{T}A\lambda$ where
$A\in\mathbb{R}^{p\times p}$ is given by

$A_{ij}=\int_{\mathbb{R}^{d}}\langle
T_{i}(x)-\operatorname*{Id}(x),T_{j}(x)-\operatorname*{Id}(x)\rangle
d\mu_{0}(x)$ (5)

and $T_{i}$ is the optimal transport map from $\mu_{0}$ to $\mu_{i}$ as in
(1). Furthermore, this is a convex function in $\lambda$.

The proof is given in the Supplementary Material Section 2.1. The expression
$T_{i}(x)-\operatorname*{Id}(x)$ is the displacement of the point $x$ when
being transported from $\mu_{0}$ to $\mu_{i}$. Under this interpretation,
$\langle T_{i}(x)-\operatorname*{Id},T_{j}(x)-\operatorname*{Id}\rangle$ can
be thought of as the angle between the displacement associated to the optimal
transport map between $\mu_{0}$ to $\mu_{i}$ with that of $\mu_{0}$ to
$\mu_{j}$. Integrating this quantity with respect to $\mu_{0}$ therefore
quantifies the average angle between displacements and if the optimal
displacement maps move in similar directions.

We now use this functional form as part of a convex program and show that it
can be used to check if a $\mu_{0}$ is a Karcher mean. Our first main result
ties Propositions 1 to the conditions for equivalence between barycenters and
Karcher means.

######  Theorem 1.

Let $\\{\mu_{i}\\}_{i=0}^{p}$ satisfy A1-A3. Then
$\mu_{0}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})$ if and only if

$\min_{\lambda\in\Delta^{p}}\lambda^{T}A\lambda=0$

where $A\in\mathbb{R}^{p\times p}$ is given by (5). Furthermore, if the
minimum value is 0 and $\lambda_{*}$ is an optimal argument, then
$\mu_{0}=\nu_{\lambda_{*}}$.

The proof is given in the Supplementary Material Section 2.3. Theorem 1 gives
a geometric characterization of the solution to (3) which relies on the
evaluation of inner-products in the tangent space
$T_{\mu_{0}}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ (noting that
$A_{ij}=\langle
T_{i}-\operatorname*{Id},T_{j}-\operatorname*{Id}\rangle_{\mu_{0}}$) and
solving a constrained quadratic program. Furthermore, this result suggests a
method for finding an approximator of the $\mu_{0}$ in the set of barycenters:
simply minimize the objective and use $\lambda_{*}$, even if the minimum value
is not zero.

Note that if $A$ has rank less than $p-1$ there may be multiple minimizers
which correspond to redundancies in the set
$\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})$; see Supplementary Material
Section 6.

###  2.2 Theorem 1 For Gaussians

In the Gaussian case there are formulas for both the optimal transport maps
and the integration involved in $A_{ij}$. Let $\mathbb{S}_{++}^{d}$ denote the
set of symmetric positive definite $d\times d$ matrices. Let
$\mu_{i}=\mathcal{N}(0,S_{i})$ with $S_{i}\in\mathbb{S}_{++}^{d}$ for
$i=0,\dots,p$. Then the optimal transport map from $\mu_{0}$ to $\mu_{i}$ is
given by $T_{i}(x)=C_{i}x$ where

$C_{i}=S_{0}^{-1/2}\left(S_{0}^{1/2}S_{i}S_{0}^{1/2}\right)^{1/2}S_{0}^{-1/2}$
(6)

(Agueh & Carlier, 2011). Furthermore, A1-A3 are satisfied for non-degenerate
Gaussians. Combining Theorem 1 with (6 gives the following.

######  Corollary 1.

For $i=1,\dots,p$, let $\mu_{i}=\mathcal{N}(0,S_{i})$ with
$S_{i}\in\mathbb{S}_{++}^{d}$. Then $\mu_{0}$ is a barycenter if and only if
$\mu_{0}=\mathcal{N}(0,S_{0})$ for some $S_{0}\in\mathbb{S}_{++}^{d}$, and the
convex program of Theorem 1 has minimum value 0, where the matrix $A$ is given
by $A_{ij}=\operatorname{Tr}\left((C_{i}-I)(C_{j}-I)S_{0}\right).$
Furthermore, if the minimum value is zero and $\lambda_{*}$ is a optimal
argument, then $\mu_{0}=\nu_{\lambda_{*}}$

The proof is deferred to the Supplementary Material Section 2.3. We note that
in the Gaussian case, one can use $||\nabla G_{\lambda}(\mu_{0})||_{\mu_{0}}$
to upper bound $W_{2}(\mu_{0},\nu_{\lambda})$ (Chewi et al., 2020; Altschuler
et al., 2021). This further justifies the choice of optimizing $||\nabla
G_{\lambda}(\mu_{0})||_{\mu_{0}}^{2}$, since it can be seen as a convex, sharp
upper bound on the distance to the set of barycenters.

###  2.3 Projections of Compatible Measures

If one imposes extra structure on the set $\\{\mu_{i}\\}_{i=0}^{p}$, then it
is possible to substantially strengthen Theorem 1.

######  Definition 2.

A family of measures $\\{\mu_{i}\\}_{i=0}^{p}$ is said to be _compatible_ if
for any $i,j,k\in\\{0,1,...,p\\}$ the optimal transport map $T_{i}^{k}$ from
$\mu_{i}$ to $\mu_{k}$ can can be expressed as

$T_{i}^{k}=T_{j}^{k}\circ T_{i}^{j}$

where $T_{i}^{j},T_{j}^{k}$ are the optimal transport maps from $\mu_{i}$ to
$\mu_{j}$ and $\mu_{j}$ to $\mu_{k}$ respectively.

Note that whenever these maps exist, the map $T_{j}^{k}\circ T_{i}^{j}$ is a
valid map from $\mu_{i}$ to $\mu_{k}$, but it need not be optimal. Under the
condition of compatibility, we have the following result.

######  Theorem 2.

Suppose that $\\{\mu_{i}\\}_{i=0}^{p}$ are compatible and let $\lambda_{*}$ be
the minimizer in Theorem 1. Then

$\lambda_{*}=\operatorname*{arg\,min}_{\lambda\in\Delta^{p}}W_{2}^{2}(\mu_{0},\nu_{\lambda}).$

The proof is deferred to the Supplementary Material Section 2.4. In other
words, when the measures are compatible one can solve (3) and find the _exact
projection_ of $\mu_{0}$ onto $\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p}).$

There are several known conditions under which the family of measures is
compatible. These include any set of continuous measures on $\mathbb{R}$, any
family of Gaussians with simultaneously diagonalizable covariance matrices,
families of radial transformations of a base measure, or families with a
"common dependence structure”; see (Panaretos & Zemel, 2020) Section 2.3 for
further details.

##  3 Finite Sample Analysis and Rates of Convergence

The results of the previous section are all reliant upon having exact
knowledge of both the underlying measure $\mu_{0}$ as well as the optimal
transport maps $T_{i}$. In practice, outside of specific parametric families
(e.g., Gaussians), it is rare that one will have either of these available.

Much more common is the setting where the $\mu_{i}$ are accessed through
i.i.d. sampling. In this setting we need to estimate the maps $T_{i}$, as well
as $\mu_{0}$ and use these estimates to compute an $\hat{A}$ whose entries are
$\hat{A}_{ij}=\langle\hat{T}_{i}-\operatorname*{Id},\hat{T}_{j}-\operatorname*{Id}\rangle_{\hat{\mu}_{0}}$,
where $\hat{T}_{i}$ and $\hat{\mu_{0}}$ are our estimates of the transport
maps and $\mu_{0}$ respectively. The quality of the estimate $\hat{A}$ will
depend on the approximations used above, and ideally these would be
accompanied by performance guarantees.

###  3.1 The Entropic Map Estimate

Let $X_{1},...,X_{n}\sim\mu_{0}$ and $Y_{1},...,Y_{n}\sim\mu_{1}$ be i.i.d.
samples. In (Pooladian & Niles-Weed, 2021) the authors propose using the
entropic map to estimate the transport maps. For $\epsilon>0$, this map is
defined by

$\hat{T}_{\epsilon}(x)\triangleq\frac{\frac{1}{n}\displaystyle\sum_{i=1}^{n}Y_{i}\exp\left(\frac{1}{\epsilon}(g_{\epsilon}(Y_{i})-\frac{1}{2}||x-Y_{i}||_{2}^{2})\right)}{\frac{1}{n}\displaystyle\sum_{i=1}^{n}\exp\left(\frac{1}{\epsilon}(g_{\epsilon}(Y_{i})-\frac{1}{2}||x-Y_{i}||_{2}^{2})\right)}$
(7)

where $g_{\epsilon}$ is defined to be the optimal variable in the dual
entropic optimal transport problem on the samples. See the Supplementary
Material 1.1 for details on $g_{\epsilon}$ and important properties of
$\hat{T}_{\epsilon}$. Practically, the calculation of $g_{\epsilon}$ has been
well studied and there exists fast algorithms for its computation (Peyré &
Cuturi, 2020). The estimate has also been shown to converge (with rates) to
the optimal transport map in $L_{2}(\mu_{0})$ under the appropriate technical
conditions.

To accompany this we also have the following result, stated informally below.
In this result $a\lesssim b$ means that there exists a constant $C>0$ such
that $a\leq Cb$. We allow this constant $C$ to depend on the size of the
support of the measures, upper and lower bounds on their densities, and the
regularity properties of the optimal transport maps $T_{i},T_{j}$, but it is
importantly independent of the sample size $n$.

######  Theorem 3.

(Informal) Let $i,j\in\\{1,...,p\\}$ and suppose that
$\mu_{i},\mu_{j},\mu_{0}$ are supported on bounded domains and that the maps
$T_{i}$ and $T_{j}$ are sufficiently regular. Let
$X_{1},...,X_{2n}\sim\mu_{0},Y_{1},...,Y_{n}\sim\mu_{i},Z_{1},...,Z_{n}\sim\mu_{j}$.
For an appropriately chosen $\epsilon$, let $\hat{T}_{i}$ and $\hat{T}_{j}$ be
the entropic maps computed using
$\\{X_{i}\\}_{i=1}^{n},\\{Y_{i}\\}_{i=1}^{n},\\{Z_{i}\\}_{i=1}^{n}$. Then we
have

$\begin{split}&\mathbb{E}\left[\left|A_{ij}-\frac{1}{n}\sum_{k=n+1}^{2n}\langle\hat{T}_{i}(X_{k})-X_{k},\hat{T}_{j}(X_{k})-X_{k}\rangle\right|\right]\\\
&\hskip
42.67912pt\lesssim\frac{1}{\sqrt{n}}+n^{-\frac{\alpha+1}{4(d^{\prime}+\alpha+1)}}\sqrt{\log
n}\end{split}$ (8)

where $d^{\prime}=2\lceil d/2\rceil$, and $\alpha\leq 3$ depends on the
regularity of optimal transport maps.

A precise statement and proof of this theorem is given in the Supplementary
Material Section 2.5. Theorem 3 tells us that in expectation, by using the
entropic map we can accurately approximate the entries of the matrix $A_{ij}$,
and do so in a numerically feasible way. We chose the entropic map as it is
computationally tractable and has competitive convergence rates when compared
to the intractable but optimal estimator of (Hütter & Rigollet, 2021). We
expect similar analysis holds for other recently proposed estimators (Deb et
al., 2021; Gunsilius, 2021; Manole et al., 2021; Muzellec et al., 2021) under
different assumptions.

Theorem 3 can be combined with a matrix perturbation analysis to demonstrate
the consistency of the estimated barycentric coordinates
$\operatorname*{\hat{\lambda}}$ from Algorithm 1 for the true ones
$\lambda_{*}$ when $\mu_{0}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i=1}^{p})$,
that is when $\lambda_{*}^{T}A\lambda_{*}=0$. This is stated more precisely as
follows.

######  Corollary 2.

Let $\operatorname*{\hat{\lambda}}$ be the random estimate obtained from
Algorithm 1. Suppose that $A$ has an eigenvalue of 0 with multiplicity 1 and
that $\lambda_{*}\in\Delta^{p}$ realizes $\lambda_{*}^{T}A\lambda_{*}=0$. Then
under the assumptions of Theorem 3,

$\operatorname*{\mathbb{E}}[\|\operatorname*{\hat{\lambda}}-\lambda_{*}\|_{2}^{2}]\lesssim\frac{1}{\sqrt{n}}+n^{-\frac{\alpha+1}{4(d^{\prime}+\alpha+1)}}\sqrt{\log
n}.$

Algorithm 1 Estimate $\lambda$

Input: i.i.d. samples
$\\{X_{1},...,X_{2n}\\}\sim\mu_{0},\\{\\{Y_{1}^{i},...,Y_{n}^{i}\\}\sim\mu_{i}:i=1,...,p\\},$
regularization parameter $\epsilon>0$.

for $i=1,...,p$ do

Set $M^{i}\in\mathbb{R}^{n\times n}$ with
$M^{i}_{jk}=\frac{1}{2}||X_{j}-Y_{k}^{i}||_{2}^{2}$.

Solve for $g^{i}$ as the optimal $g$ in
$\begin{split}&\max_{f,g\in\mathbb{R}^{n}}\frac{1}{n}\sum_{j=1}^{n}f_{j}+\frac{1}{n}\sum_{k=1}^{n}g_{k}\\\
&\hskip
42.67912pt-\frac{\epsilon}{n^{2}}\sum_{j,k}^{n}\exp\left((f_{j}+g_{k}-M_{jk}^{i})/\epsilon\right).\end{split}$

Define $\hat{T}_{i}$ through 7 with $g_{\epsilon}=g_{i}$ and
$\\{Y_{1}^{i},...,Y_{n}^{i}\\}$.

end for

Set $\hat{A}\in\mathbb{R}^{p\times p}$ to be the matrix with entries
$\hat{A}_{ij}=\frac{1}{n}\sum_{k=n+1}^{2n}\langle\hat{T}_{i}(X_{k})-X_{k},\hat{T}_{j}(X_{k})-X_{k}\rangle.$

Return
$\hat{\lambda}=\displaystyle\operatorname*{arg\,min}_{\lambda\in\Delta^{p}}\lambda^{T}\hat{A}\lambda$.

The proof is deferred to the Supplementary Material Section 2.6. Note that the
implicit constant in the inequality of Corollary 2 depends on $p$ in addition
to the dependencies of listed above. Corollary 2 ensures that in the large
sample limit, the entropic regularization parameter may be chosen to guarantee
precise estimation of the true barycentric coordinates. Note that the rate of
convergence in Corollary 2 depends crucially on the smoothness of the
transport maps between the reference measures ($\alpha$) and the
dimensionality of the space in which they are supported ($d^{\prime}$).

In the next section we illustrate the utility of the BCM and propose novel
approaches based on Theorem 1 and Algorithm 1 for several applications.

##  4 Applications

###  4.1 Barycenter Parameterized Covariance Estimation

Extending the set-up in (Musolas et al., 2021), we first consider the case
where $\mu_{i}=\mathcal{N}(0,S_{i}),S_{i}\in\mathbb{S}^{d}_{++}$ for
$i=0,\dots,p$. We will use $S_{i}$ both as a matrix and to refer to $\mu_{i}$;
similarly, let $S_{\lambda}$ denote $\nu_{\lambda}$. Let
$\\{x_{i}\\}_{i=1}^{n}$ be i.i.d. samples with empirical covariance $\hat{S}$.
In this setting, we can use the formulas laid out in Corollary 1.

For comparison, we also consider maximum likelihood estimation (MLE) for the
parameter $\lambda$ in the BCM. Let $P_{\lambda}$ denote the probability
density function of $S_{\lambda}$. The MLE is
$\operatorname*{arg\,max}_{\lambda\in\Delta^{p}}\prod_{i=1}^{n}P_{\lambda}(x_{i}),$
which is equivalent to minimizing the KL-divergence
$D_{KL}(\hat{S},S_{\lambda})$. However, this problem may be non-convex and
difficult to optimize. We solve it numerically using auto-differentiation
(Bartholomew-Biggs et al., 2000) and use the coordinates recovered by this
method, which may not be the true MLE.

Experimental Setup: We consider the problem of estimating the covariance
matrix $S_{0}$ from i.i.d. samples
$\\{x_{i}\\}_{i=1}^{n}\sim\mathcal{N}(0,S_{0})$ when $S_{0}$ is known to be a
barycenter of $S_{1},...,S_{p}$. To do so we use the following procedure:

  1. 1.

Choose $\lambda$ and $\\{S_{i}\\}_{i=1}^{p}$: Select $\lambda\in\Delta^{p}$
and the reference measures $\\{S_{i}\\}_{i=1}^{p}$ either by hand or at
random.

  2. 2.

Generate $S_{\lambda}$: Using Algorithm 1 of (Chewi et al., 2020), find the
true barycenter, $S_{\lambda}$.

  3. 3.

Sample: Sample $\\{x_{i}\\}_{i=1}^{n}$ i.i.d. from
$\mathcal{N}(0,S_{\lambda})$. Compute the empirical covariance
$\hat{S}=\frac{1}{n}\sum_{i=1}^{n}x_{i}x_{i}^{T}.$

  4. 4.

Estimate $\hat{\lambda}$: Perform the optimization in Corollary 1, or use MLE
with $S_{0}=\hat{S}$ to recover an estimate $\hat{\lambda}$.

  5. 5.

Compute $S_{\hat{\lambda}}$: Again using Algorithm 1 of (Chewi et al., 2020)
find $S_{\hat{\lambda}}$ and use that as the estimate of $S_{\lambda}$.

The results are in Figure 3. These show that using our approach (denoted
_gradient norm_) we can estimate the underlying covariance matrix much more
accurately than the empirical covariance. Our optimization is competitive with
the auto-differentiation approach to MLE, while being significantly more
stable and much faster. Indeed, computing all the estimates $\hat{\lambda}$ to
generate Figure 3 takes our method less than 1 second compared to over 3 hours
and 40 minutes for MLE. For reference it takes on average roughly 0.01 seconds
to perform empirical covaraiance estimation222All reported times are obtained
using 20 Intel Xeon CPU E5-2660 V4 cores at 2.00 GHz.. Further details are in
the Supplementary Material Section 3.

Figure 3: $W_{2}$ between true covariance and estimates. We set $p=6,d=10$,
$\lambda\sim\text{Unif}(\Delta^{p})$, and
$\mu_{i}\sim\text{Wish}(I_{d})+0.5I_{d}$ with results averaged over 250
trials, and 1 standard deviation shaded.

We now consider applications where it is necessary to directly estimate the
optimal transport map from samples (unlike the previous case where we estimate
it via sample-covariance estimators) and between measures with non-uniform
weights on the observed support. To that end, we first modify Algorithm 1 for
practical purposes.

###  4.2 Adapting Algorithm 1 for Applications

In practice, Algorithm 1 has a few mild requirements: that we use the same
number of samples for each reference measure; that the weight on each sample
is the same; and access to an extra set of points for evaluating the inner
products. Algorithm 2 relaxes these requirements and is in terms of matrix-
vector operations, making it suitable for practical implementation. We note
that if one uses Algorithm 2 with the same constraints as in Algorithm 1, then
both produce the exact same $\hat{T}_{i}$ on the samples $X_{1},...,X_{n}$
(Pooladian & Niles-Weed, 2021). The main difference between the two approaches
is that the dual potentials in Algorithm 1 make an out-of-sample extension
possible which may be of its own interest.

Algorithm 2 Estimate $\lambda$ on Point Clouds

Input: PMFs $p\in\mathbb{R}^{n_{0}}$, $q^{i}\in\mathbb{R}^{n_{i}}$, support
matrices $X\in\mathbb{R}^{n_{0}\times d}$,$Y^{i}\in\mathbb{R}^{n_{i}\times
d}$, regularization parameter $\epsilon>0$.

for $i=1,...,p$ do

Set $M^{i}\in\mathbb{R}^{n_{0}\times n_{i}}$ with
$M^{i}_{jk}=||X_{j}-Y_{k}^{i}||_{2}^{2}$.

Solve for the entropic assignment matrix $\pi^{i}$ as the optimal matrix in
$\min_{\begin{subarray}{c}\pi\in\mathbb{R}^{n_{0}\times n_{i}}_{+}\\\
\pi\bm{1}=p\\\
\pi^{T}\bm{1}=q^{i}\end{subarray}}\sum_{j=1}^{n_{0}}\sum_{k=1}^{n_{i}}M^{i}_{jk}\pi_{jk}+\epsilon\pi_{jk}\log\pi_{jk}.$

Compute the approximate transport matrix
$\hat{T}_{i}\in\mathbb{R}^{n_{0}\times d}$ as
$\hat{T}_{i}=\text{diag}(1/p)\pi^{i}Y^{i}.$

end for

Set $\hat{A}\in\mathbb{R}^{p\times p}$ to be the matrix with entries
$\hat{A}_{ij}=\operatorname{Tr}\left(\text{diag}(p)(\hat{T}^{i}-X)(\hat{T}^{j}-X)^{T}\right).$

Return
$\hat{\lambda}=\displaystyle\operatorname*{arg\,min}_{\lambda\in\Delta^{p}}\lambda^{T}\hat{A}\lambda$.

###  4.3 Image Inpainting and Denoising

We consider the problem of recovering an image in the presence of corruption,
interpreting the image as a probability measure. We consider two specific
models of corruption: a. additive _noise_ and b. _occlusion_ of a portion of
the image. Our experimental procedure for both of these problems is outlined
below taking as experimental data the MNIST dataset of $28\times 28$ pixel
images of hand-written digits (LeCun, 1998). Additional details are in
Supplementary Material Section 4. Note that after normalization, these can be
treated as measures supported on a $28\times 28$ grid.

Experimental Setup:

Figure 4: _Left to right_ : corrupted image; recovery by linear projection;
recovery using $\hat{\lambda}$ from Algorithm 2; original image. _Top_ : white
noise added to the image as in a. _Bottom_ : occlusion of the image as in b.
We see the linear reconstruction fails, while Algorithm 2 recovers well,
albeit with some blurring that can be interpreted as a consequence of entropic
regularization.

  1. 1.

Select $\mu_{0}$: Select a digit $\mu_{0}$ a. generate a white noise image
$\zeta$. Set $\tilde{\mu}_{0}$ to be
$\tilde{\mu_{0}}=(1-\alpha)\mu_{0}+\alpha\zeta$. b. Set $\tilde{\mu}_{0}$ to
be $\mu_{0}$ with the central $8\times 8$ square removed, and renormalized.

  2. 2.

Select $\\{\mu_{i}\\}_{i=1}^{p}$: Select a set of images of the same digit as
$\mu_{0}$ to serve as the reference measures. a. Let
$\tilde{\mu}_{i}=(1-\alpha)\mu_{0}+\alpha\mathbb{E}[\zeta]$. b. Let
$\tilde{\mu}_{i}$ be $\mu_{i}$ with the central $8\times 8$ square removed,
and renormalized

  3. 3.

Estimate $A$, $\lambda$: Use Algorithm 2 to compute the approximate Gram
matrix $\hat{A}$ using $\tilde{\mu}_{0},\tilde{\mu_{1}},...\tilde{\mu_{p}}$
and then compute estimated coordinates $\hat{\lambda}$.

  4. 4.

Compute $\hat{\nu_{\lambda}}$: Output $\hat{\mu}_{0}=\nu_{\hat{\lambda}}$,
where the barycenter is reconstructed from $\mu_{1},...,\mu_{p}$.

An illustration of the procedure is shown in Figure 4. We also show for
comparison digits reconstructed using a linear method, in which the corrupted
image is recovered by computing the Euclidean projection onto the convex hull
of the reference measures (see the Supplement 4.3 for further details). To
recover the barycenter we use the method of (Benamou et al., 2015). As a
competitive baseline we compare to a well-known existing method for histogram
regression (Bonneel et al., 2016) in this setting. Our results are summarized
in Tables 1 and 2. We note that our method is over an order of magnitude
faster than (Bonneel et al., 2016) on this dataset and achieves competitive
results. This is particularly remarkable because (Bonneel et al., 2016) is
specifically adapted to measures with structured support such as grids and
meshes, while our Algorithm 2 is much more general.

| Bon., $\epsilon=10$ | Alg. 2, $\epsilon=10$ | Alg. 2, $\epsilon=0$
---|---|---|---
Occlusion | 2.5101 (2228s) | 2.5488 (3.371s) | 2.5287 (1.062s)
Noise ($\alpha=0.5$) | 2.4058 (2391s) | 2.6797 (86.32s) | 2.3787 (50.67s)
Table 1: Average quality of recovery of $\mu_{0}$ measured in $W_{2}^{2}(\mu_{0},\hat{\mu}_{0})$ when reconstructing 500 random 4’s using a barycenter constructed from 10 random 4’s, as well as run times of each method. Algorithm 2 affords comparably accurate reconstructions in orders-of-magnitude faster time than the state-of-the-art (Bonneel et al., 2016). | (Alg. 2, $\epsilon=10$) - (Bon., $\epsilon=10$) | (Alg. 2, $\epsilon=0$) - (Bon., $\epsilon=10$) | (Alg. 2, $\epsilon=10$) - (Alg. 2, $\epsilon=0$)
---|---|---|---
Occlusion | $0.03872\pm 0.3103$ | $0.01868\pm 0.2934$ | $0.02004\pm 0.2145$
Additive | $0.2739\pm 0.3564$ | $-0.02706\pm 0.2558$ | $0.3009\pm 0.2486$
Table 2: Average and standard deviation of the pairwise-difference in
reconstruction quality measured in $W_{2}^{2}(\mu_{0},\hat{\mu}_{0})$ in the
same setting as Table 1. The average pairwise-difference is smaller than the
standard deviation in all cases involving (Bonneel et al., 2016), suggesting
there is no statistically significant difference. Our approach should thus be
preferred to (Bonneel et al., 2016), on the grounds that it is considerably
more efficient while achieving essentially the same accuracy.

###  4.4 Document Classification

Finally, we consider the task of identifying the topic of a document using its
word embedding representation as the empirical distribution. We assume that
there are $t$ topics and that we have $p$ reference documents about each
topic. Each document can be represented as a high-dimensional empirical
measure using a combination of a bag-of-words representation and a node2vec
(Mikolov et al., 2013) word-embedding, and is paired with topic label. We use
the publicly available dataset provided by (Huang et al., 2016). The task is
to use a small number of labelled documents from each topic to predict the
topic of a test document.

We consider four predictors: (1) 1-Nearest Neighbor (1NN) which classifies
using the topic of the $W_{2}^{2}$-nearest reference document; (2) Minimum
Average Distance which selects the topic with reference documents on average
$W_{2}^{2}$-closest to the test document; (3) Minimum Barycenter Loss which
runs Algorithm 2 using the references from each topic separately and selects
the topic with the smallest loss in the quadratic form; (4) Maximum Coordinate
which runs Algorithm 2 using all the reference documents and chooses the topic
which receives the most mass from $\operatorname*{\hat{\lambda}}$.

Figure 5: Document topic prediction accuracy as a function of number of
reference documents in each class. _Top:_ BBC Sport dataset (5 classes).
_Bottom:_ News20 dataset (20 classes).

Our findings on the BBCSport (5 topics) and News20 (20 topics) datasets are
reported in Figure 5. To generate these figures we randomly sample $k$
reference documents per topic and a test set of 100 random documents. We apply
the four methods listed above and compute the accuracy on the test set. This
procedure is repeated 50 times for each choice of $k$ and the average accuracy
is plotted. See Supplementary Material Section 5.2 for details.

We see that even with a very small amount of training data, our BCM approaches
perform well. Importantly, the BCM is able to represent unseen documents using
these small training sets, giving it a clear advantage over simply classifying
based on which document class it is $W_{2}^{2}$-nearest (1NN) or
$W_{2}^{2}$-nearest in an aggregate sense (Min. Avg. Dist.). This suggests
that even when $p$ is small, the BCM has high modeling capacity and that the
learned barycentric coordinates encode important information.

##  5 Conclusion

We have proposed a new method for computing coordinates under the BCM,
together with guarantees on its solution via a convex program (Theorem 1) with
closed-form coordinates in the Gaussian case (Corollary 1). We further
developed an algorithm for estimating the coordinates under the BCM when all
measures are accessible only via i.i.d. samples (Algorithm 1) which enjoys a
natural smoothness and dimension-dependent rate of convergence to the true
parameters (Theorem 3, Corollary 2). The BCM paradigm affords significant
gains in runtime and robustness for barycenter parameterized Gaussian measure
estimation, and provides an effective approach to image reconstruction and
document classification when interpreting these data as measures. The results
in this paper suggest the efficiency and effectiveness of the BCM as a broadly
viable modelling tool.

Future Work and Open Questions: Section 3 provides a consistency analysis for
$A$ and $\lambda$ by estimating the transport maps between each measure. But,
all that is needed is an _estimate on the inner product_ between the
displacement maps, which in principle could admit more sample efficient
approaches that avoid explicitly estimating the maps.

Section 4 suggests that unseen data can be well-represented in the BCM for
randomly chosen reference measures. An interesting theoretical problem is to
understand the representational capacity of the BCM using particular reference
measures (e.g., random ones, optimally chosen ones, those in certain
parametric families). This naturally touches on the question of the smoothness
of the synthesis map $\lambda\mapsto\nu_{\lambda}.$

Computationally, all of our real data experiments use randomly chosen
reference measures. It is of interest to develop efficient procedures as in
(Schmitz et al., 2018) for _learning_ reference measures that induce low
reconstruction error. Beyond that, one could regularize the analysis problem
(3) to encourage sparse or otherwise structured coordinates as is well-studied
in Euclidean settings (Aharon et al., 2006).

## Acknowledgements

MW is supported by NSF CCF-1553075. RJ is supported by NSF DRL 1931978. SA
acknowledges support by NSF CCF-1553075, NSF DRL 1931978, NSF EEC 1937057, and
AFOSR FA9550-18-1-0465. JMM acknowledges support from the Dreyfus Foundation
and the NSF through grants DMS-1912737, DMS-1924513. All authors acknowledge
support through the Tufts TRIPODS Institute, supported by the NSF under grant
CCF-1934553. We acknowledge the reviewers for their helpful comments which
improved the paper considerably.

## References

  * Agueh & Carlier (2011) Agueh, M. and Carlier, G.  Barycenters in the Wasserstein space.  _SIAM Journal on Mathematical Analysis_ , 43(2):904–924, 2011.
  * Aharon et al. (2006) Aharon, M., Elad, M., and Bruckstein, A.  K-SVD: An algorithm for designing overcomplete dictionaries for sparse representation.  _IEEE Transactions on Signal Processing_ , 54(11):4311–4322, 2006.
  * Altschuler et al. (2021) Altschuler, J. M., Chewi, S., Gerber, P., and Stromme, A. J.  Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent.  In _Advances in Neural Information Processing Systems_ , volume 34, 2021.
  * Álvarez-Esteban et al. (2016) Álvarez-Esteban, P. C., Del Barrio, E., Cuesta-Albertos, J., and Matrán, C.  A fixed-point approach to barycenters in Wasserstein space.  _Journal of Mathematical Analysis and Applications_ , 441(2):744–762, 2016.
  * Ambrosio et al. (2005) Ambrosio, L., Gigli, N., and Savare, G.  _Gradient Flows: In Metric Spaces and in the Space of Probability Measures_.  Lectures in Mathematics. ETH Zürich. Birkhäuser Basel, 2005.
  * Bartholomew-Biggs et al. (2000) Bartholomew-Biggs, M., Brown, S., Christianson, B., and Dixon, L.  Automatic differentiation of algorithms.  _Journal of Computational and Applied Mathematics_ , 124:171–190, 12 2000.
  * Benamou et al. (2015) Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyré, G.  Iterative Bregman projections for regularized transportation problems.  _SIAM Journal on Scientific Computing_ , 37(2):A1111–A1138, 2015.
  * Bonneel et al. (2015) Bonneel, N., Rabin, J., Peyré, G., and Pfister, H.  Sliced and Radon Wasserstein barycenters of measures.  _Journal of Mathematical Imaging and Vision_ , 51(1):22–45, 2015.
  * Bonneel et al. (2016) Bonneel, N., Peyré, G., and Cuturi, M.  Wasserstein barycentric coordinates: histogram regression using optimal transport.  _ACM Transactions on Graphics_ , 35(4):71:1–71:10, 2016.
  * Brenier (1991) Brenier, Y.  Polar factorization and monotone rearrangement of vector-valued functions.  _Communications on Pure and Applied Mathematics_ , 44(4):375–417, 1991.
  * Cheng et al. (2021) Cheng, K., Aeron, S., Hughes, M. C., and Miller, E.  Dynamical Wasserstein barycenters for time-series modeling.  In _Advances in Neural Information Processing Systems_ , volume 34, 2021.
  * Chewi et al. (2020) Chewi, S., Maunu, T., Rigollet, P., and Stromme, A. J.  Gradient descent algorithms for Bures-Wasserstein barycenters.  In _Conference on Learning Theory_ , pp. 1276–1304. PMLR, 2020\.
  * Chizat et al. (2020) Chizat, L., Roussillon, P., Léger, F., Vialard, F.-X., and Peyré, G.  Faster Wasserstein distance estimation with the Sinkhorn divergence.  _arXiv preprint arXiv:2006.08172_ , 2020.
  * Claici et al. (2018) Claici, S., Chien, E., and Solomon, J.  Stochastic Wasserstein barycenters.  In _International Conference on Machine Learning_ , pp. 999–1008. PMLR, 2018.
  * Colombo et al. (2021) Colombo, P., Staerman, G., Clavel, C., and Piantanida, P.  Automatic text evaluation through the lens of Wasserstein barycenters.  _arXiv preprint arXiv:2108.12463_ , 2021.
  * Cuturi (2013) Cuturi, M.  Sinkhorn distances: Lightspeed computation of optimal transport.  _Advances in Neural Information Processing Systems_ , 26:2292–2300, 2013.
  * Cuturi & Doucet (2014) Cuturi, M. and Doucet, A.  Fast computation of Wasserstein barycenters.  In _International Conference on Machine Learning_ , pp. 685–693. PMLR, 2014.
  * Deb et al. (2021) Deb, N., Ghosal, P., and Sen, B.  Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Donoho (2006) Donoho, D. L.  Compressed sensing.  _IEEE Transactions on Information Theory_ , 52(4):1289–1306, 2006.
  * Fréchet (1948) Fréchet, M.  Les éléments aléatoires de nature quelconque dans un espace distancié.  _Annales de l’institut Henri Poincaré_ , 10(4):215–310, 1948.
  * Genevay (2019) Genevay, A.  _Entropy-regularized optimal transport for machine learning_.  PhD thesis, Paris Sciences et Lettres (ComUE), 2019.
  * Gunsilius (2021) Gunsilius, F. F.  On the convergence rate of potentials of brenier maps.  _Econometric Theory_ , pp. 1–37, 2021.
  * Huang et al. (2016) Huang, G., Quo, C., Kusner, M. J., Sun, Y., Weinberger, K. Q., and Sha, F.  Supervised word mover’s distance.  In _Advances in Neural Information Processing Systems_ , pp. 4869–4877, 2016.
  * Hütter & Rigollet (2021) Hütter, J.-C. and Rigollet, P.  Minimax estimation of smooth optimal transport maps.  _The Annals of Statistics_ , 49(2):1166–1194, 2021.
  * Karcher (1977) Karcher, H.  Riemannian center of mass and mollifier smoothing.  _Communications on Pure and Applied Mathematics_ , 30(5):509–541, 1977.
  * Kroshnin et al. (2019) Kroshnin, A., Tupitsa, N., Dvinskikh, D., Dvurechensky, P., Gasnikov, A., and Uribe, C.  On the complexity of approximating Wasserstein barycenters.  In _International Conference on Machine Learning_ , pp. 3530–3540. PMLR, 2019.
  * Le Gouic & Loubes (2017) Le Gouic, T. and Loubes, J.-M.  Existence and consistency of Wasserstein barycenters.  _Probability Theory and Related Fields_ , 168(3):901–917, 2017.
  * LeCun (1998) LeCun, Y.  The MNIST database of handwritten digits.  _http://yann. lecun. com/exdb/mnist/_ , 1998.
  * Mallat (1999) Mallat, S.  _A wavelet tour of signal processing_.  Elsevier, 1999.
  * Manole et al. (2021) Manole, T., Balakrishnan, S., Niles-Weed, J., and Wasserman, L.  Plugin estimation of smooth optimal transport maps.  _arXiv preprint arXiv:2107.12364_ , 2021.
  * McCann (1997) McCann, R. J.  A convexity principle for interacting gases.  _Advances in Mathematics_ , 128(1):153–179, 1997\.
  * Mikolov et al. (2013) Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., and Dean, J.  Distributed representations of words and phrases and their compositionality.  In _Advances in Neural Information Processing Systems_ , pp. 3111–3119, 2013.
  * Montesuma & Mboula (2021) Montesuma, E. F. and Mboula, F. M. N.  Wasserstein barycenter for multi-source domain adaptation.  In _IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 16785–16793, 2021.
  * Murphy (2012) Murphy, K. P.  _Machine learning: a probabilistic perspective_.  MIT press, 2012.
  * Musolas et al. (2021) Musolas, A., Smith, S. T., and Marzouk, Y.  Geodesically parameterized covariance estimation.  _SIAM Journal on Matrix Analysis and Applications_ , 42(2):528–556, 2021.
  * Muzellec et al. (2021) Muzellec, B., Vacher, A., Bach, F., Vialard, F.-X., and Rudi, A.  Near-optimal estimation of smooth transport maps with kernel sums-of-squares.  _arXiv preprint arXiv:2112.01907_ , 2021.
  * Panaretos & Zemel (2020) Panaretos, V. M. and Zemel, Y.  _An invitation to statistics in Wasserstein space_.  Springer Nature, 2020.
  * Peyré & Cuturi (2020) Peyré, G. and Cuturi, M.  Computational optimal transport, 2020.
  * Pooladian & Niles-Weed (2021) Pooladian, A.-A. and Niles-Weed, J.  Entropic estimation of optimal transport maps.  _arXiv: 2109.12004_ , 2021.
  * Popoviciu (1935) Popoviciu, T.  Sur les équations algébriques ayant toutes leurs racines réelles.  _Mathematica_ , 9:129–145, 1935.
  * Rabin et al. (2011) Rabin, J., Peyré, G., Delon, J., and Bernot, M.  Wasserstein barycenter and its application to texture mixing.  In _International Conference on Scale Space and Variational Methods in Computer Vision_ , pp. 435–446. Springer, 2011.
  * Redko et al. (2019) Redko, I., Courty, N., Flamary, R., and Tuia, D.  Optimal transport for multi-source domain adaptation under target shift.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 849–858. PMLR, 2019.
  * Santambrogio (2015) Santambrogio, F.  _Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling_.  Progress in Nonlinear Differential Equations and Their Applications. Springer International Publishing, 2015.
  * Schmitz et al. (2018) Schmitz, M. A., Heitz, M., Bonneel, N., Ngole, F., Coeurjolly, D., Cuturi, M., Peyré, G., and Starck, J.-L.  Wasserstein dictionary learning: optimal transport-based unsupervised nonlinear dictionary learning.  _SIAM Journal on Imaging Sciences_ , 11(1):643–678, 2018.
  * Schwerdtfeger (1961) Schwerdtfeger, H.  _Introduction to linear algebra and the theory of matrices_.  P. Noordhoff, 1961.
  * Singh et al. (2020) Singh, S. P., Hug, A., Dieuleveut, A., and Jaggi, M.  Context mover’s distance & barycenters: Optimal transport of contexts for building representations.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 3437–3449. PMLR, 2020.
  * Tošić & Frossard (2011) Tošić, I. and Frossard, P.  Dictionary learning.  _IEEE Signal Processing Magazine_ , 28(2):27–38, 2011.
  * Villani (2003) Villani, C.  _Topics in Optimal Transportation_.  Graduate studies in mathematics. American Mathematical Society, 2003.
  * Xu et al. (2018) Xu, H., Wang, W., Liu, W., and Carin, L.  Distilled Wasserstein learning for word embedding and topic modeling.  In _Advances in Neural Information Processing Systems_ , pp. 1716–1725, 2018.
  * Yang et al. (2021) Yang, L., Li, J., Sun, D., and Toh, K.-C.  A fast globally linearly convergent algorithm for the computation of Wasserstein barycenters.  _Journal of Machine Learning Research_ , 22(21):1–37, 2021.

Supplementary Material:
_Measure Estimation in the Barycentric Coding Model_

##  1 Precise Statements of Technical Results

For completeness, we replicate the theorem describing conditions under which
Karcher means are guaranteed to be barycenters.

######  Theorem 4.

((Panaretos & Zemel, 2020), Theorem 3.1.15) Let $\lambda\in\Delta^{p}$ with
$\lambda_{i}\neq 0$ for all $i$. Let $\Omega\subset\mathbb{R}^{d}$ be open and
convex and let
$\mu_{1},\dots,\mu_{p}\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$
supported on $\Omega$ with bounded, strictly positive densities
$g_{1},\dots,g_{p}$. Suppose that $\mu_{0}$ is an absolutely continuous
Karcher mean for $\lambda$ and is supported on $\Omega$ with bounded strictly
positive density $f$ there. Then $\mu_{0}$ is the Wasserstein barycenter for
coordinates $\lambda$ if one of the following holds:

  1. (1)

$\Omega=\mathbb{R}^{d}$ and the densities $f,g_{1},\dots,g_{p}$ are locally
Hölder continuous.

  2. (2)

$\Omega$ is bounded and the densities $f,g_{1},\dots,g_{p}$ are bounded below
on $\Omega$.

Next we include a technical theorem which we will leverage in our proof of the
convergence of the entries $A_{ij}$. Before doing so, we introduce the
necessary background.

######  Theorem 5.

((Brenier, 1991)) Let $\mu\in\mathcal{P}_{2,\emph{ac}}(\Omega)$ and
$\nu\in\mathcal{P}(\Omega)$. Then

  1. 1.

There exists a solution $T_{0}$ to (1), with $T_{0}=\nabla\varphi_{0}$, for a
convex function $\varphi_{0}$ solving

$\inf_{\varphi\in L^{1}(\mu)}\int\varphi d\mu+\int\varphi^{*}d\nu$

where $\varphi^{*}$ is the convex conjugate of $\varphi$.

  2. 2.

If in addition $\nu\in\mathcal{P}_{2,\emph{ac}}(\Omega),$ then
$\nabla\varphi_{0}^{*}$ is an admissible optimal transport map from $\nu$ to
$\mu$.

Let $\mathcal{C}^{\alpha}$ denote the space of functions possessing
$\lfloor\alpha\rfloor$ continuous derivatives and whose
$\lfloor\alpha\rfloor^{\text{th}}$ derivative is
$(\alpha-\lfloor\alpha\rfloor$) Hölder smooth. Using these conventions, we
require the technical conditions:

A4: $\mu,\nu\in\mathcal{P}_{2,\text{ac}}(\Omega)$ for a compact convex set
$\Omega$, with densities satisfying $p(x),q(x)\leq M$ and $q(x)\geq m>0$ for
all $x\in\Omega$.

A5: $\varphi_{0}\in\mathcal{C}^{2}(\Omega)$ and
$\varphi^{*}_{0}\in\mathcal{C}^{\alpha+1}(\Omega)$ for $\alpha>1$.

A6: There exist $l,L>0$ such that $T_{0}=\nabla\varphi_{0}$, with
$lI\preceq\nabla^{2}\varphi_{0}(x)\preceq LI$ for all $x\in\Omega$.

This first assumption is a specialization of A1 and A2, restricting to the
case where $\Omega$ is compact. The latter two can be thought of as regularity
conditions on the optimal transport maps between base and reference measures.
A5 asserts that the optimal transport maps have smooth derivatives and A6
controls the derivatives from both above and below.

###  1.1 Properties of The Entropic Map

Let $X_{i}\sim\mu_{0},Y_{i}\sim\mu_{1}$ for $i=1,...,n$ be i.i.d. samples from
their respective distributions. From these samples construct the empirical
measures $\hat{\mu}_{0,n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{X_{i}}$ and
$\hat{\mu}_{1,n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{Y_{i}}$.

For $\epsilon>0$ and $n$ samples from each measures, the discrete entropically
regularized OT problem can be written

$\min_{\begin{subarray}{c}\pi\in\mathbb{R}^{n\times n}_{+}\\\
\pi\bm{1}=(1/n)\bm{1}\\\
\pi^{T}\bm{1}=(1/n)\bm{1}\end{subarray}}\sum_{j,k=1}^{n}\pi_{jk}||X_{j}-Y_{k}||_{2}^{2}+\epsilon\pi_{jk}\log\pi_{jk}$

This problem has a dual formulation (Genevay, 2019) given by

$\max_{f,g\in\mathbb{R}^{n}}\frac{1}{n}\sum_{i=1}^{n}f_{i}+g_{i}-\epsilon\langle
e^{f/\epsilon},Ke^{g/\epsilon}\rangle$

where $K_{i,j}=\exp(-||X_{i}-Y_{i}||_{2}^{2}/\epsilon)$. $g_{\epsilon}$ is
defined to be the optimal $g$ in the maximization above, parameterized by the
setting of $\epsilon>0.$ For the method proposed it is important that the
$f_{\epsilon}$ and $g_{\epsilon}$ are chosen such that for all
$x,y\in\mathbb{R}^{d}$ we have

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\exp\left(\frac{1}{\epsilon}\left[f_{\epsilon}(x)+g_{\epsilon}(Y_{i})-\frac{1}{2}||x-Y_{i}||_{2}^{2}\right]\right)$
$\displaystyle=1$
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\exp\left(\frac{1}{\epsilon}\left[f_{\epsilon}(X_{i})+g_{\epsilon}(y)-\frac{1}{2}||X_{i}-y||_{2}^{2}\right]\right)$
$\displaystyle=1.$

This is always possible, and can be easily done in practice (Pooladian &
Niles-Weed, 2021).

One of the primary motivations for selecting the entropic map as our map
estimate is that $g_{\epsilon}$ can be efficiently computed (Peyré & Cuturi,
2020), and it comes with the performance guarantee stated below.

In the following theorem the constants may depend on the dimension $d$, the
diameter of the support $|\Omega|$, $M,m,l,L$, and
$||\varphi_{0}^{*}||_{\mathcal{C}^{\alpha+1}}$.

######  Theorem 6.

((Pooladian & Niles-Weed, 2021) Theorem 3) Under A4,A5,A6, the entropic map
$\hat{T}=T_{\epsilon,(n,n)}$ from $\hat{\mu}_{n}$ to $\hat{\nu}_{n}$ with
regularization parameter $\epsilon\asymp n^{-1/(d^{\prime}+\bar{\alpha}+1)}$
satisfies

$\mathbb{E}||\hat{T}-T_{0}||^{2}_{L^{2}(\mu)}\lesssim(1+I_{0}(\mu,\nu))n^{-\frac{(\bar{\alpha}+1)}{2(d^{\prime}+\bar{\alpha}+1)}}\log
n$ (9)

where $d^{\prime}=2\lceil d/2\rceil$ and $\bar{\alpha}=\min(\alpha,3)$ and
$I_{0}(\mu,\nu)$ is the integrated Fisher information between $\mu$ and $\nu$.

For more information on $I_{0}(\mu,\nu)$ see (Chizat et al., 2020). We remark
that when $\alpha\geq 2$ then under A4,A5,A6, it has been shown
$I_{0}(\mu,\nu)\leq C$ for a positive constant $C$ (Chizat et al., 2020). The
notation $a\asymp b$ means that there exists positive constants $c,C$ such
that $ca\leq b\leq Ca$ and $a\lesssim b$ means that there is a positive
constant $C$ such that $a\leq Cb$. Note that this result still holds without
the convexity condition in A4. We include it for ease of exposition below and
to ensure that the interior of an $\Omega$ that satisfies A4 is convex and
therefore satisfies A3.

##  2 Proofs of Theoretical Results

###  2.1 Proof of Proposition 1

###### Proof.

We first apply the definitions given in Section 2 and leverage (4):

$\displaystyle||\nabla G_{\lambda}(\mu_{0})||_{\mu_{0}}^{2}$
$\displaystyle=||-\sum_{i=1}^{p}\lambda_{i}(T_{\mu_{0}\rightarrow\mu_{i}}-\operatorname*{Id})||_{\mu_{0}}^{2}$
$\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}\left\langle\sum_{i=1}^{p}\lambda_{i}(T_{\mu_{0}\rightarrow\mu_{i}}-\operatorname*{Id}),\sum_{j=1}^{p}\lambda_{j}(T_{\mu_{0}\rightarrow\mu_{j}}-\operatorname*{Id})\right\rangle
d\mu_{0}$ $\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}\sum_{i,j=1}^{p}\lambda_{i}\lambda_{j}\langle
T_{\mu_{0}\rightarrow\mu_{i}}-\operatorname*{Id},T_{\mu_{0}\rightarrow\mu_{j}}-\operatorname*{Id}\rangle
d\mu_{0}$ $\displaystyle=$
$\displaystyle\sum_{i,j=1}^{p}\lambda_{i}\lambda_{j}A_{ij}=\lambda^{T}A\lambda.$

To establish the convexity of the function in terms of $\lambda$, it is
sufficient to show that $A$ is symmetric positive semi-definite. Observe that
$A$ is a Gram matrix of the set
$\\{T_{\mu_{0}\rightarrow\mu_{i}}-\operatorname*{Id}\\}_{i=1}^{p}\subset
T_{\mu_{0}}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ with the inner
product in $T_{\mu_{0}}\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$. It
is classical that all Gram matrices are symmetric positive semi-definite
(Schwerdtfeger, 1961) and this completes the proof. ∎

###  2.2 Proof of Theorem 1

###### Proof.

The convexity of the objective is established in Proposition 1, and the
feasible set $\Delta^{p}$ is trivially convex.

By Theorem 4, when assumptions A1-A3 hold, all Karcher means are barycenters,
and therefore $\mu_{0}=\nu_{\lambda}$, and it is always the case that
barycenters are Karcher means. Therefore it is sufficient to check that the
minimum value is 0 if and only if $\mu_{0}$ is a Karcher mean.

In the case that the minimum is zero, there is some $\lambda_{*}$ such that
$(\lambda_{*})^{T}A\lambda_{*}=0$ and we can apply Proposition 1 to show that
$\mu_{0}$ is a Karcher mean for $\lambda_{*}$. Similarly if $\mu_{0}$ is a
Karcher mean for some $\lambda_{*}$ by definition we have $||\nabla
G_{\lambda_{*}}(\mu_{0})||_{\mu_{0}}^{2}=0$ and again by Proposition 1 we have
that $\lambda_{*}$ achieves the minimum of zero in the optimization. ∎

###  2.3 Proof of Corollary 1

###### Proof.

A1,A2,A3 are clearly satisfied in the Gaussian case with
$\Omega=\mathbb{R}^{d}$ since the pdf of a non-degenerate multivariate
Gaussian is positive everywhere and Lipschitz continuous. Furthermore the
barycenter of zero-mean Gaussians is itself a zero-mean Gaussian ((Agueh &
Carlier, 2011) Theorem 6.1).

Therefore we can apply Theorem 1 and all that remains is to calculate $A_{ij}$
as follows. Noting that the $C_{i}$ are all symmetric (Agueh & Carlier, 2011)
we have

$\displaystyle A_{ij}$
$\displaystyle=\int_{\mathbb{R}^{d}}\left\langle(T_{i}(x)-\operatorname*{Id}(x)),(T_{j}(x)-\operatorname*{Id}(x))\right\rangle
dS_{0}(x)$
$\displaystyle=\int_{\mathbb{R}^{d}}\left\langle(C_{i}-I)x,(C_{j}-I)x\right\rangle
dS_{0}(x)$ $\displaystyle(\text{Eq. \ref{eq:map_mat}})$
$\displaystyle=\underset{X\sim\mathcal{N}(0,S_{0})}{\mathbb{E}}\left[X^{T}\left(C_{i}-I\right)^{T}\left(C_{j}-I\right)X\right]$
$\displaystyle=\underset{X\sim\mathcal{N}(0,S_{0})}{\mathbb{E}}\left[X^{T}\left(C_{i}-I\right)\left(C_{j}-I\right)X\right]$
$\displaystyle(C_{i}\text{ all sym.})$
$\displaystyle=\operatorname{Tr}\left(\left(C_{i}-I\right)\left(C_{j}-I\right)S_{0}\right).$

The last equality comes from the identity
$\underset{X\sim\mathcal{N}(0,S)}{\mathbb{E}}\left[X^{T}BX\right]=\operatorname{Tr}(BS)$
which holds for all $B\in\mathbb{R}^{d\times d}$ and
$S\in\mathbb{S}^{d}_{++}.$ ∎

###  2.4 Proof of Theorem 2

###### Proof.

In this setting, the barycenter $\nu_{\lambda}$ can be expressed as

$\nu_{\lambda}=\left[\sum_{i=1}^{p}\lambda_{i}T_{0}^{i}\right]\\#\mu_{0}$

((Panaretos & Zemel, 2020) Theorem 3.1.9). Together with Proposition 1, this
implies (recalling $T_{i}=T_{0}^{i}$)

$W_{2}^{2}(\mu_{0},\nu_{\lambda})=\int||x-\sum_{i=1}^{p}\lambda_{i}T_{i}(x)||_{2}^{2}d\mu_{0}(x)=||\nabla
G_{\lambda}(\mu_{0})||_{\mu_{0}}^{2}=\lambda^{T}A\lambda$

and therefore

$\operatorname*{arg\,min}_{\lambda\in\Delta^{p}}W_{2}^{2}(\mu_{0},\nu_{\lambda})=\operatorname*{arg\,min}_{\lambda\in\Delta^{p}}\lambda^{T}A\lambda.$

∎

###  2.5 Exact Statement and Proof of Theorem 3

######  Theorem 7.

(Formal Version of Theorem 3, Convergence Rate of $A_{ij}$) Let $T_{1},T_{2}$
be the optimal transport maps from $\mu_{0}$ to $\mu_{1}$, $\mu_{2}$
respectively. Suppose that A4, A5, A6 are satisfied for both pairs
$(\mu_{0},\mu_{1})$ and $(\mu_{0},\mu_{2})$. Let
$X_{1},...,X_{2n}\sim\mu_{0},Y_{1},...,Y_{n}\sim\mu_{1},Z_{1},...,Z_{n}\sim\mu_{2}$.
Let $\hat{T}_{1}$ and $\hat{T}_{2}$ be the entropic maps from $\mu_{0}$ to
$\mu_{1}$ and $\mu_{2}$ respectively, both with $\epsilon\asymp
n^{-1/(d^{\prime}+\bar{\alpha}+1)}$ and computed using $X_{1},...,X_{n}$. Then
we have

$\displaystyle\mathbb{E}\left[\left|\int\langle
T_{1}-\operatorname*{Id},T_{2}-\operatorname*{Id}\rangle
d\mu_{0}-\frac{1}{n}\sum_{i=n+1}^{2n}\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle\lesssim$
$\displaystyle\frac{1}{\sqrt{n}}+n^{-\frac{\bar{\alpha}+1}{4(d^{\prime}+\bar{\alpha}+1)}}\sqrt{\log
n}\sqrt{1+I_{0}(\mu_{0},\mu_{1})+I_{0}(\mu_{0},\mu_{2})}$

where $d^{\prime}=2\lceil d/2\rceil$, $\bar{\alpha}=\alpha\wedge 3$ and
$I_{0}(\mu_{0},\mu_{1})$ is the integrated Fisher information along the
Wasserstein geodesic between $\mu$ and $\mu_{1}$.

Before we begin the proof, we remark that by writing $\mu_{i},\mu_{j}$ instead
of $\mu_{1},\mu_{2}$, the random variable in the expectation is just
$|A_{ij}-\hat{A}_{ij}|$ and therefore this result allows us to control the
entry-wise deviations of $\hat{A}_{ij}$ from their true values.

###### Proof.

Throughout, unless otherwise noted, the expectation is with respect to all of
$X_{1},...,X_{2n},Y_{1},...,Y_{n}$ and $Z_{1},...,Z_{n}$. We also use the
notations $X^{n}$ to denote the set of variables $(X_{1},...,X_{n})$, and
similarly for $Y^{n}$ and $Z^{n}$.

First observe that we can assume without loss of generality that $0\in\Omega$;
see (Peyré & Cuturi, 2020) Remark 2.9 for invariance of $T_{1},T_{2}$ under
translation. For the translation invariance of $\hat{T}_{1},\hat{T}_{2}$,
observe that both obtaining $g_{\epsilon}$, and evaluating
$\hat{T}_{1},\hat{T}_{2}$ at fixed points, requires only the distances
$||X_{i}-Y_{j}||_{2}^{2}=||(X_{i}-t)-(Y_{j}-t)||_{2}^{2}$. We calculate:

$\displaystyle\mathbb{E}\left[\left|\int\langle
T_{1}-\operatorname*{Id},T_{2}-\operatorname*{Id}\rangle
d\mu_{0}-\frac{1}{n}\sum_{i=n+1}^{2n}\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle=$ $\displaystyle\mathbb{E}\left[\left|\int\langle
T_{1}-\operatorname*{Id},T_{2}-\operatorname*{Id}\rangle
d\mu_{0}-\frac{1}{n}\sum_{i=n+1}^{2n}\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle\right.\right.$
$\displaystyle\hskip
28.45274pt+\left.\left.\frac{1}{n}\sum_{i=n+1}^{2n}\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle-\frac{1}{n}\sum_{i=n+1}^{2n}\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle\leq$ $\displaystyle\mathbb{E}\left[\left|\int\langle
T_{1}-\operatorname*{Id},T_{2}-\operatorname*{Id}\rangle
d\mu_{0}-\frac{1}{n}\sum_{i=n+1}^{2n}\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle\hskip
28.45274pt+\mathbb{E}\left[\left|\frac{1}{n}\sum_{i=n+1}^{2n}\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle-\frac{1}{n}\sum_{i=n+1}^{2n}\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$

We consider the first and second terms separately. For brevity, let $h$ be the
function defined by

$h(x)=\langle T_{1}(x)-x,T_{2}(x)-x\rangle.$

Note that $h$ is bounded on $\Omega$ since

$|h(x)|=|\langle
T_{1}(x)-x,T_{2}(x)-x\rangle|\leq||T_{1}(x)-x||_{2}||T_{2}(x)-x||_{2}\leq(2|\Omega|)(2|\Omega|)=4|\Omega|^{2},$

which implies by Popoviciu’s inequality (Popoviciu, 1935) that
$\operatorname*{Var}[h(X)]\leq 16|\Omega|^{4}$. From this it follows

$\displaystyle\mathbb{E}\left[\left|\int\langle
T_{1}-\operatorname*{Id},T_{2}-\operatorname*{Id}\rangle
d\mu_{0}-\frac{1}{n}\sum_{i=n+1}^{2n}\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle\right|\right]$ $\displaystyle=$
$\displaystyle\mathbb{E}\left[\left|\mathbb{E}_{X\sim\mu_{0}}[h(X)]-\frac{1}{n}\sum_{i=n+1}^{2n}h(X_{i})\right|\right]$
$\displaystyle\leq$
$\displaystyle\sqrt{\frac{\operatorname*{Var}[h(X)]}{n}}\leq\sqrt{\frac{16|\Omega|^{4}}{n}}\lesssim\frac{1}{\sqrt{n}}$

where we have used that for all i.i.d. random variables $U,U_{1},...,U_{n}$
with finite variance

$\mathbb{E}[|\mathbb{E}[U]-\frac{1}{n}\sum_{i=1}^{n}U_{i}|]\leq\sqrt{\frac{\operatorname*{Var}[U]}{n}}.$

We now handle the second term:

$\displaystyle\mathbb{E}\left[\left|\frac{1}{n}\sum_{i=n+1}^{2n}\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle-\frac{1}{n}\sum_{i=n+1}^{2n}\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle\leq$
$\displaystyle\frac{1}{n}\sum_{i=n+1}^{2n}\mathbb{E}\left[\left|\langle
T_{1}(X_{i})-X_{i},T_{2}(X_{i})-X_{i}\rangle-\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle=$ $\displaystyle\mathbb{E}\left[\left|\langle
T_{1}(X_{n+1})-X_{n+1},T_{2}(X_{n+1})-X_{n+1}\rangle-\langle\hat{T}_{1}(X_{n+1})-X_{n+1},\hat{T}_{2}(X_{n+1})-X_{n+1}\rangle\right|\right]$
$\displaystyle=$ $\displaystyle\mathbb{E}\left[\left|\langle
T_{1}(X_{n+1}),T_{2}(X_{n+1})\rangle-\langle\hat{T}_{1}(X_{n+1}),\hat{T}_{2}(X_{n+1})\rangle\right.\right.$
$\displaystyle\hskip
85.35826pt+\left.\left.\langle\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1}),X_{n+1}\rangle+\langle\hat{T}_{2}(X_{n+1})-T_{2}(X_{n+1}),X_{n+1}\rangle\right|\right]$
$\displaystyle\leq$ $\displaystyle\mathbb{E}\left[\left|\langle
T_{1}(X_{n+1}),T_{2}(X_{n+1})\rangle-\langle\hat{T}_{1}(X_{n+1}),\hat{T}_{2}(X_{n+1})\rangle\right|\right]$
$\displaystyle\hskip
85.35826pt+\mathbb{E}\left[\left|\langle\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1}),X_{n+1}\rangle\right|\right]+\mathbb{E}\left[\left|\langle\hat{T}_{2}(X_{n+1})-T_{2}(X_{n+1}),X_{n+1}\rangle\right|\right]$

Here we can control the three terms separately. We first take care of the
middle term and the last term follows by an identical argument.

$\displaystyle\mathbb{E}\left[\left|\langle\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1}),X_{n+1}\rangle\right|\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[\left|\langle\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1}),X_{n+1}\rangle\right|\right]$
$\displaystyle\leq$
$\displaystyle\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[||\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1})||_{2}||X_{n+1}||_{2}\right]$
$\displaystyle(\text{Cauchy-Schwarz})$ $\displaystyle\leq$
$\displaystyle|\Omega|\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[||\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1})||_{2}\right]$
$\displaystyle(||X_{n+1}||\leq|\Omega|)$ (10) $\displaystyle=$
$\displaystyle|\Omega|\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[\mathbb{E}_{X_{n+1}}||\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1})||_{2}\right]$
$\displaystyle=$
$\displaystyle|\Omega|\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[\mathbb{E}_{X_{n+1}}\sqrt{||\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1})||_{2}^{2}}\right]$
$\displaystyle\leq$
$\displaystyle|\Omega|\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[\sqrt{\mathbb{E}_{X_{n+1}}||\hat{T}_{1}(X_{n+1})-T_{1}(X_{n+1})||_{2}^{2}}\right]$
$\displaystyle(\text{Jensen})$ $\displaystyle=$
$\displaystyle|\Omega|\mathbb{E}_{X^{n},Y^{n}}\left[||\hat{T}_{1}-T_{1}||_{L^{2}(\mu_{0})}\right]$
$\displaystyle\leq$
$\displaystyle|\Omega|\sqrt{\mathbb{E}_{X^{n},Y^{n}}\left[||\hat{T}_{1}-T_{1}||_{L^{2}(\mu_{0})}^{2}\right]}$
$\displaystyle\lesssim$
$\displaystyle\sqrt{(1+I_{0}(\mu_{0},\mu_{1}))n^{-\frac{\bar{\alpha}+1}{2(d^{\prime}+\bar{\alpha}+1)}}\log(n)}.$
$\displaystyle(\text{Theorem \ref{thm:PNW}})$

As mentioned above the third term is exactly the same except replacing $T_{1}$
and $\hat{T}_{1}$ with $T_{2}$ and $\hat{T_{2}}$ as well as $Y^{n}$ with
$Z^{n}$.

To control the first term we start with a different trick and then end up
following the same pattern as in the bound above.

$\displaystyle\mathbb{E}\left[\left|\langle
T_{1}(X_{n+1}),T_{2}(X_{n+1})\rangle-\langle\hat{T}_{1}(X_{n+1}),\hat{T}_{2}(X_{n+1})\rangle\right|\right]$
$\displaystyle=$ $\displaystyle\mathbb{E}\left[\left|\langle
T_{1}(X_{n+1})-\hat{T}_{1}(X_{n+1}),T_{2}(X_{n+1})\rangle+\langle\hat{T}_{1}(X_{n+1}),T_{2}(X_{n+1})-\hat{T}_{2}(X_{n+1})\rangle\right|\right]$
$\displaystyle\leq$ $\displaystyle\mathbb{E}\left[\left|\langle
T_{1}(X_{n+1})-\hat{T}_{1}(X_{n+1}),T_{2}(X_{n+1})\rangle\right|\right]+\mathbb{E}\left[\left|\langle\hat{T}_{1}(X_{n+1}),T_{2}(X_{n+1})-\hat{T}_{2}(X_{n+1})\rangle\right|\right]$
$\displaystyle\leq$
$\displaystyle\mathbb{E}\left[||T_{1}(X_{n+1})-\hat{T}_{1}(X_{n+1})||_{2}||T_{2}(X_{n+1})||\right]+\mathbb{E}\left[||\hat{T}_{1}(X_{n+1})||_{2}||T_{2}(X_{n+1})-\hat{T}_{2}(X_{n+1})||_{2}\right]$
$\displaystyle\leq$
$\displaystyle|\Omega|\mathbb{E}\left[||T_{1}(X_{n+1})-\hat{T}_{1}(X_{n+1})||_{2}\right]+|\Omega|\mathbb{E}\left[||T_{2}(X_{n+1})-\hat{T}_{2}(X_{n+1})||_{2}\right]$
$\displaystyle=$
$\displaystyle|\Omega|\mathbb{E}_{X^{n},Y^{n},X_{n+1}}\left[||T_{1}(X_{n+1})-\hat{T}_{1}(X_{n+1})||_{2}\right]+|\Omega|\mathbb{E}_{X^{n},Z^{n},X_{n+1}}\left[||T_{2}(X_{n+1})-\hat{T}_{2}(X_{n+1})||_{2}\right].$
(11)

Above we have made use of the fact that both $T_{2}(X_{1})$ and
$\hat{T_{1}}(X_{1})$ are contained in $\Omega$ which implies that
$||T_{2}(X_{1})||_{2}\leq|\Omega|$. $T_{2}(X_{1})\in\Omega$ follows from the
fact that $T_{2}$ is a map from $\mu_{0}$ to $\mu_{2}$. To see that
$||\hat{T}_{1}(X_{1})||\leq|\Omega|$, note that from equation (7)
$\hat{T}_{1}(X_{1})$ is a convex combination of $Y_{1},...,Y_{n}\in\Omega$ and
each of which satisfies $||Y_{i}||_{2}\leq|\Omega|$.

Observe that both terms in (11) have already been controlled, starting from
Equation (10) (the latter requires replacing $T_{1},\hat{T}_{1}$, and $Y^{n}$
with $T_{2},\hat{T}_{2}$, and $Z^{n}$ respectively). Applying the bound
derived there and combining terms terms we have the following result

$\displaystyle\mathbb{E}\left[\left|\int\langle
T_{1}-\operatorname*{Id},T_{2}-\operatorname*{Id}\rangle
d\mu_{0}-\frac{1}{n}\sum_{i=n+1}^{2n}\langle\hat{T}_{1}(X_{i})-X_{i},\hat{T}_{2}(X_{i})-X_{i}\rangle\right|\right]$
$\displaystyle\lesssim\frac{1}{\sqrt{n}}+2\sqrt{(1+I_{0}(\mu_{0},\mu_{1}))n^{-\frac{\bar{\alpha}+1}{2(d^{\prime}+\bar{\alpha}+1)}}\log(n)}+2\sqrt{(1+I_{0}(\mu_{0},\mu_{2}))n^{-\frac{\bar{\alpha}+1}{2(d^{\prime}+\bar{\alpha}+1)}}\log(n)}$
$\displaystyle=\frac{1}{\sqrt{n}}+n^{-\frac{\bar{\alpha}+1}{4(d^{\prime}+\bar{\alpha}+1)}}\sqrt{\log
n}\left(\sqrt{1+I_{0}(\mu_{0},\mu_{1})}+\sqrt{1+I_{0}(\mu_{0},\mu_{2})}\right)$
$\displaystyle\lesssim\frac{1}{\sqrt{n}}+n^{-\frac{\bar{\alpha}+1}{4(d^{\prime}+\bar{\alpha}+1)}}\sqrt{\log
n}\sqrt{1+I_{0}(\mu_{0},\mu_{1})+I_{0}(\mu_{0},\mu_{2})}.$

∎

###  2.6 Proof of Corollary 2

###### Proof.

Let $B_{n}$ denote the entrywise bound in Theorem 3. Noting that
$\hat{\lambda}^{T}\hat{A}\operatorname*{\hat{\lambda}}\leq\lambda_{*}^{T}\hat{A}\lambda_{*}$
by construction, we estimate

$\displaystyle\operatorname*{\mathbb{E}}[\hat{\lambda}^{T}A\operatorname*{\hat{\lambda}}]$
$\displaystyle=\operatorname*{\mathbb{E}}\left[\hat{\lambda}^{T}(A-\hat{A})\operatorname*{\hat{\lambda}}\right]+\operatorname*{\mathbb{E}}[\hat{\lambda}^{T}\hat{A}\hat{\lambda}]$
$\displaystyle\leq\operatorname*{\mathbb{E}}\left[|\hat{\lambda}^{T}(A-\hat{A})\operatorname*{\hat{\lambda}}|\right]+\operatorname*{\mathbb{E}}[\lambda_{*}^{T}\hat{A}\lambda_{*}]$
$\displaystyle=\operatorname*{\mathbb{E}}\left[|\hat{\lambda}^{T}(A-\hat{A})\operatorname*{\hat{\lambda}}|\right]+\operatorname*{\mathbb{E}}[\lambda_{*}^{T}(\hat{A}-A)\lambda_{*}]\,\,$
$\displaystyle(\lambda_{*}^{T}A\lambda_{*}=0)$
$\displaystyle=\operatorname*{\mathbb{E}}\left[\bigg{|}\sum_{i,j=1}^{p}(\hat{\lambda})_{i}(\hat{\lambda})_{j}(A-\hat{A})_{ij}\bigg{|}\right]+\operatorname*{\mathbb{E}}\left[\bigg{|}\sum_{i,j=1}^{p}(\lambda_{*})_{i}(\lambda_{*})_{j}(A-\hat{A})_{ij}\bigg{|}\right]$
$\displaystyle\leq\operatorname*{\mathbb{E}}\left[\sum_{i,j=1}^{p}(\hat{\lambda})_{i}(\hat{\lambda})_{j}|A_{ij}-\hat{A}_{ij}|\right]+\operatorname*{\mathbb{E}}\left[\sum_{i,j=1}^{p}(\lambda_{*})_{i}(\lambda_{*})_{j}|A_{ij}-\hat{A}_{ij}|\right]$
$\displaystyle(\hat{\lambda},\lambda_{*}\in\Delta^{p},\text{ triangle ineq.})$
$\displaystyle\leq
2\operatorname*{\mathbb{E}}\left[\sum_{i,j=1}^{p}|A_{ij}-\hat{A}_{ij}|\right]$
$\displaystyle=2\sum_{i,j=1}^{p}\operatorname*{\mathbb{E}}[|A_{ij}-\hat{A}_{ij}|]$
$\displaystyle\lesssim 2p^{2}B_{n}$ $\displaystyle(\text{Theorem
\ref{thm:convergence}})$

Since $A$ is positive semidefinite and by assumption
$\lambda_{*}\in\Delta^{p}$ satisfies $\lambda_{*}^{T}A\lambda_{*}=0$, it
follows that $\lambda_{*}$ is an eigenvector of $A$ with eigenvalue 0. Let
$0<\alpha_{2}\leq\dots\leq\alpha_{p}$ be the non-zero eigenvalues of $A$ with
associated orthonormal eigenvectors $v_{2},\dots,v_{p}$. Orthogonally
decompose
$\operatorname*{\hat{\lambda}}=\hat{\beta}\lambda_{*}+\operatorname*{\hat{\lambda}}_{\perp}$,
where $\hat{\beta}\in\mathbb{R}$ and $\operatorname*{\hat{\lambda}}_{\perp}$
is in the span of $\\{v_{2},\dots,v_{p}\\}$. Note that $\hat{\beta}$ and
$\hat{\lambda}_{\perp}$ are random. Then,

$\displaystyle\operatorname*{\mathbb{E}}[\|\operatorname*{\hat{\lambda}}-\hat{\beta}\lambda_{*}\|_{2}^{2}]=$
$\displaystyle\operatorname*{\mathbb{E}}[\|\hat{\lambda}_{\perp}\|_{2}^{2}]$
$\displaystyle=$
$\displaystyle\operatorname*{\mathbb{E}}\left[\sum_{i=2}^{p}|v_{i}^{T}\hat{\lambda}_{\perp}|^{2}\right]$
$\displaystyle\leq$
$\displaystyle\frac{1}{\alpha_{2}}\operatorname*{\mathbb{E}}\left[\sum_{i=2}^{p}\alpha_{i}|v_{i}^{T}\hat{\lambda}_{\perp}|^{2}\right]$
$\displaystyle=$
$\displaystyle\frac{1}{\alpha_{2}}\operatorname*{\mathbb{E}}[|(\hat{\lambda}_{\perp})^{T}A\hat{\lambda}_{\perp}|]$
$\displaystyle=$
$\displaystyle\frac{1}{\alpha_{2}}\operatorname*{\mathbb{E}}[|\hat{\lambda}^{T}A\operatorname*{\hat{\lambda}}|]$
$\displaystyle\lesssim$ $\displaystyle\frac{2p^{2}}{\alpha_{2}}B_{n}.$

Summing both sides of the equation
$\hat{\lambda}=\hat{\beta}\lambda_{*}+\hat{\lambda}_{\perp}$ and recalling
$\lambda_{*},\hat{\lambda}\in\Delta^{p}$ yields

$\displaystyle 1$
$\displaystyle=\hat{\beta}+\sum_{j=1}^{p}(\hat{\lambda}_{\perp})_{j}$
$\displaystyle\leq\hat{\beta}+\|\hat{\lambda}_{\perp}\|_{1}$
$\displaystyle\leq\hat{\beta}+\sqrt{p}\|\hat{\lambda}_{\perp}\|_{2},$

which implies that $\mathbb{E}[(1-\hat{\beta})^{2}]\leq
p\mathbb{E}[\|\hat{\lambda}_{\perp}\|_{2}^{2}]\lesssim\frac{2p^{3}}{\alpha_{2}}B_{n}$.

Finally, we use the fact that
$\operatorname*{\hat{\lambda}}-\hat{\beta}\lambda_{*}=\hat{\lambda}_{\perp}$
and $(\hat{\beta}-1)\lambda_{*}$ are orthogonal to bound:

$\displaystyle\operatorname*{\mathbb{E}}[\|\operatorname*{\hat{\lambda}}-\lambda_{*}\|_{2}^{2}]$
$\displaystyle=\operatorname*{\mathbb{E}}[\|\operatorname*{\hat{\lambda}}-\hat{\beta}\lambda_{*}\|_{2}^{2}+\|(\hat{\beta}-1)\lambda_{*}\|_{2}^{2}]$
$\displaystyle=\operatorname*{\mathbb{E}}[\|\operatorname*{\hat{\lambda}}-\hat{\beta}\lambda_{*}\|_{2}^{2}]+\operatorname*{\mathbb{E}}[\|(\hat{\beta}-1)\lambda_{*}\|_{2}^{2}]$
$\displaystyle\lesssim\frac{2p^{2}}{\alpha_{2}}B_{n}+\operatorname*{\mathbb{E}}[(\hat{\beta}-1)^{2}\|\lambda_{*}\|_{2}^{2}]$
$\displaystyle\leq\frac{2p^{2}}{\alpha_{2}}B_{n}+\operatorname*{\mathbb{E}}[(\hat{\beta}-1)^{2}]$
$\displaystyle\lesssim\frac{2p^{2}}{\alpha_{2}}B_{n}+\frac{2p^{3}}{\alpha_{2}}B_{n}$

as desired.

∎

###  2.7 Karcher Means are Barycenters under A4-A6

Corollary 2 (which holds under A4-A6) can be combined with Theorem 1 (which
holds under A1-A3) as follows. Suppose $||G_{\lambda}(\nu)||_{\nu}^{2}=0$ for
some $\lambda\in\Delta^{p}$ and reference measures $\\{\mu_{i}\\}_{i=1}^{p}$,
and assume that the collection $\\{\nu\\}\cup\\{\mu_{i}\\}_{i=1}^{p}$ satisfy
assumptions A4-A6. We aim to show that this implies that $\nu$ is a barycenter
for the measures $\\{\mu_{i}\\}_{i=1}^{p}$.

Let $\tilde{\mu}_{i}$ be the measure with density

$d\tilde{\mu}_{i}(x)=\begin{cases}d\mu_{i}(x)&x\in\Omega^{\circ},\\\
0&\text{otherwise}\end{cases}$

where $\Omega^{\circ}$ is the interior of $\Omega$. Define $\tilde{\nu}$
similarly.

We first remark that

$1=\mu_{i}[\Omega]=\mu_{i}[\Omega^{\circ}]+\mu_{i}[\partial\Omega]=\tilde{\mu}_{i}[\Omega^{\circ}]+0=\tilde{\mu}_{i}[\Omega^{\circ}]$

and therefore $\tilde{\mu}_{i}$ (and analogously $\nu$) are valid measures
(note we have used that $\mu_{i}[\partial\Omega]=0$, which follows from
$\mu_{i}$ being a.c. and $\partial\Omega$ having Lebesgue measure 0). It
follows from A4 that the collection $\\{\nu\\}\cup\\{\mu_{i}\\}_{i=1}^{p}$
satisfies A2 and A3. Furthermore, the set $\Omega^{\circ}$ is open and convex
and therefore A1 is also satisfied. Therefore if $\tilde{\nu}$ is a Karcher
mean for $\\{\tilde{\mu}_{i}\\}$ then it is a barycenter for them as well.

We next show that $\tilde{\nu}$ is a Karcher mean of
$\\{\tilde{\mu}_{i}\\}_{i=1}^{p}$ with coordinate $\lambda$. Consider the maps
$\tilde{T}_{i}:\Omega^{\circ}\rightarrow\Omega^{\circ}$ which are defined by
$\tilde{T}_{i}(x)=T_{i}(x)$ and are undefined outside $\Omega^{\circ}$ where
$T_{i}$ is the optimal transport map from $\nu\rightarrow\mu_{i}$. We first
require that $\tilde{T}_{i}$ is in fact optimal for $\tilde{\nu}$ and
$\tilde{\mu}_{i}$ and that it is well-defined $\tilde{\nu}$-a.e.

To show the latter, note that $T_{i}$ is a map from $\nu$ to $\mu_{i}$ and
therefore $\nu[T_{i}^{-1}(\partial\Omega)]=\mu_{i}[\partial\Omega]=0$ which
implies that $T_{i}(x)\in\Omega^{\circ}$ for $\nu$-a.e. $x$. Furthermore
$\tilde{\nu}[\Omega^{\circ}]=\nu[\Omega^{\circ}]=1$ and therefore the map
$\tilde{T}_{i}$ is well-defined $\tilde{\nu}$-a.e.

To show that the map $\tilde{T}_{i}$ is optimal, one only needs to remark that
for any measurable $B\subset\Omega$ that $T_{i}$ transports mass optimally
between $B$ and $T_{i}(B)$ (since otherwise $T_{i}$ would fail to be optimal).
Applying this fact to $\Omega^{\circ}$ and using the considerations above,
shows that $\tilde{T}_{i}$ is optimal for transporting between $\tilde{\nu}$
and $\tilde{\mu}_{i}$.

Having established that the $\tilde{T}^{i}$ are optimal, we can now calculate
$||G_{\lambda}(\tilde{\nu})||_{\tilde{\nu}}^{2}$:

$\displaystyle||G_{\lambda}(\tilde{\nu})||_{\tilde{\nu}}^{2}$
$\displaystyle=\sum_{i,j=1}^{p}\lambda_{i}\lambda_{j}\int_{\Omega^{\circ}}\left\langle\tilde{T}_{i}(x)-x,\tilde{T}_{j}(x)-x\right\rangle
d\tilde{\nu}(x)$
$\displaystyle=\sum_{i,j=1}^{p}\lambda_{i}\lambda_{j}\int_{\Omega^{\circ}}\left\langle
T_{i}(x)-x,T_{j}(x)-x\right\rangle d\nu(x)$
$\displaystyle=\sum_{i,j=1}^{p}\lambda_{i}\lambda_{j}\int_{\Omega^{\circ}}\left\langle
T_{i}(x)-x,T_{j}(x)-x\right\rangle
d\nu(x)+\lambda_{i}\lambda_{j}\int_{\partial\Omega}\left\langle
T_{i}(x)-x,T_{j}(x)-x\right\rangle d\nu(x)$
$\displaystyle=\sum_{i,j=1}^{p}\lambda_{i}\lambda_{j}\int_{\Omega}\left\langle
T_{i}(x)-x,T_{j}(x)-x\right\rangle d\nu(x)$
$\displaystyle=||G_{\lambda}(\nu)||_{\nu}^{2}=0$

which shows that $\tilde{\nu}$ is a Karcher mean for $\\{\tilde{\mu}_{i}\\}$,
and is also a barycenter of these measures. To conclude, we only require the
fact that if two measures
$\xi,\tilde{\xi}\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ differ
only on a set of measure zero, then for any
$\gamma\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}$ we have
$W_{2}(\xi,\gamma)=W_{2}(\tilde{\xi},\gamma)$. Applying this fact repeatedly
we have

$\displaystyle\min_{\xi}\frac{1}{2}\sum_{i=1}^{n}\lambda_{i}W_{2}^{2}(\xi,\mu_{i})$
$\displaystyle=\min_{\xi}\frac{1}{2}\sum_{i=1}^{n}\lambda_{i}W_{2}^{2}(\xi,\tilde{\mu}_{i})$
$\displaystyle=\frac{1}{2}\sum_{i=1}^{n}\lambda_{i}W_{2}^{2}(\tilde{\nu},\tilde{\mu}_{i})$
$\displaystyle=\frac{1}{2}\sum_{i=1}^{n}\lambda_{i}W_{2}^{2}(\nu,\tilde{\mu}_{i})$
$\displaystyle=\frac{1}{2}\sum_{i=1}^{n}\lambda_{i}W_{2}^{2}(\nu,\mu_{i})$

which shows that $\nu$ is indeed a barycenter of $\mu_{i}$, as desired.

##  3 Covariance Estimation Experiment Details

###  3.1 Methods for Maximum Likelihood Estimation

In order to solve the maximum likelihood estimation problem we differentiate
through a truncated version of (Chewi et al., 2020) Algorithm 1 and perform
projected gradient descent. This is implemented using the auto-differentiation
library PyTorch. The procedure is summarized in Algorithm 3.

Algorithm 3 MLE

Input: $\\{S_{i}\\}_{i=0}^{p}$, $\eta>0$, MaxIters > 0, FPIters > 0, SQIters >
0

$i\leftarrow 0,\lambda\leftarrow(1/p)\bm{1}$

while Not Converged and $i<$ MaxIters do

$i\leftarrow i+1$

$\nabla\mathcal{L}(\lambda)\leftarrow$ BackPropLoss($\\{S_{j}\\}_{j=0}^{p}$,
$\lambda$, FPIters, SQIters)

$\lambda\leftarrow$ SimplexProject($\lambda-\eta\nabla\mathcal{L}(\lambda)$)

end while

Return $\lambda$

Here SimplexProject$(x)$ is defined as

$\text{SimplexProject}(x)=\operatorname*{arg\,min}_{y\in\Delta^{p}}||x-y||_{2}.$

which enforces the constraints on $\lambda$ at each iteration.

To compute $\nabla\mathcal{L}(\lambda)$ we use auto-differentiation to obtain
a gradient of the procedure given in Algorithm 4. In this procedure the square
root of a matrix is computed using SQIters number of Newton-Schulz iterations.
The forward pass of the loss computation is given in Algorithm 4, and the
gradient is obtained by back propagation through it.

Algorithm 4 ComputeLoss

Input: $\\{S_{i}\\}_{i=0}^{p}$, $\lambda$, FPIters > 0, SQIters > 0

BC $\leftarrow S_{0}$

for $j=1,...,$ FPIters do

BC_root = SquareRoot(BC, SQIters)

BC_root_inv = Invert(BC_root)

BC = BC_root_inv $\left(\sum_{i=1}^{p}\lambda_{i}\text{SquareRoot( BC\\_root,
SQIters,}S_{i}\text{ BC\\_root })\right)$ BC_root_inv

end for

Return Tr($\text{BC}^{-1}S_{0}$) + log det BC

The parameters that we chose in our experiments are given in Table 3.

Parameter | Value
---|---
$\eta$ | 0.0003
MaxIters | 500
FPIters | 10
SQIters | 10
Table 3: Parameters used when performing the maximum likelihood estimation.

These parameters were sufficient to ensure that both the matrix square roots
and fixed point iterations converged, and $\lambda$ always converged before
the final iteration was reached.

###  3.2 Considerations for Instability of MLE

Due to the requirement of differentiating through a matrix inverse in
Algorithm 3 it is possible that the procedure may fail to numerically
converge. In our trials, this happened in less than 0.5% of all attempts. When
it does happen we discard the result, and these numerical failures impact
neither the results of Figure 3, nor the timings of the trials.

##  4 MNIST Experiment Details

###  4.1 The MNIST Dataset

The MNIST Dataset (LeCun, 1998) consists of a collection of $28\times 28$
black and white images of hand-written digits. The brightness (or intensity)
of each pixel is an integer between 0 and 255 with 0 being black and 255 being
white.

Each image is initially a matrix $I^{m}\in\mathbb{R}^{28\times 28}_{+}$, which
is normalized to sum to one,
$P^{m}_{ij}=I^{m}_{ij}/(\sum_{i^{\prime},j^{\prime}}I_{i^{\prime},j^{\prime}})$
and then converted into a point cloud

$\mu^{m}=\sum_{i,j=1}^{28}P^{m}_{ij}\delta_{(i,j)}.$

###  4.2 White Noise Model

In order to apply additive noise we must generate white-noise images $\zeta$.
We generate $\zeta$ as a random matrix $\mathbb{R}^{28\times 28}$ with each
entry $\zeta_{ij}\sim\text{Unif}([0,1])$. This random matrix is then converted
into a point cloud using the same process as is used to turn the image matrix
into a point cloud.

###  4.3 Linear Recovery

The linear reconstruction in Figure 4 method is done by the following
procedure.

  1. 1.

Given a corrupted image (which we will treat as a matrix) $\tilde{\mu}_{0}$
and a collection of references $\tilde{\mu}_{1},...,\tilde{\mu}_{p}$ (also
treated like matrices) estimate $\hat{\lambda}$ as

$\hat{\lambda}=\operatorname*{arg\,min}_{\lambda\in\Delta^{p}}||\tilde{\mu}_{0}-\sum_{i=1}^{p}\lambda_{i}\tilde{\mu}_{i}||_{F}^{2}.$

This corresponds to projecting $\tilde{\mu}_{0}$ onto the convex hull of
$\tilde{\mu}_{1},...,\tilde{\mu}_{p}$ with respect to the Euclidean
(equivalently Frobenius) norm.

  2. 2.

Recover $\hat{\mu}_{0}$ as

$\hat{\mu}_{0}=\sum_{i=1}^{p}\hat{\lambda}_{i}\mu_{i}.$

##  5 Document Classification Experiment Details

###  5.1 Datasets

Dataset | Description | B.O.W. Dimension | Average Words | $\mathcal{C}$
---|---|---|---|---
BBCSPORT | BBC sports articles labeled by sport | 13243 | 117 | 5
20NEWS | canonical news article dataset | 29671 | 72 | 20
Table 4: Dataset characteristics.

We show the characteristics of the datasets in Table 4. Bag of words (B.O.W.)
dimension is the number of the unique words contained in each dataset. Average
words is the average number of the unique words in a document and
$\mathcal{C}$ is the number of classes of each of the datasets.

Both the BBCSport and News20 datasets that we used are those made publicly
available from (Huang et al., 2016). Both datasets consist of two components:

  1. 1.

The datasets employ a word2vec (Mikolov et al., 2013) embedding to represent
the words. Each word is represented as a point $w$ on the unit sphere in
$\mathbb{R}^{300}$ and the semantic relationships between words are encoded in
the geometry of the embeddings. Call the set of words $\mathcal{W}$, then the
word2vec embedding is the matrix $W\in\mathbb{R}^{|\mathcal{W}|\times 300}$.

  2. 2.

Each document $D^{m}$ with label $y^{m}$ is represented by a set of tuples
$\\{(c^{m}_{i},\bm{w}^{m}_{i})\\}_{i=1}^{n}$, where $n$ is the number unique
words in $D^{m}$, $c^{m}_{i}$ is the number of times $\bm{w}^{m}_{i}$ shows in
$D^{m}$.

We use the following procedure to convert a document into an empirical
measure.

  1. 1.

Given a document $D^{m}=\\{(c^{m}_{i},\bm{w}^{m}_{i})\\}_{i=1}^{n}$, let
$\bm{p}^{m}\in\mathbb{R}^{n}$ be the probability vector with entries
$p^{m}_{i}=c^{m}_{i}/(\sum_{k}c^{m}_{k})$.

  2. 2.

The document empirical measure is then

$\mu^{m}=\sum_{i}p^{m}_{i}\delta_{\bm{w}^{m}_{i}}.$

For convenience we filter out the points on which $p^{i}$ places no support.

###  5.2 Procedure for Making Figure 5

As discussed in the main body, for $k\geq 2$ we select $k$ documents from each
topic and 100 test documents, and then repeat 50 times and average. In both
plots in Figure 5 we iterate over different choices of $k$. To reduce
computational load we reuse the test sets and reference documents within each
trial across choices of $k$. For example, all of the documents for $k=4$ are
used for $k\geq 5$. This allows the Wasserstein distances and estimates
$\hat{A}_{ij}$ to be cached avoiding further computation.

We note that this procedure produces plots in which each of the points
individually has the same sampling pattern as if there was no re-use, it is
only across samples that there is correlation. However, this correlation is
shared across the prediction methods consider which preserves the fairness of
the comparison of the methods.

##  6 Possible Rank Deficiency of A

In general, there may exist a set $\Lambda\subset\Delta^{p}$ such that
$|\Lambda|>1$ and for all $\lambda\in\Lambda$ we have $\mu_{0}=\nu_{\lambda}$.
Indeed, suppose there are index sets $I_{1},I_{2}$ such that $I_{1}\cap
I_{2}=\varnothing$ and $I_{1}\cup I_{2}=\\{1,2,\dots,p\\}$ but
$\operatorname*{Bary}(\\{\mu_{i}\\}_{i\in
I_{1}})\cap\operatorname*{Bary}(\\{\mu_{i}\\}_{i\in I_{2}})\neq\varnothing$.
If $\mu_{0}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i\in
I_{1}})\cap\operatorname*{Bary}(\\{\mu_{i}\\}_{i\in I_{2}})$ then it can be
expressed as $\nu_{\lambda^{1}}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i\in
I_{1}})$ and $\nu_{\lambda^{2}}\in\operatorname*{Bary}(\\{\mu_{i}\\}_{i\in
I_{2}}).$ Note that we may interpret $\lambda^{1}$ and $\lambda^{2}$ as
vectors in $\Delta^{p}$ with supports on $I_{1},I_{2}$, respectively. Now let
$\alpha\in[0,1]$ and let $\lambda=(1-\alpha)\lambda^{1}+\alpha\lambda^{2}$
(with the indices modified as needed). Then we have

$\operatorname*{arg\,min}_{\nu\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}}\sum_{i=1}^{p}\frac{\lambda_{i}}{2}W_{2}^{2}(\nu,\mu_{i})=\operatorname*{arg\,min}_{\nu\in\operatorname{\mathcal{P}_{2,ac}(\mathbb{R}^{d})}}(1-\alpha)\left(\sum_{i\in
I_{1}}\frac{\lambda^{1}_{i}}{2}W_{2}^{2}(\nu,\mu_{i})\right)+\alpha\left(\sum_{i\in
I_{2}}\frac{\lambda^{2}_{i}}{2}W_{2}^{2}(\nu,\mu_{i})\right)=\mu_{0}$

where the last equality follows from the fact that if $x$ minimizes both
$f(x)$ and $g(x)$ then it must also minimize $\alpha f(x)+(1-\alpha)g(x)$ for
$\alpha\in[0,1]$, and the assumption that $\mu_{0}$ is the barycenter. This
demonstrates that $\mu_{0}$ actually corresponds to all
$\lambda=\lambda=(1-\alpha)\lambda^{1}+\alpha\lambda^{2}$ for all
$\alpha\in[0,1]$.

Below, we will give a geometric characterization for the rank deficiency of
$A$ and how that relates to the optimality of the quadratic objective in
Theorem 1. We first recall the quadratic problem over the probability simplex:
$\min_{\lambda\in\Delta^{p}}\lambda^{T}A\lambda$. Since $A$ is a symmetric and
positive semidefinite matrix, a Rayleigh quotient analysis yields that the
minimum possible value of the objective without simplex constraints is zero,
which would be attained at an eigenvector corresponding to a zero eigenvalue.
Given this fact, finding an exact solution to the above quadratic problem
amounts to determining whether the probability simplex intersects the
eigenspace of $A$ associated with the zero eigenvalue. Hereafter, we will
denote this space by $E_{0}$.

######  Proposition 2.

If $E_{0}\cap\Delta^{p}=\emptyset$, there is no $\lambda\in\Delta^{p}$ such
that $\lambda^{T}A\lambda=0$.

###### Proof.

Any $\lambda\in\Delta^{p}$ realizing $\lambda^{T}A\lambda=0$ is an eigenvector
corresponding to the zero eigenvalue. The result follows. ∎

Proposition 2 shows in particular that exact minimization is only possible
when $\textrm{rank}(A)<p$. Next, we state conditions under which there is a
unique solution to the minimization program.

######  Proposition 3.

If $\textrm{rank}(A)=p-1$ and $E_{0}\cap\Delta^{p}\neq\emptyset$, then there
is a unique $\lambda\in\Delta^{p}$ such that $\lambda^{T}A\lambda=0$.

###### Proof.

If $\textrm{rank}(A)=p-1$, $E_{0}=\\{\alpha v\ |\ \alpha\in\mathbb{R}\\}$ for
an eigenvector $v\neq 0$ corresponding to the zero eigenvalue. Let $\lambda\in
E_{0}\cap\Delta^{p}$. Then $\lambda=\alpha v$ for some $\alpha\in\mathbb{R}$.
The constraint that $\lambda\in\Delta^{p}$ determines $\alpha$ uniquely, which
gives the result. ∎

Note that if $E_{0}$ non-trivially intersects the probability simplex and
$\textrm{rank}(A)=p-1$, the condition for uniqueness is equivalent to the
statement that the origin can be written as a convex combination of the
columns of $A$. We next consider the remaining case when
$\textrm{rank}(A)<p-1$; for simplicity of exposition, we focus on the specific
case $\textrm{rank}(A)=p-2$.

######  Proposition 4.

Let $\textrm{rank}(A)=p-2$ and $v_{0}$ and $v_{1}$ be distinct eigenvectors
corresponding to the zero eigenvalue of $A$. If $v_{0}\in\Delta^{p}$ and
$v_{1}\in\Delta^{p}$, then there are infinitely many solutions to
$\lambda^{T}A\lambda=0$ with $\lambda\in\Delta^{p}$.

###### Proof.

By construction, $Av_{0}=0$ and $Av_{1}=0$. Let $\alpha\in[0,1]$ and define
$v_{\alpha}=\alpha v_{0}+(1-\alpha)v_{1}$. Since $v_{1}$ and $v_{2}$ are in
$\Delta^{p}$, it follows that $v_{\alpha}\in\Delta^{p}$. In addition,
$Av_{\alpha}=\alpha Av_{1}+(1-\alpha)Av_{2}=0$ so that $v_{\alpha}\in E_{0}$.
Since $\alpha\in[0,1]$ was arbitrary, the result follows. ∎

The main conclusion from the above Propositions is the following: if
$\textrm{rank}(A)<p-1$, the only way to mandate a unique solution is to
require that only one eigenvector corresponding to the zero eigenvalue
intersects with probability simplex. By convex geometry, if two eigenvectors
corresponding to the zero eigenvalue intersect the probability simplex, there
will be infinitely many solutions to the quadratic program.

##  7 Example of a Karcher Mean that is not a Barycenter

As mentioned in Section 2, special care must be made to ensure that any
Karcher mean is also a barycenter. In this section we give a simple example
which illustrates this point.

Consider $S^{1}=\\{x\in\mathbb{R}^{2}:||x||_{2}=1\\}$ with distance given by
$\rho(x,y)=\cos^{-1}\langle x,y\rangle$. Let $z_{1}\neq z_{2}\in S^{1}$ with
$z_{1}\neq-z_{2}$, and $\lambda=(1/2,1/2)$. Then the barycenter is given by

$x_{\lambda}=\frac{z_{1}+z_{2}}{||z_{1}+z_{2}||_{2}}.$

However, the point $-\frac{z_{1}+z_{2}}{||z_{1}+z_{2}||_{2}}$ is not the
barycenter but it is a Karcher mean. This is illustrated in Figure 6. To show
this rigorously one need only compute the Euclidean gradient and demonstrate
that it is orthogonal to the tangent space.

Figure 6: _Red_ : The points $z_{1},z_{2}$. _Green_ : The true barycenter
$x_{\lambda}$. _Blue_ : A Karcher mean which is not the barycenter. _Gold_ :
The tangent space at the blue point.