AWS 기술 블로그

한국어 Reranker를 활용한 검색 증강 생성(RAG) 성능 올리기

검색 증강 생성 (Retrieval-Augmented Generation, RAG)은 효율적인 데이터 검색과 대규모 언어 모델 (Large Language Model, LLM) 을 결합하여 정확하고 관련성 높은 응답을 생성하는 AI 기술로 부상했습니다. 특히, RAG 방식은 최신 정보를 반영하여 답변의 부정확성이나 환각을 줄인다는 면에서 많은 사용자들의 관심을 받고 있습니다. 그러나 RAG 시스템이 보다 복잡한 실제 애플리케이션에 적용 됨에 따라, 시멘틱 벡터 검색 및 답변 생성과 같은 기본적인 RAG 파이프라인 구성 만으로는 프로덕션 수준의 요구 사항을 만족 시키기에 어려울 수도 있습니다.

RAG를 개선하기 위한 고급 방법 중 하나는 순위 재정렬 (Reranking) 입니다. 간단히 말해, Reranking은 RAG가 생성한 후보 문서들에 대해 질문에 대한 관련성 및 일관성을 판단하여 문서의 우선 순위를 재 정렬 하는 것입니다. 즉, 질문과 관련성 있는 문서들을 컨텍스트의 상위권에 위치 시킴으로써, 답변의 정확도를 올리는 방법입니다.

이 글은 Reranker를 기반으로 RAG 성능을 향상시키는 방법에 대해서 다루고 있습니다. 첫 번째로 문제 정의 에서는 Reranker가 필요한 배경, 두 번째 해결전략 (Solution strategy) 섹션에서는 문제를 해결할 수 있는 Reranker의 원리, 세번째 방법론에서는 기존 RAG Pipeline과 Reranker의 통합 및 한글에 대한 Amazon SageMaker 기반 파인 튜닝 방법을 다룹니다. 마지막으로 평가 섹션에서는 평가용 데이터에 대한 성능체크를 통해 Reranker의 효과를 확인 합니다.

이 글은  Reranker의 개념 및 필요성을 이해하고, 또한 파인 튜닝 기반의 커스텀 모델을 통한 도메인 적응문제 해결에 도움이 될 수 있습니다.

문제 정의

먼저 현재 RAG의 문제점에 대해서 살펴보도록 하겠습니다. RAG는 수많은 텍스트 문서에서 의미론적 검색 (Semantic search) 과정을 수행합니다. 의미론적 검색에는 일반적으로 벡터 검색을 활용 되는데, 이 과정에서 두 가지 정보 손실이 발생하게 됩니다. 첫 번째는 문서의 임베딩 벡터 변환 과정에서의 손실입니다. 임베딩 벡터 변환이란 문서를 n개의 숫자로 재 표현하는 것을 의미하는데, 정보의 손실은 문서가 긴 경우에 정해진 벡터의 차원으로 표현하기 어려울 때 발생하게 됩니다. 두 번째는 검색 과정에서의 손실입니다. RAG는 검색 시간 단축을 위해 Approximate Nearnest Neighbor search (ANNs) 기술을 활용합니다. 이 방법은 질문과 문서 사이의 관련성 체크 횟수를 현저히 줄임으로써 검색속도를 높일 수 있습니다. 하지만 이 과정에서 발생하는 관련성 정확도 하락이 정보의 손실로 나타나게 됩니다. 이러한 정보 손실로 인해 LLM으로 전달되는 컨텍스트의 상위 k개 이내에서 질문에 대한 관련 정보가 누락되는 경우가 발생합니다.

이러한 문제는 검색 후 반환되는 문서수를 늘림 (k 증가) 으로써 간단히 해결할 수 있습니다. 하지만 이 방법은 LLM에 전달하는 컨텍스트가 늘어나기 때문에 비용 효율적이지 않습니다. 그 뿐만이 아닙니다. 최근 출판된 논문[1]에 따르면 RAG의 정확도는 관련정보의 컨텍스트 내 존재 유무가 아닌, 순서라는 것을 발견하였습니다. 즉, 관련 정보가 컨텍스트 내 상위권에 위치하고 있을 때 좋은 답변을 얻을 수 있다는 뜻 입니다

그림 1: Lost in Middle: 질문에 대한 관련 문서가 컨텍스트 중간에 위치할 경우, LLM 응답 정확도가 낮아진다.

해결 전략

앞선 내용을 통해 RAG의 성능을 향상 시키기 위해서는 질문에 대한 관련 문서가 컨텍스트에 존재 할 뿐만 아니라, 컨텍스트 내에서 그 순서 또한 상위권에 위치하고 있어야 함을 알 수 있었습니다. 이에 대한 해결책으로 Reranker를 사용할 수 있습니다.

Reranker는 질문과 문서 사이의 유사도를 측정하는 것을 목표로 합니다. 이것은 기존의 벡터 검색의 목적과 동일합니다. 하지만 질문과 문서에 대한 독립적인 임베딩을 활용하는 Bi-encoder 형태 (그림 2-a)의 벡터 검색과는 다르게 Reranker는 질문과 문서를 하나의 인풋으로 활용하는 Cross-encoder 형태 (그림 2-b)라는 점에서 차이를 보이고 있습니다. 즉, Reranker는 질문과 문서를 동시에 분석 (Self-attention) 함으로써, 독립적인 임베딩 벡터 기반의 Bi-encoder 방식에 비해 더욱 정확한 유사도 측정이 가능하다는 장점을 가지고 있습니다. 우리는 이러한 장점을 활용하여 질문-문서 사이의 관련성을 더욱 정교하게 측정할 수 있습니다.


그림2: Encoder 종류

 여기까지만 보면 Reranker가 임베딩 모델을 대체할 수 있다고 생각할 수 있습니다. 하지만 그렇치 않습니다. 질문과 문서를 동시에 인풋으로 활용하는 Reranker의 특성 상, 사용자가 질문을 하는 시점에서 모든 문서들에 대한 질문 과의 관련성을 측정이 요구됩니다. 하지만 이 연산은 굉장히 오래 걸릴 수 밖에 없습니다. 이와 관련하여 Sentence-BERT [2]에서는 4천만개 문서에 대해 V100 GPU기반으로 연관성 측정을 수행했을 때 약 50시간이 소요 된다고 보고하고 있습니다.

이러한 문제점을 극복하기 위해 두 단계 (two-stage) 전략을 적용할 수 있습니다 (그림 3). 즉, 첫번째 단계에서는 기존의 벡터 검색 방식으로 대규모 문서 중 질문과 관련성이 높을 것 같은 후보군을 검색하고, 두번째 단계에서는 검색된 문서들에 대해 reranker기반으로 관련성을 재 측정 하게 됩니다. 이렇게 되면 문석의 검색 속도를 높이면서도 (첫번째 단계) 질문과의 관련성을 정확히 측정 (두번째 단계) 하는 것이 가능해 집니다.


그림 3 : 두 단계 검색 전략 (벡터 검색 기반 1차 필터링 및 Reranker 기반 순위 재 정렬)

방법론

Reranker 모델은 현재 CohereAI, Beijing Academy of Artificial Intelligence (BAAI) 에서 제공되고 있습니다. CohereAI의 rerank-multilingual-v2.0 모델은 한국어를 지원하고 있지만 유료로 제공하고 있고, BAAI의 bge-reranker-large 모델은 무료이지만 공식적으로 영어와 중국어만을 지원하고 있습니다. (학습 데이터 셋에 한국어도 소량 포함되어 있습니다.)

1) 한국어 파인튜닝

한국어에 대응하기 위해서는 파인튜닝 과정이 필요합니다. 한국어 파인튜닝을 위한 리소스는 아래와 같습니다.

  • 학습 데이터: msmarco-triplets
    • msmarco-triplets 데이터를 Amazon Translate로 번역하여 활용 (샘플 코드) 하였으며, 데이터 셋은 “query”, “pos”, “neg” 세 가지 타입 구성이 필요합니다.
      • format: {“query”: str, “pos”: List[str], “neg”: List[str]}
      • examples: {“query”: “대한민국의 수도는?”, “pos”: [“미국의 수도는 워싱턴이고, 일본은 도교이며 한국은 서울이다.”], “neg”: [“미국의 수도는 워싱턴이고, 일본은 도교이며 북한은 평양이다.”]}
  • 베이스 모델: BAAI/bge-reranker-large 
  • 파인 튜닝 도구: Amazon SageMaker 기반 pytorch distributed training framework (Sample code)
    • Amazon SageMaker의 Hugging Face Deep Learning Container (DLC)를 활용하면, Hugging Face와의 연계를 손쉽게 할 수 있습니다.

학습은 ml.g5.12xlarge instance 1개가 활용 되었고, 3 epoch 기준, 36시간 소요 되었습니다. 학습된 모델은 Hugging Face Dongjin-kr/ko-reranker 를 통해 접근 가능하며, 파인 튜닝 전체 소스코드는 Git repo를 통해 확인 할 수 있습니다.

2) Reranker 호스팅

학습된 모델은 Amazon S3에 저장이 되며, Amazon SageMaker를 이용하여 배포할 수 있습니다 (Sample code). 만약 추가 재학습 과정 없이 Hugging Face에 등록된 모델을 활용하실 경우, 아래의 코드로 배포 가능합니다.


import boto3
import sagemaker
from sagemaker.huggingface import HuggingFaceModel

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

# Hub Model configuration. https://huggingface.co/models
hub = {
    'HF_MODEL_ID':'Dongjin-kr/ko-reranker',
    'HF_TASK':'text-classification'
}

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    transformers_version='4.28.1',
    pytorch_version='2.0.0',
    py_version='py310',
    env=hub,
    role=role, 
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type='ml.g5.large'
)

runtime_client = boto3.Session().client('sagemaker-runtime')
payload = json.dumps(
    {
        "inputs": [
            {"text": "나는 너를 싫어해", "text_pair": "나는 너를 사랑해"},
            {"text": "나는 너를 좋아해", "text_pair": "너에 대한 나의 감정은 사랑 일 수도 있어"}
        ]
    }
)

response = runtime_client.invoke_endpoint(
    EndpointName="<endpoint-name>",
    ContentType="application/json",
    Accept="application/json",
    Body=payload
)

## deserialization
out = json.loads(response['Body'].read().decode()) ## for json
print (f'Response: {out}')

3) RAG pipeline과 Reranker 통합 하기

3-1) 아키텍쳐

아래 그림은 Reranker가 포함된 RAG 아키텍쳐 입니다. 질문 및 문서에 대한 임베딩은 Amazon Bedrock의 Titan Text Embeddings, 답변을 위한 추론은 Amazon Bedrock의 Claude v2.1, 의미 (semantic) 및 어휘 (lexical) 검색은 Amazon OpenSearch를 활용하며, 전체적인 Task orchestration은 Amazon EC2를 통해 구성합니다.


그림 4: RAG 아키텍쳐

해당 아키텍쳐 기반 RAG 시스템의 작업 흐름은 크게 사전 작업 (흰색)과 검색 작업 (노란색)으로 나눠 집니다. 먼저 사전작업에서는 사용자의 문서에 대한 임베딩 벡터 변환 및 벡터 스토어 저장 작업을 수행합니다. 이 과정은 흰색 원의 1 ~ 3 에 해당합니다. 다음은 검색 과정입니다. 사용자가 의해 발생된 질문은 (1), 임베딩 벡터로 변환 (2)이 되고, 질문의 텍스트와 변환된 임베딩 정보를 기반으로 Amazon OpenSearch를 통해 hybrid (의미 + 어휘) 검색을 수행합니다 (3). Reranker는 검색된 문서들에 대해 질문 과의 관련성을 재 측정 하게 되고 (4), 측정된 관련성 기준으로 컨텍스트 내 문서의 순위를 정렬하여 LLM에게 인풋으로 제공합니다 (5). LLM은 제공받은 정보를 바탕으로 사용자의 질문에 대한 답을 생성하여 전달 하는 것을 끝으로 작업이 종료됩니다 (6). 제시 된 아키텍처가 구현된 샘플 코드는 여기를 통해 확인하실 수 있으며 각 절차 별 샘플 코드는 아래와 같습니다.

3-2) Reranker retriever 구성 하기 (Hybrid search + reranker)

RAG의 작업 흐름 관리를 위해 LLM 어플리케이션 프레임워크인 LangChain을 활용 하였습니다. 특히 LangChain의 RetrievalQA 체인을 이용하면 질의 응답 관련 작업을 손쉽게 구현할 수 있습니다.

from langchain.chains import RetrievalQA

qa = RetrievalQA.from_chain_type(
    llm= "<your llm>",
    chain_type="stuff",
    retriever= "<your retriever>",
    return_source_documents=True,
    chain_type_kwargs={
        "prompt": "<your QA prompt>",
        "verbose": False,
    },
    verbose=False
)

RerievalQA 체인은 사용할 LLM, retriever, 그리고 prompt를 정의 해 주어야 합니다. 각 항목 정의를 위한 샘플 코드는 “Reranker 기반 RAG 시스템 샘플 코드 ” 에서 확인할 수 있습니다. 이 중에서 reranker기반의 retriever는

  1. 질문에 대한 의미 및 어휘 기반 검색 결과를 전달하던 기본적인 Retriever에 Reranking 작업이 추가 되어야 하고
  2. LangChain과의 통합을 위해 LangChain이 정의하고 있는 Retriever의 구성요소를 갖추고 있어야 합니다.

이를 위해 LangChain의 BaseRetriever를 상속받아 구현하였습니다. 샘플 코드에서의 reranking retriever는 “OpenSearchHybridSearchRetriever”로 정의되어 있으며 여기서 “OpenSearchHybridSearchRetriever”로 검색 후 확인 할 수 있습니다. 이를 바탕으로 RerievalQA 체인의 “Retriever”를 아래와 같이 정의할 수 있습니다.

from langchain.chains import RetrievalQA
from utils.rag import OpenSearchHybridSearchRetriever

opensearch_hybrid_retriever = OpenSearchHybridSearchRetriever(
    # necessary
    os_client=os_client,
    vector_db=vector_db,
    index_name=index_name,
    llm_emb=llm_emb,
    llm_text=llm_text,

    # option for lexical
    minimum_should_match=0,
    filter=[],

    # option for rank fusion
    fusion_algorithm="RRF", # ["RRF", "simple_weighted"], rank fusion 방식 정의
    ensemble_weights=[.51, .49], # [for semantic, for lexical], Semantic, Lexical search 결과에 대한 최종 반영 비율 정의
    reranker=True, # enable reranker with reranker model
    reranker_endpoint_name=endpoint_name, # endpoint name for reranking model
    
    # option for async search
    async_mode=True,

    # option for output
    k=5, # 최종 Document 수 정의
    verbose=False,
)

qa = RetrievalQA.from_chain_type(
    llm= "<your llm>",
    chain_type="stuff",
    retriever= opensearch_hybrid_retriever,
    return_source_documents=True,
    chain_type_kwargs={
        "prompt": "<your QA prompt>",
        "verbose": False,
    },
    verbose=False
)

성능 평가

1) 평가 데이터 만들기

 평가를 위해 질문-문서-정답에 대한 데이터가 필요합니다. 이 데이터는 데이터에서 질문과 그에 해당하는 문서의 집합으로 볼 수 있습니다. 만약 이 과정을 인간이 수동으로 한다면 꽤 많은 시간과 비용이 발생할 수 있습니다. 따라서 우리는 이 과정을 LLM을 통해 대체하고자 합니다. 전체 과정은 그림 5와 같습니다.

그림 5: LLM기반 평가 데이터 셋 만들기

 먼저 전체 문서(1)에 대해 LLM으로 질문을 생성(2)합니다. 이때 사용된 프롬프트는 아래와 같습니다.

from langchain.prompts import PromptTemplate

retriever_prompt_template = """
\n\nHuman: Here is the context information, inside <context></context> XML tags.

<context>{context}</context>Given the context information and not prior knowledge.
generate only questions based on the below query.
You are a Professor. Your task is to setup \
{num_questions_per_chunk} questions for an upcoming \quiz/examination.
The questions should be diverse in nature \across the document.
The questions should not contain options, start with "-"
Restrict the questions to the context information provided.
Write in Korean.

\n\nAssistant:"""

PROMPT_RETRIEVER = PromptTemplate(
    template=retriever_prompt_template,
    input_variables=["context", "num_questions_per_chunk"]
)

그 다음 주어진 문서와 생성된 질문으로 실제 정답을 생성(3)합니다. LLM 모델 평가의 편향을 제거하기 위해, 우리는 평가에 활용하지 않는 다른 LLM (AI21Labs – Jurrasic)을 사용하였습니다. 이때 사용한 프롬프트는 아래와 같습니다.

generation_prompt_template = """
Here is the context, inside <context></context> XML tags.
<context>
{context}
</context>
Only using the context as above, answer the following question with the rules as below:
    - Don't insert XML tag such as <context> and </context> when answering.
    - Write as much as you can
    - Be courteous and polite
    - Only answer the question if you can find the answer in the context with certainty.
    - Skip the preamble
    - Use three sentences maximum and keep the answer concise.
    - If the answer is not in the context, just say "Could not find answer in given contexts."
    - Answer in Korean.
Question:
{question}
Answer:"""

PROMPT_GENERATION = PromptTemplate(
    template=generation_prompt_template,
    input_variables=["context", "question"]
)

질문 및 답변 생성을 위한 프롬프트는 여기를 참고하였으며 전체 코드는 여기에서 확인 가능합니다. 최종적으로 135개의 질문-문서-정답 데이터 셋을 생성하였습니다.

2) 결과

다음으로 Reranker의 효과에 대해 확인해 보도록 하겠습니다. “해결전략” 섹션을 통해 Reranker는 질문과 문서를 동시에 분석함으로써, 질문에 대한 문서의 관련성을 더욱 정확히 판단할 수 있음을 알 수 있었습니다. 이를 확인하기 위해 Hit Rate 와 Mean Reciprocal Rank (MRR) 을 성능 지표로 활용 하였습니다. Hit Rate는 컨텍스트 내 정답 존재 여부를 나타내고, MRR은 컨텍스트 내 문서 중에서 질문과 관련성이 높은 문서의 순위가 높을 수록 (1 위에 가까울 수록) 더욱 높은 값을 가지기 때문에 Reranker의 효과 측정에 적합하다고 판단 하였습니다.
또한 LLM의 답변 성능을 측정하기 위해서 “Reranker 유사도”를 활용 하였습니다. Reranker는 입력으로 사용 되는 질문과 문서가 관련성이 높으면 높은 값을 반환합니다. 만약 우리가 Reranker의 입력을 질문에 대한 실제 정답과, LLM이 제공한 정답을 동시 넣는다고 가정해 보겠습니다. LLM이 제공한 정답이 실제 정답과 유사하다면 높은 점수를, 그렇지 않다면 낮은 점수가 반환될 것 입니다. 이러한 특성을 활용하면 LLM 답변의 성능을 확인 할 수 있습니다. 하지만 이 수치는 간접적인 수치라는 것에 주의해야 합니다. 정확한 수치가 필요하다면 휴먼 평가가 반드시 필요합니다.

테이블 1: Reranker 성능

 위의 테이블은 Reranker의 존재 여부에 따른 효과를 Retriever, Generator 관점에서 각각 나태내고 있습니다. “w/o reranker”를 베이스라인으로 보았을 때, Reranker가 존재 할 때 모든 수치가 향상 됨을 알 수 있습니다. 먼저 Retriever의 Hit Rate 및 MRR이 상승했다는 뜻은 Reranker를 활용할 때 질문에 대한 정답이 들어 있는 문서가 컨텍스트 내에서 더욱 상위권에 위치한다는 뜻입니다. 그리고 이 것은 Reranker similarity의 상승 즉, LLM의 답변 성능 향상으로 이어진 다는 것을 알 수 있습니다.

결론

이 글은 Reranker를 사용하여 RAG의 답변 성능을 향상 시키는 방법에 대해서 다루고 있습니다. 먼저 우리는 Reranker가 질문과 문서의 동시 분석함으로써 기존의 임베딩 유사도 대비 질문에 대한 관련 문서를 더욱 정확히 찾을 수 있다는 것을 확인하였습니다. 또한 이러한 점은 양질의 컨텍스트 구성을 가능하게 함으로써 LLM 답변 정확도 향상을 이끌어 낼 수 있음을 보였습니다.
즉, Reranker RAG 파이프라인에서 Retrieval 성능을 개선할 수 있는 방법 중 하나이며, 프로덕션 수준의 높은 정확도가 필요하시다면 꼭 한번 적용해 보시길 바랍니다.

함께 읽으면 도움이 되는 참조 블로그

아래의 링크에서 실습 소스 파일 및 기계 학습(ML)과 관련된 자료를 확인하실 수 있습니다.

Reference

  • [1] Lost in the Middle: How Language Models Use Long Contexts, F.Liu et al., 2023
  • [2] Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks, Reimers et al., 2019
Dongjin Jang, Ph.D.

Dongjin Jang, Ph.D.

장동진 AIML 스페셜리스트 솔루션즈 아키텍트는 데이터 사이언티스트 경험을 바탕으로 고객의 머신러닝 기반 워크로드를 도와드리고 있습니다. 추천 시스템, 이상탐지 및 수요 예측과 같은 다양한 분야에 대한 고민을 고객과 함께 해결 하였고, 최근에는 생성형 AI을 통해 고객의 혁신을 지원하고 있습니다.