Title: Scalable In-context Ranking with Generative Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: large language model

information retrieval

efficient fine-tuning

efficient inference
Abstract: In-context Ranking (ICR) is an emerging paradigm for Information Retrieval (IR), which leverages contextual understanding of LLMs by directly incorporating the task description, candidate documents, and the query into the model's input prompt and tasking the LLM to identify relevant document(s). While it is effective, efficiency is a significant challenge in this paradigm, especially as the candidate list grows due to quadratic/super-linear scaling of attention operation with context length. To this end, this paper first identifies inherent and exploitable structures in the attention of LLMs finetuned for ICR: (1) inter-document block sparsity: attention is dense within each document block but sparse across different documents in the context; and (2) query-document block relevance: the attention scores from certain query tokens to a document block in middle layers strongly correlate with that document's actual relevance. Motivated by these observations, we introduce BlockRank (Blockwise In-context Ranking), a novel method that adapts the attention operation in an LLM by (a) architecturally enforcing the observed inter-document block sparsity, reducing attention complexity from quadratic to linear without loss in performance, and (b) optimizing query-document block relevance for true relevant documents during fine-tuning using an auxiliary contrastive training objective, improving retrieval in attention. Experiments on BEIR, MSMarco and NQ with Mistral-7B demonstrate that BlockRank Mistral matches or outperforms existing SOTA listwise rankers and controlled fine-tuned baseline while being significantly more efficient at inference (4.7x for 100 MSMarco documents in context) and scaling gracefully to long-context shortlists, around 500 documents in-context (approximately 100K context length) within a second, presenting a scalable and effective solution for ICR.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents ReTuning, a fine-tuning method for improving the efficiency and effectiveness of large language models (LLMs) on In-context Retrieval (ICR) tasks. The authors identify structured attention patterns in LLMs fine-tuned for retrieval and propose two key innovations: (1) a structured sparse attention mechanism that reduces inference complexity, and (2) an auxiliary contrastive loss that explicitly aligns query-token attention with relevant documents. The proposed method achieves competitive retrieval accuracy while being significantly faster than standard fine-tuning baselines on benchmarks such as MS MARCO and Natural Questions.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n1. The paper provides a clear empirical analysis of attention patterns in fine-tuned LLMs, revealing structured sparsity and signal-carrying query tokens.\n\n2. The proposed attention-based inference mechanism avoids auto-regressive decoding. This idea is interesting.\n\n3. ReTuning outperforms or matches strong baselines on ICR tasks with notable improvements in inference efficiency, especially as the number of candidate documents scales up.\n\n**Weaknesses:**\n1. The structured sparse attention design is similar to block-wise attention patterns used in prior works such as \"Block Attention for Efficient Prefilling\", which the paper does not acknowledge, potentially overclaiming novelty.\n\n2. The method relies on pre-defined signal tokens (e.g., “:”, “[”) for attention supervision, but it remains unclear how the system handles general queries without such tokens.\n\n3. The attention-based scoring mechanism depends on extracting attention weights from intermediate layers, making it incompatible with techniques like FlashAttention, which could significantly hinder scalability in long-context scenarios.\n\n4. The use of raw summed attention over document tokens may introduce length bias, favoring longer documents regardless of relevance.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose ReTuning (Retrieval-Tuning) to address the scalability and efficiency challenges of In-context Retrieval (ICR), a paradigm where LLMs directly identify relevant documents from a list provided in the input prompt. While ICR allows LLMs to make listwise relevance judgments using their generative capabilities, it suffers from high computational cost due to the quadratic complexity of attention over long input sequences. To address these problem, ReTuning introduces two key innovations:\n\n1. Structured Sparse Attention: each document token attends only to its own content and shared instructions while query tokens attend to the entire prompt, reducing attention computation from quadratic to linear, significantly improving efficiency.\n\n2. Auxiliary Contrastive Attention Loss: enhance attention from specific query tokens (e.g., “:”) to relevant documents using a contrastive loss.\n\nExperiments on MSMARCO and NQ demonstrate that ReTuned models match or outperform fully fine-tuned LLMs in retrieval effectiveness.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n1. This paper tackles a timely and Important problem that how to adapt the causal attention mechanism in LLMs to better serve as in-context retrievers. This line of work is essential for scaling retrieval-augmented systems and making LLMs more retrieval-aware.\n\n2. The paper is well-written, logically organized, and easy to follow. The observations, motivation, methodology, and experiments are coherently laid out, making the contributions accessible and compelling.\n\n3. The authors conduct detailed attention pattern analysis that reveals structural sparsity and retrieval signals. These observations are well-motivated and directly inform the design of the proposed method. The accompanying ablation studies effectively validate the contribution of each component.\n\n---\n\n### Weaknesses\n\n1. **Limited Model and Dataset Scope**: The experimental evaluation is limited to Mistral-7B, without testing on other widely used LLMs such as LLaMA, Mixtral, or larger-scale models. This raises questions about the generalizability of the proposed method across architectures and scales. In addition, the datasets only use MSMARCO and NQ, which do not provide broad coverage. Including more diverse and challenging benchmarks such as BEIR or MTEB would strengthen the empirical claims.\n\n2. **Outdated or Incomplete Baselines**: Some of the compared baselines are relatively dated or limited in scope. Recent methods that adapt LLMs for retrieval by modifying the attention mask, such as NV-Embed, are not included in the comparison. Incorporating these stronger baselines would provide a more convincing evaluation of the proposed approach's effectiveness.\n\ntypo at line 25: neutral -> neural\n\nQUESTIONS:\n1. Does the reported latency in Figure 4 include the time spent on constructing the structured attention masks used in ReTuning? Could the authors clarify whether this overhead is significant compared to the rest of the forward pass, and what is the time scale?\n\n2. I am confused about the Full-FT setting.\n\n> Our primary comparison is with Full Fine-tuning (Full-FT), where Mistral-7B-v0.3 is fine-tuned on the same ICR training data using standard self-attention and inferred via auto-regressive decoding to generate 290 the relevant document identifier.\n\nDoes this mean that during training, the model uses a bidirectional (BERT-style) attention mask, while during inference it switches to a causal (decoder-style) mask?\n\n3. In Table 2, both Full-FT and ReTuning (full) achieve the same Precision@1 (28.7) under the Decode inference mode. This is somewhat surprising, as I would expect ReTuning (full) to experience some degradation under decoding, due to the training–inference mismatch. Do the authors have any insight into why the performance remains identical?\n\nLIMITATIONS:\nyes, in Supplementary Material", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates how LLMs handle in-context retrieval (ICR) - ranking a list of candidate passages presented directly in the prompt.\nIt finds that fine-tuned models exhibit two patterns: document tokens mostly attend only to their own content, while a few “signal” query tokens sharply attend to the truly relevant document. Leveraging these insights, the authors introduce ReTuning, a fine-tuning strategy that (i) hard-codes a matching sparse-structured attention mask to reduce complexity from quadratic to linear, and (ii) adds a contrastive auxiliary loss that explicitly pushes those signal query tokens to concentrate attention on the correct passage.  At inference time, the learned attention scores can be read off mid-layer to rank passages without autoregressive decoding. Experiments with Mistral-7B on MS MARCO and NQ show ReTuned models match or surpass standard SFT in accuracy while running roughly 2-3× faster during training and up to 4.7× faster at inference as the number of in-context documents grows.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1. The paper provides a valuable analysis of fine-tuned models’ attention mechanisms, revealing patterns such as self-focused attention on document tokens. These insights effectively motivate the proposed structured sparse-attention design.\n\n2. Empirical results show strong performance on standard benchmarks, including MS MARCO and Natural Questions (NQ).\n\n3. Ablation studies are thorough and demonstrate the contribution of each component to the overall performance.\n\nWeaknesses\n\n1. The proposed structured attention mechanism lacks novelty, as similar designs have been explored in prior work. For example, LiT5 (https://arxiv.org/pdf/2312.16098) employs a comparable sparse attention structure for passage reranking.\n\n2. The evaluation is limited to in-distribution settings (e.g., training and testing on MS MARCO). Incorporating out-of-distribution evaluations—such as BEIR or the TREC-DL 2021/2022 test sets (used by LiT5)—would strengthen the paper’s generalization claims.\n\n3. Baseline comparisons are somewhat narrow, omitting recent state-of-the-art rerankers such as RankZephyr. Including these would provide a more comprehensive assessment of the proposed method.\n\nQUESTIONS:\nIn Figure 2, why does the query token almost not attend to other document tokens?\n\nHow would the model perform if a different first-stage retriever were used?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the problem of in-context retrieval, where the query and the documents are directly fed into LLMs to generate a ranked list of documents. The authors first observe interesting findings that the attention of each document tends to focus primarily on itself, and some special token of the query, like \":\", always serve as the retrieval role by interacting with all document tokens. Based on these findings, they introduce ReTuning by (1) employing a sparse attention structure; (2) applying a tailored contrastive learning loss to the special query token by treating the attention scores between this token and each document as similarities. Experiments on MSMARCO and NQ demonstrate the effectiveness of ReTuning. Further ablation shows the effectiveness of each component.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n1. In-context retrieval is an important problem with significant practical implications, as it enables retrieval to be more directly tailored to the user’s instruction by processing all documents simultaneously within an LLM.\n\n2. The findings that the attention of each document tends to focus primarily on itself, and some special token of the query, like \":\", always serves as the retrieval role, are both interesting and new.\n\n3. Detailed ablations are conducted to demonstrate the effectiveness of each component.\n\n4. The paper is well-written, although the structure could be improved.\n\n**Weakness:**\n\n1. Unrealistic experiment settings. The experiments rerank only 30 documents for NQ, which is substantially different from prior work [1], which conducts ICR over hundreds/thousands of documents.\n\n2. Missing baselines. The authors claim the experiments are second-stage reranking, while most of the baselines are still first-stage dense retrievers. Moreover, the rerankers compared are somewhat old and weak, and there are strong rerankers (pointwise/listwise) like RankVicuna and FIRST available.\n\n3. The metric of precison@1 for NQ is also not a realistic metric for retrieval tasks. Typically, Top-5, 20, and 100 accuracy are reported for NQ.\n\n4. The connection between ReTuning and dense retrievers (e.g., ColBERT) is unclear. As each document only attends to itself and the instruction, these documents can be seen as the same as the independent encoding process of dense retrievers. Moreover, the retrieval step is simply calculating the similarity scores between the query and each document via attention scores, which is also quite similar to ColBERT.\n\n5. The structure of the paper can be improved by reducing frequent references such as “see Table 5”, as these hurt the reading flow. \n\n[1] Can Long-Context Language Models Subsume Retrieval, RAG, SQL, and More? Arxiv. 2024.\n\nQUESTIONS:\n(1) Stricly enforcing the attention across documents may restrict the intra-document interactions, would this lead to performance drop?\n\n(2) Which middle layer is selected?\n\n(3) Where are the latency results in Table 3?\n\n(4) Many citations, e.g., [11-15], are without their proceedings information.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}