RAG再進化?基于長期記憶的檢索增強生成新范式-MemoRAG 原創
前言
RAG現在工作很多,進化的也很快,再來看看一個新的RAG工作-MemoRAG。
文章提出,RAG在減少大模型對于垂類知識的問答幻覺上取得了不錯的效果,也成為私域知識問答的一種范式。然而,傳統RAG系統主要適用于明確信息需求的問答任務,但在處理涉及模糊信息需求或非結構化知識的復雜任務時表現不佳。因為,現實世界中的許多問題信息需求是模糊的,外部知識是非結構化的,例如理解書籍中主要角色之間的相互關系。
因此,研究難點在于:
- 如何有效處理模糊的信息需求
- 如何從非結構化知識中提取有用信息
- 如何在長文本上下文中進行有效的信息檢索和生成
本文介紹的MemoRAG,一種基于長期記憶的檢索增強生成新范式。
方法
MemoRAG
- 傳統RAG
MemoRAG:MemoRAG提出了一個雙系統架構,采用了一個輕量級但長上下文的LLM來形成數據庫的全局記憶,并在任務呈現時生成草稿答案,提示檢索工具在數據庫中定位有用信息。另一方面,它利用一個能力較強的LLM,根據檢索到的信息生成最終答案。
MemoRAG的核心是引入了一個記憶模塊:
記憶模塊設計
1.輸入
輸入序列 X 包含 n 個標記,表示為X1,...,Xn。
2.標注注意力機制
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(cnotallow=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(cnotallow=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
本文轉載自公眾號大模型自然語言處理 作者:余俊暉
原文鏈接:??https://mp.weixin.qq.com/s/f5GiIcjtEuxtOitPgfsjcQ??
