Sakana AI's UnMaskFork Gives Diffusion Language Models Test-Time Scaling Power
Sakana AI's UnMaskFork uses MCTS and multi-model collaboration to unlock test-time scaling for diffusion language models, beating all baselines on coding benchmarks without any extra training.
- What happened: Sakana AI's UnMaskFork (UMF) introduces test-time scaling for masked diffusion language models via multi-model collaboration.
- Core insight: Instead of temperature-based diversity (which breaks MDLMs), UMF creates diversity by switching between models mid-generation, using their shared intermediate token representation.
- Search mechanism: Monte Carlo Tree Search navigates the space of model-switching patterns, with caching of intermediate states to avoid redundant computation.
- Results: UMF scores 28.0 on LiveCodeBench with 2 models and 32.0 with 3 models at equal compute, beating all baselines including multi-model AB-MCTS (21.0).
- No training required: Works purely at inference time using pre-trained MDLMs; adding more diverse models directly improves performance.
- Limitation: Requires a verifier (test cases or reward model) to score answers, so it's currently limited to tasks with checkable correctness like coding and math.
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.