PaperlayerView exact revision on arXiv

Source-linked AI summary

Interpretable Contrastive Monte Carlo Tree Search Reasoning

Zitian Gao, Boye Niu, Xuzheng He, Haotian Xu, Hongzhang Liu, Aiwei Liu, Xuming Hu, Lijie Wen

arXiv:2410.01707v3cs.CLcs.AI

TL;DR

LLM MCTS reasoning faces limited reward-model analysis and lower speed than CoT. SC-MCTS* addresses these issues with interpretable contrastive rewards, improved search components, and speculative decoding, outperforming o1-mini by 17.4% on average on Blocksworld with Llama-3.1-70B.

  • Problem

    MCTS reasoning depends heavily on reward models and is substantially slower than CoT, while prior work provides limited component-level quantitative and interpretability analysis.

  • Method

    SC-MCTS* combines contrastive-decoding reward models, statistically informed reward combination, refined UCT selection and backpropagation, and speculative decoding.

  • Results

    17.4% average improvement over OpenAI’s o1-mini was achieved using Llama-3.1-70B on the Blocksworld dataset.

  • Takeaways & Limitations

    SC-MCTS* integrates interpretability, efficiency, and generalizability for multi-step LLM reasoning within the evaluated Blocksworld setting.

  • Takeaways & Limitations

    The Multi-RM method focuses on clearly distinct reward modes, leaving general cases for future work.

Abstract

from arXiv · show

We propose SC-MCTS*: a novel Monte Carlo Tree Search (MCTS) reasoning algorithm for Large Language Models (LLMs), significantly improves both reasoning accuracy and speed. Our motivation comes from: 1. Previous MCTS LLM reasoning works often overlooked its biggest drawback--slower speed compared to CoT; 2. Previous research mainly used MCTS as a tool for LLM reasoning on various tasks with limited quantitative analysis or ablation studies of its components from reasoning interpretability perspective. 3. The reward model is the most crucial component in MCTS, however previous work has rarely conducted in-depth study or improvement of MCTS's reward models. Thus, we conducted extensive ablation studies and quantitative analysis on components of MCTS, revealing the impact of each component on the MCTS reasoning performance of LLMs. Building on this, (i) we designed a highly interpretable reward model based on the principle of contrastive decoding and (ii) achieved an average speed improvement of 51.9% per node using speculative decoding. Additionally, (iii) we improved UCT node selection strategy and backpropagation used in previous works, resulting in significant performance improvement. We outperformed o1-mini by an average of 17.4% on the Blocksworld multi-step reasoning dataset using Llama-3.1-70B with SC-MCTS*. Our code is available at https://github.com/zitian-gao/SC-MCTS.

1 INTRODUCTION

SC-MCTS* targets two MCTS reasoning challenges: reward models that are difficult to design and MCTS’s slower speed than Chain of Thought. It combines interpretable reward modeling with component improvements and speculative decoding, outperforming o1-mini by 17.4% on average using Llama-3.1-70B on Blocksworld.

  • Reward-model quality largely determines MCTS reasoning ability, while prior approaches incur extra models, training, tools, or dataset dependencies.
  • MCTS reasoning is slower than CoT because multi-layer search trees and repeated iterations require many LLM node visits.MCTS trees typically have 2–10 layers, and better performance commonly requires 2–10 iterations.
  • SC-MCTS* redesigns rewards with contrastive decoding and improves UCT selection and backpropagation to favor steadily improving paths.
  • Speculative decoding speeds MCTS reasoning by an average of 52%, while extensive ablations analyze component effects and reward interpretability.
  • 17.4% average improvement over OpenAI’s o1-mini was achieved using Llama-3.1-70B on the Blocksworld multi-step reasoning dataset.

2 RELATED WORK

Prior work applies MCTS and decoding strategies to LLM reasoning, but SC-MCTS* combines action-level contrastive decoding with speculative decoding for MCTS guidance. This differs from token-level speculative contrastive decoding and targets multi-step reasoning.

  • Chain-of-Thought reasoning helps LLMs decompose complex tasks into sequences of manageable steps.
  • SC-MCTS* combines contrastive and speculative decoding without additional cost, using action-level rather than token-level contrastive decoding to guide MCTS.

3 PRELIMINARIES

The paper models multi-step reasoning as sequential decision-making and introduces MCTS as a four-phase search procedure. It then reviews contrastive and speculative decoding, which support action-level rewards and accelerated reasoning.

  • 3.2 MONTE CARLO TREE SEARCH: MCTS iteratively selects, expands, simulates, and backpropagates through a search tree to improve reasoning through trials and heuristics.Selection favors promising and novel nodes; expansion samples actions, simulation rolls out trajectories, and backpropagation updates visited-node values.
  • 3.3 CONTRASTIVE DECODING: Contrastive decoding improves generation by contrasting expert and amateur model distributions while filtering implausible tokens with an adaptive plausibility constraint.Its logits combine expert and amateur information, and the α-mask defines a valid vocabulary before generation.
  • 3.3 CONTRASTIVE DECODING: The proposed action-level contrastive decoding averages confidence differences across an entire reasoning action rather than individual tokens.This action-level reward is intended to capture expert–amateur confidence differences more robustly than vanilla token-level contrastive decoding.
  • 3.4 SPECULATIVE DECODING AS "FREE LUNCH": Speculative decoding generates candidate tokens with a smaller model and accepts those consistent with the target model, resampling rejected tokens from an adjusted distribution.The method uses a target distribution p and an approximation distribution q to accelerate generation while preserving target-model consistency.
  • 3.4 SPECULATIVE DECODING AS "FREE LUNCH": Because contrastive and speculative decoding use the same smaller models, the paper combines them to obtain speculative-decoding acceleration without additional model requirements.The paper characterizes this acceleration as a “free lunch.”

4 METHOD

SC-MCTS* combines interpretable reward models with revised reward normalization, node selection, and backpropagation for MCTS reasoning. Its reward design uses action-level contrastive divergence alongside loglikelihood and self-evaluation signals, while its search updates address reward-scale mismatch, UCT sensitivity, and path progression.

  • 4.1 MULTI-REWARD DESIGN: SC-MCTS* combines action-level contrastive JS divergence, loglikelihood, and self-evaluation rewards through statistically informed normalization.Reward factors are normalized using fine-grained prior statistics, with online updates intended to address distribution shift during reasoning.
  • 4.1 MULTI-REWARD DESIGN: The proposed contrastive reward computes Jensen-Shannon divergence between expert and amateur logits over complete action-token sequences rather than individual tokens.This action-level formulation captures model behavior across the entire sequence at once, contrasting with serial token-level methods.
  • 4.1 MULTI-REWARD DESIGN: Multi-RM clusters reward distributions into regions and normalizes each region separately before combining reward models.The method uses region-specific means and standard deviations; standard normalization is recovered when only one region exists.
  • 4.1 MULTI-REWARD DESIGN: The Multi-RM design is limited to cases with clearly distinct reward modes and updates mode statistics online to mitigate distribution shift.The authors leave general cases without clearly distinct modes for future work.
  • 4.2 NODE SELECTION STRATEGY: UCT selects actions by balancing average reward with an exploration term controlled by C, whose value requires quantitative investigation because default settings may fail.The selection formula uses parent and child visit counts together with the exploration constant C.
  • 4.3 BACKPROPAGATION: The revised backpropagation values paths by their incremental progress, clipping and downweighting negative changes while penalizing longer paths.This favors smoothly improving paths that approach the final goal over less stable paths.

5 EXPERIMENTS

Experiments evaluate SC-MCTS* on Blocksworld across accuracy, speed, parameter sensitivity, ablations, and reward interpretability. SC-MCTS* achieves strong accuracy gains, speculative decoding accelerates node reasoning, and its reward values align with goal progress.

  • 5.2 MAIN RESULTS: SC-MCTS* significantly outperforms RAP-MCTS and 4-shot CoT across easy and hard modes, and Llama-3.1-70B with SC-MCTS* surpasses 4-shot CoT using Llama-3.1-405B in easy mode.The comparison uses Blocksworld across all steps and difficulty levels, with Llama-3-70B and Llama-3.1-70B for the evaluated MCTS methods.
  • 5.3 REASONING SPEED: Approximately 100% speedup is achieved by pairing Llama-3.1-405B with Llama-3.1-8B, while Llama-3.1-70B with Llama-3.2-1B increases reasoning speed by 51.9% over vanilla decoding.The results show that speculative decoding with small language models can substantially improve average node-level MCTS reasoning speed, although overly small draft models may reduce speed.
  • 5.4 PARAMETERS: The UCT constant C affects exploration differently across methods: RAP-MCTS uses a value too small to function effectively, whereas SC-MCTS* has a better-suited value for its reward scale.The experiment isolates C by using MCTS base with the common RLL reward model for both methods.
  • 5.4 PARAMETERS: Accuracy increases steadily through the first seven MCTS iterations, after which additional iterations yield relatively smaller improvements under the depth-limited setting.The later diminishing returns are attributed to exponentially growing exploration nodes in later iterations.
  • 5.5 ABLATION STUDY: 80.92% accuracy, a 25% improvement over base MCTS, results from combined reward-model, UCT, and backpropagation enhancements in the difficult Step 6 setting.The ablation begins at 55.92%; RJSD, RLL, and RSE add 14.47%, Multi-RM adds 3.29%, UCT optimization adds 5.27%, and backpropagation refinement adds 1.97%.
  • 5.6 INTERPRETABILITY STUDY: SC-MCTS* reward values correlate significantly with action progress ∆a, and reward-bin colors consistently track the proportion of positive ∆a.The alignment indicates that the reward model captures progress toward the goal and provides interpretable signals for MCTS action selection.

6 CONCLUSION

SC-MCTS* improves LLM multi-step reasoning by combining enhanced reward modeling, node selection, and backpropagation with speed-oriented decoding. On Blocksworld, it outperforms o1-mini by 17.4% on average using Llama-3.1-70B.

  • 17.4% average improvement over o1-mini is achieved on the Blocksworld dataset using Llama-3.1-70B with SC-MCTS*.
  • SC-MCTS* combines improvements in reward modeling, node selection strategy, and backpropagation to enhance LLM reasoning accuracy and speed.
  • Action-level variables aggregate all tokens in a reasoning step, whereas token-level variables operate at a lower decoding granularity.
  • Traditional sequence-aggregated logit-difference contrastive rewards are less stable than JS-divergence rewards because logit differences are unbounded.
  • Prior MCTS reasoning systems used external verification tools, self-training, critic models, or mutual consistency to construct reward signals.
  • Speculative decoding accelerates autoregressive generation by computing multiple tokens in parallel without retraining or changing model structure.

C REWARD FUNCTIONS CORRELATION

The reward-correlation analysis supports combining reward functions, while the SC-MCTS* procedure uses contrastive reward evaluation, online statistics updates, UCT selection, expansion, and greedy rollout.

  • Reward functions have pairwise correlations below 0.15 in absolute value, supporting their combination in the Multi-RM method.
  • SC-MCTS* first obtains prior reward-value data, then uses two SLMs for contrastive rewards and speculative-decoding speedup during MCTS reasoning.
  • Node selection chooses a child using the UCT value combining child value with an exploration term.
  • The algorithm evaluates each expanded child with contrastive rewards and updates reward-factor statistics before selecting the highest-valued child for rollout.
  • Online incremental updates adjust reward means and standard deviations because normalization can encounter distribution shift during reasoning.
  • Reward factors are assigned to clusters before their values and statistics are updated.

E EXPERIMENTAL SETTINGS

Experiments use quantized checkpoints and completion-mode inference, with checkpoint and language-model hyperparameter details provided in the accompanying tables.

  • 4-bit quantized checkpoints reduce memory usage and improve inference speed while causing around 2% performance loss.
  • The experiments use completion mode to format outputs into single reasoning steps that can be converted into MCTS nodes.
  • Table 3 lists the checkpoints used in the experiments and their repository links.
  • Table 4 reports the LLM hyperparameters and EOS tokens used in the experiments.

F BLOCKSWORLD DATASET

The Blocksworld dataset contains 600 planning instances spanning simple to complex configurations with varying block counts, goals, obstacles, and plan lengths.

  • 600 Blocksworld instances cover planning problems ranging from 3–5 blocks to cases with up to 25 blocks.
  • More complex instances introduce additional goals and obstacles, creating a range of problem difficulties for evaluating planning algorithms.

F.1 DIFFICULTY SETTINGS

The paper divides Blocksworld into Easy and Hard settings with different demonstration-selection procedures, then examines how UCT parameter choices affect exploration.

  • Difficulty Settings: The Blocksworld dataset is divided into Easy and Hard settings from the original 600 instances.
  • Easy Blocksworld: Easy setting demonstrations match each test problem’s required minimum step count, with 10 randomly selected cases per group.
  • Hard Blocksworld: Hard setting demonstrations are sampled globally without matching the test problem’s minimum action count, leaving 590 test cases.
  • UCT Parameter Analysis: The optimized UCT parameter makes node selection prioritize unexplored nodes over repeatedly following old paths that may lead to dead ends.

H OPENAI API DATA

The paper reports resource usage for its main experiments and identifies broader reward-model integration and cross-dataset step-splitting as future directions.

  • Resource Usage: SC-MCTS∗ uses approximately 280 GPU hours with Llama-3 70B and 300 GPU hours with Llama-3.1 70B across difficulty modes.
  • Resource Usage: RAP-MCTS uses approximately 420 GPU hours with Llama-3 70B and 450 GPU hours with Llama-3.1 70B across difficulty modes.
  • Future Work: Future work includes combining metrics-based reward models with LM-based reward models such as Critic LLM and Eurus.
  • Future Work: The paper identifies more general step-splitting methods for other tasks and datasets as an open direction for improving MCTS reasoning generalization.
Loading 2410.01707v3…