Source-linked AI summary
MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval Augmentation
Hongjin Qian, Zheng Liu, Peitian Zhang, Kelong Mao, Defu Lian, Zhicheng Dou, Tiejun Huang
TL;DR
Long-context processing remains costly and difficult, while conventional RAG depends on explicit queries and structured knowledge that general tasks often lack. MemoRAG builds global memory, generates draft-answer clues for retrieval, and uses retrieved evidence for final generation; experiments report superior performance across diverse long-context tasks.
Problem
Long-context LLM processing can remain insufficient and computationally expensive, while conventional RAG requires explicit search intent and well-structured knowledge that general long-context tasks often lack.
Method
MemoRAG builds global memory with a light long-range system, generates draft answers as retrieval clues, and uses an expressive system to produce final answers from retrieved evidence.
Results
MemoRAG significantly outperforms standard RAG systems, advanced RAG systems, and long LLMs across two long-context benchmarks and varied real-world applications.
Takeaways & Limitations
MemoRAG extends global-memory-enhanced retrieval augmentation beyond traditional QA to non-QA and complex QA long-context tasks.
Takeaways & Limitations
The light global-memory design is inferior to compact global memory because native context limits, sparse attention, and full key-value-cache storage constrain it.
Abstract
from arXiv · showhide
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.
1 Introduction
MemoRAG addresses costly and insufficient long-context processing by combining global memory with clue-guided retrieval and expressive final generation. It is designed for unstructured tasks lacking explicit search intent and performs across both QA and non-QA settings.
- Motivation: Long-context processing remains computationally expensive despite longer-capacity LLMs, motivating retrieval methods that focus computation on relevant context.Direct processing incurs substantial inference-time and GPU-memory costs.
- Motivation: Traditional RAG struggles on general long-context tasks because they often lack explicit search intent and well-structured inputs for indexing.Examples include summarizing a book’s characters and processing long text files or multi-year financial reports.
- Framework: MemoRAG uses a light, long-range memory system to form global context memory and an expressive system to generate the final answer from retrieved information.The memory module generates draft answers as retrieval clues, which are then used to locate relevant source information.
- Memory design: MemoRAG’s memory module is made length-scalable, retentive, and instructive through configurable KV compression and reinforcement from generation-quality feedback.RLGF rewards clues that support high-quality answers and reinforces the memory module to generate such clues.
- Experiments: MemoRAG is evaluated on QA and non-QA tasks from LongBench, InfiniteBench, and a 20-domain long-document understanding benchmark.The evaluation includes tasks favorable and unfavorable to traditional RAG methods.
- Contributions: The authors report that MemoRAG extends RAG beyond traditional QA to non-QA and complex QA tasks.This broadens the reported applicability of retrieval augmentation across long-context scenarios.
2 Method
MemoRAG addresses long-context processing by combining global-memory clues with retrieval and expressive answer generation. Its memory design compresses long inputs, supports scalable context handling, and is trained to produce clues that improve retrieval and final answers.
- 2.2 MemoRAG: MemoRAG uses a global memory to generate task-specific draft answers that guide retrieval before an expressive model produces the final response.The memory model stores the long context, generates answer clues for each task, and the final system uses retrieved evidence to answer.
- 2.3 Memory Model: The memory module is designed to be length-scalable, retentive, and instructive for handling long contexts and generating useful retrieval clues.Length scalability and retention come from model design, while instructiveness is developed through multi-stage, data-driven training.
- 2.3.1 Memory Model Design: KV compression progressively converts raw tokens into fewer memory tokens while preserving essential semantics, reducing memory use by approximately β× at compression ratio β.The reported compression ratios are β ∈ [4, 8, 16, 32, 64].
- 2.3.1 Memory Model Design: With β=64, a 128K-context LLM can process up to an 8M-token context because the compressed memory contains far fewer tokens than the raw input.The compact representation extends the effective context length beyond the model’s native window.
- 2.3.2 Memory Model Training: RLGF trains the memory model to prefer clues that improve evidence retrieval and final answer generation rather than merely matching supervised clue labels.Its preference-based ranking loss uses rewards for preferred and non-preferred outputs.
3 Experiment
Experiments compare MemoRAG with standard and advanced RAG systems and long-context LLMs across conventional, complex, non-QA, and super-long-context tasks. MemoRAG consistently outperforms the baselines, while ablations support its design choices and efficiency analyses characterize its trade-offs.
- 3.3 Main Experiments: MemoRAG handles straightforward QA, complex QA, and non-QA summarization tasks, including scenarios where traditional RAG struggles with implicit information needs.Its memory-generated clues support evidence retrieval for complex questions and summarization tasks without explicit queries.
- 3.3 Main Experiments: On UltraDomain’s 18 diverse datasets, MemoRAG surpasses both standard baselines and direct full-context processing when most inputs exceed the generator’s 128K-token limit.Figure 3 reports this evaluation across contexts reaching up to one million tokens.
- 3.3 Main Experiments: MemoRAG consistently outperforms standard and advanced RAG systems and long-context LLMs across the evaluated datasets.The main experiments use LongBench, InfiniteBench, and broader UltraDomain evaluations.
- 3.4 Ablation Study: Ablations show that memory-model design, training stages, foundation-model choice, and generator choice each contribute to MemoRAG’s effectiveness.Removing technical designs causes degradation, while improvements remain consistent with Qwen2-7B-Instruct and across three alternative generators.
- 3.4 Ablation Study: As compression ratio β increases, performance declines but stabilizes at β = 32, while MemoRAG remains above standard RAG across tested values.Higher β improves efficiency but reduces semantic richness; the tested range is β∈[4, 8, 16, 32, 64].
- 3.4 Ablation Study: MemoRAG’s indexing and retrieval are slower than standard RAG but remain more efficient than long-LLM pre-filling because of its optimized memory model.GraphRAG is reported as the slowest in the efficiency comparison.
4 Related Work
Prior work addresses long-context processing through longer context windows, efficient attention, input compression, and retrieval. However, standard RAG remains limited for complex tasks with implicit information needs and unstructured inputs.
- Long Context: Long-context LLM processing is computationally expensive, motivating efficient attention and input-compression approaches.Longer contexts also present performance challenges when relevant information appears in the middle of the context.
- RAG: Standard RAG retrieves precise evidence from large inputs but struggles when search intent is implicit or the source context is not well structured.Effective conventional RAG generally depends on explicit information needs and structured external datasets suitable for indexing.
- RAG: The literature motivates more general-purpose RAG systems that can support complex long-context tasks beyond conventional explicit-query settings.This need follows from the limited generalization of standard RAG to tasks involving implicit information needs.
- RAG: Recent methods modify retrieval through hypothetical documents, query rewriting, or knowledge graphs, but these approaches retain important limitations.HyDE and RQ-RAG rely on internal model knowledge, while GraphRAG incurs substantial construction costs.
5 Conclusion
MemoRAG builds global memory from an entire long context and uses draft answers as retrieval clues for final generation. Experiments report advantages over standard and advanced RAG systems and long-context LLMs across aggregation-heavy and traditional tasks.
- 5 Conclusion: MemoRAG significantly outperforms standard RAG systems, advanced RAG systems, and long LLMs across extensive long-context evaluations.The reported evaluations span two long-context benchmarks and various real-world applications.
- 5 Conclusion: MemoRAG excels at tasks requiring high-level information aggregation while also improving traditional tasks commonly handled by previous RAG systems.Its global-memory clues guide retrieval toward precise evidence before final response generation.
A Implementation Details
The implementation trains and optimizes a compressed memory model while keeping the underlying model frozen, then evaluates it under standardized retrieval and generation settings. The setup uses scalable context handling and dedicated efficiency accelerators.
- A Implementation Details: The memory model is pretrained on 2 billion RedPajama tokens with randomly selected compression ratios from 4 to 64.Pretraining uses a 2048-token memory context window, one epoch, batch size 8, and learning rate 5e-5.
- A Implementation Details: Supervised fine-tuning uses 17,116 samples spanning 4K to 64K tokens for two epochs.The SFT batch size is 8 and the learning rate is 1e-5.
- A Implementation Details: RLGF optimization ranks 2,000 generated clue answers by their contribution to end-to-end performance and labels them preferred or rejected.These rankings provide the preference data used during the optimization stage.
- A Implementation Details: Training freezes the underlying model and updates only newly initialized memory-model parameters, reducing the need for full parameter fine-tuning.With Qwen2-7B-instruct, the newly initialized parameters total approximately 1.1 billion.
- A Implementation Details: The light global-memory system extends context windows with SelfExtend and accelerates prefill with MInference.These components support long-range memory processing in the light system.
- A Implementation Details: Main experiments use compression ratio β=4, BGE-M3 retrieval with three hits, 512-token chunks, official task prompts, and task-specific shared generation settings.Training and evaluation run on eight NVIDIA A800-80G GPUs.
A.1 Case Study
A 56.6K-token legal-contract case study contrasts query-only retrieval with MemoRAG's clue-guided evidence search. MemoRAG locates relevant passages and produces a more comprehensive and precise response about the contract-specific meaning of “Outside Date.”
- A.1 Case Study: In a 56.6K-token legal contract, standard RAG retrieves poorly because the query does not explicitly express the significance sought for “Outside Date.”Its answer therefore gives a general definition rather than explaining the term’s significance within the contract.
- A.1 Case Study: MemoRAG accurately locates relevant evidence passages, leading to a more comprehensive and precise response.The case illustrates the benefit of clue-guided retrieval for a high-level legal-document understanding question.
B More details of Dataset Construction
MemoRAG’s training data combines diverse long-context sources with generated high-level questions, while RLGF ranks clue combinations by their end-to-end retrieval performance.
- SFT dataset: Long contexts from novels, academic papers, news, financial reports, and legal contracts are sampled up to 80K tokens for SFT data construction.Strong LLMs generate high-level questions and answers from these contexts.
- RLGF dataset: For RLGF data, each clue retrieves its top-3 evidence, and greedy evaluation selects the best and worst combinations of three or more clues.These combinations become the preferred and rejected training answers.
C More details of UltraDomain
UltraDomain evaluates MemoRAG on specialized financial and legal contexts alongside broad cross-disciplinary materials, using manually reviewed high-level QA samples.
- Dataset composition: The benchmark includes financial reports, legal contracts, 428 college textbooks across 18 domains, and a miscellaneous dataset for broader coverage.These sources test domain-specific comprehension and adaptability beyond finance and law.
- QA construction: Evaluation samples use text spans up to 128K tokens, with GPT-4 generating high-level question-answer pairs that require comprehensive context understanding.Graduate students selected questions that were not directly searchable and evaluated answer quality.
- Benchmark statistics: The resulting benchmark contains 3,240 evaluation samples and is documented with dataset statistics in Tables 3 and 4.The tables provide statistical information for the utilized and out-of-domain evaluation datasets.