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

Abstract: Built upon the existing analysis of retrieval heads in LLMs, 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 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 (e.g., 4B 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 that assesses the capabilities of 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.

• The paper introduces QRRanker, a generation-free listwise reranker that trains middle-layer attention heads to assign continuous relevance scores across up to 50 retrieved candidates.
• QRRanker-4B achieves 77.37 R@5 on MuSiQue, 96.90 on HotpotQA, and a 72.13 average retrieval score, outperforming larger rerankers such as GroupRank-32B.
• On LoCoMo, reranking only the top three dialogue chunks reaches about 57.3 Overall F1 while using roughly 854 tokens, indicating that stronger retrieval can outperform more complex memory systems efficiently.

Motivation and positioning

The paper addresses a structural limitation of embedding-based retrieval: fixed-dimensional vectors cannot fully encode the combinatorial complexity of query–document interactions, and cosine similarity is ill-suited to relations such as causality or analogy. Rerankers mitigate this, but existing LLM-based rerankers face a trade-off. Pointwise cross-encoders score documents independently and lose the global view of the candidate shortlist; listwise generative rerankers see the whole list but must emit discrete Likert-scale judgments through next-token prediction, which constrains both score granularity and the pool of usable training data. The authors build on prior mechanistic analyses of retrieval heads and Query-focused Retrieval (QR) heads—attention heads whose query-to-context attention weights correlate with passage relevance—and move from probing these heads to explicitly training them for ranking.

Method

QRRanker is a listwise, generation-free reranker built on Qwen3-4B-Instruct-2507. Sixteen QR heads are precomputed by QR score on 1,000 NarrativeQA samples; notably, all selected heads lie in middle layers (17–24). At inference, the prompt concatenates up to 50 retrieved candidates with the question; the relevance score of chunk $c_i$ is the sum over QR heads of the mean attention mass from query tokens to $c_i$ tokens. Because scores are continuous attention aggregates rather than generated tokens, the model trains on arbitrary binary-relevance data without Likert supervision, and inference requires only a prefill pass.

Training uses a group contrastive loss: per-sample scores are max–min normalized (with scale factor 8) to stabilize against attention-sink sensitivity in instruction tokens, and all positives in the top-50 candidate set contribute simultaneously rather than one at a time. Training instances combine MuSiQue (official supporting facts) and NarrativeQA (silver evidence following Mindscape), with candidates from Qwen3-Embedding-8B. An optional memory prefix prepends summaries mapped from the retrieved chunks—block-based summaries for books, event-centric structured summaries with source-utterance links for dialogues—capped at 512 tokens at inference.

Main results

Across five datasets spanning Wikipedia multi-hop QA (MuSiQue, HotpotQA), long narrative QA (NarrativeQA, DetectiveQA), and dialogue memory (LoCoMo), QRRanker-4B outperforms pointwise and listwise baselines of comparable or larger size:

Method	MuSiQue R@5	HotpotQA R@5	NarrativeQA R@10	DetectiveQA R@10	Overall avg@10
Qwen3-Embedding-8B	62.55	89.05	32.39	31.17	57.80
Qwen-Reranker-4B (trained)	69.71	93.95	49.62	51.25	68.82
GroupRank-32B	65.08	90.60	48.83	51.38	66.95
QRHeads-4B (out-of-box)	71.22	94.80	48.89	45.58	67.59
QRRanker-4B	77.37	96.90	54.93	53.76	72.13

Two claims stand out. First, a 4B attention-based ranker exceeds GroupRank-32B, an RL-trained reasoning reranker eight times its size, which the authors attribute to better modeling of inter-passage dependencies within the shortlist. Second, on Wikipedia QA it surpasses HippoRAG-v2's graph-based pipeline (e.g., 77.37 vs. 74.70 R@5 on MuSiQue), suggesting that simple retrieve-and-rerank can displace elaborate graph construction when the reranker is strong.

Downstream QA confirms the retrieval gains transfer: on NarrativeQA, F1 rises from 30.51 (trained Qwen3-Reranker-4B) to 33.61; DetectiveQA accuracy improves from 62.85 (SFT-Embedding-8B) to 67.25.

LoCoMo and the memory-system comparison

The most consequential result concerns LoCoMo. Using only raw dialogue chunks—reranking the top-50 from SFT-Embedding-8B and feeding just the top-3 (~854 tokens)—QRRanker achieves 57.03 Overall F1 with GPT-4o-mini and 57.32 with GPT-5-mini, exceeding dedicated memory frameworks including MemoryOS, Zep, Mem0, Nemori, TiMem, Synapse, Membox, CompassMem, ES-Mem, and SimpleMem. Several of these systems consume 2,000–20,000 tokens per query to maintain explicit memory stores or graphs. Retrieval-side, QRRanker reaches 87.34 R@3 versus 85.93 for untrained QR heads and 76.01 for SFT-Embedding-8B (+11.33). The implication is direct: a stronger search over minimally processed history can beat sophisticated memory management architectures, at least on this benchmark. This claim carries a caveat—the evaluation uses only top-3 chunks as generator context, so it tests evidence localization rather than full conversational state tracking—but the token-efficiency advantage is substantial regardless.

Contextual information and head selection

The summary-prefix extension helps where evidence is diffuse but not where it is localized: Recall@3 improves by +2.67 on DetectiveQA, +1.02 on NarrativeQA, and +0.70 on LoCoMo, yet slightly degrades HotpotQA (−0.30) and Musique (−0.03). A hierarchical clustering-tree prefix over Wikipedia passages yields no gains, indicating abstract global summaries add little when gold evidence is highly localized.

A layer-level ablation shows the framework is robust to head choice. Variants selecting 16 trainable heads from layers 17–24 or 28–35 match full QRRanker on NarrativeQA (R@10 of 54.28 and 54.65 vs. 54.44), while low layers (10–17) drop substantially (49.91). Strikingly, overlap between trained variant heads and the original QR heads is low, implying QR training activates latent retrieval capability in non-QR heads. Because trained heads concentrate in layers 17–24, truncating the model after layer 24 preserves performance while cutting cost.

Efficiency

On 20-query latency measurements, QRRanker achieves 1095 ms P50 versus 1222 ms for Qwen3-Reranker-4B (batch=50), with lower compute (82.74 vs. 115.69 TFLOPs/query). The truncated middle-layer variant reaches 910 ms P50 and 69.83 TFLOPs with no reported performance loss, making it the lightest configuration evaluated.

Limitations and open questions

Several constraints deserve note. The framework depends on a bi-encoder first stage; its behavior when the initial top-50 misses gold evidence entirely is not characterized. Silver evidence construction for NarrativeQA inherits noise from the Mindscape procedure, and the LoCoMo comparison restricts generation context to three chunks, which may understate what memory systems offer for questions requiring aggregated state. The summary-prefix benefit is shown to be domain-dependent, but no principled criterion decides when to apply it. Finally, the finding that arbitrary middle-layer heads can be trained into retrievers raises an open question about whether explicit QR-head identification is necessary at all, or whether any mid-layer initialization suffices.

Conclusion

QRRanker demonstrates that directly optimizing a small set of attention heads for ranking yields a listwise reranker with continuous scores, freedom from Likert-scale supervision, and generation-free inference. With a 4B backbone it sets new results across Wikipedia, narrative, and dialogue benchmarks—including the best reported Overall F1 on LoCoMo using fewer than 900 tokens of context—and supports efficiency-oriented truncation to middle layers. Its central empirical message is that retrieval quality, rather than memory architecture complexity, may be the binding constraint in long-context dialogue systems.