No Compute Left Behind: Rethinking Reasoning and Sampling with Masked Diffusion Models

TL;DR

This paper examines masked diffusion language models (MDLMs) that infill randomly masked tokens and questions the practical gains of any-order and multi-token decoding on math and coding tasks. It introduces reasoning-as-infilling, a framework that pre-fills a reasoning template to separate reasoning from the final answer, enabling early exits, posterior reasoning traces, and post-training, alongside Multi-token Entropy Decoding (MED) for adaptive parallel decoding. The authors demonstrate that MED yields 2–3x speedups with minimal performance loss and that posterior reasoning traces can significantly boost fine-tuning performance, sometimes matching improvements from human-annotated traces. Overall, the work shows that the extra inference and training capabilities unlocked by MDLMs can be harnessed for efficient inference, data generation, and model improvement, motivating new directions for future MDLM architectures and tooling.

Abstract

Masked diffusion language models (MDLMs) are trained to in-fill positions in randomly masked sequences, in contrast to next-token prediction models. Discussions around MDLMs focus on two benefits: (1) any-order decoding and 2) multi-token decoding. However, we observe that for math and coding tasks, any-order algorithms often underperform or behave similarly to left-to-right sampling, and standard multi-token decoding significantly degrades performance. At inference time, MDLMs compute the conditional distribution of all masked positions. A natural question is: How can we justify this additional compute when left-to-right one-token-at-a-time decoding is on par with any-order decoding algorithms? First, we propose reasoning-as-infilling. By using MDLMs to infill a reasoning template, we can structure outputs and distinguish between reasoning and answer tokens. In turn, this enables measuring answer uncertainty during reasoning, and early exits when the model converges on an answer. Next, given an answer, reasoning-as-infilling enables sampling from the MDLM posterior over reasoning traces conditioned on the answer, providing a new source of high-quality data for post-training. On GSM8k, we observe that fine-tuning LLaDA-8B Base on its posterior reasoning traces provides a performance boost on par with fine-tuning on human-written reasoning traces. Additionally, given an answer, reasoning-as-infilling provides a method for scoring the correctness of the reasoning process at intermediate steps. Second, we propose multi-token entropy decoding (MED), a simple adaptive sampler that minimizes the error incurred by decoding positions in parallel based on the conditional entropies of those positions. MED preserves performance across benchmarks and leads to 2.7x fewer steps. Our work demonstrates that the training and compute used by MDLMs unlock many new inference and post-training methods.

Figures (6)

• Figure 1: learn the conditional distributions at each masked token position. A) We reframe reasoning as infilling a prompted reasoning template, which enables directly modeling answer token probabilities during reasoning. This provides several benefits, like B) enabling early exits or post-hoc reasoning given a pre-filled answer. C) We also utilize the entropy of these distributions to adaptively set the number of tokens decoded at each step.
• Figure 2: Left-to-right sampling with is a competitive sampling algorithm for reasoning and coding. When performing entropy decoding ye2025dream, we observe that full any-order sampling results in poor performance on all tasks but Sudoku. Left-to-right block decoding is required to make any-order sampling performant, and left-to-right sampling (block size $=1$) is always within a few percent of the best configuration. We also observe in \ref{['tab:decoding_strategies']} in \ref{['appsec:observations']} that performant block any-order configurations sample a large portion of tokens left-to-right. We consider sequences of length $128$.
• Figure 3: enables parallel token decoding without any loss in performance. We compare decoding with $\lambda \in \{0.1, 0.2, 0.3\}$ to entropy decoding with a fixed number of tokens $k \in \{1, 2\}$. We observe that significantly reduces the number of NFEs while matching accuracy and maintaining a low KL.
• Figure 4: enable scoring their own reasoning process without an external process verifier. We score gsm8k reasoning traces, generated left-to-right with LLaDA, at intermediate steps using (a) gold answer log probabilities, (b) the answer block entropy bound, and (c) an 8B process reward model (PRM) zhang2025lessonsdevelopingprocessreward. Gold answer probability at intermediate steps is more predictive of final correctness than PRM scores. Even without gold labels, the answer block entropy is weakly correlated with correctness.
• Figure 5: Distribution of answer block log probability scores for posterior samples. To score the $1419$ posterior reasoning traces, we iteratively unmask each chain left-to-right, and compute the average answer block log probabilities across all timesteps. Posterior chains rated as correct by GPT4-o tend to have higher scores.