PaperlayerView exact revision on arXiv

Source-linked AI summary

Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning

Yuxi Xie, Anirudh Goyal, Wenyue Zheng, Min-Yen Kan, Timothy P. Lillicrap, Kenji Kawaguchi, Michael Shieh

arXiv:2405.00451v2cs.AIcs.LG

TL;DR

The paper addresses sparse instance-level preference supervision and static reward modeling by using MCTS for iterative, step-level preference learning with online policy updates. The method improves arithmetic and commonsense reasoning performance, while analysis identifies tradeoffs involving reasoning-chain length and iterative training dynamics.

  • Problem

    Instance-level preference data provides sparse supervision, while conventional RLHF often relies on an offline, static reward model, motivating iterative preference learning for LLM reasoning.

  • Method

    MCTS uses the current policy to generate preference data, decomposes instance-level signals into step-level preferences, combines outcome correctness with self-evaluation, and updates the policy with DPO.

  • Results

    The method produces substantial improvements across arithmetic and commonsense reasoning tasks, including 81.8% on GSM8K, 34.7% on MATH, and 76.4% on ARC-C over the Mistral-7B SFT baseline.

  • Takeaways & Limitations

    Online sampling is important for improving the policy toward optimal alignment, and MCTS-based step-level supervision can refine reasoning across diverse tasks.

  • Takeaways & Limitations

    Longer reasoning chains can introduce errors in commonsense tasks, while continuous iterative updates can cause periodic knowledge loss through insufficient optimization.

Abstract

from arXiv · show

We introduce an approach aimed at enhancing the reasoning capabilities of Large Language Models (LLMs) through an iterative preference learning process inspired by the successful strategy employed by AlphaZero. Our work leverages Monte Carlo Tree Search (MCTS) to iteratively collect preference data, utilizing its look-ahead ability to break down instance-level rewards into more granular step-level signals. To enhance consistency in intermediate steps, we combine outcome validation and stepwise self-evaluation, continually updating the quality assessment of newly generated data. The proposed algorithm employs Direct Preference Optimization (DPO) to update the LLM policy using this newly generated step-level preference data. Theoretical analysis reveals the importance of using on-policy sampled data for successful self-improving. Extensive evaluations on various arithmetic and commonsense reasoning tasks demonstrate remarkable performance improvements over existing models. For instance, our approach outperforms the Mistral-7B Supervised Fine-Tuning (SFT) baseline on GSM8K, MATH, and ARC-C, with substantial increases in accuracy to $81.8\%$ (+$5.9\%$), $34.7\%$ (+$5.8\%$), and $76.4\%$ (+$15.8\%$), respectively. Additionally, our research delves into the training and inference compute tradeoff, providing insights into how our method effectively maximizes performance gains. Our code is publicly available at https://github.com/YuxiXie/MCTS-DPO.

1 Introduction

The paper addresses limitations of static, instance-level preference learning by using MCTS to generate granular step-level preferences online. It combines self-evaluation with DPO to iteratively improve LLM reasoning, reporting gains across arithmetic and commonsense tasks.

  • Motivation: Iterative learning repeatedly uses the current policy to collect and analyze new preference data before updating the policy.This contrasts with RLHF settings where the reward model is often trained offline and remains static.
  • Motivation: Instance-level preference data provides sparse supervision and may not fully exploit MCTS for improving LLMs.The paper identifies preference granularity and reliance on critics or learned reward functions as central challenges.
  • Approach: MCTS rollouts provide step-level guidance, while self-evaluation lets the model act as both policy and critic during preference collection.The approach targets more granular and efficient policy improvement than instance-level supervision.
  • Approach: MCTS-enhanced iterative preference learning uses DPO to update the LLM policy from preferences generated on the fly.The algorithm begins with an initial SFT policy, iteratively samples data with MCTS, and applies preference learning updates.
  • Results: 81.8% on GSM8K, 34.7% on MATH, and 76.4% on SciQ are reported for the proposed method, exceeding the Mistral-7B SFT baseline.The corresponding improvements are +5.9%, +5.8%, and +15.8%, respectively.

2 MCTS-Enhanced Iterative Preference Learning

The method iteratively improves an LLM by using its current policy to generate search trees, extracting step-level preferences with MCTS and self-evaluation, and updating the policy with DPO. Search values guide preference labels, while adaptive smoothing and online policy updates address noisy labels and theoretical failure cases.

  • Iterative framework: Each iteration samples prompts with the current policy, generates responses, extracts dynamically evolving preference data, and updates the policy.The updated policy is then used in the next data-collection iteration.
  • Step-level MCTS: MCTS decomposes reasoning into discrete token-sequence steps and uses look-ahead to estimate expected future rewards.Stepwise self-evaluation refines these predictions to improve process consistency and decision accuracy.
  • Step-level MCTS: MCTS alternates selection, expansion, and backup while balancing exploration and exploitation through action values, visit counts, and policy probabilities.The sampling breadth is initialized at b1 and annealed to a smaller b2 < b1 for subsequent steps.
  • Preference construction: Higher-Q candidate steps are preferred, with outcome correctness and self-evaluation combined to compute rewards during expansion.Correct terminal, incorrect terminal, and unfinished intermediate states receive outcome values 1, −1, and 0, respectively.
  • Preference construction: For tasks where the initial policy performs poorly, ground-truth reasoning steps are additionally included during MCTS preference-data collection.The paper gives MATH as an example of this setting.
  • Preference learning: DPO updates the policy from MCTS preferences using conservative optimization and visit-count-based adaptive label smoothing.The preferred and dispreferred responses are represented as yw and yl, while β scales the KL constraint.

3 Theoretical Analysis

The analysis compares offline and online preference sampling, showing that stale sampling policies can fail while latest-policy sampling avoids the identified failure case under the abstract model. This motivates the paper’s online DPO formulation.

  • The framework covers offline sampling from a fixed distribution and online sampling from the latest policy at each iteration.
  • Offline preference learning can fail with high probability when its sampling policy differs substantially from the current policy.The theorem constructs a case where the learned policy assigns the optimal outcome probability at most 1 − c with probability at least 1 − 2ϵ^M.
  • If policy mismatch is extreme, the offline setting can assign the optimal outcome near-zero probability with probability one, motivating on-policy sampling.The discussion describes the problematic regime as ϵ = 0 and c approximately 1.0.
  • Under the online setting, sampling from the latest policy guarantees optimal-outcome probability one after at least n + 1 iterations in the abstract model.The guarantee applies for any global optimizer of the preference losses.
  • Because the optimal policy is inaccessible in practice, the method uses online DPO with preference samples drawn from the latest available policy.

4 Experiments

Experiments evaluate MCTS-enhanced iterative preference learning across arithmetic and commonsense reasoning, then analyze compute scaling, self-evaluation, ablations, and qualitative behavior. The method improves several benchmarks, while results expose task-dependent benefits and limitations of extended intermediate reasoning.

  • Arithmetic Reasoning: 81.8% on GSM8K and 34.7% on MATH mark substantial arithmetic gains from the Mistral-7B baseline.The method increases GSM8K from 75.9% to 81.8% and MATH from 28.9% to 34.7%.
  • Commonsense Reasoning: 2.5%, 3.0%, and 2.1% absolute gains occur on ARC-C, AI2Sci-M, and SciQ, respectively, over direct tuning.On OBQA and CSQA, the method improves substantially over SFT but remains less efficient than direct tuning.
  • Training- vs. Test- Time Compute Scaling: Training-time compute scaling gives the method a higher pass-rate scaling exponent, especially on unseen SciQ, while inference-time ceilings vary by task.On SciQ, performance reaches 88.6% versus roughly 84% for self-consistency; on MATH, sampling exceeds 35% while post-training performance is around 32.2%.
  • Functions of Self-Evaluation Mechanism: Self-evaluation accuracy improves significantly across GSM8K, MATH, and ARC-C when ground-truth example answers are included.The ablation compares self-evaluation score distributions and discriminative ability with and without the “EXAMPLE ANSWER.”
  • Ablation Study: 76.4% on ARC-C and 88.5% on SciQ outperform offline and instance-level variants, supporting step-level online preference learning.On MATH without ground-truth supervision, differences are smaller: 32.2% for the method versus 31.9% offline and 32.9% instance-level.
  • Qualitative Analysis: Longer MCTS reasoning chains can introduce commonsense errors in OBQA, whereas the method corrects mistakes and forms accurate extended chains on MATH.The qualitative analysis compares intermediate reasoning correctness across OBQA, AI2Sci, and MATH.

5 Related Work

Self-improvement research collects training data from model generations, using static reward heuristics, self-rewarding, or exploration strategies to guide online data collection.

  • Static reward heuristics guide the collection of high-quality positive examples from model generations.
  • Continuously updated self-rewarding collects both positive and negative data for preference learning.
  • Exploration via rejection sampling is another strategy for online data collection.

6 Conclusion

The paper presents MCTS-enhanced iterative preference learning, combining step-level preference sampling with DPO to improve LLM alignment. Its analysis emphasizes online sampling and highlights tradeoffs involving search strategies, historical data, and offline versus online learning.

  • MCTS-enhanced iterative preference learning uses MCTS as a policy improvement operator for step-level preference learning.The approach is designed to enhance LLM alignment and produce high-quality training data through quality exploitation and diversity exploration.
  • Online sampling is theoretically identified as key to improving the LLM policy toward optimal alignment.
  • The paper proposes exploring search strategies and historical data or policies to augment and diversify training examples.
  • The analysis highlights a tradeoff between offline and online learning because online learning can exhibit cyclic performance changes.

A Related Work

Related work spans iterated learning, self-training, preference learning, and guided reasoning search. The paper distinguishes its approach by combining outcome assessment and self-evaluation, then decomposing feedback into fine-grained MCTS signals for step-level preference learning.

  • Iterated Learning: Iterated learning lets an apprentice self-play, learn from expert feedback, and replace the current expert in the next iteration.Prior work applies expert iteration to formal mathematical reasoning and proof search.
  • Novelty: The proposed method combines outcome assessment and LLM self-evaluation, decomposing their feedback through MCTS into step-level preference signals.
  • Self-Training: Self-training assigns pseudo-labels to unlabeled data using a learned labeler or other labeling strategy.Recent approaches use static reward heuristics, confidence filtering, or continuously updated self-rewarding to curate training examples.
  • Preference Learning: Preference learning commonly uses reward models for online RL, while newer methods avoid separate reward or value networks through alternative preference-learning procedures.
  • Guided Search for Reasoning: Guided-search methods improve reasoning by eliciting intermediate chains, decomposing them into steps, and applying process supervision.Process supervision has been used to train more reliable reward models than outcome supervision in mathematical reasoning.

B Theoretical Analysis of Online DPO

The theoretical analysis compares offline and online DPO and shows why fixed sampling can fail to correct suboptimal outputs. It motivates on-policy sampling and sufficient per-iteration optimization to avoid cyclic behavior and improve convergence.

  • Motivation: DPO uses a fixed policy to collect preference data, which can diverge from the current policy and complicate learning from policy-generated data.The paper identifies this sampling-policy discrepancy as a problem for preference learning.
  • Framework: The abstract formulation covers offline settings with a fixed sampling distribution and online settings that sample from the latest policy.
  • Theoretical Failure Case: A fixed sampler may never sample a suboptimal output that remains probable under the current policy, allowing the loss to minimize without eliminating that output.
  • Theoretical Failure Case: The probability of never sampling such an output over m iterations is bounded below by (1 − 2ϵ)^m ≥ 1 − 2ϵm.
  • Online Dynamics: Under online updates, previously dispreferred outputs can be excluded from later sampling, increasing the number of suboptimal samples assigned zero probability.
  • Implications: The analysis motivates online DPO because jointly sampling from the current and optimal policies would avoid the fixed-sampling failure, although the optimal policy is unavailable in practice.
  • Implications: Insufficient DPO minimization per online iteration can produce cyclic behavior, motivating several rounds of loss minimization at each iteration.

C Implementation Details

The experiments use Mistral-7B with specified SFT and DPO settings, while MCTS configurations vary across arithmetic and commonsense reasoning tasks. Hyperparameter analysis favors broader initial searches followed by narrower subsequent searches to reduce compute while maintaining performance.

  • Mistral-7B experiments use up to 4 × 40GB NVIDIA A100 GPUs.
  • DPO uses a 1e-6 learning rate, batch size 32, β = 0.1, 64-token maximum steps, and K = 5 MCTS iterations.SFT uses a 5e-6 learning rate and batch size 128; both use cosine scheduling.
  • Arithmetic reasoning: Arithmetic preference learning combines GSM8K and MATH training data, using b1 = 5, b2 = 3, and maximum search depth d = 4.The combined prompt dataset contains 24K samples, and collecting preferences takes about 2 minutes per sample.
  • Commonsense reasoning: Commonsense preference learning combines ARC, AI2Science, OBQA, and CSQA, using b1 = 4, b2 = 2, and maximum search depth d = 3.The combined dataset contains 12K samples, and early stopping is used around 50% of training progress.
  • MCTS hyperparameter tuning: O2 performs better by using a broader initial search space and a smaller subsequent search space, reducing compute while maintaining good performance.The comparison is reported in Table 5.

D Further Analysis

Further analyses show that combining outcome correctness with self-evaluation improves search and learning performance, while qualitative trees illustrate how MCTS and policy training affect reasoning paths and diversity.

  • Reward criteria in MCTS: 83.0% on ARC-C, 90.5% on AI2Sci-M, and 85.8% on GSM8K are achieved after integrating outcome correctness and self-evaluation.The corresponding baselines are 60.6%, 70.9%, and 75.9%.
  • Reward criteria in MCTS: 76.4% accuracy on ARC-C follows learning with greedy decoding, compared with a 9.1% increase in MCTS performance.The reported comparison indicates a larger learning gain on a task where the initial policy is weaker.
  • Reward criteria in MCTS: Self-evaluation consistently improves MCTS performance before and after learning when included with other reward components.
  • Qualitative analysis on collected preferences: MCTS identifies the correct science answer through broad search and collects steps along the successful path as positive preference-learning samples.The example ground-truth option is D, and the search tree reports Q values for preference pairs.
  • Qualitative analysis on collected preferences: The trained policy produces more successful predictions of the ground-truth option D and tends to generate longer reasoning chains under fine-grained process supervision.
  • Qualitative analysis on collected preferences: Higher generation diversity lets MCTS explore more alternative correct solutions, producing more paths with correct predictions.

E Extended Experiments

Extended experiments examine alternative loss functions and base models, while qualitative examples compare MCTS search trees across breadth settings and policies trained at different stages.

  • Loss function: IPO achieves similar performance to DPO while providing more stable performance on the held-out dataset.The paper presents IPO as an alternative preference-learning loss addressing DPO overfitting.
  • Base model: Llama2-13B is used to validate the approach's generalizability on arithmetic reasoning with SFT on Arithmo and DPO on GSM8K and MATH.The experiment uses up to 2 × 80GB NVIDIA A100 GPUs.
  • Qualitative search-tree examples: The science-question search trees compare MCTS settings with b1 = 4, b2 = 2; b1 = 3, b2 = 3; and a policy after preference learning.Figures 5–7 use the same question, whose ground-truth answer is option D.
  • Qualitative search-tree examples: The GSM8K search-tree examples compare an initial policy with a policy tuned after 3 epochs using the same search breadth.The example asks how many loaves remain after sales and returned loaves are accounted for.
Loading 2405.00451v2…