Skip to main content

프롬프트해커 대니의 최신 소식 구독하기!

Rerank란? RAG 성능을 높이는 5가지 방법 총정리

검색 증강 생성(Retrieval-Augmented Generation, RAG)은 대규모 언어 모델(LLM)의 성능을 크게 향상시킬 수 있는 강력한 기술입니다. 하지만 RAG 시스템이 복잡해질수록 기본적인 검색과 생성만으로는 높은 수준의 성능을 달성하기 어려워집니다. 이때 Rerank라는 기술을 활용하면 RAG의 성능을 한 단계 더 끌어올릴 수 있습니다. 이 글에서는 Rerank의 개념과 RAG 성능 향상을 위한 5가지 주요 방법에 대해 자세히 알아보겠습니다.

Rerank란 무엇인가?​

Rerank는 검색 결과의 순위를 재조정하는 과정을 말합니다. RAG 시스템에서 Rerank는 초기 검색 결과에서 가져온 문서들의 순위를 다시 매기는 역할을 합니다. 이를 통해 사용자의 질문과 가장 관련성 높은 문서들을 상위에 배치하여 LLM이 더 정확한 답변을 생성할 수 있도록 돕습니다.

Rerank의 핵심은 Cross-encoder 구조를 활용한다는 점입니다. 일반적인 벡터 검색에서 사용하는 Bi-encoder와 달리, Cross-encoder는 질문과 문서를 동시에 입력으로 받아 더 정확한 관련성 점수를 계산할 수 있습니다. 이를 통해 단순한 벡터 유사도 비교보다 훨씬 정교한 순위 조정이 가능해집니다.

RAG 성능을 높이는 5가지 Rerank 방법​

1. 한국어 특화 Reranker 활용​

한국어 RAG 시스템의 성능을 높이기 위해서는 한국어에 특화된 Reranker를 사용하는 것이 중요합니다. 예를 들어, BAAI/bge-reranker-large 모델을 한국어 데이터로 파인튜닝한 'Dongjin-kr/ko-reranker' 모델을 활용할 수 있습니다. 이 모델은 한국어의 특성을 잘 반영하여 더 정확한 순위 조정을 수행할 수 있습니다.

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

model_path = "Dongjin-kr/ko-reranker"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

def rerank(query, documents):
    pairs = [[query, doc] for doc in documents]
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    with torch.no_grad():
        scores = model(**inputs, return_dict=True).logits.view(-1,).float()
    return sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)

2. 다국어 지원 Reranker 적용​

글로벌 서비스를 제공하는 경우, 다양한 언어를 지원하는 Reranker를 사용하는 것이 효과적입니다. 'BAAI/bge-reranker-v2-m3' 모델은 여러 언어를 지원하면서도 빠른 연산 속도를 자랑합니다. 이 모델을 활용하면 다국어 환경에서도 일관된 성능의 Rerank를 수행할 수 있습니다.

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

model_path = "BAAI/bge-reranker-v2-m3"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

def rerank_multilingual(query, documents):
    pairs = [[query, doc] for doc in documents]
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    with torch.no_grad():
        scores = model(**inputs, return_dict=True).logits.view(-1,).float()
    return sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)

3. 두 단계 검색 전략 구현​

Rerank의 연산 비용을 고려하여 두 단계 검색 전략을 구현하는 것이 효율적입니다. 첫 번째 단계에서는 빠른 벡터 검색을 통해 후보 문서를 추출하고, 두 번째 단계에서 Reranker를 적용하여 순위를 정교하게 조정합니다. 이 방식을 통해 검색 속도와 정확도 사이의 균형을 맞출 수 있습니다.

from sentence_transformers import SentenceTransformer, util

# 1단계: 벡터 검색
embedder = SentenceTransformer('distilbert-base-nli-mean-tokens')
corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)

def vector_search(query, top_k=100):
    query_embedding = embedder.encode(query, convert_to_tensor=True)
    cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
    top_results = torch.topk(cos_scores, k=top_k)
    return [corpus[idx] for idx in top_results[1]]

# 2단계: Rerank
def two_stage_search(query, corpus, top_k=10):
    candidates = vector_search(query, top_k=100)
    reranked = rerank(query, candidates)
    return reranked[:top_k]

4. 컨텍스트 최적화​

Rerank 결과를 LLM에 전달할 때, 컨텍스트 윈도우 크기를 고려한 최적화가 필요합니다. 가장 관련성 높은 정보를 컨텍스트의 앞부분에 배치하고, 중복되거나 불필요한 정보는 제거하여 LLM이 효과적으로 정보를 활용할 수 있도록 합니다.

def optimize_context(reranked_docs, max_tokens=2000):
    context = ""
    for doc, score in reranked_docs:
        if len(context) + len(doc) > max_tokens:
            break
        context += doc + "\n\n"
    return context.strip()

def rag_with_rerank(query, corpus):
    reranked = two_stage_search(query, corpus)
    context = optimize_context(reranked)
    llm_response = generate_llm_response(query, context)
    return llm_response

5. 동적 Reranker 선택​

질문의 특성에 따라 다양한 Reranker를 동적으로 선택하는 전략을 구현할 수 있습니다. 예를 들어, 일반적인 질문에는 경량화된 모델을, 전문적인 질문에는 더 복잡한 모델을 사용하는 방식입니다. 이를 통해 계산 리소스를 효율적으로 사용하면서도 높은 성능을 유지할 수 있습니다.

def select_reranker(query):
    if is_complex_query(query):
        return complex_reranker
    else:
        return simple_reranker

def dynamic_rerank(query, documents):
    reranker = select_reranker(query)
    return reranker(query, documents)

결론​

Rerank 기술은 RAG 시스템의 성능을 크게 향상시킬 수 있는 강력한 도구입니다. 한국어 특화 모델 활용, 다국어 지원, 두 단계 검색 전략, 컨텍스트 최적화, 그리고 동적 Reranker 선택 등의 방법을 통해 RAG 시스템의 정확도와 효율성을 높일 수 있습니다. 이러한 기술들을 적절히 조합하고 최적화하면, 사용자의 질문에 대해 더욱 정확하고 관련성 높은 답변을 제공할 수 있을 것입니다.

Rerank 기술의 적용은 단순히 검색 결과의 순위를 바꾸는 것 이상의 의미를 갖습니다. 이는 LLM이 더 나은 컨텍스트를 바탕으로 답변을 생성할 수 있게 하여, 궁극적으로는 AI 시스템의 전반적인 성능과 신뢰성을 향상시키는 데 기여합니다. 따라서 RAG 시스템을 개발하거나 개선할 때 Rerank 기술의 도입을 적극적으로 고려해볼 만합니다.

참고 자료​