Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training (2309.17179v2)

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, LLMs 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.

• The paper presents TS-LLM, an AlphaZero-inspired tree-search framework that integrates a learned value function to guide LLM decoding and training.
• It employs dual-purpose methods by combining policy distillation and value function learning to iteratively enhance reasoning and output accuracy.
• Empirical results demonstrate significant performance gains in deep planning tasks, outperforming CoT baselines on benchmarks like RLHF alignment.

Alphazero-like Tree-Search can Guide LLM Decoding and Training

Introduction

Recent advancements in guiding LLMs through tree-search algorithms highlight the potential for enhanced reasoning capabilities. Traditional approaches like Tree-of-Thought (ToT) and Reasoning via Planning (RAP) have demonstrated the ability to boost performance using tree-search methods such as BFS/DFS and MCTS. However, these methods are constrained by shallow search depths, typically limited to 10 or fewer steps, impeding effectiveness in tasks requiring deeper planning.

This study introduces an AlphaZero-like framework (TS-LLM) that leverages a learned value function to expand the applicability of tree-search algorithms in LLM decoding and training across various problem domains with greater search depths.

Key Innovations

1. Learned Value Function: TS-LLM utilizes a value function adapted from LLMs to provide more reliable evaluations than prompt-based self-assessment techniques.
2. Dual Purpose - Training and Inference: Unlike methods solely focused on inference, TS-LLM integrates tree-search into the training process of LLMs, encouraging iterative enhancement through policy distillation and value function learning.

   Figure 1: Overview of TS-LLM showing sentence-level and token-level node expansion paradigms for tree-search integration.

Methodology

Tree-Search Algorithm

TS-LLM adopts AlphaZero-like tree-search algorithms to guide the decision-making in LLMs during both inference and training phases. The tree is constructed with node expansion governed by either a sentence or token-level basis, depending on the nature of the task. This setup facilitates deep search capabilities.

• Search Algorithm Variants:
  • BFS-V and DFS-V: These algorithms utilize value-based pruning during breadth and depth traversal respectively.
  • MCTS and MCTS-$\alpha$: Combining Monte Carlo methods with value function approximations ensures robust search through potential outputs, optimizing cumulative rewards.

Evaluation and Training

The learned value function and reward models are developed using training datasets with rewards labelled based on task-specific outcomes. TS-LLM's novel training paradigm iteratively refines LLM performance:

• Policy Improvement: Tree-search enhances the generation dataset.
• Policy Distillation: Supervised learning from augmented data.
• Policy Evaluation: Continual adaptation of value functions using augmented samples.

Empirical Analysis

Performance Metrics

Results demonstrate TS-LLM's superiority in deep planning tasks compared to traditional methods. Evaluation metrics across diverse tasks, including reasoning and alignment, illustrate significant improvements in accuracy and reward optimization.

• Comparison of Path@1: TS-LLM consistently outperforms CoT baselines, particularly in complex tasks like Chess Endgame and RLHF alignment, capable of handling trees with depth up to 64.

  Figure 2: Aggregated task results showing the progression in performance with increasing tree-search depth.

Scalability and Efficiency

The integration of tree-search in both training and inference ensures scalable application across tasks of varying complexity. However, computational overhead, especially during node expansion, suggests a need for further optimizations.

Conclusion

TS-LLM represents a significant step forward in integrating advanced tree-search methodologies with LLMs, supporting improvements in both performance and training efficiency. Future exploration may focus on addressing computational burdens and expanding the framework's applicability to broader domains, potentially transforming practices in LLM-based decision-making and reasoning tasks.

This work lays the foundation for continued advancements in AI-driven LLM optimization, fostering robust and adaptive AI systems for complex problem-solving scenarios.