Can I have your order? Monte-Carlo tree search for slot filling ordering in diffusion language models
File(s)
Accepted version
Author(s)
Type
Conference Paper
Abstract
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 base-lines 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 MCD-
IFFUSE 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.
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 base-lines 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 MCD-
IFFUSE 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.
Date Acceptance
2026-05-09
Copyright Statement
© The Author(s). This paper is embargoed until publication.
Source
43rd International Conference on Machine Learning (ICML 2026)
Publication Status
Accepted
Start Date
2026-07-06
Finish Date
2026-07-11
Coverage Spatial
Seoul, South Korea