k-Sample inference
via Multimarginal Optimal Transport

Let $\mu_{1},\ldots,\mu_{k}$ be Borel probability measures supported on $\mathcal{X}\subseteq\mathbb{R}^{d}$, $d\geq 1$. The Multimarginal Optimal Transport ($MOT$) problem (equation (4.3) of [1]) is the optimization problem

where $\mathcal{C}(\mu_{1},\ldots,\mu_{k})$ is the set of Borel probability measures on the product space $\mathcal{X}\times\cdots\times\mathcal{X}$ with marginals $\mu_{1},\ldots,\mu_{k}$. Different choices for the cost function $c(x_{1},\ldots,x_{k})$ are possible; throughout the paper, we fix the choice to be

Under this choice, the $MOT$ problem is equivalent (Proposition 3.1.2 of [60] to the Wasserstein barycenter problem (equation 2.2 of [1])

where $W_{2}(\cdot,\cdot)$ denotes 2-Wasserstein distance. By equivalence here we mean that the optimal values of both programs are equal, and the optimal solutions $\pi^{*}$ of (2) and $\nu^{*}$ of (4) are related by $\nu^{*}=M_{\#}\pi^{*}$ where $M$ is the map that averages a given $k$-tuple of points from the supports of the $k$ measures ($M\#\pi^{*}$ stands for pushforward of a measure $\pi^{*}$ by the map $M$). We remark here that when the measures are discrete, the barycenter problem (4) generally has more than one optimal solution $\nu^{*}$. This presents challenges for statistical inference concerning barycenter solutions [52], but does not impede the analysis of the optimal value of barycenter or $MOT$ problems (recalling that all optimal solutions result in the same optimal value).

Let $MOT(\mu_{1},\ldots,\mu_{k})$ denote the optimal value of the $MOT$ program (2). Observe that $MOT(\mu_{1},\ldots,\mu_{k})=0$ if and only if the $k$ measures $\mu_{1},\ldots,\mu_{k}$ are all the same. Indeed, if $\nu$ is (any) optimal solution to the barycenter problem (4), then $MOT(\mu_{1},\ldots,\mu_{k})=0$ is equivalent to zero optimal value in the barycenter program (4), i.e. $W_{2}^{2}(\mu_{i},\nu)=0$ for all $i=1,\ldots,k$. Due to metric properties of $2$-Wasserstein distance (Theorem 7.3 of [80]), this is equivalent to all the $k$ measures being the same (and equal to $\nu$).

This observation suggests that testing $H_{0}$ in (1) can be addressed via testing for $MOT(\mu_{1},\ldots,\mu_{k})=0$. To this end, suppose that the data on $k$ samples $(X^{1}_{i})_{i=1}^{n_{1}},\cdots,(X^{k}_{i})_{i=1}^{n_{k}}$ of sizes $n_{1},\cdots,n_{k}$, respectively, is available to estimate the underlying measures $\mu_{1},\cdots,\mu_{k}$ by the empirical measures $\hat{\mu}_{1}:=\frac{1}{n_{i}}\sum_{i=1}^{n_{1}}\delta_{x^{1}_{i}},\cdots,\hat{\mu}_{k}:=\frac{1}{n_{k}}\sum_{i=1}^{n_{k}}\delta_{x^{k}_{i}}$. To test for $MOT(\mu_{1},\ldots,\mu_{k})=0$ based on the data, we consider the asymptotic distribution $\mathcal{D}_{0}$ of the empirical estimator $MOT(\widehat{\mu}_{1},\ldots,\widehat{\mu}_{k})$ under $H_{0}$ and reject $H_{0}$ when the estimator value is large.

More generally, once the asymptotic distribution $\mathcal{D}$ of empirical $MOT$ is known, one can estimate various functionals of $\mathcal{D}$, such as Confidence Regions (CR’s) for a true $MOT$ value either $H_{0}$ or $H_{a}$ in (1). Inference of this type requires knowledge of the asymptotic distribution of empirical version of the $MOT$ value in (2). Our derivation of these distributions leverages rich literature on asymptotic theory for the Wasserstein distance, whose main results we briefly review below.

The squared 2-Wasserstein distance $W_{2}^{2}(\mu_{1},\mu_{2})$ is the optimal value of the problem

with $c(x_{1},x_{2})=\|x_{1}-x_{2}\|^{2}$, which can be viewed as a particular case of the $MOT$ problem (2) with $k=2$ measures. Being a true metric on a space of probability measures on a given metric space ([81]), the 2-Wasserstein distance $W_{2}$ (and, more generally, the $p$-Wasserstein distance $W_{p}$) provides a natural way to compare probability measures while respecting the geometry of the supporting metric space. Under this framework, the true measures are estimated by their empirical counterparts, and statistical inference is conducted using limiting laws for the empirical Wasserstein distance [59, 63].

The forms of the weak limits depend on two main factors: dimensionality of the support and the nature of the measures (where the cases $\mu=\nu$ and $\mu\neq\nu$ may have different limits). Letting $OT_{c}$ denote the optimal value in (5) (with possibly different costs $c$), the limiting laws have general form of

with $\rho_{n}=\sqrt{\frac{n_{1}n_{2}}{n_{1}+n_{2}}}$ (when only $\mu$ is estimated from the data while $\nu$ is not, the “one-sample” version with $\rho_{n}=\sqrt{n}$ is considered). When measures are supported on $\mathbb{R}$, the limits $L$ can be Gaussian, with variance that depending on the truth $\mu,\nu$ under the “alternative” assumption $\mu\neq\nu$ [57, 22], and are non-Gaussian under the “null” assumption $\mu=\nu$ [22, 21]. When measures are supported on $\mathbb{R}^{d}$, $d>1$ and are absolutely continuous, the curse of dimensionality takes place: the empirical Wasserstein distance converges in expectation to to the true one too slowly [26, 30]. It is still possible, however, to obtain convergence statement similar to (6) in any dimension $d\geq 1$ by replacing the centering true value $OT_{c}(\mu,\nu)$ with expectation of the empirical value $\mathbb{E}\left(OT_{c}(\widehat{\mu}_{n_{1}},\widehat{\nu}_{n_{2}}\right)$ [23]. The limit $L$ is Gaussian when $\mu\neq\nu$ and is degenerate (i.e. limiting random variable has zero variance) when $\mu=\nu$.

Favorable situation arises when measures are supported on a finite space $\mathcal{X}=\{x_{1},\ldots,x_{N}\}\subseteq\mathbb{R}^{d}$: [72] show that the limit law of the form (6) hold for the $W_{p}$ distance in any dimension $d\geq 1$ and use resulting laws to construct statistical inference under $H_{0}$ and $H_{a}$ (the case of countable support is treated in [75]). The laws under either $H_{0}$ or $H_{a}$ are non-degenerate and given by

where $G_{1},G_{2}$ are the weak limits of multinomial processes $\sqrt{n_{1}}(\widehat{\mu}_{n_{1}}-\mu)$ and $\sqrt{n_{2}}(\widehat{\nu}_{n_{2}}-\nu)$, and $\Phi^{*}$ is set of optimal solutions to the dual of the $W_{p}^{p}$ program (5). The results are extended in [46] to general measures supported on $\mathbb{R}^{d}$, $d=1,2,3$ and general costs $c$ (with discussion of limitations in higher dimensions $d\geq 4$), thus providing a unified approach to weak limits of empirical OT costs centered by the true population value.

The starting point for theoretical results of this paper is the weak limit (7) [72]. Inspired by this result, we establish the limits of the form (6), where $OT_{c}$ is now the optimal value of the $MOT$ program (2) with $k\geq 2$ measures supported on a finite space $\mathcal{X}=\{x_{1},\ldots,x_{N}\}\subseteq\mathbb{R}^{d}$ for any $d\geq 1$. The implications of these results to $k$-sample inference and further theoretical findings related to our limits are summarized below.

Asymptotic distributions of $MOT$: We provide asymptotic distribution of $MOT$ cost on finite spaces by establishing Hadamard directional differentiability of the $MOT$ functional and combining it with functional Delta method [9, 29, 46, 70, 69, 65, 72, 75]. The resulting limit is a Hadamard directional derivative of $MOT$ at the true $\mu$ in the direction of the limit $G$ of the empirical process $\rho_{n}\left(\widehat{\mu}_{n}-\mu\right)$ for suitably defined rate $\rho_{n}$ (Theorem 2.2(a)). For $k=2$ measures, our limit recovers the one for the Wasserstein distance on finite spaces obtained in [72]; for $k>2$ measures, our limit allows to construct novel inference procedures for the $k$-sample problem using $MOT$ (Section 2.2). We specify the structure of the limit under the assumption of $H_{0}$ and construct a low complexity stochastic upper bound on the null distribution (Theorem 2.2(b)) that is used to efficiently approximate the limit under $H_{0}$ (Section 3.2). The bound is tight for $k=2$. We specify the structure of the limit under $H_{a}$ and provide sufficient conditions for the limit to be Gaussian by leveraging the results from geometry of multitransportation polytopes from [28]. When the limits are not Gaussian, we construct the Normal lower bounds on the alternative limiting distribution (Theorem 2.2(c)). Our stochastic bounds on the null and alternative distributions provide an analytically tractable way to analyze the power of the Optimal Transport based tests (Section 2.4), which, to the best of our knowledge, was not yet considered in the literature.

Consistency and power: We provide a novel power analysis for Optimal Transport based tests of hypotheses (1) that encompasses both our test (12) and the two-sample test in [72] and can potentially be applied to tests based on limiting laws in [75] and [46]. We show consistency of the test under fixed alternatives (Proposition 2.8), as well as uniform consistency in a certain broad class of alternatives (Theorem 2.9 and Proposition 2.11). We illustrate theoretical power results in $k=2$ case by providing a lower bound on the power function that explicitly relates sample size and the effect size (Corollary 2.10, Figure 2). We also quantify how the population version of our statistic changes with number of measures for certain sequences of alternatives (Lemmas 2.13 and 2.14), suggesting potential power advantages in these cases. For the case of small sample sizes, we provide a permutation version of the $MOT$ based test (Section 3.3). Comparison with with state-of-the-art tests of [47, 50, 64] shows strong finite sample power performance of our tests (Figure 5).

Computational complexity results: Leveraging recent complexity result for $MOT$/barycenter program [2], we prove polynomial time complexity of the derivative bootstrap that consistently estimates asymptotic distribution of $MOT$ under $H_{0}$ (Lemma 3.2); polynomial complexity of m-out-of-n bootstrap and permutation procedure follow directly from [2] (Table 2). We demonstrate that the null upper bound of Theorem 2.2(b) can efficiently approximate the null distribution when the cardinality of the support $N=|\mathcal{X}|$ is large (Figure 4).

Applications to real data and software: We illustrate performance of $MOT$ based $k$-sample inference on two synthetic datasets showing strong power performance when testing $H_{0}$ and the ability to produce meaningful and interpretable confidence regions under $H_{a}$ (Section 4.1). Further, we apply our methodology to real data on cancers in the United States populations to confirm two claims in cancer studies that were previously shown using different methodologies (Section 4.2, Figures 7 and 8). Current version of the software that implements our methods is available at https://github.com/kravtsova2/mot.

Denote the vector of $k$ true measures supported on $\mathcal{X}=\{x_{1}\ldots,x_{N}\}\subseteq\mathbb{R}^{d}$ by

and the vector of the empirical counterparts $\widehat{\mu}_{i}:=\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}\delta_{x_{j}}$ by

with sample sizes

where $n\to\infty$ to be interpreted as each sample size tending to infinity.

Let $MOT(\mu)$ be the optimal value of the program (2), which on the finite space $\mathcal{X}$ becomes the finite-dimensional linear program

The optimization variable $\pi\in\mathbb{R}^{N^{k}}$ is a column vector representing joint probability distribution with marginals $\mu_{1},\ldots,\mu_{k}$ (frequently called multicoupling), a matrix $A\in\mathbb{R}^{kN\times N^{k}}$ encodes the constraints for $\pi$ to be a multicoupling (i.e. that summing certain entries of $\pi$ gives the marginals $\mu_{1},\ldots,\mu_{k}$), and a cost column vector $c\in\mathbb{R}^{N^{k}}$ contains Euclidean distances between measure support points to their averages given by (3)¹¹1The linear program formulation of MOT problem is discussed, for example, on p.3 of [54]..

For reader’s convenience, Example 1 below illustrates the structure of the linear program (8) in a case of three measures, each supported on two points in $\mathbb{R}^{1}$:

The dual program of (8) is given by

(the derivation of the dual follows from the standard theory of linear programming, e.g. Section 4.1 of [5]). A column vector $u:=\begin{pmatrix}u_{1}\\ \vdots\\ u_{k}\end{pmatrix}\in\mathbb{R}^{kN}$ contains dual variables, one for each measure, and the objective of (10) can be thought of summing the contributions $\langle u_{1},\mu_{1}\rangle+\cdots+\langle u_{k},\mu_{k}\rangle$.

Let $\Phi^{*}$ denote the set of dual optimal solutions to (10). This set consists of all vectors $u$ that result in the maximum value of the dual objective (= minimum value of the primal objective by strong duality, e.g. Theorem 4.4 of [5]) and satisfy the dual constraints, i.e.

We consider the asymptotic behavior of scaled and centered empirical estimator $MOT(\widehat{\mu}_{n})$ by establishing the weak limit

as $n\to\infty$ where $X\overset{d}{=}X_{0}$ under $H_{0}$ and $X\overset{d}{=}X_{a}$ under $H_{a}$. The set $\Phi^{*}$ will be needed to define the limit $X$.

Consider the statistic

where $MOT(\mu)=0$ under $H_{0}$ and $MOT(\mu)>0$ under $H_{a}$.

An $\alpha$-level test of $H_{0}$ would reject $H_{0}$ if $\rho_{n}MOT(\widehat{\mu}_{n})$ is large, i.e. if $T_{n}$ exceeds a $(1-\alpha)$th quantile of its null distribution $\mathcal{D}_{0}$. However, as Theorem 2.2 shows, $\mathcal{D}_{0}$ depends on the unknown true $\mu$, and hence care must be taken to ensure that the estimated cut-off used for the test still results in the (asymptotic) level $\alpha$.

To this end, we consider a consistent bootstrap estimator of $\mathcal{D}_{0}$ given in Proposition 3.1 and denote its $(1-\alpha)$-th quantile by $c_{\alpha,\mathcal{D}_{0}}$. Consistency of the bootstrap is shown using results of [29], and by Corollary 3.2 of the same work such bootstrap based cut-off gives an asymptotic level $\alpha$ test of $H_{0}$. Using this cut-off, we define the asymptotic test of $H_{0}$ as a map

Similarly, the distribution of $T_{n}$ under $H_{a}$ is consistency estimated by bootstrap in Proposition 3.1, with resulting $(\alpha/2)$-th and $(1-\alpha/2)$-th quantiles denoted by $c_{\alpha/2,\mathcal{D}_{a}}$ and $c_{1-\alpha/2,\mathcal{D}_{a}}$, respectively. The asymptotic $(1-\alpha)\%$ Confidence Region for $MOT(\mu)$ under $H_{a}$ is given by

Theorem 2.2(a) provides the general form of the asymptotic distribution of $MOT(\widehat{\mu}_{n})$ on finite spaces. This distribution is given by the Hadamard directional derivative of the $MOT$ functional, which is an optimal value of a linear program with a feasible set consisting of dual optimal solutions $\Phi^{*}$ (11). If $\Phi^{*}$ is a singleton, the limit in Theorem 2.2(a) is a linear combination of Gaussians, and hence is also Gaussian. If not, the limit is the maximum (taken over the feasible set $\Phi^{*}$) of such linear combinations.

By the theory of linear programming, it is possible to assess whether the set of dual optimal solutions $\Phi^{*}$ is a singleton or not based on the corresponding set of basic optimal solutions to the primal program (we use Theorem 5.6.1 of [71], and Chapters 4 and 5 of [5] for the general linear programming results employed below). In our case, the basic optimal solutions to the primal program (8) are the vertices of a multitransportation polytope $P(\mu_{1},\ldots,\mu_{k}):=\{\pi\geq 0:A\pi=\mu\}$ given by multicouplings $\pi^{*}$. These vertices contain $\leq\text{rank}(A)=kN-k+1$ positive entries, and a vertex is termed degenerate if it contains strictly less (p. 366 of [28]).

The dual optimal set $\Phi^{*}$ cannot be a singleton if an optimal solution to the primal $MOT$ program is unique and degenerate. This is always the case under $H_{0}$: the unique optimal solution $\pi^{*}$ is given by the “identity” multicoupling (with $\mu_{1}$ in the entries with the same tuple indices and zeros otherwise - so it is degenerate). Hence, the asymptotic distribution of $MOT$ under $H_{0}$ is never a Gaussian (Theorem 2.2(b)).

The dual optimal set $\Phi^{*}$ is a singleton if there exists a non-degenerate primal optimal vertex, i.e. an optimal multicoupling $\pi^{*}$ with $kN-k+1$ positive entries. This can be possible under certain $H_{a}$’s; in particular, this is possible if a multitransportation polytope $P(\mu_{1},\ldots,\mu_{k}):=\{\pi\geq 0:A\pi=\mu\}$ contains no degenerate vertices. In this case, $\Phi^{*}$ is a singleton, and the corresponding asymptotic distribution of $MOT$ is Gaussian (Theorem 2.2(c)). We use a result in discrete geometry from [28] to provide a sufficient condition (A1) that leads to Gaussian limits under $H_{a}$:

Using the above results, Proposition 2.6 defines an upper bound on all test cut-offs $c_{\alpha,\mathcal{D}_{0}}$, which is independent of the nature and the number of measures in $\mu(k)=(\mu_{1},\ldots,\mu_{k})$. This bound is used to prove consistency of the test (12) with cut-offs $c_{\alpha,\mathcal{D}_{0}}$ uniformly over $\mu$ (Theorem 2.9) and $k$ (Proposition 2.11).