RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs
이 글에서는 “RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs” 논문을 간단히 정리한다. 해당 논문은 기존 RAG(Retrieval-Augmented Generation)에 별도 랭킹 모델을 사용하지 않고, 하나의 LLM만으로 질문과 문서 간의 적합도를 판단해 상위 문서를 선별(reranking)하고 답변까지 생성하는 새로운 방법을 제안한다.
1. 배경과 문제 설정
대형 언어 모델(LLM)은 거대한 파라미터로 다양한 질의에 답변할 수 있지만, 모든 지식을 파라미터에 내재화하기는 현실적으로 어렵다. 시시각각 업데이트되는 정보나 희소(Long-tail)한 지식을 다루려면, RAG가 주로 활용된다. RAG는 크게 아래 과정을 따른다.
- Retrieve: 외부 검색(예: BM25, DPR, Dragon 등)을 통해 질문에 대한 잠재적으로 관련 있는 문서를 상위 k개 찾는다.
- Generate: 검색된 문서를 LLM에 입력하여 최종 답변을 생성한다.
문제는 retriever가 완벽하지 않기에, k가 너무 작으면 중요한 정보를 놓칠 수 있고, k가 커지면 노이즈가 많아져 답변이 흐려질 수 있다는 점이다. 따라서, 기존에는 랭킹 모델을 따로 두어 검색된 Top-N 문서 중 진짜로 중요한 k개만 재선정(rerank)하는 방식이 사용되곤 했다. 그러나 별도 모델 관리가 복잡하며, 새로운 도메인에 대한 일반화가 부족할 수 있다는 단점이 존재한다.
2. RankRAG란 무엇인가
RankRAG는 “LLM 하나”로 랭킹과 생성 능력을 동시에 수행하도록 학습시킨 프레임워크다. 별도의 랭킹 전용 모델 없이, LLM에 ‘질문-문단 적합도’를 학습하게 하여, 유의미한 문맥만 최종 답변 생성에 사용하도록 한다. 이렇게 하면 파이프라인이 단순해지면서도, 기존 랭킹 모델이 제공하던 문서 선별 기능을 LLM 내부로 통합할 수 있다.
3. 학습 구조
RankRAG는 크게 두 단계의 Finetuning을 거친다.
3.1 Stage-I: Supervised Finetuning
먼저, 일반적인 대화/QA/CoT 등의 데이터로 LLM을 학습한다. 이렇게 하면 전반적인 질문 응답, Instruction-Following 능력이 향상된다.
3.2 Stage-II: 랭킹 + RAG 통합 파인튜닝
이후, 두 번째 단계에서 랭킹 태스크를 포함한 다양한 데이터셋을 합쳐서 재학습한다. 예를 들어:
- Context-rich QA: 단일 문단과 질문/답변이 있는 전통 QA (SQuAD, DROP 등)
- Retrieval-augmented QA: 실제 검색 결과(정답 문단 + Hard-Negative 문단)를 여러 개 포함해 답변을 생성하도록 요구
- Context Ranking: “이 문단이 해당 질문에 얼마나 관련이 있는지”를 True/False로 레이블링
- Retrieval-augmented Ranking: 여러 문단 중 정답과 관련된 문단만 선택하도록 지시
이처럼 QA와 랭킹을 함께 학습하면, 모델이 ‘중요 문맥을 선별하는 능력’과 ‘문맥에 기반해 답변을 만드는 능력’을 동시에 습득하게 된다.
3.3 추론 시 동작 방식
Stage-II까지 학습이 끝난 RankRAG는 추론 시 아래와 같은 과정을 밟는다.
- Retrieve: 우선, retriever로 Top-N 문서를 가져온다.
- Re-rank: RankRAG가 각 문서의 “질문과의 관련도”를 스스로 추정해 점수를 매기고, 상위 k개만 최종 선택한다.
- Generate: 선택된 k개 문서를 LLM 입력으로 사용해, 최종 답변을 생성한다.
4. 주요 기여(Main Contribution)
- 하나의 LLM으로 랭킹 + 생성 통합: 별도 랭킹 모델을 유지할 필요 없이, 통합된 파이프라인으로 간단하고 효율적인 RAG를 구현한다.
- 소량의 랭킹 데이터만으로도 높은 성능: MS MARCO 같은 대규모 랭킹 셋이 없어도, 적은 데이터로 전통 랭커보다 나은 정확도와 Recall을 달성한다.
- 통일된 인스트럭션 포맷으로 데이터 집약: (질문 x, 문맥 c, 답 y) 구조로 QA/랭킹을 동시에 학습해, 상호 학습 시너지가 생긴다.
- 여러 도메인과 벤치마크에서 우수한 결과: 일반 위키, Multi-Hop QA, 대화형 QA, 바이오메디컬(의료) QA 등 다양한 작업에서 GPT-4 수준에 근접하거나 능가하는 성능이 보고된다.
5. 한계(Weakness)
- 추론 지연: 모델 크기가 클수록 문서별 랭킹 계산 시간이 증가한다. Top-N이 늘어날수록 추가 비용이 커진다.
- Retriever 품질 의존: 초기 검색이 부실하면, 최종적으로 선택할 문서가 다 엉뚱할 수 있다.
- 큰 모델을 이용한 랭킹: 랭킹조차 대형 LLM이 담당하므로, 모델을 효율적으로 운영하기 위한 인프라와 비용이 요구된다.
6. 장점(Strength)
- 파이프라인 단순화: “LLM + 전용 랭커”를 각각 관리하지 않고 하나의 모델로 해결한다.
- 도메인 제로샷 성능: 의료/바이오처럼 전문 분야에서도, 추가 학습 없이도 준수한 결과를 낸다.
- Multi-Hop, 대화형 QA에 유리: 랭킹 단계에서 노이즈를 효과적으로 제거해 답변 정확도를 높인다.
7. 간단한 예시
“1930년에 개최된 월드컵 우승국은 어디인가?”라는 질문을 생각해보자(답: 우루과이).
- retriever가 먼저 Top-30 문서를 가져왔으나, 중요한 문서는 일부에 불과하다.
- RankRAG가 문서별 “True/False(질문과의 관련성)” 확률을 계산해 상위 k=5만 남긴다.
- 그 안에 우루과이에 관한 정확한 근거가 담겨 있으면, 최종 답변이 제대로 나온다.
8. 결론
RankRAG는 “랭킹 모델을 따로 두지 않고, 적은 랭킹 데이터만으로 하나의 LLM이 Relevance 파악 + 최종 생성까지 해내는” RAG 통합 기법이다. 효율이나 구현 난이도 측면에서 꽤 유용한 아이디어로 볼 수 있다. 추론 비용은 다소 늘어나지만, 도메인 적응력이나 Multi-Hop QA 성능 면에서 장점이 크다.
Reference:
Yue Yu, Wei Ping, Zihan Liu, Boxin Wang, Jiaxuan You, Chao Zhang, Mohammad Shoeybi, Bryan Catanzaro.
“RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs.”
arXiv:2407.02485
