UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching
Kou Misaki, Takuya Akiba
Abstract
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on16
- Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code GenerationJiawei Liu, Chunqiu Steven Xia, Yuyao Wang, Lingming ZhangNeurIPS 2023 · 2,317 citations
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
- Large Language Diffusion ModelsShen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang et al.NeurIPS 2025 · 949 citations
- Simple and Effective Masked Diffusion Language ModelsSubham S. Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan et al.NeurIPS 2024 · 929 citations
- Argmax Flows and Multinomial Diffusion: Learning Categorical DistributionsEmiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré et al.NeurIPS 2021 · 782 citations
Related papers
- DiffuReason: Enhancing Reasoning Ability for Diffusion Language Models via Monte Carlo Tree SearchYIPING SONG, Jinyu You, Zhiliang Tian, Jinsong Su et al.ICML 2026
- Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language ModelsJoshua Ong, Yu Zhao, Mihaela C. Stoian, Wenda Li et al.ICML 2026
- On Powerful Ways to Generate: Autoregression, Diffusion, and BeyondChenxiao Yang, Cai Zhou, David Wipf, Zhiyuan LiICLR 2026 · 7 citations
- MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward OptimizationChenglong Wang, Yang Gan, Hang Zhou, Chi Hu et al.NeurIPS 2025 · 4 citations
- Plan for Speed: Dilated Scheduling for Masked Diffusion Language ModelsOmer Luxembourg, Haim Permuter, Eliya NachmaniICML 2026 · 30 citations
