PaperlayerView exact revision on arXiv

Source-linked AI summary

Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training

Xidong Feng, Ziyu Wan, Muning Wen, Stephen Marcus McAleer, Ying Wen, Weinan Zhang, Jun Wang

arXiv:2309.17179v2cs.LGcs.AIcs.CL

TL;DR

Existing tree-search methods depend on prompted value functions and shallow search, limiting their applicability across models and long-horizon tasks. TS-LLM uses AlphaZero-like search with learned value functions to guide both LLM decoding and training. Empirical evaluations show broad task coverage, improved decoding, and tree-search depth up to 64.

  • Problem

    Existing tree-search methods rely on prompted value functions and focus on shallow search, limiting applicability across models and long-horizon tasks.

  • Method

    TS-LLM combines AlphaZero-like tree search with learned LLM-based value and outcome reward models to guide inference and iterative training.

  • Results

    Empirical evaluations across reasoning, planning, alignment, and decision-making show that TS-LLM enhances decoding and supports tree search up to depth 64.

  • Takeaways & Limitations

    TS-LLM provides a versatile framework for guiding LLM decoding and training across different tasks and model scales.

  • Takeaways & Limitations

    The scalability of tree-search aggregation remains an open question for future work.

Abstract

from arXiv · show

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.

1. Introduction

Existing tree-search methods improve LLM reasoning but rely on prompted value functions and limited search depths. TS-LLM introduces an AlphaZero-like framework with learned value functions that guides inference and training across tasks and model sizes.

  • Prompt-based value functions make existing tree-search methods dependent on strong LLMs and carefully designed prompts.
  • TS-LLM uses AlphaZero-like deep tree search with a learned LLM-based value function to guide both inference and training.
  • Its iterative pipeline uses tree search to generate improved trajectories, then policy distillation and value-function learning to improve the model.
  • TS-LLM applies across various tasks and LLMs from 125M to 7B parameters, without requiring advanced large-scale models or specialized prompts.
  • TS-LLM extends LLM generation tree search to depth 64, beyond the depths reported for ToT and RAP.
  • Evaluations cover reasoning, planning, alignment, and decision-making tasks, analyzing TS-LLM’s design elements, advantages, and limitations.

2. Related Work

Prior work improves multistep reasoning through prompting, evaluation, and tree search, while TS-LLM extends learned-value tree search to decoding and training across broader task settings.

  • Multistep reasoning research spans base-model improvements, step-by-step prompting, learned reward models, and self-evaluation.
  • Recent search-guided reasoning methods use non-linear trees with beam search, depth-/breadth-first search, or Monte Carlo Tree Search.
  • TS-LLM differs by using a learned value function to guide decoding and training across reasoning and RLHF alignment tasks.
  • Data augmentation methods improve LLMs through rejection sampling or reinforcement learning on augmented datasets.
  • TS-LLM uses tree search as a policy-improvement operator to generate augmented samples for training both LLMs and value functions.

3. Enhancing LLMs with Tree Search

TS-LLM is a versatile AlphaZero-like framework that uses learned value and reward models to guide LLM decoding through tree search and extend that guidance to iterative training. It supports different language-task action spaces and search algorithms, while exposing trade-offs between sentence-level breadth and token-level depth.

  • Problem formulation: Language generation is formulated as a multistep MDP in which the LLM policy samples token sequences as actions and task rewards provide feedback.Intermediate rewards are often sparse, with the final generation receiving the task reward.
  • Action spaces: Sentence-level action nodes suit step-structured tasks, whereas token-level nodes suit tasks without explicit intermediate steps such as RLHF.Sentence-level expansion samples subsequent intermediate steps; token-level expansion treats each token as a discrete action.
  • Search-space trade-offs: Sentence-level search yields shallower trees but requires width-limited subsampling, while token-level search avoids that discrepancy at the cost of much greater depth.The tree max width w and tree max depth d characterize the search space, and the two action designs create different computational challenges.
  • Learned guidance: TS-LLM uses a learned LLM-based value function and outcome reward model to make tree-search guidance more generally applicable than prompt-based approaches.The value function and reward model provide the main guidance for search and can be trained from sampled trajectories and ground-truth rewards.
  • Tree-search algorithms: TS-LLM combines value-guided BFS, DFS, and MCTS with AlphaZero-like variants, including MCTS with value-function approximation.BFS-V prunes by breadth, DFS-V prunes by depth, and MCTS-α evaluates leaf nodes with the learned value function before backup.
  • Training with tree search: During training, tree search improves generations, policy distillation imitates positive trajectories, and policy evaluation retrains the value and outcome-reward models in cycles.The cyclic process iteratively refines the policy, value function, and outcome reward model on tree-search-augmented data.

4. Experiments

Experiments evaluate TS-LLM across tasks, search algorithms, computation budgets, aggregation, iterative training, and value-function training. Results support learned value guidance, AlphaZero-like back-propagation, and iterative updates, while revealing fairness and scalability caveats.

  • Experiment Setups: Five tasks span mathematical, logical, alignment, and chess settings with varied search widths and depths.The evaluation includes GSM8K, Game24, PrOntoQA, synthetic RLHF data, and chess endgames.
  • Performance of Different Algorithms: Learned LLaMA2-V dominates prompt-based GPT-3.5/LLaMA2-7B values across Game24 and GSM8K BFS Path@1 comparisons.This holds even when the few-shot GPT-3.5 policy is out of distribution for LLaMA2-V evaluation.
  • Performance of Different Algorithms: MCTS-α and MCTS-Rollout significantly outperform baselines on long-horizon RLHF and Chess Endgame tasks while remaining comparable on shallow-search tasks.The results are reported for Path@1; shallow settings include GSM8K, Game24, and PrOntoQA.
  • Performance of Different Algorithms: Equal-token comparisons reduce TS-LLM’s advantage against CoT-SCORM, although most tree-search algorithms remain dominant on four tasks with larger search spaces.On GSM8K, only BFS greedy value search is best under the reported equal-token comparison.
  • Performance of Different Algorithms: MCTS is nearly best in both performance and computation cost among BFS-V, DFS-V, and MCTS, indicating the importance of value back-propagation.AlphaZero-style intermediate value back-propagation is also reported as necessary for deep-search problems.
  • Search Aggregation: Aggregation generally benefits TS-LLM, but its gains are smaller than CoT-SC’s in small-scale problems, leaving aggregation scalability open.TS-LLM variants also consume much more computation than CoT, making Path@1 comparisons unfair without equal-token controls.
  • TS-LLM for Training LLM: TS-LLM iteratively improves policy, value, and ORM models through tree-search-generated data, while policy performance slightly exceeds RFT but remains below PPO-based policy.The experiments cover iterative updates on GSM8K and RLHF.

5. Conclusion

TS-LLM is an AlphaZero-like framework for tree-search-guided LLM decoding and training, using a learned value function across diverse tasks and language models. The paper presents it as a versatile framework while identifying scalability and aggregation as areas for further work.

  • TS-LLM guides both LLM decoding and training with AlphaZero-like tree search and a learned value function.
  • The framework is evaluated as a method for guiding LLM decoding and training, with further discussion of limitations and future work deferred to Appendix B.
  • The framework is designed for reasoning tasks and other scenarios, including RLHF alignment, with comparisons across tree-search approaches and computation costs.
  • TS-LLM is presented as generally applicable to different tasks and language models, while supporting iterative improvement of the language model and value function.

B. Limitation and future work

The framework uses AlphaZero-inspired MCTS variants to balance exploration and exploitation during deep search. Its current limitations include computation burdens, incomplete action-space designs, and weak gains from large MCTS aggregation.

  • TS-LLM cannot yet scale to really large-scale scenarios because node expansion and value evaluation add computation burdens.
  • Future work includes key-value caching, broader mixed sentence- and token-level action-space designs, more diverse MCTS aggregation, and multi-iteration scaling.
  • TS-LLM uses MCTS variants because balancing exploration and exploitation becomes difficult when the search tree is sufficiently deep.
  • The search performs selection, leaf expansion and evaluation, then backs up visit counts and action values along the selected path.
  • The framework includes MCTS-α, traditional-style MCTS, and MCTS-Rollout, which differ in their relation to AlphaZero, traditional MCTS, and best-first search.
  • Unlike BFS-V and DFS-V, the MCTS variants propagate value and visit-history information through the search process.

D.1. Different Value Training of iterative update

The iterative-update experiments compare two value-function and ORM training paradigms using different mixtures of tree-search data. Training from initial weights with old and new data performs better on both old and new policies.

  • Table 7 compares value and ORM training from initial weights using mixed old and new tree-search data against continual training using only new data.
  • The mixed-data training paradigm, {v, r̂}ϕ1, outperforms {v, r̂}RLϕ1 on both old and new policies during tree search.

D.2. Results of different node expansion on tasks

Larger tree-search spaces generally improve performance but increase token consumption. Search-algorithm differences are modest on shallow tasks and more pronounced on deeper ones, where MCTS variants gain clearer advantages.

  • The node-expansion experiments report Path@1 results across GSM8K, Game24, and PrOntoQA, with additional ORM-based aggregation baselines averaged over three seeds.
  • Tree-max-width is varied to examine trends in performance and computation consumption, with larger settings especially beneficial for Game24.
  • Larger search spaces generally improve tree-search performance, but wider trees also increase generated-token counts.
  • On GSM8K and Game24, BFS performs well for shallow search, while MCTS-α and MCTS-Rollout show limited gains that are mainly attributed to extra token consumption.
  • For deeper PrOntoQA and RLHF searches, performance gaps among search algorithms become more pronounced across expansion widths.

D.3. Wall-time and Engineering challenges

The reported tree-search implementation incurs substantial runtime overhead relative to greedy decoding, while several engineering strategies could reduce repeated computation and improve parallelism. The authors characterize the current implementation as an unoptimized prototype with substantial acceleration potential.

  • Wall-time results: 2–3 times: Most TS-LLM tree-search algorithms require two to three times CoT greedy decoding’s wall-time, while MCTS-Rollout on GSM8k runs much longer.TS-LLM is also less computationally efficient than CoT-SC because of complicated search procedures and intermediate value-function calls.
  • Wall-time results: Tree-search wall-time comparisons are consistent with comparisons based on the number of generated tokens.
  • Engineering challenges: The current implementation performs repeated child-value computations, leaving substantial potential for accelerating the tree-search process.The authors describe the implementation as an algorithm prototype without specific engineering optimization.
  • Engineering challenges: Sharing policy and value decoders could reduce value-evaluation overhead by reusing policy-rollout computations and replacing additional decoding with MLP computation.
  • Engineering challenges: KV-cache reuse can preserve ancestor-node calculations, while large-batch vectorization could batch inference across multiple parent nodes when resources permit.
  • Engineering challenges: Multi-GPU parallel tree-search remains an engineering need because LLM inference on GPUs is the main bottleneck in the current implementation.

D.4. Discussion about Shared LLM decoder for both policy and critic.

The shared policy/value decoder is explored as an efficiency-oriented alternative to separate decoders. Preliminary Game24 results indicate much faster value estimation, but policy/value stability still requires further work.

  • Training setting: The shared-decoder setup combines supervised-finetuning and value-training losses as Ltotal = LSFT + 0.5 · LValue.
  • Performance comparison: On Game24, shared-decoder CoT performance increases from 12.7 to 16.3, while generated tokens per problem increase from 76 to 166.The authors attribute the longer outputs to violations of Game24’s four-step calculation rules.
  • Computation efficiency: 20x and 9x: A shared decoder increases value-estimation efficiency by 20x for token-level nodes and 9x for sentence-level nodes.The comparison uses token-level width w = 50 and sentence-level width w = 20.
  • Limitations: The shared structure shows computational-efficiency potential, but more work is needed to improve policy/value performance stability.

E.7. Details of aggregation experiments

The aggregation experiments vary search structure, sampling, and answer-aggregation procedures, including intra-tree and inter-tree MCTS-α variants. Iterative updates add tree-search-generated data to policy, value, and reward-model training.

  • Search variants: Inter-tree Search builds a new tree for each search, increasing search-space diversity with computation proportional to the number of searches.The experiments compare MCTS-α intra-tree and inter-tree settings.
  • Search variants: Except for BFS-V, multiple searches are conducted sequentially; for BFS-V, the search count corresponds to beam size.
  • MCTS-α sampling: MCTS-α aggregation adds Dirichlet noise with η ∼ Dir(0.3) and ε = 0.25, then samples actions according to visit counts with τ = 1.
  • Search variants: For path@1, the experiments use MCTS-α without sampling; aggregation experiments use MCTS-α intra-tree or inter-tree variants.
  • Iterative updates: TS-LLM iteratively improves policy, value function, and ORM by adding tree-search data to the training buffer after direct-decoding initialization.
  • Iterative update data: GSM8k tree sampling produces 78.7k distinct answers after deduplication, of which 73.2% are correct.

F. Qualitative Results

The paper presents qualitative generation examples for Game24, GSM8k, PrOntoQA, and the RLHF alignment task. These examples are organized by method and task in the corresponding result tables.

  • Qualitative examples: Qualitative examples cover Game24, GSM8k, PrOntoQA, and the RLHF alignment task.
  • Qualitative examples: Table 18 specifically presents qualitative examples of Game24 outputs.

G.1. Results per task per aggregation

This section organizes qualitative and detailed results across GSM8k, Game24, PrOntoQA, and RLHF alignment. Because of computational-resource limits, most results use one seed, except path@1 results.

  • Results per task per aggregation: Detailed results are provided separately for GSM8k, Game24, PrOntoQA, and RLHF alignment.The corresponding results appear in Tables 22–25.
  • Results per task per aggregation: Most reported results use 1 seed because of computational-resource limits, except path@1 results.The passage explicitly identifies path@1 as the exception.
  • Results per task per aggregation: Qualitative output examples are presented for GSM8k, ProntoQA, and RLHF.These examples are listed in Tables 19–21.
Loading 2309.17179v2…