Source-linked AI summary
Toward Self-Improvement of LLMs via Imagination, Searching, and Criticizing
Ye Tian, Baolin Peng, Linfeng Song, Lifeng Jin, Dian Yu, Haitao Mi, Dong Yu
TL;DR
LLMs remain unreliable on complex reasoning and planning, and existing prompting, fine-tuning, and self-critique approaches face data or feedback limitations. ALPHALLM integrates prompt synthesis, language-tailored MCTS, and critic models into an annotation-free self-improvement loop, improving mathematical reasoning performance and reaching GPT-4-comparable results in the reported experiments.
Problem
LLMs still struggle with complex reasoning and planning, while self-critique effectiveness in these settings remains uncertain and supervised improvement depends on high-quality data.
Method
ALPHALLM combines synthetic prompt generation, ηMCTS for efficient language-task search, and three critic models to guide policy improvement without additional annotations.
Results
57.8 to 92.0 on GSM8K and 20.7 to 51.0 on MATH were achieved from Llama-2-70b and WizardMath-70B-V1.0, performing comparably to GPT-4.
Takeaways & Limitations
The results demonstrate the potential for an effective self-improvement loop that enhances LLM mathematical reasoning without extra data annotations.
Takeaways & Limitations
Synthetic prompt generation is relatively simple, and greedy-sampling performance is substantially inferior to decoding with ηMCTS.
Abstract
from arXiv · showhide
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.
1 Introduction
LLMs remain challenged by complex reasoning and planning, while self-improvement through self-critique lacks reliable evidence in such settings. ALPHALLM combines prompt synthesis, efficient MCTS, and critic-guided trajectory selection to improve LLMs without additional annotations.
- Complex reasoning and strategic planning remain difficult for LLMs despite their broad capabilities and advanced prompting methods.
- Self-correction and self-learning refine responses or learn from self-judged rewards, but their effectiveness for intricate reasoning and planning remains uncertain.
- The framework adapts MCTS to language-task challenges involving scarce data, vast search spaces, and subjective feedback.
- ALPHALLM synthesizes prompts, applies ηMCTS to search language trajectories, and uses critic signals to improve the policy.
- 57.8 to 92.0 on GSM8K and 20.7 to 51.0 on MATH were achieved from Llama-2-70b and WizardMath-70B-V1.0, performing comparably to GPT-4.
2 Related Work
Related work covers effective search strategies for complex reasoning and planning and self-improvement methods that rely on internal supervision and critique.
- Search with LLM: Search methods for reasoning tasks include dynamically pruned beam search and tree- or graph-based representations of possible reasoning paths.
- LLM Self-improving: LLM self-improvement seeks alignment with human preferences and values using knowledge inside the model, making reliable critique signals a central concern.
3 Preliminaries
The paper formulates autoregressive text generation as an MDP and applies MCTS to optimize policies over response trajectories. Self-improvement iteratively refines the policy using searched responses and rewards.
- Problem Formulation: An LLM generates response tokens autoregressively, with each token conditioned on the preceding context.
- Problem Formulation: The generation process is modeled as an MDP whose states encode trajectory context, actions are sampled tokens, and transitions concatenate actions to states.
- Problem Formulation: Self-improvement iteratively refines a policy from an initial expert-generated dataset using synthesized prompts and searched responses that maximize expected reward.
- Monte Carlo Tree Search: MCTS repeatedly performs selection, expansion, evaluation, and backpropagation while sampling a search tree for policy optimization.
- Monte Carlo Tree Search: UCB balances exploration and exploitation using node values, visit counts, parent visit counts, and a tuning hyperparameter.
4 ALPHALLM
ALPHALLM builds a self-improvement pipeline around prompt synthesis, ηMCTS, and critic models. Its search design reduces language-search complexity through option-level actions, adaptive branching, state merging, rollout policies, and process supervision.
- ALPHALLM Architecture: ALPHALLM comprises imagination for prompt synthesis, ηMCTS for high-quality trajectory search, and critics that provide reliable guidance.
- Imagination: The synthesizer transforms one or more initial examples into prompts intended to increase training-data diversity and complexity.
- Option-level MCTS: Option-level MCTS treats a search node as a sequence of tokens or sentences, reducing search depth relative to token-level exploration.
- Importance-Based Adaptive Branching: A fixed branching factor can make option-level search too shallow or too thin because the option space contains unlimited token combinations.
- Importance-Based Adaptive Branching: Theorem 4.1 sets the optimal branching factor so mt−1 is proportional to node importance I(st), allocating more children to states with larger value deviations.
- State Merging: State merging groups similar options to increase diversity among explored groups and cover more of the problem space with limited rollouts.
- Simulation: Rollout simulation projects future trajectories to improve exploration, exploitation, and reward estimation, while multiple trajectories reduce estimation variance.
- Critics: The critic system includes a value function, a process reward model for immediate option rewards, and additional critic models designed to guide search.
5 Experiments
Experiments evaluate ALPHALLM on GSM8K and MATH using synthetic prompts, ηMCTS, and ablations of its search components. Results show gains from self-improvement, adaptive search design, and increased rollout computation.
- Experimental Setup: ALPHALLM is evaluated on GSM8K and MATH using final-answer annotations, synthetic prompts, and comparisons with proprietary and open-source models.The evaluation uses chain-of-thought prompting and measures average rollouts as a computational-efficiency indicator.
- Ablation Study: 79.5% accuracy with vanilla MCTS increases to 84.9% with adaptive branching, 85.9% with PRM, 86.5% with ORM and fast rollout, and 87.0% with state merging on GSM8K.Increasing the number of rollouts alongside the other components produces the best performance in the ablation.
- Ablation Study: On MATH, ηMCTS reaches 45.4 accuracy with 148 rollouts, versus 44.1 with 198 rollouts when option formulation is removed.Using only intrinsic knowledge for ORM lowers accuracy to 38.8, indicating the importance of tool augmentation for difficult math assessment.
- Self-Improvement Analysis: Training on reranked or ηMCTS trajectories improves the initial policy across two rounds, while ηMCTS yields better accuracy and efficiency than reranking.The comparison reports greedy decoding and ηMCTS with small or large rollout budgets.
- Hyperparameter Analysis: Increasing fast-rollout trajectories improves performance, whereas varying state-merge heuristic functions has limited impact.The experiments use n = 4 as a performance-efficiency trade-off for fast-rollout estimation.
6 Conclusion
ALPHALLM combines prompt synthesis, efficient ηMCTS search, and critic models to improve LLMs without additional annotations. On mathematical reasoning tasks, it substantially improves performance and can perform comparably to GPT-4 when decoded with ηMCTS.
- Conclusion: ALPHALLM addresses data scarcity, large language search spaces, and subjective feedback through prompt synthesis, optimized ηMCTS, and three critic models.The framework integrates MCTS with LLMs to create a self-improving loop without additional annotations.
- Conclusion: On mathematical reasoning tasks, ALPHALLM significantly boosts LLM performance without requiring extra data annotations.The reported findings support self-improvement through synthetic prompts, search, and criticism.
- Conclusion: When decoded with ηMCTS, ALPHALLM performs comparably to GPT-4.This conclusion is reported for the paper’s mathematical reasoning experiments.
A.1 Imagination, Searching, Criticizing and Learning Loop
Algorithm 1 defines an LLM self-improving loop that synthesizes prompts, collects trajectories with reward-guided search, constructs a dataset, and updates the policy. The algorithm is presented as the paper’s central iterative procedure.
- Loop overview: The self-improving loop begins with an initial dataset and iterates for K training loops.Each iteration produces updated policy parameters θk.
- Procedure: Algorithm 1 is the stated implementation of the LLM self-improving loop.The text directs readers to the algorithm for the complete procedure.
- Imagination: Synthetic prompts are generated from the previous policy and dataset before each policy update.The generated prompts form the basis for constructing the next training dataset.
- Searching: The method collects trajectories using a search algorithm such as MCTS guided by a reward model.Search is applied under the current policy and synthetic-prompt dataset.
- Learning: Collected trajectories are used to construct Dk, after which the policy is updated by minimizing a loss against the previous policy.The update maps the synthesized dataset into new policy parameters θk.
A.2 Option-level MCTS
Option-level MCTS searches by repeatedly selecting nodes, expanding them with policy-sampled options, and evaluating rewards and simulated future trajectories. Its operations define how language-task search proceeds over options.
- Search operations: Option-level MCTS consists of selection, expansion, and simulation operations.These operations organize the search from the root through newly generated option states.
- Selection and expansion: Selection iteratively chooses a child from the root, while expansion generates a new node from the parent’s previous state.The new option is sampled using policy π and terminated according to β.
- Simulation: Simulation evaluates the scaled reward of the expanded node together with simulated future trajectories.This provides the search with reward information beyond the newly expanded node.
A.3 Importance-Based Adaptive Branching Under Uniform Distribution
Under a uniform value distribution, the analysis derives an importance-based branching rule for tree search. The resulting theorem states that the branching factor grows with node importance.
- Uniform-distribution analysis: The analysis models mt values as uniformly distributed between vmax and vmin and bounds the expected minimum distance for a new value.The bound is expressed using the range vmax−vmin.
- Importance relation: The expected minimum-distance constraint is connected to node importance through vmax−vmin = 2I(st).This yields the bound Eϕ(t) ≤ I(st).
- Adaptive branching: Theorem 4.1 sets the optimal branching factor so that mt−1 is proportional to node importance I(st).The theorem states this as the optimal rule under the specified condition on I(st).
- Proof strategy: The proof formulates an optimization problem, introduces Lagrange multipliers, and solves the resulting stationarity conditions.The derivation proceeds by differentiating the Lagrangian with respect to mt and λt.
A.4 Importance-Based Adaptive Branching Under Gaussian Distribution
Under Gaussian assumptions, the branching analysis uses approximations for minimum absolute differences and sample ranges to derive a constraint on mt. The resulting analysis finds a positive correlation between branching and node importance.
- Gaussian model: The Gaussian analysis treats relevant differences as independent and identically distributed Gaussian random variables.It then studies the minimum of mt Gaussian differences.
- Minimum-difference approximation: The expected minimum absolute difference is approximated using order statistics and scales inversely with the square root of the sample count.For mt samples, the approximation captures dependence on the number of samples.
- Range approximation: The expected range of mt normal samples is approximated through Gaussian extreme-value percentiles and its dominant term.The range approximation uses z0.9995 and z0.0005 percentiles before simplifying to the primary term.
- Optimization: The optimization rewrites the constraint using g(mt) = mt ln(mt) and seeks the minimum mt satisfying g(mt) ≥ I2(st).For mt > e−1, the analysis sets the constraint to equality to solve for mt.
- Conclusion: The analysis concludes that mt and node importance I(st) have a positive correlation.The equation cannot be solved directly, but the directional relationship remains observable.
A.5 Prompt Templates
The appendix specifies distinct critic prompts for process-level and outcome-level assessment, alongside policy prompts used for self-improvement experiments.
- A.5.1 PRM: The PRM prompt asks the model to verify whether the last reasoning step is correct and return a textual reward.
- A.5 Prompt Templates: The PRM assesses local step correctness, whereas the ORM assesses trajectory-level solution quality.
- A.5 Prompt Templates: Both critic templates produce textual rewards that can guide assessment during self-improvement.
- A.5.2 ORM: The ORM prompt evaluates the complete solution, including method, calculations, use of problem information, and final answer.
- A.5.3 Policy Finetuning: For MATH, policy fine-tuning uses WizardMath V1.0 70B with its proposed self-improvement system prompt.
- A.5.3 Policy Finetuning: The GSM8K policy prompt frames interaction as a helpful, detailed, and polite assistant responding to a user query.
A.6 MCTS Details
The appendix details search comparisons, critic evaluation, implementation resources, and study limitations relevant to interpreting ALPHALLM's MCTS results.
- A.7 Additional Ablations: Using Llama-2-70b as the fast-rollout model improves performance by reducing bias, whereas Abel-002-7B is faster at similar computational resources.
- A.8 Search Comparison: Reranking with ORM consistently outperforms self-consistency across response counts from 10 to 50.
- A.8 Search Comparison: ηMCTS achieves better MATH results than reranking with only half as many rollouts, while BFS reaches 80.6 on GSM8K with 87.9 rollouts.
- A.9 Critic Evaluation: Value function and PRM evaluation on GSM8K compares precision, recall, and ECE: the value function has higher precision and better calibration, while PRM has higher recall.
- Implementation: The experiments used four GPUs for 70B model serving, one GPU for 7B serving, and 64 GPUs to train 70B models.
- A.12 Limitations and Future Work: The authors identify simple synthetic-prompt generation, weaker greedy-decoding performance than ηMCTS, and static critics as limitations requiring further exploration.