[논문 리뷰] Chain-of-Retrieval Augmented Generation

Liang Wang, Haonan Chen외 4명|ArXiv.org|2025. 01. 24.

Speech and dialogue systems인용 수 3

한 줄 요약

CoRAG은 정보를 체인으로 순차적으로 검색하고 추론하도록 학습시키며, 다중-힙 QA 및 지식 집약적 작업의 성능을 향상시키고, 계산량을 제어하기 위한 테스트 시 전략을 제공합니다.

ABSTRACT

This paper introduces an approach for training o1-like RAG models that retrieve and reason over relevant information step by step before generating the final answer. Conventional RAG methods usually perform a single retrieval step before the generation process, which limits their effectiveness in addressing complex queries due to imperfect retrieval results. In contrast, our proposed method, CoRAG (Chain-of-Retrieval Augmented Generation), allows the model to dynamically reformulate the query based on the evolving state. To train CoRAG effectively, we utilize rejection sampling to automatically generate intermediate retrieval chains, thereby augmenting existing RAG datasets that only provide the correct final answer. At test time, we propose various decoding strategies to scale the model's test-time compute by controlling the length and number of sampled retrieval chains. Experimental results across multiple benchmarks validate the efficacy of CoRAG, particularly in multi-hop question answering tasks, where we observe more than 10 points improvement in EM score compared to strong baselines. On the KILT benchmark, CoRAG establishes a new state-of-the-art performance across a diverse range of knowledge-intensive tasks. Furthermore, we offer comprehensive analyses to understand the scaling behavior of CoRAG, laying the groundwork for future research aimed at developing factual and grounded foundation models.

연구 동기 및 목표

• 단일 검색 단계가 아니라 반복적 검색과 추론을 가능하게 하여 RAG를 개선하려는 동기를 부여한다.
• 반복추출(rejection sampling)을 사용하여 중간 검색 체인을 포함한 QA 데이터세트를 보강한다.
• 검색-생성 체인에서 다음 동작을 예측하도록 LLM을 학습한다.
• 체인 기반 검색을 위한 계산량 확장을 위한 테스트 타임 디코딩 전략을 연구한다.
• 일반화성과 확장성을 평가하기 위해 CoRAG를 multi-hop QA 데이터셋과 KILT 벤치마크에서 평가한다.

제안 방법

• 서브쿼리와 서브-답변의 시퀀스로 QA 데이터셋을 보강하기 위해 rejection sampling으로 검색 체인을 생성한다.
• 서브쿼리, 서브-답변, 최종 답변 예측을 포함하는 다중 작업 목표로 오픈 소스 LLM을 미세 조정한다.
• 각 서브쿼리에 대해 상위-k 문서를 얻기 위해 리트리버를 사용하고 체인의 로그가능도로 체인 품질을 평가한다.
• 토큰 소비를 제어하기 위한 탐욕적, best-of-N 샘플링, 트리 탐색을 포함한 테스트-타임 디코딩 전략을 제공한다.
• 데이터세트, 리트리버, 일반화 시나리오 전반에 걸친 확장성 동작 및 강건성을 분석한다.
• 선택적으로 현재 정보가 충분한지 예측하여 테스트 시 체인을 중단하는 것을 학습한다.

(a) Test-time scaling behavior of CoRAG.

실험 결과

연구 질문

• RQ1 단일 단계 RAG와 비교했을 때 반복적 검색 및 추론이 복잡하고 다중-홉 QA 작업의 성능을 향상시키는가?
• RQ2 테스트 시 계산량(체인 길이와 체인 수)이 성능과 효율성에 어떤 영향을 미치는가?
• RQ3 검색 연쇄가 QA를 넘어 다양한 지식 집중적 작업에 걸쳐 일반화되는가?
• RQ4 약한 리트리버와 강한 리트리버 및 모델 사용이 CoRAG의 효과에 미치는 영향은 무엇인가?
• RQ5 추론 시기에 검색 체인을 조기에 중지하는 메커니즘을 학습하는 것이 유익한가?

주요 결과

• CoRAG는 다중-홉 QA 데이터세트에서 강력한 기준선을 크게 능가하며, 여러 디코딩 전략에서 EM/F1의 주목할 만한 향상을 보인다.
• KILT 벤치마크에서 CoRAG는 다양한 작업에서 최첨단 성능을 달성하나, 가장 큰 FEVER 작업에서 예외가 있다.
• 짧은 체인에서 시작할 때 더 긴 검색 체인은 성능을 향상시키지만, 체인 길이가 늘어나면 이득이 감소한다.
• 테스트 시 확장성은 여러 데이터세트에서 토큰 소비와 성능 간의 로그-선형 관계를 따른다.
• 강한 리트리버를 사용할 때의 이점과 약한 리트리버에서도 일부 이득이 나타나는 등 강건성 실험은 일반화가 작업 유형 전반에 걸쳐 우호적임을 보여준다.
• 특성 제거 실험은 반복적 학습이 혼합된 결과를 보이며, 지시-미세조정된 LLM이 종종 고품질의 검색 체인을 생성한다는 것을 시사한다.

(b) An example of CoRAG in action.