Monte Carlo Tree Search를 통한 검색 증강 생성 (RAG) 강화

Monte Carlo Tree Search를 통한 검색 증강 생성 (RAG) 강화

검색 증강 생성 (RAG)을 Monte Carlo Tree Search (MCTS)와 통합함으로써 지식 집약적 작업에서 소형 언어 모델 (SLMs)의 추론 능력을 향상시키는 새로운 접근 방식입니다.

우리는 검색 증강 생성 (RAG)을 Monte Carlo Tree Search (MCTS)와 통합하여 지식 집약적 작업에서 소형 언어 모델 (SLMs)의 추론 능력을 향상시키는 새로운 접근 방식인 MCTS-RAG를 소개합니다.

검색과 생성을 순차적으로 수행하는 기존의 RAG 방식과 달리, MCTS-RAG는 반복적인 의사 결정 과정을 통해 검색과 추론을 모두 동적으로 개선합니다. 모델은 먼저 MCTS를 활용하여 여러 추론 경로를 탐색하며, 사실적 정확도를 높이기 위해 전략적으로 검색을 통합합니다. 검색된 정보는 이후 추론 과정 내에서 평가되어, 높은 신뢰도를 가진 지식만이 최종 답변에 기여하도록 보장합니다. 이러한 구조적 추론과 적응형 검색의 통합은 환각 (hallucinations)을 줄이고 응답 일관성을 향상시킵니다. ComplexWebQA, GPQA, 그리고 FoolMeTwice에 대한 실험 결과는 MCTS-RAG가 7B 규모의 모델로도 GPT-4o와 같은 프런티어 LLM에 필적하는 성능을 달성할 수 있음을 보여주며, 소규모 모델의 추론에 대한 새로운 표준을 제시합니다.

• Python 3.10
• CUDA 12
• 최신 PyTorch
• 최신 transformers
• 최신 vllm

다음은 MCTS-RAG 생성기를 실행하는 예시입니다:

bash scripts/run_gsm8k_generator.sh

run_gsm8k_generator.sh 스크립트에는 다음과 같은 여러 설정 가능한 파라미터가 포함되어 있습니다:

--dataset_name : 데이터셋 이름.
--test_json_filename : 테스트 JSON 파일명 (기본값: test_all).
--model_ckpt : 모델 체크포인트 경로.
--note : 출력 폴더를 표시하기 위한 추가 메모. 추가 인자가 없으면 생성기 출력 폴더는 ./run_outputs/<dataset_name>/<model_ckpt>/<execute_time>---[<note>]가 됩니다.
--num_rollouts : 롤아웃 (rollouts) 횟수 (기본값: 16).

이러한 파라미터들을 귀하의 요구 사항에 맞춰 조정하십시오.

다음은 rStar 판별기 (discriminator)를 실행하는 예시입니다:

bash scripts/run_gsm8k_discriminator.sh

run_gsm8k_discriminator.sh 스크립트에는 다음과 같은 여러 설정 가능한 파라미터가 포함되어 있습니다:

--root_dir : 평가 결과 폴더의 경로.
--dataset_name : 데이터셋 이름.
--note : 출력 폴더를 표시하기 위한 추가 메모. 추가 인자가 없으면, 판별기 출력 폴더는 <root_dir>/dis_<execute_time>---<note>가 됩니다.

실험 결과, 우리의 접근 방식은 LLaMA 3.1-8B를 사용하여 CWQA에서 20% 이상, GPQA에서 15% 이상, FMT에서 10% 이상의 성능 향상을 달성하며 CWQA, GPQA, FMT에서 베이스라인 (baselines)을 능가함을 보여주었습니다. 검색 (retrieval) 및 검증 (verification)을 정교화함으로써, 이는 사실적 정확도를 높이고 환각 (hallucinations)을 최소화하며 복잡한 추론 (reasoning) 작업에 대한 새로운 벤치마크를 설정합니다.

저희의 작업이 도움이 되었다면, 다음을 인용해 주시기 바랍니다:

@misc{hu2025mctsragenhancingretrievalaugmentedgeneration,
title={MCTS-RAG: Enhancing Retrieval-Augmented Generation with Monte Carlo Tree Search},
author={Yunhai Hu and Yilun Zhao and Chen Zhao and Arman Cohan},
...