PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable In-context Ranking with Generative Models

Nilesh Gupta, Chong You, Srinadh Bhojanapalli, Sanjiv Kumar, Inderjit Dhillon, Felix Yu

arXiv:2510.05396v2cs.IRcs.LG

TL;DR

As candidate lists grow, attention makes LLM-based ICR computationally expensive. BlockRank combines task-specific sparse attention with a contrastive auxiliary objective, and experiments report strong ranking effectiveness with greater inference efficiency.

  • Problem

    As candidate lists grow, ICR prompts make inference computationally expensive and memory-intensive because attention scales quadratically or super-linearly with context length.

    Source: §1 · p. 2

  • Method

    BlockRank enforces task-specific structured sparse attention for linear complexity and uses a contrastive auxiliary loss to optimize attention from query signals toward relevant documents.

    Source: §1 · p. 2

  • Results

    Across BEIR, MSMarco, and NQ, BlockRank matches or outperforms strong baselines and the fully fine-tuned model while being more efficient at inference.

    Source: §1 · p. 3, §6 · p. 12

  • Takeaways & Limitations

    The authors present BlockRank as a scalable and effective approach for LLM-based ICR, with findings demonstrated primarily on one model architecture.

    Source: §6 · p. 12

    Excluding the query from the instruction prefix enables query-independent document representations to be cached offline, but the authors report a noticeable performance drop.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

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.

1. Introduction

In-context ranking uses an LLM to judge a query against candidate documents in one prompt, but growing candidate lists make attention costly. The paper identifies exploitable attention structure and introduces BlockRank, which uses structured sparsity and an auxiliary attention objective; its experiments report strong ranking performance with substantially lower inference cost.

1. Introduction

  • In-context ranking places the task description, candidate documents, and query in one LLM prompt, but growing candidate lists make attention computationally expensive.The attention operation scales quadratically or super-linearly with context length.Source: §1 · p. 2
  • Selected query tokens, including the final query token and “:”, develop strong attention toward relevant document tokens, particularly in middle layers.The passage relates this pattern to Chen et al. (2025) and Wu et al. (2024).Source: §1 · p. 2
  • BlockRank imposes structured sparse attention for linear complexity and contrastively trains query-token attention toward relevant documents, also enabling prefill-stage relevance inference.The method builds on attention patterns in which documents mainly attend locally and selected query tokens focus on relevant documents in middle layers.Source: §1 · p. 2
  • Figure 2 depicts prompt chunking, structured attention, and an auxiliary middle-layer attention loss that sharpens attention on relevant documents.Documents attend to themselves and the instruction segment, while the query attends to the full prompt.Source: Figure 2 · p. 3
  • Across BEIR, MSMarco, and NQ, BlockRank matches or outperforms strong baselines and a fully fine-tuned model while using an order of magnitude less inference compute.The reported experiments use standard retrieval benchmarks.Source: §1 · p. 3

2. Problem Setup and Related Work

This section defines in-context ranking as identifying relevant items from a query-conditioned candidate list, then situates it among neural retrieval, listwise reranking, and efficient attention methods. The paper focuses on hard candidates from a first-stage retriever and proposes task-specific attention structure for this setting.

2.1. Problem Formulation: In-context Ranking

  • In the paper’s setup, the full corpus is reduced to a smaller candidate list retrieved by a first-stage model because processing the entire corpus in one prompt is infeasible.The experiments use about 100 candidates.Source: §2.1 · p. 3
  • Including the query in the instruction prefix improves performance in the authors’ experiments, although it prevents query-independent document representations from being cached offline.The authors hypothesize that early query conditioning helps focus document representations on the specific information need.Source: §2.1 · p. 4
  • Given a query and candidate list, in-context ranking trains or uses an LLM to identify relevant document IDs, typically by predicting a ranked permutation and taking the top k.The task balances prediction effectiveness and computational efficiency as candidate-list size increases.Source: §2.1 · p. 4

2.2. Related Work

  • Unlike dual-encoder and cross-encoder retrieval methods, in-context ranking processes the query and candidate list together in one LLM context, supporting full contextualization and instruction-following.The passage contrasts this setup with dense dual encoders, cross encoders, and late-interaction models.Source: §2.2 · p. 5
  • Prior listwise LLM rerankers include prompted proprietary models and adapted open-source models, while BlockRank adds task-specific attention restructuring and explicit training of attention patterns.The passage also notes efficiency work using Seq2Seq architectures and single-token decoding.Source: §2.2 · p. 5
  • This paper ranks hard top-k candidates from a first-stage retriever rather than random corpus subsets, a setting the authors describe as more difficult and more practically faithful.The cited prior study evaluated random subsets, while the paper emphasizes semantically similar candidates.Source: §2.2 · p. 5
  • BlockRank applies structured sparsity to the logical prompt structure—separating instructions, documents, and query—rather than using general-purpose long-context patterns.The passage positions it within efficient structured-attention approaches such as sliding-window and global-local patterns.Source: §2.2 · p. 5

3. Emergence of Structured Attention in In-context Ranking

The attention analysis identifies sparse interactions across document blocks and query tokens that focus on relevant documents. These relevance signals strengthen in middle layers.

Analysis Setup

  • The analysis uses a Mistral-7B-v0.3 model fine-tuned on MSMarco ICR examples with the standard next-token prediction objective.The attention analysis is anchored on a random sample, with additional results noted in the appendix.Source: §3 · p. 6

Observation 1: Inter-document Block Sparsity

  • Document-token attention is concentrated within each document rather than distributed uniformly across the context, motivating sparse attention over document blocks.The analysis reports stronger attention within a document than toward other context regions.Source: §3 · p. 6

Observation 2: Query-document Block Relevance

  • Certain query tokens, including delimiters and end-of-prompt tokens, attend more strongly to the ground-truth relevant document than to irrelevant documents.The example visualizes this pattern at Layer 18.Source: §3 · p. 6
  • The query-to-document discriminative signal is weak initially, strengthens around layers 8–24, and persists or slightly diffuses in final layers.Source: §3 · p. 6

4. BlockRank: Blockwise In-context Ranking

BlockRank combines structured sparse attention, a contrastive objective that tunes query-to-document attention, and attention-based inference. Its components exploit document-block sparsity and relevance signals identified in the attention analysis.

4. BlockRank: Blockwise In-context Ranking

  • BlockRank combines a structured attention mechanism, an auxiliary attention loss, and an alternative attention-based inference method.Source: §4 · p. 7

4.1. Blockwise Structured Attention

  • Document tokens attend only within their own document and to instruction tokens, while query tokens attend across the prompt; this structure reduces attention complexity from quadratic to linear.Instruction tokens use causal attention within the instruction segment.Source: §4.1 · p. 7
  • Shared local document positions promote order-invariant document representations, while query positions use a fixed offset of 8192.Instruction positions remain sequential; each document starts its positions immediately after the instruction segment.Source: §4.1 · pp. 7–8

4.2. Auxiliary Attention Loss (Laux)

  • The auxiliary InfoNCE loss at a selected middle layer encourages signal-carrying query tokens to assign higher attention scores to relevant than irrelevant documents.The combined training objective uses next-token prediction loss plus the auxiliary loss, with λ=0.1 in the experiments.Source: §4.2 · p. 8

4.3. Efficient Attention-Based Inference

  • At inference, BlockRank can score candidate documents from a partial forward pass to the target middle layer and return the top-ranked document identifiers.It uses the same signal-carrier attention scoring procedure as the auxiliary loss.Source: §4.3 · p. 9

5. Experimental Results

The experiments assess BlockRank’s retrieval effectiveness, efficiency, and scalability against state-of-the-art rerankers and a controlled Full-FT baseline. Ablations examine its training losses and inference methods.

5.1. Experimental Setup

  • The evaluation covers 11 BEIR datasets with top-100 Contriever candidates and in-domain MSMarco and NQ lists sized from N=10 to 500.Effectiveness is measured by nDCG@10 on BEIR and P@1 and MRR@10 on MSMarco and NQ; efficiency is end-to-end inference latency.Source: §5.1 · p. 9, §5 · p. 9

5.2. Main Performance Comparison

  • On BEIR, BlockRank Mistral scores 54.8 versus FIRST’s 54.3; in controlled in-domain tests, it consistently outperforms Full-FT Mistral.The BEIR result is reported for MSMarco-trained BlockRank; Table 2 compares BlockRank and Full-FT trained on the same data.Source: Table 1 · p. 10, §5.2 · pp. 10–11, Table 2 · p. 11
  • At N=100, BlockRank is approximately 4.7× faster; its latency reaches 1.15s at N=500 while P@1 remains 28.7%, versus Full-FT’s approximately 26.7%.BlockRank P@1 peaks around 29.2% at N=200; the passage reports latency scaling linearly with N.Source: Figure 4 · p. 11, §5.2 · p. 11

5.3. Ablation Studies

  • Adding L_aux raises BlockRank Attn Prec@1 from 27.8 to 29.1, while removing L_NTP drops Decode Prec@1 to 15.8.The full configuration achieves the highest reported Attn Prec@1.Source: §5.3 · pp. 11–12, Table 3 · p. 11
  • At N=50, BlockRank attention-based inference achieves MRR@10 42.0; autoregressive decoding yields P@1 28.7 for both BlockRank and Full-FT.The comparison covers effectiveness and latency on MSMarco.Source: §5.3 · p. 12, Table 4 · p. 12

6. Conclusion

The paper presents BlockRank as an efficient, scalable ICR approach, while noting limits to the evidence and potential societal risks. The authors call for further investigation of direct attention inference across diverse tasks.

6. Conclusion

  • BlockRank combines task-specific sparse attention for linear complexity with a contrastive auxiliary loss that optimizes internal attention signals for relevance.The conclusion reports that experiments on MSMarco and NQ show Mistral-7B matching or surpassing standard fine-tuning effectiveness while being more efficient at inference and training.Source: §6 · p. 12
  • The evidence is primarily for one model architecture; robustness of learned attention signals for direct inference across highly diverse tasks needs more investigation.Source: §6 · p. 12

A. Societal Impact

  • The authors suggest that more efficient retrieval could reduce energy use and support research, education, decision-making, and access in resource-constrained settings.These are proposed potential benefits, not reported experimental outcomes.Source: A. Societal Impact · p. 15
  • LLMs may inherit or amplify training-data biases, motivating transparency and mechanisms to identify and mitigate harmful or misleading content.Source: A. Societal Impact · p. 15

B. Dataset and Hyperparameter Details

The appendix describes dataset and hyperparameter details for reproducibility.

B.1. Datasets

The appendix specifies the retrieval datasets and fine-tuning settings for BlockRank and Full-FT, including BlockRank-specific parameters.

B.1. Datasets

  • MSMarco v1 has 8.8M passages and about 500K training queries; NQ320K has about 320K passages and 300K training queries.The listed validation sets contain 6,980 MSMarco queries and 7,830 NQ queries.Source: §B.1 · p. 15

B.2. Fine-tuning Details (BlockRank and Full-FT)

  • Both models use Adafactor at 3 × 10^-7, global batch size 32, and one epoch; Full-FT uses next-token prediction loss on answer tokens.The schedule uses linear warmup followed by cosine decay.Source: §B.2 · p. 15

B.3. BlockRank Specific Hyperparameters

  • BlockRank uses auxiliary-loss weight λ=0.1, InfoNCE temperature τ=0.05, signal tokens [":", "[‘"], and middle layer l*=20.Chunk length is 160 for MSMarco and 384 for NQ.Source: §B.3 · p. 16

C. Attention Complexity Analysis

BlockRank’s structured attention limits each document’s attention to its own segment and the instruction, while query tokens can attend across the full context. Under the section’s equal-length segment assumption, total attention complexity scales linearly with the number of candidate documents.

C. Attention Complexity Analysis

  • With each segment treated as length L_chunk, document tokens attend within their segment and to the instruction, while query tokens attend to all documents.Instruction, document, and query segments are each assumed to have effective length L_chunk; query attention spans a context of (N + 2)L_chunk.Source: C. Attention Complexity Analysis · p. 16
  • The total cost sums instruction, document, and query attention, with the document-segment cost repeated across all N documents.The document segments contribute N·O(2L_chunk^2·d).Source: C. Attention Complexity Analysis · p. 16
  • O(N·L_chunk^2·d) is BlockRank’s dominant per-layer attention cost, compared with O(N^2·L_chunk^2·d) for a standard Transformer on a comparable long sequence.The section identifies linear scaling in N for BlockRank and quadratic scaling for the standard Transformer.Source: C. Attention Complexity Analysis · p. 17

D.1. Calibration Problem in Beam Decoding with Full-FT Model

The Full-FT model’s beam-decoded document IDs show lower digit entropy than random IDs, indicating less diversity in the generated prediction list. The analysis uses ten unique predictions per query for candidate lists of 100 documents.

D.1. Calibration Problem in Beam Decoding with Full-FT Model

  • Across 10 beam-decoded predictions, Full-FT had lower average entropy than random for id0 (2.28 ± 0.43 vs. 2.55 ± 0.25) and id1 (2.19 ± 0.46 vs. 2.66 ± 0.24).This indicates less diversity in constituent digits than random chance would suggest.Source: Table 5 · p. 17, §D.1 · p. 17
  • For 100 randomly shuffled candidates, each query yielded ten unique two-digit ID predictions, and entropy was compared with ten random unique IDs.Because IDs are assigned independently of document content, the analysis treats lower entropy as a sign of concentration in generated digits.Source: §D.1 · p. 17

D.2. Analysis of Retrieval Signals in Attention Patterns of Full-FT Mistral

Attention-based inference in Full-FT Mistral finds retrieval signals associated with specific query tokens and strongest in middle layers. The appendix also reports that duplicating the query in the prompt prefix improves Precision@1 for both Full-FT and BlockRank on MSMarco.

D.2. Analysis of Retrieval Signals in Attention Patterns of Full-FT Mistral

  • Certain query tokens, especially end-of-query or delimiter tokens, yield higher attention-based P@1, while attention-derived P@1 and MRR@10 are strongest in middle layers.Figures 5 and 6 compare query-token choices and Transformer layers; the cited analysis reports declining performance in earlier and later layers.Source: §D.2 · p. 18, Figure 5 · p. 18, Figure 6 · p. 18

D.3. Layerwise Emergence of Retrieval Signals and Choice of 𝑙∗

  • Layers 12 through 24 develop stronger attention-based Precision@1 during training, and the experiments use l*=20, with similar performance at other reasonable middle layers.Figure 7 tracks layerwise Precision@1 on a held-out MSMarco training subset.Source: §D.3 · p. 18, Figure 7 · p. 18

D.4. Impact of Including Query in Prompt Prefix

  • With the query in both the prefix and at the end, Full-FT Precision@1 rose from 27.2 to 28.7 (+1.5), while BlockRank rose from 24.2 to 28.1 (+3.9).The query-prefix format was used for the other reported experiments.Source: Table 6 · p. 19, §D.4 · p. 19

D.5. Cross-dataset Generalization

BlockRank performs best on its in-domain test set; on an unseen dataset, it remains above the zero-shot baseline but falls below its in-domain scores.

D.5. Cross-dataset Generalization

  • On out-of-domain tests, BlockRank P@1 is above the No Training baseline but considerably below its in-domain performance.Table 7 reports cross-dataset P@1 for models fine-tuned on MSMarco or NQ and tested on both test sets, alongside zero-shot Mistral-7B-instruct.Source: Table 7 · p. 19, §D.5 · p. 19
Loading 2510.05396v2…