Monte Carlo Tree Diffusion — System 2 Planning with Diffusion Models

Paper. Jaesik Yoon, Hyeonseo Cho, Doojin Baek, Yoshua Bengio, Sungjin Ahn. Monte Carlo Tree Diffusion for System 2 Planning. ICML 2025 Spotlight. KAIST · SAP · Mila · NYU. [arXiv] · [project page] · [code]

0. The Picture in One Paragraph

Diffusion models have emerged as capable offline planners — they can generate whole trajectories toward a goal by running a denoising process over a trajectory representation. But they are fundamentally one-shot samplers: giving a diffusion planner more inference-time compute (more samples, more steps) doesn’t reliably improve outcomes. Monte Carlo Tree Search (MCTS) has the opposite property — it was built to improve with compute by systematically exploring a search tree. MCTD bridges the two by reconceptualizing the denoising process as a tree-structured rollout. Partial denoising states become tree nodes; binary guidance schedules (guided vs. unguided sampling) become meta-actions that MCTS selects among; fast DDIM jumps stand in for cheap rollout simulations; and standard UCT backpropagation updates node values after each simulated plan is scored by a reward model. The result is a diffusion planner that can actually use extra compute — performance scales monotonically with the number of tree iterations, and on long-horizon maze tasks where all diffusion baselines collapse, MCTD achieves near-perfect success.

1. The Problem — Diffusion Planners Can’t Scale with Compute

Offline diffusion planners (Diffuser and descendants) are trained on trajectory datasets to denoise random noise into a complete trajectory from current state to goal. This is powerful: you get behavioral diversity for free, planning is just sampling, and the whole thing is tractable to train.

The weakness appears on hard long-horizon tasks. There are two obvious ways to give a diffusion planner more inference-time compute:

Strategy	What it does	Does it help?
More denoising steps	Finer denoising trajectory	Marginal or no gain
Random search	Sample N trajectories, pick best	No systematic improvement
Replanning	Re-run the planner periodically	Minor; not well-coordinated

None of these strategies are adaptive — they don’t use the result of one computation to guide the next. In MCTS terms, they explore blindly rather than accumulating evidence about which branches are promising. For easy tasks this doesn’t matter (the single shot works). For hard tasks — long-horizon navigation through a complex maze — it means performance plateaus at the quality of the first sample.

The “System 2” framing in the title is precise: System 1 thinking is fast and intuitive (single-shot diffusion), System 2 is slow and deliberate (MCTS-guided search). Diffusion planners have been stuck in System 1 mode. MCTD adds the System 2 loop.

2. The Move — Instantiate MCTS Inside the Denoising Process

MCTS is defined by four steps: Select → Expand → Simulate → Backpropagate. MCTD maps each step onto something the diffusion planner can do natively.

The key insight: the trajectory is divided into a sequence of temporally contiguous subplans. A partially denoised state — where some subplans have been denoised and others remain noisy — is a natural tree node. Children of a node are different ways to continue denoising (from the same partial state).

3. The Three Core Mechanisms

3.1 Subplan decomposition → tree nodes

Divide the full planned trajectory into $K$ subplan segments in time. Each tree node corresponds to a specific combination of (subplan index, partial denoising state). Deeper nodes have more subplans resolved; leaf nodes have a fully denoised trajectory.

This decomposition is what makes tree branching meaningful: you can branch on which future you commit to for the first subplan, hold that fixed, and then explore alternatives for the second subplan. Without decomposition, branching mid-denoising would produce incoherent trajectories.

3.2 Guidance levels → meta-actions

At each expansion, MCTD must decide how to continue denoising the next subplan. The available meta-actions are:

Guided (1): denoise under the goal-conditioned guidance signal → exploitation, pulls toward the known goal.
Unguided (0): denoise from the unconditional prior → exploration, diversifies the candidate futures.

These guidance levels form the action space MCTS selects among. Node selection follows standard UCT:

\[a^* = \arg\max_a \left[ Q(s, a) + C \sqrt{\frac{\ln N(s)}{N(s,a)}} \right],\]

where $Q(s, a)$ is the mean reward of subtrees reached via action $a$ from node $s$, $N(s)$ is the visit count of $s$, and $C$ is the exploration constant. High UCT value = promising branch with remaining uncertainty.

The binary guidance meta-action is the minimum useful abstraction: it’s coarse enough for MCTS to reason over efficiently, but it produces meaningfully different trajectories (exploitation vs. exploration) at each step. The full guidance schedule — $(g_1, g_2, \ldots, g_K)$ — is what MCTS is really searching over.

3.3 Jumpy denoising → simulation oracle

After expanding a node, MCTD needs to evaluate it — to score “how good is this partial plan likely to be?” Running the full denoising chain to completion is too slow to do at every node. Instead, MCTD uses DDIM (which allows arbitrary step-size skips) to jump from the current partial denoising state to a fully denoised trajectory in a fraction of the steps.

This “jumpy” simulation is an approximation, but it’s cheap and good enough for backpropagation — the reward model evaluates the fast-simulated complete trajectory and returns a scalar score.

3.4 Backpropagation

After simulation, the reward is backpropagated up the tree to update $Q(s, a)$ estimates for all ancestor nodes. The guidance schedule along the winning path accumulates positive weight; schedules along low-reward paths are down-weighted. Future selection iterations draw on this accumulated evidence.

This is standard MCTS backpropagation — the contribution is that the tree structure and action space are well-defined enough for it to work inside a diffusion process.

4. What MCTD Actually Buys

The argument for the method can be stated in one sentence:

Diffusion planners were failing on hard tasks because good plans require coordinated decisions across subplan segments — decisions that can’t be made well in a single shot. MCTD gives the planner a way to try a commitment for the first subplan, evaluate consequences via fast rollout, and revise if the consequences look bad — the same feedback loop MCTS uses for any sequential decision problem.

The analogy to chess engines is exact: a single-shot diffusion sample is like playing the first move that looks locally good. MCTD is like running a minimax search to see how the game evolves before committing.

5. Experiments

5.1 Benchmarks

MCTD is evaluated on OGBench (Offline Goal-Conditioned RL Benchmark) across:

PointMaze — navigation for a point-mass robot in mazes: Medium, Large, Giant.
AntMaze — same layouts with a complex ant robot (higher DOF).
Visual PointMaze — same task but with RGB image observations instead of state, testing robustness to partial observability.

The backbone model is Diffuser (Janner et al., 2022) — a transformer diffusion model trained offline on trajectory datasets. The MCTS layer is inference-time only; no additional training is needed.

5.2 Baselines

Baseline	Description
Diffuser	Single-shot denoising plan
Diffuser-Replanning	Re-runs Diffuser at fixed intervals
Diffuser-Random Search	Samples N random trajectories, picks best by reward
Diffusion Forcing	Autoregressive sequential diffusion model

Diffuser-Random Search is the most important baseline: it uses the same total denoising budget as MCTD but allocates it as parallel random samples rather than a guided tree search. Its performance determines whether the gain is from more compute or better use of compute.

5.3 Results

PointMaze and AntMaze:

All standard diffusion baselines (Diffuser, Diffuser-Replanning, Diffusion Forcing) suffer dramatic performance drops as maze size increases from Medium → Large → Giant. Diffuser-Random Search performs no better than single-shot Diffuser, confirming that the issue is not compute quantity but compute quality.

MCTD achieves near-perfect success across all maze sizes and robot morphologies. On PointMaze-Giant — the hardest task — MCTD reaches ~100% success while all baselines fail.

Inference-time scaling:

MCTD success rate increases monotonically with the number of tree iterations. This is the key property: unlike every diffusion baseline, MCTD has a genuine “use more compute → get better plans” lever. This is what the paper means by System 2 planning.

Visual PointMaze:

MCTD outperforms both Diffuser and Diffusion Forcing on the image-observation variant, showing that the tree-search structure adds robustness even when the input is partial/visual.

5.4 Planning time cost

The compute cost is real. Vanilla MCTD requires 8–40× more wall-clock time than a single-shot Diffuser, depending on maze size. On PointMaze-Giant, this reaches ~264 seconds per plan.

This is addressed by a follow-up (Fast-MCTD, NeurIPS 2025) that introduces parallel tree rollouts with delayed updates and trajectory coarsening, recovering ~100× speedup while maintaining 100% success. That’s a downstream engineering result, not in this paper, but it shows the approach is practically viable.

6. The Taxonomy Connection

MCTD fits into the growing literature on inference-time scaling for decision-making:

Paper/system	How it uses extra compute
Best-of-N sampling	Random independent draws, pick max-reward
Process Reward Models	Score intermediate steps, rerank
Tree-of-Thought / MCTS for LLMs	Token-level tree search with language priors
MCTD	Tree search inside the denoising trajectory

The distinctive feature of MCTD is that the “tree” is not over language tokens or high-level plans — it’s over partially denoised trajectory states. This is a lower level of abstraction, which makes each node geometrically meaningful and enables the fast-DDIM simulation trick.

7. Why It Matters

Three reasons:

It breaks the “one-shot” ceiling on diffusion planners. Until MCTD, running a diffusion planner with more compute didn’t help beyond a point. The inference-time scaling property is a new lever that diffusion-based planning didn’t have.
The MCTS-to-diffusion mapping is clean and reusable. The three mechanisms — subplan nodes, guidance meta-actions, jumpy simulation — are not specific to maze navigation. Any offline diffusion planner with a reward model can potentially apply the same wrapper.
It validates the System 2 framing for generative planners. The pattern — slow deliberate search over a space of possible futures, guided by evaluation of fast cheap simulations — is the same pattern AlphaGo used for Go and LLM tree-of-thought methods use for reasoning. MCTD is the trajectory-generation version of that pattern.

8. Limitations Worth Knowing

Reward model dependency. Simulation scoring depends on a pre-trained reward model. If the reward model is inaccurate, backpropagation corrupts the search tree. There’s no mechanism to detect or correct for reward-model errors.
Offline training only. The diffusion backbone is trained on a fixed offline dataset; MCTD does not interact with the environment during search. Distributional shift between the offline data and the actual environment is not addressed.
Coarse guidance action space. Binary guided/unguided is the minimum meaningful action space. More expressive guidance schedules (continuous guidance strength, per-timestep variation) might improve performance but would also make the MCTS branching factor harder to manage.
Approximation error in jumpy denoising. DDIM simulation introduces approximation error relative to full denoising. This is acceptable in practice but can cause mispricing of nodes in subtler tasks.
No online replanning integration. MCTD plans a full trajectory offline before execution begins. There’s no mechanism for closed-loop real-time replanning as the environment evolves (though this is an obvious future extension).

9. The Takeaway for a First Reader

If you remember three things:

Diffusion planners are one-shot samplers. More compute (random sampling, replanning) doesn’t help because there’s no adaptive feedback loop. MCTD adds one by embedding MCTS inside the denoising process: partially denoised subplans are tree nodes, binary guidance schedules are meta-actions, and fast DDIM jumps are the cheap simulation oracle.
UCT governs node selection, backpropagation updates value estimates from simulated plan rewards, and performance scales monotonically with inference-time compute — the defining property of System 2 planning that all diffusion baselines lack.
On hard long-horizon maze tasks, MCTD achieves ~100% success where all diffusion baselines fail. Diffuser-Random Search — the same compute budget, allocated randomly — shows no improvement over single-shot Diffuser, confirming the gain is from structured search, not more samples.

That’s the arc: diffusion planners lack inference-time scaling → MCTS has it by construction → embed MCTS inside denoising via subplan nodes + guidance actions + jumpy simulation → unlock systematic compute scaling for diffusion-based planning.

References

Yoon, J., Cho, H., Baek, D., Bengio, Y., & Ahn, S. (2025). Monte Carlo Tree Diffusion for System 2 Planning. ICML 2025 Spotlight. arXiv:2502.07202.
Project page: https://sites.google.com/view/mctd-s2planning/home
Code: https://github.com/ahn-ml/mctd
Janner, M., Du, J., Tenenbaum, J., & Levine, S. (2022). Planning with Diffusion Models. NeurIPS 2022. (Diffuser backbone.)
Song, J., Meng, C., & Ermon, S. (2021). Denoising Diffusion Implicit Models. ICLR 2021. (DDIM simulation oracle.)
Fast-MCTD: arXiv:2506.09498 (follow-up, ~100× speedup via parallel rollouts and coarsening, NeurIPS 2025).