[LLM] 자기 추론 프레임워크

바이두에서 발표한 '자기 추론을 통한 검색 증가 언어모델의 향상(Improving Retrieval Augmented Language Model with Self-Reasoning)'이라는 논문이 아카이브에 게재되었다.

LLM의 성능을 개선하기 위해서 도입된 RAG(Retrieval-Augmented Generation)의 경우, 대규모 언어 모델의 출력을 최적화하여 응답을 생성하기 전에 학습 데이터 소스 외부의 신뢰할 수 있는 지식 베이스를 참조하도록 하는 프로세스인데 이 과정에서 LLM에 주어진 쿼리와 연관된 지식이 지식 베이스에 없거나, 찾는 과정에서 잘못된 지식을 찾을 경우 응답을 생성하는 과정에서 환각 현상이 발생할 가능성이 크다.

 

해당 논문에서는 이러한 환각 현상을 줄이는 것을 목표로 'LLM 자체에서 생성된 추론 경로'를 활용하였다고 한다.

여기서 말하는 'LLM 자체에서 생성된 추론 경로'란 1) 관련성 인식 프로세스(relevance-aware process) 2) 증거 인식 선택 프로세스(evidence-aware selective process) 3) 경로 분석 프로세스(trajectory analysis process) 등 3단계 추론 경로를 말하는데 아래의 그림을 참조하여 자세히 설명하면 

1) 관련성 인식 프로세스(relevance-aware process)
우선 DPR 및 Contriever을 기본 검색기 R로 선택하여 상위 k 관련 문서를 회수한 이후, 이 k개의 문서에 대해서 질문과의 연관성 및 그 이유에 대해서 서술하도록 지시한다. 

2) 증거 인식 선택 프로세스(evidence-aware selective process) 

위에서 연관성이 있다고 생각된 LLM에 관련 문서를 선택하고 선택된 문서의 주요 문장 스니펫을 자동으로 선택하도록 지시하는데 이 때도 역시 LLM에 선택된 스니펫이 질문에 답변할 수 있는 이유를 출력하도록 요청한다. 

3) 경로 분석 프로세스(trajectory analysis process)

마지막으로 이전 프로세스의 모든 자기 추론 궤적을 통합하여 추론 스니펫 체인을 형성하여 검색 증강 생성의 전반적인 성능을 향상시킵니다.특히 LLM에 자체 내에서 추론 궤적을 분석하고 궁극적으로 간결한 분석과 짧은 답변을 출력하도록 요청한다.

 

이를 간단하게 구현한 코드는 아래와 같다. 

 

from transformers import DPRQuestionEncoder, DPRContextEncoder
from sentence_transformers import SentenceTransformer
import torch
import faiss
from langchain_huggingface import HuggingFaceEndpoint,

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-27b-it")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-27b-it",
    device_map="auto",
    torch_dtype=torch.bfloat16
)

# DPR model load
question_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
context_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

# Contriever load
contriever = SentenceTransformer('facebook/contriever')

# 문서 데이터베이스 (예시)
documents = [
    "DPR은 Dense Passage Retrieval의 약자입니다.",
    "Contriever는 대조 학습을 사용하는 검색 모델입니다.",
    "검색 시스템은 정보 접근성을 향상시킵니다."
]

dpr_doc_embeddings = context_encoder(documents, return_tensors="pt").pooler_output
contriever_doc_embeddings = contriever.encode(documents, convert_to_tensor=True)

# FAISS 인덱스 생성
dimension = dpr_doc_embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
index.add(dpr_doc_embeddings.cpu().numpy())

# 해당 논문에 DPR과 Contriever를 사용한다고 기술되어 있기는 하나 정확한 방식에 대해서 기재되어 있지 않아, score를 결합하여 종합적으로 판단함

def hybrid_retrieve(query, top_k=5):
    # DPR 검색
    q_embedding_dpr = question_encoder(query, return_tensors="pt").pooler_output
    dpr_scores, dpr_indices = index.search(q_embedding_dpr.cpu().numpy(), top_k)
    
    # Contriever 검색
    q_embedding_contriever = contriever.encode(query, convert_to_tensor=True)
    contriever_scores = torch.matmul(q_embedding_contriever, contriever_doc_embeddings.T)
    
    # 스코어 정규화
    dpr_scores = (dpr_scores - dpr_scores.min()) / (dpr_scores.max() - dpr_scores.min())
    contriever_scores = (contriever_scores - contriever_scores.min()) / (contriever_scores.max() - contriever_scores.min())
    
    # 스코어 결합 (예: 평균)
    combined_scores = (dpr_scores + contriever_scores.numpy()) / 2
    
    # 최종 순위 결정
    final_indices = combined_scores.argsort()[::-1][:top_k]
    
    return [documents[i] for i in final_indices]

def format_retrieved_docs(docs):
    formatted_docs = []
    for i, doc in enumerate(docs, start=1):
        formatted_docs.append(f"문장{i}: {doc}")
    return "\n".join(formatted_docs)
    
qeury = '검색시스템은 무엇인가요?'

task_1 = f'''query: {query} 
		    documents : {format_retrieved_docs(hybrid_retrieve(query))}
            다음 query와 documents를 보고, query와 연관성이 있는 것으로 보이는 document를 고르고 
            그 이유에 대해서 서술하시오.
'''

input_ids = tokenizer(task_1, return_tensors="pt").to("cuda")

with torch.no_grad():
	outputs = model.generate(**input_ids)
output = tokenizer.decode(outputs[0])

task_2 = f'''query: {query} 
		    output : {output}
            query와 output을 보고 ouput의 연관성 있는 문장으로 부터 
			query의 답변을 하는데 유의미한 문장을 추출하고 그 이유를 설명하라
'''

input_ids = tokenizer(task_2, return_tensors="pt").to("cuda")

with torch.no_grad():
	outputs = model.generate(**input_ids)
output_2 = tokenizer.decode(outputs[0])

task_2 = f'''query: {query} 
		    추론 과정 1: {output}
            추론 과정 2: {output2}
            위의 추론 과정을 보고 query에 대하여 적절한 응답을 하시오 
'''

input_ids = tokenizer(task_2, return_tensors="pt").to("cuda")

with torch.no_grad():
	outputs = model.generate(**input_ids)
output3 = tokenizer.decode(outputs[0])