AttentionRetriever: Attention Layers are Secretly Long Document Retrievers
2026-02-12 • Information Retrieval
Information RetrievalArtificial Intelligence
AI summaryⓘ
The authors noticed that current methods for finding information in long documents don't work very well because they miss important details and connections. They created AttentionRetriever, a new tool that uses attention and looks at key entities to understand and find the right parts of long texts. Their experiments showed that AttentionRetriever finds information more accurately and efficiently than older methods. This helps large language models handle long documents better.
Retrieval Augmented GenerationLarge Language ModelsLong Document RetrievalAttention MechanismEntity-based RetrievalContext-AwarenessDense RetrievalEmbeddings
Authors
David Jiahao Fu, Lam Thanh Do, Jiayu Li, Kevin Chen-Chuan Chang
Abstract
Retrieval augmented generation (RAG) has been widely adopted to help Large Language Models (LLMs) to process tasks involving long documents. However, existing retrieval models are not designed for long document retrieval and fail to address several key challenges of long document retrieval, including context-awareness, causal dependence, and scope of retrieval. In this paper, we proposed AttentionRetriever, a novel long document retrieval model that leverages attention mechanism and entity-based retrieval to build context-aware embeddings for long document and determine the scope of retrieval. With extensive experiments, we found AttentionRetriever is able to outperform existing retrieval models on long document retrieval datasets by a large margin while remaining as efficient as dense retrieval models.