Large-scale similarity search with Optimal Transport (2023)

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.

## introduction information retrieval is the process of extracting relevant information from large collections of data. it has applications in search engines, recommendation systems, and content analyses. an information retrieval task can be formulated as a k-nearest neighbor search problem. one major challenge in k-nearest neighbor search is the scalability of document retrieval tasks, where data can range from millions to billions of documents. the bag-of-words (bow) model is a widely used approach for document representation in natural language processing and information retrieval. it represents a document as a vector of word frequencies, disregarding the order and context of the words. this approach is simple and efficient; however, it 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, leading to poor retrieval performance when dealing with documents that are dissimilar in content. word mover's distance (wmd) @xcite and its variants @xcite are document dissimilarity measures that can handle overlapping problems arising in the bow model. wmd calculates the distance between two documents by measuring the minimum cost of transforming the word embeddings of one document into those of another @xcite . 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, a tree-based approximation of wmd known as tree-wasserstein distance (twd) can be used. quadtree @xcite and clus-tertree @xcite are techniques that enable tree construction for twd and can be computed in linear time. recently, @xcite applied quadtree to large-scale nearest neighbor searches. however, their performance is limited. moreover, @xcite proposed flowtree, an algorithm that has improved performance but increased 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. in this study, we introduce an ot-based nearest neighbor search for large-scale problems. specifically, we combine three methods: the l1regularized technique for learning the weights of the edges in a tree @xcite , the sliced variant of twd @xcite , and the nearest neighbor search based on faiss @xcite . this combination enable the construction of a computationally efficient ot-based nearest neigh-bor search that can be scaled to a large dataset. through benchmark experiments, we demonstrated that the proposed framework can compute the nearest neighbors three orders of magnitude faster than the vanilla wasserstein distance. contribution: the contribution of this paper is summarized below: ## proposed framework in this section, we propose the nearest neighbor search method based on twd @xcite . ## experiments in this section, we assess the effectiveness of our proposed method using twitter, bbcsport, amazon, and classic datasets. ## conclusion in this study, we introduced a methodological improvement aimed at enhancing the computational speed of the wasserstein distance approximation with trees. specifically, we combined the l1regularized technique to learn the weights of the edges in a tree @xcite , the sliced-twd @xcite , and the efficient nearest neighbor search @xcite . to further reduce the computational time, we disregarded the leaf nodes of the trees, as they can significantly reduce the dimensions of the l1 embedding vectors. the sliced twd allowed us to deal with the decrease in representation power resulting from the leaf loss. through benchmark experiments, we demonstrated that the proposed method can reduce computation time over 10 3 times compared with the wmd method while maintaining a reasonable level of accuracy.