Interpretable Contrastive Monte Carlo Tree Search Reasoning

2024-10-02

arXiv.org

AI Summary

To address the slow inference speed, low accuracy, and poor interpretability of reward models in Monte Carlo Tree Search (MCTS) for multi-step reasoning with large language models (LLMs), this paper proposes SC-MCTS*, an interpretable contrastive MCTS reasoning framework tailored for LLMs. Our key contributions are: (1) a novel contrastive-decoding-based interpretable reward model; (2) an enhanced UCT policy and backpropagation mechanism integrated with speculative decoding, accelerating node expansion by 51.9%; and (3) the first systematic, quantitative interpretability analysis and ablation study of individual MCTS components. Evaluated on the Blocksworld dataset using Llama-3.1-70B, SC-MCTS* achieves a 17.4% average accuracy improvement over o1-mini. The implementation is publicly available.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Multiagent Systems: Multiagent Learning

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

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.

Problem

• Monte Carlo Tree Search (MCTS)
• Language Model Inference
• Reward Model Enhancement

Innovation

• SC-MCTS*
• Contrastive Decoding
• Predictive Decoding