Exact Statistical Inference for the Wasserstein Distance by
Selective Inference

Regarding the high-level conceptual contribution, we introduce a novel approach to explicitly characterize the selection bias of the data-driven coupling problem inspired by the concept of conditional Selective Inference (SI). In regard to the technical contribution, we provide an exact (non-asymptotic) inference method for the Wasserstein distance. To our knowledge, this is the first method that can provide valid CI, called selective CI, for the Wasserstein distance that guarantees the coverage property in finite sample size. Another practically important advantage of this study is that the proposed method is valid when the Wasserstein distance is computed in multi-dimensional problem, which is impossible for almost all the existing asymptotic methods since the limit distribution of the Wasserstein distance is only applicable for univariate data. We conduct experiments on both synthetic and real-world datasets to evaluate the performance of the proposed method.

In traditional statistics, reliability evaluation with Wasserstein distance has been based on asymptotic theory, i.e., sample size $\rightarrow\infty$. In the univariate case, instead of solving the optimization problem, the Wasserstein can be described by using an inverse of the distribution function. For example, let $F^{-1}$ be the quantile function of the data and $F_{n}^{-1}$ be the empirical quantile function of the generated data, the Wasserstein distance with $\ell_{2}$ distance is computed by $\int_{0}^{1}(F^{-1}(t)-F_{n}^{-1}(t))^{2}\,dt$. Based on the quantile function, several studies (Del Barrio et al., 1999, 2019; Ramdas et al., 2017) derived the asymptotic distribution of the Wasserstein distance. Obviously, these methods can not guarantee the validity in finite sample size. Moreover, since the quantile function is only available in univariate case, these methods can not be extended to multivariate cases which are practically important.

Recently, Imaizumi et al. (2019) has proposed an approach on multidimensional problems. However, it is important to clarify that this study does not provide statistical inference for the “original” Wasserstein distance. Instead, the authors consider an approximation of the Wasserstein distance, which does not require solving a LP. Besides, this method also relies on asymptotic distribution of the test statistic which is approximated by the Gaussian multiplier bootstrap. Therefore, to our knowledge, statistical inference method for the Wasserstein distance in multi-dimensional problems is still a challenging open problem.

Conditional SI is a statistical inference framework for correcting selection bias. Traditionally, there are mainly two types of approaches for selection bias correction. The first approach is family-wise error rate (FWER) control, which includes standard multiple comparison methods such as traditional Bonferroni correction. However, the FWER control is too conservative and it is difficult to utilize it for selection bias correction in complex adaptive data analysis. Another approach is false discovery rate (FDR) control, in which the target is to control the expected proportion of discoveries that are false at a given significance level. The FDR control is less conservative than FWER control, and it is used in many high-dimensional statistical inference problems.

Conditional SI is the third approach for selection bias correction. The basic idea of conditional SI was known before, but it becomes popular by the recent seminal work proposed by Lee et al. (2016). In that study, exact statistical inference on the selected features by Lasso was considered. Their basic idea is to employ the sampling distributions of the selected parameter coefficients conditional on the selection event that a subset of features is selected by the Lasso. By using the conditional distribution in statistical inference, the selection bias of the Lasso (i.e., the fact that the Lasso selected the features based on the data) can be corrected. Their contribution was to show that, even in complex data analysis methods such as Lasso, the exact sampling distribution can be characterized if appropriate selection events are considered.

After the seminal work (Lee et al., 2016), many conditional SI approaches for various feature selection methods were proposed in the literature (Loftus and Taylor, 2015; Yang et al., 2016; Tibshirani et al., 2016; Suzumura et al., 2017). Furthermore, theoretical analyses and new computational methods for conditional SI are still being actively studied (Fithian et al., 2014; Le Duy and Takeuchi, 2021; Duy and Takeuchi, 2021; Sugiyama et al., 2021a). However, most of conditional SI studies are focused on feature selection problems. Although there have been applications to several problems such as change point detection (Hyun et al., 2018; Duy et al., 2020b; Sugiyama et al., 2021b), outlier detection (Chen and Bien, 2019; Tsukurimichi et al., 2021), and image segmentation (Tanizaki et al., 2020; Duy et al., 2020a), these problems can also be interpreted as feature selection in a broad sense. Our novelty in this study is to first introduce conditional SI framework for statistical inference on the Wasserstein distance, which is a data-dependent adaptive distance measure. Our basic idea is based on the facts that the Wasserstein distance is formulated as the solution of a linear program (LP), and the optimal solution of an LP is characterized by the selected basic variables. In this study, we consider the sampling distribution of the Wasserstein distance conditional on the selected basic variables, which can be interpreted as considering the selection event on the optimal coupling between the two distributions.

We define the cost matrix $C(\bm{X},\bm{Y})$ of pairwise distances ($\ell_{1}$ distance) between elements of $\bm{X}$ and $\bm{Y}$ as

We can vectorize $C(\bm{X},\bm{Y})$ in the form of

where ${\rm{vec}}(\cdot)$ is an operator that transforms a matrix into a vector with concatenated rows. The matrix $\Theta$ is defined as

where the operator $\circ$ is element-wise product, ${\rm{sign}}(\cdot)$ is the operator that returns an element-wise indication of the sign of a number, $\bm{1}_{m}\in\mathbb{R}^{m}$ is the vector of ones, $\bm{0}_{m}\in\mathbb{R}^{m}$ is the vector of zeros, and $I_{m}\in\mathbb{R}^{m\times m}$ is the identity matrix.

To compare two empirical measures $P_{n}$ and $Q_{m}$ with uniform weight vectors $\bm{1}_{n}/n$ and $\bm{1}_{m}/m$, we consider the following Wasserstein distance, which is defined as the solution of a linear program (LP),

Given $\bm{X}^{\rm{obs}}$ and $\bm{Y}^{\rm{obs}}$ respectively sampled from models (1) and (2) ¹¹1To make a distinction between random variables and observed variables, we use superscript ^obs, e.g., $\bm{X}$ is a random vector and $\bm{X}^{\rm{obs}}$ is the observed data vector., the Wasserstein distance in (6) on the observed data can be re-written as

where $\bm{t}={\rm{vec}}(T)\in\mathbb{R}^{nm}$, $\bm{c}(\bm{X}^{\rm{obs}},\bm{Y}^{\rm{obs}})\in\mathbb{R}^{nm}$ is defined in (4), $S=\left(M_{r}~{}M_{c}\right)^{\top}\in\mathbb{R}^{(n+m)\times nm}$ in which

that performs the sum over the rows of $T$ and

that performs the sum over the columns of $T$, and $\bm{h}=\left(\bm{1}_{n}/n~{}\bm{1}_{m}/m\right)^{\top}\in\mathbb{R}^{n+m}$ ²²2 We note that there always exists exactly one redundant equality constraint in linear equality constraint system in (7). This is due to the fact that sum of all the masses on $\bm{X}^{\rm{obs}}$ is always equal to sum of all the masses on $\bm{Y}^{\rm{obs}}$ (i.e., they are all equal to 1). Therefore, any equality constraint can be expressed as a linear combination of the others, and hence any one constraint can be dropped. In this paper, we always drop the last equality constraint (i.e., the last row of matrix $S$ and the last element of vector $\bm{h}$) before solving (7). .

Let us denote the set of basis variables (the definition of basis variable can be found in the literature of LP, e.g., Murty (1983)) obtained when applying the LP in (7) on $\bm{X}^{\rm{obs}}$ and $\bm{Y}^{\rm{obs}}$ as

We would like to note that the identification of the basic variables can be interpreted as the process of determining the optimal coupling between the elements of $\bm{X}^{\rm obs}$ and $\bm{Y}^{\rm obs}$ in the optimal transport problem for calculating the Wasserstein distance. Therefore, ${\mathcal{M}}_{\rm{obs}}$ in (8) can be interpreted as the observed optimal coupling obtained after solving LP in (7) on the observed data ³³3We suppose that the LP is non-degenerate. A careful discussion might be needed in the presence of degeneracy.. An illustration of this interpretation is shown in Figure 1. We also denote by ${\mathcal{M}}_{\rm{obs}}^{c}$ a set of non-basis variables. Then, the optimal solution of (7) can be written as

where $S_{:,{\mathcal{M}}_{\rm{obs}}}$ is a sub-matrix of $S$ made up of all rows and columns in the set ${\mathcal{M}}_{\rm{obs}}$.

In the literature of LP, the matrix $S_{:,{\mathcal{M}}_{\rm{obs}}}$ is also referred to as a basis. After obtaining $\hat{\bm{t}}$, the Wasserstein distance can be re-written as

where $\Theta$ is defined in (5), and $\Theta_{{\mathcal{M}}_{\rm{obs}},:}$ is a sub-matrix of $\Theta$ made up of rows in the set ${\mathcal{M}}_{\rm{obs}}$ and all columns.

The goal is to provide a confidence interval (CI) for the Wasserstein distance. The $W(P_{n},Q_{m})$ in (9) can be written as

where $\bm{\eta}=\Theta_{{\mathcal{M}}_{\rm{obs}},:}^{\top}\hat{\bm{t}}_{{\mathcal{M}}_{\rm{obs}}}$ is the test-statistic direction. It is important to note that $\bm{\eta}$ is a random vector because ${\mathcal{M}}_{\rm{obs}}$ is selected based on the data. For the purpose of explanation, let us suppose, for now, that the test-statistic direction $\bm{\eta}$ in (10) is fixed before observing the data; that is, non-random. Let us define

Thus, we have $\bm{\eta}^{\top}(\bm{X}~{}\bm{Y})^{\top}\sim\mathbb{N}\left(\bm{\eta}^{\top}(\bm{\mu}_{\bm{X}}~{}\bm{\mu}_{\bm{Y}})^{\top},\bm{\eta}^{\top}\tilde{\Sigma}\bm{\eta}\right)$. Given a significance level $\alpha\in[0,1]$ (e.g., 0.05) for the inference, the naive (classical) CI for

with $(1-\alpha)$-coverage is obtained by

where $\sigma^{2}=\bm{\eta}^{\top}\tilde{\Sigma}\bm{\eta}$ and $F_{w,\sigma^{2}}(\cdot)$ is the c.d.f of the normal distribution with a mean $w$ and variance $\sigma^{2}$. With the assumption that $\bm{\eta}$ in (10) is fixed in advance, the naive CI is valid in the sense that

However, in reality, because the test-statistic direction $\bm{\eta}$ is actually not fixed in advance, the naive CI in (12) is unreliable. In other words, the naive CI is invalid in the sense that the $(1-\alpha)$-coverage guarantee in (13) is no longer satisfied. This is because the construction of $\bm{\eta}$ depends on the data and selection bias exists.

In the next section, we introduce an approach to correct the selection bias which is inspired by the conditional SI framework, and propose a valid CI called selective CI ($C_{\rm{sel}}$) for the Wasserstein distance which guarantees the $(1-\alpha)$-coverage property in finite sample (i.e., non-asymptotic).

In this section, we propose an exact (non-asymptotic) selective CI for the Wasserstein distance by conditional SI. We interpret the computation of the Wasserstein distance in (7) as a two-step procedure:

• •

  Step 1: Compute the cost matrix in (3) with the $\ell_{1}$ distance and obtain the vectorized form of the cost matrix in (4).

• •

  Step 2: Solving the LP in (7) to obtain ${\mathcal{M}}_{\rm{obs}}$ which is subsequently used to construct the test-statistic direction $\bm{\eta}$ in (10) and calculate the distance.

Since the distance is computed in data-driven fashion, for constructing a valid CI, it is necessary to remove the information that has been used for the initial calculation process, i.e., steps 1 and 2 in the above procedure, to correct the selection bias. This can be achieved by considering the sampling distribution of the test-statistic $\bm{\eta}^{\top}(\bm{X}~{}\bm{Y})^{\top}$ conditional on the selection event; that is,

where ${\mathcal{S}}(\bm{X},\bm{Y})$ is the sign vector explained in the construction of $\Theta$ in (5).

Interpretation of the selection event in (14). The first and second conditions in (14) respectively represent the selection event for steps 1 and 2 in the procedure described in the beginning of this section. The first condition ${\mathcal{S}}(\bm{X},\bm{Y})={\mathcal{S}}(\bm{X}^{\rm{obs}},\bm{Y}^{\rm{obs}})$ indicates that ${\rm{sign}}(x_{i}-y_{j})={\rm{sign}}(x^{\rm{obs}}_{i}-y_{j}^{\rm{obs}})$ for $i\in[n],j\in[m]$. In other words, for $i\in[n],j\in[m]$, we condition on the event whether $x^{\rm{obs}}_{i}$ is greater than $y^{\rm{obs}}_{j}$ or not. The second condition ${\mathcal{M}}(\bm{X},\bm{Y})={\mathcal{M}}_{\rm{obs}}$ indicates that the set of selected basic variables for random vectors $\bm{X}$ and $\bm{Y}$ is the same as that for $\bm{X}^{\rm{obs}}$ and $\bm{Y}^{\rm{obs}}$. This condition can be interpreted as conditioning on the observed optimal coupling ${\mathcal{M}}_{\rm{obs}}$ between the elements of $\bm{X}^{\rm{obs}}$ and $\bm{Y}^{\rm{obs}}$, which is obtained after solving the LP in (7) on the observed data (see Figure 1).

Selective CI. To derive exact statistical inference for the Wasserstein distance, we introduce so-called selective CI for $W^{\ast}=W^{\ast}(P_{m},Q_{m})=\bm{\eta}^{\top}\left(\bm{\mu}_{\bm{X}}~{}\bm{\mu}_{\bm{Y}}\right)^{\top}$ that satisfies the following $(1-\alpha)$-coverage property conditional on the selection event:

for any $\alpha\in[0,1]$. The selective CI is defined as

where the pivotal quantity

is the c.d.f of the truncated normal distribution with a mean $w\in\mathbb{R}$, variance $\sigma^{2}=\bm{\eta}^{\top}\tilde{\Sigma}\bm{\eta}$, and truncation region ${\mathcal{Z}}$ (the detailed construction of ${\mathcal{Z}}$ will be discussed later in §3.2) which is calculated based on the selection event in (17). The $\bm{q}(\bm{X},\bm{Y})$ in the additional third condition is the sufficient statistic of nuisance parameter defined as

in which $\bm{c}=\tilde{\Sigma}\bm{\eta}(\bm{\eta}^{\top}\tilde{\Sigma}\bm{\eta})^{-1}$ with $\tilde{\Sigma}$ is defined in (11). Here, we note that the selective CI depends on $\bm{q}(\bm{X},\bm{Y})$ because the pivotal quantity in (17) depends on this component, but the sampling property in (15) is kept satisfied without this additional condition because we can marginalize over all values of $\bm{q}(\bm{X},\bm{Y})$. The $\bm{q}(\bm{X},\bm{Y})$ corresponds to the component $\bm{z}$ in the seminal paper of Lee et al. (2016) (see Section 5, Eq. 5.2 and Theorem 5.2). We note that additionally conditioning on $\bm{q}(\bm{X},\bm{Y})$ is a standard approach in the SI literature and it is used in almost all the SI-related works that we cited in this paper.

To obtain the selective CI in (16), we need to compute the quantity in (17) which depends on the truncation region ${\mathcal{Z}}$. Therefore, the remaining task is to identify ${\mathcal{Z}}$ whose characterization will be introduced in the next section.

We define the set of $(\bm{X}~{}\bm{Y})^{\top}\in\mathbb{R}^{n+m}$ that satisfies the conditions in Equation (17) as

According to the third condition $\bm{q}(\bm{X},\bm{Y})=\bm{q}(\bm{X}^{\rm{obs}},\bm{Y}^{\rm{obs}})$, the data in ${\mathcal{D}}$ is restricted to a line in $\mathbb{R}^{n+m}$ as stated in the following Lemma.

Lemma 1 indicates that we do not have to consider the $(n+m)$-dimensional data space. Instead, we only to consider the one-dimensional projected data space ${\mathcal{Z}}$ in (23), which is the truncation region that is important for computing the pivotal quantity in (17) and constructing the selective CI $C_{\rm{sel}}$ in (16).

Characterization of truncation region ${\mathcal{Z}}$. We can decompose ${\mathcal{Z}}$ into two separate sets as ${\mathcal{Z}}={\mathcal{Z}}_{1}\cap{\mathcal{Z}}_{2}$, where

We first present the construction of ${\mathcal{Z}}_{1}$ in the following Lemma.

Next, we present the construction of ${\mathcal{Z}}_{2}$. Here, ${\mathcal{Z}}_{2}$ can be interpreted as the set of values of $z$ in which we obtain the same set of the selected basic variables ${\mathcal{M}}_{\rm{obs}}$ when applying the LP in (7) on the prametrized data $\bm{a}+\bm{b}z$.

Once ${\mathcal{Z}}_{1}$ and ${\mathcal{Z}}_{2}$ are identified, we can compute the truncation region ${\mathcal{Z}}={\mathcal{Z}}_{1}\cap{\mathcal{Z}}_{2}$. Finally, we can use ${\mathcal{Z}}$ to calculate the pivotal quantity in (17) which is subsequently used to construct the proposed selective CI in (16). The details of the algorithm is presented in Algorithm 1.

In this section, we evaluate the performance of the proposed selective CI in terms of coverage guarantee, CI length and computational cost. We also show the results of comparison between our selective CI and the naive CI in (12) in terms of coverage guarantee. We would like to note that we did not conduct the comparison in terms of CI length because the naive CI could not guarantee the coverage property. In statistical viewpoint, if the CI is unreliable, i.e., invalid or does not satisfy the coverage property, the demonstration of CI length does not make sense. Besides, we additionally compare the performance of the proposed method with the latest asymptotic study (Imaizumi et al., 2019).