前言
RAG现在工作很多,进化的也很快,再来看看一个新的RAG工作-MemoRAG 。
文章提出,RAG在减少大模型对于垂类知识的问答幻觉上取得了不错的效果,也成为私域知识问答的一种范式。然而,传统RAG系统主要适用于明确信息需求的问答任务,但在处理涉及模糊信息需求或非结构化知识的复杂任务时表现不佳。因为,现实世界中的许多问题信息需求是模糊的,外部知识是非结构化的,例如理解书籍中主要角色之间的相互关系。
因此,研究难点 在于:
- 如何有效处理模糊的信息需求
- 如何从非结构化知识中提取有用信息
- 如何在长文本上下文中进行有效的信息检索和生成
本文介绍的MemoRAG,一种基于长期记忆的检索增强生成新范式。
方法
MemoRAG
- 传统RAG
这里, 表示最终答案, 表示输入查询,表示从相关数据库 中检索到的上下文, 和 分别表示生成模型和检索模型, 和 表示模型参数。
- MemoRAG :MemoRAG提出了一个双系统架构,采用了一个轻量级但长上下文的LLM来形成数据库的全局记忆,并在任务呈现时生成草稿答案,提示检索工具在数据库中定位有用信息。另一方面,它利用一个能力较强的LLM,根据检索到的信息生成最终答案。
MemoRAG的核心是引入了一个记忆模块 :
记忆模型的作用 :记忆模型 的设计目的是建立数据库 的全局记忆,并生成有助于检索的线索 。
记忆模型选型 :任何能够有效处理超长上下文的语言模型都可以作为记忆模型。文章也实现了两个记忆模型(memorag-qwen2-7b-inst和memoragmistral-7b-inst )。
- 表示由记忆模型 生成的中间答案,用作检索线索。
- 这个中间答案 帮助检索模型 从数据库 中检索最相关的上下文 。
记忆模块设计
- 输入
输入序列 X 包含 n 个标记,表示为 2. 标注注意力机制
3. 短期记忆到长期记忆的转换 为了将短期记忆转换为长期记忆,引入了记忆标记 作为LLMs中长期记忆的信息载体。假设底层LLM 的工作上下文窗口长度为 ,在每个上下文窗口后,附加 个记忆标记:
新的注意力变成:
通过多个注意力过程,原始标记被编码成隐藏状态 ,其中 表示原始标记的隐藏状态, 表示记忆标记的隐藏状态。 4. 记忆模块训练 :
- 训练过程
记忆模块的训练分为两个阶段 :
- 预训练:使用来自RedPajama数据集的随机抽样长上下文对模型进行预训练,使记忆模块能够从原始上下文中学习如何形成记忆。
- 指令微调(SFT):使用特定任务的SFT数据,使MemoRAG能够基于形成的记忆生成特定任务的线索。
- 训练目标
这个公式表示训练过程中的目标是最大化给定先前记忆标记 和最近原始标记 的情况下,下一个标记 的生成概率。
通过这种设计,记忆模块能够有效地将大量原始上下文压缩成少量的记忆标记,同时保留关键的语义信息,从而在处理长上下文和高层次查询时提供显著的优势。
基本使用
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)
摘要任务
res = pipe(context=context, task_type="summarize", max_new_tokens=512)
print(f"MemoRAG summary of the full book:\n {res}")
实验
参考文献
- paper:MEMORAG: MOVING TOWARDS NEXT-GEN RAG VIA MEMORY-INSPIRED KNOWLEDGE DISCOVERY,https://arxiv.org/pdf/2409.05591v2
- code:https://github.com/qhjqhj00/MemoRAG
往期相关
【LLM & RAG】英伟达“ChatQA 2”训练策略概述
【LLM & KBQA】FlexKBQA:一种结合LLM的KBQA框架
【LLM & RAG & text2sql】大模型在知识图谱问答上的核心算法详细思路及实践
【开源分享】KBQA核心技术及结合大模型SPARQL查询生成问答实践
【文档智能 & RAG】RAG增强之路:增强PDF解析并结构化技术路线方案及思路
【RAG】一种新颖的Agentic RAG系统,Golden-Retriever解决工业领域知识问答
【RAG】WeKnow-RAG:融合Web搜索与知识图谱的自适应检索增强生成方法