源码角度看MemoRAG框架是怎么实现的？实现逻辑及优缺点分析

最新推荐文章于 2025-03-02 10:34:41 发布

AI大模型-搬运工

最新推荐文章于 2025-03-02 10:34:41 发布

阅读量967

点赞数 31

文章标签：人工智能 MemoRAG LLM RAG 大语言模型自然语言处理 AI大模型

本文链接：https://blog.csdn.net/2401_84208172/article/details/142375364

版权

这篇技术报告介绍了一个名为MemoRAG的新型检索增强生成系统。

其思想在于，通过长期记忆启发的知识发现来解决传统RAG系统在处理长文本和复杂任务时的局限性，底层逻辑是：从全局记忆中生成精确的线索，将原始输入与答案连接起来，并从复杂数据中解锁隐藏的洞察，这个思路其实不太新。

一句话说，就是，在接收用户查询后，MemoRAG先基于记忆生成模糊答案和答案线索，然后通过检索器补充答案细节，最终生成完整答案。

今天来看具体看看MemoRAG的一些具体的逻辑，看看这个是否真的是那么回事儿。

先说一个具体的结论：

这个是专门用来做单个长文本问答的。

给定一个长文本，先对他做编码，做成KQV，然后做存储。也就是里面的memory.bin。来一个query后，为了可以跟整个文档做交互，所以模型进行query编码时，先加载之前长文本的kqv权重，然后再进行编码。也就相当于拼接了之前的记忆。接着加上记忆之后，根据下游prompt，做rewrite，线索生成任务。

此外，为了后面做召回检索，还是用的传统RAG做的切分，embedding，并建立索引，放在faiss里面，也就是index.bin，后续流程就是上面所说的问答后续流程了。

因此，我们可以做个总结：

在优点方面：

这个场景，针对的是长文本问答，比如kimi输入一个文档，做各种问题。

有了记忆模型之后，相当于看过了整个文档。所以效果会更好。所以，核心点在于，每次上传文档后，需要先把文档读一遍，然后建memory索引。

当多个用户看一篇文档时，可以做cache加速，不用重复建memory索引。

其论文中也说了，主要擅长做摘要和qa任务。

在缺点方面：

来回建，每个文档都得事先存一个，10000个，就得存10000个memory.bin，这个空间占据很大，例如，上图所示，21w字，memory索引要11gb。

此外，文档太长，也搞不定，超出了最大长文长度，效果也很难讲，并且会很慢。

此外，其对长文本模型能力的要求是很高的。

所以，综合来看，这个工作更多的是长文本的演示demo，是个trick的工作，并做不了大规模的落地。

不过，针对长文本做历史kqv-存储是发展趋势，对长文本进行压缩，query与全局文档的语义交互这些，都很有趣。‍

供大家一起参考并思考，持续前进，保持自律，踏实成长。

一、Memo-RAG索引阶段解读

这段代码块展示了如何使用MemoRAG（Memory Reasoning about Gist）流水线中的memorize函数，通过记忆和推理来改善文本检索和理解的能力。

加载文本文件：代码首先加载一个文本文件的内容，这里的例子中使用的文件是harry_potter.txt。

处理文本：通过memorize函数处理加载的文本。这个过程可能包括文本的解析、编码和存储，以便模型可以更好地理解和检索信息。

存储结果：处理后的文本被存储在指定的目录中。在这个目录中，会创建三个关键文件：

memory.bin：这个文件存储了记忆模型的键值（KV）缓存，它允许快速检索之前处理过的信息。这可以看作是一种快速访问存储在模型中的知识的机制。
index.bin：这个文件包含了文本语料库的密集嵌入（dense embeddings），这有助于高效地检索相关段落。嵌入是一种将文本转换为数学表示（通常是向量）的方法，以便机器学习模型可以更好地处理和理解文本。
chunks.json：这个文件包含了从输入上下文中提取的段落或块（chunks），这些在检索过程中使用。这些块是文本的一部分，它们被用来在需要时快速检索信息。

保存预处理数据：如果设置了save_dir参数，该方法会将预处理的数据（即记忆、嵌入和块）保存到磁盘。这样做可以在未来的操作上大大提高效率，因为从磁盘加载缓存的数据比从头开始重新处理和编码上下文要快得多。
缓存机制：这种缓存机制在处理大型文本或频繁访问的数据集时特别有用，因为它减少了重复处理相同数据的需要，从而节省了时间和计算资源。

其中：

1、memory.bin的获取，直接使用的一个长文本模型，获取模型最后的kv进行存储。

2、index.bin的获取，采用的是传统的embedding方案，使用faiss进行存储，并对文档进行切片，并对文档进行doc-embedding。

这个可以从https://github.com/qhjqhj00/MemoRAG/blob/main/memorag/retrieval.py中找到对应的实现细节。

针对所有自学遇到困难的同学们，我帮大家系统梳理大模型学习脉络，将这份 LLM大模型资料 分享出来：包括LLM大模型书籍、640套大模型行业报告、LLM大模型学习视频、LLM大模型学习路线、开源大模型学习教程等, 😝有需要的小伙伴，可以 扫描下方二维码领取🆓↓↓↓

👉[CSDN大礼包🎁：全网最全《LLM大模型入门+进阶学习资源包》免费分享（安全链接，放心点击）]()👈

二、Memo-RAG推理阶段解读

这里，涉及到对第一步构建好的memory.bin进行加载。

可以从https://github.com/qhjqhj00/MemoRAG/blob/main/memorag/memorag.py中找到具体的实现逻辑：

1、总体实现逻辑

def _handle_rag(self, query: str, prompt_template: str, max_new_tokens: int, use_memory_answer: bool):
        ## 1、召回线索文本
        text_spans = self.mem_model.recall(query)
        ## 2、生成线索问题
        surrogate_queries = self.mem_model.rewrite(query)
        ## 3、生成召回问题（线索文本+线索问题+草稿答案并做长度过滤，<3去掉）
        retrieval_query, potential_answer = self._prepare_retrieval_query(query, text_spans, surrogate_queries, use_memory_answer)
       ## 4、召回相关文档
        retrieval_results = self._retrieve(retrieval_query)
        if potential_answer:
            retrieval_results.append(f"The answer might be {potential_answer}.")
        ## 5、文档+草稿答案拼接prompt
        knowledge = "\n\n".join(retrieval_results)
        ## 6、生成答案
        return self._generate_response("qa_gen", query, knowledge, prompt_template, max_new_tokens)

可以先看一些特定的prompt，这个是核心点之一

对应prompt:https://github.com/qhjqhj00/MemoRAG/blob/main/memorag/prompt.py

其中，维护了一个prompt的字典，来存储不同的任务，比如线索生成、片段生成、答案生成等。

2、具体实现步骤

step1：召回线索文本

text_spans = self.mem_model.recall(query)

执行原理：

def recall(
        self,
        query, max_new_tokens=128) -> str:
        return self.generate(self.prompts["span"], query, max_new_tokens=max_new_tokens)[0]

调用prompt:

"span": """
你会得到一个与文章相关的问题。为了有效地回答这个问题，你需要回想文章中的具体细节。你的任务是识别并提取文章中与问题相关的一个或多个具体线索文本。

### 问题：{question}
### 指示：
1. 你对文章有一个大致的理解。你的任务是生成一个或多个具体的线索，帮助查找文章中的支持证据。
2. 线索应以文本片段的形式呈现，这些片段将有助于回答问题。
3. 只输出线索。如果有多个线索，请用换行符分隔。
4. 请用中文回答。""",

step2：生成线索问题

 surrogate_queries = self.mem_model.rewrite(query)

执行原理：

def rewrite(
        self,
        query, max_new_tokens=128) -> str:
        return self.generate(self.prompts["sur"], query, max_new_tokens=max_new_tokens)[0]

调用prompt:

"sur": """
你会得到一个与文章相关的问题。为了有效地回答这个问题，你需要回想文章中的具体细节。你的任务是生成精确的线索问题，帮助找到文章中必要的信息。


### 问题：{question}
### 指示：
1. 你对文章有一个大致的理解。你的任务是生成一个或多个具体的线索，帮助查找文章中的支持证据。
2. 线索应以精确的替代问题形式呈现，澄清原问题。
3. 只输出线索。如果有多个线索，请用换行符分隔。
4. 请用中文回答。""",

step3：生成召回问题（线索文本+线索问题+草稿答案并做长度过滤，<3去掉）

  retrieval_query, potential_answer = self._prepare_retrieval_query(query, text_spans, surrogate_queries, use_memory_answer)

执行原理：

step4：召回相关文档

       retrieval_results = self._retrieve(retrieval_query)
        if potential_answer:
            retrieval_results.append(f"The answer might be {potential_answer}.")

执行原理：

def answer(
        self,
        query, max_new_tokens=128) -> str:
        return self.generate(self.prompts["qa"], query, max_new_tokens=max_new_tokens)[0]

调用prompt:

"qa": """   你会得到一个与文章相关的问题。你的任务是直接回答这个问题。      ### 问题：{question}   ### 指示：   基于文章的内容，直接回答问题。不要包含除答案之外的任何额外内容。""",

step5：文档+草稿答案拼接prompt

knowledge = "\n\n".join(retrieval_results)

step6：生成答案

  return self._generate_response("qa_gen", query, knowledge, prompt_template, max_new_tokens)

总结

本文主要介绍了MemoRAG: Moving towards Next-Gen RAG Via Memory-Inspired Knowledge Discovery》，https://arxiv.org/pdf/2409.05591的具体代码，结论就是，这个跟graphrag一样，是个特定的trick，用空间来换性能，所以会受到很大的落地挑战。

所以，无论是graphrag，还是memorag，都不用吹，没什么太落地的东西，trick思维居多，rag落地，还是要老老实实地走传统那一套，做知识库等的优化。