MemoRAG：重新定义长期记忆的AI问答模型

MemoRAG模型是如何实现长记忆的？

©作者|Blaze

来源|神州问学

引言

随着人工智能的发展，AI问答模型在各种应用场景中表现出色，尤其是在信息检索和知识问答领域。传统的RAG模型通过结合外部知识库的实时检索与生成模型，极大地提升了回答的准确性。然而，这类模型仍然面临一个重要挑战：无法有效处理长期信息，尤其是在需要持续记忆和动态更新知识的场景中表现不佳。

在此背景下，MemoRAG应运而生。MemoRAG不仅继承了RAG模型的检索增强能力，还通过引入“长期记忆”机制，解决了AI模型在复杂问题和持续任务中对历史信息的依赖问题。通过将信息存储为“记忆”，MemoRAG可以动态更新和检索长期信息，使其在个性化问答、知识管理以及任务跟踪等方面展现出显著优势。

本文将深入探讨MemoRAG的核心技术原理、应用场景以及它如何重新定义AI问答模型中的长期记忆能力。同时，我们将对比MemoRAG与其他RAG模型的不同之处，分析其在未来人工智能发展中的潜力。

MemoRAG的核心概念

MemoRAG（Memory-Augmented Retrieval-Augmented Generation）是一种在传统RAG模型基础上引入“长期记忆”机制的创新AI模型。RAG模型通过检索外部知识库中的相关信息来增强生成式模型的回答能力，能够处理开放领域的复杂问答。然而，RAG模型的核心弱点在于其只能依赖于即时的检索，无法对长期上下文或历史信息进行有效处理。这种缺陷使得它在处理需要长期记忆的场景（如复杂任务、个性化对话和知识积累）时存在局限。

在处理需要对整个数据库进行深度理解的查询时，可以通过《哈利·波特》书籍作为数据库对标准RAG和MemoRAG进行对比。左侧图展示了由于查询隐含的信息，标准RAG难以精确定位所需的证据，导致生成的答案不够准确。右侧图展示了MemoRAG通过构建覆盖整个数据库的全局记忆，在接收到查询后，首先回忆出相关线索，从而检索到更有用的信息，提供更加准确和全面的答案。

MemoRAG 的设计灵感源于人类的记忆机制。就像一个人读完《哈利·波特》系列小说后，虽然能大致记住故事主线、主要角色及其关系，但不可能准确记住所有细节。当被问到“邓布利多与格林德沃的关系如何？”时，可能只能给出一个模糊的概述，但通过快速查阅相关章节，可以找到具体的细节，从而归纳出一个完整的答案。MemoRAG的“记忆-回忆-检索-生成”过程，正是对这种人类记忆机制的仿生模拟！

MemoRAG的特点

动图封面

MemoRAG的创新在于为AI模型添加了“记忆增强”模块。这一模块允许模型像人类一样，不仅能从即时信息中学习，还能保存过去的关键信息，并在未来的交互中动态检索和应用这些记忆。这使得MemoRAG能够：

存储与更新长期记忆：MemoRAG引入了持久化存储机制，能够记住与用户交互中积累的重要信息，并在回答问题时结合这些历史数据。这种记忆并不仅限于一次性的任务完成，而是可以跨越多次会话保持一致性和连贯性。

强化生成的上下文意识：通过结合即时检索信息与长期记忆，MemoRAG能够生成更具上下文意识和个性化的回答。例如，在面对多轮对话或复杂任务时，MemoRAG能够利用之前的回答或用户历史偏好，更加准确地提供所需信息。

应对复杂、长期任务：MemoRAG在处理复杂问答、持续的项目管理或长期知识积累时表现尤为出色。通过记住用户或任务的特定需求，它可以持续优化其响应，提升用户体验。

记忆与检索的融合：MemoRAG结合了RAG模型的实时检索优势和长期记忆的积累，使得生成的回答更加准确、信息更为丰富，并能更好地满足动态变化的需求。

MemoRAG使用指南

要使用 MemoRAG，您需要安装 Python 以及相关库，以下是快速入门步骤：

1. 环境安装

安装依赖：

pip install torch==2.3.1
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

安装MemoRAG：

pip install memorag

2. MemoRAG基础用法

MemoRAG 使用起来非常简单，可以直接与 HuggingFace 的模型进行初始化。以下是使用 MemoRAG 的几个关键步骤和参数说明：

记忆构建：通过调用 MemoRAG.memorize() 方法，记忆模型能够在长输入上下文上构建全局记忆。

默认上下文处理能力：

●使用默认参数设置时，TommyChien/memorag-qwen2-7b-inst 能够处理多达 40 万个 tokens 的上下文。

●TommyChien/memorag-mistral-7b-inst 可以处理最多 12.8 万个 tokens 的上下文。

扩展上下文处理能力：

通过增加 beacon_ratio 参数，模型可以处理更长的上下文。例如，TommyChien/memorag-qwen2-7b-inst 在 beacon_ratio=16 时，最多可以处理 100 万个 tokens。

from memorag import MemoRAG

# Initialize MemoRAG pipeline
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-mistral-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3", 
    gen_model_name_or_path="mistralai/Mistral-7B-Instruct-v0.2", # Optional: if not specify, use memery model as the generator
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    access_token="hugging_face_access_token",  # Optional: Hugging Face access token
    beacon_ratio=4
)

context = open("examples/harry_potter.txt").read()
query = "How many times is the Chamber of Secrets opened in the book?"

# Memorize the context and save to cache
pipe.memorize(context, save_dir="cache/harry_potter/", print_stats=True)

# Generate response using the memorized context
res = pipe(context=context, query=query, task_type="memorag", max_new_tokens=256)
print(f"MemoRAG generated answer: \n{res}")

运行上述代码时，编码的键值（KV）缓存、Faiss 索引以及分块的段落会存储在指定的 save_dir 目录中。之后，如果再次使用相同的上下文，数据可以从磁盘中快速加载：

pipe.load("cache/harry_potter/", print_stats=True)

通常，加载缓存的权重非常高效。例如，使用 TommyChien/memorag-qwen2-7b-inst 作为记忆模型，对一个包含 20 万个 token 的上下文进行编码、分块和索引大约需要 35 秒，而从缓存文件加载则只需 1.5 秒。

3.摘要任务

要执行摘要任务，可以使用以下脚本：

res = pipe(context=context, task_type="summarize", max_new_tokens=512)

print(f"MemoRAG summary of the full book:\n {res}")

关于MemoRAG的应用还有很多，具体可以参考MemoRAG的项目主页。

MemoRAG评测

为了展示MemoRAG特点，在三个基准测试上对MemoRAG 进行了测试，以下是结合三种生成模型的记忆模型实验结果。

实验结果显示，MemoRAG 在所有数据集上的表现都超过了所有基线模型，展示了其强大的领域泛化能力。

MemoRAG应用场景

MemoRAG 的应用场景广泛，尤其适用于需要在大规模数据库或长上下文中进行复杂信息检索和回答生成的任务。以下是一些典型的 MemoRAG 应用场景：

1. 法律文件分析

法律文件通常篇幅较长，涉及复杂的法律条款和判例。MemoRAG 可以在处理大量法律文本时，通过记忆和回忆机制，快速从全局上下文中检索相关法律条文或案例，提供准确的法律建议或判例参考。

2.学术研究和文献综述

研究人员经常需要对海量文献进行分析并生成总结。MemoRAG 可以通过其记忆模型，处理大量学术文章或书籍，提供更全面的文献综述，并帮助研究人员快速找到相关的研究成果。

3. 大型产品知识库查询

企业的产品数据库通常包含大量信息，MemoRAG 可以处理来自不同来源的产品文档，快速定位产品特性、规格和使用方法，帮助客户支持或销售团队快速回应客户的复杂询问。

4. 历史事件分析

当涉及复杂的历史事件或时间线时，MemoRAG 可以记住整个历史资料库的内容，并在查询时回忆相关细节，从而生成准确的历史事件分析或回答。

5.医疗记录和诊断支持

在医疗领域，医生需要访问患者的详细病历和治疗历史。MemoRAG 可以处理长病历记录，并帮助医生快速找到患者病史中的关键信息，支持诊断和治疗决策。

6.小说和长篇文学分析

MemoRAG 能够处理长篇小说或文学作品，帮助读者或研究人员在分析角色关系、情节发展或主题时，从记忆中检索到相关情节或细节，为文学分析提供有力支持。

这些应用场景展示了 MemoRAG 在处理长上下文和复杂信息检索任务中的优势，特别是在需要跨多个来源或长时间跨度的数据中找到关键信息的场景中。

总结

MemoRAG 是一种基于高效、超长记忆模型的创新型 RAG框架，重新定义了 AI 在处理长上下文和复杂信息检索时的能力。通过结合记忆和生成机制，MemoRAG 能够在面对庞大的数据集时构建全局记忆，从而实现更精准的证据检索与生成。凭借其在多个基准测试中的优异表现，MemoRAG 展现出在法律、医疗、文献分析等多种场景下的巨大潜力。尽管项目仍处于初期阶段，MemoRAG 已通过项目代码实现了快速部署，并支持中文语料处理，未来将持续优化模型的轻量化、记忆机制的多样性以及应用广泛性，推动 AI 问答技术的进一步发展。

参考文献

https://github.com/qhjqhj00/MemoRAG?tab=readme-ov-file

https://arxiv.org/abs/2409.0559