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

Abstract: Efficient k-nearest neighbor search is a fundamental task, foundational for many problems in NLP. When the similarity is measured by dot-product between dual-encoder vectors or L2-distance, there already exist many scalable and efficient search methods. But not so when similarity is measured by more accurate and expensive black-box neural similarity models, such as cross-encoders, which jointly encode the query and candidate neighbor. The cross-encoders’ high computational cost typically limits their use to reranking candidates retrieved by a cheaper model, such as dual encoder or TF-IDF. However, the accuracy of such a two-stage approach is upper-bounded by the recall of the initial candidate set, and potentially 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 retrieval, relying solely on the cross-encoder. Retrieval is made efficient with CUR decomposition, a matrix decomposition approach that approximates all pairwise cross-encoder distances from a small subset of rows and columns of the distance matrix. Indexing items using our approach is computationally cheaper than training an auxiliary dual-encoder model through distillation. 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.

## introduction finding top-k scoring items for a given query is a fundamental sub-routine of recommendation and information retrieval systems @xcite . 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 similarity between a query and an item in such applications @xcite @xcite . in this work, we focus on efficient k-nearest neighbor 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 @xcite @xcite as compared to dual-encoder foot_0 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 separate retrieval model such as a dual-encoder or a tf-idf-based model @xcite @xcite . the accuracy of such a two-stage approach is upper bounded by the recall of relevant items by the initial retrieval model. much of recent work either attempts to distill information from an expensive but more expressive cross-encoder model into a cheaper student model such as a dual-encoder @xcite @xcite @xcite , or focuses on cheaper alternatives to the cross-encoder model while attempting to capture fine-grained interactions between the query and the item @xcite @xcite . in this work, we tackle the fundamental task of efficient k-nearest neighbor search for a given query according to the cross-encoder. our proposed approach, anncur, uses cur decomposition (mahoney and @xcite , a matrix factorization approach, to approximate cross-encoder scores for all items, and retrieves k-nearest neighbor items while only making a small number of calls to the cross-encoder. our proposed method selects a fixed set of anchor queries and anchor items, and uses scores between anchor queries and all items to generate latent embeddings for indexing the item set. at test time, we generate latent embedding for the query using cross-encoder scores for the test query and anchor items, and use it to approximate scores of all items for the given query and/or retrieve top-k items according to the approximate scores. in contrast to distillation-based approaches, our proposed approach does not involve any additional compute-intensive training of a student model such as dual-encoder via distillation. in general, the performance of a matrix factorization-based method depends on the rank of the matrix being factorized. in our case, the entries of the matrix are cross-encoder scores for queryitem pairs. to further improve rank of the score matrix, and in-turn performance of the proposed matrix factorization based approach, we propose [emb]-ce which uses a novel dot-product based scoring mechanism for cross-encoder models (see figure 1a ). in contrast to the widely used [cls]-ce approach of pooling query-item representation into a single vector followed by scoring using a linear layer, [emb]-ce produces a score matrix with a much lower rank while performing at par with [cls]-ce on the downstream task. we run extensive experiments with crossencoder models trained for the downstream task of entity linking. the query and item in this case correspond to a mention of an entity in text and a document with an entity description respectively. for the task of retrieving k-nearest neighbors according to the cross-encoder, our proposed approach presents superior recall-vs-computational cost trade-offs over using dual-encoders trained via distillation as well as over unsupervised tf-idf-based methods ( §3.2). we also evaluate the proposed method for various indexing and test-time cost budgets as well as study the effect of various design choices in §3.3 and §3.4. ## conclusion in this paper, we proposed, anncur, a matrix factorization-based approach for nearest neighbor search for a cross-encoder model without relying on an auxiliary model such as a dual-encoder for retrieval. anncur approximates the test query's scores with all items by scoring the test-query only with a small number of anchor items, and retrieves items using the approximate scores. empirically, for @xmath0, our approach provides test-time recallvs-computational cost trade-offs superior to the widely-used approach of using cross-encoders to rerank items retrieved using a dual-encoder or a tf-idf-based model. this work is a step towards enabling efficient retrieval with expensive similarity functions such as cross-encoders, and thus, moving beyond using such models merely for re-ranking items retrieved by auxiliary retrieval models such as dual-encoders and tf-idf-based models. by the national science foundation under grant no. nsf1763618, in part by the chan zuckerberg initiative under the project "scientific knowledge base construction", in part by international business machines corporation cognitive horizons network agreement number w1668553, and in part using high performance computing equipment obtained under a grant from the collaborative r&d fund managed by the massachusetts technology collaborative. rico angell was supported by the nsf graduate research fellowship under grant no. 1938059. any opinions, findings and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect those of the sponsor(s).