Query-Focused Retrieval Heads Improve Long-Context Reasoning and Re-ranking

要約

最近の研究では、ヘイスタックの針タスクでのコピーパステの動作によって測定されるように、長いコンテキスト言語モデル（LMS）の顕著な情報を取得する責任のある注意ヘッドのサブセットである検索ヘッド（Wu et al。、2025b）が特定されています。
このペーパーでは、長いコンテキストから検索を強化する注意ヘッドの改善されたセットであるQRHead（クエリ中心の検索ヘッド）を紹介します。
実際のタスク（長いコンテキストQAなど）のいくつかの例を使用して、入力クエリに関して注意スコアを集約することによりQRHEADを識別します。
さらに、QR-Headの蓄積された注意質量を検索スコアとして使用する効率的かつ効果的なレトリバーであるQR- Retrieverを紹介します。
最高の検索スコアを持つ最も関連性の高いパーツを選択することにより、長いコンテストの推論にQR-レトリバーを使用します。
Multi-Hopの推論タスクLongmemevalとClipperでは、これにより、完全なコンテキストで10％以上のパフォーマンスが得られ、強い密なレトリバーよりも優れています。
また、QRETRIEVERをBeirベンチマークの再ランカーとして評価し、RankGPTなどの他のLLMベースの再ランカーを上回る強力なゼロショットパフォーマンスを達成することを発見しました。
さらなる分析により、QuryContextの注意スコアリングとタスク選択の両方が、QRHEADを強力な下流のユーティリティで識別するために重要であることが示されています。
全体として、私たちの仕事は汎用レトリバーに貢献し、LMSの長いコンテスト機能に関する解釈可能性の洞察を提供します。

要約(オリジナル)

Recent work has identified retrieval heads (Wu et al., 2025b), a subset of attention heads responsible for retrieving salient information in long-context language models (LMs), as measured by their copy-paste behavior in Needle-in-a-Haystack tasks. In this paper, we introduce QRHEAD (Query-Focused Retrieval Head), an improved set of attention heads that enhance retrieval from long context. We identify QRHEAD by aggregating attention scores with respect to the input query, using a handful of examples from real-world tasks (e.g., long-context QA). We further introduce QR- RETRIEVER, an efficient and effective retriever that uses the accumulated attention mass of QRHEAD as retrieval scores. We use QR- RETRIEVER for long-context reasoning by selecting the most relevant parts with the highest retrieval scores. On multi-hop reasoning tasks LongMemEval and CLIPPER, this yields over 10% performance gains over full context and outperforms strong dense retrievers. We also evaluate QRRETRIEVER as a re-ranker on the BEIR benchmark and find that it achieves strong zero-shot performance, outperforming other LLM-based re-rankers such as RankGPT. Further analysis shows that both the querycontext attention scoring and task selection are crucial for identifying QRHEAD with strong downstream utility. Overall, our work contributes a general-purpose retriever and offers interpretability insights into the long-context capabilities of LMs.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google