Posso Anotar Seu Pedido? Busca em Árvore de Monte Carlo para Ordenação de Preenchimento de Slots em Modelos de Linguagem de Difusão

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

Resumo

Embora a decodificação por planeamento e preenchimento em Modelos de Difusão Mascarada (MDMs) mostre potencial para raciocínio matemático e de código, o desempenho permanece altamente sensível à ordem de preenchimento das lacunas, frequentemente resultando em uma variação substancial da saída. Apresentamos o McDiffuSE, uma estrutura que formula a seleção de lacunas como um processo de tomada de decisão e otimiza as ordens de preenchimento através da Busca em Árvore de Monte Carlo (MCTS). O McDiffuSE utiliza simulações de antecipação para avaliar completamentos parciais antes do comprometimento, explorando sistematicamente o espaço combinatório de ordens de geração. Experimentos mostram uma melhoria média de 3,2% em relação a linhas de base autoregressivas e de 8,0% sobre a linha de base de planeamento e preenchimento, com ganhos notáveis de 19,5% no MBPP e 4,9% no MATH500. Nossa análise revela que, embora o McDiffuSE siga predominantemente uma ordenação sequencial, a incorporação de geração não sequencial é essencial para maximizar o desempenho. Observamos que constantes de exploração maiores, em vez de um aumento no número de simulações, são necessárias para superar os vieses de confiança do modelo e descobrir ordenações eficazes. Estas descobertas estabelecem o planeamento baseado em MCTS como uma abordagem eficaz para melhorar a qualidade da geração em MDMs.

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.