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

arXiv: https://arxiv.org/abs/2602.04344 · HTML: https://mlsanity.com/paper/2602.04344/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 26 of its 30 references that are in the map.
- First posted on arXiv 4 Feb 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 26 of 26)

- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [The Lessons of Developing Process Reward Models in Mathematical Reasoning](https://mlsanity.com/paper/2501.07301/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Models](https://mlsanity.com/paper/2505.19223/) (2025)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)
- [DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation](https://mlsanity.com/paper/2506.20639/) (2025)
- [Dream-Coder 7B: An Open Diffusion Language Model for Code](https://mlsanity.com/paper/2509.01142/) (2025)
- [Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search](https://mlsanity.com/paper/2503.04412/) (2025)
- [Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models](https://mlsanity.com/paper/2506.20701/) (2025)
- [Monte Carlo Tree Diffusion for System 2 Planning](https://mlsanity.com/paper/2502.07202/) (2025)
- [Stable-DiffCoder: Pushing the Frontier of Code Diffusion Large Language Model](https://mlsanity.com/paper/2601.15892/) (2026)
- [Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement](https://mlsanity.com/paper/2507.08390/) (2025)
- [Diffusion Language Model Inference with Monte Carlo Tree Search](https://mlsanity.com/paper/2512.12168/) (2025)
- [Tree Reward-Aligned Search for TReASURe in Masked Diffusion Language Models](https://mlsanity.com/paper/2509.23146/) (2025)

## Built on by (most cited first, 2 of 2)

- [VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing](https://mlsanity.com/paper/2606.28301/) (2026)
- [CaRE Compute-aware Remasking Evaluation Protocol for Masked Diffusion Language Models](https://mlsanity.com/paper/2607.24763/) (2026)

## Abstract (arXiv)

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.