No Compute Left Behind: Rethinking Reasoning and Sampling with Masked Diffusion Models (2510.19990v1)

Abstract: Masked diffusion LLMs (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.

• The paper introduces reasoning-as-infilling and Multi-token Entropy Decoding (MED) to enhance inference speed with minimal accuracy loss in math and coding tasks.
• It empirically shows that left-to-right decoding remains competitive, while any-order and multi-token decoding degrade performance except in highly structured tasks like Sudoku.
• Fine-tuning on posterior reasoning traces boosts performance on benchmarks, demonstrating the potential for self-improving masked diffusion models.

Rethinking Reasoning and Sampling with Masked Diffusion LLMs

Introduction and Motivation

This paper critically examines the practical benefits and limitations of Masked Diffusion LLMs (MDLMs) for reasoning and code generation tasks, contrasting them with traditional next-token prediction (NTP) models. While MDLMs are often promoted for their any-order and multi-token decoding capabilities, empirical results demonstrate that, for math and coding benchmarks, left-to-right, one-token-at-a-time decoding is highly competitive, and naive multi-token decoding degrades performance. The authors propose new inference and post-training strategies that leverage the unique properties of MDLMs, specifically their ability to in-fill arbitrary masked positions and provide conditional distributions over all masked tokens.

Figure 1: MDLMs learn conditional distributions at each masked token position, enabling structured infilling, early exits, and adaptive multi-token decoding.

Empirical Analysis of Decoding Strategies

The authors conduct a comprehensive evaluation of decoding strategies on reasoning (gsm8k, math500) and coding (HumanEval, Sudoku) benchmarks using state-of-the-art MDLMs (Dream-7B, LLaDA-8B). The findings are as follows:

• Any-order decoding: On text-based tasks, any-order sampling either underperforms or behaves similarly to left-to-right decoding. Only in highly structured tasks like Sudoku does any-order provide a significant benefit.
• Multi-token decoding: Decoding even two tokens in parallel leads to substantial accuracy drops across all tasks except Sudoku. The joint distribution learned by MDLMs is not well-approximated by the product of marginals, resulting in high KL divergence from the true distribution.

  Figure 2: Left-to-right sampling with MDLMs is a competitive sampling algorithm for reasoning and coding; any-order decoding is only advantageous for tasks like Sudoku.

These results challenge the prevailing narrative that MDLMs' flexibility in decoding order and parallelism is universally beneficial for language tasks.

Reasoning-as-Infilling: A New Paradigm

The paper introduces reasoning-as-infilling, a prompting and inference paradigm that exploits MDLMs' ability to in-fill arbitrary masked positions. By pre-filling a reasoning template that explicitly separates reasoning and answer tokens, several new capabilities are unlocked:

• Early exits: By monitoring the entropy of the answer block during generation, the model can terminate reasoning early when answer uncertainty is low, reducing inference cost without sacrificing accuracy.
• Posterior sampling of reasoning traces: Given a question-answer pair, MDLMs can sample from the posterior over reasoning traces conditioned on the answer, providing high-quality synthetic data for post-training.
• Scoring intermediate reasoning: The model can score partial reasoning traces for their likelihood of leading to the correct answer, enabling dense feedback for fine-tuning without external verifiers.

  Figure 3: MDLMs enable scoring their own reasoning process without an external process verifier.

Multi-token Entropy Decoding (MED)

To address the inefficiency of single-token decoding while avoiding the pitfalls of naive multi-token sampling, the authors propose Multi-token Entropy Decoding (MED). MED adaptively selects masked positions to decode in parallel based on their conditional entropy, only decoding multiple tokens when their uncertainty is below a threshold. This approach provides a principled upper bound on the KL divergence between the true joint and the factorized distribution, ensuring that parallel decoding does not significantly degrade sample quality.

Figure 4: MED enables parallel token decoding without any loss in performance.

Empirical results show that MED achieves $2$–$3\times$ speedups in inference with negligible loss in accuracy on gsm8k and HumanEval. When combined with early exits, total speedup reaches $3.3\times$ with only a $0.1\%$ drop in accuracy.

Post-training with Posterior Reasoning Traces

A key contribution is the demonstration that MDLMs can generate high-quality reasoning traces by sampling from the posterior $p_\theta(r \mid c, a)$, where $r$ is the reasoning trace, $c$ is the context, and $a$ is the answer. Fine-tuning on these posterior traces yields performance improvements on par with, or exceeding, fine-tuning on human-annotated traces.

• On gsm8k, fine-tuning LLaDA-8B Base on posterior reasoning traces improves test accuracy by $+14.9\%$, outperforming fine-tuning on gold traces.
• Posterior traces judged by GPT-4o and Qwen2.5-Math-PRM are rated as correct in $\sim40\%$ of cases where the original model failed, indicating the utility of this approach for data augmentation.

  Figure 5: Distribution of answer block log probability scores for posterior samples; correct traces have higher scores.

  

  Figure 6: Answer block log probability scores predict posterior reasoning trace quality, enabling filtering of low-quality samples.

Scoring and Filtering Reasoning Traces

The ability to compute answer block log probabilities for partial reasoning traces allows MDLMs to provide dense, model-internal feedback for the quality of reasoning, without requiring external process reward models. This enables:

• Early termination of low-quality chains
• Filtering of synthetic data for post-training
• Potential for intermediate reward shaping in RLHF-style fine-tuning

Theoretical and Practical Implications

The findings have several implications:

• Theoretical: The results highlight a gap between the expressive power of MDLMs and the practical utility of their any-order/multi-token decoding for text. The proposed entropy-based bounds provide a principled framework for safe parallel decoding.
• Practical: Reasoning-as-infilling and MED offer concrete methods to reduce inference cost and improve data efficiency in post-training. The ability to generate and score posterior reasoning traces opens new avenues for self-improving models and dense supervision.

Limitations and Future Directions

MDLMs incur higher computational cost than NTP models, scaling as $\mathcal{O}(\ell^2)$ for sequence length $\ell$, and current architectures lack efficient caching for any-order decoding. Further research is needed to make MDLMs practical for long-context tasks and to develop architectures that retain their unique capabilities with lower overhead.

Conclusion

This work provides a rigorous empirical and theoretical analysis of MDLMs for reasoning and code generation. While the anticipated benefits of any-order and multi-token decoding are limited for text, the unique properties of MDLMs enable new inference and post-training strategies—reasoning-as-infilling and MED—that yield substantial practical gains. The ability to sample and score posterior reasoning traces without external verifiers is particularly promising for scalable, self-improving LLMs. Future work should focus on architectural innovations to mitigate computational costs and further exploit the distinctive capabilities of MDLMs.