🤖 AI

Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language Models

This paper introduces McDiffuSE, a framework that leverages Monte Carlo Tree Search to optimize slot infilling orders in Masked Diffusion Models, significantly improving generation quality on mathematical and code reasoning tasks by effectively exploring non-sequential completion strategies.

Original authors: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake a complex cake, but instead of following a recipe step-by-step (mix, then bake, then frost), you have a magical oven that lets you fill in the cake in any order you want. You could frost the top first, then bake the middle, then add the sprinkles. This is how Masked Diffusion Models (MDMs) work: they are powerful AI tools that can generate text (like code or math solutions) by filling in "slots" of missing information in any order.

However, there's a catch. If you frost the cake before the batter is even in the pan, the whole thing collapses. Similarly, if the AI fills in the wrong "slot" first, the rest of the text might become nonsense, even if the AI is very confident about that specific piece.

The paper introduces a new method called MCDIFFUSE to solve this ordering problem. Here is how it works, using simple analogies:

1. The Problem: The "Confident but Wrong" Chef

Current AI models often act like a chef who is very confident about the first step they take. If the model thinks, "I'm 90% sure this sentence should start with a function definition," it will do that immediately. But sometimes, starting with the definition causes a mess later because the model didn't realize it needed to write a comment or a specific syntax marker first to make the code work.

The paper calls this the "plan-and-infill" problem: The AI needs to decide which part of the text to write next. If it picks the wrong order, the final result is poor.

2. The Solution: The "What-If" Simulator (MCTS)

The authors introduce MCDIFFUSE, which uses a technique called Monte Carlo Tree Search (MCTS).

Think of MCTS as a super-planner or a simulator. Before the AI commits to writing a specific part of the text, this planner runs thousands of "what-if" scenarios in its head:

Scenario A: What if I write the function definition first? (Simulates the rest of the code... oh no, it crashes later).
Scenario B: What if I write the syntax marker first? (Simulates the rest... ah, the code flows perfectly).

The planner doesn't just look at the immediate confidence (how sure the AI is about the next word); it looks at the long-term outcome. It asks, "If I take this path, will the whole cake turn out good?"

3. How It Decides: The "Exploration" vs. "Simulation" Balance

The paper makes a fascinating discovery about how to tune this planner. Usually, people think you need to run more simulations to get a better answer. But MCDIFFUSE found that's not the most important thing.

The Old Way: Run 100 simulations of the same path to be sure.
The MCDIFFUSE Way: Run fewer simulations, but be willing to try weird, unlikely paths that the AI usually ignores.

The Analogy: Imagine you are looking for a hidden treasure.

High Simulation, Low Exploration: You dig 100 times in the exact same spot because the map says "treasure here." You find nothing, but you are very sure you looked hard enough.
High Exploration (MCDIFFUSE): You dig a few times in the obvious spot, but you also decide to dig in the weird, overgrown bush next to it because the map might be wrong. It turns out the treasure was in the bush!

The paper found that giving the planner a "large exploration constant" (telling it to be brave and try unlikely orders) was much more effective than just running more simulations. This helps the AI escape the trap of its own overconfidence.

4. The Results: Better Code and Math

When they tested this on tasks like writing Python code and solving math problems:

It works: MCDIFFUSE significantly outperformed previous methods. On some coding tests, it improved accuracy by nearly 20%.
It's efficient: Even though it does extra planning, it actually produces shorter, more concise answers than standard AI models. Standard models often ramble or get stuck in long loops; MCDIFFUSE finds the direct path.
The "Secret Sauce": The AI mostly follows a normal, left-to-right order (like a human reading), but it strategically breaks that rule when necessary. It's like a writer who usually writes chronologically but knows exactly when to jump back and insert a crucial detail earlier to make the story make sense.

Summary

MCDIFFUSE is a smart "traffic controller" for AI text generation. Instead of letting the AI blindly fill in text based on what feels right in the moment, it uses a simulation engine to test different orders of operations. By daring to try "weird" orders that the AI usually ignores, it avoids mistakes and produces cleaner, more accurate code and math solutions.

Technical Summary: MCDIFFUSE – MCTS for Slot Filling Ordering in Diffusion Language Models

1. Problem Statement

Masked Diffusion Models (MDMs) offer a promising alternative to autoregressive models (ARMs) for reasoning tasks by enabling non-sequential, iterative denoising. However, MDMs often underperform ARMs due to the sensitivity of output quality to the order of slot infilling. While the "plan-and-infill" framework (e.g., ReFusion) addresses local dependencies within slots, it relies on heuristic planners (often confidence-based) to determine the global order of slot generation. These heuristics frequently fail to account for long-range dependencies, leading to error propagation and incoherent global outputs. The core challenge is the vast combinatorial space of possible slot permutations, making it difficult to identify an ordering that maximizes global coherence without exhaustive search.

2. Methodology: MCDIFFUSE

The authors introduce MCDIFFUSE, a training-free framework that formulates slot selection as a sequential decision-making problem solved via Monte Carlo Tree Search (MCTS). The approach does not require additional model training; instead, it optimizes the inference-time generation order.

2.1 Problem Formulation as MDP

The slot planning process is modeled as a deterministic Markov Decision Process (MDP) $\langle S, A, T, R \rangle$ :

State ( $S$ ): Encodes the set of generated slots and the current partially filled text sequence. Unfilled slots are masked.
Action ( $A$ ): Selecting the index of a masked slot to infill next.
Transition ( $T$ ): Deterministic. Selecting an action fills the corresponding slot using the model's argmax decoding.
Reward ( $R$ ): Defined as the slot-level confidence, calculated as the mean probability the model assigns to the tokens generated in that specific slot.

2.2 MCTS Adaptation

MCDIFFUSE adapts MCTS with specific mechanisms to handle the unique constraints of diffusion generation:

Prior-Guided Expansion: The model's intrinsic confidence scores are normalized to form a prior probability $P(a|s)$ for the PUCT (Predictor-Upper Confidence Tree) selection criterion. This guides the search toward semantically plausible actions while allowing exploration.
Hybrid Value Estimation: To avoid the myopia of greedy local confidence, the value of a state-action pair $V(s, a)$ $V (s, a)$ is a weighted combination of:
1. Immediate Reward: The confidence of the current slot.
2. Rollout Return: A stochastic simulation (rollout) that estimates the long-term trajectory coherence by filling remaining slots until completion.
- Formula: $V(s, a) = \lambda \cdot R(s, a) + (1 - \lambda) \cdot E_{\pi_{roll}}[G]$ , where $\lambda$ balances immediate confidence against long-term planning.
Selection Strategy: The algorithm uses PUCT to balance exploration and exploitation during the search. The final action selection at the root relies on the robust child criterion (highest visit count), ensuring stability.

3. Key Contributions

MCDIFFUSE Framework: A training-free method that tailors MCTS specifically for searching slot orderings in MDMs, effectively closing the performance gap with autoregressive models.
Strategic Non-Sequential Generation: Empirical analysis reveals that while MCDIFFUSE predominantly follows sequential (left-to-right) ordering, strategic deviations to non-sequential generation are critical for maximizing performance on complex reasoning tasks.
Exploration vs. Simulation Depth: The paper demonstrates that increasing the number of simulations ( $N_{sim}$ ) does not consistently improve performance. Instead, a large exploration constant ( $c_{uct}$ ) is necessary to overcome the model's biased confidence priors and discover effective, non-greedy slot orderings.

4. Experimental Results

MCDIFFUSE was evaluated on six reasoning benchmarks (GSM8K, MATH500, ARC, GPQA-Diamond, MBPP, HumanEval) against autoregressive baselines (Qwen2.5, Qwen3) and diffusion baselines (LLaDA, Dream, ReFusion).

Performance Gains: MCDIFFUSE achieved an average improvement of 3.2% over autoregressive baselines and 8.0% over the baseline plan-and-infill method (ReFusion).
Specific Benchmarks: Notable gains included 19.5% on MBPP (code generation) and 4.9% on MATH500.
Efficiency: MCDIFFUSE generated significantly more compact reasoning sequences. For instance, on MATH500, it produced sequences with an average length of 152.2 tokens compared to 436.0 for Qwen2.5, while achieving higher accuracy.
Coding vs. Reasoning: The method showed substantial gains on code generation tasks (structured dependencies) compared to multiple-choice reasoning, suggesting that MCTS is particularly effective where slot ordering critically affects semantic validity.

5. Significance and Claims

The paper establishes that strategic slot planning is a critical factor in the performance of Masked Diffusion Models. The authors claim that:

MCTS is Effective: MCTS-based planning is a viable and effective approach for enhancing generation quality in MDMs without requiring model retraining.
Non-Sequentiality is Essential: Contrary to the assumption that diffusion models should strictly follow sequential trajectories, the ability to deviate from left-to-right ordering is essential for resolving complex constraints and maximizing accuracy.
Exploration Breadth over Depth: The primary challenge in slot planning is not search depth (number of simulations) but exploration breadth. Overcoming the model's local confidence biases requires a high exploration constant to ensure the search visits low-prior but globally coherent paths.

The work positions MCDIFFUSE as a step toward bridging the gap between diffusion and autoregressive models, particularly in domains requiring structured, multi-step reasoning where global coherence is paramount.