Mag ik uw bestelling? Monte-Carlo Tree Search voor het bepalen van de invulvolgorde bij diffusie-taalmodellen.
Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language Models
February 13, 2026
Auteurs: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia
cs.AI
Samenvatting
Hoewel plan-and-infill-decodering in Gemaskeerde Diffusiemodellen (MDM's) veelbelovend is voor wiskundig en code-redeneerwerk, blijft de prestaties sterk afhankelijk van de volgorde van invullen, wat vaak aanzienlijke outputvariantie oplevert. Wij introduceren McDiffuSE, een raamwerk dat slotselectie formuleert als besluitvorming en invulvolgordes optimaliseert door middel van Monte Carlo Tree Search (MCTS). McDiffuSE gebruikt look-ahead-simulaties om gedeeltelijke voltooiingen te evalueren vóór definitieve keuze, waarbij de combinatorische ruimte van generatievolgordes systematisch wordt verkend. Experimenten tonen een gemiddelde verbetering van 3,2% ten opzichte van autoregressieve baseline-modellen en 8,0% ten opzichte van baseline plan-and-infill, met opmerkelijke winsten van 19,5% op MBPP en 4,9% op MATH500. Onze analyse toont aan dat hoewel McDiffuSE voornamelijk een sequentiële volgorde volgt, het incorporeren van niet-sequentiële generatie essentieel is voor het maximaliseren van de prestaties. Wij observeren dat grotere exploratieconstanten, in plaats van een toename van simulaties, nodig zijn om bias door modelvertrouwen te overwinnen en effectieve volgordes te ontdekken. Deze bevindingen vestigen MCTS-gebaseerde planning als een effectieve aanpak voor het verbeteren van de generatiekwaliteit in MDM's.
English
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.