Scalable In-context Ranking with Generative Models

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.

---

Novelty vs. Block Attention (Weakness 1)

We 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.

Moreover, we provide contributions beyond structured attention:

• 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.
• 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.

---

Robustness and Technical Limitations (Weaknesses 2, 3, 4)

(Weakness 2) Generalizability of Signal Tokens

The 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.

(Weakness 3) Incompatibility with FlashAttention

We thank the reviewer for this subtle but important practical point. Please find below our response:

• 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.
• 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.
• 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).

(Weakness 4) Potential for Length Bias

We 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.

---

We 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.

import jax
import jax.numpy as jnp
from flax import linen as nn
from einops import rearrange

# Shapes: N=docs, L=chunk_length, D=hidden_dim, H=num_heads, M=num_signal_tokens

def standard_attn(q, k, v):
    """Standard attention computation"""
    logits = jnp.einsum('qhd,khd->hqk', q, k) / jnp.sqrt(q.shape[-1])
    probs = nn.softmax(logits, axis=-1)
    return jnp.einsum('hqk,khd->qhd', probs, v)

def compute_attn_rel_probs(q, k_docs):
    """Compute attention relevance probabilities from signal query tokens to docs"""
    logits = jnp.einsum('mhd,nlhd->mnlh', q, k_docs) / jnp.sqrt(q.shape[-1])
    probs = nn.softmax(rearrange(logits, 'm n l h -> m h (n l)'), axis=-1)
    doc_probs = jnp.mean(rearrange(probs, 'm h (n l) -> m h n l', n=k_docs.shape[0]), axis=(-1, 0, 1)) # (N,)
    return doc_probs # (N,)

def retuning_attention_layer(
    inst_h: jax.Array,   # (L, D)
    docs_h: jax.Array,   # (N, L, D)
    query_h: jax.Array,  # (L, D)
    params: dict,        # W_q, W_k, W_v, W_o weights
    signal_indices: jax.Array = None # Optional, for relevance scores
) -> tuple[jax.Array, jax.Array, jax.Array | None]:
    """Computes one layer of ReTuning's structured sparse attention and optionally relevance scores."""
    # 1. Project inputs to Q, K, V and reshape for multi-head attention
    q_inst, k_inst, v_inst = [rearrange(inst_h @ params[w], 'l (h d) -> l h d', h=H) for w in ('W_q', 'W_k', 'W_v')]
    q_docs, k_docs, v_docs = [rearrange(docs_h @ params[w], 'n l (h d) -> n l h d', h=H) for w in ('W_q', 'W_k', 'W_v')]
    q_query, k_query, v_query = [rearrange(query_h @ params[w], 'l (h d) -> l h d', h=H) for w in ('W_q', 'W_k', 'W_v')]

    # 2. Instruction Attention
    updated_inst_h = standard_attn(q_inst, k_inst, v_inst)  # (L, H, d_head)

    # 3. Document Attention Function (attends to instruction and self)
    def doc_attn_fn(q_doc, k_doc, v_doc):
        k_context = jnp.concat([k_inst, k_doc], axis=0) # (2*L, H, d_head)
        v_context = jnp.concat([v_inst, v_doc], axis=0) # (2*L, H, d_head)
        return standard_attn(q_doc, k_context, v_context)

    updated_docs_h = jax.vmap(doc_attn_fn)(q_docs, k_docs, v_docs) # (N, L, H, d_head)

    # 4. Query Attention (attends to all context)
    k_all_docs = rearrange(k_docs, 'n l h d -> (n l) h d')
    v_all_docs = rearrange(v_docs, 'n l h d -> (n l) h d')
    k_query_context = jnp.concat([k_inst, k_all_docs, k_query], axis=0)
    v_query_context = jnp.concat([v_inst, v_all_docs, v_query], axis=0)
    
    updated_query_h = standard_attn(q_query, k_query_context, v_query_context) # (L, H, d_head)

    # 5. Project outputs back to hidden dimension D
    updated_inst_h = rearrange(updated_inst_h, 'l h d -> l (h d)') @ params['W_o']
    updated_docs_h = rearrange(updated_docs_h, 'n l h d -> n l (h d)') @ params['W_o']
    updated_query_h = rearrange(updated_query_h, 'l h d -> l (h d)') @ params['W_o']

    # 6. Optional calculate relevance score
    relevance_scores = None
    if signal_indices is not None:
        relevance_scores = compute_attn_rel_probs(q_query[signal_indices], k_docs)  # (N, )

    return updated_inst_h, updated_docs_h, updated_query_h, relevance_scores

We acknowledge that while documents do have varying lengths, but would like to note that:

1. This is a common challenge for any IR algorithm that processes multiple candidates in parallel, such as ColBERT or Cross-Encoders, which typically employ a similar truncation strategy. Moreover, our method is also compatible with popular techniques from the literature:

   • Multiple Chunks: Very long documents can be split into multiple chunks and treated as independent candidates. The final document score can then be aggregated from its constituent chunk scores (e.g., by taking the max/sum).
   • Dynamic Binning: Documents can be grouped into buckets based on their length (e.g., powers of 2). This allows for more efficient padding/truncation within each bucket, trading a degree of parallelism for reduced computational waste.

2. Furthermore, majority of documents are within a manageable length and can be predetermined for a given dataset. For instance, in MSMARCO and NQ, 95% of passages are within 160 and 384 tokens, respectively. In all of our experiments, including new results on 11 BEIR dataset (please see Table 8 in response to other reviews), we do not see any tangible performance increase when using chunk sizes beyond 512 tokens, which empirically justifies our choice.

In summary, our fixed-length chunking is an empirically justified and pragmatic choice for benchmark datasets, while the ReTuning framework remains fully compatible with more advanced strategies for real-world deployments with highly irregular document lengths.

---

We thank the reviewer for the suggestion of kernelized implementation. While we believe that developing a production-grade, fused Triton/CUDA kernel is a considerable engineering effort that we view as an exciting direction for future work, nonetheless, we attempted to write a custom Triton kernel for the compute_attn_rel_probs function described above. Following are the benchmarking results on synthetic Q (shape = [M, H, D]) and K_docs (shape = [N, L, H, D]) tensors (we vary N i.e. number of docs to test scalability and keep the rest of the shape parameters similar to our qualitative experiments with Mistral model i.e. H=32, D=128, L=384):

Table 9. Custom Kernel Time Benchmark (ms)

Table 10. Custom Kernel Peak Memory Benchmark (GB)

As the results show, our custom kernel is approximately 2x faster and uses ~40-50% less peak memory than the standard PyTorch implementations at scale (N=512, 1024). This empirically validates our assertion that an I/O-aware, fused kernel provides significant practical benefits for our attention-based inference.

Such a kernel serves as an additional, lower-level optimization on top of the foundational architectural improvements we have already demonstrated and empirically validated. Please find below the code to the custom triton kernel.

@triton.jit
def _attn_rel_probs_kernel(
    Q, K_docs, Out, M, L, H,
    stride_q_m, stride_q_h, stride_k_n, stride_k_l, stride_k_h,
    stride_out_h, stride_out_n,
    sm_scale,
    BLOCK_L: tl.constexpr, BLOCK_M: tl.constexpr,
    D_HEAD: tl.constexpr, N: tl.constexpr,
):
    off_h = tl.program_id(0)
    offs_m = tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, D_HEAD)

    # Load Q matrix into SRAM, padding M to BLOCK_M
    q_ptrs = Q + off_h * stride_q_h + offs_m[:, None] * stride_q_m + offs_d[None, :]
    q_mask = offs_m[:, None] < M
    q_mat = tl.load(q_ptrs, mask=q_mask, other=0.0)

    # SRAM accumulators
    m_vec = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
    l_vec = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc_mat = tl.zeros([BLOCK_M, N], dtype=tl.float32)

    # Single Pass over K tensor
    for n_idx in range(N):
        for l_start in range(0, L, BLOCK_L):
            offs_l = l_start + tl.arange(0, BLOCK_L)
            k_offset = n_idx * stride_k_n + off_h * stride_k_h
            k_ptrs = K_docs + k_offset + offs_l[:, None] * stride_k_l + offs_d[None, :]
            k_mask = offs_l[:, None] < L
            k_block = tl.load(k_ptrs, mask=k_mask, other=0.0)

            logits = tl.dot(q_mat, tl.trans(k_block)) * sm_scale
            logits = tl.where(q_mask, logits, -float('inf'))
            
            m_block_max = tl.max(logits, axis=1)
            m_new = tl.maximum(m_vec, m_block_max)
            
            alpha = tl.exp(m_vec - m_new)
            l_vec = l_vec * alpha
            acc_mat = acc_mat * alpha[:, None]
            
            p_numerators = tl.exp(logits - m_new[:, None])
            l_vec += tl.sum(p_numerators, axis=1)
            
            p_sum_for_block = tl.sum(p_numerators, axis=1)
            n_col_mask = tl.arange(0, N) == n_idx
            acc_mat += tl.where(n_col_mask[None, :], p_sum_for_block[:, None], 0.0)
            
            m_vec = m_new

    # Final Normalization and Store
    final_probs_per_mn = acc_mat / l_vec[:, None]
    final_probs_per_mn = tl.where(q_mask, final_probs_per_mn, 0.0)
    output_per_head_n = tl.sum(final_probs_per_mn, axis=0)
    
    out_ptrs = Out + off_h * stride_out_h + tl.arange(0, N) * stride_out_n
    tl.store(out_ptrs, output_per_head_n)

def compute_attn_rel_probs_triton(q: torch.Tensor, k_docs: torch.Tensor) -> torch.Tensor:
    M, H, D_HEAD = q.shape
    N, L, _, _ = k_docs.shape
    q, k_docs = q.contiguous(), k_docs.permute(0, 2, 1, 3).contiguous()
    
    temp_out = torch.empty((H, N), device=q.device, dtype=torch.float32)
    sm_scale = 1.0 / (D_HEAD ** 0.5)
    grid = (H,)
    
    _attn_rel_probs_kernel[grid](
        q, k_docs, temp_out, M, L, H,
        q.stride(0), q.stride(1), 
        k_docs.stride(0), k_docs.stride(2), k_docs.stride(1),
        temp_out.stride(0), temp_out.stride(1),
        sm_scale, D_HEAD=D_HEAD, N=N
    )
    
    doc_probs = torch.sum(temp_out, dim=0)
    return (doc_probs / (M * H * L))

def compute_attn_rel_probs_pytorch(q, k_docs):
    """
    PyTorch reference implementation for computing attention relevance probabilities.
    """
    # float32 for numerical stability
    q_f32, k_docs_f32 = q.float(), k_docs.float()

    sm_scale = 1.0 / (q_f32.shape[-1] ** 0.5)
    logits = torch.einsum('mhd,nlhd->mnlh', q_f32, k_docs_f32) * sm_scale

    M, N_p, L_p, H_p = logits.shape
    logits_rearranged = rearrange(logits, 'm n l h -> m h (n l)', n=N_p, l=L_p)
    probs = F.softmax(logits_rearranged, dim=-1)

    probs_rearranged = rearrange(probs, 'm h (n l) -> m h n l', n=N_p, l=L_p)
    doc_probs = torch.mean(probs_rearranged, dim=(0, 1, 3))
    return doc_probs

compiled_pytorch_fn = torch.compile(compute_attn_rel_probs_pytorch)

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.

---

Evaluation Scope and Baselines (Weaknesses 1&2)

Limited Dataset Scope and Baselines: 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. 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.

Table 8. nDCG@10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model

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.

Limited Model Scope: We 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.

---

Questions

Q1: Does the reported latency in Figure 4 include the time spent on constructing the structured attention masks?

Yes, 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.

Q2: Confusion regarding Full-FT setting...

Mistral-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.

• Full-FT uses the model's standard full causal attention mask, where each token can attend to all preceding tokens in the sequence.
• ReTuning uses our proposed sparse causal attention mask, where tokens have a restricted view of preceding tokens.

There 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.

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?

We 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.

• 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.
• 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.

---

References

[1] FIRST: Faster Improved Listwise Reranking with Single Token Decoding, Reddy et al, EMNLP 2024

[2] RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!, Pradeep et al, 2023

[3] RankVicuna: Zero-Shot Listwise Document Reranking with Open-Source Large Language Models, Pradeep et al, 2023

[4] Retrieval Head Mechanistically Explains Long-Context Factuality, Wu et al, ICLR 2025

---

We 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.

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.

The 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.

---

Novelty: Differentiating ReTuning from LiT5 (Weakness 1)

We 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.

Architectural Design and Expressiveness: While both methods incorporate a similar inductive bias (treating documents as self-contained), the implementations are distinct.

• 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.
• 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.

Contributions Beyond Structured Attention: Our paper's contribution is significantly broader than just a sparse attention mask.

• 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.
• 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.

Efficiency of the Training Signal: A subtle but critical distinction is the type of supervision required.

• LiT5-Distill is trained on rich, listwise ranking data generated by a powerful GPT-4 teacher model.
• 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.

In 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.

---

Evaluation Scope and Baselines (Weaknesses 2 & 3)

Our 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.

• 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.

Table 8. nDCG@10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model

• 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.

---

Questions

Q1: In Figure 2, why does the query token almost not attend to other document tokens?

We 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).

Q2: How would the model perform if a different first-stage retriever were used?

Our 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.

---

References

[1] FIRST: Faster Improved Listwise Reranking with Single Token Decoding, Reddy et al, EMNLP 2024

---

We 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.

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. 

---

Experimental Realism and Baselines (Weaknesses 1, 2, 3)

Weak experimental setting

• 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.

• 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).

Table 8. nDCG@10 on BEIR benchmark, all re-ranker rank top-100 documents retrieved from Contriever retrieval model

Furthermore, we want to address other key points on realism:

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.

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.

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.

---

Connection to Dense Retrievers and ColBERT (Weakness 4)

The 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:

• Full Contextualization: The query representation is influenced by the full set of candidate documents, and document representations are conditioned on the shared instructions.

• 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.

While our structured attention limits direct cross-document interaction, the query module attends to all documents, serving as a global information aggregator.

---

Questions

Q1: Would strictly enforcing attention not lead to a performance drop?

As 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.

Q2: Which middle layer is selected?

We 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.

Q3: Where are the latency results in Table 3?

We 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.

Q4: Missing proceedings information in citations.

Thank 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.

---

We 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.