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. For this reason, existing approaches perform k-NN search with cross-encoders using heuristic retrieve-and-rerank approaches where retrieval is done 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 a given cross-encoder model. 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/train queries. With the goal of reducing the indexing complexity in order to scale to millions of items, we propose to learn query/item embeddings using sparse-matrix factorization-based approaches to approximate the cross-encoder. We perform k-NN search for a given test-query using the approximate scores with an adaptive search strategy that performs retrieval over multiple rounds and uses feedback on retrieved items to improve the cross-encoder approximation and hence retrieval in subsequent rounds. Empirically, our proposed approaches provide significant improvements in recall-vs-latency tradeoffs over existing retrieve-and-rerank pipelines on zero-shot entity linking and information retrieval benchmarks.
In the final chapter, we propose various loss functions and strategies for training cross-encoder models to improve k-NN search performance of the proposed methods by making the resulting cross-encoder score matrix easier to approximate without affecting the accuracy of cross-encoder similarity on downstream tasks. We also use proposed k-NN search methods to dynamically sample negative items/examples while training cross-encoders to improve the robustness of cross-encoder models on downstream tasks.
Advisor: Andrew McCallum