Large-scale similarity search with Optimal Transport
Authors: Cléa Mehnia Laouar
Yuki Takezawa
Makoto Yamada
Venue: EMNLP 2023 Main (2023)
Keywords: Optimal transport
Wasserstein distance
Document Classification

Abstract: Wasserstein distance is a powerful tool for comparing probability distributions and is widely used for document classification and retrieval tasks in NLP. In particular, it is known as the word mover's distance (WMD) in the NLP community. WMD exhibits excellent performance for various NLP tasks; however, one of its limitations is its computational cost and thus is not useful for large-scale distribution comparisons. In this study, we propose a simple and effective nearest neighbor search based on the Wasserstein distance. Specifically, we employ the L1 embedding method based on the tree-based Wasserstein approximation and subsequently used the nearest neighbor search to efficiently find the $k$-nearest neighbors. Through benchmark experiments, we demonstrate that the proposed approximation has comparable performance to the vanilla Wasserstein distance and can be computed three orders of magnitude faster than the vanilla Wasserstein distance.

Large-Scale Similarity Search with Optimal Transport
Cléa Laouar', Yuki Takezawa Makoto Yamada'
OOknaww Institute of Science and Technology. "Kyoto University C. w..aaarrriist... jp, yuki-takezawadml. ist. i. 1...ooo..aac...... ac. jp, jc, makoto. yamadaeoitt............ jp
Abstract
leading to poor retrieval performance when dealing with documents that are dissimilar in content.
Wasserstein distance is powerful tool for com- paring probability distributions and is widely used for document classification and retrieval tasks NLP. In particular, it is known as the word mover's distance (WMD) the NI.P com- munity. WMD exhibits excellent performance for various NLP tasks; however, one of lim- itations is its computational cost and thus is not useful for large-scale distribution compar- isons. In this study, we propose simple and effective nearest neighbor search based on the Wasserstein distance. Specifically, we employ the L1 embedding method based on the tree- based Wasserstein approximation and subse- quently used the nearest neighbor search to ef- ficiently find the k-nearest neighbors. Through benchmark experiments, we demonstrate that the proposed approximation has comparable performance to the vanilla Wasserstein distance and can be computed three orders of magnitude faster than the vanilla Wasserstein distance.
Word mover's distance (WMD) (Kusner et al., 2015) and its variants (Yokoi et al., 2020; Huang et al., 2016) are document dissimilarity measures that can handle overlapping problems arising the BOW model. WMD calculates the distance be- tween two documents by measuring the minimum cost of transforming the word embeddings of one document into those of another (Kusner et al., 2015; Sato et al., 2022). However, WMD and its variants have a computational complexity of the cubic order, making them slow and impractical for large-scale information retrieval tasks involving millions or billions of documents.
To decrease computational complexity, tree-based approximation of WMD known tree- wasseessnnnttin distance (TWD) can be used. QuadTree (Indyk and Thaper, 2003) and Clus- terTree (L.e et al., 2019) are techniques that enable tree construction for TWD and can be computed in linear time. Recently, Backurs et al. (2020) applied Quad'Tree to large-scale nearest neighbor searches. However, their performance is limited. Morcover, Backurs et al. (2020) proposed FlowTrree an al- gorithm that has improved performance but in- creased computation time. Additionally, existing approaches are based on pairwise comparisons and cannot directly compute one-to-many distributions. Hence, there is no optimal transport (OT) method that can be used for practical large-scale nearest neighbor search with high accuracy.
1 Introduction
Information retrieval the process extracting relevant information from large collections data. It has applications in search engines, recommen- dation systems, and content analyses. An informa- tion retrieval task can be formulated as k-nearest neighbor search problem. One major challenge in k-nearest neighbor search is the scalability of doc- ument retrieval tasks, where data can range from millions to billions documents.
The bag-of-words. (BOW) model is a widely used approach for document representation in natu- ral language processing and information retrieval. It represents document as vector of word fre- quencies, disregarding the order and context of the words. This approach is simple and efficient; however, suffers when two documents do not share common words. This is because the BOW model does not capture the semantic relationships between words and their context in the document,
In this study, we introduce an OT-based near- est neighbor search for large-scale problems. Specifically, we combine three methods: the L1- regularized technique for learning the weights of the edges in (Yamada al., 2022), the sliced variant of TWD (Le et al., 2019), and the nearest neighbor search based on FAISS (Johnson et al., 2019). This combination enable the construction a computationally efficient based nearest neigh-

---

bor search that can be scaled large dataset. Through benchmark experiments, we demonstrated that the proposed framework can compute the near- est neighbors three orders of magnitude faster than the vanilla Wasserstein distance.
2003). However, its performance is limited because of poor approximation of the original Wasserstein distance.
2.2
Word Mover's Distance Approximation with Tree
Contribution: The contribution of this paper is summarized below:
In this section, we explain how to approximate WMD using a tree. Note that WMD is a Wasser- stein distance that uses word vectors to compute the distance d(æ, x').
large-scale nearest neighbor search method based on OT is proposed.
A simple yet memory-efficti representation of I.I embedding presented.
The difference between WMD and TWD is the distance metrie. Specifically, WMD uses L2 distance d(æ, x') = a 2'112, whereas TWD uses a tree metric (Le et al., 2019; Yamada et al., 2022). Thus, if d(xr, x') = (æ x'), we can approximate WMD with TWD (i.e., WI(H, W (1,
The proposed algorithm is demonstrated to be several orders of magnitude faster than WMD.
2 Proposed Framework
this section, we propose the nearest neighbor search method based on TWD (Indyk and Thaper, 2003; Le et al., 2019).
Proposition 1 (Yamada et al., 2022) We denote B E (0, 10, = [bi, b bNugl and b, C {0, 13... The shortest path distance between leaves i and j can be represented as
2.1 Nearest Neighbor Search with Tree Wasserstein Distance
d d7 10y) w
+ TT(be
(b:
2b, o b;). (1)
The TWD between the two discrete measures μ ENi azda, and
Ejara bydmy, with a 1 =
We stated the edge-weight estimation problem using the closed-fore expression in Proposition Our approach assumed that the tree was con- structed using tree construction algorithms, such as Quad'Tree (Indyk and Thaper, 2003) and Clus- terTree (Le et al., 2019). and we held the tree struc- ture constant (fixed B). Subsequently, the weight estimation problem can be expressed (Yamada et al., 2022)
and b 1 = 1, can be expressed by L1 embedding (Takezawa et al., 2021):
Wr v) = =|uaa ublli:
where ua diag(w)Ber Ba and us diag(w) Bb. E R. is the edge weight of tree, diag(w () E R. NxN is the diagonal matrix whose diagonal cl- ements are w. and B E {0, NxNNx is a tree parameter. If the node is the ancestor node of the leaf node j. [B] = 1; othwerise, [B]s 0. In addition, N is the total number of nodes of tree and Niaa is the number of leaf nodes.
W argmin Σ WER (IJ)en
(d(æ;;);) w 21,1)² + A|| w|ll,,
where ži b by 2b o bj. and Q2 denotes a set of indices. In this study, we randomly subsampled the indices to reduce computational cost.
Using TWD, the nearest neighbor search from {px] 1==1 with μ = E:a, ag dæz can be written as
=;
K argmin (E..........)
| uaa
2.3 TWD with efficient NN
In this study, we aimed to propose large-scale similarity search; however, a simple ncarest neigh- bor search method cannot effectively scale to large datasets. Thus, we employed an off-the-shelf efficient nearest neighbor search method called FAISS' Specifically, we proposed using the L1 version of the FAISS function, TWD is defined as the LI distance. However, the dimensionality
Hence, the nearest neighbor search using TWD can be formulated as the nearest neighbor search with a LI distance. In this study, we proposed combination of the LI embedding and the nearest neighbor methods. Specifically, for efficient nearest neighbor search, we used the IndexIFlat and GPUIndexFlat functions in the FAISS package.
LI embedding for TWD is a well-known result in theoretical computer science (Indyk and Thaper,
'tt http::///tttbbocooooeecedtttters

---

@1
22
W2
wi
W
W4
W5
6
1 X2)
frz, X (omoodoing)
Jembooising
X2 (ombeddting (embedaing) (ambodoing)
X3
X4 (ambedding)
X1 (embeddingg
X2
X
X4 (embediingl
(embeddinng fembedding!
Figure I: Illustration of without leaf nodes.
of substantially high because the dimensional- ity of the LI embedding vector z depends on the number of word vectors, which is large. Thus, di- rectly using the tree embedding z = diag(w) Bo. is inefficient.
3.1 Setup
We evaluated our proposed method for document classification tasks using the k-nearest neighbors (kNN) search. Our primary objective is to assess the accuracy and computation time of the tree- based method and compare them with those of the Wasserstein distance method. To calculate the Wasserstein distance, we used the Python optimal transport package² and considered the resulting ac- curacy and computation time as baselines. The (BOW) results are given for compar-
To address this issue, proposed the use of the ClusterTree (Le et al., 2019) and disregard the leaf nodes (see Figure 1). This implies using only the internal nodes and is independent of the number of leaf nodes (i.e., the number of word vectors). For example, if we use ClusterTree with classes and depth D, the maximum number of internal nodes is K + k + +... KP. As rule of thumb, we can set K E 63,4,5,6) and D E {3, 4, 5,66.. 6}.
ison.
We specifically evaluated the ClusterTree ap- proach (Le et al., 2019) for the tree-based meth- ods. The number of clusters and the maximum depth for all experiments were set to K = 5 and D 6, respectively. Our initial step involved constructing a tree using entire word embedding (word2vec) vector X, followed by the computation of the Wasserstein distance with the tree. We set the regularization parameter
Using only the internal nodes can significantly reduce the number of dimensions However, the representation power significantly reduced. To address this issue, we proposed the use of multi- ple trees. Specifically, let B, denote ClusterTree parameter without leaf nodes; the sliced embed- ding is expressed as
10-2 and explored different numbers of slices T (1,3, 5, 10). To solve the Lasso-based regression problem asso- ciated with the weight estimation problem, the SPAMS library² was used. In addition, the weight of the leaf nodes was set to zero for the leafless tree method (1l-TWD). For your reference, we con- ducted additional evaluation using the QuadTree approach, and you can find the corresponding sults in the appendix.
diag (1 Bia diag( (w2) Bra
z
diag( wn Bga
where the weight for the t-th sliced weight WZ is trained as follows:
B. = argmin (d(xs, - T == +|||w|h,
WCRN (NJ)R
To compare the performance of all methods. used the kNN search for varying values of {1,5, 10, 15) and used GPU FAISS index with the LI metric for the tree-based methods. By changing the random seed, the datasets were split five times into training and testing subsets such that 70% of each dataset was used for training. Twitter,
and 2. ((0) + 6/0) 2(( o b]0)) E N. with N, = K + K2 +... KD. This modification can
(0)
significantly reduce the training time.
3 Experiments
this section, we assess the effectiveness of our proposed method using Twitter, BBCSport, Ama- zon, and Classic datasets.
P hnttps://pythontt.ittt withh/hthhhttt.tt........h. io/index.htme ittt:://thoth.iirialpe.. http::/hwhh....mapassssrereraaas fr/ppaople/mmiir/l/ggaass

---

Twitter
BBCSport
Amazzn
Classic
Methods
10
10
10
5
WMD
248.7
250.2
250.9
283.9
275.5
267.5
5382.8
5380.3
5380,5
5093.1
5079,4
S097 9
0.06
0
0.010
0.695
0.692
0.321
Faiss TWD
0.042
0.043
0.043
0.011
0.011
0.011
0.784
0.779
0.779
0.390
0.389
0.389
IDDWD
0.020
0.020
0.020
0.006
0.006
0.006
0.104
0.101
0.101
0.083
0.081
0.082
n. [.llec=]
Faiss Tl-WW
0.244
0.239
0.167
n. n_Slice=3
Paiss II-TW
n 1_slice=55
0.044
0.045
0.010
0.378
0.370
0.257
0.258
Faiss I1TWD
a-licc=10
0.075
0.072
0.071
0.017
0.017
0.017
0.752
0.745
0.748
0:493
0.486
0.486
Table 1: Average testing time of the datasets used in the experiments for k (1,5,10).
1.0
1.0
0.9
0.9
0.8
0.8
0.7
0.7-
0.6
0.6
0.5
0.5
Number of nearest neighbers k
Number of nearest neighbors k
I-TWD slicem1 falss (ITWO nuslic--3
(1TW slicem10
ILTWD siioe=
I/TW siice100
WMD
falss TWD nusice-3
(a) Amazon.
(h) Classic.
Figure 2: Average accuracy of the datasets used in the experiments.
BBCSport, Amazon, and Classic datasets used for the evaluation presented in Table 2.
ing the CPU FAISS index, and you can review the results in Table 3 located in the appendix. For m {1,5), Figure gives the proportion of the first nearest neighbors of a given method that are actually the m first nearest neighbors when using WMD. WMD is considered as a ground truth. As an illustration, within the Classic database, 46% of the data points from the testing set have the tree Wasserstein method with leaf nodes (faiss TWD method) sharing the same nearest neighbor as WMD.
3.2 Results
The document classification experiments aimed to evaluate the effectiveness of tree-based methods in reducing the time required to classify data. To this end, we used four datasets and performed five different train-test splits on each dataset. For each split, we constructed a tree for the TWD-based methods and measured the accuracy and time re- quired to classify the data from the testing set. The process of constructing a tree is a one-time task that enables the handling of all queries. The aver- age time required to construct the tree using each method is shown in Figure 4.
For both Figures 2 and 3, the results presented are for Amazon and Classic datasets. A compre- hensive overview of outcomes across all datasets is available in the appendix. The average time and curacy and the top-m score were calculated based on the five-time split for each dataset.
The results of the experiments are presented in Table 1, Figure 2 and Figure 3. Table shows the average time required to classify 30% of the data from each dataset for the k-nearest neighbors for k E (1, 5, 10), and Figure 2 illustrates the corre- sponding average accuracy. For information, conducted an analysis of speed performance us-
We observed that tree- based methods can reduce the classification time by factor of 10 compared to the WMD method. The fastest method is II- TWD with n sliie increasing with n n_slice. The faiss TWD method is faster than the leafless method when n_ plice be- comes large. In contrast, the accuracy improves
with the time required
chttps://githhn hwt.iuu.cooggggaaaaae Somm/gahhuagg/S-ww

---

0.7
0.7
0.6
0.6
88 micer
100 miscoe
0.4
0.4
0.3
0.3
0.2
0.2
0.1
0.1
0.0
0.0
talot TMD
hales I:TW diceu tass (TD n_Dinsscces
milee I:TV0 diseo10
reles TWD
I:TWO d.diea3
hile HITWD r.cken1 ciceee1
Fala KTWD FTWD=_slte- 1
BOW
falns I-WW3 n, site-1
Tissll/T s_3koo-5
BOW
Amazon.
(b) Classic.
Figure 3: Top-1 and Top-5 neighbors scores, with WMD as ground truth, for the datasets used in the experiments.
for the leafless-tree methods when n_slice is large, well as the top-m score. For example, ob- served that the fastest method among all the tested methods is Il-TWD with n ssiceeee 1. However, exhibits the lowest accuracy among all the methods presented in Figure 2, and the lowest top-m score Figure 3, BOW apart. As a result, a trade-off between test time and accuracy must be achieved.
Limitation
The proposed approach is based on the approxima- tion of the Wasserstein distance; the performance is generally lower than that of the vanilla Wasserstein distance. Furthermore, while we have focused on improving WMD's computational efficiency, we acknowledge that there are contexts where WMD may not be the most suitable measure. Our aim in this work has been offer more efficient ap- proach for computing WMD, especially in appli- cations that prioritize vocabulary associations over syntactic structure, like content-based recommen- dation systems, plagiarism detection, and informa- tion retrieval. Eventually, the potential scalability of our method relies on both the scalability of the tree embedding technique and the chosen nearest neighbor package. One limitation in the major- ity of NN packages is lack of support for sparse formats. Given that our tree embedding produces highly sparse vectors, addressing this limitation is crucial for handling larger datasets.
In our experiments, the 11-TWD method with o_siice- 3 balances between test time and accu- racy. Consequently, it offers suitable trade-off solution for classification tasks. Overall, our exper- iment demonstrates the effectiveness of tree-based methods in reducing the classification time while maintaining an acceptable level of accuracy.
4 Conclusion
this study, we introduced a methodological im- provement aimed at enhancing the computational speed of the Wasserstein distance approximation with trees. Specifically, we combined the L1- regularized technique to learn the weights of the edges in tree (Yamada al., 2022), the sliced- TWD (Le et al., 2019), and the efficient nearest neighbor search (Johnson et al., 2019). To further reduce the computational time, we disregarded the leaf nodes of the trees, as they can significantly reduce the dimensions of the LI embedding vec- tors. The sliced TWD allowed us to deal with the decrease in representation power resulting from the leaf loss. Through benchmark experiments, demonstrated that the proposed method can reduce computation time over 10' times compared with the WMD method while maintaining a reasonable level of accuracy.
Ethics/Broader impact
This is the first practical approach for nearest neigh- bor search using Wasserstein distance. Thus, by replacing a simple NN-based vector search with the proposed method, we can achieve significant gain over simple vector-based NN methods. Thus, notwithstanding the simplicity of the approach, it improve many NLP applications based near- neighbor search. There are no ethical issues associated with this study.
Acknowledgement
Makoto Yamada was supported by MEXT KAK-
ENHI Grant Number 20H04243.

---

References
A Related work
Arturs Backurs, Yihe Dong, Piotr Indyk, Ilya Razen- shteyn, and Tal Wagner. 2020. Scalable nearest neigh- bor search for optimal transport. In ICMI.
In this study, we considered an optimal transport- based nearest neighbor search for a large-scale dataset. Thus, we described the optimal transport (Wasserstein distance) and the nearest neighbor search (NN).
Jon Louis Bentley. 1975. Multidimensional binary search trees used for associative searching. Com- munications of the ACM, 18(9):509-517.
Optimal transport and Wasserstein distance: Optimal transport is widely used in machine learn- ing, computer vision, and natural language process- ing tasks. The Monge formulation was originally proposed (Monge, 1781), in which it moves mass from one place to another. In NLP communities, the Kantorovich formulation is more widely used than the Monge formulation (Kantorovich, 1942; Peyré et al., 2019). In this study, we will focus on the Kantorovich formulations.
Marco Cuturi. 2013. Sinkhorn distances: Lightspeed computation of optimal transport. In NIPS.
Mayur Datar, Nicole Immorlica, Piotr Indyk, and Va- hab S Mirrokni. 2004. -ecality-sensitive hashing scheme based on p-stable distributions. In Proceed- ings of the twentieth annual symposium on Computa- tional geometry, pages 253-262.
Gao Huang, Chuan Guo, Matt Kusner, Yu Sun, Fei Sha, and Kilian Q Weinberger. 2016. Supervised word mover's 's distance. NIPS.
Kusner al. (2015) proposed the word mover's distance (WMD) for NLP tasks. WMD is given as
Piotr Indyk and Nitin Thaper. 2003. Fast image retrieval via embeddings. In 3rd international workshop on statistical and computational theories of vision, vol- ume 2. page 5. Nice, France.
m m' ΣΣ
WI(p,V) = min neu(a,aa i=1 ==1
Fidd(z; xzj),
Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 7(3):535-547.
where μt = E3 azdn; and odd==a In 1= and a aT =
addition, x E Ied is a word vector, II a transport plan, and U denotes a set of joint probability distributions.
L Kantorovich. 1942. On the transfer of masses (in russian). Doklady Akademii Nauk, volume 37, pages 227-229.
WMD is highly used in various natural language processing tasks, including document classification (Kusner et al., 2015; Sato et al., 2022), text gen- eration, and machine translation. However, the computation of WMD is cubic with respect the number of data points. Therefore, it cannot effec- tively scale to large datasets.
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Wein- berger. 2015. From word embeddings to document distances. ICML.
Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. 2019. Tree-sliced approximation wasser- stein distances. NeurIPS.
Gaspard Monge. 1781. Mémoire sur la théorie des déblais des remblais. Mem. Math. Phys. Acad. Royale Sei., pages 666-704.
To speed up the computation of WMD (or Wasserstein distance), we can use the Sinkhorn algorithm (Cuturi, 2013), which solves entropic regularized optimal transport problems. Owing to the Sinkhorn algorithm, WMD can be efficiently solved in quadratic time.
Gabriel Peyré, Marco Cuturi, al. 2019, Computa- tional optimal transport. Foundations and Trends in Machine Learning, 155-6):355-677
Ryoma Sato, Makoto Yamada, and Hisashi Kashima. 2022. Re-evaluating word mover's distance. ICML.
To further speed up the computation of Wasser- stein distance, Wasserstein distance on a tree is useful (Indyk and Thaper, 2003). Specifically, it first embeds data points in a tree and computes the optimal transport on the tree. The advantage of tree-based optimal transport is that the Wasserstein distance can be obtained analytically and computed in linear cost. Le et al. (2019) proposed tree sliced variants of Wasserstein distance and showed that the performance significantly be improved by using multiple trees. Morcover, TWD can be represented by the LI distance between two embedded vectors
Yuki Takezawa, Ryoma Sato, and Makoto Yamada. 2021. Supervised wrsswsssernnein distance. In ICML.
Makoto Yamada, Yuki Takezawa, Ryoma Sato, Han Bao, Zornitsa Kozareva, and Sujith Ravi. 2022. Approximating I-wasserstein distance with trees. TMLR.
Sho Yokoi, Ryo Takahashi, Reina Akama, Jun Suzuki, and Kentaro Inui. 2020. Word rotator's distance. EMNLP.

---

(Takezawa al., 2021). Thus, it can be efficiently computed using GPUs. The tree-based method is computationally efficient; however, it does not approximate the Wasserstein distance using an ar- bitrary metric. To solve this issue, Yamada et al. (2022) recently proposed regression-basee ap- proach to approximate Wasserstein distance using TWD.
B Additional results
Figure shows the tree construction time that is performed only once for each method. The smaller the n_slice, the faster the construction time. Table gives the average time required to classify 30% of the data from each dataset for the k-nearest neighhors for & (1, 5, 10} when using CPU FAISS index. The smaller the n slsce the faster the construction time.
Nearest Neighbor Search: Nearest neighbor search is fundamental technique used in vari- ous fields, including natural language processing (NLP), machine learning, data mining, informa- tion retrieval, and computer vision. The nearest neighbor search is the key technique in document retrieval, which is typical application of NLP. It is also used for image retrieval in computer vision applications.
The average accuracy for all datasets is presented in Figure 5.
Figure 6 reports the results for all datasets of an- other performance indicator top-m neighbors score. Top-m neighbors score, here for m { 5), gives the proportion of the first nearest neigh- bors of given method that are actually the first nearest neighbors when using WMD. WMD is con- sidered as ground truth.
One of the key challenges the nearest neighbor search is the computational cost. Because the num- ber documents increases dramatically, a vanilla nearest neighbor search is not applicable to such large-scale data. To handle the scalability issue, an approximate nearest neighbor search, which effi- ciently finds the nearest neighbors by approximat- ing the distance computation, is heavily used.
Ultimately, Figures 7, 8 and Table 4 present the tree construction time, the average accuracy and the av- crage time required to classify 30% of the data from each database generated using the Quad'Tree approach. These results underscore that the Clus- terTrre approach for the tree construction gives better performance.
One standard approach to use kd-tree to parti- tion the space and quickly find nearest neighbors (Bentley, 1975). The other techniques are based on secssity hashing (Datar et al., 2004). Re- cently, FAISS has been proposed for finding nearest neighbors from billion-scale datasets using GPUs (Johnson et al., 2019). Although approximate near- est neighbor search algorithms are heuristic ap- proaches, they can efficiently search for nearest neighbors. Recently, several useful ANN packages have been developed, including the FAISS AN- NOY 6 and ?
Number of documents n
Number of unique words Nicof
Twitter
3108
4489
BBCsport 737
10103
Amazon 8000
30249
Classic 7093
18080
Table 2: Number of documents and unique words of the datasets used for the experiments.
https///gtthuu..cmm//aceeerereee
'https://githul hwcphhuuu.com/eeeee hooglee-rreaarch//re//ssssterss

---

Algorithm Sliced weight estimation with trees.
1: Input: The matrix X, the regularization parameter A 2 0, and set of indices Q
2: for 1.2...., do
3: (andom.seed(i)
4: B, = ClusterTree(X)
5: B: = RemoveLeaf Br)
Compute 2/3 from Bt and d(a; mz), (1,5) EΩ
argmin WER Eiiee (d(;,, - w =! + A|w||1
7: B; ==
8: end for
9: return {(B,, wz
200
100
150
seconds
80
Seconds
60
100
40
50
20
TWD
n_sice=1 n_sice33 n_siic=55 n_sice-10
TWD
n_sike=1 n.slie-33 n_slice=5 si
(a) Twitter.
(b) BBCSport.
500
300
400
Seconds
200
150
200
100
50
TM/
n_slice=1 n,slice=3 n,slicesS nuslice=10
TWD
nusiie=1 n.slic=33 n,slice=5 nulice==10
(c) Amazon.
(d) Classic.
Figure 4: One-time tree construction time (expressed in seconds) of the datasets used in the experiments.

---

BpCSptot
Amnuoo
Classic
Methods
1
5
10
10
S
10
5
10
TWD
3,969
4.033
169.923
81.603
Faiss II-TWD
n_slice=1
0.430
0,428
0.449
0.044
0.044
0.045
14.749
15.26 26.
15.228
10.330
9.860
9.997
Faiss Il ID-TWD
2.224
2.319
2.336
0.156
0.159
0.139
47.057
47.426
47.380
30.582
30.345
30.283
n_ n_slicem3
IL-TWD
4.205
4.331
4.298
0.325
0.307
0.308
78.731
80.303
80.349
51.749
51.749
$1.580
nualice=5
Paiss II-TW)
n_slice=100
9.171
8,849
9.135
0.790
0.778
0.809
163.348
160.211
159.874
102.815
102.325
102.379
Table 3: Average testing time of the datasets used the experiments for E {1,55 10) using CPU faiss index.
1.0
1.0
0.9
0.9
Accuracy
0.8
0.8
0.7
0.7
0.6
0.6
0.5
0.5
Number of nearest neighbors k
Number of nearest neighbors
I/TWO rustice-5
TWD rusticees
tudiceu1 farss (ITWO n_slic-33
(1TWD s00ice110
n_stoce2 falss ITDD nusicce33
ITWD silcem0
(a) Twitter.
(b) BBCSport.
1.0
1.0
0.9
Accuracy
0.8
Accuracy
0.7
0.6
0.6
0.5
0.5
10
15
10
15
Number
neighborrs
neighbors
Iaos TW3
faiss (HTWD suslicent
now WXO
Sares TWD
nalee I-WWD _ssleens
now
fasss I1TWO sikee1
fuiss II-WD slicem10
faiss 11-1 Sitcel
falss I-WD siice10o
WMD
IITWD rusiic=a3
IIT n_sice=3
(c) Amazon.
(d) Classic.
Figure 5: Average accuracy of the datasets used in the experiments.

---

0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
Me TWC
18Ba1-TWO Jilooen
taxa STVD n,38co-12
ta TVB
Tale FTWE aulccn3
Ta 1-WV Nowe10
1a 1100 2.110=1
Tosal 1SST-TWO rralioe3
50W
Tass 1a59:/1W3 kem
Teta 37WD naalsoe)
BOW
(a) Twitter.
(b) BBCSport.
0.7
0.7
0.5
0.5
0.3
0.3
0.1
0.0
1,3kc=13
rusksced
rudkiows
(c) Amazon.
(d) Classic.
Figure 6: Top-1 and Top-5 neighbors scores, with WMD ground truth, for the datasets used in the experiments.
100
200
150
Seconds
60
Seconds
100
20
0
TWD
n_sicem1
rusiee1 rusike=10
(a) Twitter.
(b) BBCSport.
800
400
700
600
300
500
Seconds
400
200
100
100
TWD
n_stice=1 n_siioma nisicoss r_sice=10
TWD
n_ske- nuskiae3 n_slicemS nusike=10
(c) Amazon.
(d) Classic.
Figure 7: One-time tree construction time (expressed in seconds) of the datasets used in the experiments for the QuadTree approach.

---

Twitter
BBCSport
Amazon
Classic 5
Methods
10
5
10
10
10
WMD
248.
250.2
250.9
2839
273.5
267.5
5382.8
5380.3
5380.5
5003.1
5079.4
S097.9
BOW
0.036
0.036
0.036
0.011
0.011
0.010
0.695
0.692
0.692
0.324
0.321
0.322
0.051
0.799
Paiss I-TWD
0.020
0.021
0.022
0.008
0.009
0.008
0.117
0.115
0.117
0.096
0.090
0.095
slicees
Faiss I--WD
slioom3
0.030
0.031
0.030
0.012
0.009
0.009
0.298
0.300
0.302
0.283
0.266
0.248
0.036
0.035
0.043
0.009
0.010
0.009
0.489
0.484
0.486
0.407
0.404
0.403
n. 5lice=5
II-TWD
n_ <lioc- 10
0.056
0.068
0.018
0.021
0.016
0.925
0.918
0.923
0.764
0.756
Table 4: Average testing time of the datasets used in the experiments for & {1,5, 10) for the QuadTree approach.
1.0
0.9
0.9
0.8
0.8
Accuracy
0.7
Accuracy
0.7
0.6
0.5
0.4
0.4
0.3
0.3
15
10
15
Number of nearest neighbors k
Number nearest neighbors
fermw (I-TWD sullicemt
TWD
I-VD h_sktees
sikce1 (+7WD susiiee
I-TWD
(I- yilooe
(a) Twitter.
(b) BBCSport.
1.0
1.0
0.9
Accuracy
0.7
Accuracy
0.7
0.6
0.6
0.5
0.5
0.4
10
15
10
15
Numbar of nearest neighbors
Number
neighborr
fares TW13
farss (I-TWD susticeuis
now wxo
Rilss TWD
ratss I-TD r_stoews
now
fwos (+TAD R.sicem1
faiss IITWD nuslice-10
fales I1-WD stowes
falss ITWD r_siice10
WMD
(I-TWD sicceed
files IITWD siswes
(c) Amazon.
(d) Classic.
Figure 8: Average accuracy of the datasets used in the experiments for the Quad'Tree approach.


Reviews:
[{"review_id": "hCXK0rlc6O", "reviewer": "Reviewer_ZksN", "summary": "", "questions": "What is speed performance when using a CPU faiss index? The GPU index may become untenable with larger data.\n\nHave you considered a continuous bag of words baseline? For example by averaging the word vectors then performing kNN? Kusner et al shows this does not work well, but it could still be interesting to sanity check and/or include as a row in the table to strengthen the motivation for using WMD.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> It would have been nice to include one or more non-WMD baselines. That being said, many such baselines are included in Kusner et al in favor of WMD (although embeddings have likely improved since then).\n\nThe results with BOW are presented in Tables 1 and 2. As you can see, the proposed method outperforms BOW. Note that the BOW is computed using CPU; the computation speed is slower than that of the proposed TWD. If we use GPUs for BOW, we can surely get almost similar computational speed of the proposed method.\n\nTable 1 : BOW average accuracy of the datasets used in the experiments for k $\\in$ \\{1,5,10\\}.\n|         |       | Twitter|       |      | BBCSport |       |      | Amazon |       |      | Classic |       |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Methods | 1     | 5       | 10    | 1     | 5        | 10    | 1     | 5      | 10    | 1     | 5       | 10    |\n| BOW     | 0.624 | 0.603   | 0.622 | 0.885 | 0.865    | 0.875 | 0.643 | 0.580  | 0.559 | 0.822 | 0.811   | 0.820 |\n\nTable 2 : Average testing time of the datasets used in the experiments for k $\\in$ \\{1,5,10\\}.\n|         |       | Twitter |       |      | BBCSport |       |      | Amazon |       |      | Classic |       |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:\n| Methods | 1     | 5       | 10    | 1     | 5        | 10    | 1     | 5      | 10    | 1     | 5       | 10    |\n| WMD | 248.7 |   250.2 |   250.9 |   283.9 |   275.5 |   267.5 |   5382.8 |   5380.3 |   5380.5 |  5093.1 |   5079.4 |   5097.9 | \n| BOW |   0.658 |   0.646 |   0.647 |   0.105 |   0.111 |   0.109 |   27.623 |   27.653 |   27.488 | 12.943 |   12.595 |   12.111 | \n|  Faiss TWD |   0.042 |   0.043 |   0.043 |   0.011 |   0.011 |   0.011 |   0.784 |   0.779 |   0.779 |  0.390 |   0.389 |   0.389 | \n| Faiss ll-TWD $n\\_{slice}=1$ | 0.020 |   0.020 |   0.020 |   0.006 |   0.006 |   0.006 |   0.104 |   0.101 |  0.101 |   0.083 |   0.081 |   0.082 | \n| Faiss ll-TWD $n\\_{slice}=3 $ |   0.032 |   0.033 |   0.033 |   0.008 |   0.008 |   0.008 |   0.244 |   0.239 | 0.239 |   0.171 |   0.167 |   0.167 | \n| Faiss ll-TWD $n\\_{slice}=5$ |   0.044 |   0.044 |   0.045 |   0.010 |   0.010 |   0.010 |   0.378 |   0.370 | 0.370 |   0.264 |   0.257 |   0.258 | \n| Faiss ll-TWD $n\\_{slice}=10$ |   0.075 |   0.072 |   0.071 |   0.017 |   0.017 |   0.017 |   0.752 |   0.745 |   0.748 |   0.493 |   0.486 |   0.486 | \n\n> The datasets explored are still fairly small. Real world datasets can easily have thousands or millions of documents.\n\n\nThank you for your response. We were able to assess the performance of our approach using small to medium-sized datasets. Our method involves combining tree embedding with readily available nearest neighbor search packages. The potential scalability of our method relies on both the scalability of the tree embedding technique and the chosen nearest neighbor package.\n\nRegarding tree embedding, our current implementation effectively scales up to handling around $10^4$ documents. However, extending this scalability to accommodate millions to billions of documents poses challenges. We are actively exploring the incorporation of efficient clustering algorithms to achieve this scalability goal. Nevertheless, refining the weight estimation process for the tree methods remains an ongoing task.\n\nAdditionally, we've encountered a limitation in the majority of nearest neighbor packages, as they lack support for sparse formats. Given that our tree embedding produces highly sparse vectors, addressing this limitation is crucial for handling larger datasets. It's worth noting that devising an efficient method for approximated nearest neighbor search with sparse data is a vibrant area of ongoing research.\n\nWe will include these discussion in the final version.\n\n> (low priority) It would be nice to see performance of WMD on additional tasks besides knn classification. For example, to retrieve documents for the BEIR tasks (perhaps using the approach from Hyde) or when retrieving passages in open domain QA.\n\nThank you very much for your suggestion. In the final version, we plan to include  nearest neighbor search experiments based on tree based methods and report top-1 and top-5 scores, where WMD is a ground truth. \n\n> What is speed performance when using a CPU faiss index? The GPU index may become untenable with larger data.\n\nThe speed performance when using a CPU faiss index is presented in the following Table 3.\n\nTable 3 : Average testing time of the datasets used in the experiments for k $\\in$ \\{1,5,10\\} using a CPU faiss index.\n|         |       | Twitter |       |      | BBCSport |       |      | Amazon |       |      | Classic |       |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:\n| Methods | 1     | 5       | 10    | 1     | 5        | 10    | 1     | 5      | 10    | 1     | 5       | 10    |\n| Faiss TWD |   3.969 |   3.917 |   4.033 |   0.389 |   0.347 |   0.351 |   174.105 |   169.923 |   165.540 |   81.241 |   80.908 |   81.603 |\n| Faiss ll-TWD $n\\_{slice}=1$ |   0.430 |   0.428 |   0.449 |   0.044 |   0.044 |   0.045 |   14.749 |   15.265 |   15.228 |   10.330 |   9.860 |   9.997 |\n| Faiss ll-TWD $n\\_{slice}=3$ |   2.224 |   2.319 |   2.336 |   0.156 |   0.159 |   0.139 |   47.057 | 47.426 |   47.380 |   30.582 |   30.345 |   30.283 |\n| Faiss ll-TWD $n\\_{slice}=5$ |   4.205 |  4.331 |  4.298 |   0.325 |   0.307 |   0.308 |   78.731 |   80.303 |   80.349 |   51.749 |   51.749 |   51.580 |\n| Faiss ll-TWD $n\\_{slice}=10$ |  9.171 |   8.849 |   9.135 |   0.790 |   0.778 |   0.809 |   163.348 |   160.211 |   159.874 |   102.815 |   102.325 | 102.379 |\n\n> Have you considered a continuous bag of words baseline? For example by averaging the word vectors then performing kNN? Kusner et al shows this does not work well, but it could still be interesting to sanity check and/or include as a row in the table to strengthen the motivation for using WMD.\n\nThank you for the suggestions. No, we have not run the continuous bag of words baseline. During the rebuttal period, we quickly run the BOW experiments, and we found that the proposed method outperforms BOW.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. WMD is useful for comparing documents but perhaps impractical, which is why it is common to resort to stuff documents into a single embedding. As NLP applications get more ambitious, documents are longer and perhaps it is better to use multiple embeddings. This paper can enable efficient use of WMD for applications like this.\n\n2. The results are encouraging for both accuracy and speed on multiple datasets. Although I'm not sure why \n\n3. The paper is also informative on the tradeoffs between accuracy and speed, which should be helpful for practitioners.", "reasons_to_reject": "1. It would have been nice to include one or more non-WMD baselines. That being said, many such baselines are included in Kusner et al in favor of WMD (although embeddings have likely improved since then).\n\n2. The datasets explored are still fairly small. Real world datasets can easily have thousands or millions of documents.\n\n3. (low priority) It would be nice to see performance of WMD on additional tasks besides knn classification. For example, to retrieve documents for the BEIR tasks (perhaps using the approach from Hyde) or when retrieving passages in open domain QA.\n\nBEIR https://arxiv.org/abs/2104.08663\n\nHyde https://arxiv.org/abs/2212.10496\n\nOpen-domain QA w/ retrieval https://arxiv.org/abs/2301.12652", "questions_to_authors": "What is speed performance when using a CPU faiss index? The GPU index may become untenable with larger data.\n\nHave you considered a continuous bag of words baseline? For example by averaging the word vectors then performing kNN? Kusner et al shows this does not work well, but it could still be interesting to sanity check and/or include as a row in the table to strengthen the motivation for using WMD.", "paper_topic_and_main_contributions": "This paper is about word mover's distance (WMD), a technique to compare two documents of varying length using their word embeddings. WMD can be used for kNN-based methods, but it is slow. The authors present an approximate WMD using faster nearest neighbor lookup (through FAISS), achieving a 1000x speed up and while retaining most of the accuracy of the exact method.", "excitement": 4, "reproducibility": 5, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "n/a", "Typos_Grammar_Style_and_Presentation_Improvements": "Maybe worth using L1 everywhere to be consistent and is easier to read than l1."}, "extra_text": {}}, {"review_id": "wuAAeYUbjd", "reviewer": "Reviewer_Dg6r", "summary": "", "questions": "1. Can FAISS-based tree Wasserstein acceleration be applied to unbalanced optimal transport document distance [1] or some already known tree approximation [2]?\n2. Can we train weights $w$ on one dataset and transfer them to the other dataset?\n\n[1] Wang, Z., Zhou, D., Yang, M., Zhang, Y., Rao, C., & Wu, H. (2020, September). Robust document distance with Wasserstein-fisher-rao metric. In Asian Conference on Machine Learning (pp. 721-736). PMLR.\n\n[2] Sato, R., Yamada, M., & Kashima, H. (2020). Fast unbalanced optimal transport on a tree. Advances in neural information processing systems, 33, 19039-19051.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The cluster-tree TWD and other TWD variants. This paper only combines the FAISS and the cluster-tree TWD. There are no comparisons with other strong TWD baselines such as directly using word embeddings in the quadtree data structure.\n\nWe quickly run the Quadtree experiments. The results with Quadtree are presented in the following Tables 1, 2 and 3. Compared to the Clustertree, we found that the performances of Quadtree are lower.\n\nTable 1 : Average accuracy of the datasets used in the experiments for k $\\in$ \\{1,5,10\\} with quadtree.\n|         |       | Twitter |       |      | BBCSport |       |      | Amazon |       |      | Classic |       |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:\n| Methods | 1     | 5       | 10    | 1     | 5        | 10    | 1     | 5      | 10    | 1     | 5       | 10    |\n|Faiss TWD | 0.669 |0.680 |0.701 |0.948 |0.958 |0.959 |0.822 |0.869 |0.886 |0.941 |0.954 | 0.956 |\n|Faiss ll-TWD $n\\_{slice}$=1 | 0.533 |0.629 |0.683 |0.824 |0.901 |0.905 |0.303 |0.487 |0.503 |0.464 |0.451 |  0.451 |\n|Faiss ll-TWD $n\\_{slice}$=3 |0.495 |0.693 |0.692 |0.931 |0.944 |0.934 |0.383 |0.511 |0.533 |0.470 |0.474 |  0.477|\n|Faiss ll-TWD $n\\_{slice}$=5 | 0.597 |0.689 |0.697 |0.925 |0.941 |0.944 |0.514 |0.612 |0.634 |0.526 |0.475 |  0.472 |\n|Faiss ll-TWD $n\\_{slice}$=10 |0.648 |0.681 |0.701 |0.950 |0.960 |0.958 |0.778 |0.826 |0.842 |0.772 |0.807 |  0.802|\n\n\nTable 2 : Average testing time of the datasets used in the experiments for k $\\in$ \\{1,5,10\\} with quadtree.\n|         |       | Twitter |       |      | BBCSport |       |      | Amazon |       |      | Classic |       |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:\n| Methods | 1     | 5       | 10    | 1     | 5        | 10    | 1     | 5      | 10    | 1     | 5       | 10    |\n| Faiss TWD |   0.045 |   0.051 |   0.051 |   0.012 |   0.012 |   0.012 |   0.799 |   0.803 |   0.799 |   0.642 |   0.647 |   0.637 | \n| Faiss ll-TWD $n_{slice}$=1 |   0.020 |   0.021 |   0.022 |   0.008 |   0.009 |   0.008 |   0.117 |   0.115 |   0.117 |   0.096 |   0.090 |   0.095 |\n| Faiss ll-TWD $n_{slice}$=3 |   0.030 |   0.031 |   0.030 |   0.012 |   0.009 |   0.009 |   0.298 |   0.300 |   0.302 |   0.283 |   0.266 |   0.248 |\n| Faiss ll-TWD $n_{slice}$=5 |   0.036 |   0.035 |   0.043 |   0.009 |   0.010 |   0.009 |   0.489 |   0.484 |   0.486 |   0.407 |   0.404 |   0.403 | \n| Faiss ll-TWD $n_{slice}$=10 |   0.056 |   0.058 |   0.068 |   0.018 |   0.021 |   0.016 |   0.925 |   0.918 |   0.923 |   0.736 |   0.764 |   0.756 | \n\nTable 3 : One-time tree construction time (expressed in seconds) with quadtree for the datasets used in the experiments.\n|         | Twitter | BBCSport | Amazon | Classic |\n|:---:|:---:|:---:|:---:|:---:\n| Faiss TWD                   |  40.2    |  93.1     |  339.1  |  190.1  |\n| Faiss ll-TWD $n_{slice}$=1  |  9.7     |  22.2     |  76.3   |  41.6    | \n| Faiss ll-TWD $n_{slice}$=3  |  30.0    |  64.9     |  256.7 |  132.4  |\n| Faiss ll-TWD $n_{slice}$=5  |  48.0    |  107.5    |  402.9  |  222.8  |\n| Faiss ll-TWD $n_{slice}$=10 |  96.6    |  218.7    |  806.0  |  439.5 | \n\n\n> Moreover, the potential impact of this work could be limited if the optimal transport-based search doesn't outperform the sentence embedding-based search. Approximate OT distances in L1 space might be useless under the hypothesis where some sort of sentence embeddings (possibly generated by LMs) can outperform OT distances, and almost surely outperform approximate OT. Though this hypothesis is not now fully validated, the current NLP community seems to prefer the LM-based approach.\n\nThank you for acknowledging the significance of the problem. Indeed, we consider the Optimal Transport (OT) based distance to be a potent metric, and it would be intriguing to ascertain its validity within the context of the existing literature on Language Models (LMs). While our current paper is focused on presenting a computationally efficient nearest neighbor search rooted in OT, we acknowledge that the broader goal of validating this approach within the LM literature remains a future direction. We intend to delve deeper into this aspect and augment our discussion in the final version.\n\nWe wish to underscore that the Wasserstein distance itself is a foundational technique with diverse applications in Natural Language Processing (NLP) tasks. Hence, despite the superior performance of sentence embeddings over the OT based method, it remains pertinent to cultivate a computationally efficient approach.\n\n> Can FAISS-based tree Wasserstein acceleration be applied to unbalanced optimal transport document distance [1] or some already known tree approximation [2]?\n\nTo compute the unbalanced OT, we need to solve optimization and cannot get analytical solution. Thus, it is not sure whether we can directly apply UOT for vector search. This is a very interesting future work. \n\n> Can we train weights  on one dataset and transfer them to the other dataset?\n\nIt is possible to train weights on one dataset and transfer them to the other dataset as long as the vocabulary of the new dataset is included in the dataset used to train weights.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The improvement of efficiency brought by FAISS is significant, while the performance drop compared to original TWD with cluster tree is somewhat not very large.", "reasons_to_reject": "Despite its technical soundness, the potential impact of this work is limited from two aspects:\n\n1. The cluster-tree TWD and other TWD variants. This paper only combines the FAISS and the cluster-tree TWD. There are no comparisons with other strong TWD baselines such as directly using word embeddings in the quadtree data structure.\n\n2. Moreover, the potential impact of this work could be limited if the optimal transport-based search doesn't outperform the sentence embedding-based search. Approximate OT distances in L1 space might be useless under the hypothesis where some sort of sentence embeddings (possibly generated by LMs) can outperform OT distances, and almost surely outperform approximate OT. Though this hypothesis is not now fully validated, the current NLP community seems to prefer the LM-based approach.", "questions_to_authors": "1. Can FAISS-based tree Wasserstein acceleration be applied to unbalanced optimal transport document distance [1] or some already known tree approximation [2]?\n2. Can we train weights $w$ on one dataset and transfer them to the other dataset?\n\n[1] Wang, Z., Zhou, D., Yang, M., Zhang, Y., Rao, C., & Wu, H. (2020, September). Robust document distance with Wasserstein-fisher-rao metric. In Asian Conference on Machine Learning (pp. 721-736). PMLR.\n\n[2] Sato, R., Yamada, M., & Kashima, H. (2020). Fast unbalanced optimal transport on a tree. Advances in neural information processing systems, 33, 19039-19051.", "paper_topic_and_main_contributions": "In this paper, the authors proposed to combine the fast estimation of L1 neighboring distance with Tree Wasserstein Distance. It is also natural to use libraries like FAISS by viewing TWD is essential L1 metric. Technically, it is a practical skill to accelerate the linear approximation of the computation. ", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "In line 105, B_{I,j} should be B_{i,j}."}, "extra_text": {"Missing_References": "[1] Wang, Z., Zhou, D., Yang, M., Zhang, Y., Rao, C., & Wu, H. (2020, September). Robust document distance with Wasserstein-fisher-rao metric. In Asian Conference on Machine Learning (pp. 721-736). PMLR.\n\n[2] Sato, R., Yamada, M., & Kashima, H. (2020). Fast unbalanced optimal transport on a tree. Advances in neural information processing systems, 33, 19039-19051."}}, {"review_id": "5NW0PWxqH7", "reviewer": "Reviewer_8vBA", "summary": "", "questions": "A. The paper mentions using \"IndexFlat\" for approximated nearest neighbor search, this index actually performs exhaustive search and would not have any speedup.\n\nAfter rebuttal: thanks the authors for the clarification on this point.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The four datasets used for evaluation are small. The largest one only has 8k documents. It is not clear whether the proposed method can scale to larger datasets.\n\nThank you for your response. We were able to assess the performance of our approach using small to medium-sized datasets. Our method involves combining tree embedding with readily available nearest neighbor search packages. The potential scalability of our method relies on both the scalability of the tree embedding technique and the chosen nearest neighbor package.\n\nRegarding tree embedding, our current implementation effectively scales up to handling around $10^4$ documents. However, extending this scalability to accommodate millions to billions of documents poses challenges. We are actively exploring the incorporation of efficient clustering algorithms to achieve this scalability goal. Nevertheless, refining the weight estimation process for the tree methods remains an ongoing task.\n\nAdditionally, we've encountered a limitation in the majority of nearest neighbor packages, as they lack support for sparse formats. Given that our tree embedding produces highly sparse vectors, addressing this limitation is crucial for handling larger datasets. It's worth noting that devising an efficient method for approximated nearest neighbor search with sparse data is a vibrant area of ongoing research.\n\nWe will include these discussion in the final version.\n\n> Word mover's distance is not a very accurate measure of document similarity. It ignores the word order and the syntactic structure of the documents. Thus, improving computation of WMD efficiency may not be very useful in many NLP applications.\n\nThank you for your comments. Introducing the order information in tree Wasserstein distance is indeed an interesting problem. However, compared to the standard Wasserstein distance formulation, it is difficult to incorporate the positional information. It needs more research to overcome this problem. We will add more discussion in the final version.\n\nMoreover, while our focus has been on improving the computational efficiency of WMD, we acknowledge that there are contexts where WMD may not be the most suitable measure due to the specific limitations you highlighted. Our intention in this work has been to offer a more efficient computation approach for WMD that could still benefit certain applications, especially those where capturing vocabulary associations is of greater importance than maintaining syntactic structure such as in content-based recommendation systems, plagiarism detection, and information retrieval. However, we will ensure to clarify these aspects and discuss the limitations of WMD in more detail in our final version.\n\n> A. The paper mentions using \"IndexFlat\" for approximated nearest neighbor search, this index actually performs exhaustive search and would not have any speedup.\n\nThank you for your comment. Faiss library allows approximated nearest neighbor search but \"IndexFlat\" does perform efficient exhaustive search. This will be modified in the paper. At the moment \"IndexFlat\" is the only index capable of accommodating the essential L1 metric needed for TWD. While it does perform exhaustive search, a data structure in RAM is built which performs efficiently the argmin search on the index and reduce the computation time. This computation time could further be reduced with the adaptation of other indexes to the L1 metric.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The motivation of the paper is clear. How to efficiently compute WMD is an important problem in many areas.\n2. The proposed method is backed by detailed mathematical derivations and analysis.\n3. Experiments on four datasets show some advantages of inference speed.", "reasons_to_reject": "1. The four datasets used for evaluation are small. The largest one only has 8k documents. It is not clear whether the proposed method can scale to larger datasets.\n2. Word mover's distance is not a very accurate measure of document similarity. It ignores the word order and the syntactic structure of the documents. Thus, improving computation of WMD efficiency may not be very useful in many NLP applications.", "questions_to_authors": "A. The paper mentions using \"IndexFlat\" for approximated nearest neighbor search, this index actually performs exhaustive search and would not have any speedup.\n\nAfter rebuttal: thanks the authors for the clarification on this point.", "paper_topic_and_main_contributions": "This paper proposes a new method to speed up the computation of the word mover's distance. The method is based on the L1 embedding method and the tree-based Wasserstein approximation. Based on the kNN classification results on four datasets, the proposed method is comparable to the vanilla Wasserstein distance in terms of performance and is much faster in inference speed.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors introduce an efficient approximation to the computation of the Word mover's distance, a distance metric used to compare distributions which has found applications in NLP through its usage for document classification.
Reviewers highlighted the efficiency brought by the use of approximate nearest neighbour search while only resulting in minor performance drop when using the original formulation. The method seems readily applicable and could open up venues to the usage of Wasserstein distances to larger datasets. At the same time, concerns were raised on the limitations of using this to only one variation of that distance, as well as the small datasets used to perform experiments. However, this work fulfils the call for a short paper which is to propose one neat idea.
Decision: Accept-Main