🔗9.3.5 MCTS-DPO

Paper link: Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning
📌
Researchers from the National University of Singapore and DeepMind proposed a new method that uses Monte Carlo Tree Search (MCTS) to generate preference data. Instead of directly using Direct Preference Optimization (DPO, see Section 3.5 DPO) to enhance large language model reasoning, inspired by AlphaZero, this work uses MCTS to decompose instance-level rewards into finer-grained step-level signals for collecting preference data, and then directly applies iterative DPO to update LLM policies.
On the GSM8K, MATH, and ARC-C datasets, this method significantly outperformed the Mistral-7B supervised fine-tuning (SFT) baseline, with accuracy improvements to 81.8% (+5.9%), 34.7% (+5.8%), and 76.4% (+15.8%), respectively.

Background

  • Limitations of LLM reasoning ability: Although language models can generate seemingly reasonable text responses, they often struggle with tasks requiring strict logical reasoning, such as mathematical reasoning or symbolic reasoning. These models frequently fail to converge or produce accurate and reliable answers, which limits their applications in real-world scenarios with high reasoning demands.
  • LLMs and preference data: The main methods of utilizing preference data are constructing reward models and directly applying preference data to update model strategies.
  • Shortcomings of existing methods to improve reasoning ability: Many current attempts to enhance LLM reasoning ability (e.g., using instance-level learning strategies or RL-based optimization) have limitations. These methods often cannot finely decompose reward signals, making it difficult to ensure the quality of intermediate reasoning steps. As a result, the use of data is inefficient, and there is still a lack of sufficient theoretical support.
  • Importance of iterative methods: Unlike traditional RLHF (Reinforcement Learning from Human Feedback), iterative methods emphasize dynamic and continuous improvement. They repeatedly collect and analyze data to update strategies, enabling models to adapt to the complexity of human decision-making and reasoning. The success of AlphaZero across different domains demonstrates the potential of iterative methods, especially when combined with search strategies like MCTS.
  • Challenges of applying MCTS in LLMs: Using MCTS to collect preference data for improving current strategies face challenges. These mainly involve determining the granularity of MCTS application (traditional instance-level preference data collection may lose information) and the reliance of MCTS on evaluative feedback (requiring meaningful feedback signals to guide strategy improvements).
  • MCTS: MCTS performs search using a tree structure. The root node is at the top of the tree, and nodes without children are called leaf nodes. Each node records visit counts and reward values. MCTS makes decisions by considering both visit counts and average rewards when selecting child nodes. MCTS is divided into four steps: Select, Expand, Simulate, and Backup.

Method

MCTS-Iterative Preference Learning mainly involves two steps:
  1. Used MCTS to collect step-level preference data
  1. Used an iterative DPO strategy to update the model.
In this work, MCTS estimates the state-action value ) to assign preferences. Steps with relatively higher or lower -values are labeled as positive and negative, respectively. These form preference data for subsequent DPO training. Then, the strategy is continuously updated through repeated iterations.
notion image
 
⭐

1.1 MCTS collects step-level preference data

  • Reasoning process decomposition and state representation: The reasoning process is decomposed into discrete steps, where each step is a sequence of tokens. The state is defined as the reasoning prefix up to that point. By adding a new reasoning step , the state transitions from to . In an LLM, corresponds to concatenating with .
      • Using the current LLM policy , candidate steps are sampled from its probability distribution , where is the task input prompt.
  • Select stage: The goal of the selection stage is to identify nodes that balance search quality and computational efficiency. Node selection uses the PUCT (Predictor + Upper Confidence bounds applied to Trees) formula:
    • : action value of taking action in state .
    • : number of visits to state .
    • : probability distribution over generating from the policy , where a length penalty is added to prevent excessively long reasoning chains.
  • Expand Stage: The expansion stage occurs when generating new nodes during the selection process, where new nodes are evaluated. The reward value of step is expressed as the difference in state rewards: , with the state reward defined as:
    • : correctness score of the result (1 = correct terminal state, -1 = incorrect terminal state, 0 = unfinished intermediate state).
    • : self-evaluation score, defined here as the probability of outputting a correctness-related token under the condition of eval prompt + task input prompt + state .
        • Let’s understand the self-evaluation formula with the code segment:
        In practice, the method takes the output token sequence and searches for tokens like , (answer options) or judgment words such as correct, incorrect, wrong, etc. It then finds the first such token, extracts its score, and adds up the probabilities of tokens representing correctness (from correct_token_ids). This sum serves as the confidence score (self-evaluation score).
        In summary, this calculates the model’s self-confidence in the reasoning step it outputs at the current state (e.g., whether it outputs “Option A correct” or “Option B wrong”).
  • Backup Stage: Backward Propagation Update
      • Once a terminal state is reached, the update propagates back from the terminal node up to the root node. The update rule is similar to RL definitions:
    • State-action value is updated through the next state’s value and the immediate reward .
    • State value is defined as the expectation (average) of all action values at that state.
    • Visit counts are incremented by 1.
⭐

1.2 Build MCTS Search Tree

For each step in the generated response, the paper constructs an MCTS search tree with K iterations.
  • To balance diversity, quality, and efficiency, the exploration constant is set to for the initial steps and to a smaller value for subsequent steps.
  • For a search tree with depth , the algorithm extracts step-level preference data:
    • First, it normalizes the visit counts of each child node relative to its parent’s total visit counts (representing both quality and diversity of generation).
    • Then, it selects the parent node with the largest normalized value.
    • Finally, among the parents’ child steps, the two steps with the highest and lowest Q-values are chosen as positive and negative preference samples.
Example:
  1. Build the MCTS tree: The root is the user’s question (prompt). Each level is one reasoning step. If we go 3 steps deep, the tree has 3 layers of nodes (Step 1 → Step 2 → Step 3).
  1. Evaluate with MCTS: Each node is expanded and simulated multiple times. Every child node records a visit count (how often MCTS chose it). A node with a higher visit count means MCTS thinks it’s a better step (higher chance of leading to a correct answer).
  1. Normalize the visit counts: For each parent, divide child visits by the parent’s total visits. This creates a distribution (like probabilities). Example:
      • Parent A → children visits [60, 20, 20] → normalized [0.6, 0.2, 0.2].
      • Parent B → children visits [34, 33, 33] → normalized [0.34, 0.33, 0.33].
  1. Pick the most “confident” parent:
      • Parent A’s distribution [0.6, 0.2, 0.2] is skewed → strong preference.
      • Parent B’s distribution [0.34, 0.33, 0.33] is flat → uncertain.
      So we prefer Parent A for training data.
  1. Extract preference samples: Within the chosen parent, compare children by Q-values (expected reward/accuracy). The highest-Q child is the positive sample. The lowest-Q child is the negative sample.
⭐

2. Iterative Preference Learning

The step-level preference data collected by MCTS is directly used to update strategies with DPO (see Section 3.5 DPO). Considering noise in preference labels, this work uses a conservative version of DPO and applies adaptive label smoothing based on MCTS visit counts:
  • In the -th iteration, represent the step-level preference data (chosen and rejected).
  • is the KL regularization scaling factor.
  • is the label smoothing variable, calculated from visit counts of the preference data in MCTS:
  • is the KL constraint part of the DPO loss.
After optimization, the updated strategy is obtained, and the process is repeated iteratively for LLM training.
Role of Conservative DPO and Smoothed Labels
  • Addressing noisy preference labels: Conservative DPO makes optimization more robust to label noise, preventing overfitting to noisy preference data and ensuring more reliable updates.
  • Improving generalization while avoiding overfitting: Label smoothing adjusts the contribution of each preference sample based on MCTS visit counts, reducing the risk of overfitting to extreme labels (e.g., “absolutely correct” or “absolutely wrong”). This helps improve generalization.
Algorithm workflow (red part has typo, where it should be and ):
notion image

Experiment

  • Datasets: Focused on arithmetic and commonsense reasoning tasks.
    • Arithmetic reasoning: GSM8K and MATH.
    • Commonsense reasoning: ARC, AI2Science, OpenBookQA (OBQA), CommonSenseQA (CSQA), and others.
  • Baselines: Compared with STaR, Crystal, Direct Tuning, LMSI, Math-Shepherd, and also with variants like instance-level iterative preference learning and offline MCTS. Base model: Mistral-7B.
🔬
Results
  • Arithmetic reasoning:
      • On GSM8K and MATH, this method significantly outperformed baselines:
    • GSM8K accuracy: improved from 75.9% to 81.8%.
    • MATH accuracy: improved from 28.9% to 34.7%.
      • Compared with Math-Shepherd (which uses a similar iterative process), this method achieved comparable performance without training a separate reward/value network, and even further improved by combining MCTS with ground-truth answers.
      notion image
  • Commonsense reasoning: Also showed improvements on tasks like ARC-Challenge (ARC-C), AI2Science-Middle (AI2Sci-M), and SciQ, accuracy increased by +2.5%, +3.0%, and +2.1% respectively, over SFT. However, on OBQA and CSQA, performance was relatively lower compared to SFT. This was likely because the base model lacked specific domain knowledge, and uncertainty in the intermediate reasoning chain led to error propagation.
      • notion image
🔬
In some task cases, such as OBQA and CSQA, the performance was lower compared to SFT. This may be due to the base model lacking domain-specific knowledge and the uncertainty introduced by intermediate reasoning chains leading to prediction errors. However, on tasks like MATH, which require long reasoning chains, the method outperformed SFT.
notion image
🔬
Example of a search tree for a science problem:
Problem: Methane gas (CH₄) reacts with oxygen. The unbalanced equation is shown below: CH₄ + x O₂ → 2CO₂ + 4H₂O. How many O₂ molecules are needed to correctly balance this equation? Options: (A) 1, (B) 2, (C) 3, (D) 4. The correct answer is (D) 4.
Search setup:
  • Search breadth parameters: = 4, = 2.
  • Each sequence starts with a number that indicates the visit count of the node.
  • The sequence ends with Q (state-action value) and P (probability), which represent the node’s expected value and probability score.
notion image

Conclusion

This paper proposes an MCTS-enhanced iterative preference learning method, where MCTS is used as a strategy improvement operator. By collecting step-level preference data, this method strengthens LLM alignment. Theoretical analysis shows that online sampling is crucial for improving LLM strategies. Future work could focus on improving both data and algorithms to address performance fluctuation issues in online learning frameworks, such as:
  • Better search strategies,
  • More effective use of historical data and strategies,
  • Balancing offline and online learning,
 
 
Prev
CivitAI’s Payment Issue
Next
Awesome-AI-Tutorials
Loading...
Article List
LLM Learning Roadmap
✨ Awesome-Anything
Lazy loaded image Awesome-AI-Tutorials
Lazy loaded image Awesome-AI-Researchs
Lazy loaded image Awesome-GitHub-Repos
Lazy loaded image Awesome-AI-Examples
🖼️ Digital Image Processing
Lazy loaded image DIP I Image Demosaicing and Histogram Manipulation
Lazy loaded image DIP II Image Denoising
Lazy loaded image DIP III Edge Detection
Lazy loaded image DIP IV Digital Half-toning
Lazy loaded image DIP V Geometric Image Modification
🍃 LLM Components
Lazy loaded image 1.1 Tokenization
Lazy loaded image 1.2 Embedding
Lazy loaded image 1.3 Attention
Lazy loaded image 1.4 FFN, Residual Addition, and LN
Lazy loaded image 1.5 Positional Encoding
Lazy loaded image 1.6 Architecture & Decoding Policy
🌱 LLM Pre-training
Lazy loaded image 2.1 Pre-training Definition
Lazy loaded image 2.2 Pre-training Data
Lazy loaded image 2.3 Pre-training Workflow
Lazy loaded image 2.4 Continued Pre-training (CPT)
Lazy loaded image 2.5 Pre-training Evaluation
☘️ LLM Post-Training
Lazy loaded image 3.1 Supervised Fine-Tuning (SFT)
Lazy loaded image 3.2 Reinforcement Learning
Lazy loaded image 3.3 RLHF: Reinforcement Learning from Human Feedback
Lazy loaded image 3.4 Optimizing PPO-Based Algorithms
Lazy loaded image 3.5 Direct Preference Optimization (DPO)
🍀 LLM Popular Models
Lazy loaded image 4.1 BERT & Variants
Lazy loaded image 4.2 PaLM Series
Lazy loaded image 4.3 GPT Series
Lazy loaded image 4.4 LLaMA Series
Lazy loaded image 4.5 GLM Series
Lazy loaded image 4.6 Qwen Series
Lazy loaded image 4.7 DeepSeek Series
Lazy loaded image 4.8 MOE Series
Lazy loaded image 4.9 Other
🪴 LLM Applications
Lazy loaded image 5.1 Prompt Engineering
Lazy loaded image 5.2 LLM-based Agents
Lazy loaded image 5.3 Retrieval-Augmented Generation (RAG)
Lazy loaded image 5.4 Enterprise-level RAG Practice (I)
Lazy loaded image 5.5 Enterprise-level RAG Practice (II)
🌿 LLM Optimization
Lazy loaded image 6.1 GPU Usage Analysis
Lazy loaded image 6.2 Latency Analysis for Inference
Lazy loaded image 6.3 FlashAttention
Lazy loaded image 6.4 PagedAttention
Lazy loaded image 6.5 Packing for Efficient Training
Lazy loaded image 6.6 PEFT
🌾 LLM Compression
Lazy loaded image 7.1 Quantization
Lazy loaded image 7.2 Pruning
Lazy loaded image 7.3 Knowledge Distillation
🌵 LLM Hands-on Practice
Lazy loaded image 8.1 nanoGPT: Training a GPT from Scratch
Lazy loaded image 8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
Lazy loaded image 8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
Lazy loaded image 8.4 OpenRLHF: Reward Model Training and Code Analysis
Lazy loaded image 8.5 OpenRLHF: PPO Tutorial and Code Analysis
🌴 LLM Must-read Papers
Lazy loaded image 9.1 Technical Reports Explained
Lazy loaded image 9.2 MedLLM
Lazy loaded image 9.3 o1 Techinical Paper
Lazy loaded image 9.4 Prompting
Lazy loaded image 9.5 LLM-as-Judge
🌳 LLM Q&A
Lazy loaded image 10.1 Why GRPO training loss starts at 0
Lazy loaded image 10.2 VLM-R1 Reproduction Results
Lazy loaded image 10.3 KL(P||Q) or KL(Q||P)?
Lazy loaded image 10.4 Left Padding or Right Padding?
Lazy loaded image 10.5 What is QK-Norm?
Lazy loaded image 10.6 Why LLM training rarely adopts Dropout?
Lazy loaded image 10.7 Dropout and DropPath
🐝 VLM Image Encoders
Lazy loaded image 11.0 Introduction
Lazy loaded image 11.1 Vision Transformer
Lazy loaded image 11.2 CLIP
Lazy loaded image 11.3 BLIP
Lazy loaded image 11.4 EVA
Lazy loaded image 11.6 BLIP-2
Lazy loaded image 11.7 SigLIP
📝 MISC.
🤔 Large Language Models Pass the Turing Test.
Lazy loaded image OpenAI ChatGPT Progression
⏳ Claude Sonnet 4 Test with Sandtris
🔍 What DeepSeek-R1's Surge Teaches Us About SEO
💰 CivitAI’s Payment Issue
🎓 arXivSub, an Academia Assistant