REX-RAG: 검색 증강 생성(RAG)을 위한 강화학습 프레임워크

Wentao Jiang¹ *, Xiang Feng¹ *, Zengmao Wang¹ †, Yong Luo¹, Pingbo Xu²,³ , Zhe Chen⁴, Bo Du¹, Jing Zhang¹ †

1 Wuhan University 컴퓨터 과학 학부, 중국,
2 Zhejiang Cancer Hospital 마취과, 중국,
3 Chinese Academy of Sciences 의학 연구소, Hangzhou, Zhejiang, China
4 La Trobe University 컴퓨터 과학 및 정보 기술 학부, 호주

¹ *, Xiang Feng
¹ *, Zengmao Wang
¹ †, Yong Luo
¹ , Pingbo Xu
² ,³ , Zhe Chen
⁴ , Bo Du
¹ , Jing Zhang
¹ †

2025.08.11

• Github에 코드를 공개했습니다!!!

REX-RAG는 혼합 샘플링 전략을 통해 추론 막다른 길(reasoning dead ends)에서 벗어나고, 원칙적인 정책 교정 메커니즘(principled policy correction mechanism)을 통해 안정적인 정책 학습을 유지하는 검색 증강 생성(Retrieval-Augmented Generation)을 위한 강화학습(RL) 프레임워크입니다. 이는 다단계 추론(multi-hop reasoning) 및 일반 QA 작업에서 상당한 성능 향상을 제공하며, 강력한 도메인 외 일반화(out-of-domain generalization)와 다양한 RL 학습 알고리즘과의 호환성을 갖추고 있습니다.

그림 1: REX-RAG 개요.

참고: bash 스크립트 내의 경로가 로컬 환경과 일치하도록 설정했는지 확인하십시오.

먼저, 필요한 의존성(dependencies)을 설치합니다. 더 빠른 설치를 위해 uv 사용을 권장합니다.

# pip 업그레이드 및 uv 설치
pip install --upgrade pip
pip install uv
...

sglang 설치에 대한 자세한 내용은 공식 문서를 참조하십시오.

리트리버(retriever)는 Wikipedia 코퍼스(corpus)가 필요합니다. 이를 수동으로 처리하거나 사전 처리된 버전을 다운로드할 수 있습니다.

옵션 A: Wikipedia를 수동으로 처리. FlashRAG Wiki Processing의 지침을 따르십시오. -
옵션 B: 사전 처리된 데이터 다운로드. Hugging Face Datasets에서 데이터를 다운로드하십시오.

데이터를 얻은 후, 검색 인덱스를 구축합니다:

bash scripts/search_engine/build_index.sh

필요한 데이터셋은 git lfs로 가져올 수 있습니다.

git lfs pull

대안으로 자체 사용자 정의 데이터셋을 사용할 수도 있습니다. Search-R1에 설명된 전처리 방법을 참조하십시오.

먼저, 리트리버 서버를 시작합니다:

bash scripts/search_engine/retrieval_server.sh

그런 다음, 메인 애플리케이션을 실행할 수 있습니다.

bash scripts/search-r1-sgl/run_grpo_sglang_fsdp.sh

그림 2: REX-RAG를 사용한 7개의 QA 벤치마크에 대한 주요 실험 결과.

그림 3: Qwen2.5-7B와 REX-RAG가 적용된 Qwen2.5-7B를 비교하는 불확실성 정량화 시각화.

그림 3은 원본 Qwen2.5-7B 모델의 추론 궤적과 REX-RAG로 향상된 동일한 모델을 비교하는 시각화 분석을 제시합니다. 이 분석에서는 LogTokU (GitHub)에서 제공하는 불확실성 정량화 방법을 사용했습니다.

그들의 프레임워크를 따라, 우리는 두 가지 유형의 불확실성을 분석합니다:

조학적 불확실성 (Aleatoric Uncertainty, AU): 본질적인 데이터 무작위성을 나타냅니다.
인식론적 불확실성 (Epistemic Uncertainty, EU): 모델 지식의 격차를 포착합니다.

이러한 값들은 토큰 수준의 신뢰도 점수를 통해 측정됩니다. 시각화 결과는 REX-RAG가 추론 토큰에 대해 훨씬 더 높은 신뢰도 점수(일반적으로 0.60.8 범위)를 달성하는 반면, 기준 모델은 낮은 신뢰도(일반적으로 0.20.4 범위)를 보인다는 것을 보여줍니다.

저희의 작업에 중요한 역할을 한 다음 오픈 소스 프로젝트들에 감사를 표하고 싶습니다:

불확실성 시각화 분야에서 뛰어난 작업을 수행한 LogTokU에 특별히 감사드립니다. 저희는 이 기능을 분석에 맞게 적용했습니다.

만약 저희의 작업이 유용하다고 생각되시면, ⭐을 주시고 다음 논문을 인용해 주시기를 부탁드립니다:

@article{jiang2025rex,
title={REX-RAG: Reasoning Exploration with Policy Correction in Retrieval-Augmented Generation},
author={Jiang, Wentao and Feng, Xiang and Wang, Zengmao and Luo, Yong and Xu, Pingbo and Chen, Zhe and Du, Bo and Zhang, Jing},
...```