A linear time approximation of W asserstein distance with word embedding selection (2023)

Abstract: Wasserstein distance, which can be computed by solving the optimal transport problem, is a powerful method for measuring the dissimilarity between documents. In the NLP community, it is referred to as word mover’s distance (WMD). One of the key challenges of Wasserstein distance is its computational cost since it needs cubic time. Although the Sinkhorn algorithm is a powerful tool to speed up to compute the Wasserstein distance, it still requires square time. Recently, a linear time approximation of the Wasserstein distance including the sliced Wasserstein and the tree-Wasserstein distance (TWD) has been proposed. However, a linear time approximation method suffers when the dimensionality of word vectors is high. In this study, we propose a method to combine feature selection and tree approximation of Wasserstein distance to handle high-dimensional problems. More specifically, we use multiple word embeddings and automatically select useful word embeddings in a tree approximation of Wasserstein distance. To this end, we approximate Wasserstein distance for each word vector by tree approximation technique, and select the discriminative (i.e., large Wasserstein distance) word embeddings by solving an entropic regularized maximization problem. Through our experiments on document classification, our proposed method achieved high performance.

## introduction wasserstein distance, which is obtained by solving the optimal transport problem, is a powerful method for measuring the distance between distributions and is used in natural language processing to measure dissimilarity between documents @xcite @xcite @xcite . @xcite proposed word mover's distance (wmd), which uses word vectors to compute the ground metric of the wasserstein distance and can measure dissimilarity between documents using wasserstein distance. more specifically, given a word embedding and the frequency of the words in each document, wmd is formulated as an optimal transport problem in the word embedding space. although the optimal transport problem can be solved using a linear programming algorithm, it requires cubic computational time with respect to the number of samples. @xcite proposed to introduce the entropy regularization to the optimal transport problem, which can be solved with a square computational time with respect to the number of samples by the sinkhorn algorithm. however, the sinkhorn algorithm is too slow for practical application in tasks such as natural language processing, that involve a large number of words. the tree-wasserstein distance (twd) is a method to measure the distance between distributions using a tree structure, which can approximate wasserstein distance with a linear computational time with respect to the number of nodes in the tree @xcite @xcite @xcite @xcite . as a method for constructing tree structures, @xcite proposed quadtree, which constructs a tree by recursively partitioning the original space. in addition, @xcite used hierarchical clustering to construct a tree. existing approaches successfully approximate the original optimal transport problems. however, when the dimensionality of the vectors (e.g., word vectors) is high, their performance can be degraded due to the curse of dimensionality. to address high-dimensional problems, robust variants of optimal transport for high-dimensional data have recently been proposed @xcite . @xcite proposed the subspace robust wasserstein distances @xcite , which projects data into a subspace and considers optimal transport in the subspace. @xcite proposed feature-robust optimal transport (frot), which measures distances between distributions with group feature selection when feature groups are given as prior information. since srw and frot alternately solve the feature selection optimization problems and the optimal transport problem with either linear programming or the sinkhorn algorithm, they do not scale to large datasets. therefore, applying these techniques to nlp tasks such as document classification that involve large amounts of data is difficult. in this study, we propose the tree-wasserstein distance with group feature selection (twd-gfs) to combine feature selection and twd to handle high-dimensional problems and scale to large datasets. we consider a setting in which some distributions are given as training data to select discriminative feature groups in advance. twd-gfs first constructs a tree structure for each of the given feature groups (i.e., word embeddings), and its output is the weighted sum of the twd calculated using each tree. the weight for each feature group can be obtained by solving an entropic regularized maximization problem using the distribution of the training data. after training, twd-gfs can measure distances with selected features much faster than srw and frot. through the synthetic experiments, we show that our proposed method can select discriminative features in a high-dimensional setting with a large number of noise features. then, through the real experiment, we consider multiple word embeddings as feature groups and confirmed the performance of twd-gfs on word embeddingbased document classification tasks with word embedding selection. contributions: our contributions are as follows: notation: in the following sections, we write 1 n for a n-dimensional vector with all ones, and σ n for a n-dimensional probability simplex. ## related works in this section, we introduce wasserstein distance, tree-wasserstein distance, and optimal transport with feature selection. ## proposed method in this section, we describe a problem setting considered in this research, our proposed method and a comparison between our approach and existing techniques. ## synthetic experiments in this section, we show that our proposed method can select the correct feature groups in highdimensional settings with a lot of noise features. in this experiment, we first created 10 distributions in two dimensions as shown in figure 1a . each distribution contained 20 points and followed a separate multidimensional normal distribution, the covariance matrix of which was an identity ma-trix and the mean of which was randomly sampled from a uniform distribution u (-20, 20) in each dimension. we can consider 10 c 2 = 45 pairs of distributions. next, we added some noise features sampled from u (-20, 20) equal to the data range. for example, when the number of the noise features is 10, the number of all features is 2+10 = 12. we calculated two lists, including a list of distances of 45 distribution pairs computed with 1-wasserstein distance for two-dimensional data, and a list of distances of 45 distribution pairs computed with twd for the data containing noise features. in figure 1b , we show the pearson correlation coefficient (pcc) between these lists when noise features were added. figure 1b shows that pcc worsened as the number of noise features increased for both quadtree and clustertree, so twd cannot approximate 1-wasserstein distance on an original dimension well in the noisy settings due to the curse of dimensionality. selecting useful features is important in the highdimensional setting with a large number of noise features as described above. we show that our proposed method can correctly select the original two-dimensional features among all features. in the following experiment, the number of noise features was 20 and the total number of all features was 22, and we let x 0 , x 1 be the original features and x 2 , ..., x 21 be the noise features. suppose that the feature groups as prior information are given as [(0, 1), (2, 3), ..., (20, 21)], where two features are given as a single group and the number of groups is 11. we learned weights using algorithm 1 and checked the values of w * . as input to algorithm 1, the data x were all 22-dimensional points belong-ing to 10 distributions, the index-pair set as training data was ω = {(0, 1), (0, 2), ..., (8, 9)}(|ω| = 45), which was all pairs of distributions, and the hyperparameter η was set to 0.1 in both of quadtree and clustertree. we chose the small hyperparameter because we know the true feature group in the synthetic data, and it is reasonable to expect better results with the small regularization parameter, which makes the learned weights closer to a onehot vector. as a result, for both methods, the value of the weight corresponding to the feature group (0, 1) became more than 0.99, and other weight values were less than 0.01. these results indicate that our proposed method correctly selected the original two-dimensional features using the training data. ## real experiments in this section, we evaluate the performance of our proposed method through word embeddingbased document classification experiments on real datasets. we used five open datasets foot_0 , including bbcsport, twitter, amazon, classic, and recipe. each dataset has five splits of training and test data. the detail of the datasets is given in appendix b. we used the k-nearest neighbor (k-nn) algorithm and evaluated its performance in terms of the average of the test error rate computed for each split. moreover, we discuss the computational time of twd-gfs and other methods with optimal transport in appendix e. ## conclusion in this study, we proposed twd-gfs to combine twd and group feature selection to handle high-dimensional problems in the tree approximation of the wasserstein distance. unlike existing feature-robust methods of optimal transport with high computational costs, twd-gfs can scale to large datasets owing to the rapid computation of twd. twd-gfs can select discriminative feature groups by solving the entropic regularized maximization problem. through synthetic experiments, we showed that our proposed method can select discriminative features in a high-dimensional setting with a large number of noise features. through real experiments, we considered multiple word embeddings as a feature group and applied twd-gfs to word embedding selection. we confirmed that twd-gfs with clustering-based trees outperformed the methods using concatenating word embeddings and achieved high performance across all datasets.