Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes UnMaskFork (UMF), an MCTS-based test-time scaling method for masked diffusion LMs that branches over deterministic unmasking actions (different MDLMs or low-temperature heuristics) with node caching, outperforming Best-of-N and DTS* on coding benchmarks under fixed NFE.

Problem

Test-time scaling for masked diffusion language models is underexplored; stochastic temperature-based sampling, the standard AR scaling approach, degrades MDLM generation quality, and existing diffusion tree search methods are sample-inefficient under fixed compute.

Approach

Formulates unmasking as a search tree where actions are deterministic partial-unmasking configurations (switching among multiple pretrained MDLMs or low-temperature remasking heuristics), runs MCTS with UCT, and caches deterministic rollouts so revisited state-action pairs cost zero NFE.

Key results

• At NFE=12288 on 100-sample subsets, UMF reaches 28.0 Pass@1 on LiveCodeBench, 88.0 on HumanEval+, and 72.0 on MBPP+, above the best BoN (19.0/75.0/66.0) and DTS* (18.0/75.0/68.0) baselines.
• On 105 MATH problems, UMF improves Pass@1 from 49.52% at NFE=768 to 60.95% at NFE=12288.
• Caching raises LiveCodeBench Pass@1 at NFE=12288 from 26.0 to 28.0 with about 55.8% cache hit rate.

Surprising findings

• Random remasking and high-temperature stochastic sampling degraded performance in masked diffusion decoding, with LLaDA random remasking reaching only 2.0 Pass@1 on LiveCodeBench under BoN.
• A single-model UMF variant (Dream-Coder only, varying temperature/remask) already outperforms multi-model BoN, DTS*, and AB-MCTS baselines.

Limitations

• Evaluated on small subsets (100 problems for coding, 105 for MATH), so variance is high.
• Relies on heuristic UCT without a learned value or policy function.
• Tokenizer switching between heterogeneous MDLMs introduces small length drift and requires text-based re-encoding.

Experimental setup

Datasets: LiveCodeBench HumanEval+ MBPP+ MATH

Benchmarks: LiveCodeBench HumanEval+ MBPP+ MATH

Metrics: Pass@1 NFE

Baselines: Best-of-N DTS* AB-MCTS ReMDM TReASURe Pair baseline

Builds on: Diffusion Tree Sampling AB-MCTS UCT LLaDA Dream

Closest papers by meaning

• Scaling up Masked Diffusion Models on Text (ICLR 2025 · similarity 0.91)
• Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models (ICML 2026 · similarity 0.91)
• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (NeurIPS 2025 · similarity 0.91)
• On the Reasoning Abilities of Masked Diffusion Language Models (ICLR 2026 · similarity 0.90)
• TEST-TIME SCALING IN DIFFUSION LLMS VIA HIDDEN SEMI-AUTOREGRESSIVE EXPERTS (ICLR 2026 · similarity 0.90)
• Lookahead Path Likelihood Optimization for Diffusion LLMs (ICML 2026 · similarity 0.89)