前言

RAG現(xiàn)在工作很多，進化的也很快，再來看看一個新的RAG工作-MemoRAG。

文章提出，RAG在減少大模型對于垂類知識的問答幻覺上取得了不錯的效果，也成為私域知識問答的一種范式。然而，傳統(tǒng)RAG系統(tǒng)主要適用于明確信息需求的問答任務(wù)，但在處理涉及模糊信息需求或非結(jié)構(gòu)化知識的復(fù)雜任務(wù)時表現(xiàn)不佳。因為，現(xiàn)實世界中的許多問題信息需求是模糊的，外部知識是非結(jié)構(gòu)化的，例如理解書籍中主要角色之間的相互關(guān)系。

因此，研究難點在于：

• 如何有效處理模糊的信息需求
• 如何從非結(jié)構(gòu)化知識中提取有用信息
• 如何在長文本上下文中進行有效的信息檢索和生成

本文介紹的MemoRAG，一種基于長期記憶的檢索增強生成新范式。

方法

MemoRAG

• 傳統(tǒng)RAG

MemoRAG：MemoRAG提出了一個雙系統(tǒng)架構(gòu)，采用了一個輕量級但長上下文的LLM來形成數(shù)據(jù)庫的全局記憶，并在任務(wù)呈現(xiàn)時生成草稿答案，提示檢索工具在數(shù)據(jù)庫中定位有用信息。另一方面，它利用一個能力較強的LLM，根據(jù)檢索到的信息生成最終答案。
MemoRAG的核心是引入了一個記憶模塊：

記憶模塊設(shè)計

1.輸入
輸入序列 X 包含 n 個標記，表示為X₁，...，X_n。

2.標注注意力機制

4.記憶模塊訓(xùn)練：

• 訓(xùn)練過程
  記憶模塊的訓(xùn)練分為兩個階段：

預(yù)訓(xùn)練：使用來自RedPajama數(shù)據(jù)集的隨機抽樣長上下文對模型進行預(yù)訓(xùn)練，使記憶模塊能夠從原始上下文中學(xué)習(xí)如何形成記憶。

指令微調(diào)（SFT）：使用特定任務(wù)的SFT數(shù)據(jù)，使MemoRAG能夠基于形成的記憶生成特定任務(wù)的線索。

• 
  
• 訓(xùn)練目標

• 

通過這種設(shè)計，記憶模塊能夠有效地將大量原始上下文壓縮成少量的記憶標記，同時保留關(guān)鍵的語義信息，從而在處理長上下文和高層次查詢時提供顯著的優(yōu)勢。

基本使用

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。之后，如果再次使用相同的上下文，則可以快速從磁盤加載數(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}")

實驗

參考文獻

• 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)載自公眾號大模型自然語言處理  作者：余俊暉

原文鏈接：??https://mp.weixin.qq.com/s/f5GiIcjtEuxtOitPgfsjcQ??