ReTuning: Towards Scalable In-context Retrieval
Nilesh Gupta, Chong You, Srinadh Bhojanapalli, Sanjiv Kumar, Inderjit Dhillon, Felix Yu
Information retrieval (IR)—finding relevant documents for a given query from a corpus—is a problem of fundamental importance in machine learning.  An emerging paradigm for IR is In-context Retrieval (ICR), which leverages large language models (LLMs) by directly incorporating candidate documents, the query, and task description into the model's input prompt for relevance identification.  While this method harnesses LLM capabilities, achieving efficient and effective ICR presents a significant challenge, especially as the candidate list expands.  To this end, this paper first identifies inherent structures in LLMs finetuned for ICR: (1) document tokens attend primarily to tokens within the same document and in the instructions when prompted properly;  (2) query tokens, particularly the last ones and certain special ones such as ":", attend strongly towards the relevant documents.  Motivated by these observations, we introduce ReTuning (Retrieval-Tuning), a novel method that adapts an LLM for ICR by (1) enforcing a structured sparse attention, which reduces attention complexity from quadratic to linear, and (2) directly optimizing attention to the relevant documents from the signal-carrying query tokens using a contrastive training objective, which improves retrieval effectiveness.  Experiments on MSMarco and NQ with Mistral-7B demonstrate that ReTuned Mistral matches or outperforms standard fine-tuned Mistral while being significantly more efficient at inference (4.7x for 100 MSMarco documents in context) and 2−3x during training.  Our approach presents a scalable and effective solution for ICR.