RAG 성능 향상을 위한 Reranker 구현

본 글은 Claude Sonnet 4로 작성한 후 수정한 글입니다.


Reranker의 이론적 원리

정보 검색에서 Reranker는 전통적인 "retrieve-then-rerank" 패러다임의 핵심 구성 요소이다. 이 접근법은 두 단계로 이루어진다: 먼저 효율적인 검색 방법으로 후보 문서들을 수집하고, 그 다음 더 정교한 모델로 이들을 재정렬한다.

초기 검색 단계에서는 벡터 유사도나 BM25와 같은 빠른 검색 방법을 사용하여 수천 개의 문서에서 수십 개의 후보를 선별한다. 이 과정은 높은 recall을 목표로 하며, 관련성이 있을 만한 문서들을 놓치지 않는 것에 중점을 둔다.

두 번째 단계인 재정렬에서는 계산 비용이 높지만 정확도가 뛰어난 모델을 사용한다. 이 모델은 각 쿼리-문서 쌍을 정밀하게 분석하여 실제 관련성을 평가하고, 최종 순위를 결정한다. 이러한 두 단계 접근법은 효율성과 정확성을 모두 확보할 수 있는 균형잡힌 해결책이다.

Bi-Encoder vs Cross-Encoder 비교

RAG 시스템에서 문서 임베딩을 생성하는 방법은 크게 두 가지로 나뉜다:

Bi-Encoder 구조

Bi-Encoder는 쿼리와 문서를 독립적으로 인코딩하여 각각의 임베딩 벡터를 생성한다. 이후 코사인 유사도와 같은 방법으로 두 벡터를 비교한다. 이 접근법의 장점은 문서들을 미리 인코딩하여 벡터 데이터베이스에 저장할 수 있다는 점이다. 검색 시에는 쿼리만 인코딩하면 되므로 매우 빠른 검색이 가능하다.

하지만 Bi-Encoder는 본질적인 한계가 있다. 쿼리와 문서를 독립적으로 처리하기 때문에 둘 사이의 복잡한 상호작용을 포착하지 못한다. 예를 들어, 쿼리의 특정 부분이 문서의 특정 부분과 어떻게 연관되는지 파악하기 어렵다.

Cross-Encoder 구조

Cross-Encoder는 쿼리와 문서를 동시에 Transformer 네트워크에 입력한다. 이를 통해 두 텍스트 간의 깊은 상호작용을 모델링할 수 있다. Transformer의 attention mechanism이 쿼리의 각 토큰과 문서의 각 토큰 사이의 관계를 학습하여 더 정교한 관련성 판단을 내릴 수 있다.

Cross-Encoder는 분류 문제로 설계되어 주어진 쿼리-문서 쌍에 대해 관련성 점수를 직접 출력한다. Cross-Encoder 모델은 데이터에 대한 벡터 임베딩을 생성하지 않고, 대신 데이터 쌍에 대한 분류 메커니즘을 사용한다.

Cross-Encoder의 동작 원리

Cross-Encoder는 다음과 같은 구조로 동작한다:

  1. 입력 구성: 쿼리와 문서를 "[CLS] 쿼리 [SEP] 문서 [SEP]" 형태로 결합
  2. 토큰화: 결합된 텍스트를 BERT 토크나이저로 처리
  3. 인코딩: Transformer 레이어들이 모든 토큰 간의 상호작용을 모델링
  4. 분류: [CLS] 토큰의 표현을 사용해 관련성 점수 예측

이 과정에서 중요한 것은 쿼리와 문서의 모든 토큰이 서로 attention을 주고받을 수 있다는 점이다. 이를 통해 단순한 의미적 유사성을 넘어서 구체적인 맥락과 의도를 파악할 수 있다.

왜 Reranker가 필요한가?

이는 정보 검색에서 흔한 기술로, 먼저 가장 관련성이 높은 문서들을 검색한 후 더 정확한 모델을 사용하여 재정렬하는 방법이다. 이러한 접근법이 필요한 이유는 다음과 같다:

  1. 의미적 정확성: 벡터 유사도만으로는 놓칠 수 있는 미묘한 관련성을 포착
  2. 컨텍스트 이해: 쿼리의 의도와 문서의 실제 내용 간의 정합성 평가
  3. 다양성 보장: 상위 결과의 정보적 다양성 확보

기본 RAG 시스템 구현

먼저 기본적인 RAG 시스템을 구축해보자:

import faiss
import numpy as np
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader
import openai
from typing import List, Tuple

class BasicRAGSystem:
    def __init__(self, openai_api_key: str):
        """
        기본 RAG 시스템 초기화
        
        Args:
            openai_api_key: OpenAI API 키
        """
        self.embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
        self.vector_store = None
        self.documents = []
        
    def build_vector_store(self, documents: List[str]):
        """
        FAISS 벡터 저장소 구축
        
        Args:
            documents: 문서 리스트
        """
        # 문서 분할
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        
        # 문서를 청크로 분할
        splits = []
        for doc in documents:
            splits.extend(text_splitter.split_text(doc))
        
        self.documents = splits
        
        # FAISS 벡터 저장소 생성
        self.vector_store = FAISS.from_texts(
            texts=splits,
            embedding=self.embeddings
        )
        
    def retrieve_documents(self, query: str, k: int = 5) -> List[Tuple[str, float]]:
        """
        문서 검색
        
        Args:
            query: 검색 쿼리
            k: 검색할 문서 수
            
        Returns:
            검색된 문서와 유사도 점수 리스트
        """
        if not self.vector_store:
            raise ValueError("벡터 저장소가 구축되지 않았습니다.")
            
        # 유사도 검색 실행
        results = self.vector_store.similarity_search_with_score(query, k=k)
        
        # 결과를 (문서, 점수) 형태로 반환
        return [(doc.page_content, score) for doc, score in results]

Cross-Encoder 기반 Reranker 구현

이제 Cross-Encoder를 사용한 Reranker를 구현해보자:

from sentence_transformers import CrossEncoder
import torch
from typing import List, Tuple

class CrossEncoderReranker:
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        """
        Cross-Encoder 기반 Reranker 초기화
        
        Args:
            model_name: 사용할 cross-encoder 모델명
        """
        self.model = CrossEncoder(model_name)
        
    def rerank(self, query: str, documents: List[str]) -> List[Tuple[str, float]]:
        """
        Cross-Encoder를 사용한 문서 재정렬
        
        Args:
            query: 검색 쿼리
            documents: 재정렬할 문서 리스트
            
        Returns:
            재정렬된 문서와 관련성 점수 리스트
        """
        if not documents:
            return []
            
        # 쿼리-문서 쌍 생성
        query_doc_pairs = [(query, doc) for doc in documents]
        
        # Cross-Encoder로 관련성 점수 계산
        # 모델이 각 쌍을 동시에 처리하여 상호작용 모델링
        scores = self.model.predict(query_doc_pairs)
        
        # 점수를 기반으로 문서 정렬 (높은 점수가 더 관련성 높음)
        scored_docs = list(zip(documents, scores))
        scored_docs.sort(key=lambda x: x[1], reverse=True)
        
        return scored_docs

class EnhancedRAGSystem(BasicRAGSystem):
    def __init__(self, openai_api_key: str, reranker_model: str = None):
        """
        Cross-Encoder Reranker가 포함된 향상된 RAG 시스템
        
        Args:
            openai_api_key: OpenAI API 키
            reranker_model: 사용할 cross-encoder 모델명
        """
        super().__init__(openai_api_key)
        self.reranker = CrossEncoderReranker(reranker_model) if reranker_model else None
        
    def retrieve_and_rerank(self, query: str, initial_k: int = 20, final_k: int = 5) -> List[Tuple[str, float]]:
        """
        문서 검색 후 Cross-Encoder로 재정렬
        
        Args:
            query: 검색 쿼리
            initial_k: 초기 검색할 문서 수 (더 많은 후보 확보)
            final_k: 최종 선택할 문서 수
            
        Returns:
            재정렬된 상위 문서들
        """
        # 1차 검색: Bi-Encoder로 빠른 후보 선별
        initial_results = self.retrieve_documents(query, k=initial_k)
        documents = [doc for doc, _ in initial_results]
        
        if not self.reranker:
            # Reranker가 없으면 기본 결과 반환
            return initial_results[:final_k]
        
        # 2차 재정렬: Cross-Encoder로 정밀한 관련성 평가
        reranked_results = self.reranker.rerank(query, documents)
        
        # 상위 final_k개 문서 반환
        return reranked_results[:final_k]

실제 사용 예시

완성된 시스템을 사용하는 예시이다:

def main():
    # API 키 설정
    openai_api_key = "your-openai-api-key"
    
    # Cross-Encoder 기반 향상된 RAG 시스템 초기화
    rag_system = EnhancedRAGSystem(
        openai_api_key=openai_api_key,
        reranker_model="cross-encoder/ms-marco-MiniLM-L-6-v2"
    )
    
    # 샘플 문서 데이터
    documents = [
        "Python은 간결하고 읽기 쉬운 프로그래밍 언어이다.",
        "머신러닝에서 Python은 가장 널리 사용되는 언어 중 하나이다.",
        "FAISS는 Facebook에서 개발한 고성능 벡터 검색 라이브러리이다.",
        "OpenAI의 임베딩 모델은 텍스트를 벡터로 변환하는 데 사용된다.",
        "Langchain은 LLM 애플리케이션 개발을 위한 프레임워크이다.",
        "Cross-Encoder는 쿼리와 문서를 동시에 처리하여 정확한 관련성을 평가한다.",
        "Transformer 모델의 attention mechanism은 토큰 간 관계를 학습한다."
    ]
    
    # 벡터 저장소 구축
    rag_system.build_vector_store(documents)
    
    # 검색 쿼리
    query = "Python 머신러닝 개발"
    
    # Bi-Encoder만 사용한 기본 검색
    print("=== Bi-Encoder 기본 검색 결과 ===")
    basic_results = rag_system.retrieve_documents(query, k=3)
    for i, (doc, score) in enumerate(basic_results, 1):
        print(f"{i}. [유사도: {score:.4f}] {doc}")
    
    print("\n=== Cross-Encoder Reranker 적용 결과 ===")
    # Cross-Encoder Reranker 적용 검색
    reranked_results = rag_system.retrieve_and_rerank(query, initial_k=5, final_k=3)
    for i, (doc, score) in enumerate(reranked_results, 1):
        print(f"{i}. [관련성: {score:.4f}] {doc}")

if __name__ == "__main__":
    main()

성능 최적화 팁

Cross-Encoder Reranker의 성능을 더욱 향상시키기 위한 몇 가지 팁이다:

1. 적절한 initial_k 값 선택

# 너무 작으면 관련 문서를 놓칠 수 있음
# 너무 크면 계산 비용이 증가함
optimal_initial_k = min(20, len(documents) // 2)

2. 배치 처리를 통한 속도 향상

class BatchCrossEncoderReranker(CrossEncoderReranker):
    def rerank_batch(self, queries: List[str], documents_batch: List[List[str]]) -> List[List[Tuple[str, float]]]:
        """
        배치 단위로 Cross-Encoder 재정렬 수행
        """
        all_pairs = []
        batch_sizes = []
        
        for query, docs in zip(queries, documents_batch):
            pairs = [(query, doc) for doc in docs]
            all_pairs.extend(pairs)
            batch_sizes.append(len(pairs))
        
        # 한 번에 모든 쿼리-문서 쌍 처리
        all_scores = self.model.predict(all_pairs)
        
        # 결과를 쿼리별로 분할
        results = []
        start_idx = 0
        for batch_size in batch_sizes:
            end_idx = start_idx + batch_size
            scores = all_scores[start_idx:end_idx]
            docs = [pair[1] for pair in all_pairs[start_idx:end_idx]]
            
            scored_docs = list(zip(docs, scores))
            scored_docs.sort(key=lambda x: x[1], reverse=True)
            results.append(scored_docs)
            
            start_idx = end_idx
            
        return results

3. 캐싱을 통한 효율성 향상

from functools import lru_cache
import hashlib

class CachedCrossEncoderReranker(CrossEncoderReranker):
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        super().__init__(model_name)
        self.cache = {}
        
    def _get_cache_key(self, query: str, documents: List[str]) -> str:
        """캐시 키 생성"""
        content = query + "||" + "||".join(sorted(documents))
        return hashlib.md5(content.encode()).hexdigest()
    
    def rerank(self, query: str, documents: List[str]) -> List[Tuple[str, float]]:
        """캐시를 활용한 재정렬"""
        cache_key = self._get_cache_key(query, documents)
        
        if cache_key in self.cache:
            return self.cache[cache_key]
        
        result = super().rerank(query, documents)
        self.cache[cache_key] = result
        
        return result

관련 연구 및 주요 논문

RAG 시스템에서 reranker는 검색된 문서들을 다시 순위화하여 가장 관련성이 높은 문서를 선별하는 중요한 역할을 한다. 최근 연구들은 다양한 reranker 설계와 학습 방법을 통해 RAG 시스템의 성능을 크게 향상시키고 있다.

1. G-RAG: 그래프 기반 Reranking을 통한 RAG 개선

출처: "Don't Forget to Connect! Improving RAG with Graph-based Reranking" (2024)

G-RAG는 문서 간의 연결성을 고려하는 그래프 신경망(GNN) 기반 reranker를 제안한다. 기존 RAG 시스템이 문서 간의 연결을 무시하는 문제를 해결하기 위해, 검색된 문서들을 그래프로 구성하고 GNN을 통해 문서 간의 관계를 모델링한다. 이를 통해 직접적으로 관련성이 낮아 보이지만 다른 관련 문서와 연결된 문서들을 효과적으로 식별할 수 있다. 실험 결과 기존 방법 대비 상당한 성능 향상을 보였다.

2. HyperRAG: KV-Cache 재사용을 통한 효율성 향상

출처: "HyperRAG: Enhancing Quality-Efficiency Tradeoffs in Retrieval-Augmented Generation with Reranker KV-Cache Reuse" (2024)

HyperRAG는 reranker의 KV-Cache를 재사용하여 RAG 시스템의 품질-효율성 트레이드오프를 개선한다. 기존의 encoder-only reranker(MiniLM-L6-v2, bge-reranker-v2-m3)와 decoder-only reranker(Gemma 2B)를 비교 분석하고, KV-Cache 재사용 메커니즘을 통해 계산 비용을 크게 줄인다. 이 방법은 성능 저하 없이 reranking 과정의 효율성을 높여 실제 운영 환경에서의 적용 가능성을 크게 향상시킨다.

3. NV-RerankQA-Mistral-4B-v3: 벤치마킹과 파인튜닝

출처: "Enhancing Q&A Text Retrieval with Ranking Models: Benchmarking, fine-tuning and deploying Rerankers for RAG" (2024)

이 연구는 Q&A 텍스트 검색을 위한 ranking 모델들의 벤치마킹과 파인튜닝을 다룬다. 특히 상업적 사용이 가능한 공개 데이터셋만을 활용하여 NV-RerankQA-Mistral-4B-v3 모델을 개발했다. 다양한 공개 ranking 모델들의 성능을 체계적으로 비교하고, 각각의 장단점을 분석한다. 이 연구는 실제 RAG 시스템에서 ranking 모델 선택을 위한 실용적인 가이드라인을 제공한다.

4. Multi-Reranker: 금융 도메인에서의 성능 최적화

출처: "Multi-Reranker: Maximizing performance of retrieval-augmented generation in the FinanceRAG challenge" (2024)

Multi-Reranker는 ACM-ICAIF '24 FinanceRAG 챌린지를 위해 개발된 고성능 금융 특화 RAG 시스템이다. 쿼리 확장과 코퍼스 정제를 통한 전처리 최적화, 그리고 다중 reranker 앙상블을 통해 성능을 향상시킨다. 특히 금융 문서와 공시 자료 분석에 특화된 reranker 설계를 제안하며, 도메인 특화 RAG 시스템의 성능 최적화 방법론을 제시한다.

5. MLLM Reranker: 멀티모달 RAG를 위한 지식 강화 Reranking

출처: "MLLM Is a Strong Reranker: Advancing Multimodal Retrieval-augmented Generation via Knowledge-enhanced Reranking and Noise-injected Training" (2024)

멀티모달 대규모 언어 모델(MLLM)을 reranker로 활용하는 연구이다. 기존 텍스트 기반 reranker의 한계를 극복하기 위해 이미지와 텍스트를 함께 처리할 수 있는 MLLM을 reranker로 사용한다. 지식 강화 reranking과 노이즈 주입 훈련을 통해 멀티모달 검색 성능을 크게 향상시킨다. 이 접근법은 visual question answering과 같은 복합 모달리티 태스크에서 특히 효과적이다.

마무리

Cross-Encoder 기반 Reranker는 RAG 시스템의 검색 품질을 크게 향상시킬 수 있는 강력한 도구이다. Bi-Encoder의 효율적인 초기 검색과 Cross-Encoder의 정밀한 재정렬을 결합함으로써, 대규모 문서 컬렉션에서도 높은 정확도를 달성할 수 있다.

핵심은 쿼리와 문서 간의 깊은 상호작용을 모델링하는 Cross-Encoder의 능력에 있다. 이는 단순한 벡터 유사도를 넘어서 실제 관련성을 평가할 수 있게 해준다. 다만 추가적인 계산 비용이 발생하므로, 적절한 배치 처리와 캐싱 전략을 통해 실제 프로덕션 환경에서의 성능을 최적화하는 것이 중요하다.

향후 연구에서는 더 효율적인 Cross-Encoder 구조나 ColBERT와 같은 late interaction 방식을 통해 정확성과 효율성을 모두 만족하는 방향으로 발전할 것으로 예상된다.