PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2602.12586v2cs.AI

TL;DR

Plan-and-infill decoding in MDMs is sensitive to slot ordering, creating a need for better planning across interdependent generation steps. McDiffuSE formulates slot selection as decision making and uses MCTS lookahead to search infilling orders. It improves performance over diffusion and autoregressive baselines, while analysis shows that selective non-sequential generation and broad exploration are important.

  • Problem

    MDM plan-and-infill performance is highly sensitive to slot ordering because inter-slot dependencies can propagate errors and undermine global coherence.

  • Method

    McDiffuSE is a training-free framework that formulates slot ordering as decision making and uses MCTS lookahead to explore generation trajectories.

  • Results

    MCDIFFUSE matches or exceeds autoregressive performance on five of six benchmarks, including a 19.45% absolute accuracy gain on MBPP.

  • Takeaways & Limitations

    Selective non-sequential generation and a large exploration constant help discover effective slot orderings beyond biased sequential or confidence-greedy choices.

  • Takeaways & Limitations

    Text-diffusion slots may not be the best speed-accuracy application, with strongest impact expected when reasoning trajectories permit revision or alternative orderings.

Abstract

from arXiv · show

While plan-and-infill decoding in Masked Diffusion Models (MDMs) shows promise for mathematical and code reasoning, performance remains highly sensitive to slot infilling order, often yielding substantial output variance. We introduce McDiffuSE, a framework that formulates slot selection as decision making and optimises infilling orders through Monte Carlo Tree Search (MCTS). McDiffuSE uses look-ahead simulations to evaluate partial completions before commitment, systematically exploring the combinatorial space of generation orders. Experiments show an average improvement of 3.2% over autoregressive baselines and 8.0% over baseline plan-and-infill, with notable gains of 19.5% on MBPP and 4.9% on MATH500. Our analysis reveals that while McDiffuSE predominantly follows sequential ordering, incorporating non-sequential generation is essential for maximising performance. We observe that larger exploration constants, rather than increased simulations, are necessary to overcome model confidence biases and discover effective orderings. These findings establish MCTS-based planning as an effective approach for enhancing generation quality in MDMs.

1. Introduction

Masked Diffusion Models remove the fixed left-to-right decoding constraint, but plan-and-infill quality remains highly sensitive to slot ordering. MCDIFFUSE addresses this planning problem with MCTS lookahead and achieves strong gains across reasoning benchmarks.

  • Motivation: MDMs permit non-sequential decoding but can underperform autoregressive models because simultaneously generated interdependent tokens lack mutual conditioning.This makes output quality sensitive to generation order.
  • Motivation: Plan-and-infill reduces inner-slot dependency difficulty, but errors in inter-slot ordering can propagate across iterations and undermine global coherence.The remaining challenge is searching a vast combinatorial space of slot permutations.
  • Approach: MCDIFFUSE formulates slot selection as decision making and uses prior-guided MCTS with lookahead simulations over possible slot orderings.Its hybrid reward combines immediate denoising quality with rollout-based long-term trajectory coherence.
  • Results: 19.45% absolute accuracy gain on MBPP and 4.9% on MATH500 demonstrate strong improvements across six reasoning benchmarks.MCDIFFUSE matches or exceeds autoregressive performance on five of six benchmarks under identical conditions.
  • Analysis: Non-sequential generation is essential for maximising performance, even though MCDIFFUSE predominantly follows sequential ordering.The search selectively departs from left-to-right generation for difficult cases.
  • Analysis: Exploration breadth matters more than simulation count: a large exploration constant helps overcome biased confidence priors and discover effective slot orderings.Increasing simulations alone does not consistently improve performance.

2. Background

The paper frames slot ordering as a sequential decision problem and introduces MCTS as a search procedure balancing exploration and exploitation. Neural policy priors guide early exploration while rollout values estimate longer-term returns.

  • Markov Decision Processes: A Markov decision process consists of states, actions, transitions, and rewards, with a policy seeking to maximise expected cumulative reward.The state captures an environment configuration and actions determine probabilistic successor states.
  • Monte Carlo Tree Search: MCTS approximates optimal policies by combining tree search with Monte Carlo simulation across selection, expansion, simulation, and backpropagation phases.Nodes represent states and edges represent actions.
  • Monte Carlo Tree Search: Selection traverses the tree using a bandit policy, expansion adds unexplored actions, simulation rolls out a trajectory, and backpropagation updates values.Repeated simulations estimate the value of alternative actions.
  • Neural-Guided MCTS: The exploration constant controls the exploration–exploitation trade-off, with larger values producing more exploration.Visit counts track state and action selection frequency during backpropagation.
  • Neural-Guided MCTS: PUCT balances exploitation of empirical mean returns with exploration guided by neural action priors.The prior’s influence diminishes as visit counts increase.

3. MCDIFFUSE

MCDIFFUSE represents slot-order selection as a deterministic MDP and adapts MCTS to choose infilling orders without auxiliary training. Its value estimates combine immediate slot confidence with stochastic rollout quality.

  • Problem Formulation: Slot ordering is represented as a permutation and optimised in a deterministic MDP whose actions select unfilled slots.The goal is to capture how generating one slot influences the others.
  • Problem Formulation: A state records the generated-slot prefix and partially filled sequence, with all ungenerated slots represented by [MASK].The initial state contains an empty ordering and K masked slots.
  • Problem Formulation: Each action selects one masked slot, while tokens within a selected slot are generated autoregressively.Slots themselves may be generated in arbitrary order.
  • Value Estimation: MCDIFFUSE uses model confidence as a reward and combines immediate slot quality with rollout-based future coherence.This avoids relying solely on greedy local confidence.
  • MCTS Procedure: MCTS repeatedly selects actions with PUCT, expands admissible slots, simulates completions, and backpropagates estimated values to update root statistics.The resulting process is repeated until the slot sequence is fully determined.
  • Value Estimation: The value function V(s_t, a) mixes immediate reward R(s_t, a) and expected rollout return using coefficient λ.Rollouts stochastically fill remaining slots until a terminal state and average accumulated rewards.

4. Experimental Setup

The evaluation covers six reasoning, code-generation, and general-knowledge benchmarks using Pass@1 and matched baseline prompting. Experiments compare MCDIFFUSE with diffusion, slot-ordering, and autoregressive baselines under specified search settings.

  • Benchmarks and Metrics: Six benchmarks span GSM8K, MATH500, MBPP, HumanEval, ARC Challenge, and GPQA-Diamond, all evaluated with Pass@1.The study re-evaluates baselines with chain-of-thought prompting for fair comparison.
  • Models and Baselines: MCDIFFUSE is based on ReFusion and is compared with eight diffusion, slot-ordering, and autoregressive baselines.Autoregressive baselines include Qwen2.5 7B and Qwen3 8B.
  • Hyperparameters: The main settings are λ = 0.3, c = 50, Nsim = 256, and τ = 0.5.Sensitivity analyses are reported separately.

5. Experimental Results

MCDIFFUSE outperforms masked diffusion and autoregressive baselines across most benchmarks, and surpasses all MDMs on every benchmark evaluated.

  • MCDIFFUSE outperforms both MDM and ARM baselines across five of six benchmarks.
  • MCDIFFUSE significantly outperforms all MDMs across all six benchmarks.
  • 25.98% improvement on HumanEval and 4.00% improvement on MATH500 versus ReFusion.

2. Coding tasks benefit more from MCTS slot plan-

McDiffuSE achieves larger gains on code-generation benchmarks than on multiple-choice reasoning benchmarks, while also producing more compact reasoning sequences on tasks requiring extensive reasoning.

  • Coding-task performance: 19.45% improvement on MBPP and 16.32% on HumanEval exceed gains on ARC Challenge (0.7%) and GPQA (4.33%).The reported pattern associates larger improvements with code-generation benchmarks than with multiple-choice reasoning benchmarks.
  • Coding-task performance: Code-generation gains are attributed to dependencies among variable declarations, function definitions, and control flow that benefit from strategic slot planning.
  • Generation ordering: Figure 2 relates MBPP accuracy to sequentiality rate and compares average trends for ReFusion and McDiffuSE with a sequential baseline.
  • Reasoning compactness: 152.2 average tokens on MATH500 versus 436.0 for Qwen2.5 7B accompanies better accuracy for McDiffuSE.
  • Reasoning compactness: Reasoning-length reductions are statistically significant across GSM8K, MATH500, and GPQA, with p < 0.001.

6. Analysis and Discussion

MCDIFFUSE improves slot planning by selectively exploring non-sequential orderings and using broad MCTS exploration to avoid locally confident choices. Its gains come with higher compute, while additional simulations alone offer limited benefit.

  • Ordering patterns: 91.1% of coding and 93.8% of mathematical slot-ordering decisions remain sequential, making left-to-right generation the default strategy.Non-sequential decisions are relatively rare across both task categories.
  • Ordering patterns: 60.7% of cases where MCDIFFUSE succeeds but the sequential baseline fails contain at least one non-sequential ordering decision.These departures concentrate in challenging cases and can resolve constraints that greedy left-to-right decoding misses.
  • Search configuration: Higher exploration constants consistently improve accuracy, indicating that stronger exploration is necessary to overcome baseline confidence bias.Changing the exploration constant affects selection without requiring additional forward passes.
  • Search configuration: At c = 2.0, increasing Nsim from 30 to 270 decreases accuracy on both MBPP and MATH500.Under insufficient exploration, additional simulations can reinforce locally confident but globally myopic orderings.
  • Efficiency: Nsim = 270 achieves the highest accuracy, but improves over Nsim = 30 by less than 2% while increasing computational cost nine-fold.Generation time scales approximately linearly with Nsim, whereas varying c has negligible cost.
  • Qualitative mechanism: In the MBPP example, look-ahead rejects an unnecessary import and prioritises the function definition to improve long-term trajectory coherence.The search evaluates future generation steps before committing to the slot order.
  • Efficiency: MCDIFFUSE uses approximately 35% more compute than ReFusion, increasing from 1.95E + 16 to 2.63E + 16 FLOPs.The reported overhead accompanies consistent accuracy gains across evaluated benchmarks.

7. Related work

Prior work established MCTS as a decision-making method for games, reasoning, and visual diffusion, while this paper applies it to slot ordering in masked diffusion text generation. The related literature positions MCDIFFUSE at the intersection of MCTS and MDM research.

  • Monte Carlo tree search: MCTS combines tree search and Monte Carlo simulation to balance exploration and exploitation in sequential decision-making.It has been used in game-playing agents and more recently for mathematical and coding reasoning.
  • Monte Carlo tree search: MCTS has also been applied to visual diffusion models to improve generation quality.The paper distinguishes this prior use from its focus on masked diffusion text generation.
  • Paper positioning: The paper presents MCDIFFUSE as the first application of MCTS to MDMs for text generation, specifically for slot ordering.Its novelty lies in adapting MCTS to the ordering problem in masked diffusion decoding.
  • Masked diffusion models: Masked diffusion models support non-autoregressive sequence generation but still struggle to match autoregressive models on complex reasoning tasks.This motivates methods that improve MDM generation while preserving flexible ordering.

8. Conclusion

MCDIFFUSE enhances masked diffusion models through training-free strategic slot selection, combining MCTS adaptations with broad exploration and selective non-sequential generation. Across six reasoning benchmarks, it improves over existing MDM baselines and performs competitively with autoregressive models.

  • MCDIFFUSE formulates slot ordering as decision-making and navigates generation-order combinations with MCTS without additional training.The framework uses confidence-aware value propagation and adaptive exploration budgets.
  • MCDIFFUSE consistently improves existing MDM baselines across six reasoning benchmarks and achieves competitive performance with autoregressive models under identical conditions.The conclusion reports particularly strong gains on code generation tasks.
  • Strategic deviations from mostly sequential generation are essential for maximising performance.The analysis identifies non-sequential generation as a critical complement to the dominant sequential pattern.
  • Broad exploration, rather than simply increasing simulations, helps overcome model priors and discover effective slot orderings.The conclusion attributes this requirement to biased confidence priors.

Impact Statement

MCDIFFUSE applies MCTS-based slot planning to masked diffusion generation without additional training, using lookahead and stochastic confidence rollouts. The method’s design balances immediate confidence with long-term rollout feedback, while its experiments tune exploration and value-estimation parameters on MATH500.

  • MCDIFFUSE applies MCTS to slot selection in masked diffusion models without requiring additional training.The method explores multiple search paths with lookahead to reduce propagation of unmodelled dependencies across iterations.
  • MCTS slot selection uses simulations, PUCT-based traversal, expansion, rollout evaluation, and backpropagation to choose the next slot.The algorithm normalises slot confidence scores into priors and returns the selected slot index as the next action.
  • Stochastic confidence rollouts sample remaining slots with probability proportional to temperature-scaled model confidence, producing an estimate of completion quality.The rollout preserves exploration among possible completion paths rather than deterministically selecting the highest-confidence slot.
  • The MCDIFFUSE configuration uses 30 MCTS simulations and exploration constant c = 10.0, while experiments tune temperature and mixing coefficient on MATH500.The reported tuning found τ = 0.5 and λ = 0.3.
  • λ = 0.3 prioritises long-term rollout feedback over immediate confidence in value estimation.The paper characterises rollout returns as more informative for complex mathematical reasoning while retaining confidence as a guiding prior.

D.1. Analysis Between Token Length and Accuracy

MCDIFFUSE improves token efficiency and accuracy relative to autoregressive baselines, with especially large gains on coding tasks. Its compact generations can preserve or improve correctness, although performance varies across benchmarks and reasoning difficulty.

  • 64.77% average token reduction accompanies a 3.38% accuracy improvement across benchmarks versus the autoregressive baseline.The reductions are statistically significant at p < 0.001.
  • 88.05% token reduction and 8.39% accuracy improvement are reported for MBPP.The paper identifies code generation as the setting with particularly strong efficiency gains.
  • 72.55% token reduction remains substantial on ARC despite a modest +0.52% accuracy improvement.This comparison indicates that token savings can remain large even when accuracy gains are limited.
  • Coding tasks show 3.13× higher improvements than reasoning and multiple-choice tasks, with the difference statistically significant.The reported statistics are t = 2.97, p = 0.041, and Cohen’s d = 3.21.
  • MCDIFFUSE uses 1,024 tokens versus 32,768 for Qwen2.5 7B and Qwen3 8B while marginally outperforming both on GSM8K.On GPQA it surpasses Qwen2.5 but underperforms Qwen3 8B, and it is weaker on MATH500.
  • Coding benchmarks show gains above 10%, while ARC-Challenge and GPQA show marginal improvements and MATH500 shows +8.2%.The paper associates stronger improvements with tasks requiring extended, multi-step reasoning.

H. Efficiency Analysis of MCDIFFUSE

MCTS slot planning introduces a computation–accuracy trade-off: simulation count increases generation time approximately linearly, while exploration breadth can improve accuracy more efficiently. Increasing simulations from 30 to 270 yields only marginal gains at nine times the computational cost.

  • Generation time scales approximately linearly with the simulation budget Nsim, whereas changing exploration constant c adds negligible overhead.The exploration constant changes the selection criterion without requiring additional forward passes.
  • Less than 2% accuracy improvement separates Nsim = 270 from Nsim = 30, despite a nine-fold increase in computational cost.The result appears in both coding and MATH500 analyses.
  • A high-exploration, low-budget setting such as c = 100 and Nsim = 30 achieves a 13.62% MBPP improvement with reasonable inference time.The paper presents this configuration as a practical efficiency–accuracy trade-off.
  • On MATH500, higher exploration constants obtain higher accuracy, while increasing simulations shows diminishing returns.The same exploration-versus-budget pattern is reported for coding tasks.
Loading 2602.12586v2…