RAG再進(jìn)化?基于長(zhǎng)期記憶的檢索增強(qiáng)生成新范式-MemoRAG 原創(chuàng)
前言
RAG現(xiàn)在工作很多,進(jìn)化的也很快,再來(lái)看看一個(gè)新的RAG工作-MemoRAG。
文章提出,RAG在減少大模型對(duì)于垂類(lèi)知識(shí)的問(wèn)答幻覺(jué)上取得了不錯(cuò)的效果,也成為私域知識(shí)問(wèn)答的一種范式。然而,傳統(tǒng)RAG系統(tǒng)主要適用于明確信息需求的問(wèn)答任務(wù),但在處理涉及模糊信息需求或非結(jié)構(gòu)化知識(shí)的復(fù)雜任務(wù)時(shí)表現(xiàn)不佳。因?yàn)?,現(xiàn)實(shí)世界中的許多問(wèn)題信息需求是模糊的,外部知識(shí)是非結(jié)構(gòu)化的,例如理解書(shū)籍中主要角色之間的相互關(guān)系。
因此,研究難點(diǎn)在于:
- 如何有效處理模糊的信息需求
- 如何從非結(jié)構(gòu)化知識(shí)中提取有用信息
- 如何在長(zhǎng)文本上下文中進(jìn)行有效的信息檢索和生成
本文介紹的MemoRAG,一種基于長(zhǎng)期記憶的檢索增強(qiáng)生成新范式。
方法
MemoRAG
- 傳統(tǒng)RAG
MemoRAG:MemoRAG提出了一個(gè)雙系統(tǒng)架構(gòu),采用了一個(gè)輕量級(jí)但長(zhǎng)上下文的LLM來(lái)形成數(shù)據(jù)庫(kù)的全局記憶,并在任務(wù)呈現(xiàn)時(shí)生成草稿答案,提示檢索工具在數(shù)據(jù)庫(kù)中定位有用信息。另一方面,它利用一個(gè)能力較強(qiáng)的LLM,根據(jù)檢索到的信息生成最終答案。
MemoRAG的核心是引入了一個(gè)記憶模塊:
記憶模塊設(shè)計(jì)
1.輸入
輸入序列 X 包含 n 個(gè)標(biāo)記,表示為X1,...,Xn。
2.標(biāo)注注意力機(jī)制
4.記憶模塊訓(xùn)練:
- 訓(xùn)練過(guò)程
記憶模塊的訓(xùn)練分為兩個(gè)階段:
預(yù)訓(xùn)練:使用來(lái)自RedPajama數(shù)據(jù)集的隨機(jī)抽樣長(zhǎng)上下文對(duì)模型進(jìn)行預(yù)訓(xùn)練,使記憶模塊能夠從原始上下文中學(xué)習(xí)如何形成記憶。
指令微調(diào)(SFT):使用特定任務(wù)的SFT數(shù)據(jù),使MemoRAG能夠基于形成的記憶生成特定任務(wù)的線(xiàn)索。
- 訓(xùn)練目標(biāo)
通過(guò)這種設(shè)計(jì),記憶模塊能夠有效地將大量原始上下文壓縮成少量的記憶標(biāo)記,同時(shí)保留關(guān)鍵的語(yǔ)義信息,從而在處理長(zhǎng)上下文和高層次查詢(xún)時(shí)提供顯著的優(yōu)勢(shì)。
基本使用
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}")
運(yùn)行上述代碼時(shí),編碼后的鍵值 (KV) 緩存、Faiss 索引和分塊段落都存儲(chǔ)在指定的 中save_dir。之后,如果再次使用相同的上下文,則可以快速?gòu)拇疟P(pán)加載數(shù)據(jù):
pipe.load("cache/harry_potter/", print_stats=True)
摘要任務(wù)
res = pipe(cnotallow=context, task_type="summarize", max_new_tokens=512)
print(f"MemoRAG summary of the full book:\n {res}")
實(shí)驗(yàn)
參考文獻(xiàn)
- paper:MEMORAG: MOVING TOWARDS NEXT-GEN RAG VIA MEMORY-INSPIRED KNOWLEDGE DISCOVERY,https://arxiv.org/pdf/2409.05591v2
- code:https://github.com/qhjqhj00/MemoRAG
本文轉(zhuǎn)載自公眾號(hào)大模型自然語(yǔ)言處理 作者:余俊暉
原文鏈接:??https://mp.weixin.qq.com/s/f5GiIcjtEuxtOitPgfsjcQ??
