Early Exit Strategies for Approximate k-NN Search in Dense Retrieval (2408.04981v1)

Abstract: Learned dense representations are a popular family of techniques for encoding queries and documents using high-dimensional embeddings, which enable retrieval by performing approximate k nearest-neighbors search (A-kNN). A popular technique for making A-kNN search efficient is based on a two-level index, where the embeddings of documents are clustered offline and, at query processing, a fixed number N of clusters closest to the query is visited exhaustively to compute the result set. In this paper, we build upon state-of-the-art for early exit A-kNN and propose an unsupervised method based on the notion of patience, which can reach competitive effectiveness with large efficiency gains. Moreover, we discuss a cascade approach where we first identify queries that find their nearest neighbor within the closest t << N clusters, and then we decide how many more to visit based on our patience approach or other state-of-the-art strategies. Reproducible experiments employing state-of-the-art dense retrieval models and publicly available resources show that our techniques improve the A-kNN efficiency with up to 5x speedups while achieving negligible effectiveness losses. All the code used is available at https://github.com/francescobusolin/faiss_pEE

• The paper introduces a patience-based early exit method that reduces the number of clusters probed in dense retrieval.
• It compares regression-based, classification, and cascade approaches, demonstrating up to 5.13× speedups with minimal loss in effectiveness.
• Extensive experiments on models like STAR and CONTRIEVER validate the efficiency gains for large-scale approximate k-NN search.

Early Exit Strategies for Approximate $k$-NN Search in Dense Retrieval

In the paper titled "Early Exit Strategies for Approximate $k$-NN Search in Dense Retrieval," Francesco Busolin et al. address the challenge of optimizing approximate $k$-nearest-neighbors ($k$-NN) search in the context of dense retrieval. Dense retrieval uses neural encoders to map queries and documents into high-dimensional vectors, enabling efficient retrieval by performing approximate $k$-NN search on these embeddings. A typical approach involves a two-level indexing system wherein the document embeddings are clustered. During query processing, the closest clusters to the query are exhaustively examined. The approach taken by Busolin et al. aims to enhance this process via early exit strategies, enabling competitive effectiveness with remarkable efficiency gains.

Methodology Overview

The authors propose an unsupervised method based on the concept of "patience" to determine when to terminate the search for the $k$-NN. They explore several adaptive strategies, including a cascade approach that initially identifies if a query's nearest neighbor can be found within a small subset of clusters. If not, further clusters are visited based on either the patience method or other state-of-the-art techniques.

Problem Statement

Given an embedded query and a precomputed set of document embeddings, the goal is to identify the $k$ most relevant documents efficiently. Traditional $k$-NN approaches probe a fixed number $N$ of clusters for every query, balancing between retrieval accuracy and computational cost. The authors argue for adaptive methods that dynamically adjust the number of clusters probed based on the query's characteristics.

Regression-Based Approach

The regression-based method, as introduced by prior work [Li et al., 2020], estimates the number of clusters to probe using a learned regression model. It chooses the number of clusters to visit based on per-query predictions informed by a regression model trained on features related to the query and the clusters' similarities.

Patience-Based Approach

The patience-based method examines the stability of the result set as more clusters are probed. If the result set stabilizes (i.e., the intersection of consecutive result sets remains relatively constant over a few iterations), the search is terminated early. This approach leverages the empirical observation that the result set saturates after examining a certain number of clusters.

Classification and Cascade Approaches

In the classification approach, queries are initially classified into two categories: those likely to find their nearest neighbors in $\tau$ clusters and those that may require more. For the latter, a cascade approach is employed where the first stage uses the classification method and subsequent stages use either the regression model or the patience-based method to decide if more clusters need to be probed.

Experimental Evaluation

The authors conduct extensive experiments using three state-of-the-art dense retrieval models—STAR, CONTRIEVER, and TAS-B—on the MS-MARCO dataset. Key findings are summarized below:

• Regression-Based Method: Achieves good efficiency but does not significantly outperform the patience-based method.
• Patience-Based Method: Demonstrates substantial speedups while maintaining competitive effectiveness. The method probes significantly fewer clusters compared to regression-based approaches.
• Classification Approach: With appropriate weighting, the classifier identifies which queries can be early-exited with minimal impact on effectiveness. However, adding a cascade stage enhances efficiency further.

Numerical Results

The experiments reveal that the patience-based method offers up to 5.13× speedups compared to the baseline $k$-NN with negligible effectiveness loss. Furthermore, the cascade approach introduces efficiency improvements, though at a minor expense to retrieval effectiveness.

Implications and Future Directions

The proposed early exit strategies have notable implications for practical retrieval systems, especially those dealing with large-scale data. The adaptability of the methods ensures that retrieval systems can minimize computational overhead without significantly compromising on relevance. Future research could extend these findings by exploring the impact of varying cluster granularity and developing more nuanced patience-based heuristics. Moreover, understanding the behavior of clustered indexes under different approximation tolerances could further refine the balance between efficiency and accuracy in dense retrieval systems.

The authors conclude that their unsupervised, patience-based methods provide a promising alternative to existing techniques, achieving higher efficiency with fewer computational resources. This work exemplifies the ongoing efforts to optimize retrieval systems, making them more adept at handling the growing volumes of data encountered in real-world applications.