Search papers, labs, and topics across Lattice.
This paper introduces a query-focused and memory-aware reranking framework that leverages attention scores from selected heads in large language models to estimate passage-query relevance in a listwise manner. By producing continuous relevance scores, the framework enables training on diverse retrieval datasets without Likert-scale supervision. Experiments show the method outperforms existing pointwise and listwise rerankers, achieving state-of-the-art results on the LoCoMo benchmark and demonstrating the benefits of contextual augmentation and attention head selection for efficiency.
Achieve state-of-the-art long-context reranking with a lightweight (4B parameter) model by cleverly focusing on query-relevant attention heads, outperforming much larger models.
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 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.
