---
title: 'UnMaskFork (UMF): Scaling MDLM Inference'
url: https://www.emergentmind.com/topics/unmaskfork-umf
type: topic
---

# UnMaskFork (UMF): Scaling MDLM Inference

UnMaskFork (UMF) is a test-time scaling framework for Masked Diffusion Language Models (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" [2602.04344]. 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 [2602.04344].

## 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 large language models, 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 [2602.04344].

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 [2602.04344].

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 [2602.04344]. 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\) with residual mask ratio
\[
\rho(z)=\frac{|\mathcal{M}(z)|}{n_g},
\]
where \(n_g\) is the length of the generation segment [2602.04344]. The root node is
\[
z_T=(x_{\text{prompt}},m,\ldots,m),\qquad \rho(z_T)=1,
\]
and terminal nodes are fully unmasked states \(z_0\) satisfying \(\rho(z_0)=0\) [2602.04344].

At a node \(N\) with state \(z\), the action set \(\mathcal{A}\) consists of inference configurations
\[
a=(\theta_a,T_a,g_a),
\]
where \(\theta_a\) selects one of \(k\) pretrained MDLMs, \(T_a\) is a temperature, and \(g_a\) is a deterministic remasking strategy such as entropy-based or low-confidence unmasking [2602.04344]. The paper states that UMF typically uses \(T_a\approx 0\) for determinism.

Executing an action induces a deterministic transition
\[
F_a:z\rightarrow z',
\]
where \(z'\) is obtained by repeatedly applying the MDLM’s single-step unmasking transition under \((\theta_a,T_a,g_a)\) until a lower target ratio is reached. The fixed ratio schedule is
\[
\rho\in\{0.9,0.8,0.7,\ldots,0.2,0\}.
\]
Each single-step unmask consumes \(1\) NFE, while repeated \((N,a)\) pairs are free once cached [2602.04344].

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 \(n_{\text{budget}}\). 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 [2602.04344].

Selection uses UCT. For each child \(N'\) reached from node \(N\) via action \(a\), UMF maintains the parent visit count \(n_{\text{par}}=\mathrm{visits}(N)\), the child visit count \(n_{\text{tgt}}=\mathrm{visits}(N')\), and cumulative reward \(\sum_i r_i\). The UCT score is
\[
\mathrm{UCT}(N')=\frac{\sum_i r_i}{n_{\text{tgt}}}
+c_{\text{exp}}\sqrt{\frac{\log n_{\text{par}}}{n_{\text{tgt}}}},
\]
with \(c_{\text{exp}}=1\) [2602.04344].

During expansion, if \((N,a)\) is already present in the cache \(C\), 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 \(z_0\), evaluates a reward, and caches the result [2602.04344]. 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 \((\theta_a,T_a,g_a)\) is fully deterministic, so
\[
\mathrm{Var}[R]=0.
\]
The paper therefore states that a single simulation yields an unbiased, zero-variance estimate \(Q(N,a)=R\), and that UMF does not learn an explicit policy prior \(\pi(a\mid N)\); instead, the prior effect is implicit in the tree structure [2602.04344]. The score decomposition is
\[
Q(N,a)=\mathrm{mean\_reward}(N,a),\qquad
U(N,a)=c_{\text{exp}}\sqrt{\frac{\log \mathrm{visits}(N)}{\mathrm{visits}(N\rightarrow a)}},
\]
and action selection maximizes \(Q+U\) [2602.04344].

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 [2602.04344].

## 4. Objective, action design, and caching behavior

UMF’s stated objective is to maximize the expected terminal reward \(R(\tau)\) subject to an NFE budget constraint:
\[
\max_{\pi}\;\mathbb{E}_{\tau\sim \pi}[R(\tau)]
\]
under
\[
\sum_t \mathrm{NFE}(\mathrm{transition}_t)\le n_{\text{budget}}.
\]
MCTS with UCT provides the exploration-exploitation mechanism, and inference stops when NFE reaches \(n_{\text{budget}}\) [2602.04344].

The action design is deliberately restricted to deterministic, high-quality branches. The paper specifies that actions are chosen from a small, discrete set \(\mathcal{A}\), for example two MDLMs \(\{\text{Dream-Coder},\text{LLaDA}\}\), each paired with a low temperature and a confidence-based remasking rule [2602.04344]. In the reported experiments, Dream-Coder-v0-Instruct-7B uses \(T=0.1\) with entropy remasking, and LLaDA-8B-Instruct uses \(T=0\) with low-confidence remasking [2602.04344].

Caching is not an auxiliary optimization but an integral part of the algorithmic design. Once a node-action pair \((z,a)\) is visited, UMF deterministically unmaskes all the way to the next ratio, caches the full sub-rollout keyed by \((z,a)\), and guarantees that revisiting the same pair reproduces the identical sub-trajectory at zero NFE [2602.04344]. 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 \(K_t^a\):
\[
\sum_t \mathbb{E}_z[\min_a \epsilon_t^a(z)]
\le
\min_a \sum_t \mathbb{E}_z[\epsilon_t^a(z)].
\]
They argue that interleaving locally best choices can therefore yield lower overall diffusion KL error than any fixed model [2602.04344]. 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-\(N\) 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 \( \mathrm{NFE}=12\,288 \) in the main code comparison [2602.04344].

| Benchmark | Baselines at \( \mathrm{NFE}=12\,288 \) | UMF |
|---|---:|---:|
| LiveCodeBench | BoN best \(\approx 19\%\); DTS* \(\approx 18\%\)–\(18\%\); AB-MCTS \(\approx 21\%\) | \(28\%\) |
| HumanEval+ | BoN \(\approx 75\%\); DTS* \(\approx 75\%\); AB-MCTS \(\approx 81\%\) | \(88\%\) |
| MBPP+ | BoN \(\approx 66\%\); DTS* \(\approx 68\%\); AB-MCTS \(\approx 68\%\) | \(72\%\) |

On MATH, using \(105\) problems, UMF improves from \(49.5\%\) pass@1 at \( \mathrm{NFE}=768 \) to \(60.9\%\) at \( \mathrm{NFE}=12\,288 \), described as an \(11.4\)-point gain over the low-budget baseline [2602.04344].

The scaling results are also explicit. UMF shows smooth, monotonic improvement with NFE up to \(24\,576\), reaching \(30\%\) pass@1 on LiveCodeBench, whereas stochastic baselines plateau or fluctuate [2602.04344]. 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 \(12\,288\) NFE from \(26\%\) to \(28\%\), with a reported cache hit rate of approximately \(56\%\) [2602.04344]. Second, mixing two models outperforms mixing temperatures or mixing remasking strategies by \(1\)–\(4\) points [2602.04344]. Third, a single UMF run interleaving both models at \( \mathrm{NFE}=12\,288 \) outperforms the "Pair" baseline that runs two single-model UMFs at \( \mathrm{NFE}=6\,144 \) each and picks the better answer, with \(24\%\) versus \(28\%\) on LiveCodeBench [2602.04344].

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 [2602.04344]. 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 [2602.04344].

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 [2602.04344].

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 [2602.04344].

Source: https://www.emergentmind.com/topics/unmaskfork-umf