[論文レビュー] Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training

Xidong Feng, Ziyu Wan+5 more|arXiv (Cornell University)|Sep 29, 2023

Natural Language Processing TechniquesComputer Science被引用数 3

ひとこと要約

この論文では、大規模言語モデル（LLM）のデコードおよび学習をガイドする学習済み価値関数を用いた、AlphaZeroに類似した木探索フレームワークであるTS-LLMを提案する。深層木探索とポリシー蒸留、価値関数学習を統合することで、推論・計画・整合性・意思決定タスクにおいて最先端の性能を達成し、探索深さ64までスケーリング可能である。これは、従来の手法が7〜10の深さに制限されていたのとは顕著に異なる。

ABSTRACT

Recent works like Tree-of-Thought (ToT) and Reasoning via Planning (RAP) aim to augment the reasoning capabilities of LLMs by using tree-search algorithms to guide multi-step reasoning. These methods rely on prompting a pre-trained model to serve as a value function and focus on problems with low search depth. As a result, these methods will not work in domains where the pre-trained LLM does not have enough knowledge to serve as an effective value function or in domains that require long-horizon planning. To address these limitations, we present an AlphaZero-like tree-search learning framework for LLMs (termed TS-LLM), systematically illustrating how tree-search with a learned value function can guide LLM decoding. TS-LLM distinguishes itself in two key ways. (1) Leveraging a learned value function and AlphaZero-like algorithms, our approach can be generally adaptable to a wide range of tasks, language models of any size, and tasks of varying search depths. (2) Our approach can guide LLMs during both inference and training, iteratively improving the LLM. Empirical results across reasoning, planning, alignment, and decision-making tasks show that TS-LLM outperforms existing approaches and can handle trees with a depth of 64.

研究の動機と目的

• 既存の木探索手法が提示する課題、例えばプロンプトベースの価値関数依存性や浅い探索深さの制限を解消すること。
• 多様なタスク、LLMのサイズ、探索深さに一般化可能でスケーラブルなフレームワークを構築すること。
• 木探索を推論およびエンドツーエンドのLLM学習に統合し、モデル性能を段階的に向上させること。
• 学習済み価値関数がプロンプトベースの自己評価を上回る信頼性と有効性を示すことを実証すること。

提案手法

• LLMが生成する軌道を教師データとして用い、Supervised Fine-Tuningにより学習された価値関数を有する、AlphaZeroを模した木探索アルゴリズムを導入する。
• モンテカルロ木探索（MCTS）を用い、学習済み価値関数で行動選択をガイドすることで、デコード中に深さ64までの探索が可能になる。
• 木探索の軌道から得た知識をLLMに転移させるためにポリシー蒸留を採用し、LLMのポリシーを向上させる。
• 木探索の軌道から真値ラベルを用いて価値関数を学習し、ポリシーおよび価値関数の反復的改善を可能にする。
• 同じパイプラインを用いて、推論（Game24）、質問応答（PrOntoQA）、RLHFの整合性（RLHF alignment）など複数のタスクに適用する。
• 探索と活用のバランスを取るために、インラインターツリーおよびインターツリーツリー探索を組み合わせたハイブリッド探索戦略を採用する。

Figure 1: (a) Left: Two node expansion paradigms on Game24: sentence-level and token-level. We adopt sentence-level setting in this task. (b) Right: TS-LLM consists of an iterative process over tree-search and training. First, TS-LLM enhances LLM inference by tree-search to obtain improved trajector

実験結果

リサーチクエスチョン

• RQ1 木探索フレームワークにおける学習済み価値関数は、LLM推論タスクにおいてプロンプトベースの自己評価を上回るか？
• RQ2 AlphaZeroに類似した木探索は、自然言語生成において深さ64の深い推論タスクにスケーリング可能か？
• RQ3 LLM学習に木探索を統合することで、推論・計画・整合性といった多様なタスクにおける性能向上が達成できるか？
• RQ4 精度およびスケーラビリティの観点から、ToT や RAP といった既存手法と比較して、TS-LLMの性能はどの程度か？
• RQ5 木探索とポリシー／価値関数の反復的改善のループは、LLMにおける一貫した性能向上をもたらすか？

主な発見

• PrOntoQAでは深さ20で100%の正答率を達成し、MCTS-Rollout や CoT-SC-Tree といったすべてのベースラインを上回った。
• Game24では、MCTS-α-intra木探索を用いたTS-LLMが深さ100で84.53%の正答率を達成し、CoT-SC（18.23%）や BFS-V（72.65%）を大きく上回った。
• RLHFの整合性タスクでは、MCTS-Rolloutを用いたTS-LLMがN=50の平均報酬2.491、最高報酬2.746を記録し、CoT（0.387）や BFS-V（-1.474）を上回った。
• このフレームワークは深さ64までスケーリング可能であり、ToT や RAP が7〜10の深さに制限されるのとは顕著に異なる。
• TS-LLMは、推論、計画、整合性、意思決定のすべての評価タスクで性能向上を示し、汎用性を実証した。
• 125Mパラメータの小さなLLMでも一貫した性能向上が得られるなど、TS-LLMにおける学習済み価値関数は、プロンプトベースの自己評価よりも信頼性が高いことが示された。

Figure 2: Aggregation results for four tasks w.r.t. number of sequences(Path@ $N$ ) on the 1st row and the number of tokens on the 2nd row. TS-LLM benefits from aggregation but struggles to scale in small-scale problems.