Sakana AI's UnMaskFork Gives Diffusion Language Models Test-Time Scaling Power

Test-time scaling — giving a model more compute at inference to make it smarter — powers reasoning models, Best-of-N sampling, and MCTS-based search. All of that work has been built around autoregressive LLMs, which generate text one token at a time, left to right. A fundamentally different class of model has been gaining ground, and until now nobody had figured out how to give it the same scaling treatment.

Sakana AI's UnMaskFork (UMF), accepted at ICML 2026, is the first test-time scaling method designed specifically for Masked Diffusion Language Models (MDLMs). Instead of generating text sequentially, MDLMs start with a fully masked sequence and iteratively fill in tokens in parallel, seeing the entire context at once. The paper shows that having multiple MDLMs collaborate on generating a single answer — each picking up where the others left off — consistently outperforms every existing test-time scaling baseline on coding benchmarks.

Why the Standard Playbook Breaks

The usual tricks don't transfer. In autoregressive models, diversity comes for free: raise the temperature (a parameter that flattens the probability distribution over tokens, making the model more willing to pick less-likely words), generate N different answers, and pick the best one. That's Best-of-N, and it's a workhorse of test-time scaling.

MDLMs like Dream-Coder and LLaDA don't work that way. The Sakana team tested Best-of-N on Dream-Coder and found that raising the temperature or randomizing the unmasking order sharply degrades answer quality. Best-of-N built on either source of randomness failed to surpass a single deterministically decoded answer. The controllable randomness that makes Best-of-N work for GPT-style models actively breaks MDLM generation. These models need to commit to high-confidence tokens first; injecting noise at that stage corrupts the whole sequence.

Diversity Through Model Switching

UnMaskFork's solution is to create diversity through model switching rather than randomness. Every MDLM operates on the same intermediate representation — a sequence of partially unmasked tokens — so one model can hand off to another mid-generation. Dream-Coder might fill in the first third of a sequence with high confidence, then LLaDA takes over for the middle section, then Dream-Coder finishes. The handoff is seamless because the partially-filled token sequence is model-agnostic.

Different MDLMs produce different predicted tokens and prediction confidences, so the generated text diverges from what any single model would have produced. Because MDLMs learn to generate from partially filled states during training, they can take over text partially generated by another MDLM and continue from there. This structural diversity replaces temperature-based diversity.

MCTS as the Search Engine

Generating every possible model-switching pattern would be computationally explosive. UMF uses Monte Carlo Tree Search (MCTS) — a search algorithm that balances exploration of new paths with exploitation of promising ones — to navigate the space efficiently. A single search iteration works as follows:

1. Select: Walk down the existing tree, guided by UCT (Upper Confidence bound for Trees), which prioritizes branches with high scores while also visiting less-explored ones.
2. Expand: At a node with unexplored actions, apply one action — let a chosen MDLM unmask the sequence one stage further — and roll out to a complete answer.
3. Backup: Score the completed answer using test cases or a reward model, then propagate that score back up the tree to inform future selections.

Each action uses near-deterministic decoding at near-zero temperature, so repeating a rollout from the same state with the same action reproduces essentially the same text. That determinism means the partially masked intermediate states traversed during rollout can be cached and reused, letting subsequent MCTS steps skip redundant computation. The caching is what makes the search tractable.

Compute budget is measured in NFE (Number of Function Evaluations) — the number of times an MDLM scores the full token sequence. This is the standard cost metric for diffusion models, analogous to the number of forward passes in an autoregressive model.

Benchmark Results

UMF was evaluated on three coding benchmarks — LiveCodeBench, HumanEval+, and MBPP+ — using Dream-Coder and LLaDA as the two base models. Results at a fixed budget of 12,288 model calls:

Adding a third model, DiffuCoder-cpGRPO, lifts LiveCodeBench from 28.0 to 32.0 and MBPP+ from 72.0 to 76.0, raising the average across the three tasks from 62.7 to 65.0. More models means richer switching patterns to explore. Extending the budget to NFE=24,576 pushes UMF to 30.0 on LiveCodeBench with two models.

UMF also outperforms Best-of-N Pair and DTS* Pair, which already incorporate a model ensemble effect. The gains come from having multiple models share the generation process, not from independently combining their outputs.

What Collaboration Looks Like in Practice

A worked example from LiveCodeBench shows how the collaboration actually unfolds. The winning solution follows a switching path of D→L→L→L→D: Dream-Coder initiates the solution by outlining implementation steps, LLaDA continues that outline, fills in the explanation at the end of the code, and begins the implementation, then Dream-Coder takes over and completes the remaining code. Neither model alone could solve the problem — the correct answer emerged only from that specific handoff sequence.

Where It Fits and What It Can't Do

UMF is most useful when:

• You're working with MDLMs on tasks where answers can be verified, such as coding or math
• You have access to multiple pre-trained MDLMs trained on different data or with different methods
• You want to improve performance without additional training or fine-tuning
• You're willing to trade inference compute for accuracy

The method has real limitations. Search cost scales with the number of models and the number of checkpoints (switch points) defined. The authors suggest a future direction of learning which model should handle which part of generation, which could reduce that cost significantly. UMF also requires a verifier — test cases or a reward model — to score candidate answers, so it won't work out of the box for open-ended generation tasks where correctness is hard to define. And the MDLM ecosystem is still young; the method's value grows as more diverse models become available.

The paper also confirms that UMF extends to block diffusion, with experiments showing substantial gains on LiveCodeBench (19.0 to 31.0), establishing broader applicability beyond standard MDLMs.

Where This Fits in the Diffusion LLM Story

Diffusion LLMs like LLaDA and Dream have reached performance on par with similarly sized autoregressive LLMs across many benchmarks, and commercial-scale diffusion models like Mercury have shown that diffusion-based code generators can rival top autoregressive code models while generating text faster. The architecture is maturing quickly. What it has lacked is the test-time scaling infrastructure that has made autoregressive models so much more capable in practice.

UMF fits into a broader research program at Sakana. The company has been exploring collective intelligence — coordinating multiple models to produce capabilities no single model has alone. Their AB-MCTS work applied this to autoregressive LLMs; their commercial Sakana Fugu product takes it to frontier model orchestration. UnMaskFork extends the same approach into the diffusion paradigm.

No code or model weights have been released alongside the paper, but the method requires only pre-trained MDLMs and a scoring function, making it straightforward to implement on top of any existing MDLM stack. As the diffusion LLM ecosystem grows, every new MDLM trained on different data becomes another branching option the search can exploit.