¿Puedo Tomar Su Pedido? Búsqueda de Árbol Monte-Carlo para la Ordenación de Relleno de Espacios en Modelos de Lenguaje de Difusión

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

February 13, 2026

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

cs.AI

Resumen

Si bien la decodificación de planificación y relleno en Modelos de Difusión Enmascarada (MDM) muestra potencial para el razonamiento matemático y de código, su rendimiento sigue siendo altamente sensible al orden de relleno de los espacios, lo que a menudo genera una variabilidad sustancial en la salida. Presentamos McDiffuSE, un marco que formula la selección de espacios como un proceso de toma de decisiones y optimiza los órdenes de relleno mediante Búsqueda en Árbol Monte Carlo (MCTS). McDiffuSE utiliza simulaciones de anticipación para evaluar completamientos parciales antes de comprometerse con ellos, explorando sistemáticamente el espacio combinatorio de órdenes de generación. Los experimentos muestran una mejora promedio del 3.2% sobre los modelos autoregresivos base y del 8.0% sobre el método base de planificación y relleno, con ganancias notables del 19.5% en MBPP y del 4.9% en MATH500. Nuestro análisis revela que, si bien McDiffuSE sigue predominantemente un orden secuencial, la incorporación de generación no secuencial es esencial para maximizar el rendimiento. Observamos que se necesitan constantes de exploración más grandes, en lugar de un aumento de simulaciones, para superar los sesgos de confianza del modelo y descubrir órdenes efectivos. Estos hallazgos establecen la planificación basada en MCTS como un enfoque efectivo para mejorar la calidad de la generación en los 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.