UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching
This work addresses the challenge of enhancing inference-time performance for diffusion-based language models, offering a novel scaling method that could benefit applications in code generation and mathematical reasoning, though it appears incremental as it builds on existing test-time scaling strategies.
The paper tackled the problem of improving the reasoning abilities of Masked Diffusion Language Models (MDLMs) by proposing UnMaskFork (UMF), a framework that uses Monte Carlo Tree Search to optimize generation paths, resulting in consistent outperformance over existing test-time scaling baselines on complex coding benchmarks and strong scalability on mathematical reasoning tasks.
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.