Могу я принять ваш заказ? Поиск по дереву Монте-Карло для определения порядка заполнения слотов в диффузионных языковых моделях

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

February 13, 2026

Авторы: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia

cs.AI

Аннотация

Хотя декодирование по принципу «план-заполнение» в маскированных диффузионных моделях (MDM) демонстрирует перспективность для математических рассуждений и анализа кода, производительность остается крайне чувствительной к порядку заполнения слотов, что часто приводит к значительной вариативности выходных данных. Мы представляем McDiffuSE — фреймворк, который формулирует выбор слота как принятие решений и оптимизирует порядок заполнения с помощью поиска в дереве Монте-Карло (MCTS). McDiffuSE использует прогнозирующее моделирование для оценки частичных завершений до фиксации выбора, систематически исследуя комбинаторное пространство порядков генерации. Эксперименты показывают среднее улучшение на 3.2% по сравнению с авторегрессионными базовыми методами и на 8.0% по сравнению с базовым методом «план-заполнение», с существенным приростом в 19.5% на наборе MBPP и 4.9% на MATH500. Наш анализ показывает, что хотя McDiffuSE в основном следует последовательному порядку, включение несеквентной генерации необходимо для максимизации производительности. Мы наблюдаем, что для преодоления смещений уверенности модели и обнаружения эффективных порядков необходимы более высокие константы исследования, а не увеличение числа симуляций. Эти результаты устанавливают планирование на основе MCTS как эффективный подход для повышения качества генерации в 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.