Fast Monte Carlo Tree Diffusion: 100× Speedup via Parallel and Sparse Planning

Jaesik Yoon, Hyeonseo Cho, Yoshua Bengio, Sungjin Ahn

Keywords: Diffusion, MCTS, Long-term Planning, Offline RL, Goal-conditioned RL, Efficient Planning, Reasoning, Inference-Time Scaling

TL;DR: Fast Monte Carlo Tree Diffusion (Fast-MCTD) achieves up to 100× speedup over MCTD through parallel rollouts and sparse trajectory planning, maintaining strong performance in complex long-horizon tasks while being computationally efficient.

Abstract: Diffusion models have recently emerged as a powerful approach for trajectory planning. However, their inherently non-sequential nature limits their effectiveness in long-horizon reasoning tasks at test time. The recently proposed Monte Carlo Tree Diffusion (MCTD) offers a promising solution by combining diffusion with tree-based search, achieving state-of-the-art performance on complex planning problems. Despite its strengths, our analysis shows that MCTD incurs substantial computational overhead due to the sequential nature of tree search and the cost of iterative denoising. To address this, we propose Fast-MCTD, a more efficient variant that preserves the strengths of MCTD while significantly improving its speed and scalability. Fast-MCTD integrates two techniques: Parallel MCTD, which enables parallel rollouts via delayed tree updates and redundancy-aware selection; and Sparse MCTD, which reduces rollout length through trajectory coarsening. Experiments show that Fast-MCTD achieves up to 100× speedup over standard MCTD while maintaining or improving planning performance. Remarkably, it even outperforms Diffuser in inference speed on some tasks, despite Diffuser requiring no search and yielding weaker solutions. These results position Fast-MCTD as a practical and scalable solution for diffusion-based inference-time reasoning.

Primary Area: Reinforcement learning (e.g., decision and control, planning, hierarchical RL, robotics)

Submission Number: 19331