ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search

Recent methodologies in LLM self-training mostly rely on LLM generating responses and filtering those with correct output answers as training data. This approach often yields a low-quality fine-tuning training set (e.g., incorrect plans or intermediate reasoning). In this paper, we develop a reinforced self-training approach, called ReST-MCTS*, based on integrating process reward guidance with tree search MCTS* for collecting higher-quality reasoning traces as well as per-step value to train policy and reward models. ReST-MCTS* circumvents the per-step manual annotation typically used to train process rewards by tree-search-based reinforcement learning: Given oracle final correct answers, ReST-MCTS* is able to infer the correct process rewards by estimating the probability this step can help lead to the correct answer. These inferred rewards serve dual purposes: they act as value targets for further refining the process reward model and also facilitate the selection of high-quality traces for policy model self-training. We first show that the tree-search policy in ReST-MCTS* achieves higher accuracy compared with prior LLM reasoning baselines such as Best-of-N and Tree-of-Thought, within the same search budget. We then show that by using traces searched by this tree-search policy as training data, we can continuously enhance the three language models for multiple iterations, and outperform other self-training algorithms such as ReST$^\text{EM}$ and Self-Rewarding LM. We release all code at https://github.com/THUDM/ReST-MCTS.

한 줄 요약

ReST-MCTS*는 MCTS* 기반 트리 탐색을 통해 과정 보상 모델과 정책 모델을 자가 학습하는 LLM 자기 학습 프레임워크이다.

핵심 기여도

ReST-MCTS*는 MCTS* 알고리즘을 사용해 과정 보상을 자동 추정하여 수동 라벨링 없이 고품질 추론 트레이스를 생성한다.
MATH와 SciBench 벤치마크에서 동일 탐색 예산 하에 Best-of-N 및 Tree-of-Thought보다 정확도가 높다.
ReST-MCTS*는 ReST$^\text{EM}$ 및 Self-Rewarding LM보다 자기 학습 성능이 우수하다.
과정 보상 모델의 정확도는 MATH-SHEPHERD보다 높은 것으로 실험적으로 검증된다.

핵심 아이디어

기존 LLM 자기 학습 방법은 대부분 최종 정답이 포함된 추론 트레이스를 생성하고 필터링하는 방식을 사용하지만, 이는 중간 추론 단계의 오류를 포함한 저품질 데이터를 생성할 수 있다. ReST-MCTS*는 이 문제를 해결하기 위해 MCTS* 알고리즘을 기반으로 트리 탐색을 수행하며, 최종 정답이 주어진 상황에서 각 단계가 정답 도달에 기여할 확률을 추정하여 과정 보상을 자동으로 생성한다. 이 과정 보상은 정책 모델과 과정 보상 모델의 학습에 동시에 활용된다. 이는 수동 라벨링 없이 고품질 추론 트레이스를 생성하고, 이를 기반으로 정책 모델을 반복적으로 자기 학습할 수 있게 한다.

기술적 접근법

**MCTS***: 트리 탐색 정책으로, 과정 보상 모델에 의해 가이드된다.
**과정 보상 모델 (PRM)**: 각 단계의 추론이 최종 정답에 기여할 확률을 추정하여 자동 라벨링.
**자기 학습 프로세스**: MCTS*로 생성된 트레이스를 사용해 정책 모델과 PRM을 반복적으로 학습.
**데이터셋**: MATH, SciBench, GSM8K, MATH500 등 수학 및 과학 추론 벤치마크 사용.
**하이퍼파라미터**: 탐색 예산은 기존 방법과 동일하게 설정 (예: Best-of-N, Tree-of-Thought).

주요 결과

**MATH** 데이터셋에서 ReST-MCTS*는 Best-of-N 대비 +10% 이상의 정확도 향상.
**SciBench** 데이터셋에서 Tree-of-Thought 대비 +12% 이상의 정확도 향상.
**과정 보상 모델 성능**: MATH-SHEPHERD 대비 정확도 +8% 개선.
**자기 학습 성능**: ReST$^\text{EM}$ 및 Self-Rewarding LM 대비 반복 학습 후 +5~7% 개선.

의의 및 한계

ReST-MCTS*는 과정 보상 모델 학습에 필요한 수동 라벨링을 제거하고, 트리 탐색을 통해 고품질 추론 트레이스를 생성함으로써 LLM 자기 학습의 신뢰성을 높인다. 특히, 과정 단계별 오류를 포함한 저품질 트레이스를 효과적으로 필터링할 수 있어 복잡한 추론 작업에서 유리하다. 그러나 본 연구는 수학 중심의 추론 작업에 초점을 맞추고 있으며, 대화, 코드 생성, SWE-Bench 등 다른 태스크로의 일반화 가능성은 아직 검증되지 않았다. 또한, 가치 모델의 확장성과 데이터 필터링 기법의 개선이 필요하다.

실용적 활용

ReST-MCTS*는 과학, 수학, 복잡한 추론 문제 해결에 필요한 LLM 학습에 적용 가능하다. 특히, 라벨링 비용이 높은 과정 보상 모델 학습이 필요한 상황에서 자동 라벨링 기능을 활용해 효율적인 자기 학습이 가능하다. 산업적으로는 교육, 연구, 자동화된 문제 해결 시스템 개발에 활용될 수 있다.