UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching

Test-time scaling strategies have effectively leveraged inference-time compute to enhance the reasoning abilities of Autoregressive Large Language Models. In this work, we demonstrate that Masked Diffusion Language Models (MDLMs) are inherently amenable to advanced search strategies, owing to their iterative and non-autoregressive generation process. To leverage this, we propose UnMaskFork (UMF), a framework that formulates the unmasking trajectory as a search tree and employs Monte Carlo Tree Search to optimize the generation path. In contrast to standard scaling methods relying on stochastic sampling, UMF explores the search space through deterministic partial unmasking actions performed by multiple MDLMs. Our empirical evaluation demonstrates that UMF consistently outperforms existing test-time scaling baselines on complex coding benchmarks, while also exhibiting strong scalability on mathematical reasoning tasks.

Test-time scaling, which gives Large Language Models (LLMs) more time to think before answering, has recently become a crucial technique for improving their performance. While highly successful for standard LLMs, it is equally important to unlock this potential for Masked Diffusion Language Models (MDLMs), a variant of LLMs that generates text by iteratively filling in blank spaces. Because MDLMs share this unique process of refining partially completed text, different MDLMs can easily collaborate on the same problem. We introduce UnMaskFork (UMF), a method that allows multiple models to work together, exploring different paths to the solution by systematically branching their decisions. We tested UMF on coding and mathematical tasks and found that it consistently outperforms existing test-time scaling methods for MDLMs. By enabling this collaborative unmasking, UMF introduces structural diversity into the problem-solving process, proving that teamwork between models is a highly effective way to scale reasoning capabilities.