Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, S. Kakade, Sitan Chen

arXiv:2502.06768 · 2026-07-27 공개 · arXiv · PDF

autoregressive-models generative-modeling inference-strategies masked-diffusion-models sudoku token-ordering logic-puzzles training-complexity

Abstract

In recent years, masked diffusion models (MDMs) have emerged as a promising alternative approach for generative modeling over discrete domains. Compared to autoregressive models (ARMs), MDMs trade off complexity at training time with flexibility at inference time. At training time, they must learn to solve an exponentially large number of infilling problems, but at inference time, they can decode tokens in essentially arbitrary order. In this work, we closely examine these two competing effects. On the training front, we theoretically and empirically demonstrate that MDMs indeed train on computationally intractable subproblems compared to their autoregressive counterparts. On the inference front, we show that a suitable strategy for adaptively choosing the token decoding order significantly enhances the capabilities of MDMs, allowing them to sidestep hard subproblems. On logic puzzles like Sudoku, we show that adaptive inference can boost solving accuracy in pretrained MDMs from $<7$% to $\approx 90$%, even outperforming ARMs with $7\times$ as many parameters and that were explicitly trained via teacher forcing to learn the right order of decoding.

한국어 요약

한 줄 요약

마스킹 확산 모델(MDM)의 토큰 순서 결정 전략이 추론 성능에 미치는 영향을 분석하고, 적응적 추론이 성능을 7%에서 90%까지 향상시킴을 보인다.

핵심 기여도

핵심 아이디어

MDM은 훈련 시 모든 가능한 마스킹 패턴을 학습해야 하므로, ARMs보다 훨씬 복잡한 서브문제를 풀어야 한다. 이는 이론적으로 계산 불가능한 수준의 복잡도를 가진다. 그러나 추론 시 토큰 순서를 자유롭게 조정할 수 있는 유연성이 MDM의 강점이다. 본 연구는 이 유연성을 활용한 적응적 추론 전략(예: Top-K 확률 마진 기반)을 제안하고, 이를 통해 훈련 시 학습하지 못한 순서 문제를 우회할 수 있음을 보인다. 특히, 수학 퍼즐(Sudoku)에서 이 접근법이 성능을 극적으로 향상시킴을 실험적으로 입증했다.

기술적 접근법

주요 결과

의의 및 한계

본 연구는 MDM의 훈련 복잡도와 추론 유연성의 균형을 분석하고, 적응적 추론 전략이 MDM의 성능을 극적으로 향상시킬 수 있음을 보인다. 이는 MDM이 ARMs의 강점을 보완할 수 있음을 시사하며, 특히 추론 단계에서의 유연성이 핵심 가치로 작용함을 강조한다. 그러나 현재 제안된 적응 전략은 Top-K와 같은 간단한 알고리즘에 의존하며, 더 복잡한 상황에서는 한계가 있을 수 있다.

실용적 활용

MDM과 적응적 추론 전략은 자연어 생성, 코드 생성, 논리 추론, 텍스트 인필링 등에서 활용 가능하다. 특히, 순서에 민감한 작업(예: 퍼즐 해결, 계획 수립)에서 MDM의 유연성이 실질적 이점을 제공할 수 있다.