[논문 리뷰] ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search

Dan Zhang, Sining Zhoubian외 4명|arXiv (Cornell University)|2024. 06. 06.

Business Process Modeling and Analysis인용 수 6

한 줄 요약

ReST-MCTS*를 소개하는 자가학습 프레임워크로, MCTS*-guided 프로세스 보상 모델이 단계별 추론을 자동으로 라벨링하고, LLM 추론 작업에서 정책 모델과 보상 모델의 상호 개선을 가능하게 한다.

ABSTRACT

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$^ ext{EM}$ and Self-Rewarding LM. We release all code at https://github.com/THUDM/ReST-MCTS.

연구 동기 및 목표

• 밀집한 인간 주석 없이 고품질의 단계별 추론 보상을 자동으로 획득하는 것.
• MCTS*-guided 탐색을 사용하여 중간 추론 흔적을 생성하고 평가한다.
• 반복적 자가학습을 통해 정책 모델과 프로세스 보상 모델을 상호적으로 정제한다.

제안 방법

• 부분 해에 대해 단계별 품질 v_k를 추론하는 프로세스 보상 모델 V_theta를 정의한다.
• 트리 탐색과 부분 백업을 안내하기 위해 v_k를 가치 타깃으로 사용하는 MCTS*를 개발한다.
• MCTS* 안내 탐색으로 정보를 얻은 추론 흔적을 생성하도록 정책 모델 pi를 학습시킨다.
• 정답에 가까운 흔적을 사용하여 V_theta와 pi를 업데이트하기 위해 MuZero-스타일의 상호 자가학습을 반복적으로 수행한다.
• 검색 트리 내의 롤아웃을 활용하여 명시적 단계별 라벨 없이도 단계별 보상을 추론한다.
• 동일한 탐색 예산 하에서 Best-of-N 및 Tree-of-Thought 기반선과 비교한다.

Figure 1 : The left part presents the process of inferring process rewards and how we conduct process reward guide tree-search. The right part denotes the self-training of both the process reward model and the policy model.

실험 결과

연구 질문

• RQ1 MCTS*-를 통한 자동 프로세스 보상 추론이 인간의 단계 주석 없이도 고품질의 중간 추론 흔적을 생성할 수 있는가?
• RQ2 PRM-가이드 MCTS*를 사용하는 것이 기존 방법(예: ReST EM, Self-Rewarding)보다 추론 벤치마크에서 정책 및 보상 모델의 자가학습을 개선하는가?
• RQ3 수학 및 과학 과제에서 고정된 탐색 예산하에 기본 추론 정책과 비교하여 ReST-MCTS*의 성능은 어떤가?

주요 결과

• ReST-MCTS*는 동일한 탐색 예산 하에서 이전의 추론 기반선보다 더 높은 정확도를 보인다.
• 정책과 프로세스 보상 모델의 상호 자가학습은 반복에 걸쳐 성능을 향상시키며, ReSTEM 및 Self-Rewarding을 능가한다.
• 추론된 단계별 보상은 트리 탐색을 효과적으로 안내하고 자가학습에 사용되는 더 높은 품질의 흔적으로 이어진다.
• Self-Consistency 및 Best-of-N과 비교하여, 서로 다른 백본에서 다수의 벤치마크에서 ReST-MCTS*와 MCTS*가 개선되었거나 경쟁적인 결과를 달성한다.
• V_theta를 통한 프로세스 보상 모형화는 일부 이전 보상 생성 방법(예: MATH-SHEPHERD)보다 더 강력한 검증 신호를 제공한다.
• SciBench 및 MATH 벤치마크에서 ReST-MCTS*는 여러 LLM 백본에 대해 강력한 향상을 보인다.

(a) Self-training of value model on MATH.