주문하시겠습니까? 확산 언어 모델에서 슬롯 채우기 순서 결정을 위한 몬테카를로 트리 탐색

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

마스크된 확산 모델(MDM)의 계획 및 채우기(plan-and-infill) 디코딩은 수학 및 코드 추론 분야에서 유망한 가능성을 보여주지만, 그 성능은 슬롯 채우기 순서에 매우 민감하여 종종 상당한 출력 변동성을 초래합니다. 본 논문에서는 슬롯 선택을 의사 결정 문제로 정식화하고 몬테카를로 트리 탐색(MCTS)을 통해 채우기 순서를 최적화하는 프레임워크인 McDiffuSE를 제안합니다. McDiffuSE는 예측(look-ahead) 시뮬레이션을 사용하여 최종 확정 전에 부분 완성본을 평가함으로써 생성 순서의 조합 공간을 체계적으로 탐색합니다. 실험 결과, 자기회귀(autoregressive) 베이스라인 대비 평균 3.2%, 기존 계획 및 채우기 베이스라인 대비 8.0%의 성능 향상을 보였으며, 특히 MBPP에서 19.5%, MATH500에서 4.9%의 괄목할 만한 성과를 거두었습니다. 분석 결과 McDiffuSE는 주로 순차적 순서를 따르지만, 성능을 극대화하기 위해서는 비순차적 생성을 통합하는 것이 필수적인 것으로 나타났습니다. 또한 모델의 확신 편향을 극복하고 효과적인 순서를 발견하기 위해서는 시뮬레이션 횟수의 증가보다는 더 큰 탐색 상수가 필요하다는 점을 확인했습니다. 이러한 결과는 MCTS 기반 계획이 MDM의 생성 품질을 향상시키는 효과적인 접근법임을 입증합니다.

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.