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

Thank you for your time and valuable feedback. We hope that our response helps to address the points you’ve raised. We would be happy to further discuss and clarify any remaining questions/concerns.

(1) This paper uses K nearest neighbor search to iteratively update query embedding and approximate cross-encoder results. However, many references of nearest neighbor search are missing

• In the related work section, we have over 25 citations outlining the relationship between our work and kNN search literature, especially in the Nearest Neighbor Search paragraph of related work. This paragraph in related work covers data structures for efficient indexing, speeding up similarity computation (e.g. using quantization), and the combination of graph-based/tree-based indexes and neural models.
• Our work focuses on approximating an expensive black-box similarity function (eg. cross-encoders) using matrix factorization approaches to perform kNN search, and we discuss related work in the Approximating Similarity Function paragraph in the Related Work section. If you have further suggestions for citations for this paragraph, we greatly appreciate them.
• Compared to work on index structures for efficient kNN search (e.g.,CoverTree, HNSW, etc), we view our work as orthogonal (e.g., the vectors produced by our method could be indexed by any such vector-based kNN index). See second paragraph Sec 2.1 of paper for more discussion.
• If, given the above comments, the reviewer feels we have inadequately placed our work in the landscape of kNN search, please do let us know. Clearly presenting our work amidst that landscape is important to us. We welcome any further content or citation suggestions that the reviewer might have. Thank you for raising this point to us.

(2) AXN utilizes sparse matrix to reduce index costs. The paper lacks detailed analysis regarding this technique's impact on results concerning varying degrees of sparsity.

• The ablation you reference is indeed already contained in our paper; it is present both in Table 3 and in Appendix B.4. These experiments vary the number of queries in the sparse matrix as well as number of items scored per query (i.e. sparsity) and measure the change in test-time Top-k-Recall and indexing latency.
• In Appendix B.4, we observe that increasing the number of items per query in the matrix (i.e. reducing sparsity), and increasing number of train queries generally leads to improvement in test-time Top-k-Recall although at increased indexing cost, as shown in Figure 10 (in updated pdf). This is an expected consequence; e.g., more observed data improves recall performance.
• In Table 3, we show a detailed breakdown of indexing cost for different configurations of the sparse matrix. Here are some key observations:
  • Computing the cross-encoder (CE) scores for query-item pairs in the sparse matrix takes a significant fraction of overall indexing time.
  • The time taken to factorize the sparse matrix is proportional to the number of query-item pairs observed in the sparse matrix. However, it remains a small fraction of the overall indexing cost.

(3) … AXN and CUR-based methods need to compute low-dimensional embeddings for queries and items. AXN uses sparse matrix to reduce the cost. This can also be applied to CUR-based methods to reduce the index time. …

• The most significant component of indexing cost is computation of cross-encoder scores between train/anchor queries and items.
• To the best of our understanding, the item indexing step for CUR-based methods must compute a dense matrix by scoring each item against all of the anchor queries while we propose use of sparse matrices to significantly bring down the number of query-item scores computed using the cross-encoder.
• The indexing cost for AdaCUR can be controlled by the varying number of anchor/train queries, and in Fig 1, we report AdaCUR results for varying number of anchor/train queries.
• As shown in Fig 1, at a given indexing cost, our sparse matrix factorization based approaches yield better test-time Top-k-Recall values than AdaCUR.
• Were there other ablations that the reviewer had in mind that might help improve our understanding of where our proposed sparse matrix based approaches outperforms CUR-based methods such as AdaCUR?

Thank you for your time and valuable feedback. We have attempted to address the concerns raised by the reviewer and we would be happy to further discuss and clarify any remaining questions/concerns.

Inference Time for AXN vs DE-based retrieve-and-rerank

The inference time and its comparison to simple DE+kNN approaches is not made clear. Would a large dual encoder with the same inference time as AXN match its recall? Some of this is tucked away in the appendix in Fig 5 and 6 which seems to suggest the margins between DE and AXN are low normalized for inference cost.

• We present a direct comparison between AXN and DE-based retrieve-and-rerank at varying inference cost budgets in Figure 6 (Fig 7 in updated pdf) where AXN consistently outperforms DE-based retrieve-and-rerank baseline.
• We measure inference cost by the number of cross-encoder (CE) calls made at test-time in Figure 6 (Fig 7 in updated pdf) as inference latency for AXN (with 5 to 10 rounds) and DE-based retrieve-and-rerank approaches is similar at a given CE call budget.
  • The reason behind similar inference latency for AXN (with 5 to 10 rounds) and DE-based retrieve-and-rerank is that the inference latency overhead incurred by AXN is less than 2 to 5% of overall time when using 5-10 rounds for inference and the majority of time is spent in computing cross-encoder scores for the retrieved items, as shown in Figure 5.
• Using a larger dual-encoder would not increase the inference time significantly as the main test-time bottleneck is computing CE scores for retrieved items.
  • We add additional results for a larger dual-encoder (Fig 27 and 28 in updated pdf) where we observe AXN continues to outperform DE-based retrieve-and-rerank baseline. See second bullet in Choice of Dual-Encoder paragraph for details.

Choice of Dual-Encoder

The choice of DE is not discussed. Is it completely irrelevant? Is it being compared to state-of-the-art encoders? would a large dual encoder with same inference time as AXN match it's recall? Or would DE saturate well before AXN.

We run experiments with different dual-encoder baselines and justify the choices below

• For ZeShEL dataset, the dual-encoder model used is representative of the state-of-the-art model.

• For BeIR benchmark, we used a dual-encoder model released publicly that is representative of the state-of-the-art model under the setting which does not involve any distillation-based training on MS-MARCO dataset.

  • We refer the reviewer to Appendix A.2.1 for more details of the dual-encoder models used in this paper.

• We report additional results in Figures 27 and 28 (in updated pdf) for another larger dual-encoder model which is representative of the state-of-the-art on MS-MARCO dataset (see here).

  • This is bert-base model finetuned on MS-MARCO using ground-truth data as well as further trained through several rounds of distillation using a cross-encoder (different from the ones used in this paper).
  • We observe similar trends such as
    • Adaptive retrieval (AXN) outperforms retrieve-and-rerank style inference, and
    • Sparse matrix factorization based approaches can offer improved performance as compared to the dual-encoder.

• Note that we additionally also compare with $DE_{Dstl}$, dual-encoder models further finetuned on the target domain via distillation using the given cross-encoder model.

  • We find that such distillation based training is compute- and resource-intensive.
  • In contrast, our proposed sparse matrix factorization based approaches can be up to 5x faster while matching or improving over the performance of distilled dual-encoder models.
  • Also, as shown in Figure 1, using AXN with $DE_{Dstl}$, outperforms retrieve-and-rerank approach (RnR) with the same distilled dual-encoder $DE_{Dstl}$ for k=100 and matches performance for k=1.

Thank you for your time and valuable feedback. We have attempted to address the concerns raised by the reviewer and we would be happy to further discuss and clarify any remaining questions/concerns.

Gains are primarily under fixed index time scenario which is usually a one-time cost

• While indexing is a one-time cost, it is a significant one. In particular, it is crucial to scaling to large numbers of items in terms of time and compute resources.

  • Existing approaches such as CUR-based methods do not scale well to millions of items and can require 1000s of GPU hours, and training parametric dual-encoder models via distillation also requires significant amount of time and multiple GPUs with large amounts of memory.
  • In contrast, our proposed matrix factorization based approaches incur up to 100x and 5x less indexing time as compared to AdaCUR or dual-encoder distillation-based approaches respectively.

• In addition to proposing time- and resource-efficient indexing methods, we propose an adaptive retrieval method (AXN) that yields significant improvement over existing dual-encoder (DE) based retrieve-and-rerank pipelines,

• Moreover, AXN can be used in combination with any existing DE model to improve test-time performance without incurring any additional indexing cost.

The proposed approach is evaluated only on zero-shot tasks, does this approach also benefits standard retrieval tasks

• We focus on zero-shot tasks as poor performance of dense-embedding models in zero-shot setting motivates the use expensive cross-encoder models.
• As shown in Table 2 in Thakur et al. (2021), there is a substantial gap of about 15 points between zero-shot performance of dual-encoder and cross-encoder models.
• As reported in Table 2 in Thakur et al. (2021), using cross-encoder models yield marginal (about 1%) improvements over dual-encoder/dense embedding models on in-domain datasets such as MS-MARCO, and the marginal gains may not justify using cross-encoders for retrieval in such in-domain settings.

Response to Questions

It is a bit surprising that RnR DE models are performing worse than TF-IDF on Hotpot, is it because this CE model was trained on triplets mined using TF-IDF?

• As shown and reported in prior work (Thakur et al., 2021), TF-IDF is a strong baseline for these zero-shot information retrieval tasks and dual-encoders typically perform worse than TF-IDF by 3-4% on average (see Table 2 in Thakur et al. (2021)).
• While TF-IDF performing better than dual-encoders on downstream task metric is a result of the nature of the task and data, we hypothesize that TF-IDF also performing better at kNN search with cross-encoder indicates that the cross-encoder is aligned with the TF-IDF similarity measure to some extent.
  • We hypothesize that this alignment could be a result of both the training strategy used for training the cross-encoder model as well as the nature of the task and data.

CE model is trained conditioned on some specific negativing mining distribution so maybe for RnR baselines we should also compare with a retrieval model which is the same as the retrieval model used for the negative mining so that the train and test-time behaviours are same

• We agree that the training strategy used for training the cross-encoder can play a critical role in the performance of various baselines.
• A detailed study of how the cross-encoder training strategy affects performance of various retrieval methods is an interesting research question but beyond the scope of this paper.
• For ZeShEL datasets, the CE model is trained using negatives retrieved using a separately trained dual-encoder (DE), and this is the same DE model that is used in the DE-based retrieve-and-rerank baseline reported in the paper. We used the CE model released by Yadav et al. (2022) in our experiments.
• For BeIR, the baseline DE model is trained on the same set of labeled triplets as the cross-encoder model.

For a given budget of X CE calls, how should one distribute X calls in the number of rounds and K CE calls inside each round in AXN inference

Distributing CE call budget over multiple rounds

• In our experiments, we distributed the CE call budget uniformly over each round.
• We observed that for a given CE call budget, increasing the number of rounds leads to improvements in Top-k-Recall metrics (as reported in Fig 6 in updated pdf) at the expense of increased overhead (as reported in Figure 5) with the Top-k-Recall saturating around 5-10 rounds.
• We use 10 rounds for BeIR domains and 5 rounds for ZeShEL domains.

References:

• Thakur, Nandan, et al. "BeIR: A heterogenous benchmark for zero-shot evaluation of information retrieval models." NeurIPS 2021 Dataset and Benchmark Track.
• Yadav, Nishant, et al. "Efficient nearest neighbor search for cross-encoder models using matrix factorization." EMNLP 2022.

Thank you for your time, thoughtful review, and support of our paper.

Figure 1's legends and corresponding subplots are hard to read. The subplots are too small, and hard to map points to the legends.

We have included a larger version of the figures in the Appendix (Figures 9 through 12 in the original submission and Figures 11 through 16 in the updated pdf) as we were able to enlarge the figure 1 in the main paper marginally given the current page limit constraints.

Response to Questions:

In the experiment section it was not clear how many rounds of updates are performed in all AXN experiments.

• As mentioned in Appendix A.4, we use 10 rounds for BeIR datasets and 5 rounds for ZeShEL datasets, unless specified otherwise.

Consider fixing max of CE calls B_{CE}, but varying the number of iterative search rounds and the number of items to retrieve in each round. Will it affect AXN's performance and total indexing time?

• We added additional results in Figure 6 which show that, for a fixed CE call budget, performance of AXN improves as we increase the number of iterative search rounds, with the performance saturating around 5-10 rounds. We divide the CE call budget uniformly over each round.
• Figure 5 shows the trend for how the overhead for AXN varies under different test-time cost budgets ( $B_{CE}$ ), and we observe that AXN incur 2 to 5% overhead when using 5-10 rounds for inference.
• The offline indexing time is independent of maximum test-time CE call budget, and depends on the number of query-item pairs scored using the cross-encoder in the sparse matrix as well as time taken to factorize the matrix during the offline indexing stage.