Puis-je prendre votre commande ? Recherche arborescente de Monte-Carlo pour l'ordonnancement du remplissage de slots dans les modèles de langage à diffusion

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

February 13, 2026

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

cs.AI

papers.abstract

Bien que le décodage par planification et remplissage dans les modèles de diffusion masquée (MDM) montre des résultats prometteurs pour le raisonnement mathématique et la génération de code, les performances restent très sensibles à l'ordre de remplissage des emplacements, produisant souvent une variance substantielle des sorties. Nous présentons McDiffuSE, un cadre qui formule la sélection des emplacements comme un processus décisionnel et optimise les ordres de remplissage par une recherche arborescente de Monte Carlo (MCTS). McDiffuSE utilise des simulations prospectives pour évaluer les complétions partielles avant engagement, explorant systématiquement l'espace combinatoire des ordres de génération. Les expériences montrent une amélioration moyenne de 3,2 % par rapport aux modèles autorégressifs de référence et de 8,0 % par rapport à la méthode planification-remplissage de base, avec des gains notables de 19,5 % sur MBPP et 4,9 % sur MATH500. Notre analyse révèle que si McDiffuSE suit principalement un ordre séquentiel, l'incorporation d'une génération non séquentielle est essentielle pour maximiser les performances. Nous observons que des constantes d'exploration plus grandes, plutôt qu'un nombre accru de simulations, sont nécessaires pour surmonter les biais de confiance du modèle et découvrir des ordonnancements efficaces. Ces résultats établissent la planification basée sur MCTS comme une approche efficace pour améliorer la qualité de génération dans les MDM.

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.