MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval Augmentation

Abstract: Processing long contexts presents a significant challenge for LLMs. While recent advancements allow LLMs to handle much longer contexts than before (e.g., 32K or 128K tokens), it is computationally expensive and can still be insufficient for many applications. Retrieval-Augmented Generation (RAG) is considered a promising strategy to address this problem. However, conventional RAG methods face inherent limitations because of two underlying requirements: 1) explicitly stated queries, and 2) well-structured knowledge. These conditions, however, do not hold in general long-context processing tasks. In this work, we propose MemoRAG, a novel RAG framework empowered by global memory-augmented retrieval. MemoRAG features a dual-system architecture. First, it employs a light but long-range system to create a global memory of the long context. Once a task is presented, it generates draft answers, providing useful clues for the retrieval tools to locate relevant information within the long context. Second, it leverages an expensive but expressive system, which generates the final answer based on the retrieved information. Building upon this fundamental framework, we realize the memory module in the form of KV compression, and reinforce its memorization and cluing capacity from the Generation quality's Feedback (a.k.a. RLGF). In our experiments, MemoRAG achieves superior performances across a variety of long-context evaluation tasks, not only complex scenarios where traditional RAG methods struggle, but also simpler ones where RAG is typically applied.

• The paper introduces a dual-system RAG framework that leverages global memory to enhance long-context query processing.
• It employs a lightweight LLM for global memory creation and a heavy LLM for answer generation, optimizing retrieval accuracy.
• Experiments on UltraDomain benchmarks show superior performance in context-rich tasks compared to standard RAG methods.

MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval Augmentation

Introduction

MemoRAG introduces a novel paradigm in Retrieval-Augmented Generation (RAG) designed to address limitations posed by long context processing in LLMs. Traditional RAG systems struggle with tasks involving ambiguous information needs and unstructured knowledge due to their dependency on relevance matching between explicit queries and structured information. MemoRAG overcomes these limitations by integrating a dual-system architecture that enhances retrieval capabilities through global memory construction and utilization. This approach enables efficient query resolution and information retrieval from extensive and complex databases.

The MemoRAG Architecture

MemoRAG employs a dual-system architecture consisting of a light-weight, long-range LLM to establish a global memory representation of the database, and a more expressive, computationally intensive LLM to generate the final answer. This design enables MemoRAG to process highly implicit queries by first generating draft answers through the global memory context, which guides the retrieval tools in identifying relevant database information.

Figure 1: Comparison between Standard RAG and MemoRAG in processing queries requiring high-level understanding across an entire database.

The global memory aids the retrieval of scattered and implicit information by recalling relevant cues, thus facilitating comprehensive answers to complex queries—a capability that is prominently displayed when juxtaposed with traditional RAG methods, as illustrated in Figure 1.

Memory Module Implementation

The memory module in MemoRAG is crafted to memorize global information and generate informative retrieval cues. It leverages a transformer-based model to compress the database's raw input tokens into memory tokens. This process maintains semantic richness while allowing the efficient handling of extensive contexts, thus translating short-term memory interactions into long-term memory representations through specialized attention mechanisms.

The architecture supports various token compression ratios and is adaptable to different context lengths. Frameworks like key-value compression are integrated to manage super-long contexts efficiently.

Performance Optimization

To enhance retrieval accuracy, MemoRAG is optimized through pre-training and supervised fine-tuning of its memory model. This training involves initializing specific weight matrices that adaptively map memory tokens to perform contextual queries that guide retrieval processes.

The memory model construction produces intermediate answers or clues, thus linking high-level query semantics to specific evidence segments within the database. These clues are pivotal when assembling responses to queries characterized by implicit information needs or requiring distributed evidence aggregation.

Experiments and Results

The effectiveness of MemoRAG is corroborated through extensive evaluation on UltraDomain benchmarks, which encompass complex tasks across diverse sectors, such as finance and legal, as well as broader academic disciplines. MemoRAG consistently outperforms traditional RAG systems in scenarios that demand a high-level understanding of comprehensive contexts.

Additionally, MemoRAG showcases its capability in managing long texts, presenting high-quality answers based on finely-tuned retrieval paths from global memory. Its integration into generative models underscores its superiority over existing long-context processing techniques prevalent in standard LLM architectures.

Conclusion

MemoRAG represents a significant advancement in the domain of retrieval-augmented LLMs. By harnessing global memory structures alongside enhanced retrieval mechanisms, MemoRAG pushes the boundaries of what can be achieved in long-context and complex-query interpretation. Future work will extend MemoRAG's applications, potentially incorporating personalized assistants and life-long conversational agents to further exploit its robust memory and retrieval systems. The release of MemoRAG’s models and frameworks will undoubtedly stimulate ongoing research and practical deployment in long-context LLM tasks.