UnMaskFork (UMF): Scaling MDLM Inference

• UnMaskFork is a test-time scaling framework for MDLMs that reformulates the unmasking process as a search tree to optimize generation through deterministic branching.
• It employs Monte Carlo Tree Search with a fixed inference budget to explore diverse inference configurations using low-temperature settings and deterministic remasking rules.
• Empirical results demonstrate significant performance gains on benchmarks like LiveCodeBench, HumanEval+, MBPP+, and MATH by leveraging zero-variance rollouts and efficient caching.

UnMaskFork (UMF) is a test-time scaling framework for Masked Diffusion LLMs (MDLMs) that formulates the unmasking trajectory as a search tree and uses Monte Carlo Tree Search (MCTS) to optimize generation under a fixed inference budget measured in Number of Function Evaluations (NFE). It was introduced in "UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching" (Misaki et al., 4 Feb 2026). The central premise is that MDLMs already generate text by iteratively unmasking a sequence of mask tokens into word tokens, and that this iterative, non-autoregressive process is naturally compatible with tree-search-based inference. In UMF, diversity is obtained not through stochastic perturbations but through deterministic branching among multiple high-quality inference configurations, such as different pretrained MDLMs, temperatures approximately equal to zero, and deterministic remasking heuristics (Misaki et al., 4 Feb 2026).

1. Problem formulation and motivation

UMF is motivated by a contrast between autoregressive test-time scaling and masked diffusion inference. Standard scaling methods for autoregressive LLMs, including Best-of-$N$, tree search, and self-consistency, typically inject stochasticity through higher temperature in order to increase diversity. In MDLMs, however, early stochastic choices propagate through later denoising steps because each committed token becomes context for all subsequent updates; the reported empirical result is that raising temperature across the entire schedule harms coding performance even though diversity increases (Misaki et al., 4 Feb 2026).

The underlying MDLM generation process is written as a reverse diffusion from a fully masked sequence

$z_T\in(V\cup\{m\})^n$

to a fully unmasked sequence

$z_0\in V^n$

in $T$ discrete time steps. At each step $t$, for every position $i$, the model predicts a categorical distribution

$$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$

and a remasking or unmasking rule chooses a subset $S_t\subseteq \mathcal{M}(z_t)$ of positions still masked to commit to either greedy tokens in the limit $T\to 0$ or sampled tokens (Misaki et al., 4 Feb 2026).

The framework therefore shifts the source of diversity from stochastic sampling to structural diversity. The paper states that MDLMs benefit more from interleaving different pretrained models or deterministic remasking heuristics, because each such configuration defines a distinct, high-quality reverse-diffusion kernel with low variance (Misaki et al., 4 Feb 2026). A plausible implication is that UMF treats inference-time compute as a mechanism for adaptive control over denoising dynamics rather than for repeated noisy sampling.

2. Search-tree representation of unmasking trajectories

UMF formalizes inference as a search tree $\mathcal{T}$ whose nodes correspond to partial states $z_T\in(V\cup\{m\})^n$0 with residual mask ratio

$z_T\in(V\cup\{m\})^n$1

where $z_T\in(V\cup\{m\})^n$2 is the length of the generation segment (Misaki et al., 4 Feb 2026). The root node is

$z_T\in(V\cup\{m\})^n$3

and terminal nodes are fully unmasked states $z_T\in(V\cup\{m\})^n$4 satisfying $z_T\in(V\cup\{m\})^n$5 (Misaki et al., 4 Feb 2026).

At a node $z_T\in(V\cup\{m\})^n$6 with state $z_T\in(V\cup\{m\})^n$7, the action set $z_T\in(V\cup\{m\})^n$8 consists of inference configurations

$z_T\in(V\cup\{m\})^n$9

where $z_0\in V^n$0 selects one of $z_0\in V^n$1 pretrained MDLMs, $z_0\in V^n$2 is a temperature, and $z_0\in V^n$3 is a deterministic remasking strategy such as entropy-based or low-confidence unmasking (Misaki et al., 4 Feb 2026). The paper states that UMF typically uses $z_0\in V^n$4 for determinism.

Executing an action induces a deterministic transition

$z_0\in V^n$5

where $z_0\in V^n$6 is obtained by repeatedly applying the MDLM’s single-step unmasking transition under $z_0\in V^n$7 until a lower target ratio is reached. The fixed ratio schedule is

$z_0\in V^n$8

Each single-step unmask consumes $z_0\in V^n$9 NFE, while repeated $T$0 pairs are free once cached (Misaki et al., 4 Feb 2026).

This search-tree construction is the formal core of UnMaskFork. It replaces the conventional view of diffusion inference as a single fixed denoising trajectory with a branching process over partial unmasking states. This suggests that UMF is not merely a reranking layer over complete generations; it operates directly on intermediate latent text states.

3. MCTS, UCT, and deterministic partial unmasking

UMF performs standard MCTS over the unmasking tree subject to an NFE budget $T$1. The procedure repeatedly selects a node, expands an untried action, rolls out deterministically to a terminal state, and backpropagates the resulting reward until the total NFE reaches the budget (Misaki et al., 4 Feb 2026).

Selection uses UCT. For each child $T$2 reached from node $T$3 via action $T$4, UMF maintains the parent visit count $T$5, the child visit count $T$6, and cumulative reward $T$7. The UCT score is

$T$8

with $T$9 (Misaki et al., 4 Feb 2026).

During expansion, if $t$0 is already present in the cache $t$1, UMF retrieves the resulting node and reward at zero NFE. Otherwise, it calls the transition routine that unmaskes the state to the next target ratio, caches every intermediate node, then continues deterministically to a fully unmasked terminal state $t$2, evaluates a reward, and caches the result (Misaki et al., 4 Feb 2026). Reported examples of reward are pass@k on code tests and exact match on MATH.

A central property of this setup is that each rollout under $t$3 is fully deterministic, so

$t$4

The paper therefore states that a single simulation yields an unbiased, zero-variance estimate $t$5, and that UMF does not learn an explicit policy prior $t$6; instead, the prior effect is implicit in the tree structure (Misaki et al., 4 Feb 2026). The score decomposition is

$t$7

and action selection maximizes $t$8 (Misaki et al., 4 Feb 2026).

The deterministic partial unmasking strategy is implemented over a small discrete action set. The paper gives as an example two MDLMs, Dream-Coder and LLaDA, each paired with a recommended low-temperature setting and confidence-based remasking. It further states that, unlike stochastic sampling baselines, deterministic branching yields zero-variance returns, so each branch is evaluated exactly once and the NFE budget is spent on exploration rather than repeated averaging (Misaki et al., 4 Feb 2026).

4. Objective, action design, and caching behavior

UMF’s stated objective is to maximize the expected terminal reward $t$9 subject to an NFE budget constraint: $i$0 under

$i$1

MCTS with UCT provides the exploration-exploitation mechanism, and inference stops when NFE reaches $i$2 (Misaki et al., 4 Feb 2026).

The action design is deliberately restricted to deterministic, high-quality branches. The paper specifies that actions are chosen from a small, discrete set $i$3, for example two MDLMs $i$4, each paired with a low temperature and a confidence-based remasking rule (Misaki et al., 4 Feb 2026). In the reported experiments, Dream-Coder-v0-Instruct-7B uses $i$5 with entropy remasking, and LLaDA-8B-Instruct uses $i$6 with low-confidence remasking (Misaki et al., 4 Feb 2026).

Caching is not an auxiliary optimization but an integral part of the algorithmic design. Once a node-action pair $i$7 is visited, UMF deterministically unmaskes all the way to the next ratio, caches the full sub-rollout keyed by $i$8, and guarantees that revisiting the same pair reproduces the identical sub-trajectory at zero NFE (Misaki et al., 4 Feb 2026). The paper emphasizes that this behavior is especially effective because deterministic actions make repeated evaluation unnecessary.

The authors also present an interpretive bound by viewing inference as adaptive selection among reverse-diffusion kernels $i$9: $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$0 They argue that interleaving locally best choices can therefore yield lower overall diffusion KL error than any fixed model (Misaki et al., 4 Feb 2026). This suggests a formal justification for multi-model branching beyond heuristic ensembling.

5. Empirical evaluation on code and mathematical reasoning

The reported evaluation covers LiveCodeBench, HumanEval+, and MBPP+ for code, plus MATH for mathematical reasoning. The compared methods are Best-of-$$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$1 at various temperatures, DTS* (diffusion tree sampling), AB-MCTS with single-model and multi-model actions, and a "Pair" baseline that splits the budget equally between single-model UMFs. All methods are matched on $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$2 in the main code comparison (Misaki et al., 4 Feb 2026).

Benchmark	Baselines at $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$3	UMF
LiveCodeBench	BoN best $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$4; DTS* $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$5–$$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$6; AB-MCTS $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$7	$$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$8
HumanEval+	BoN $$p_{\theta,i}(x\mid z_t)=\mathrm{softmax}(\ell_{\theta,i}(x\mid z_t)/T),$$9; DTS* $S_t\subseteq \mathcal{M}(z_t)$0; AB-MCTS $S_t\subseteq \mathcal{M}(z_t)$1	$S_t\subseteq \mathcal{M}(z_t)$2
MBPP+	BoN $S_t\subseteq \mathcal{M}(z_t)$3; DTS* $S_t\subseteq \mathcal{M}(z_t)$4; AB-MCTS $S_t\subseteq \mathcal{M}(z_t)$5	$S_t\subseteq \mathcal{M}(z_t)$6

On MATH, using $S_t\subseteq \mathcal{M}(z_t)$7 problems, UMF improves from $S_t\subseteq \mathcal{M}(z_t)$8 pass@1 at $S_t\subseteq \mathcal{M}(z_t)$9 to $T\to 0$0 at $T\to 0$1, described as an $T\to 0$2-point gain over the low-budget baseline (Misaki et al., 4 Feb 2026).

The scaling results are also explicit. UMF shows smooth, monotonic improvement with NFE up to $T\to 0$3, reaching $T\to 0$4 pass@1 on LiveCodeBench, whereas stochastic baselines plateau or fluctuate (Misaki et al., 4 Feb 2026). This reported monotonicity is consistent with the deterministic-search premise: when branch values are zero-variance and reusable, additional compute can be allocated to exploration in a controlled manner.

6. Ablations, limitations, and scope

The ablation study in Section 5.3 isolates three components. First, caching improves pass@1 at $T\to 0$5 NFE from $T\to 0$6 to $T\to 0$7, with a reported cache hit rate of approximately $T\to 0$8 (Misaki et al., 4 Feb 2026). Second, mixing two models outperforms mixing temperatures or mixing remasking strategies by $T\to 0$9–$\mathcal{T}$0 points (Misaki et al., 4 Feb 2026). Third, a single UMF run interleaving both models at $\mathcal{T}$1 outperforms the "Pair" baseline that runs two single-model UMFs at $\mathcal{T}$2 each and picks the better answer, with $\mathcal{T}$3 versus $\mathcal{T}$4 on LiveCodeBench (Misaki et al., 4 Feb 2026).

The paper also states several limitations. UMF’s performance hinges on having diverse high-quality models and heuristics; when only a single MDLM is available, the gain reduces to mixing temperatures and remasking strategies, which still help but less so (Misaki et al., 4 Feb 2026). Practical overheads include tokenizer mapping when switching between models and the need for a fast cache, although the reported implementation found these negligible compared with MDLM forward cost (Misaki et al., 4 Feb 2026).

A common misconception would be to interpret UMF primarily as a stochastic search method for diffusion models. The paper directly argues the opposite: its distinctive feature is the replacement of stochastic noise with deterministic, high-quality branching actions. Another possible misconception is that UMF simply reranks independently generated final answers. The search-tree construction, partial-ratio schedule, and caching over intermediate states indicate that UMF operates over the internal unmasking trajectory itself (Misaki et al., 4 Feb 2026).

In that formulation, UnMaskFork occupies a specific position within test-time scaling research: it adapts MCTS and UCT to masked diffusion text generation by exploiting properties that are atypical for autoregressive decoding, namely iterative unmasking, deterministic low-temperature transitions, and reusable partial trajectories. The reported significance of the method is therefore not only its benchmark gains but also the claim that MDLM inference is inherently amenable to advanced search strategies (Misaki et al., 4 Feb 2026).