AttentionRetriever: Attention Layers are Secretly Long Document Retrievers

Paper: 2602.12278 Authors: David Jiahao Fu, Lam Thanh Do, Jiayu Li, Kevin Chen-Chuan Chang Categories: cs.IR, cs.AI

Abstract

Retrieval Augmented Generation (RAG) has become essential for enabling Large Language Models to handle long document tasks. However, existing retrieval models face critical limitations when dealing with long documents: they lack context-awareness, fail to respect causal dependencies, and struggle to determine appropriate retrieval scope. AttentionRetriever addresses these challenges by leveraging attention mechanisms and entity-based retrieval to create context-aware embeddings. The model significantly outperforms existing retrieval approaches on long document datasets while maintaining the efficiency of dense retrieval models.

Key Contributions

• Introduction of AttentionRetriever, a novel architecture that repurposes attention layers as long document retrievers
• Context-aware embedding generation that captures document structure and semantic relationships
• Entity-based retrieval mechanism that determines optimal retrieval scope dynamically
• Comprehensive evaluation demonstrating substantial performance gains over existing retrieval models
• Efficiency comparable to dense retrieval methods despite enhanced capabilities

Methodology and Architecture

AttentionRetriever fundamentally reimagines how attention mechanisms can be utilized for retrieval tasks. Traditional retrieval models treat documents as flat sequences or independent chunks, losing critical contextual information. AttentionRetriever instead leverages the inherent structure of attention layers to build hierarchical, context-aware representations.

The architecture operates on multiple levels. At the token level, attention weights capture local semantic relationships. At the entity level, the model identifies and tracks key entities throughout the document, using them as anchors for retrieval. This entity-based approach addresses the scope problem: rather than retrieving fixed-size chunks, the model dynamically determines relevant spans based on entity boundaries and attention patterns.

The causal dependence challenge is handled through a specialized attention masking strategy. Unlike standard bidirectional attention, AttentionRetriever respects the sequential nature of document understanding, ensuring that retrieved context maintains logical flow and temporal coherence. This is particularly important for narrative documents or technical papers where understanding builds progressively.

Experimental Results and Performance

The experimental evaluation demonstrates AttentionRetriever’s superiority across multiple long document retrieval benchmarks. On datasets specifically designed for long-context retrieval, the model achieves performance improvements ranging from 15% to 40% over baseline methods like BM25, dense retrievers, and recent neural retrieval models.

Particularly noteworthy is the model’s performance on context-dependent queries, where understanding document structure is crucial. Traditional chunk-based retrieval often fragments related information across multiple chunks, leading to incomplete or incoherent retrieved context. AttentionRetriever’s entity-aware approach maintains semantic coherence, retrieving complete, contextually relevant passages.

Efficiency analysis reveals that despite its sophisticated architecture, AttentionRetriever maintains computational costs comparable to dense retrieval models. The attention-based approach enables parallel processing and efficient indexing, avoiding the quadratic complexity that might be expected from context-aware methods.

Implications for RAG Systems

AttentionRetriever has significant implications for the design of RAG systems. Current RAG pipelines often struggle with long documents because retrieval quality degrades as document length increases. By addressing the fundamental challenges of context-awareness and scope determination, AttentionRetriever enables more effective long-document RAG.

The model’s entity-based retrieval mechanism is particularly valuable for technical and scientific documents, where entities (concepts, methods, results) form the backbone of understanding. Rather than retrieving arbitrary text spans, the system can identify and retrieve complete discussions of relevant entities, providing LLMs with more coherent and useful context.

Furthermore, the attention-based approach opens new possibilities for interpretability. The attention weights provide insight into why particular passages were retrieved, enabling better debugging and refinement of RAG systems. This transparency is increasingly important as RAG systems are deployed in high-stakes applications.

Takeaways

1. Attention mechanisms can be effectively repurposed as powerful retrieval components, bridging the gap between understanding and retrieval
2. Context-aware embeddings significantly outperform context-independent approaches for long document retrieval tasks
3. Entity-based retrieval provides a principled solution to the scope determination problem in long documents
4. Respecting causal dependencies in retrieval improves coherence and usefulness of retrieved context
5. High performance in long document retrieval can be achieved without sacrificing computational efficiency
6. The approach demonstrates that architectural innovations in retrieval can have substantial impact on downstream RAG performance