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

Test-Time Scaling in Masked Diffusion Language Models

Test-Time Scaling (TTS) is a well-established approach to enhancing the capabilities of LLMs on tasks requiring reasoning, such as coding and mathematics, without the need for additional training. There are various TTS methods, such as allocating more test-time compute for the LLM to think longer or scaling up the number of generated candidate answers. Among these, Best-of-N is a simple yet effective method widely adopted for autoregressive LLMs (AR-LLMs). Best-of-N is applicable to tasks where the correctness of the answer can be verified, such as coding.

Best-of-N leverages the property that increasing the temperature parameterTemperature $T$ is a parameter that smooths the token probability distribution output by the model. As $T$ increases, the distribution becomes flatter, making low-probability tokens more likely to be sampled. Conversely, in the limit of $T \to 0$, generation becomes deterministic, always selecting the highest-probability token. allows LLMs to generate diverse answers without significantly compromising quality. The model generates N answers for the same problem at a high temperature, the outputs are scored with a verifier, and the highest-scoring candidate is selected. This seemingly simple method has been empirically shown to be highly effective for AR-LLMs.

So, is Best-of-N also effective for Masked Diffusion Language Models (MDLMs), which, unlike AR-LLMs, do not generate text from left to right? We tested this with Dream-Coder, an MDLM with strong coding performance. In addition to temperature, MDLMs have a second natural source of diversity: injecting randomness into the unmasking orderIn standard MDLM decoding, the model's prediction confidence at each mask position is measured using metrics such as the entropy of the predicted distribution or the probability assigned to the proposed token, and masked positions are unmasked starting from those with the highest confidence.. As shown below, Best-of-N built on either source of randomness failed to surpass the single answer obtained by standard deterministic decoding.

How can we achieve TTS in MDLMs without relying on temperature or unmasking order? To answer this question, we developed UnMaskFork (UMF).

The Idea Behind UnMaskFork

The idea behind UnMaskFork is as follows: rather than relying on randomness from temperature or unmasking order to produce diverse answers, it creates diversity by having multiple MDLMs share the generation of an answer and by varying how that generation is divided among them.

To understand how multiple MDLMs can collaboratively generate text, let us first consider an intermediate state during generation. In this state, mask tokens and tokens already filled with text coexist. When a single MDLM generates text, masked positions are unmasked in descending order of the model's prediction confidence. An important point here is that this intermediate representation has the same structure regardless of which MDLM is used.When two MDLMs use different tokenizers, the already-generated text is converted between them by decoding it with the source tokenizer and re-encoding it with the target tokenizer; special tokens such as masks are mapped directly between the models. Because of this shared structure, it is possible to hand the sequence to another MDLM, which then continues generating the answer.

From that point on, since different MDLMs have different predicted tokens and prediction confidences, the generated text will generally differ from what a single model would have produced. Since MDLMs acquire the ability to generate text from partially filled states during their training process, they can take over text partially generated by another MDLM and continue generating. This enables multiple MDLMs to share the generation process.

Exploring Effective Generation-Sharing Strategies for MDLMs

In UMF, rather than simply generating a large number of answers randomly as in Best-of-N, we achieve efficient TTS by using Monte Carlo Tree Search (MCTS) to preferentially explore promising sharing strategies. Below, we explain how this search is conducted.

For example, consider the case of two MDLMs, Dream-Coder and LLaDA. First, starting from a fully masked sequence, we fix in advance the checkpoints at which the model in charge can be switched. Suppose we set two checkpoints, at mask ratios of 66% and 33%. In this case, various unmasking sequence patterns can be considered, such as D→D→L or L→D→L, and these patterns can be represented as a tree structure as shown in the figure above. The state of each node is a token sequence in which some of the masks have been filled in with text, and the actions selectable from a node correspond to the choice of MDLM. An action advances the unmasking of the token sequence by the chosen MDLM and transitions to a new node.

TTS via UMF is achieved by starting from a tree containing only the root node and building this tree using MCTS. Each MCTS iteration consists of three steps: (1) Select: walking down the tree, guided by a rule called UCT that prioritizes promising directions, to a node with unexplored actions, (2) Expand: applying one of those actions to add a new child node, immediately rolling the unmasking out to a complete answer, and evaluating it, and (3) Backup: propagating the resulting score up the tree to inform subsequent selections.

Furthermore, an important feature of UMF is that each action uses near-deterministic decoding at near-zero temperature: repeating a rollout from the same state with the same action reproduces essentially the same text. Thanks to this determinism, the partially masked intermediate states traversed during the rollout stage can be cached and reused, allowing subsequent MCTS steps to skip redundant computation.

UMF Consistently Outperforms Other TTS Methods

We experimentally confirmed that UMF indeed outperforms other TTS approaches. When comparing different TTS methods, we evaluate the performance achieved on a given task under the same inference budget. For diffusion models, this budget is commonly measured by the Number of Function Evaluations (NFE). NFE is defined as the number of times the MDLM predicts token probabilities for the entire token sequence, meaning a higher NFE corresponds to a larger computational cost.

The figure above illustrates the performance of UMF across three coding tasks. We can see that UMF consistently outperforms other standard TTS methods. While other TTS approaches that require raising the temperature suffer from a significant performance drop in the low-NFE regime, UMF avoids increasing the temperature, allowing it to maintain consistently high performance even in that regime. Furthermore, by leveraging the diversity originating from multiple MDLMs, UMF continues to improve its performance as NFE increases.

UMF Is More Than a Simple Model Ensemble

To understand the gains from UMF, we compared UMF against other TTS methods that incorporate model ensembling. A simple model ensemble effect works as follows: if, say on LiveCodeBench, some problems can only be solved by Dream-Coder and others only by LLaDA, performance improves just by letting the two models solve each problem independently and keeping the better-scoring of the two answers. The table below compares UMF against non-UMF TTS methods augmented with this ensemble effect.

"Best-of-N Pair" and "DTS* Pair" incorporate the aforementioned model ensemble effect into TTS methods like Best-of-N and DTS*. Since UMF outperforms these methods, we can see that the performance improvement is not driven merely by the effect of independently ensembling models. Instead, having multiple models share the generation process creates structural diversity in the answers, and leveraging this diversity is what improves performance.

Furthermore, UMF substantially outperforms AB-MCTS, which also uses multiple models. This indicates that even among TTS methods utilizing structural diversity, UMF can explore answers more efficiently by restricting branching to a small set of deterministic actions and by skipping redundant computation through caching.

Increasing the Number of Models Can Further Improve Performance

The diversity introduced by sharing the generation process in UMF can be further expanded by increasing the number of participating MDLMs. The table below shows the performance when DiffuCoder-cpGRPO is added as a third model to the action set of the two-model UMF consisting of Dream-Coder and LLaDA.

In the three-model UMF with DiffuCoder-cpGRPO added, LiveCodeBench improves from 28.0 to 32.0 and MBPP+ from 72.0 to 76.0 (with HumanEval+ remaining mostly flat), raising the average across the three tasks from 62.7 to 65.0. This suggests that as the number of available models increases, the sharing patterns that can be explored become richer, which can further boost the performance of UMF.

Each MDLM Contributes Where It Is Confident

The figures above show an example of the search tree and the generated text when a problem is successfully solved by UMF on LiveCodeBench. As can be seen from the search tree, the problem is solved by following a complex sharing path of D→L→L→L→D, while every other explored branch reaches a much lower reward. Such difficult problems become solvable not by relying on either model alone, but by having each model generate the parts it is most confident about and performing TTS using structural diversity. Looking at the generated text in the second figure, we can see that the answer is not simply written from beginning to end. Instead, Dream-Coder initiates the solution by outlining the implementation steps in the first stage; LLaDA then continues that outline, fills in the explanation at the end of the code, and begins the implementation; and Dream-Coder finally takes over and completes the remaining code.

Conclusion and Future Work

Because MDLMs behave differently from autoregressive LLMs, they call for TTS methods that exploit those differences, and UMF is one such method. UMF opens up a new direction for TTS in MDLMs by showing that we can leverage the structural diversity achieved by having multiple MDLMs share the answer generation process. This takes advantage of the fact that the intermediate states of MDLM generation share the same structure across models and can be handed from one model to another. Furthermore, the paper shows that UMF also applies to block diffusion, with experiments confirming substantial gains (19.0 → 31.0 on LiveCodeBench), establishing it as a highly versatile method.

Moreover, UMF requires no additional training or internal modifications to the models; it only needs pre-trained MDLMs and a way to score candidate answers at inference time, such as test cases or a reward model. Differences between models trained on different data or with different methods directly translate into a richer diversity of sharing patterns that can be explored. Therefore, we believe the value of UMF will continue to grow as a wider variety of MDLMs emerge. To further advance the multi-model collaboration demonstrated by UMF, a potential future direction is to reduce the computational cost required for the search, for instance, by learning which model should be responsible for which part.

This work is part of our broader research pursuing the "collective intelligence of AI," much like AB-MCTS and Sakana Fugu, which coordinate multiple LLMs. UMF shows that this philosophy is effective not only for AR-LLMs but also within the new paradigm of MDLMs. We will continue to advance research that transforms model diversity into strength.