Rethinking Reasoning with MDLMs: Early Exits, Post-hoc Reasoning, and Beyond

Abstract

The reasoning paradigm, where language models reason before answering, has enabled breakthroughs on tasks such as mathematical problem-solving. While current tooling for reasoning is built around next-token prediction trained models, recent works introduce an alternative choice: masked diffusion language models (MDLMs). MDLMs are trained to infill positions in randomly masked sequences. We introduce reasoning-as-infilling, a prompting technique that pre-fills tokens to explicitly delimit reasoning and answer regions, unlocking a unified set of capabilities for MDLM reasoning. Because answer positions are explicitly designated, the model's conditional distributions over answer tokens are directly accessible during generation. This enables early exits when the model is certain of its answer. The same framework supports post-hoc reasoning: given question-answer pairs, MDLMs can sample high-quality reasoning traces from their posterior, a distribution that is intractable for autoregressive models. On GSM8k, fine-tuning LLaDA-8B-Base on these posterior traces improves accuracy by +14.9%, matching gains from human-written traces (+13.4%). Finally, given a reference answer, the answer region distributions enable scoring partial reasoning traces at intermediate steps, providing intermediate rewards that are more strongly correlated with correctness than scores from a specialized process reward model. At intermediate steps, answer-likelihood scores from autoregressive models are significantly less predictive of correctness than those from MDLMs. Our results demonstrate that the MDLM training objective provides promising benefits for reasoning.