MemoRAG는 경량 메모리 모델(Memory Model)이 전체 데이터베이스에 대한 글로벌 이해를 형성하고, 이를 기반으로 초안 답변을 생성한 뒤, 그 초안에서 검색 단서(retrieval clue)를 추출하여 더 정확한 검색을 수행하는 이중 시스템 RAG 아키텍처입니다. 최종 답변은 강력한 생성 모델(Generation Model)이 검색된 문서를 기반으로 생성합니다.
기존 RAG는 사용자 질문을 그대로 검색 쿼리로 사용합니다. 그러나 질문이 모호하거나 추상적일 때(“이 데이터셋의 핵심 인사이트는?”), 질문 자체로는 적절한 문서를 검색하기 어렵습니다.MemoRAG는 메모리 모델이 전체 코퍼스를 “읽은” 상태에서 초안을 생성하고, 이 초안에서 구체적인 검색 단서를 추출하여 검색 품질을 향상시킵니다.
시스템
모델
역할
Memory Model
경량 모델 (light LM)
전체 코퍼스에 대한 글로벌 이해 형성, 초안 답변 생성
Generation Model
강력한 모델 (strong LM)
검색된 문서를 기반으로 정확한 최종 답변 생성
Memory Model의 글로벌 메모리란?
Memory Model은 전체 코퍼스의 텍스트를 요약/압축한 형태의 “글로벌 메모리”를 내부적으로 보유합니다. 원 논문에서는 긴 컨텍스트를 처리할 수 있는 경량 모델을 사용하여, 전체 데이터베이스의 내용을 하나의 모델에 기억시킵니다.이 모델은 정확한 답변을 생성하는 것이 목적이 아니라, 대략적인 방향성을 제시하는 역할을 합니다. 초안의 정확도가 낮더라도 검색 단서를 추출하는 데는 충분합니다.
from langchain.chat_models import init_chat_modelfrom langchain_core.prompts import ChatPromptTemplatefrom langchain_core.output_parsers import StrOutputParser# Memory Model: 경량 모델 (글로벌 메모리 역할)memory_llm = init_chat_model("gpt-4o-mini", temperature=0.3)# Generation Model: 강력한 모델generation_llm = init_chat_model("gpt-4o", temperature=0)def generate_memory_draft(state: MemoRAGState) -> MemoRAGState: """Memory Model이 글로벌 지식을 기반으로 초안 답변을 생성합니다.""" question = state["question"] prompt = ChatPromptTemplate.from_messages([ ("system", ( "당신은 전체 문서 데이터베이스를 읽은 메모리 모델입니다.\n" "질문에 대해 기억하는 내용을 바탕으로 대략적인 답변 초안을 작성하세요.\n" "정확하지 않더라도 관련된 키워드, 개념, 방향성을 최대한 포함하세요.\n" "이 초안은 정확한 문서를 검색하기 위한 단서로 활용됩니다." )), ("human", "{question}"), ]) chain = prompt | memory_llm | StrOutputParser() draft = chain.invoke({"question": question}) return {"memory_draft": draft}
Copy
def extract_clues(state: MemoRAGState) -> MemoRAGState: """초안 답변에서 검색 단서(구체적인 검색 쿼리)를 추출합니다.""" question = state["question"] draft = state["memory_draft"] prompt = ChatPromptTemplate.from_messages([ ("system", ( "원래 질문과 초안 답변을 분석하여, 정확한 문서를 검색하기 위한 " "구체적인 검색 쿼리를 2~3개 생성하세요.\n" "각 쿼리는 서로 다른 측면을 커버해야 합니다.\n" "한 줄에 하나씩 출력하세요." )), ("human", "원래 질문: {question}\n\n초안 답변: {draft}"), ]) chain = prompt | memory_llm | StrOutputParser() result = chain.invoke({"question": question, "draft": draft}) clues = [line.strip() for line in result.strip().split("\n") if line.strip()] return {"retrieval_clues": clues}
Copy
from langchain_openai import OpenAIEmbeddingsfrom langchain_chroma import Chromaembeddings = OpenAIEmbeddings()vectorstore = Chroma(collection_name="documents", embedding_function=embeddings)retriever = vectorstore.as_retriever(search_kwargs={"k": 3})def retrieve(state: MemoRAGState) -> MemoRAGState: """검색 단서를 사용하여 관련 문서를 검색합니다.""" clues = state["retrieval_clues"] # 각 검색 단서로 검색 후 결과 병합 all_docs = [] seen_contents = set() for clue in clues: docs = retriever.invoke(clue) for doc in docs: if doc.page_content not in seen_contents: seen_contents.add(doc.page_content) all_docs.append(doc) return {"documents": all_docs}
Copy
def generate_final(state: MemoRAGState) -> MemoRAGState: """Generation Model이 검색된 문서를 기반으로 최종 답변을 생성합니다.""" question = state["question"] documents = state["documents"] prompt = ChatPromptTemplate.from_messages([ ("system", ( "다음 문서를 참고하여 질문에 정확하게 답변하세요.\n" "문서에 근거한 사실만 작성하세요.\n\n" "{context}" )), ("human", "{question}"), ]) chain = prompt | generation_llm | StrOutputParser() context = "\n\n".join(doc.page_content for doc in documents) answer = chain.invoke({"context": context, "question": question}) return {"answer": answer}
from langgraph.graph import StateGraph, START, ENDworkflow = StateGraph(MemoRAGState)# 노드 추가workflow.add_node("generate_memory_draft", generate_memory_draft)workflow.add_node("extract_clues", extract_clues)workflow.add_node("retrieve", retrieve)workflow.add_node("generate_final", generate_final)# 엣지 연결workflow.add_edge(START, "generate_memory_draft")workflow.add_edge("generate_memory_draft", "extract_clues")workflow.add_edge("extract_clues", "retrieve")workflow.add_edge("retrieve", "generate_final")workflow.add_edge("generate_final", END)# 컴파일 및 실행app = workflow.compile()result = app.invoke({"question": "이 데이터셋에서 가장 중요한 발견은 무엇인가?"})print(result["answer"])
Memory Model의 글로벌 메모리는 전체 코퍼스 내용의 압축된 표현입니다. 원 논문에서는 긴 컨텍스트를 지원하는 경량 모델에 전체 텍스트를 입력하지만, 위 예제에서는 시스템 프롬프트로 글로벌 메모리 역할을 근사합니다. 실제 구현에서는 코퍼스 요약이나 메모리 임베딩을 활용하세요.
