Adaptive Retrieval and Scalable Indexing for k-NN Search with Cross-Encoders

TL;DR

This paper tackles the challenge of performing accurate k-NN search with cross-encoders (CE) at scale, where exact CE scoring is prohibitively expensive. It introduces Axn, an adaptive, multi-round retrieval framework that uses offline sparse matrix factorization to learn item embeddings and a test-time linear-regression-based update to the query embedding, enabling approximate CE scores with significantly reduced CE calls. Off-line, it learns latent query and item representations by factorizing a sparse CE score matrix G, using either transductive MFTrns or inductive MFInd approaches initialized from a base dual-encoder, thereby avoiding expensive CE fine-tuning. On-line, Axn iteratively refines the query representation, selecting new candidates via inner products and finally ranking retrieved items by exact CE scores within a fixed cross-encoder budget, achieving up to 5% recall improvement at k=1 and 54% at k=100 while offering substantial indexing-time speedups (up to 100x) over CUR-based methods and distillation-based DE approaches. Empirically, across ZeShEL and BeIR benchmarks, Axn outperforms retrieve-and-rerank baselines and maintains robust downstream task performance, demonstrating practical gains in scalable CE-based retrieval without extensive fine-tuning or prohibitive offline costs.

Abstract

Cross-encoder (CE) models which compute similarity by jointly encoding a query-item pair perform better than embedding-based models (dual-encoders) at estimating query-item relevance. Existing approaches perform k-NN search with CE by approximating the CE similarity with a vector embedding space fit either with dual-encoders (DE) or CUR matrix factorization. DE-based retrieve-and-rerank approaches suffer from poor recall on new domains and the retrieval with DE is decoupled from the CE. While CUR-based approaches can be more accurate than the DE-based approach, they require a prohibitively large number of CE calls to compute item embeddings, thus making it impractical for deployment at scale. In this paper, we address these shortcomings with our proposed sparse-matrix factorization based method that efficiently computes latent query and item embeddings to approximate CE scores and performs k-NN search with the approximate CE similarity. We compute item embeddings offline by factorizing a sparse matrix containing query-item CE scores for a set of train queries. Our method produces a high-quality approximation while requiring only a fraction of CE calls as compared to CUR-based methods, and allows for leveraging DE to initialize the embedding space while avoiding compute- and resource-intensive finetuning of DE via distillation. At test time, the item embeddings remain fixed and retrieval occurs over rounds, alternating between a) estimating the test query embedding by minimizing error in approximating CE scores of items retrieved thus far, and b) using the updated test query embedding for retrieving more items. Our k-NN search method improves recall by up to 5% (k=1) and 54% (k=100) over DE-based approaches. Additionally, our indexing approach achieves a speedup of up to 100x over CUR-based and 5x over DE distillation methods, while matching or improving k-NN search recall over baselines.

Figures (16)

• Figure 1: Top-1-Recall and Top-100-Recall at inference cost budget ($m$) of 100 and 500 CE calls respectively versus indexing time for various approaches. Matrix factorization approaches (MF) can be significantly faster than $\textsc{adaCUR}_{}\xspace$ and training DE via distillation ($\textsc{DE}_{\textsc{dstl}}$). The proposed adaptive $k$-NN search method (Axn,) provides consistent improvement over corresponding retrieve-and-rerank style inference (RnR).
• Figure 2: Breakdown of indexing latency of MF and $\textsc{DE}_{\textsc{dstl}}$ into various steps with training time shown on the right of each bar for different values of $\lvert \mathcal{Q}_\text{train} \rvert$ and no. of items scored per query ($k_d$).
• Figure 3: Downstream task performance versus indexing time for proposed and baseline approaches on different domains. All methods use a fixed inference cost budget of 100 cross-encoder calls.
• Figure 4: Score distribution for cross-encoder (CE) and dual-encoder (DE) models on SciDocs for BeIR and YuGiOh from ZeShEL. For each domain, we use cross-encoder and dual-encoder models trained on the corresponding task. See §\ref{['apndx_subsec:cross_encoder_training']} for details on cross-encoder training and §\ref{['apndx_subsec:dual_encoder_training']} for dual-encoder training.
• Figure 5: Breakdown of inference latency for $\textsc{adaCUR}_{\textsc{DE}_{\textsc{src}}\xspace}\xspace$ and Axn$\textsc{DE}_{\textsc{src}}$,$\textsc{DE}_{\textsc{src}}$ under different test-time CE call budgets for domain=Hotpot-QA. See §\ref{['apndx_subsec:adaptive_search_overhead']} for detailed discussion.