[論文レビュー] ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search

Dan Zhang, Sining Zhoubian+4 more|arXiv (Cornell University)|Jun 6, 2024

Business Process Modeling and Analysis被引用数 6

ひとこと要約

ReST-MCTS*を導入する自己学習フレームワーク。MCTS*-ガイド付きのプロセス報酬モデルが各ステップの推論を自動でラベル付けし、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 searchを用いて中間推論の軌跡を生成・評価する。
• 反復的な自己学習を通じて方針モデルとプロセス報酬モデルを相互に精練する。

提案手法

• 部分解に対する各ステップの品質 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 と比較して、MCTS* を用いる ReST-MCTS* は、異なるバックボーンに跨る複数のベンチマークで改善または競合的な結果を達成。
• V_theta によるプロセス報酬モデリングは、従来の報酬生成手法の中にはより強い検証信号を提供する（例：MATH-SHEPHERD）。
• SciBench および MATH ベンチマークで、ReST-MCTS* は複数のLLMバックボーン全体で堅牢な改善を示す。

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