MemoRAG：突破性记忆驱动的检索增强生成系统，重塑复杂任务处理

检索增强生成利用检索工具访问外部数据库，从而通过优化上下文来提升大型语言模型（LLMs）的生成质量。然而，现有的检索方法本质上受到限制，因为它们只能在明确提出的查询和结构化知识之间进行相关性匹配，无法处理涉及模糊信息需求或非结构化知识的任务，即：现有的 RAG 系统主要在简单的问答任务中有效。本文提出了一种新的检索增强生成范式MemoRAG，它由长期记忆驱动。MemoRAG 采用双系统架构。一方面，它使用一个轻量但长程的 LLM 来形成数据库的全局记忆。一旦任务被提出，它生成草稿答案，提示检索工具在数据库中定位有用的信息。另一方面，它利用一个昂贵但表达力强的 LLM，根据检索到的信息生成最终答案。

基于这一通用框架，本文进一步通过增强提示机制和记忆能力来优化 MemoRAG 的性能。在实验中，MemoRAG 在各种评估任务中表现出色，包括传统 RAG 失败的复杂任务和 RAG 常用的简单任务。MemoRAG 正在密集开发中，其原型和资源将不断在项目仓库中发布。

传统的 RAG 系统通常需要明确的信息需求和良构的知识。因此，它们的应用大多局限于简单的问答任务。然而，对于许多现实世界的问题来说，信息需求是模糊的，外部知识是非结构化的。例如，一本书的读者可能想了解主要角色之间的相互关系。为了解决这个问题，系统需要首先识别主要角色的名字，然后找到相应名字共存的部分，从中推断出它们的相互关系。换句话说，这需要在有效检索相关信息之前，基于书的上下文知识理解信息需求。

为了应对上述挑战，本文提出了一个名为MemoRAG的新框架。该框架引入了一个智能接口，将任务与数据库中的相关知识连接起来。对于每个呈现的任务，MemoRAG 会提示其记忆模块生成检索线索。这些线索本质上是基于数据库的压缩表示（即记忆）起草的答案。尽管可能存在错误细节，但这些线索明确揭示了所呈现任务的潜在信息需求。此外，它们还可以直接对应于现实中的源信息。通过使用这些线索作为查询，MemoRAG 可以有效地从数据库中检索必要的知识。

根据上述机制，记忆模块需要具备以下特性：1）记忆性：记住整个数据库的全局信息 2）指导性：提供有用的线索，基于这些线索可以全面检索所需的所有知识。

因此，本文提出了以下设计来优化 MemoRAG 的性能。首先，本文引入了双系统架构，使用一个轻量级 LLM 作为记忆，另一个重量级 LLM 执行检索增强生成。轻量级 LLM 必须具有成本效益和长上下文能力，能够在有限的计算预算内容纳整个数据库。其次，本文对记忆进行微调，以便生成的线索能够实现优化的检索质量。

为了评估 MemoRAG 的有效性，本文开发了一个名为 ULTRADOMAIN 的综合基准测试集，该测试集包含来自不同领域（如法律、金融、教育、医疗、编程）的复杂 RAG 任务，具有长输入上下文。ULTRADOMAIN 包括以下特征的任务：1）涉及隐含信息需求，2）需要分布式证据收集，3）需要对整个数据库的高层次理解。这些查询的一个共同挑战是，相关知识不能通过简单搜索直接检索到。虽然传统的 RAG 方法在处理此类复杂任务时表现不佳，但 MemoRAG 通过基于精确检索的知识生成高质量答案，表现出色。此外，在传统的问答任务中，MemoRAG 仍然具有很强的竞争力，其基于线索的检索方法继续提供显著优势。

标准的 RAG 框架可以简洁地表示为：，其中， 和 分别表示生成模型和检索模型。 表示输入查询， 是从相关数据库 中检索到的上下文， 是最终答案。在许多实际场景中，输入查询 通常携带隐含的信息需求意图，这对于通常依赖词汇或语义匹配的标准检索器来说，可能难以完全理解。这一限制突显了设计一个中间模块以弥合许多实际场景中出现的语义差距的必要性。

本文提出了 MemoRAG，它利用记忆模型 作为输入查询 和相关数据库 之间的语义桥梁。形式上，这一过程可以表示为：，这里， 表示一个可能不完整或缺乏细节的阶段性答案，作为一组答案线索，指导从 中检索最相关的上下文。记忆模型 旨在建立数据库 的全局记忆。在实践中，任何能够高效处理超长上下文的语言模型都可以作为记忆模型。例如，一个结合了键值压缩技术的 7B 语言模型可能是一个合适的选择。

虽然这样的模型可能无法生成一个完全详细和准确的答案，但它可以生成一个粗略的轮廓，有助于定位正确的答案。 阶段性答案 的形式是根据每个任务的具体要求量身定制的。例如，在输入查询隐含的问答任务中，阶段性答案 可能包括中间步骤，如生成更明确和消歧的替代查询，以及来自数据库的具体文本证据，这些证据有助于最终答案。此外，对于不涉及明确查询的任务，如摘要生成，阶段性答案可能包括从上下文中提取的关键点或概念，这对于组装一个连贯和准确的摘要至关重要。

MemoRAG 的系统实现可在此代码库中 https://github.com/qhjqhj00/MemoRAG 找到。目前，本文发布了两个记忆模型：memorag-qwen2-7b-inst 和 memorag-mistral-7b-inst，分别基于 Qwen2-7B-Instruct 和 Mistral-7B-Instruct-v0.2。

在 MemoRAG 中，记忆模块 可以是任何设计用于高效处理超长上下文的模型，而目前系统采用了一种 token 压缩技术，实现了对扩展上下文的高效处理，具体细节将在后文详细解释。这两个记忆模型支持 2 到 16 的压缩比，允许它们管理不同长度的上下文。例如，当应用 16 的压缩比时，memorag-qwen2-7b-inst 可以处理最多 128K×16 个 token。在实践中，memorag-mistral-7b-inst 在上下文长度达到 150K 个 token 时表现良好，而 memorag-qwen2-7b-inst 在上下文长度达到 600K 个 token 时表现良好。当处理更长的上下文时，它们的性能会略有下降。

MemoRAG 可以将各种检索方法集成到系统中，包括稀疏检索、密集检索和重排序。当前实现默认使用密集检索。在未来的版本中，本文计划提供灵活的检索方法和易于使用的接口。

同样，MemoRAG 可以将任何生成语言模型集成为生成器。当前实现支持从 HuggingFace 生成模型或通过商业 API（如 Azure 和 OpenAI）初始化生成器。由于记忆模型依赖于其底层模型，MemoRAG 使用记忆模块的底层模型作为默认生成器。例如，当使用 memorag-mistral-7b-inst 作为记忆模型时，MemoRAG 默认使用 Mistral-7B-Instruct-v0.2作为生成模型。

MemoRAG 系统可以部署在各种 GPU 上，从更经济实惠的选项（如 NVIDIA T4 16GiB）到高端选项（如 NVIDIA A100 80GiB）。使用 NVIDIA T4 16GiB 时，MemoRAG 可以处理上下文长度为 68K 个 token 的数据库。使用 NVIDIA A100 80GiB 时，系统能够处理上下文长度达到一百万个 token 的数据库。

▶Memory Module

本文提出了一种灵活的模型架构，专门设计用于促进记忆的形成。该记忆模型逐步将原始输入 token 压缩成显著较少的记忆 token，同时保留基本的语义信息。具体来说，假设输入 包含 个 token，，并由基于 transformer 的模型 处理。每一层的注意力交互可以表示为：

其中 、 和 分别是 query、key 和 value 投影的权重矩阵， 是 key 向量的维度。经过多层 Transformer 的深度注意力交互后， 的输入序列得到了全面的理解。这类似于人类的短期记忆，准确但只能包含最近感知的内容。我们将这个过程表示为 ，其中 指的是输入序列 的隐藏状态， 可以是任何预训练的 LLM。

为了实现从短期记忆到长期记忆的转换，本文引入了记忆 token ，作为 LLM 中长期记忆的信息载体。具体来说，假设底层 LLM 的工作上下文窗口长度为 ，在每个上下文窗口之后，我们附加 个记忆 token，即：

在注意力交互过程中，我们为记忆形成的特殊目的初始化了另一组权重矩阵 、 和 。因此，我们有：

其中 、 和 分别指的是记忆 token 的 query、key 和 value， 和 指的是先前记忆 token 的 KV 缓存。我们将记忆 token 表示为 ，将转换过程表示为 。对于 个原始 token ，经过多次注意力过程后，它们被编码为隐藏状态 ，其中 表示原始 token 的隐藏状态， 表示记忆 token 的隐藏状态。在记忆形成之后， 个原始 token 的 KV 缓存被丢弃，类似于人类记忆中的遗忘过程。经过 个上下文窗口后，MemoRAG 逐步将 中的所有原始 token 转换为记忆 token。因此，我们有 ，这表示由输入 形成的全局记忆 。

▶Training For Memory Module

如上所述，本文初始化了另一组权重矩阵 、 和 ，用于将记忆 token 映射到 query、key 和 value 向量（如公式 7 所述）。新初始化的权重矩阵在训练过程中会被更新，而底层 LLM 的参数保持不变。本文在两个阶段训练新初始化的参数：（1）预训练：使用 RedPajama 数据集中随机抽取的长上下文对模型进行预训练，使 MemoRAG 的记忆模块能够从原始上下文中学习如何形成记忆；（2）监督微调（SFT）：使用特定任务的 SFT 数据，使 MemoRAG 能够基于形成的记忆生成特定任务的线索。

尽管具有挑战性，但准确记住记忆中的细节仍然是任何人类记忆增强训练的最终目标。MemoRAG 中记忆模型的训练目标也追求这一目标，可以表述为：

▶The MemoRAG Framework

在前一节中，输入序列 被转换为紧凑的记忆表示 ，该表示从全局角度封装了高级语义信息。利用这个记忆 的一个直接方法是提示它生成特定任务的答案，即 ，其中 代表任务描述（例如查询或摘要指令）。虽然这种方法是可行的，但由于 是原始输入序列 的高度浓缩形式，可能会缺乏信息的准确性。这类似于人类可能难以从记忆中回忆详细信息，但可以生成一个草稿答案，然后通过重新访问和查找相关证据来进行完善。

在 MemoRAG 中，全局记忆 被用来生成特定任务的线索 。这些线索有助于勾勒出预期的答案 ，有效地弥合了原始输入上下文和真实答案之间的差距。基于这些记忆生成的线索，可以使用任何独立的检索器来定位输入序列中的精确证据文本。随后，基于检索到的证据文本生成最终答案 ，即 ，其中 包括输入查询或任务指令 和检索到的上下文 。

生成模型可以是任何生成式语言模型，默认情况下，MemoRAG 使用记忆模型的底层语言模型进行生成，避免了加载额外模型参数的需要。

▶Applicability

模糊信息需求 对于标准 RAG 来说，处理模糊信息需求是具有挑战性的，因为用户的意图通常没有明确说明，需要更深层次的上下文理解和推理能力。为了解决这个问题，MemoRAG 在相关数据库中创建了一个全局记忆，使其能够推断隐式查询的潜在意图。通过生成阶段性答案，例如更具体的答案线索，MemoRAG 弥合了隐式信息需求与检索过程之间的差距。这显著拓宽了标准 RAG 系统在涉及隐式查询的任务中的适用范围。

上表展示了 MemoRAG 如何处理隐式查询。例如，输入查询“这本书如何传达爱的主题？”与相关数据库中的内容缺乏直接的语义连接，因为“爱的主题”在原始文本中没有明确说明。MemoRAG 在整个数据库中形成了一个全局记忆，并生成了关键答案线索，帮助从数据库中检索相关内容。

分布式证据查询的信息检索 处理需要分布式证据收集的查询为标准 RAG 系统引入了额外的复杂性，因为这些查询需要整合跨多个步骤或非结构化数据库的知识。大多数检索方法在处理这些查询时都很困难，因为它们需要对相互关联的数据点进行连贯的理解，通常跨越不同的上下文。

为了解决这一挑战，MemoRAG 利用其全局记忆能力连接和整合数据库中多个步骤的相关信息。通过生成指导检索相互关联数据点的阶段性答案，MemoRAG 有效地管理了多跳查询的复杂性。这种方法使 MemoRAG 能够显著提升标准 RAG 系统在涉及多跳推理任务中的性能。

上表展示了 MemoRAG 如何处理这样的查询。例如，输入查询“过去三年中哪一年收入最高？”需要分析跨多个年份的财务数据。MemoRAG 从过去十年的大公司财务报告中形成全局记忆，将多跳查询重新表述为几个具体查询，并整合这些信息以确定收入最高的年份。

信息聚合 信息聚合任务，例如总结长文档，需要能够将大量非结构化数据浓缩成简明且连贯的输出。标准 RAG 系统通常在这些任务中表现不佳，因为它们依赖于检索离散的信息片段，而没有有效的机制将其结合并总结成全面的概述。MemoRAG 通过利用其全局记忆来捕捉和综合整个数据集的关键点来解决这一挑战。通过这一过程，MemoRAG 能够生成代表内容基本要素的中间阶段性答案，然后用这些答案从原始内容中检索详细信息。所有这些信息被聚合以生成最终的总结。

上表展示了 MemoRAG 如何处理信息聚合任务。例如，任务是总结一份关于城市建设的政府报告。MemoRAG 首先从报告中提取关键点，如基础设施发展、预算分配和未来规划，然后检索详细内容，并聚合这些信息以生成报告的全面总结。

当前版本的 MemoRAG 主要针对上述应用场景。以下是本文计划的 MemoRAG 应用场景，这些场景可以通过进一步在特定任务训练数据上进行训练来实现。

个性化助手 个性化助手任务，例如根据用户的偏好推荐歌曲，需要深入了解用户的独特特征和历史。这是因为个性化信息需求通常是模糊的，受到用户个人特质的强烈影响。标准 RAG 系统可能在这些任务中表现不佳，因为它们通常依赖于一般的相关性匹配，而不是基于特定用户数据个性化结果。

MemoRAG 通过利用全局记忆来分析和理解用户的对话历史，从而增强个性化。这使得 MemoRAG 能够识别和利用关键线索，如用户的音乐偏好、知识背景、年龄以及从过去互动中推断出的其他相关因素。通过综合这些信息，MemoRAG 能够生成高度个性化的推荐，与用户的品味紧密契合。

上表展示了 MemoRAG 如何处理个性化推荐查询。例如，当被问到“你能推荐一首歌给我吗？”时，MemoRAG 分析用户的对话历史，识别对某些音乐类型、艺术家或年代的偏好，并利用这些信息建议符合用户档案的歌曲。

终身对话搜索 对话搜索经常涉及语义省略的查询，其中用户的意图依赖于先前互动的上下文。查询重写是一种广泛采用的技术来解决这一挑战。在终身对话搜索中，当前查询的语义依赖可能延伸到更早的互动，因此准确识别极长互动历史中的相关上下文至关重要。标准 RAG 系统可能在这项任务中表现不佳，因为查询中呈现的信息需求存在模糊性。这主要是因为标准 RAG 系统缺乏有效跟踪和整合不断演变的对话上下文的能力，往往导致检索不完整或不准确。

MemoRAG 通过利用其全局记忆来维护和利用对话历史的完整上下文，从而解决这一挑战。这使得系统能够通过参考相关的先前交流并填补意义上的空白来解释语义省略的查询。因此，MemoRAG 能够准确解释和响应依赖于先前对话上下文的后续查询。例如，考虑在讨论某篇研究论文后提出的查询“它有任何弱点吗？”标准检索系统可能难以理解“它”指的是什么，可能会检索到不相关的信息。相比之下，MemoRAG 会回顾之前的对话，识别出“它”指的是正在讨论的研究论文，然后检索关于该论文弱点的信息。

上表展示了 MemoRAG 如何在对话搜索上下文中处理语义省略的查询。通过利用完整的对话历史，MemoRAG 确保了准确且上下文适当的响应。

为了评估 MemoRAG 在不同复杂任务中的泛化能力，本文在 UltraDomain 基准上对 MemoRAG 进行了评估。UltraDomain 中的大多数查询涉及模糊的信息需求或非结构化的知识检索挑战。UltraDomain 由两种类型的数据集组成。第一种类型包括三个上下文大小在 100K tokens 以下的数据集，这些数据集与本文的训练数据集分布相同。本文将这三个数据集称为 in-domain 数据集。第二种类型包括从 428 本英文大学教科书中提取的 18 个数据集，上下文长度可达一百万 tokens，涵盖哲学（如《尼采与希腊思想》）和数学（如《数学家的量子理论》）等学科。由于第二种类型的数据集的数据分布与训练数据不同，本文将其称为 out-of-domain 数据集。实验结果总结在表6中，本文得出以下结论：

1. MemoRAG 在所有数据集上均优于所有基线模型，展示了其强大的领域泛化能力。
2. 直接将完整上下文输入 LLM 通常比其他 RAG 方法（BGE-M3、Stella-v5 和 HyDE）表现更好。这个发现揭示了标准 RAG 系统在处理长上下文和高级问题时的困难。
3. 相比之下，MemoRAG 始终优于直接使用完整上下文的表现，说明其在处理超长上下文和解决复杂任务方面的有效性。
4. MemoRAG 在三个 in-domain 数据集上表现出色，表明其潜力可以通过更多样化的训练数据进一步增强。

上表展示了在三个基准上的实验结果，从中我们可以得出结论，MemoRAG 在所有数据集上通常都优于所有基线，除了一个异常值。首先，对于开放域 QA 任务，MemoRAG 在所有数据集上都优于所有基线，除了使用 Llama3 作为生成器的 en.qa。这验证了在标准 RAG 的舒适区内，大多数查询具有明确的信息需求，MemoRAG 能够更好地在原始上下文中定位预期的证据，这要归功于记忆生成的线索。其次，大多数以前的 RAG 方法在不涉及查询的任务（如摘要任务，例如 Multi-News、GovReport 和 en.sum）上表现不佳。本文的 MemoRAG 使 RAG 系统能够从输入上下文中生成关键点，并检索更多细节以形成全面的摘要。第三，对于特定领域的任务（例如金融和法律），MemoRAG 显示出显著的改进，表明 MemoRAG 在处理具有长上下文的复杂任务方面的优越性。

总之，结果表明 MemoRAG 在各种数据集和查询类型上显著提升了标准 RAG 方法和其他基线的性能。MemoRAG 在有效处理复杂和长上下文任务方面的能力突显了其优势，特别是在标准 RAG 系统表现不佳的场景中。这种在不同生成器上的一致性强调了 MemoRAG 的鲁棒性和广泛适用性。

本文介绍了 MemoRAG，这是一种新颖的检索增强生成（RAG）系统，集成了全局上下文感知，以应对涉及长输入上下文的复杂任务所带来的挑战。MemoRAG 具有一个内存模块，该模块在整个数据库中构建了一个紧凑的全局内存，促进了上下文相关线索的生成，有效地将知识数据库与准确答案所需的精确信息联系起来。大量关于知识密集型问答、摘要以及涉及长文档的实际应用的实验表明，MemoRAG 显著优于传统的 RAG 系统。它在需要高级信息聚合的任务中表现出色，并在处理大规模文本（如教科书、财务报告和法律合同）时表现出卓越的鲁棒性和多功能性，能够处理多达一百万个 token 的上下文，并以优异的准确性解决复杂查询。