Adaptive Retrieval and Scalable Indexing for k-NN Search with Cross-Encoders
Authors: Nishant Yadav, Nicholas Monath, Manzil Zaheer, Rob Fergus, Andrew McCallum
ICLR 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We perform an empirical evaluation of our method using cross-encoder models trained for the task of entity linking and information retrieval on ZESHEL (Logeswaran et al., 2019) and BEIR (Thakur et al., 2021b) benchmark respectively. |
| Researcher Affiliation | Collaboration | Nishant Yadav1 , Nicholas Monath2, Manzil Zaheer2, Rob Fergus2, Andrew Mc Callum1 1 University of Massachusetts Amherst, 2 Google Deep Mind |
| Pseudocode | Yes | Algorithm 1 AXN Test-time k-NN Search Inference |
| Open Source Code | No | The paper provides links to pre-trained model checkpoints (e.g., 'https://huggingface.co/nishantyadav/emb-crossenc-zeshel', 'www.sbert.net/docs/pretrained-models/msmarco-v2.html') but does not explicitly state that the source code for the proposed methodology (sparse-matrix factorization and AXN) is openly available or provide a direct link to it. |
| Open Datasets | Yes | We perform an empirical evaluation of our method using cross-encoder models trained for the task of entity linking and information retrieval on ZESHEL (Logeswaran et al., 2019) and BEIR (Thakur et al., 2021b) benchmark respectively. |
| Dataset Splits | Yes | Table 1: Statistics on number of items (I), number of queries in train (Qtrain) and test (Qtest) splits for each domain. For BEIR datasets, we tune AXN weight parameter λ (in eq 7) on the dev set. |
| Hardware Specification | Yes | On an Nvidia 2080ti GPU with 12 GB memory using batch size=50. We train dual-encoder models on Nvidia RTX8000 GPUs with 48 GB GPU memory. We train both transductive (MFTRNS) and inductive (MFIND) matrix factorization models on Nvidia 2080ti GPUs with 12 GB GPU memory for all datasets with the exception that we trained MFTRNS for Hotpot-QA on Nvidia A100 GPUs with 80 GB GPU memory. |
| Software Dependencies | No | The paper mentions software components like 'Scikit-learn (Pedregosa et al., 2011)' and optimizers such as 'Adam W (Loshchilov & Hutter, 2019)', but it does not provide specific version numbers for the key ancillary software dependencies (e.g., Python, PyTorch, TensorFlow, etc.) used for the experiments. |
| Experiment Setup | Yes | We train dual-encoder models on Nvidia RTX8000 GPUs with 48 GB GPU memory. The DE models are trained using cross-entropy loss to match the DE score distribution with the CE score distribution. We train DEDSTL using Adam W (Loshchilov & Hutter, 2019) optimizer with learning rate 1e-5 and accumulating gradient over 4 steps. We trained for 10 epochs when using top-100 items per query and for 4 epochs when using top-1000 items per query. For BEIR datasets, we use BM25 with parameters as reported in Thakur et al. (2021b) and for ZESHEL, we use TF-IDF with default parameters from Scikit-learn (Pedregosa et al., 2011). We use R = 10 for domains in BEIR and R = 5 for domains in ZESHEL unless stated otherwise. For BEIR datasets, we tune AXN weight parameter λ (in eq 7) on the dev set. |