Scalable In-context Ranking with Generative Models
Authors: Nilesh Gupta
Chong You
Srinadh Bhojanapalli
Sanjiv Kumar
Inderjit S Dhillon
Felix X. Yu
Venue: NeurIPS 2025 poster ()
Keywords: large language model
information retrieval
efficient fine-tuning
efficient inference
Primary area: deep_learning
TL;DR: This work introduces BlockRank, a method designed to make LLMs efficient and scalable for in-context ranking by aligning their internal attention and inference mechanisms with the structure of the in-context ranking task.

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.

Scalable In-context Ranking with Generative Models
Nilesh Gupta'
Chong You
Srinadh Bhojanapalli Sanjiv Kumar
Inderjit Dhillon
1.2 Felix Yu?
University Texas at Austin Google Google Deepmind
com httpa///..oooolleee//////ooaaaaa
Abstract
In-context Ranking (ICR) is an emerging paradigm for Information Retrieval (IR), which leverages contextual understanding LLMs by directly incorporating the task description, candidate documents, and the query into the model's input prompt and tasking the I.I.M to identify relevant document(s). While is effective, effi- ciency is significant challenge in this paradigm, especially 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 doc- uments in the context; and (2) sucryeddcument block relevance the attention scores from certain query tokens document block in middle layers strongly correlate with that document's actual relevance. Motivated by these observations, introduce BlockRank (Blockwise In-context Ranking), novel method that adapts the attention operation an LLM by (a) architecturally enforcing the ob- served inter-document block sparsity, reducing attention complexity from quadratic to linear without loss in performance, and (b) optimizing query-document block relevance for relevant documents during fine-tuning using auxiliary con- trastive training objective, improving retrieval in attention. Experiments on BEIR, MSMarco and NQ with Mistral-7F demonstrate that BlockRank Mistral matches or outperforms existing SOTA listwise rankers and controlled fine-tuned base- line while being significantly more efficient at inference (4.7x for 100 MSMarco documents in context) and scaling gracefully to long-conted shortlists around 500 documents in-context (~ 100K context length) within a second, presenting a scalable and effective solution for ICR.
Introduction
Information retrieval (IR) is the problem of finding relevant content from a large document corpora. While sparse retrieval methods based word-level matching have existed for decades [Robertson et al. 2009 Formal et al.] (2021), modern IR systems increasingly leverage deep neutral network based representations, which achieve their success through superior ability capture deep semantic relationships [Karpukhin et al. 2020). Recently, generative large language models (LLMs) (Team) al. 2023 Achiam et al. 2023) have emerged revolutionary paradigm that transforms many sub-fields of machine learning, including IR. Through pre-training on the web, LLMs absorb an enormous amount of world knowledge and demonstrate remarkable capabilities in dialogue, question answering, reasoning, and beyond [Wei et al. 2022).
The powerful capabilities of LLM open up novel approaches for IR well. One emerging paradigm the In-context Ranking (ICR) [Lee et al. 2024 Ma et al., 2023), which directly leverages an LLM's contextual understanding capabilities. In this setup, query and list of candidate documents
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

SERAA VROIINN SPETIO
Oviy
Figure 1: Analysis of attention patterns in Mistral-7B performing In-context Ranking (ICR) MSMarco. (left) Attention averaged over middle layers 16-21 reveals structural sparsity - a strong diagonal (intra-document attention needed for local context processing) and significant attention to the first row (focus on the query tokens to document segments. Certain tokens (the last token, "!") attend primarily to the relevant document only (i.e., Doc24, highlighted in green). (right) Attention from final query tokens across layers shows retrieval signals strengthening in middle layers. These patterns motivate our BlockRank approach.
instruction). (middle) Attention in Layer 18 from individual
formatted together within the LLM's input prompt (see Figure B). tasking the model identify the most relevant document(s), often through the generative decoding process. ICR holds the promise of considering the query and all candidates simultaneously while performing relevance judgements.
Despite this promise, LLM-based ICR introduces significant efficiency challenges. As the num- ber of candidate documents increases, the input context length grows rapidly, making inference computationally expensive and memory-intensive, due to quadratic/super-I
complexity of the attention mechanism. Current methods (Lee et al. 2024 Sun et al. 2023 Pradeep et al. 2023a b) typically treat the LLM as a black-box or do not fully utilize the structure of the ICR task i.e. the input prompt is composed of sequence of potentially independent candidate documents conditioned on a shared instruction prompt. Moreover, as we discuss in Section 5.3 and Section [DI] of Appendix, auto-regressive decoding is not best suited for this task when decoding multiple predictions from the fine-tuned model (see T))le(4.
Paper Contributions. To this end, we first investigate how standard LLMs process information within the specific task structure. We conduct an analysis of the attention patterns of fine-tunay Mistral-7B model when prompted on ICR examples derived from MSMarco dataset (see Section 3 for details and Figure for visualizations). This analysis reveals two structural properties: (1) inter-document block sparsity most document tokens focusing locally (primarily within their own document, instructions, or or other documents), rather than attending densely across all candidate documents. (2) query-document block relevance - similar to the findings Wu et al. (222)). Chen al. (2025), find that last and some specific query tokens like "." (that signal start of the potential document generation process) develop strong attention weights towards relevant document tokens, particularly in the model's 's middle layers.
Building up on these insights, we propose BlockRank (Blockwise In-context Ranking), an efficient in-context ranking method. BlockRank introduces standard LLM architecture and fine- toning: (1) architecturally, it imposes a structured sparse attention in which document tokens attend only causally to their content and shared instruction tokens, reducing attention complexity from quadratic to linear; and (2) it incorporates a contrastive learning objective explicitly optimize internal attention from signal-carrying query tokens toward relevant documents, which helps the BlockRank model in two fronts: (a) attend strongly to the relevant document in context, improving retrieval quality (see Table B): (b) the ability reliably infer relevance based on attention concentration during the prefill stage, leading to further speedups in inference compared to iterative decoding (see Section 4.3) To summarize, the main contributions of this work are:
modifications (visualized in Figure 2 to

---

Standard Fins luning
Blookaant Fine suriing
Unanbedd up lyve
NTP
Put Prompt
Prompt regmentation Chunking
Leaar
Atlention Ralevance Scares
Youtt paen a int. 10:1
PMD 1:1
comblow lnfo Fi FeolAreAnn
O(angin))
Struttred Abtention
Figure 2: BlockRank starts with chunking the full prompt into segments and then processes it using structured attention, where the documents only attend to themselves and the instruction segment, while the query segment attends to the full prompt. also incorporates an auxiliary attention loss (feix) from a middle layer (I") that increases sharpness of attention on the relevant documents and enables an alternate inference mechanism using attention scores derived from l*.
analysis that characterizes attention patterns in LLMs fine-tuned for ICR, identifying key sparsity structures and latent retrieval signal carriers (specific query tokens middle layers).
an efficient approach BlockRank for In-context Ranking that enforces a structured sparse attention and contrastive training objective internal attention.
extensive experiments on standard retrieval benchmarks (BEIR, MSMarco and NQ) demonstrating that BlockRank achieve strong ICR performance, matching or outperforming strong baselines as well as full fine-tuned model (see Table D 2 while being order of magnitude efficient at inference (see Figure D.
The remainder of this paper is organized as follows: Section describes the problem setup and discusses related work. Section details our analysis of LLM attention in ICR. Section presents the BlockRank methodology. Section 5 reports experimental results, and Section 6 concludes the paper.
2 Problem Setup and Related Work
This section formally defines the ICR task addressed in this paper. We also review relevant prior work that uses LLM for IR, and position our BlockRank method in the context of the literature.
2.1 Problem Formulation: In-context Ranking
{ d and a query q, the goal of IR is to return subset of DD that are relevant to q. In this paper, we consider documents and queries in the form of text that can be parsed by an LLM, though the discussion may also apply visual and audio data when the LLM is multi-modal. Following standard practices [Lee et al., [024], we define an ICR prompt as a composition of the list of documents together with the query as following:
Given a collection of n documents
prompt(q, D) = "(Insr). {di).... .... (d))-
In practice, processing the entire corpus DD (where n can be millions) within a single prompt is infeasible due LLM context length limitations. Therefore, the ICR task we consider in this paper operates on a smaller candidate list D, (d1,.. dx) C D, where N is the number of candidates (N O(100) in our experiments) retrieved by a first-stage retrieval model (c.g. dual-cncoder). The prompt in (T) is thus applied to this candidate list Dy. Furthermore, each document representation { within the prompt often includes structured formatting beyond just the raw text c, such as its unique identifier id, We adopt the format from [Lee et al. [202]] also illustrated in Figure 3 explicitly demarcating document start, content, and end, along with identifiers (e.g., ID: id, CONTENT: C; END ID: id,).
Inst is a description of the retrieval task and can also include the query 4. While excluding the query from the instruction prefix is desirable from an efficiency standpoint as it would allow for the mnery-independent representations of documents to be pre-processed and cached offline - we find this leads to noticable drop in performance in our experiments (see Table 6) We hypothesize.
3

---

You will be given a query and a 1ist of docunents. Each document will be fornatted as ID; <id> I CONTENT: <content>> | END ID: <id>. You need read carefully and understand all of then. The query in: shich clansification group contains the nost organiems, and your geal in to find all (acument(s) from the list that can help ansver the query.
CONTENT; Thie diverse group orgenisne. includee plants, sninsls..
EXD
Documents in be-betveen onitted for brevity] ...
ID: CONTENT: Organians conpoaed of sukaryotic colla are divided into main.. EXD ID: N
start!
Which document in most relevant to answer the query? Print out the ID of the document. Query: which classification group contains most organisns. folloving docunent (s) can help ansver the query:
Final Anaver: [200]]
Figure 3: Example structure of the prompt template used in experiments, showing query- -ased instruction, abbreviated document list, and the final query section.
including the query Inst allows the model to condition each document's representation on the specific information need from the outset, enabling it to better focus on query-relevant facts and signals within each document during processing. We note that existing listwise LLM re-rankers (Sun) et al. 2023 Pradeep et al. 2023b) also apply similar formatting where the query appears before the documents. Our preliminary experiments show that one can replace the query with a similar-looking document from the corpus, suggesting that future work can potentially cxplore conditioning document representations within clusters to alleviate the need for query-dependent processing.
The objective of In-context Ranking is then formally defined as follows: given the prompt D₂) constructed from the query q and the candidate list Dy. train or utilize an LLM fo. with @ being the model weight, effectively identify and output the identifiers id' corresponding to the d' E D, deemed relevant to q. Typically it is achieved by predicting a ranked permutation of Dr, and taking the top & elements. The central challenge addressed in this paper is to develop methods to train fø and perform this prediction both effectively (high accuracy) and efficiently (low computational cost), particularly the candidate list size N increases.
2.2 Related Work
Our work builds upon research in neural retrieval, the rapidly evolving field of using LLMS for IR, and efficient attention mechanisms.
Neural Re-ranking and Retrieval Models. Prior LLMs, neural information retrieval saw significant progress with methods like dense dual-encoder retrieval (e.g., DPR (Karpukhin et al.) 2020), ANCE [Xiong et al., 22200) offering efficient first-stage filtering, and cross-encoder mod- els (e.g., monoBERT [Nogueira and Cho 2020), mono'T5 (Nogueira et al. 2020 providing high re-ranking effectiveness through deep query-document interaction. Late interaction models like ColBERTv2 [Santhanam et al. 2022) aimed to balance the trade-off between dual and cross encoders. Our work is architecturally distinct from traditional neural IR methods as it operates within the in-context ranking paradigm, where single LLM processes the query and the entire candidate list simultaneously in one context window allowing full contextualization (query and output representa- tions conditioned the full set of candidate documents) and complex instruction instruction-followingg following (e.g.. "Find documents that disagree with...").
LLMs as Listwise Re-rankers. The ability of LLMs to process and reason over long contexts spurred their application to listwise re-ranking [Ma et al. 2023) (or In-context Ranking). where multiple candidates are processed simultaneously. Initial successes often involved prompting large proprietary models like GPT-3.5/4 (Sun et al., 2023 in zero-shot or few-shot settings. While effective, these approaches typically incur high computational costs and often rely on auto-regressive generation to output rankings or relevance scores, adding latency. More recent work focuses on adapting open- source LLMs (c.g., Llama, Mistral, Zephyr, Vicuna) for this task [Pradeep et al.) 2023b
Zhang

---

et al. 2023 and improving efficiency, for instance using Seq2Seq architectures Tambeer et al. 2023
and using single-token decoding [Reddy al. [024). Recent papers have also shown the existence of retrieval heads (attention heads that carry strong retrieval signals) in many modern LLMs (Wu al. 2024) and their usefulness in inferring retrieval signals [Chen al. [025]. Our work differs from these methods by introducing effective task specific restructuring of the attention architecture for efficiency and an explicit fine-tuning objective to directly train the model's 's attention patterns for
the ranking task.
In-context Retrieval / Ranking Lee et al.
2024 studied the In-context Retrieval (ICR) paradigm for various frontier LLMss, demonstrating that long-context models can match the performance of specialized retrieval systems when processing corpora of up to few thousand documents. Our work builds on this paradigm but addresses a more challenging setting. The evaluation in that study is performed on random subset of documents from the full corpus. In contrast, our experiments focus on ranking the top-k hard candidates returned by a strong first-stage retriever i.e. In-context Ranking task. This task is arguably more difficult, as the model must distinguish between many semantically similar documents to identify the correct answer. We argue that this hard negative setting is a more faithful simulation of the practical application of LLMS in retrieval pipelines. Moreover, processing long lists of documents within the L.MM context remains challenging, with studies highlighting difficulties in effectively utilizing long-range information [Goldman et al. 2024
Efficient and Structured Attention. The quadratic complexity of the standard self-attention has spurred extensive research into more efficient attention approximations. Many successful approaches enforce structured sparsity the attention matrix, reducing complexity from quadratic to sub- linear. Notable examples include methods based on sliding windows (e.g., Longformer (Beltagy) et al. 2020]), global-local patterns (e.g., BigBird (Zaheer al. 2020 and other block-wise structures. While these methods are designed for general long-context processing, BlockRank structured attention can be seen as a task- task-specific instance of this paradigm. BlockRank's sparsity is semantically informed by the logical structure of the in-context ranking task itself sepear-spaarating instructions, documents, and the query. This content-awe aware structuring allows for a highly efficient architecture tailored for ICR.
3 Emergence of Structured Attention in In-context Ranking
Before introducing our method, we first analyze the attention mechanisms of a standard LLM when performing the ICR task defined in Section 2.1 The analysis below is anchored on Figure [ which is based on a random sample, we provide more results in the Appendix Section D
Analysis Setup. We conduct our analysis using a Mistral-7B-ve model [Jiang et al. 2023 fine- tuned on the ICR task with data derived from MSMarco (as described in Section 5.1). In particular, our fine-tuning objective is the standard Next Token Prediction (NTP) loss, without any modifications. We feed this model prompt(g, g, Dz). Let the resulting input token sequence be T = .t))
Our analysis focuses on the attention probabilities computed within the transformer layers. Given a layer index E (1...., ] where Lmode! is the total number of layers, and an attention head index h we denote the attention probability from a query token t; to a key/value token i Oij
(L.k)
(L.N) We We often consider the attention averaged across all H heads in layer: a и E Oij
(I)
examine interactions between different types of tokens by partitioning the token indices {1...., L} into sets corresponding to the instructions (Tinst), the query (T,), and each document (Tá, for k {1.....N] We visualize these interactions using heatmaps, with representative examples shown in Figure !
Observation 1: Inter-document Block Sparsity Our first key observation is that the attention patterns exhibited by document tokens are structured and sparse, rather than uniformly dense. This is clearly visible in Figure [a), which shows the segment-wise attention in the middle layers. The heatmap is dominated by the diagonal, indicating strong intra-docuss documed attention: for token f, E (1) Ta. the sum of attention probabilities towards other tokens within the same document, Σr,eTa, dij is significantly higher than attention towards other parts of the context.
This observed structured sparsity implies that computing full attention matrix might be largely redundant for this task. A significant portion of the computation could potentially be saved by
5

---

enforcing an attention pattern that focuses on local (intra-document) and instructional context, directly motivating the structured sparse attention employed in BlockRank.
Observation 2: Query-document Block Relevance Our second key observation is that certain tokens within the query 7 attends primarily to relevant documents only, particularly in the middle layers of the transformer.
Figure bb) illustrates this Layer 18. It maps the attention from individual query tokens (X-axis) to the different document segments (y-axis), We observe that certain tokens, such as delimiters (":²) and end of prompt tokens, exhibit distinet sharp attention distributions. These specific "signal carrier" tokens attend more strongly towards the segment corresponding to the ground -roundd-tuuth relevant document d' (i.e., Doc24, highlighted in the figure) compared to irrelevant documents dg (k *). Formally, let A LiyeTT, aij be the total attention from query token I; to document d2 at layer l. For
(I)
specific t, T. identified as signal carriers and middle layers L, we observe
A for most h: * x. We hypothesize that such structural tokens carry strong retrieval signals as they often precede Ta- (by design during fine-tuning but also during pre-training), hence their attention gets biased towards the in-context d" segment in order to predict the succeeding T tokens.
Furthermore, the layer depth plays critical role in the emergence of these signals. Figure (c) tracks the attention A d- from final query tokens i; to all document segments d₂ across all layers I E {1,..., Lmodet): }. The plot shows that the discriminative signal weak in the initial layers, emerges and strengthens significantly in the middle layers (approximately layers 8-24), and persists or slightly diffuses in the final layers.
4 BlockRank: Blockwise In-context Ranking
Motivated by the attention analysis presented in Section 3 we propose Blockwise In-context Rank- ing (BlockRank), an efficient in-context ranking method. BlockRank comprises of following components (see Figure 2): a structured attention mechanism enforcing sparsity, an auxiliary at- tention loss to enhance retrieval signals in attention operation, and an alternative iieettionbbsse inference method. We detail each component below.
4.1 Blockwise Structured Attention
The core of BlockRank's efficiency during fine-tuning and inference stems from restructuring of attention mechanism designed to enforce the sparse patterns observed in Section
Enforcing inter-document block sparsity. we modify attention operation such that:
Document Tokens (t: E Ta, for k fI...., N)) only attend to tokens within their own document chunk (t, E Tap) and tokens within the instruction chunk (t, E Tinst).
Query Tokens (4: E T.): attend to all tokens in the prompt (t, E T = UpTi) to gather context for identifying the relevant document(s).
Instruction Tokens (t, E Imat): attend causally within the instruction segment itself.
Instead of constructing large, explicit sparse attention masks, we implement this structured attention efficiently using the chunked representation defined as follows: the long prompt is first segmented into its logical components Sn Inst, S₂ = dy for & (1,..., NN. and Sx+1 g. Each segment Se is then processed (via standard sequence length chunking or padding) to form fixed-length chunks, typically of length Lekunn tokens. Let the token sequence corresponding to chunk S₁ be Tr CT= (41....,4z) lz), where 7' is the (potentially virtual) concatenation of all chunk sequences.
Each chunk S₁ can be processed largely in parallel (c.g., distributed along the batch dimension). Let Q!9 k((!)) ((()) be the query, key, and value matrices for chunk S₁ at layer The attention output for a token 4; in chunk S₁ is computed as follows:
The attention output is computed using self-attention within the chunk and cross-attention only to the keys and values from the instruction chunk: Attention( (1) [K(()) k (t y()) vlla)) Attention to other document chunks Smyt and the
S& is a document chunk (k E {I....,NV N):
query chunk Sy is effectively masked out.

---

If Sy. is the query chunk (k = N + 1): The attention output is computed using self- attention within the chunk and cross-attention to the keys and values from all other chunks: Attention(Qa) (Q)) [Kg] (1) K(t)
Instr K()
( dy [Vg] v( v(t
V2 Insts dj 1.... dy
Instruction chunk attention (k 0) standard causal self-attention.
This computes only the necessary attention scores, drastically reducing the computational cost, con- verting quadratic attention to linear. Please see Appendix Section C for more details and complexity analysis.
Permutation-invarian Position Embedding. To complement the structured attention, we employ specialized position embedding that reinforces the logical separation of the prompt's components. This also helps the model learn position-invariant representations for documents [Tang et al.) 2023] and distinguish the query's unique role. Specifically, tokens in the instruction segment (T are assigned standard sequential positions starting from 0. For all document segments (Td,), ). we use a shared local position space. Each document's tokens assigned positions beginning immediately after the instruction segment, as if were the only document present. For example, if the instruction has length Inst, the first token of every document dg is assigned the position Linst- This encourages the model to apply a consistent, order-invariant function to each document, mitigating any bias from its absolute position in the candidate list. Finally, distinctly separate the query from the document corpus, its tokens (Tg) are assigned positions starting from a large, fixed offset. In our experiments, we use an offset of 8192, the query tokens receive positions [8192, 8193,... This large gap ensures that the relative positional encodings between any query token and any document token are significantly different.
4.2 Auxiliary Attention Loss (Caux)
To explicitly optimize query-document block relevance for relevant documents during fine-tuning, we introduce an auxiliary loss Laux applied a specific middle layer !" (determined empirically, see Section D.3 in the Appendix). This loss encourages "signal-carrier" query tokens to attend more strongly to the relevant document.
More specifically, let sigggga C Tg be the set of indices for the identified signal-caryyig query tokens. Based
prompt template and empirical analysis we set T sigiignl = [":", "[""]. J. Let
Tdocs = U Taş be the set of indices for all tokens belonging to any document segment. For each signal token 1 E sigiin at layer I", we compute attention scores towards document tokens ty E Taocs as following: 1. Obtain query vectors (((") for tú E Tgg.iggal and key vectors Ky"" /ddz for all t, E Tooos 3. Compute normalized attention probabilities only over the document tokens: d'y softmax, (=tg) where the softmax is computed across all j such that t, E Tdocs. This normalization focuses the probability mass exclusively on the candidate documents, ignoring instructions and query tokens. 4. Aggregate these probabilities to compute an attention mass score for each document d3: S(g, dy.) EneT, ETa, o'y (Alternative: could use mean aggregation over f,). This score S(q, de) quantifies the relevance signal from the carrier tokens towards document dg. 5. Apply a contrastive loss using these scores. We use the InfoNCE loss with temperature T:
for t, E Tdocs 2. Compute raw attention logits 2yj = (x("")///Vdd
Laux LinfONCE(S d²), (S(g, du dz))ss=i7) = - log exp(S(q, dz)T)
(2)
Ek exp(S(g, dx)/7)
where d' the ground -tuth relevant document. This loss encourages the score d() for the relevant document to be higher than scores for irrelevant documents.
Overall Training Objective. The BlockRank model is fine-tuned by minimizing combined loss function that includes both the standard next-token prediction objective and our auxiliary attention loss:
Crotal = CNTP + XLax
Here, LNTP the cross-entropy loss calculated on the answer tokens (similar to standard instruction tuning) based on the model's prediction of the next token in the sequence, computed using the
7

---

final hidden states which are generated respecting the structured attention masks defined in Sec- tion 4.1] Laux is the auxiliary InfoNCE loss defined in Equation 2 applied only at layer 1" À is hyperparameter balancing the two losses (we use À = 0.1 in our experiments).
4.3 Efficient Beention-Based Inference
An advantage of BlockRank is that the auxiliary loss explicitly optimizes the attention scores S(q, de) to reflect relevance. This allows for an alternate efficient inference mechanism during the prefill stage of the context processing. can bypass the iterative auto-regressive decoding process, and even the full forward pass (depending on the choice of (*). The inference mechanism can be defined as follows: 1. Given a prompt(qq D). perform a partial forward pass of the BlockRank model up to the target middle layer (". 2. Compute the document relevance scores S(g, de) for all candidate documents & E {I...., N}} using the exact same procedure as described for the auxiliary loss calculation (Section 422 steps 1-4), utilizing the signal carrier tokens T g,iignal and performing the
softmax over document tokens Tdoes only. 3. Identify the index ie of the document with the highest score: & = arg max; Sig, dx). 4. Output the corresponding document identifier idg, for top-K predictions output arg (g, dx)
5 Experimental Results
This section empirically evaluates the proposed BlockRank method. We conduct two sets of experiments: first, an evaluation on the BEIR benchmark to assess zero-shot generalization against state-of-the-art re-rankers, and second, controlled in-domain evaluation to analyze effectiveness, efficiency, and scalability. We aim to answer the following research questions: (RQ1) How does BlockRank compare against strong baselines in terms of retrieval effectiveness, both in zero-shot generalization and in-domain settings? (RQ2) What are the efficiency benefits of BlockRank compared standard fine-tuning, particularly when scaling the number of in-context documents? (RQ3) What is the contribution of BlockRank's core components (structured sparse attention, auxiliary attention loss, and attention-base -based inference) to its overall performance?
5.1
Experimental Setup
Goal & Task. Given a query and a list of candidate documents retrieved by an initial, potentially weaker retriever, the goal is to identify the most relevant document(s) from within that list by processing the entire list in the L.M''s context.
Datasets & Formatting. For assessing zero-shot generalization, we use 11 diverse datasets from the BEIR benchmark (Takur al. 2021 replicating Table in Reedy al. [0024] In this setting,
the task is to rerank the top- 100 documents provided by Contriever Izacard et al. 2021 model.
For in-domain analysis, we two standard passage retrieval benchmarks: MSMarco Passage
Ranking Bajaj et al. 2018) and Natural Questions (NQ) [Kwiatkowski et al. 2019). During training,
we construct candidate lists for each query by retrieving an initial of 30 passages using a pre- trained sentence transformer model with teacher-forcing (i.e. always adding ground-truth documents). This list is then formatted into the prompt structure shown in Figure During in-domain evaluation, we construct lists of varying sizes (N 10 to 500) to test scalability. More details can be found in Appendix B
Evaluation. We evaluate model performance on two primary aspects: effectiveness and efficiency, For BEIR, effectiveness is measured using nDCG@ 10. For in-domain experiments on MSMarco and NQ, we report Precision@ @1 and Mean Reciprocal Rank @ 10 (MRR @ 10). Efficiency quantified by Inference Latency, the end-to-end wall-clock time per query.
Baselines. We compare BlockRank against a comprehensive set of baselines tailored to each experimental setting. For the BEIR generalization benchmark, we compare against contempo- rary listwise re-rankers, including a strong cross-encoder, Rank Vicuna [Pradeep et al.) 2023a]. RankZephyr (Pradeep et al. 2023b), and the recent state-of-the-a 2024]. For in-domain analysis, our primary comparison is with Full Fine-tuning (Full-FT) (full causal attention with only NTP loss) of the same base model and the same training data. We also include results from zero-shot LLMs Baataal-B-Instrrre Gemini-2.00 2.0.fah.. For broader context, we include
model, FIRST [Reddy al.
8

---

Table I: nDCG@ 10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model. Bold indicates the best numbers.
Rerunker
Train Data
Avg.
Climate- FEVER
DB- Pedia
FEVER
FIQA
Hotpot QA
MS Marco
NF- Corpus
NO
Sci- docs
Sci- fact
COVID
Tree-
None (Contriever)
MS Marco
15.9
23.7
41.3
75.8
32.9
63.8
10.7
32.8
-99.8
16.5
67.7
59.6
Marco
81.9
47.0
Rank Vicuna
GPT 3.5
50.7
28.2
50.0
81.0
35.9
73.5
36.7
33.1
58.6
18.4
70.5
71.3
Rank Zepbyr
GPT 3,5
53.7
25.6
50.0
80.1
42.2
71.6
42.7
37.7
65.6
20.5
76.7
78.4
FIRST
GPT-4
54.3
26.7
50.9
81.7
74.2
44.4
37.4
66.4
20.4
74.6
78.8
BlockRank
Marco
Traditional Retrieval Models such as the lexical baseline BM25, the dense retriever GTR Ni et al. 2021]. ColBERTv2 [Santhanam et al. 2022]. and best performing Sentence Transformer Encoders specific to each dataset
for MSMarco and 411-MiniIM-L12-W
for NQ). Furthermore, we consider pairwise cross-encoder baselines like monoBERT [Nouueira and
Cho 2020 and improved versions of monoT5
[Nogueira et al. 2020
Implementation Details. BlockRank and the Full-FT baseline utilize Mistral-7B-vOl as the base model. For fine-tuning both models, we employ the Adafactor optimizer [Shazeer and Stern 2018) with a learning rate of x 10-7 and global batch size of 32 (accumulated across replicas). Each model is trained for 1 epoch with a linear warmup followed by cosine decay. For BlockRank, the auxiliary loss weight
set to 0.1, and T set to 0.05. Unless stated otherwise, BlockRank results employ the proposed attention-basee -aaedd inference. Decoding based experiments with BlockRank, LLM baselines (Full-FT Mistral and Zero-Shot LLMs) utilize greedy decoding to generate the relevant document identifier(s): to get multiple predictions (for MRR @ 10 evaluation) we use constrained beam decoding with beam-size set to 10, where only valid outputs are generated. All LLM fine-tuning and inference experiments were conducted using JAX Google Cloud TPUs (specifically, 8 chip v6e configuration), and reported efficiency metrics correspond to this setup as well.
5.2 Main Performance Comparison
Generalization to Diverse Tasks (RQ1) The results in Table [ show that MSMarco-trained BlockRank Mistral (54.8) outperforms FIRST (54.3), RankZephyr (53.7), and demonstrating strong out-of-distributee generalization. Importantly, BlockRank achieves strong results with the significant efficiency gains (Figure 4 presenting compelling combination of effectiveness and scalability. Furthermore, it gets the strong performance by processing the entire list of 100 candidate documents in single forward pass instead of multiple sliding-window forward passes over the candidate set which is required for other listwise ranking models. These results also indicate that BlockRank is not sensitive to the first-stage retriever, as it effectively ranks candidates from Contriever despite its training data being constructed with different retrieval model.
(50.7),
In-Domain Performance (RQI) Table summarizes the quality comparisons on the NQ and MSMarco in a controlled environment where both BlockRank and Full-FT Mistral are trained on the same training data and evaluated on in-domain data, for broader comparison we also provide results for additional baselines trained on the same data. Our proposed BlockRank consistently outperforms its direct counterpart, Full-FT Mistral (7B).
Scalability (RQ2). Figure underscores BlockRank's substantial inference efficiency advantage over the Full- FT baseline as the number of in-context documents (N) increases. BlockRank model consistently exhibits lower latency; at N 100, it is approximately 4.7x faster. More critically, its latency scales linearly with N, reach- ing 1.15s at N = 500. Furthermore, BlockRank model maintains P@ (peaking around 29.2% for N = 200 and remaining at 28.7% for N 500), whereas Full-FT's P@ sharply degrades beyond N 100 (dropping to 26.7% at N 500).
Figure 4: P@ and Latency (annotated) of BlockRank vs Full-FT Mistral, scaling N on MSMarco.
20.5
20.5
26.5
374
Pul.77
1/
Number of In-contso Documents (N)

---

Table 2: Comparison on MSMarco and NQ datasets in controlled settings. Encoder methods are evaluated on the full corpus while the rest of the baselines are evaluated on shortlist. Best results are highlighted in Bold.
NQ
MSMarco
Category
Method
Model Size
Precision
Precision@)
MRR @
Sparse Retrieval
18.4
Sentence transformer [Reimers and Gurevych)
2019
66M
58.8
24.8
37.2
Fneo-FFn
GTR-XXI. Ni 2021
4.8B
38.8
ColBERTvZ
2022
110M
39.7
nonoBERT
Nogueira and Cho
2020
110M
38.2
Cross Encoder
Noguetaa al.
Zero-Shot LLM
Mistrll7B-vv0.--t Jiang et
2023
7B
43.5
13.1
Gemini-2.0-fan Team et al.
2025
65.1
16.9
Fine-tuned LLM
Full-FT Mistral
75.5
28.7
38.3
BlockRank Mistral (Ours)
7B
76.2
29.1
42.0
5.3
Ablation Studies
Table 3: Impact of training loss on Attention- based (Att)) and Decoding (Decode) Inference.
To understand contribution of components of BlockRank (RQ3), we perform several ablation experiments, primarily the MSMarco dataset with N 50. More ablation is provided in Section of Appendix.
Precision @
Training Configuration
Decode
Attn
Full-FT
28.7
27.6
Full-FT (w/ aux)
28.7
28.1
Impact of Training Loss Table abbates the contribu- tions of CNTP and Louz to P@ evaluated with both auto-regressive and attention- baased inference. Introduc- ing Lauæ consistently enhances performance for attention-
BlockRank (wlo nip)
15.8
28.6
BlockRank (w/o aux)
28.4
27.8
BlockRank (full)
28.7
29.1
based inference, and for BlockRank (which incorporates structured attention), it increases from 27.8 to 29.1. As expected, the CNTP objective crucial for generative decoding performance, as seen by the sharp drop in Decode Prec@ for 'BlockRank (w/o ntp)' to 15.8. Notably, the full BlockRank configuration achieves the highest Attn Prec@1 (29.1), demonstrating that Lawz effectively optimizes attention scores for direct retrieval, making attention-based inference the preferred mode for method.
Impact of Inference Method Table ablates the in- ference, comparing decoding against our attention- based approach P@ and MRR@ @ 10; for Full-FT, Decode MRR@ @ 10 uses a beam size of 10. The results show that while auto-regressive decoding yields comparable P@1 for both Full-FT (28.7) and BlockRank (28.7) models, it significantly less effective producing strong ranked list for MRR @ 10. In contrast, BlockRank with attention- based inference performs best, achieving notably better MRR @ 10 (42.0). BlockRank's attention-basse inference, optimized via its auxiliary loss, is more calibrated at assigning relevance across multiple predictions.
Table 4: Ablation: Inference Method Effectiveness & Latency (MSMarco, N=50).
Model
Inference Method
P@1
MRR@10
FuILFT
Decode
28.7
38.4
Full-FT
BlockRank BlockRank
Decode
28.7
40.0
Attn
29.1
42,0
Conclusion
This work addresses the efficiency challenge in In-Context Retrieval (ICR) by analyzing LLM atten- tion, identifying structured sparsity and query-token retrieval signals. We introduced BlockRank, method that enforces this task-specific sparsity for linear complexity and uses a contrastive auxiliary loss directly optimize these internal attention signals for relevance. Experiments on MSMareo and NQ show BlockRank (Mistral-7B) matches or surpasses standard fine-tuning effectiveness while be- ing significantly more efficient at inference and training. This offers scalable and effective approach for M-b-ed ICR. However, we acknowledge our current findings are primarily demonstrated on specific model architecture, and the robustness of the learned attention signals for direct inference across highly diverse tasks needs more investigation.
10

---

References
Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Bajaj, D. Campos, N. Craswell, L. Deng. J. Gao, X. Liu, R. Majumder, A. McNamara, B. Mitra, T. Nguyen, M. Rosenberg, X. Song, A. Stoica, S. Tiwary, and T. Wang. Ms marco: A human generated machine reading comprehension dataset, 2018. URL http::/t:///// /tt///xxvvvv.. http//axvvvv.aaaaaaa// wwr.vv.orggbbb//6111 09268
I. Beltagy, M. E. Peters, and Cohan. Longformer: The long-document transformer. arXiv preprint 0:Xiv:2004.05/550 2020.
S. Chen, B. Gutiérrez, and Y. Su. Attention in large language models yields efficient zero-shot re-rankers, 2025, https://: org/ Saarg/abs//00 /4100. .22622
T. Formal, B. Piwowarski, and S. Clinchant. Splade: Sparse lexical and expansion model for first stage ranking. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 2288-2292, 2021.
O. Goldman, Jacovi, A. Slobodkin, A. Maimon, Dagan, and Tsarfaty. Is it really long context if all you need is retrieval? towards genuinely difficult long context nlp, 2024. URL https://arre
hwwggabs/24000 00402
G. Izacard, M. Caron, L. Hosseini, S. Riedel, P. Bojanowski, A. Joulin, and E. Grave. Unsupervised dense information retrieval with contrastive learning. arXiv preprint 1rXiv:2/12.091 2021.
A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. d. I. Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, et al. Mistral 7b. arXiv preprint 2xiv:2310006822 2023.
V. Karpukhin, B. Oguz, S. Min, P. Lewis, L. Wu, S. Edunov, D. Chen, and W. T. Yih. Dense passage retrieval for open-domain question answering. In 2020 Conference Empirical Methods in Natural Language Processing, EMNLP 2020, pages 6769-6781. Association for Computational Linguistics (ACL), 2020.
K Kwiatkowski, J. Palomaki, Redfield, M. Collins, Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, K. Toutanova, I.. Jones, M. Kelcey, M.-W. Chang, A. M. Dai, J. Uszkoreit, Q. Le, and Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466 2019. doi: 1,.1162/tacla_0000 URL https: http:ss//aaaaatthllln
org (ttg/Q19-1001 1026/
J. Lee, A. Chen, Z. Dai, D. Dua, D. S. Sachan, M. Boratko, Y. Luan, S. M. Arold, V. Perot, S. Dalmia, et al. Can long-context language models subsume retrieval, rag, sql, and more? arXiv preprint arXiv:2406. 13121, 2024.
X. Ma, X. Zhang, R. Pradeep, and Lin. Zero-shot listwise document reranking with large language model, 2023. URL hhttps://arate http/xxxivv.ooo orggabs//3/50 02156
J. Ni, C. Qu, J. Lu, Z. Dai, G. H. Ábrego, J. Ma, Y. Zhao, Y. Luan, K. B. Hall, M.-W. Chang, and Y. Yang. Large dual encoders are generalizable retrievers, 2021. URL https:://aa:::: org/ 07899
R. Nogueira and K. Cho. Passage re-ranking with bert, 2020. URL httpps://arxiv.org/// 04085
R. Nogueira, Z. Jiang, and J, Lin. Document ranking with a pretrained sequence-to-sequenee model, 2020. URL URL https://aaxzzi..orrr httg:/ab/2200..000
R. Pradeep. S. Shanifymoghudlam. and J. Lin. Rankvicuna: Zero-shot listwise document reranking
with open-source large language models, 2023a. URL Utt. abtpaaab.////..//e/ 15088
R. Pradeep, Sharifymoghaddam, and J. Lin. Rankzephyr: Effective and robust zero-shot listwise reranking is breeze!. 2023b. URL https: arxiv wrtg//b////111.222224

---

R. G. Reddy, J. Doo, Y. Xu, M. A. Sultan, D. Swain, A. Sil, and H. Ji. First: Faster improved listwise reranking with single token decoding. arXiv preprint arXiv:2406. 2024.
N. Reimers and I. Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks, 2019. URL bottp:://arxiv.org/..a..
10084
S. Robertson, H. Zaragoza, et al. The probabilistic relevance framework: Bm25 and beyond. Foundations and Trends@ in Information Retrieval, 3(4):333-389, 2009.
K. Santhanam, O. Khattab, J. Saad-Falcon, C. Potts, and M. 7aharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction, 2022. URL htl..atpps//araav.orr/aaa//////
01488
N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost. Proceedings of the 35th International Conference on Machine Learning (ICML), 80:4596-4604, 2018.
W. Sun, L. Yan, X. Ma, S. Wang, P. Ren, Z. Chen, D. Yin, and Z. Ren. Is chatgpt good at search? investigating large language models as re-ranking agents, arXiv preprint arXiv:2304.09542, 2023.
M. S. Tamber, R. Pradeep, and J. Lin. Scaling down, litting up: Efficient zero-shot listwise reranking with seq2seq encoder-decoder models, 2023. URL hta.attp:///rrraa..rr//////22222
R. Tang, X. Zhang, X. Ma, J. Lin, and F. Ture. Found in the middle: Permutation self-consistency improves listwise ranking in large language models. arXiv preprint arXiv:2310.07712, 2023.
G. Team, R. Amil, S. Borgeaud, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, K. Millican, al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312. /8805 2023.
N. Thakur, N. Reimers, A. Rücklé, A. Srivastava, and I. Gurevych. Beir: A heterogenous benchmark for shro evaluation information retrieval models. arXiv preprint arXiv:2 (.XXxx:204408660 2021.
J. Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, et al. Emergent abilities of large language models. Transactions Machine Learning Research, 2022.
W. Wa, Y. Wang, G. Xiao, H. Peng, and Y. Fu. Retrieval head mechanistically explains long-context factuality, 2024. URL batts:://arxiit wwweiiiorrg/e abs /404 15574.
L. Xiong, C. Xiong, Li, K.-F. Tang, J. Liu, Bennett, J. Ahmed, and A. Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval, 2020. URL https:// wwxiv.org/ae. aa///bb//// 2007 00808
M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q Wang, L. Yang, al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17299 2020.
X. Zhang, S. Hofstätter, P. Lewis, R. Tang, and J. Lin. Rank-without-gpt: Building gpt-independent listwise rerankers on open-source large language models, 2023. Uttpst htapss://arxiv www.xiv..o//bb...... /bs// 2312.02969
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims clearly written in abstract and introduction and justified using appropriate experiments.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have written limitation of our work in Section A of the Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
13

---

Justification:
No theory
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We discuss all experimental details in Section 5.1 and Section
B
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer:
[Yes]
Justification:
All dataset used is already open-source. Code will be released at here
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https guwdes/dde/eeee Cooeeeoodeesssseeeeee //nips.com / ssiniiiiicypyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Justification:
We discuss all experimental details in Section 5.1 and Section B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Due to scale of the experiments it is computationally challenging to do multiple evaluations to get reliable error estimates.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Described Section 5.1 and Section
B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question:
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Conform with Code of Ethics
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We don't 't release any data or model
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly cite all existing works, pre-trained models and datasets used.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No crowdsourcing
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No crowdsourcing or human subject
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
18

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: no core LLM use
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy fattps://neurda: /ttpaarrpp..ccooooffeeeeeeee///s for what should or should not be described.
19

---

A Societal Impact
The BlockRank methodology, by enhancing the efficiency and scalability of In-context Retrieval (ICR) in Large Language Models (LLMs), makes advanced semantic retrieval more computationally tractable and can democratize access to powerful information discovery tools. This could acceler- ate research, improve educational outcomes by providing more relevant information quickly, and empower individuals and organizations with better decision-making capabilities. Furthermore, the increased efficiency directly translates to reduced energy consumption for eetricval-intesiin LLM applications, contributing to more environmentally sustainable Al development and deployment. By enabling effective ICR on potentially smaller or more optimized models, BlockRank could also broaden the reach of these technologies in resource-constrained environments.
However, like many advancements in AL, more efficient information retrieval also presents challenges. The underlying LLMs can inherit and potentially amplify societal biases present in their training data. Therefore, continued research in this area should accompanied by a strong emphasis on transparency, and the development of robust mechanisms to identify and mitigate the spread of harmful or misleading content.
B Dataset and Hyperparameter Details
This section provides a detailed description of the dataset and hyperparameters used this study to ensure reproducibility.
B.1 Datasets
We use two standard passage retrieval benchmarks:
MSMarco Passage Ranking [Bajaj et al. 2018): We use MSMarco v1 passage retrieval dataset, it has total 8.8M passages, 500K training queries and 6980 validation queries. We directly utilize the hard negatives collection from huggingface for training. Dur- ing we retrieve the top-N passages using asmarrcoddistilber--dddtt transformer.
sentence-
Natural Questions (NQ320K) (wiitkowskki et al. 2019] We use NQ320K passage retrieval dataset which has ~ 320K passages in the corpus, 300 training queries and 7830 validation queries. For NQ, we collect hard negatives using all 111-MiniLM-L12-v2 centence-rrnssfommer model for training. We use the same model during inference well to retrieve top-N passages.
B.2 Fine-tuning Details (BlockRank and Full-FT)
The following fine-tuning settings were used for both BlockRank and Full-FT Mistral-7B:
Optimizer: Adafactor (Shazeer and Stern 2018 with 31 0.9.
Learning Rate: 3 x 10-?.
Learning Rate Schedule: Linear warmup for 50 steps followed by a cosine decay.
Batch Size: A global batch size of 32, accumulated across replicas.
Number of Epochs: epoch.
Weight Decay: No weight decay.
Gradient Clipping: gradient norm clipped to 1.0.
Loss for Full-FT: Standard Next Token Prediction (NTP) cross-entropy loss, calculated on the answer tokens (i.e., the ID of the relevant document).
20

---

B.3
BlockRank Specific Hyperparameters
In addition to the general fine-tuning settings, the following hyperparameters are specific to the BlockRank method:
Auxiliary Loss Weight (A): The hyperparameter balancing the NTP loss and the auxiliary attention loss (Laur) was set to À 0.1 this ensures that both loss have the same scale. InfoNCE Temperature (7): The temperature parameter for the InfoNCE loss (Laur) was set to 0.05.
Signal-Carrying Query Tokens (T sisigaa:::: Based on our prompt template and empirical analysis (Section D). the set of tokens for signal-carrying query tokens was siginll
Middle Layer for Auxiliary Loss (I'): The auxiliary loss Laue was applied at a specific middle layer 1 20, determined empirically described in Section D.3
Chunk Length (Lehunk): The fixed length for chunks used in the structured attention mechanism. We set Lehunk 160 for MSMarco and Lehank = 384 for NQ, this ensures that ~ 95% of the passages get full represented in Lehunk sequence length.
C Attention Complexity Analysis
This section provides a analysis of the computational complexity of the structured attention mech- anism within single layer of the BlockRank model architecture. Our aim to clearly illustrate the scalability benefits of BlockRank, particularly its linear scaling with respect to the number of candidate documents, N. We define Lehunk as the fixed characteristic length (number of tokens) for segments after processing, and d as the hidden dimension of the model. For this analysis, we assume that the instruction segment, each of the N document segments, and the query segment have effective lengths List = Lehuukk Ldoo Lichank, and Ly Lechnnk respectively, when their attention com- putations are considered. This section focuses exclusively on the attention component's complexity, as this is where BlockRank introduces its primary architectural modification for efficiency.
The BlockRank model implements a structured sparse attention mechanism, as detailed in Section 4.1 of the main paper, where different parts of the input prompt adhere to distinct attention patterns. The instruction segment, with its effective length of Lehunks performs causal self-attention, leading to complexity of Iattn.Intt OLL chrun d). For the N document segments, each also of effective length Lehunt tokens attend both within their own segment and to tokens within the instruction segment. This means the effective context length for a token in any given document segment is Lehhnkk Tichunk = 2Techunk- Consequently, the attention complexity for a single document segment is O[Lehunk 2Lchank d) 0(21 d). Summing across all N document segments, their total attention complexity is Cattn = N. O(2L d).
The query segment, also with an effective length of Lehunk, has the full attention scope. It attends its own tokens, tokens from the instruction segment, and tokens from all document segments. The total context length for these query tokens becomes T.chunk + Techunk (N T.chunk) (N + Tettukk The attention complexity for the query segment is therefore Cattn,Quryy O (Lehankk (N Lohuunk d) = O(NN~ + L2 d).
Summing the complexities of these components gives the total attention complexity per layer for the BlockRank model, BlockRakk
Cattm.BlockRank = Cattn, Inst + Cattn, Doc + Cattn Query
Cattn BlockBlkkkRakk = (Lehhundd) + N (2 d) + O((N + 2) This simplifies to 23 + 3.N 2 chuungd)) which is O((N + 1)
d). The dominant term
thus yields a total attention complexity of Cattn, BlockRank = O(N L d). This result clearly shows that the attention complexity in the BlockRank architecture scales linearly with N, the number of documents.
In contrast, standard Transformer model processing sequence of comparable total length S m (N + 2) Lebund would exhibit an attention complexity of Cattn,St = O(S" d). For large N, this is approximately O((N + Te d) = O(N2. 122 T22 d), which quadratic with respect to N.
21

---

D Additional Results
D.1 Calibration Problem in Beam Decoding with Full-FT Model
To analyze the behavior of the standard fine- tuned (Full-FT) model when generating multi- ple distinet predictions via beam decoding, we conducted entropy analysis the individ- ual tokens of the predicted document identifiers. This experiment was designed assess the di- versity of predictions for structured identifiers (two-digit IDs from 0-99, given N 100 can- didate documents). For each query in the test set, we generated 10 unique document ID pre- dictions from the Full-FT model. We then com- puted the entropy of the distribution of the first
Table 5: Entropy of predicted document ID digits (idg and id1) over 10 beam-decoded predictions for the Full-FT model versus random predictions. Lower entropy indicates less diversity in the gener- ated digits across the prediction list.
Prediction Model Entropy ido
Entropy id
Full-FT
2.28 1 0.43
2.19 + 0.46
BlockRank
2.54 + 0.24
2.67 + 0.24
Random
2.35 = 0.25
2.66 + 0.24
digit (idg) and the second digit (id1) across these 10 predictions. This was compared against the entropy derived from 10 randomly drawn unique two-digit IDs. Because the candidate list randomly shuffled and the IDD assigned to each document is completely independent from its content, lower entropy would indicate a undesirable concentration of predicted digits, suggesting a lack of diversity in the generated list beyond the top few candidates.
The results, summarized in Table [5] show that the Full-FT model exhibits lower average entropy for both idg (2.28 + 0.43) and id, (2. 2.19 + 0.46) compared to the random baseline (2.55 + 0.25 for idg and 2.66 + 0.24 for id,) ). This decreased entropy indicates that the sequence of document identifiers generated by the Full-FT model via beam decoding tends to be less diverse in its constituent digits than random chance would suggest. To give an example, we observe that let's say the model predicts 73 as it's top prediction with high confidence, then there is a high likelihood that it will predict other IDs either starting with or ending with 3. Such concentration implies that while the model may identify strong top candidate, its ability to produce well-calibrated and varied set of subsequent predictions is limited, due to the nature of auto-regressive ae-probability distributions. This observation supports the main paper's 's discussion (Section 5.3 Table on the sub-optimality of beam decoding for generating ranked lists for ICR.
D.2 Analysis of Retrieval Signals in Attention Patterns of Full-FT Mistral
To substantiate the claims made in Section 3 of the main paper regarding the presence of retrieval signals within the internal attention patterns of a standard fine-tuned (Full-FT) language model, conducted series of analytical experiments. These experiments, detailed below, confirm the characteristics of such signals using attention- inference on the MSMarco dev dataset.
First, we investigate the specific carriers of the re- trieval signals. Figure 5 presents the P@1 perfor- mance of attention-based inference when attention scores are extracted from different query tokens within the prompt. This analysis reveals that cer- tain query tokens, particularly those located towards the end of the query or specific delimiter tokens such as ":" and terminal prompt markers, serve as strong "signal carriers," yielding significantly higher @@ I when their attention patterns are used to predict the relevant document.
Full-TT Attertion Accuracy Query Token
Query Ralcts
Figure 5: Performance of Full-FT model's attention-based inference vs the query token for which attention scores are extracted from.
Complementing this, F evaluates both P@ and Mean Reciprocal Rank @ 10 for attention- based inference function of the Transformer layer index from which attention scores are derived. This experiment confirms that the retrieval signal is most prevalent in the middle layers of the Full-FT model, with performance declining in earlier and later layers. Collectively, these empirical findings
22

---

Accuracy Layer
MRR@10
20 Precisian
15
10
Transformer Layer
Figure 6: Performance of Full-FT model's attention-based inference as function of the Transformer Layer Index from which attention scores are extracted (MSMarco).
demonstrate that standard LLMs fine-tuned for In- cn-tontex Retrieval exhibits latent retrieval signals within their attention mechanisms. These signals are characterized by their preferential emergence in middle layers and their association with specific query tokens.
D.3 Layerwise Emergence of Retrieval Signals and Choice of 1!
Figure illustrates the evolution of layerwise Pre- cision @ derived from attention scores on held- subset of MSMarco training data the Full-FT model undergoes training. It observed that effective retrieval signals, as measured by the attention-P@r 1 metric, do not develop uniformly across all layers. In- stead, they emerge more prominently and strengthen considerably in the middle layers of the transformer (layers 12 through 24) as training progresses, while shallower and deeper layers exhibit comparatively weaker signal strength. Based on this we set the 1' 20 for all of our BlockRank experiments, Al- though, we find that the choice of in BlockRank not very sensitive to this specific layer, any reason- able middle layer gives similar performance.
0.8
Precision
0.6
0.2
0.0
Training Steps
Figure Layerwise Attention Precision held-out subset of MSMarco training data vs training steps for Full-FT model
D.4 Impact Including Query in Prompt Prefix
We investigated whether providing the query context upfront, in addition its standard posi- tion at the end of the prompt, impacts retrieval performance. comper]s the Precision @ results MSMareo for the Full-FT baseline and our BlockRank model using the standard format (query only at end; denoted by X in the table) versus the query-pur mnery-prefix format (query at beginning and end; denoted by v). Including the query redundantly in the prefix (v) improved performance over the standard format (X) for
Table 6: Ablation on Including Query in Prompt Prefix. Comparison of Precision @ on MSMarco (N=100) for Full-FT and BlockRank models.
Model
Query in Prefix
Prec @ 1
Full-FT-Mistral Full-FT-Mistral
X 27.2
/
28.7
BlockRank-Misttall
X
24.2
BlockRank-Mistral
29.1
both models. The Full-FT model's 's Prec@1 increased from 27.2 to 28.7 (+1.5), while our BlockRank model more substantial increase from 24.2 to 28.1 (43.9). This suggests that priming the model with the query context before processes the candidate documents is beneficial, perhaps allowing attention mechanisms, particularly the specialized ones in BlockRank, on query-relevant information throughout the sequence. Given this clear advantage, we utilize the prompt format that includes the query in the prefix for all other reported experiments.
focus more effectively
23

---

D.5 Cross-dataset Generalization
To assess the generalization of the BlockRank models, we evaluated BlockRank Mistral mod- els trained on one dataset and tested on another, unseen dataset. Specifically, models were fine- tuned separately on the MSMarco and Natural Questions (NQ) training sets, and their Preci- @ (P@ 1) performance was subsequently measured the test sets of both NQ and MS- Marco. For reference, we also include the perfor- mance of zero-shot Mistral-7B-instruuc model (denoted as No Training). The results of this cross-datase evaluation are presented in Table 7 As expected, BlockRank models achieve their best performance when evaluated on the in-domain test When evaluated out-of-domain datasets, the performance, shows positive transfer above the No Training baseline but considerably lower than in-domain scores.
Table 7: Cross-dataset generalization performance of BlockRank Mistral models. P@ scores re- ported on the NQ and MSMarco test sets for mod- els with no training (zero-shot Mistral-7B-ing BBinstrrut)) fine-tuned on NQ, and fine-tuned on MSMarco.
Training Data
NQ P@
MSMarco P@1
No Training
43.5
13.1
NQ MSMarco
76.2
18.2
62.0
29.1
24


Reviews:
[{"review_id": "zPF7CrCxlh", "reviewer": "Reviewer_UGWw", "summary": "This 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.", "questions": "(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.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Overall, I believe the paper has good practical value in making attention-score-based, LLM-based rerankers more efficient, and I would therefore recommend a weak acceptance. However, in the revised version, please reframe the work as in-context re-ranking (rather than retrieval) and provide a detailed discussion of the key differences between this approach and previous reranker models.", "rebuttal": "We thank the reviewer for their detailed feedback and insightful questions. We are encouraged that the reviewer found our analysis of attention patterns 'interesting and new' and appreciated our detailed ablations. The reviewer’s primary concerns relate to experimental realism and the strength of our baselines. We believe the reviewer may have missed some key results in our appendix which already address some of these points. To further and more comprehensively address these valid concerns, we have conducted an extensive evaluation on the BEIR benchmark, where our method achieves better or comparable performance than existing listwise re-rankers while being much more efficient. We detail these new results and clarify our other points below. \n\n\n---\n\n\n\n## Experimental Realism and Baselines (Weaknesses 1, 2, 3)\n\n**Weak experimental setting**\n\n- We note that our original Appendix D.6 (Table 7) already contained comparisons to RankZephyr and RankVicuna on the TREC DL benchmark, where ReTuned model performs competitively. We will make these results more prominent. Our motivation for primarily benchmarking on in-distribution NQ and MSMarco test sets was so that we can compare baselines in a controlled setup without the confounding of the training data used. \n\n- To directly address the reviewer's primary concern about comparing against strong, contemporary rerankers, we evaluated our approach on the BEIR benchmark. We followed the exact experimental protocol from the recent SOTA listwise reranker, FIRST (also mentioned by the reviewer), reranking the top-100 documents from a Contriever baseline across 11 BEIR datasets. The results show that our MSMarco-trained ReTuned-Mistral (54.8), outperforms FIRST (54.3), RankZephyr (53.7), and RankVicuna (50.7).\n\n*Table 8. nDCG@10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model*\n\n| Reranker | Training Data | Average | Climate-FEVER | DBPedia | FEVER | FIQA | HotpotQA | MS Marco | NFC-orpus | NQ | Sci-docs | Sci-fact | Trec-COVID |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| None (Contriever baseline) | MS Marco | 45.9 | 23.7 | 41.3 | 75.8 | 32.9 | 63.8 | 40.7 | 32.8 | 49.8 | 16.5 | 67.7 | 59.6 |\n| Cross-Encoder | MS Marco | 50.7 | 25.5 | 47.0 | 81.9 | 35.6 | 71.8 | 47.0 | 34.5 | 57.6 | 17.0 | 69.1 | 71.0 |\n| Rank Vicuna | GPT 3.5 | 50.7 | **28.2** | 50.0 | 81.0 | 35.9 | 73.5 | 36.7 | 33.1 | 58.6 | 18.4 | 70.5 | 71.3 |\n| Rank Zephyr | GPT 3.5 +4 | 53.7 | 25.6 | 50.0 | 80.1 | 42.2 | 71.6 | 42.7 | **37.7** | 65.6 | **20.5** | **76.7** | 78.4 |\n| FIRST | GPT-4 | 54.3 | 26.7 | **50.9** | 81.7 | 42.2 | 74.2 | 44.4 | 37.4 | **66.4** | 20.4 | 74.6 | **78.8** |\n| ReTuned Mistral | MS Marco | **54.8** | 26.8 | 49.7 | **87.3** | **44.9** | **75.5** | **48.6** | 36.6 | 62.4 | 18.7 | **76.5** | 76.2 |\n\n\n\nFurthermore, we want to address other key points on realism:\n\n\n\n**Difficulty of the Task vs. [1]**: The reviewer points to [1] (Lee et al., 2024), which performs ICR over thousands of documents. We appreciate the reference, but we must highlight a critical difference in experimental design. The work in [1] evaluates retrieval from a random subset of documents from the corpus. In contrast, our experiments (both in the original submission and the new BEIR evaluation) focus on reranking the top-k hard candidates returned by a strong first-stage retriever. This is an objectively more difficult and realistic task, as the model must distinguish between many semantically similar documents to find the correct answer. We argue our \"hard negative\" setting is a more faithful simulation of a real-world second-stage retrieval method.\n\n\n\n**Scalability to Large Candidate Sets**: Regarding the concern about using only N=30 documents for NQ, we clarify that this was for our default experiments. To explicitly test the scalability and performance in more realistic, larger-candidate settings, we conducted the analysis presented in Figure 4. This experiment evaluates performance and latency for up to N=500 in-context documents, shwoing that ReTuning's performance remains robust (peaking at N=200 and staying high at N=500) , while the Full-FT baseline's performance degrades sharply beyond N=100. This directly demonstrates our method's suitability for the large-scale re-ranking scenarios. Moreover, the BEIR results presented above require ranking 100 documents for each dataset, showing robust generalization of our proposed approach in ranking O(100) passages even across out-of-distribution datasets.\n\n\n\n**P@1 Metric for NQ**: Our choice was motivated by our primary training objective, which is to identify the single best passage. P@1 is the most direct evaluation of this specific task. However, we agree that Top-k accuracy is also valuable for retrieval. For MSMarco, we did report the more standard MRR@10, where our method shows a significant improvement over the baseline (42.0 vs 38.3). We will add nDCG@10 numbers as well for NQ in the final version to provide a more comprehensive picture.\n\n\n\n---\n\n\n\n## Connection to Dense Retrievers and ColBERT (Weakness 4)\n\n\n\nThe reviewer raises an interesting point about the connection to dense retrievers like ColBERT. While there are similarities (query-document interaction), we wish to clarify the fundamental architectural difference. ReTuning operates within the In-context Retrieval (ICR) paradigm, where a single LLM processes the query and *the entire list of candidate documents simultaneously* in one context window. This is fundamentally different from ColBERT's dual-encoder like pointwise retrieval framework, which encodes the query and documents independently and computes a similarity score via late interaction. The key advantages of our in-context approach are:\n\n- Full Contextualization: The query representation is influenced by the full set of candidate documents, and document representations are conditioned on the shared instructions.\n\n- Instruction Following: The ICR setup allows the model to respond to complex, dynamic instructions included in the prompt (e.g., \"Find the document that disagrees with the following statement...\"), a capability not present in standard dense retriever architectures.\n\nWhile our structured attention limits direct cross-document interaction, the query module attends to all documents, serving as a global information aggregator.\n\n\n\n---\n\n\n\n## Questions\n\n\n\n### Q1: Would strictly enforcing attention not lead to a performance drop?\n\nAs shown in Table 1 & 2, our ReTuned model with structured sparse attention outperforms the Full-FT model that uses full attention (e.g., 29.1 vs 28.7 P@1 on MSMarco, 76.2 vs 75.5 P@1 on NQ). We hypothesize this is because (1) as shown in our analysis (Section 3 and Appendix D.2, D.3) that even standard models naturally develop a sparse attention pattern for this task. ReTuning's structured attention, therefore, does not remove a critical signal needed for the decoding prediction but rather enforces an already emergent and effective behavior (2) our auxiliary loss ($L_{aux}$ ) provides effective retrieval signals to the model which helps in explicit placement of attention mass to correct segments.\n\n\n\n### Q2: Which middle layer is selected?\n\nWe provide a full analysis of this in Appendix D.3 and Figure 7. Based on our analysis of where retrieval signals emerge during training, we chose layer $l^∗ =20$ for our experiments. We also note that performance is not overly sensitive to this specific choice, with other middle layers yielding similar results.\n\n\n\n### Q3: Where are the latency results in Table 3?\n\nWe apologize for the potentially confusing caption for Table 3. The table itself focuses on effectiveness metrics (P@1, MRR@10). The detailed latency results and scalability analysis are presented in Figure 4, which shows the end-to-end latency per query for both ReTuned and Full-FT models across varying numbers of documents. For instance, at N=100, ReTuning is 4.7x faster than the Full-FT baseline.\n\n\n\n### Q4: Missing proceedings information in citations.\n\nThank you for pointing this out. We will carefully review and update all citations to include the proper proceedings information in the final camera-ready version.\n\n--- \n\nWe hope our response and more experimental results have addressed the reviewer's concerns. We thank them again for their valuable feedback and are happy to provide any further clarifications/results during the discussion period."}, {"review_id": "44TS5RQW1r", "reviewer": "Reviewer_2XjM", "summary": "The 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.", "questions": "In 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?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The new results on BEIR appear strong and address my concern about the out-of-distribution generalization of the proposed method. I have increased my score accordingly.\n\nNonetheless, I still have reservations about the novelty of the method, as it represents an incremental adaptation of FiD/LiT5 to a decoder-only model. Given the widespread success of FiD/LiT5-style block attention in QA and retrieval tasks, along with abundant work on sparse attention, it is not surprising that this approach also works for decoder-only LM rerankers. The adoption appears incremental. Thus, my overall judgment still tends to be negative due to the limited technical originality of the proposed method.", "rebuttal": "We thank the reviewer for their constructive feedback and for identifying relevant prior work. We are glad they found our attention analysis valuable and our ablations thorough.\n\nThe reviewer's primary concerns are (1) the novelty of our method in light of LiT5, and (2) the scope of our evaluation. Below we first  clarify the fundamental distinctions between our work and LiT5, which we believe establish our novelty. We then present results on the suggested BEIR benchmark and redirect to some of our results in Appendix, which directly addresses the concerns about evaluation scope and baselines. \n\n---\n\n## Novelty: Differentiating ReTuning from LiT5 (Weakness 1)\nWe thank the reviewer for pointing to LiT5. While both our work and LiT5 aim for efficient listwise ranking in a way, we wish to clarify that our ReTuning method is quite different in its architectural design, technical insights, and training paradigm.\n\n**Architectural Design and Expressiveness**: While both methods incorporate a similar inductive bias (treating documents as self-contained), the implementations are distinct.\n- **Applicability to Modern LLMs**: ReTuning is designed for decoder-only architectures, which represent the vast majority of modern LLMs. In contrast, LiT5 is built for the encoder-decoder paradigm. This makes our approach more general and directly applicable to current state-of-the-art models.\n- **Interaction Expressiveness**: This architectural difference has significant implications for expressiveness. In LiT5's encoder-decoder setup, the encoder is forced to compress each document into a fixed-dimensional final representation. ReTuning, by using a single-pass sparse attention mask, allows query tokens to interact directly with all internal token representations of the documents across all layers and heads. This provides a much richer and more flexible interaction space, avoiding the information bottleneck of a final encoder state.\n\n**Contributions Beyond Structured Attention**: Our paper's contribution is significantly broader than just a sparse attention mask.\n- **Analysis of Emergent Retrieval Signals**: As the reviewer noted, our paper provides a valuable analysis (Section 3) of how decoder-only LLMs develop latent \"retrieval heads\" for the ICR task. This analysis itself is a contribution that augments the emerging literature on understanding the internal mechanisms of LLMs.\n- **Training Objective**: this analysis directly informs the design of our auxiliary contrastive loss ($L_{aux}$). To the best of our knowledge, this is a novel technique for directly optimizing and calibrating the internal self-attention scores of an LLM for retrieval.\n\n**Efficiency of the Training Signal**: A subtle but critical distinction is the type of supervision required.\n- LiT5-Distill is trained on rich, listwise ranking data generated by a powerful GPT-4 teacher model.\n- ReTuning, in contrast, is trained using much sparser, pointwise relevance data (i.e., a single positive document paired with hard negatives).As our BEIR results (below) show, ReTuning achieves similar if not better performance, surpassing models trained on stronger listwise supervision. \n\nIn summary, ReTuning introduces an analysis-driven approach for ICR in modern decoder-only LLMs, introduces a contrastive learning objective of internal attention signals and gets effective performance even with sparse pointwise retrieval signals.\n\n---\n\n## Evaluation Scope and Baselines (Weaknesses 2 & 3)\nOur motivation for primarily benchmarking on in-distribution NQ and MSMarco test sets was to compare baselines in a controlled setup without the confounding of the training data used. Nonetheless, we agree with the reviewers that demonstrating broad generalization is essential.\n- To address the concerns about our evaluation being limited to in-distribution datasets and narrow baselines, we evaluated on the suggested BEIR benchmark. We followed the exact experimental protocol from the recent SOTA listwise reranker, FIRST [1], reranking the top-100 documents from Contriever baseline across 11 BEIR datasets. The results show that our MSMarco-trained ReTuned-Mistral (54.8), outperforms FIRST (54.3), RankZephyr (53.7), and RankVicuna (50.7). This demonstrates both the effectiveness and strong out-of-distribution generalization of our method. Crucially, it achieves these results with the significant efficiency gains of ReTuning (Figure 4), presenting a compelling combination of effectiveness and scalability.\n\n*Table 8. nDCG@10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model*\n\n| Reranker | Training Data | Average | Climate-FEVER | DBPedia | FEVER | FIQA | HotpotQA | MS Marco | NFC-orpus | NQ | Sci-docs | Sci-fact | Trec-COVID |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| None (Contriever baseline) | MS Marco | 45.9 | 23.7 | 41.3 | 75.8 | 32.9 | 63.8 | 40.7 | 32.8 | 49.8 | 16.5 | 67.7 | 59.6 |\n| Cross-Encoder | MS Marco | 50.7 | 25.5 | 47.0 | 81.9 | 35.6 | 71.8 | 47.0 | 34.5 | 57.6 | 17.0 | 69.1 | 71.0 |\n| Rank Vicuna | GPT 3.5 | 50.7 | **28.2** | 50.0 | 81.0 | 35.9 | 73.5 | 36.7 | 33.1 | 58.6 | 18.4 | 70.5 | 71.3 |\n| Rank Zephyr | GPT 3.5 +4 | 53.7 | 25.6 | 50.0 | 80.1 | 42.2 | 71.6 | 42.7 | **37.7** | 65.6 | **20.5** | **76.7** | 78.4 |\n| FIRST | GPT-4 | 54.3 | 26.7 | **50.9** | 81.7 | 42.2 | 74.2 | 44.4 | 37.4 | **66.4** | 20.4 | 74.6 | **78.8** |\n| ReTuned Mistral | MS Marco | **54.8** | 26.8 | 49.7 | **87.3** | **44.9** | **75.5** | **48.6** | 36.6 | 62.4 | 18.72 | **76.5** | 76.2 |\n\n- We also note that our original Appendix D.6 (Table 7) already contains comparisons to RankZephyr and RankVicuna on the TREC DL20/22 benchmark, where our model performs competitively. Moreover, Appendix D.8 (Table 6) provides cross-dataset performance comparison between ReTuned models on NQ and MSMarco. We will make these results more prominent in the final version.\n\n---\n\n## Questions\n\n### Q1: In Figure 2, why does the query token almost not attend to other document tokens?\n\nWe believe the reviewer is referring to the left plot in Figure 2, which shows segment-wise attention. In that plot, the query segment's attention can appear diffused across all other segments. This is partly a visual effect of heatmap scaling - because the query attends to the most segments in the prompt (all documents and instructions), its attention mass is more distributed, making the specific attention to an individual segment appear less dense; moreover majority of tokens in the query segment are prompt instruction tokens (e.g. \"Which document is most relevant...\") that are not specifically related to any particular document hence aggregating attention mass over all tokens belonging to query segment diffuses the attention mass over all documents and concentrates more on the intra-segment tokens (shown by strong diagonal element). The middle plot of Figure 2 provides a more granular view - this plot specifically isolates the attention from individual query tokens to the document segments in a key middle layer (Layer 18). As this plot shows actual query tokens and some specific 'signal-carrying' tokens (like the final token and ':') learn to attend very strongly to the single ground-truth relevant document (highlighted in green).\n\n### Q2: How would the model perform if a different first-stage retriever were used?\nOur evaluation on the BEIR benchmark provides strong evidence here. The BEIR evaluation uses a Contriever first-stage retriever, which is different from the SentenceTransformer models used in our original experiments. The fact that our MSMarco-trained model generalizes so effectively to rerank Contriever's output across 11 diverse datasets strongly suggests that ReTuning is not overly sensitive to the specific first-stage retriever.\n\n--- \n\n## References\n\n[1] FIRST: Faster Improved Listwise Reranking with Single Token Decoding, Reddy et al, EMNLP 2024\n\n---\n\nWe hope our response and more experimental results have addressed the reviewer's concerns. We thank them again for their valuable feedback and are happy to provide any further clarifications/results during the discussion period."}, {"review_id": "DHZOdC70ic", "reviewer": "Reviewer_y5o2", "summary": "The 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.", "questions": "1. 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?", "limitations": "yes, in Supplementary Material", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### 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", "was_revised": true, "final_justification": "The authors have addressed most of my concerns. However, my main reservation is whether this method can transfer to other existing LLMs or be universally applied to any LLM, which prevents me from raising my score further. My current score and confidence reflect my overall positive impression of the paper.", "rebuttal": "We sincerely thank the reviewer for their positive and encouraging feedback, and for their thoughtful questions. We are glad they found our work \"timely and important\" and our analysis \"well-motivated.\" We provide clarifications and more results below that we hope will address the concerns and solidify the reviewer's positive assessment.\n\n---\n\n## Evaluation Scope and Baselines (Weaknesses 1&2)\n**Limited Dataset Scope and Baselines:**\nOur motivation for primarily benchmarking on in-distribution NQ and MSMarco test sets was so that we can compare baselines in a controlled setup without the confounding of the training data used. Nonetheless, we agree with the reviewers that demonstrating broad generalization is essential. To address this we evaluated our method on the suggested BEIR benchmark. We followed the exact experimental protocol from the recent SOTA listwise reranker, FIRST [1], reranking the top-100 documents from Contriever baseline across 11 BEIR datasets. The results show that our MSMarco-trained ReTuned-Mistral (54.8), outperforms FIRST (54.3), RankZephyr (53.7), and RankVicuna (50.7). This demonstrates both the effectiveness and strong out-of-distribution generalization of our method. Crucially, our method achieves *comparable or better results with the significant efficiency gains of ReTuning (Figure 4)*, presenting a compelling combination of effectiveness and scalability.\n\n*Table 8. nDCG@10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model*\n\n| Reranker | Training Data | Average | Climate-FEVER | DBPedia | FEVER | FIQA | HotpotQA | MS Marco | NFC-orpus | NQ | Sci-docs | Sci-fact | Trec-COVID |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| None (Contriever baseline) | MS Marco | 45.9 | 23.7 | 41.3 | 75.8 | 32.9 | 63.8 | 40.7 | 32.8 | 49.8 | 16.5 | 67.7 | 59.6 |\n| Cross-Encoder | MS Marco | 50.7 | 25.5 | 47.0 | 81.9 | 35.6 | 71.8 | 47.0 | 34.5 | 57.6 | 17.0 | 69.1 | 71.0 |\n| Rank Vicuna | GPT 3.5 | 50.7 | **28.2** | 50.0 | 81.0 | 35.9 | 73.5 | 36.7 | 33.1 | 58.6 | 18.4 | 70.5 | 71.3 |\n| Rank Zephyr | GPT 3.5 +4 | 53.7 | 25.6 | 50.0 | 80.1 | 42.2 | 71.6 | 42.7 | **37.7** | 65.6 | **20.5** | **76.7** | 78.4 |\n| FIRST | GPT-4 | 54.3 | 26.7 | **50.9** | 81.7 | 42.2 | 74.2 | 44.4 | 37.4 | **66.4** | 20.4 | 74.6 | **78.8** |\n| ReTuned Mistral | MS Marco | **54.8** | 26.8 | 49.7 | **87.3** | **44.9** | **75.5** | **48.6** | 36.6 | 62.4 | 18.7 | **76.5** | 76.2 |\n\n**Regarding the NV-Embed baseline**, we thank the reviewer for the suggestion. NV-Embed's goal is to produce a single dense vector representation of a document for efficient indexing and first-stage retrieval. Our work, in contrast, focuses on second-stage ranking using in-context capabilities of LLMs, where the model operates over a shortlisted list of documents in the prompt. We hope the broader point of comparing against recent comparable baselines has been addressed by the BEIR evaluation above. We will add a discussion of NV-Embed and clarify this distinction in our related work section.\n\n**Limited Model Scope**:\nWe agree with the reviewer that testing on more model families is an important direction that will add to the robustness of the claims. Certain licensing restrictions limited our ability to experiment with some other popular models for this submission. We would also like to note, similar recent works which adapt LLMs for in-context ranking [1,2,3] also tend to focus on a single, powerful open-source model to allow for a deep and controlled analysis of the proposed techniques. Moreover, one of the key contributions of our paper is utilizing internal attention patterns as retrieval signals, recent work such as [4] has shown that such patterns are a universal property found across various LLM architectures and sizes which further suggest underlying principles of ReTuning are robust and likely to generalize to other decoder-only models.\n\n---\n\n## Questions\n\n### Q1: Does the reported latency in Figure 4 include the time spent on constructing the structured attention masks?\nYes, the reported end-to-end latency in Figure 4 includes the time for mask construction. The overhead for this step is negligible. Creating the sparse mask involves simple indexing and logical operations, which are orders of magnitude faster than the computationally intensive matrix multiplications (the $Q⋅K^T$ operation) in the attention mechanism itself. We will add a note to the paper to clarify this.\n\n### Q2: Confusion regarding Full-FT setting...\nMistral-7B is a decoder-only model, so all attention is causal (left-to-right). Our description of the \"Full Fine-tuning (Full-FT)\" baseline intended to contrast it with our sparse attention method.\n- Full-FT uses the model's standard full causal attention mask, where each token can attend to all preceding tokens in the sequence.\n- ReTuning uses our proposed sparse causal attention mask, where tokens have a restricted view of preceding tokens.\n\nThere is *no* switch from bidirectional to causal attention. We apologize for the ambiguity in our initial description and will revise the text to be more explicit (e.g., using \"full causal attention\" for the baseline) to avoid this confusion.\n\n### Q3: In Table 2, both Full-FT and ReTuning (full) achieve the same Precision@1 (28.7) under the Decode inference mode... Do the authors have any insight into why the performance remains identical?\nWe believe the similar decoding performance is due to two factors: the presence of the generative training objective in both setting and the inherent sparsity of attention in this task.\n- First, the model's decoding capability is driven by the Next-Token Prediction ($\\mathcal{L}_{NTP}$) loss, which is present in both Full-FT and ReTuning losses. The critical role of this objective is confirmed in **Table 2**; removing the NTP loss (`ReTuning w/o ntp`) causes decoding Precision@1 to collapse from 28.7 to 15.8. This shows that the powerful NTP signal is essential for a ReTuned model to perform well at decoding.\n- Second, our analysis (Section 3 and Appendix D.2, D.3) shows that even standard models naturally develop a sparse attention pattern for this task. ReTuning's structured attention, therefore, does not remove a critical signal needed for the decoding prediction but rather enforces an already emergent and effective behavior. Because the crucial information flow is preserved and both models are governed by the same generative loss, their decoding performance for the top-1 document is similar.\n\n--- \n\n## References\n\n[1] FIRST: Faster Improved Listwise Reranking with Single Token Decoding, Reddy et al, EMNLP 2024\n \n[2] RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!, Pradeep et al, 2023\n\n[3] RankVicuna: Zero-Shot Listwise Document Reranking with Open-Source Large Language Models, Pradeep et al, 2023\n\n[4] Retrieval Head Mechanistically Explains Long-Context Factuality, Wu et al, ICLR 2025\n\n\n---\n\nWe hope our response and more experimental results have addressed the reviewer's concerns. We thank them again for their valuable feedback and are happy to provide any further clarifications/results during the discussion period."}, {"review_id": "N0QJxs5djX", "reviewer": "Reviewer_iUgW", "summary": "This 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.", "questions": "See weaknesses.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The response addresses most of my concerns, and I keep my positive assessment.", "rebuttal": "We sincerely thank the reviewer for their feedback and questions. We are glad they found our analysis and motivation clear. We provide clarifications below that we hope will address the concerns and solidify the reviewer's positive assessment.\n\n---\n\n## Novelty vs. Block Attention (Weakness 1)\nWe would like to note that while 'block attention' and other block-wise sparse methods are known techniques for general efficiency in LLMs, our work's novelty lies in its *task-specific design and integration into a complete training and inference system for In-context Retrieval*.\n\nMoreover, we provide contributions beyond structured attention:\n- **Analysis of Emergent Retrieval Signals**: Our paper provides a valuable analysis (Section 3 and Appendix D.2, D.3, D.4) of how LLMs develop latent \"retrieval heads\" for the ICR task. This analysis itself is a contribution that augments the emerging literature on understanding the internal mechanisms of LLMs.\n- **Training Objective**: this analysis directly informs the design of our auxiliary contrastive loss ($L_{aux}$) and alternate retrieval inference mechanisms. To the best of our knowledge, this is a novel technique for directly optimizing and calibrating the internal self-attention scores of an LLM for retrieval.\n\n--- \n## Robustness and Technical Limitations (Weaknesses 2, 3, 4)\n\n**(Weakness 2) Generalizability of Signal Tokens**\n\nThe specific signal-carrying tokens, such as \":\" and \"[\", are not required to be present in the raw user query. Instead, they are a deliberate part of the prompt template that structures the input for the model, as shown in Figure 1. By explicitly including these tokens in the prompt template, we provide a consistent and reliable structural \"anchor\" for the model to attach its retrieval signal to during fine-tuning. This ensures that such signal carriers are always present for the model to utilize, regardless of the specific content of the user's query.\n\n**(Weakness 3) Incompatibility with FlashAttention**\n\nWe thank the reviewer for this subtle but important practical point. Please find below our response:\n- **Fundamental Complexity Advantage**: ReTuning's primary efficiency gain comes from changing the attention complexity from quadratic, $O(N^2)$, to linear, $O(N)$. While FlashAttention optimizes the quadratic operation, it doesn't change its underlying complexity, making our method fundamentally more scalable for very long contexts.\n- **Potential for Custom Kernels**: Our attention-based inference does not require the entire attention matrix, it only needs the block-summed probabilities for a few specific query tokens. This computational pattern is highly amenable to a custom GPU kernel—similar in spirit to FlashAttention—that could compute these specific scores directly without materializing the intermediate matrix. Developing such a kernel is a promising direction for future engineering work to gain even further speedups.\n- **Flexible and Compatible Inference**: Our method offers two inference pathways. While our attention-based inference is incompatible with standard FlashAttention, our standard decoding path is fully compatible with it and achieves similar performance to the baseline (as shown in Table 2).\n\n**(Weakness 4) Potential for Length Bias**\n\nWe acknowledge that summing attention scores could potentially introduce length bias. Our framework, however, can easily mitigate this. We primarily used sum aggregation in our experiments and found it performed well on our benchmarks, which have natural passage length variations. However, as noted in the paper (line 233), our relevance score calculation can be trivially modified to use mean aggregation over a document's tokens instead of a sum. This provides explicit length normalization. We will expand on this point in the final version and add an ablation study comparing mean vs. sum aggregation to demonstrate this robustness.\n\n\n---\n\nWe hope our response has addressed the reviewer's concerns. We thank them again for their valuable feedback and are happy to provide any further clarifications/results during the discussion period."}]

Author rebuttal:

Meta-review: The paper builds off of insights regarding attention sparsity for in-context reranking to propose a model architecture that enforces structured sparse attention, which reduces the computational complexity of attention from quadratic to linear.  The paper also provides fairly thorough experimentation in the main paper and Appendix that has been further augmented in the rebuttal.  I appreciate the directness of the authors in responding to reviewer concerns and acknowledging deficiencies.

A few issues stood out to me during the author-reviewer discussion period:
- Concerns about terminology of ICR.  The authors have agreed to correct this on revision to in-context reranking.
- Limited Model Scope.  The reviewer relaxes their concerns on this issue, but the AC maintains that it is a critical concern.  If the proposed insights and methodology are specific to a single architecture, then it would limit the general applicability of the findings in the work.  However, I will acknowledge that BERT has been a significant single model workhorse of much existing neural retrieval so there is an established literature in this area focusing on improving individual models.
- Concerns about OOD generalization.  Rebuttal results on BEIR have addressed reviewer concerns.
- Concerns about reproducibility arising from questions about implementation details spread across multiple sections in the discussion.  I remark that code has not been released (cf. checklist).  While not required, I believe a lack of code release significantly limits replication, reuse, comparison to the proposed work.  I appreciate the development and release of code of the custom Triton kernel during discussion, but I would strongly encourage the authors to release *all* code on publication.

While the final reviewer consensus puts this paper on the borderline of the acceptance threshold, I believe the paper is notable for both its analytical/scientific insight as well as its technical contributions and thorough empirical evaluation.  On this basis, I recommend the paper for acceptance to NeurIPS in good faith that the authors will do their best to address reviewer concerns (and strongly consider the possibility of code release) in their final revised version.
Decision: Accept (poster)