Monte Carlo Tree Diffusion for System 2 Planning

2025-02-11

AI Summary

Diffusion models exhibit strong generative capabilities for planning tasks but suffer from non-scalable test-time computation (TTC): performance plateaus rather than improving monotonically with increased inference budget. To address this, we propose Diffusion-MCTS—the first framework integrating Monte Carlo Tree Search (MCTS) into the diffusion paradigm. It reformulates the denoising process as a tree-structured search, incorporating value-guided node selection, conditional resampling, and backtracking from suboptimal branches to enable iterative evaluation, pruning, and refinement. This design supports dynamic exploration-exploitation trade-offs, overcoming the fundamental TTC bottleneck of conventional diffusion-based planners. Empirically, on long-horizon planning tasks, solution quality improves monotonically with computational budget—significantly outperforming diffusion baselines—and demonstrates both TTC scalability and robustness.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Planning with Language Models
• Multiagent Systems: Multiagent Planning

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Diffusion models have recently emerged as a powerful tool for planning. However, unlike Monte Carlo Tree Search (MCTS)-whose performance naturally improves with additional test-time computation (TTC), standard diffusion-based planners offer only limited avenues for TTC scalability. In this paper, we introduce Monte Carlo Tree Diffusion (MCTD), a novel framework that integrates the generative strength of diffusion models with the adaptive search capabilities of MCTS. Our method reconceptualizes denoising as a tree-structured process, allowing partially denoised plans to be iteratively evaluated, pruned, and refined. By selectively expanding promising trajectories while retaining the flexibility to revisit and improve suboptimal branches, MCTD achieves the benefits of MCTS such as controlling exploration-exploitation trade-offs within the diffusion framework. Empirical results on challenging long-horizon tasks show that MCTD outperforms diffusion baselines, yielding higher-quality solutions as TTC increases.

Problem

• Enhance diffusion-based planning scalability
• Integrate MCTS with diffusion models
• Improve long-horizon task solution quality

Innovation

• Integrates diffusion models with MCTS
• Reconceptualizes denoising as tree-structured
• Selectively expands promising trajectories