AttentionRetriever: Attention Layers are Secretly Long Document Retrievers

TL;DR

AttentionRetriever introduces an attention-based, entity-aware retrieval model that significantly improves long document retrieval accuracy for LLMs, addressing key challenges like context-awareness and scope of retrieval.

Contribution

The paper presents a novel attention and entity-based retrieval model specifically designed for long documents, outperforming existing models in accuracy and efficiency.

Findings

Outperforms existing retrieval models on long document datasets

Maintains efficiency comparable to dense retrieval models

Effectively addresses context-awareness and scope challenges

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.