Partial Wasserstein Covering

As discussed in Sect. 1, our goal is to fill in the gap between the application dataset $\mathcal{D}_{\text{app}}$ by adding some data from a candidate dataset $\mathcal{D}_{\text{cand}}$ to a small dataset $\mathcal{D}_{\text{dev}}$. We consider $\mathcal{D}_{\text{cand}}=\mathcal{D}_{\text{app}}$ (i.e., we copy some data in $\mathcal{D}_{\text{app}}$ and add them into $D_{\text{dev}}$) in the above-mentioned scenarios, but we herein consider the most general formulation. We model this task as a discrete optimization problem called the partial Wasserstein covering problem.

Given a dataset for development $\mathcal{D}_{\text{dev}}=\{\mathbf{y}^{(j)}\}_{j=1}^{N_{\text{dev}}}$, a dataset obtained from an application $\mathcal{D}_{\text{app}}=\{\mathbf{x}^{(i)}\}_{i=1}^{N_{\text{app}}}$, and a dataset containing candidates for selection $\mathcal{D}_{\text{cand}}=\{\mathbf{s}^{(j)}\}_{j=1}^{N_{\text{cand}}}$, where $N_{\text{app}}\geq N_{\text{dev}}$, the PWC problem is defined as the following optimization:³³3We herein consider the partial Wasserstein divergence between the unlabled datasets, because the application and candidate datasets are usually unlabeled. If they are labeled, we can use the labels information as in [2].

We select a subset $S$ from the candidate dataset $\mathcal{D}_{\text{cand}}$ under the budget constraint $|S|\leq K~{}(\ll N_{\text{cand}})$, and then add that subset to the development data $\mathcal{D}_{\text{dev}}$ to minimize the partial Wasserstein divergence between the two datasets $\mathcal{D}_{\text{app}}$ and $S\cup\mathcal{D}_{\text{dev}}$. The second term is a constant with respect to $S$, which is included to make the objective non-negative.

In Eq. (2), we must specify the probability mass (i.e., $\mathbf{a}$ and $\mathbf{b}$ in Eq. (1)) for each data point. Here, we employ a uniform mass, which is a natural choice because we do not have prior information regarding each data point. Specifically, we set the weights to $a_{i}=1/N_{\text{app}}$ for $\mathcal{D}_{\text{app}}$ and $b_{j}=1/N_{\text{dev}}$ for $S\cup\mathcal{D}_{\text{dev}}$. With this choice of masses, we can easily show that $\mathcal{PW}^{2}(\mathcal{D}_{\text{app}},S\cup\mathcal{D}_{\text{dev}})=\mathcal{W}^{2}(\mathcal{D}_{\text{app}},\mathcal{D}_{\text{dev}})$ when $S=\emptyset$, where $\mathcal{W}^{2}(\mathcal{D}_{\text{app}},\mathcal{D}_{\text{dev}})$ is the Wasserstein distance. Therefore, based on the monotone property (Sect. 3.2), the objective value is non-negative for any $S$.

The PWC problem in Eq.(2) can be written as a mixed integer linear program (MILP) as follows. Instead of using the divergence between $\mathcal{D}_{\text{app}}$ and $S\cup\mathcal{D}_{\text{dev}}$, we consider the divergence between $\mathcal{D}_{\text{app}}$ and $\mathcal{D}_{\text{cand}}\cup\mathcal{D}_{\text{dev}}$. Hence, the mass $\mathbf{b}$ is an $(N_{\text{cand}}+N_{\text{dev}})$-dimensional vector, where the first $N_{\text{cand}}$ elements correspond to the data in $\mathcal{D}_{\text{cand}}$ and the remaining elements correspond to the data in $\mathcal{D}_{\text{dev}}$. In this case, we never transport to data points in $\mathcal{D}_{\text{cand}}\setminus S$, meaning we use the following mass that depends on $S$:

As a result, the problem is an MILP problem with an objective function $\langle\mathbf{P},\mathbf{C}\rangle$ w.r.t. $S\subseteq\mathcal{D}_{\text{cand}}$ and $\mathbf{P}\in U(\mathbf{a},\mathbf{b}(S))$ with $|S|\leq K$ and $a_{i}=1/N_{\text{app}}$.

One may wonder why we use the partial Wasserstein divergence in Eq.(2) instead of the standard Wasserstein distance. This is because the asymmetry in the conditions of the partial Wasserstein divergence enables us to extract only the parts with a density lower than that of the application dataset, whereas the standard Wasserstein distance becomes large for parts with sufficient density in the development dataset. Furthermore, the guaranteed approximation algorithms that we describe later can be utilized only for the partial Wasserstein divergence version.

In this section, we prove the following theorem to guarantee the approximation ratio of the proposed algorithms.

To prove Theorem 1, we reduce our problem to the partial maximum weight matching problem [12], which is known to be submodular. First, we present the following lemmas.

Proof of Lemma 1. It is well known [10] that any extreme point of the polytope is obtained by (i) choosing $n$ inequalities from $\mathbf{A}\mathbf{x}\leq\mathbf{b}$ and (ii) solving the corresponding $n$-variable linear system $\mathbf{B}\mathbf{x}=\tilde{\mathbf{b}}$, where $\mathbf{B}$ and $\tilde{\mathbf{b}}$ are the corresponding sub-matrix/vector (when replacing $\leq$ with $=$). As all elements in $\mathbf{B}^{-1}$ and $\tilde{\mathbf{b}}$ are rational, ⁴⁴4This immediately follows from the fact that $\mathbf{B}^{-1}=\tilde{\mathbf{B}}/\det{\mathbf{B}}$, where $\tilde{\mathbf{B}}$ is the adjugate matrix of $\mathbf{B}$ and $\det\mathbf{B}$ is the determinant of $\mathbf{B}$. Since $\mathbf{B}\in\mathbb{Q}^{n\times n}$, $\tilde{\mathbf{B}}$ and $\det\mathbf{B}$ are both rational. we conclude that any extreme points are rational. ∎

Proof of Lemma 2. For any fixed $S\subseteq Z$, there is an optimal solution $\mathbf{P}^{*}$ of $\psi(S)$, where all elements $P_{ij}^{*}$ are rational because all the coefficients in the constraints are in $\mathbb{Q}$ ($\because$ Lemma 1). Thus, there is a natural number $M\in\mathbb{Z}_{+}$ that satisfies $\psi(S)=\frac{1}{mlM}\max_{\mathbf{P}\in U_{\mathbb{Q}}(S)}\left<\mathbf{R},\mathbf{P}\right>$ where the set $U_{\mathbb{Q}}(S)$ is the set of possible transports when the transport mass is restricted to the rational numbers $\mathbb{Q}$.

The problem of $\psi(S)$ can be written as a network flow problem on a bipartite network following Sect 3.4 in [9]. Since the elements of $\mathbf{P}$ are integers, we can reformulate the problem of $\psi(S)$ using maximum-weight bipartite matching (i.e., an assignment problem) on a bipartite graph whose vertices are generated by copying the vertices $lM$ and $mM$ times with modified edge weights. Therefore, $G(V\cup W,E)$ is a bipartite graph, where $V=\{v_{1,1},\dots,v_{1,lM},\dots,v_{m,1},\dots,v_{m,lM}\}$ and $W=\{w_{1,1},\dots,w_{1,mM},\dots,w_{n,1},\dots,w_{n,mM}\}$. For any $i\in[\![m]\!],j\in[\![n]\!]$, the transport mass $P_{ij}\in\mathbb{Z}_{+}$ is encoded by using the copied vertices $\{v_{i,1},\dots,v_{i,lM}\}$ and $\{w_{j,1},\dots,w_{j,mM}\}$, and the weights of the edges among them are $\frac{1}{mlM}R_{ij}>0$ to represent $R_{ij}$. The partial maximum weight matching function for a graph with positive edge weights is a submodular function [12]. Therefore, $\psi(S)$ is a submodular function. ∎

Proof of Lemma 3. Given an optimal solution $\mathbf{P}^{*}$, suppose there exists $i^{\prime}$ satisfying $\sum_{j=1}^{n}P^{*}_{i^{\prime}j}<\frac{1}{m}$. This implies that $\sum_{i,j}P^{*}_{ij}<1$. We denote the gap as $\delta\coloneqq\frac{1}{m}-\sum_{j=1}^{n}P_{i^{\prime}j}^{*}>0$. Based on the assumption of $|S|\geq l$, we have $\sum_{j=1}^{n}b_{j}(S)\geq 1$. Hence, there must exist $j^{\prime}$ such that $\sum_{i=1}^{m}P_{ij^{\prime}}^{*}<b_{j^{\prime}}$. We denote the gap as $\delta^{\prime}\coloneqq b_{j^{\prime}}-\sum_{i=1}^{m}P_{ij^{\prime}}^{*}>0$. We also define $\tilde{\delta}=\min\{\delta,\delta^{\prime}\}>0$. Then, a matrix $\mathbf{P}^{\prime}=(P_{ij}^{*}+\tilde{\delta}I[i=i^{\prime},j=j^{\prime}])_{ij}$ is still a feasible solution. However, $R_{ij}>0$ leads to $\left<\mathbf{R},\mathbf{P}^{\prime}\right>>\left<\mathbf{R},\mathbf{P}^{*}\right>$, which contradicts $\mathbf{P}^{*}$ being the optimal solution. Thus, $\sum_{j=1}^{n}P_{ij}^{*}=\frac{1}{m},\forall i$. ∎

Proof of Theorem 1. Given $C_{\text{max}}=\max_{i,j}C_{ij}+\gamma$, where $\gamma>0$, the following is satisfied:

where $\mathbf{R}=(C_{\text{max}}-C_{ij})_{ij}>0$, $m=N_{x}$, $n=N_{s}+N_{y}$, and $l=N_{y}$. Here, $|S\cup Y|\geq N_{y}=l$ and Lemma 3 yield $\phi(S)=\psi(S\cup Y)-\psi(Y)$. Since $\phi(S)$ is a contraction of the submodular function $\psi(S)$ (Lemma 2), $\phi(S)=-\mathcal{PW}^{2}(X,S\cup Y)+\mathcal{PW}^{2}(X,Y)$ is a submodular function.

We now prove the monotonicity of $\phi$. For $S\subseteq T$, we have $U(\mathbf{a},\mathbf{b}(S))\subseteq U(\mathbf{a},\mathbf{b}(T))$ because $\mathbf{b}(S)\leq\mathbf{b}(T)$. This implies that $\phi(S)\leq\phi(T)$. ∎

Finally, we prove Proposition 1 to clarify the computational intractability of the PWC problem.

Proof of Proposition 1. The decision version of the PWC problem, denoted as $\mathtt{PWC}(\mathbf{a},\mathbf{b}(S),C,K,x)$, asks whether a subset $S,|S|\leq K$ such that $\phi(S)\leq x\in\mathbb{R}$ exists. We discuss a reduction to the PWC from the well-known NP-complete problems called set cover problems [13]. Given a family $\mathcal{G}=\{G_{j}\subseteq\mathcal{I}\}_{j=1}^{M}$ from the set $\mathcal{I}=[\![N]\!]$ and an integer $k$, $\mathtt{SetCover}(\mathcal{G},\mathcal{I},k)$ is a decision problem for answering $\mathtt{Yes}$ iff there exists a sub-family $\mathcal{G^{\prime}}\subseteq\mathcal{G}$ such that $|\mathcal{G^{\prime}}|\leq k$ and $\cup_{G\in\mathcal{G^{\prime}}}G=\mathcal{I}$. Our reduction is defined as follows. Given $\mathtt{SetCover}(\mathcal{G},\mathcal{I},k)$ for $i\in[\![N]\!]$ and $j\in[\![M]\!]$, we set $a_{i}=1,b_{j}(S)=|G_{j}|\cdot I[j\in S],C_{ij}=I[i\in G_{j}],K=k$, and $x=0$. Based on this polynomial reduction, the given $\mathtt{SetCover}$ instance is Yes iff the reduced $\mathtt{PWC}$ instance is $\mathtt{Yes}$. This reduction means that the decision version of PWC is NP-complete and now Proposition 1 holds.

First, we compare our algorithms, namely, the greedy-based PW-greedy-*, sensitivity-based PW-sensitivity-*, and random selection for the baseline. For the evaluation, we generated two-dimensional random values following a normal distribution for $\mathcal{D}_{\text{app}}$ and $\mathcal{D}_{\text{dev}}$, respectively. Figure 2 (left) presents the partial Wasserstein divergences between $\mathcal{D}_{\text{app}}$ and $S\cup\mathcal{D}_{\text{dev}}$ while varying the number of the selected data when $N_{\text{app}}=N_{\text{dev}}=30$ and Fig. 2 (right) presents the computational times when $K=30$ and $N_{\text{app}}=N_{\text{dev}}=N_{\text{cand}}$ varied. As shown in Fig. 2 (left), PW-greedy-LP, and PW-sensitivity-LP select exactly the same data as the global optimal solution obtained by PW-MILP. As shown in Fig. 2 (right), considering the logarithmic scale, PW-sensitivity-* significantly reduces the computational time compared with PW-MILP, while the naïve greedy algorithms do not scale. In particular, PW-sensitivity-ent can be calculated quickly as long as the GPU memory is sufficient, whereas its CPU version (PW-sensitivity-ent(CPU)) is significantly slower. PW-sensitivity-Ctrans is the fastest among the methods without GPUs. It is 8.67 and 3.07 times faster than PW-MILP and PW-sensitivity-LP, respectively. For quantitative evaluation, we define the empirical approximation ratio as $\phi(S_{K})/\phi(S^{*}_{K})$, where $S^{*}_{K}$ is the optimal solution obtained by PW-MILP when $|S|=K$. Figure 3 shows the empirical approximation ratio ($N_{\text{app}}=N_{\text{dev}}=30,K=15$) for 50 different random seeds. Both PW-*-LP achieve approximation ratios close to one, whereas PW-*-ent and PW-sensitivity-Ctrans have slightly lower approximation ratios⁵⁵5An important feature of PW-*-ent is that they can be executed without using any commercial solvers..

For the quantitative evaluation of the missing pattern findings, we herein consider a scenario in which a category (i.e., label) is less included in the development dataset than in the application dataset⁶⁶6Note that in real scenarios, missing patterns do not always correspond to labels.. We employ subsets of the MNIST dataset [40], where the development dataset contains $0$ labels at a rate of 0.5%, whereas all labels are included in the application data in equal ratios (i.e., 10%). We randomly sampled 500 images from the validation split for each $\mathcal{D}_{\text{app}}$ and $\mathcal{D}_{\text{dev}}$. We employ the squared $L^{2}$ distance on the pixel space as the ground metric for our method. For baseline methods, we employed LOF [4] novelty detection provided by scikit-learn [41] with default parameters. For the LOF, we selected top-$K$ data according to the anomaly scores. We also employed active learning approaches based on entropies of the prediction [35, 36] and coreset with the $k$-center and $k$-center greedy [5, 42]. For the entropy-based method, we train a LeNet-like neural network using the training split, and then select $K$ data with high entropies of the predictions from the application dataset. LOF and coreset algorithms are conducted on the pixel space.

We conducted 10 trials using different random seeds for the proposed algorithms and baselines, except for the PW-greedy-* algorithms, because they took over 24h. Figure 4 presents a histogram of the selected labels when $K=30$, where $x$-axis corresponds to labels of selected data (i.e., $0$ or not $0$) and $y$-axis shows relative frequencies of selected data. The proposed methods (i.e., PW-*) extract more data corresponding to the label $0$ (0.71 for PW-MILP) than the baselines. We can conclude that the proposed methods successfully selected data from the candidate dataset $\mathcal{D}_{\text{cand}}(=\mathcal{D}_{\text{app}})$ to fill in the missing areas in the distribution.

Finally, we demonstrate PWC in a realistic scenario, using driving scene images. We adopted two datasets, BDD100K [43] and KITTI (Object Detection Evaluation 2012) [44] as the application and development datasets, respectively. The major difference between these two datasets is that KITTI ($\mathcal{D}_{\text{dev}}$) contains only daytime images, whereas BDD100k ($\mathcal{D}_{\text{app}}$) contains both daytime and nighttime images. To reduce computational time, we randomly selected 1,500 data points for the development dataset from the test split of the KITTI dataset and 3,000 data points for the application dataset from the test split of BDD100k. To compute the transport costs between images, we calculated the squared $L_{2}$ norms between the feature vectors extracted by a pretrained ResNet with 50 layers [45] obtained from Torchvision [46]. Before inputting the images into the ResNet, each image was resized to a height of 224 pixels and then center-cropped to a width of 224, followed by normalization. As baseline methods, we adopted the LOF [4] and coreset with the $k$-center greedy [5]. Figure 5 presents the obtained top-3 images. One can see that PWC (PW-sensitivity-LP) selects the major pattern (i.e., nighttime scenes) that is not included in the development data, whereas LOF and coreset mainly extracts specific rare scenes (e.g., a truck crossing a street or out-of-focus scene). The coreset selects a single image of nighttime scenes; however, we emphasize that providing multiple images is essential for the developers to understand what patterns in the image (e.g., nighttime, type of cars, or roadside) are less included in the image.

The above results indicate that the PWC problem enables us to accelerate updating ML-based systems when the distributions of application and development data are different because our method does not focus on isolated anomalies but major missing patterns in the application data, as shown in Fig. 1 and Fig. 5. The ML workflow using our method can efficiently find such patterns and allow developers to incrementally evaluate and update the ML-based systems (e.g., test the generalization performance, redevelop a new model, and designing a special subroutine for the pattern). Identifying patterns that are not taken into account during development and addressing them individually can improve the reliability and generalization ability of ML-based systems, such as image recognition in autonomous vehicles.