Efficient Nearest Neighbor Search for Cross-Encoder Models using Matrix Factorization

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2171 - 2194 December 7-11, 2022 ©2022 Association for Computational Linguistics Efficient Nearest Neighbor Search for Cross-Encoder Models using Matrix Factorization Nishant Yadav†, Nicholas Monath♢, Rico Angell†, Manzil Zaheer♢, and Andrew McCallum† † University of Massachusetts Amherst ♢Google Research {nishantyadav, rangell, mccallum}@cs.umass.edu {nmonath,manzilzaheer}@google.com Abstract Efficient k-nearest neighbor search is a fun- damental task, foundational for many prob- lems in NLP. When the similarity is measured by dot-product between dual-encoder vectors or ℓ2-distance, there already exist many scal- able and efficient search methods. But not so when similarity is measured by more accurate and expensive black-box neural similarity mod- els, such as cross-encoders, which jointly en- code the query and candidate neighbor. The cross-encoders’ high computational cost typi- cally limits their use to reranking candidates retrieved by a cheaper model, such as dual en- coder or TF-IDF. However, the accuracy of such a two-stage approach is upper-bounded by the recall of the initial candidate set, and poten- tially requires additional training to align the auxiliary retrieval model with the cross-encoder model. In this paper, we present an approach that avoids the use of a dual-encoder for re- trieval, relying solely on the cross-encoder. Re- trieval is made efficient with CUR decomposi- tion, a matrix decomposition approach that ap- proximates all pairwise cross-encoder distances from a small subset of rows and columns of the distance matrix. Indexing items using our ap- proach is computationally cheaper than training an auxiliary dual-encoder model through distil- lation. Empirically, for k > 10, our approach provides test-time recall-vs-computational cost trade-offs superior to the current widely-used methods that re-rank items retrieved using a dual-encoder or TF-IDF. 1 Introduction Finding top-k scoring items for a given query is a fundamental sub-routine of recommendation and information retrieval systems (Kowalski, 2007; Das et al., 2017). For instance, in question answering systems, the query corresponds to a question and the item corresponds to a document or a passage. Neural networks are widely used to model the sim- ilarity between a query and an item in such ap- plications (Zamani et al., 2018; Hofstätter et al., 2019; Karpukhin et al., 2020; Qu et al., 2021). In this work, we focus on efficient k-nearest neigh- bor search for one such similarity function – the cross-encoder model. Cross-encoder models output a scalar similarity score by jointly encoding the query-item pair and often generalize better to new domains and unseen data (Chen et al., 2020; Wu et al., 2020; Thakur et al., 2021) as compared to dual-encoder 1 models which independently embed the query and the item in a vector space, and use simple functions such as dot-product to measure similarity. However, due to the black-box nature of the cross-encoder based similarity function, the computational cost for brute force search with cross-encoders is prohibitively high. This often limits the use of cross-encoder models to re-ranking items retrieved using a sepa- rate retrieval model such as a dual-encoder or a TF- IDF-based model …

Content truncated to 500 words