주문을 받아드릴까요? 확산 언어 모델에서 슬롯 채우기 순서 결정을 위한 몬테카를로 트리 탐색 (또는 보다 학술적인 맥락에서:) 주문을 받아드릴까요? 확산 언어 모델의 슬롯 채움 순서 결정을 위한 몬테카를로 트리 탐색

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)에서 계획-채움 디코딩 방식은 수학 및 코드 추론 분야에서 잠재력을 보이지만, 슬롯 채움 순서에 대한 민감도가 높아 결과 변동성이 큰 한계가 있습니다. 본 연구에서는 슬롯 선택 문제를 의사 결정 과정으로 공식화하고 몬테카를로 트리 탐색(MCTS)을 통해 채움 순서를 최적화하는 McDiffuSE 프레임워크를 제안합니다. McDiffuSE는 확정 전 부분 완성본을 평가하는 선행 시뮬레이션을 활용하여 생성 순서의 조합 공간을 체계적으로 탐색합니다. 실험 결과, 자동회귀 기준 모델 대비 평균 3.2%, 기준 계획-채움 방식 대비 8.0%의 성능 향상을 보였으며, MBPP에서 19.5%, MATH500에서 4.9%의 두드러진 개선 효과를 확인했습니다. 분석 결과, 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.