[Paper Review] Toward Self-Improvement of LLMs via Imagination, Searching, and Criticizing

Ye Tian, Baolin Peng+5 more|arXiv (Cornell University)|Apr 18, 2024

Artificial Intelligence in Law4 citations

TL;DR

This paper introduces AlphaLLM, a self-improvement framework for large language models (LLMs) that integrates Monte Carlo Tree Search (MCTS) with LLMs through an imagination-searching-criticizing loop. By synthesizing prompts, using an optimized MCTS with hierarchical options, and employing a trio of critic models for precise feedback, AlphaLLM boosts mathematical reasoning performance without additional annotations, achieving near-GPT-4 levels when decoded with MCTS.

ABSTRACT

Despite the impressive capabilities of Large Language Models (LLMs) on various tasks, they still struggle with scenarios that involves complex reasoning and planning. Recent work proposed advanced prompting techniques and the necessity of fine-tuning with high-quality data to augment LLMs' reasoning abilities. However, these approaches are inherently constrained by data availability and quality. In light of this, self-correction and self-learning emerge as viable solutions, employing strategies that allow LLMs to refine their outputs and learn from self-assessed rewards. Yet, the efficacy of LLMs in self-refining its response, particularly in complex reasoning and planning task, remains dubious. In this paper, we introduce AlphaLLM for the self-improvements of LLMs, which integrates Monte Carlo Tree Search (MCTS) with LLMs to establish a self-improving loop, thereby enhancing the capabilities of LLMs without additional annotations. Drawing inspiration from the success of AlphaGo, AlphaLLM addresses the unique challenges of combining MCTS with LLM for self-improvement, including data scarcity, the vastness search spaces of language tasks, and the subjective nature of feedback in language tasks. AlphaLLM is comprised of prompt synthesis component, an efficient MCTS approach tailored for language tasks, and a trio of critic models for precise feedback. Our experimental results in mathematical reasoning tasks demonstrate that AlphaLLM significantly enhances the performance of LLMs without additional annotations, showing the potential for self-improvement in LLMs.

Motivation & Objective

• To address the limitations of data scarcity and poor reasoning in LLMs for complex tasks.
• To enable self-improvement in LLMs without relying on external annotations or high-quality supervised data.
• To adapt MCTS for natural language tasks by overcoming vast search spaces and subjective feedback.
• To design a scalable, self-supervised loop where LLMs generate, search, and critique their own reasoning trajectories.
• To demonstrate that MCTS-enhanced LLMs can achieve performance comparable to GPT-4 in mathematical reasoning.

Proposed method

• Introduces an imagination component that synthesizes strategic prompts to alleviate data scarcity.
• Proposes η-MCTS, a hierarchical MCTS variant that treats text generation as options over a Markov Decision Process (MDP), enabling multi-level planning.
• Employs a trio of critic models—reasoning, tool-augmented, and intrinsic knowledge critics—to provide precise, multi-faceted feedback on response quality.
• Uses self-play trajectories from MCTS rollouts to fine-tune the policy network, enabling iterative self-improvement.
• Applies state merging and increased rollouts to improve search efficiency and convergence.
• Integrates a reward model (ORM) that combines intrinsic knowledge and external tools to assess complex reasoning steps.

Figure 1: Imagination-Searching-Criticizing self-improvement loop: Imagination component synthesizes prompts as new learning examples, with MCTS searching better trajectories guided by signals from critics for policy improving.

Experimental results

Research questions

• RQ1 Can MCTS be effectively adapted to natural language generation tasks to enable self-improvement in LLMs?
• RQ2 How can LLMs generate high-quality synthetic data to overcome data scarcity in reasoning tasks?
• RQ3 Can a multi-critic framework provide more accurate and reliable feedback than single-critic or heuristic methods?
• RQ4 To what extent can MCTS-guided search improve reasoning accuracy without additional annotations?
• RQ5 Does iterative self-improvement via MCTS lead to continual performance gains in LLMs?

Key findings

• AlphaLLM achieves 87.0% accuracy on the GSM8K mathematical reasoning benchmark, significantly outperforming the base model.
• With 148 rollouts, η-MCTS reaches 45.4% accuracy on the MATH dataset, outperforming sentence-level MCTS (44.1%) with fewer rollouts.
• The inclusion of tool-augmented critics increases performance from 38.8% to 45.4%, demonstrating the critical role of external tools in feedback.
• Models trained on trajectories from η-MCTS outperform those trained on reranked trajectories and match GPT-4 performance when decoded with η-MCTS.
• Iterative self-improvement over two rounds leads to consistent performance gains, confirming the scalability of the self-improving loop.
• The ablation study confirms that option formulation and external tools are essential for high performance, with performance dropping significantly when either is removed.

Figure 2: An overview of the four operations of $\eta$ Mcts . A node is selected, expanded, simulated with fast rollout policy until a terminal node is reached, then the signals from value function, PRM and ORM are backpropagated.