The goal of k-nearest neighbor (k-NN) search is to find the top-k similar items for a given query. This task is a widely used subroutine in search, recommendation, question-answering and many other applications in data mining, machine learning, and information retrieval systems. For many of these tasks, the state-of-the-art query-item similarity models are black-box neural similarity functions such as cross-encoders that jointly encode a query-item pair to directly output a scalar similarity. However, unlike vector-based similarity functions (e.g., inner product), computing a single query-item score using cross-encoders can be computationally expensive as cross-encoders are typically parameterized using deep neural models such as transformers. For this reason, existing approaches perform k-NN search with cross-encoders using heuristic retrieve-and-rerank approaches where retrieval with a separate model (such as dual-encoder or BM25) followed by re-ranking using the cross-encoder.
In this thesis, we propose efficient matrix factorization-based approaches to support k-NN search with cross-encoders. First, we introduce, ANNCUR, a CUR-decomposition-based method for finding (approximate) k-nearest neighbor items for a test-query. ANNCUR approximates the cross-encoder scores of all items for the test-query by comparing the test-query with only a small subset of adaptively-chosen anchor items, and performs retrieval using the approximate scores.
The indexing step for CUR-based approaches computes a dense matrix by scoring all items against a set of anchor/training queries. With the goal of reducing the indexing complexity in order to scale to a millions of items, we explore methods for factorizing sparse matrices and efficient approaches to leverage existing dual-encoder models while avoiding expensive distillation-based training.
We next propose to study whether we can combine graph-based/tree-based k-NN search approaches with node/edge filtering heuristics based on approximate cross-encoder scores from the matrix factorization-based approaches. We will continue to evaluate our new methods on zero-shot entity linking and zero-shot information retrieval benchmarks using k-NN recall as well as indexing and inference latency metrics. We also propose to explore avenues for training the cross-encoder model such that the resulting cross-encoder score matrix is easier to factorize without affecting the performance of cross-encoder similarity on downstream tasks as well as improving robustness of cross-encoder models by training them using negative examples which are dynamically sampled using the proposed k-NN search methods.
Advisor: Andrew McCallum