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

TL;DR

MemoRAG tackles the challenge of long-context processing by introducing a global memory augmented retrieval framework. It combines a light memory module that generates draft clues with a traditional retrieval-augmented generation pipeline and a separate expressive generator, enhanced by a KV-compressible memory and reinforcement learning with generation feedback (RLGF). Empirical results across long-context benchmarks and the UltraDomain suite show MemoRAG outperforms standard and advanced RAG methods and long-context LLMs, including in non-QA and complex QA tasks, while maintaining favorable efficiency. This approach broadens the applicability of RAG to very long inputs and domain-general tasks, offering scalable memory for evidence retrieval and high-quality generation.

Abstract

Processing long contexts presents a significant challenge for large language models (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.

Figures (5)

• Figure 1: Comparison of MemoRAG with Standard RAG and human cognition of a long document. Figure (a) shows standard RAG, where retrieval and generation take place in a sequential pipeline. Figure (b) illustrates how humans tackle a task about the document: 1. going through the document and forming the memory, 2. thinking about the clues to the presented task (i.e., recalling), checking the document for needed details (i.e., retrieving), 3. making a response to the task based on the memory-enhanced retrieval result. Inspired by the human cognition process, Figure (c) demonstrates MemoRAG, which creates a global memory of the long context, recalling useful clues based on memory, and retrieving information based on the clues to generate a high-quality response.
• Figure 2: Illustration of (a) task background, (b) framework comparison, and (c) application scenarios. When processing long inputs like the entire Harry Potter series, most LLMs struggle with million-token contexts. Standard RAG methods also face challenges with queries unsuitable for direct searching. MemoRAG overcomes these limitations by constructing a global memory that generates clues, guiding the retrieval of relevant evidence and enabling more accurate and comprehensive answers.
• Figure 3: Experiment results on the UltraDomain benchmark. These datasets feature contexts of up to one million tokens, covering a wide range of subjects. See more details about the benchmark in Appendix \ref{['sec:ultra']}.
• Figure 4: Ablation study. Figure (a) and (b) show the performance of different LLMs and optimization strategies. The Pretrain, SFT, and RLGF settings refer to the training stages. The Light setting uses the light memory model, introduced in Section \ref{['sec:mem']}. The Zero setting uses native LLMs without prior training. Figure (c) shows the outcomes of using different models as the generator.
• Figure 5: Analysis on the model efficiency (left) and the impact of the choice of the compression ratio $\beta$ (right).