mlsanityData to 7 Oct 2026

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

arXiv 2602.04344 ↗First posted 4 Feb 2026cs.LG cs.AIcited by 2

On the map: LLM reasoning → Diffusion reasoning and test-time scaling

also in 1 more
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.

Key facts. Cited by 2 papers in the map (data through 7 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.

Builds on

26 of 30 references.

Closest ancestors · 3

Same direction, newest first

  1. 2025Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement11 Jul
  2. 2025Remasking Discrete Diffusion Models with Inference-Time Scaling1 Mar
  3. 2025Monte Carlo Tree Diffusion for System 2 Planning11 Feb

Other references · 23

  1. 2022Self-Consistency Improves Chain of Thought Reasoning in Language Modelscited by 5,601
  2. 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
  3. 2024LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Codecited by 1,906
  4. 2021Structured Denoising Diffusion Models in Discrete State-Spacescited by 1,665
  5. 2023Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generationcited by 1,332
  6. 2024Large Language Monkeys: Scaling Inference Compute with Repeated Samplingcited by 946
  7. 2025Large Language Diffusion Modelscited by 872
  8. 2024Simple and Effective Masked Diffusion Language Modelscited by 800
  9. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distributioncited by 712
  10. 2024Simplified and Generalized Masked Diffusion for Discrete Datacited by 606
  11. 2021Argmax Flows and Multinomial Diffusion: Learning Categorical Distributionscited by 566
  12. 2025Dream 7B: Diffusion Large Language Modelscited by 515
11 more
  1. 2025The Lessons of Developing Process Reward Models in Mathematical Reasoningcited by 402
  2. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Modelscited by 385
  3. 2025LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Modelscited by 273
  4. 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusionscited by 208
  5. 2025DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generationcited by 184
  6. 2025Dream-Coder 7B: An Open Diffusion Language Model for Codecited by 81
  7. 2025Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Searchcited by 48
  8. 2025Diffusion Tree Sampling: Scalable inference-time alignment of diffusion modelscited by 31
  9. 2026Stable-DiffCoder: Pushing the Frontier of Code Diffusion Large Language Modelcited by 12
  10. 2025Diffusion Language Model Inference with Monte Carlo Tree Searchcited by 3
  11. 2025Tree Reward-Aligned Search for TReASURe in Masked Diffusion Language Modelscited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editingcited by 1
  2. 2026CaRE Compute-aware Remasking Evaluation Protocol for Masked Diffusion Language Modelscited by 1

Similar, not cited

Close in wording, no citation either way

  1. 2025Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models
  2. 2026Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models
  3. 2025Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models
  4. 2025Learning Unmasking Policies for Diffusion Language Models
  5. 2026Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Models
  6. 2024Scaling up Masked Diffusion Models on Text
  7. 2026Mask-Aware Policy Gradients for Diffusion Language Models
  8. 2026Scaling Beyond Masked Diffusion Language Models