RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs

Large language models (LLMs) typically utilize the top-k contexts from a retriever in retrieval-augmented generation (RAG). In this work, we propose a novel instruction fine-tuning framework RankRAG, which instruction-tunes a single LLM for the dual purpose of context ranking and answer generation in RAG. In particular, the instruction-tuned LLMs work surprisingly well by adding a small fraction of ranking data into the training blend, and outperform existing expert ranking models, including the same LLM exclusively fine-tuned on a large amount of ranking data. For generation, we compare our model with many strong baselines, including GPT-4-0613, GPT-4-turbo-2024-0409, and ChatQA-1.5, an open-sourced model with the state-of-the-art performance on RAG benchmarks. Specifically, our Llama3-RankRAG significantly outperforms Llama3-ChatQA-1.5 and GPT-4 models on nine knowledge-intensive benchmarks. In addition, it also performs comparably to GPT-4 on five RAG benchmarks in the biomedical domain without instruction fine-tuning on biomedical data, demonstrating its superb capability for generalization to new domains.

한 줄 요약

RankRAG는 단일 LLM을 사용해 컨텍스트 랭킹과 답변 생성을 동시에 학습하여 기존 RAG 파이프라인의 성능을 크게 개선한다.

핵심 기여도

RankRAG는 컨텍스트 랭킹과 답변 생성을 동시에 학습하는 instruction fine-tuning 프레임워크로, 기존 전문 랭킹 모델을 초과함.
Llama3-RankRAG는 9개 일반 도메인 및 5개 의료 도메인 RAG 벤치마크에서 ChatQA-1.5 및 GPT-4를 상회함.
랭킹 데이터가 훈련에 포함된 비율이 작아도, 10배 더 많은 랭킹 데이터로 훈련된 동일 LLM보다 성능이 우수함.
의료 도메인에서 별도의 도메인 훈련 없이도 GPT-4의 98% 수준 성능을 달성함.

핵심 아이디어

기존 RAG 파이프라인은 별도의 레트리버와 랭커를 사용하지만, RankRAG는 단일 LLM을 사용해 랭킹과 생성을 동시에 학습하는 새로운 접근법을 제안한다. 이는 LLM이 질문과 컨텍스트 간 관련성을 판단하는 능력과 답변 생성 능력이 상호 보완적이라는 통찰에서 비롯된다. RankRAG는 훈련 시 랭킹 데이터를 포함한 QA 및 RAG 데이터셋을 사용하여 LLM이 컨텍스트 필터링과 답변 생성을 동시에 학습하도록 유도한다. 특히, RQA(Retrieval-Augmented QA)와 RAR(Retrieval-Augmented Ranking) 데이터를 추가하여 모델이 관련 컨텍스트를 명시적으로 식별하도록 돕는다.

기술적 접근법

**모델**: Llama3-8B 및 Llama3-70B를 사용.
**훈련 데이터**: 기존 instruction tuning 데이터에 컨텍스트-부유한 QA, RAG, 랭킹 데이터를 추가.
**랭킹 데이터 비율**: 전체 훈련 데이터 중 랭킹 데이터는 소수 비율(구체적 비율은 명시되지 않음)이지만, 성능에 큰 영향을 미침.
**추론 단계**: LLM이 먼저 레트리버에서 가져온 상위 N개 컨텍스트를 재랭킹한 후, 상위 k개(예: 5개)를 기반으로 답변 생성.
**비교 모델**: ChatQA-1.5, GPT-4-0613, GPT-4-turbo-2024-0409 등.

주요 결과

**일반 도메인**: Llama3-RankRAG-8B는 9개 RAG 벤치마크에서 Llama3-ChatQA-1.5-8B를 평균 7.8% 상회함.
**의료 도메인**: Llama3-RankRAG-8B는 Meditron-70B를 6.3% 상회하며, Llama3-RankRAG-70B는 GPT-4의 98% 성능 달성.
**레트리버 독립성**: DPR 및 Contriever-MS MARCO 레트리버 모두에서 RankRAG가 ChatQA-1.5보다 평균 10% 이상 우수함.
**모델 크기 독립성**: 7B, 13B, 70B 버전 모두에서 6.3~7.8%의 성능 향상 관찰됨.

의의 및 한계

RankRAG는 단일 LLM을 사용해 랭킹과 생성을 통합함으로써 RAG 파이프라인의 복잡도를 줄이고, 뛰어난 일반화 능력을 보인다. 특히, 랭킹 데이터가 적어도 성능이 우수한 점은 훈련 데이터 효율성 향상에 기여한다. 그러나, 특정 도메인(예: 의료)에서 별도의 도메인 훈련 없이도 높은 성능을 내는 것은 훈련 데이터의 범용성에 의존하며, 이는 모든 도메인에 일반화될 수 있는지에 대한 추가 연구가 필요하다. 또한, RankRAG는 레트리버와 별도로 사용되며, 레트리버의 품질에 따라 최종 성능이 영향을 받을 수 있다.

실용적 활용

RankRAG는 단일 모델로 랭킹과 생성을 처리하므로, RAG 기반의 고객 지원 시스템, 의료 QA 플랫폼, 법률 검색 시스템 등 다양한 도메인에서 즉시 적용 가능하다. 특히, 별도의 랭킹 모델 없이도 높은 성능을 내므로, 모델 배포 및 유지 관리 비용을 절감할 수 있다.