UnMaskFork scales masked diffusion LLMs by switching models mid-generation

Most AI text generators write left to right, one token at a time. Masked diffusion language models work differently: they start with an answer made entirely of blanks and fill them in over many steps, committing first to the positions they are most confident about. A standard way to squeeze better answers out of ordinary models is Best-of-N, which generates many answers with added randomness and keeps the one that scores best on a check such as unit tests. Sakana AI researchers found that this works poorly for the diffusion coding model Dream-Coder. Raising its randomness, or shuffling the order in which it fills blanks, damaged answers so much that the best of many failed to beat a single answer generated without randomness, as their project page shows.

UnMaskFork, an ICML 2026 paper, gets variety another way: by handing a half-finished answer from one diffusion model to another. A partly filled sequence has the same structure for any such model, so Dream-Coder can fill some blanks and LLaDA can continue from there, and because the two models are confident about different things, each hand-off leads to a different answer. The paper fixes checkpoints where control can switch, treats each sequence of model choices as a branch in a tree, and uses Monte Carlo tree search, which spends more effort on branches whose finished answers score well, to find a good path. Each model decodes with almost no randomness, so a given branch always produces the same text, and partial results can be cached and reused. On LiveCodeBench, about half of the search's rollouts reused cached work once the budget reached 3,072 model calls.

On 100-problem subsets of three coding benchmarks, with every method given the same 12,288 model calls and scored on held-out tests, UnMaskFork with Dream-Coder and LLaDA, open models of 7 and 8 billion parameters, scored 28.0 on LiveCodeBench, 88.0 on HumanEval+ and 72.0 on MBPP+. Best-of-N, run on each model separately with the better answer kept, scored 19.0, 75.0 and 66.0. The gain is not only from having two models that solve different problems: running single-model UnMaskFork on each model independently, with the same total budget, also did worse than letting the models share one answer. Adding a third model, DiffuCoder-cpGRPO, raised the three-benchmark average from 62.7 to 65.0, and the paper reports gains on a 105-problem sample of the MATH dataset, where a separate reward model scored candidate answers.

The method needs a way to score finished answers during the search, such as test cases for code or a reward model for math, and it spends extra computation at answer time, which the authors note raises cost and latency. It needs no retraining, so each new diffusion model trained on different data adds options for the search. Sakana makes the same argument for Fugu, its service that coordinates several large models behind one API.