PaperlayerView exact revision on arXiv

Source-linked AI summary

Query-focused and Memory-aware Reranker for Long Context Processing

Yuqing Li, Jiangnan Li, Mo Yu, Guoxuan Ding, Yanyu Chen, Zheng Lin, Wei Zhang, Jie Zhou

arXiv:2602.12192v3cs.CL

TL;DR

Embedding and conventional reranking methods have limited query-document interaction modeling or lose global shortlist context. QRRanker trains selected QR attention heads as a lightweight listwise scorer with continuous relevance values, and it achieves strong cross-domain reranking and downstream QA results while supporting memory context and efficient inference. Its behavior on larger models remains unexplored, and some supervision uses silver evidence that may introduce label noise.

  • Problem

    Embedding similarity has a geometric bottleneck, while pointwise and listwise rerankers respectively lose shortlist context or face score-prediction limitations.

  • Method

    QRRanker trains selected query-focused retrieval heads to produce continuous passage relevance scores through listwise attention-based reranking.

  • Results

    QRRanker consistently improves reranking quality and downstream QA performance across five datasets spanning Wikipedia QA, long-context story QA, and dialogue memory.

  • Takeaways & Limitations

    The framework remains practical with small backbones, avoids generation during inference, and supports summary prefixes and mid-layer heads for added context and efficiency.

  • Takeaways & Limitations

    Behavior on larger-scale models remains unexplored, and some training supervision relies on silver evidence that may introduce label noise.

Abstract

from arXiv · show

Built upon the existing analysis of retrieval heads in large language models, we propose an alternative reranking framework that trains models to estimate passage-query relevance using the attention scores of selected heads. This approach provides a listwise solution that leverages the holistic information within the entire candidate shortlist during ranking. At the same time, it naturally produces continuous relevance scores, enabling training on arbitrary retrieval datasets without requiring Likert-scale supervision. Our framework is lightweight and effective, requiring only small-scale models, such as 3B parameters, to achieve strong performance. Extensive experiments demonstrate that our method outperforms existing state-of-the-art pointwise and listwise rerankers across multiple domains, including Wikipedia and long narrative datasets. It further establishes a new state-of-the-art on the LoCoMo benchmark, which assesses dialogue understanding and memory usage. We further demonstrate that our framework supports flexible extensions. For example, augmenting candidate passages with contextual information further improves ranking accuracy, while training attention heads from middle layers enhances efficiency without sacrificing performance.

1 Introduction

The paper addresses limitations of embedding and reranking methods by training selected retrieval heads to score passages listwise. QRRanker provides continuous relevance scores, supports small backbones and contextual extensions, and performs strongly across diverse retrieval settings.

  • Embedding vectors face a geometric bottleneck and limited similarity inductive bias for fine-grained query-document relationships such as causality, association, and analogy.
  • Pointwise rerankers provide scores but lose the shortlist’s global view, whereas generation-based listwise methods retain holistic context but face next-token prediction limitations.
  • QRRanker trains a small set of retrieval heads to optimize passage ranking, producing a listwise reranker with continuous relevance scores and no Likert-scale supervision requirement.
  • QRRanker can train effectively with backbones as small as 3B parameters, incorporate shared global context, and use middle-layer heads to reduce latency without performance loss.
  • Across narrative QA, dialogue memory, Wikipedia multi-hop QA, and reasoning-intensive retrieval, QRRanker outperforms strong general-purpose and domain-specific rerankers.

2 Related Work

Prior reranking systems trade representation quality, computational cost, shortlist context, and score flexibility. QRRanker instead uses query-focused retrieval heads for explicit listwise scoring, complementing existing retrieval and memory approaches.

  • Reranking: Bi-encoders efficiently retrieve candidates but cannot fully model fine-grained query-document interactions, while cross-encoders improve interaction modeling at higher cost.
  • Reranking: LLM-based rerankers are commonly pointwise or listwise, with pointwise scoring documents independently and listwise methods jointly considering candidate passages.
  • Retrieval heads: QRRanker trains selected query-focused retrieval heads as an explicit listwise scorer that produces fine-grained document scores from query-positive-negative supervision.
  • Memory utilization: Memory-utilization research constructs summaries, graphs, trees, and memory systems to retrieve relevant content from long stories and dialogues.
  • Retrieval heads: Retrieval and QR scores in Figure 1 are computed from attention scores of a QR attention head, with Doc2 shown as the gold chunk.

3 Preliminaries: QR-head

QR-heads are attention heads whose question-to-context attention reflects passage relevance. The paper scores and selects heads using their attention to gold chunks, then aggregates selected heads for passage scoring.

  • QR-heads are self-attention heads that attend strongly to context spans relevant to a question during question encoding.
  • For question Q and context chunks C, the QR score aggregates an attention head’s attention from question tokens to gold chunks G.
  • A higher QR score indicates stronger attention to gold chunks, so average scores on a seed set rank heads for selection.
  • The method selects the top 16 QR-heads for Qwen3-4B using 1,000 NarrativeQA training samples and applies the procedure to Llama-3.2-3B on MuSiQue.
  • A candidate chunk score is computed by aggregating its question-to-chunk attention over the selected QR-heads.

4 Method

QRRanker reranks all retrieved candidates in one pass by aggregating selected QR-head attention, and trains those heads with listwise contrastive supervision. Optional summary prefixes supply global context for long-context understanding.

  • QRRanker: QRRanker processes all candidate documents in one inference pass and derives scores from attention patterns instead of generating ranking text.
  • Training data: Training combines MuSiQue and NarrativeQA, using official supporting facts for MuSiQue and silver evidence chunks for NarrativeQA.
  • Training data: For each question, the system retrieves a top50 candidate set and labels evidence-matching candidates positive and the remaining candidates negative.
  • Summary construction: An optional summary prefix is prepended to candidate chunks, supplying coarse-grained global context alongside fine-grained evidence.
  • Summary construction: Narrative summaries use blocks of 20 consecutive chunks, while dialogue summaries extract structured events linked to source utterances.
  • Scoring: The model extracts question-to-candidate attention from each selected QR head and sums head-level retrieval scores to obtain final passage scores.
  • Scoring: Max-pooling attention can also compute candidate scores and achieves similar performance to the aggregation approach.
  • QR training: Max-min normalization stabilizes sample-dependent score ranges, while group contrastive loss jointly optimizes multiple positive documents.

5 Experimental Setup

The evaluation spans Wikipedia multi-hop QA, long-context story QA, dialogue memory, and additional retrieval-agent comparisons, using broad baseline sets and fixed candidate pools. Experiments also specify the model backbone, benchmark protocols, and Recall@k table conventions.

  • Datasets: Experiments cover Wikipedia multi-hop QA, long-context story QA, and long-term dialogue memory across multiple benchmark datasets.The setup includes HotpotQA, MuSiQue, NarrativeQA, DetectiveQA, and LoCoMo.
  • Evaluation Protocol: Wikipedia QA uses HippoRAG corpus and test splits with consistent candidate passage pools for fair comparison.The Wikipedia benchmarks are HotpotQA and MuSiQue.
  • Evaluation Protocol: NarrativeQA contains 1,272 questions with documents reaching 518k tokens, while DetectiveQA averages over 100k tokens and requires scattered-evidence localization.These datasets evaluate reasoning over extended contexts.
  • Evaluation Protocol: LoCoMo contains 50 multi-session dialogues across 10 user groups, averaging approximately 9,000 tokens per dialogue.Performance is reported for single-hop, multi-hop, temporal-reasoning, and open-domain categories.
  • Baselines: The study compares QRRanker with embedding models, general-purpose rerankers, domain-specific memory methods, and retrieval agents.The comparisons include Qwen3-based models, HippoRAG, GroupRank, ReasonRank, LoCoMo memory systems, DCI-Agent, and Nemo Retriever Agent.
  • Metrics and Implementation: Table 1 reports Recall@k after reranking top-50 candidates, with first-stage R@50 ceilings varying across Musique, HotpotQA, NarrativeQA, and DetectiveQA.Wikipedia candidates come from Qwen3-Embedding-8B, story candidates from SFT-Embedding-8B, and DetectiveQA scores average English and Chinese sets.

6 Results

QRRanker consistently improves retrieval and downstream question-answering across Wikipedia, long-story, and dialogue-memory settings. It also benefits from contextual prefixes and supports efficient middle-layer truncation while retaining competitive performance.

  • Main Results: QRRanker achieves the best overall results across five English and Chinese datasets, improving both retrieval quality and downstream task performance.The main evaluation uses Qwen3-4B-Instruct-2507, with an additional Llama-3.2-Instruct variant reported for backbone robustness.
  • Main Results: QRRanker establishes a new state-of-the-art and surpasses Qwen-Reranker-4B while improving average recall across retrieval benchmarks.It also outperforms graph-based HippoRAG on MuSiQue and HotpotQA and remains competitive with much larger 32B rerankers.
  • Long-context Story QA Performance: 33.61 F1 on NarrativeQA exceeds Qwen3-Reranker-4B's 30.51, while DetectiveQA rises from 62.85 with SFT-Embedding-8B to 67.25 with QRRanker.The reported gains indicate that selected evidence is aligned with reasoning needed for answer generation.
  • Dialogue Memory Performance: 57.03 Overall F1 with GPT-4o-mini and 57.32 with GPT-5-mini are achieved using only 854 average tokens from the raw dialogue history.QRRanker reranks 50 retrieved dialogue chunks and passes only the top three to the generator, without additional memory mechanisms.
  • Results with Contextual Information: A summary prefix consistently improves ranking on long-dialogue and long-context story benchmarks by complementing retrieved chunks with global contextual guidance.The same strategy provides no gains and can degrade performance on Wikipedia multi-hop QA, where evidence is highly localized.
  • Heads from Different Layer-Levels: Middle-layer heads from layers 17–24 preserve performance close to QRRanker, whereas lower-layer heads from layers 10–17 cause a clear drop.The middle-layer variant motivates truncation after layer 24, which achieves the lowest reported latency and resource cost.

7 Conclusion

QRRanker is a lightweight listwise reranking framework that trains selected Query-focused Retrieval heads to score candidates from attention patterns. It improves reranking and downstream QA across diverse datasets while supporting efficient extensions.

  • QRRanker trains selected Query-focused Retrieval heads to produce real-valued relevance scores in a single-pass listwise reranking framework.It reranks without generation and avoids generation-format errors.
  • QRRanker consistently improves reranking quality and downstream QA performance across five datasets spanning Wikipedia QA, long-context story QA, and dialogue memory.
  • QRRanker remains practical with a small backbone and offers inference-efficiency benefits.
  • Optional summary prefixes provide global context, while mid-layer head selection offers further efficiency improvements.

Limitations

The paper identifies two main limitations: evaluation is restricted to two backbone architectures, and some supervision uses silver rather than fully human-annotated evidence.

  • QRRanker’s behavior on larger-scale models such as 14B+ remains unexplored.Larger models may provide more retrieval-sensitive head candidates, but this is left for future work.
  • Some training supervision relies on silver evidence because narrative QA benchmarks lack fine-grained human evidence annotations.Partially relevant passages may therefore be absent from the constructed evidence set, introducing label noise.
  • QRRanker remains reasonably robust under the weak-supervision setting despite this potential label noise.

B Position Bias Analysis

The position-bias analysis tests whether QRRanker relies on candidate order rather than semantic relevance. Shuffling candidates leaves QRRanker robust and produces stable rankings across permutations.

  • Position Bias Analysis: The experiment compares original retriever order, randomly shuffled order, and QRRanker reranking under both input orders.Each condition uses the same top-50 candidate set for a query.
  • Position Bias Analysis: The random-order baseline collapses to near-chance levels after shuffling, whereas QRRanker remains robust.
  • Position Bias Analysis: 3.40 R@3 is the maximum QRRanker drop on HotpotQA after shuffling.
  • Position Bias Analysis: After shuffling, QRRanker still substantially outperforms the original-order retriever, indicating content-level relevance matching rather than positional reliance.
  • Position Bias Analysis: Kendall’s τ = 0.638 ± 0.100 and Spearman’s ρ = 0.808 ± 0.096 across LoCoMo permutations indicate high ranking agreement.

D One-step QRRanker Reranking v.s. Advanced Retrieval Agents

The paper compares QRRanker with iterative retrieval agents and presents it as an efficient one-step reranking alternative. QRRanker is competitive across task groups, while the analysis identifies complementary strengths and future integration opportunities.

  • One-step QRRanker Reranking v.s. Advanced Retrieval Agents: QRRanker uses a standard one-step embedding retrieval → rerank pipeline instead of iterative agent retrieval.It can serve as a tool-calling step without requiring repeated interaction.
  • One-step QRRanker Reranking v.s. Advanced Retrieval Agents: QRRanker achieves the best average results on BRIGHT and StoryQA against the compared retrieval agents without LLM interactive synergy.
  • One-step QRRanker Reranking v.s. Advanced Retrieval Agents: DCI-Agent performs strongly on WikipediaQA but shows a disparity across task types, partly because grep can precisely match explicit entity names.
  • One-step QRRanker Reranking v.s. Advanced Retrieval Agents: The comparison covers StoryQA, WikipediaQA, dialogue memory, and BRIGHT dataset groups, with resource consumption summarized separately.
  • One-step QRRanker Reranking v.s. Advanced Retrieval Agents: Local search is identified as a future complement to QRRanker for deeper, fine-grained evidence mining.The paper notes that grep verifies localized spans rather than whole chunks, trading recall for finer evidence exploitation.
  • One-step QRRanker Reranking v.s. Advanced Retrieval Agents: The authors state that QRRanker should enhance rather than replace retrieval agents by supporting retrieval coverage and deeper evidence digging.

G Cross-Architecture Results with Llama-3.2-3B-Instruct

QRRanker transfers effectively to the Llama-3.2-3B-Instruct architecture, improving retrieval across three out-of-domain datasets and outperforming pretrained QRHeads.

  • QRRanker on Llama-3.2-3B-Instruct consistently outperforms the base retriever and pretrained QRHeads on HotpotQA, LoCoMo, and DetectiveQA.The gains are especially large on LoCoMo and DetectiveQA, where initial retrieval quality is weaker.
  • QRRanker surpasses QRHeads at all recall cutoffs, indicating that the QR training objective transfers across model architectures.The transfer improves both multi-hop retrieval and long-context narrative retrieval under out-of-domain evaluation.

H Validation of QR Head Selection

The validation compares informed and random head selection, showing that discovered QR heads already have stronger retrieval behavior before training and improve further after QR training.

  • Layer selection study: The semi-auto selection study trains QRRanker and a NarrativeQA-only variant to evaluate which layers can provide QR-like heads.The experiment explores layer ranges for selecting and training QR-like heads.
  • Head selection validation: 95.05 R@3 is reached after QR training, demonstrating complementary benefits from informed head selection and subsequent training.The validation compares head-selection strategies before QR training and reports the further increase after training.
  • Semi-auto head selection: The semi-auto variant selects 16 heads across a chosen layer range by using a gate that chooses n heads for each sample.The selection constraint sets n = 16/(le − ls) heads per layer, with 0 < ls < le ≤ 36.
  • Semi-auto head selection: The gate computes head scores from concatenated query states, selects the top-n heads, and uses their normalized scores in retrieval-score computation.The selected heads participate in retrieval scoring, while the head scores support backward gradients during QR training.
Loading 2602.12192v3…