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

AI summary

Overview

Research area: Machine learning / natural language processing — masked diffusion language models (MDLMs) applied to multi-step reasoning, covering inference-time acceleration and post-training.

Technical level: Advanced. The paper assumes familiarity with autoregressive language modeling, masked diffusion / infilling objectives, decoding-order-based sampling, entropy, and process reward models.

Scope: The paper introduces "reasoning-as-infilling," a prompting framework that pre-fills a reasoning template separating reasoning and answer regions, and shows it enables early exiting, posterior sampling of reasoning traces from a known answer, and intermediate scoring of partial reasoning traces for MDLMs such as LLaDA-8B and Dream-7B.

What This Paper Is About

Most reasoning tooling for language models is built around next-token prediction (NTP) models, which generate left to right and cannot condition on tokens that appear later in the sequence. MDLMs are trained instead to infill randomly masked positions, so their conditional distributions at every masked position are directly accessible. The paper asks what this training objective buys for reasoning, and answers with a single prompting change — pre-filling an explicit reasoning template that delimits reasoning and answer regions — that unlocks early exits, post-hoc (posterior) reasoning generation, and step-level scoring.

Key Contributions

1. Reasoning-as-infilling. A prompting technique for MDLMs that pre-fills tokens to explicitly delimit a reasoning block and an answer block (separated by a user-specified delimiter, e.g. "The answer is: " for math or function definitions for coding), making the model's conditional distributions over answer tokens accessible during generation.

2. Early exits. Because answer positions are explicitly designated, answer-region uncertainty can be measured mid-generation using an upper bound on answer-block joint entropy, computed as the sum of marginal entropies of answer tokens. On GSM8k with one-token decoding, this yields a 1.3× acceleration with less than a 1 percentage point reduction in accuracy, and it composes with parallel decoding algorithms and variable-length diffusion models.

3. Post-hoc reasoning. Pre-filling the answer block lets an MDLM sample directly from the posterior p_θ(r | c, a) over reasoning traces, which the paper notes is intractable for NTP models. On GSM8k, fine-tuning LLaDA-8B-Base on these posterior traces improves accuracy by +14.9%, comparable to the +13.4% from the human-annotated GSM8k reasoning traces.

4. Scoring partial reasoning traces. Given a reference answer, the answer-block distributions score partial reasoning traces at intermediate steps (an intermediate score denoted φ_iprm), providing intermediate rewards more strongly correlated with correctness than scores from a specialized 7-billion-parameter process reward model.

Main Findings

• Early exits trade little accuracy for speed. On GSM8k with LLaDA-8B-Instruct and one-token entropy decoding, a no-exit baseline reaches 79.4% accuracy; with a threshold γ = 0.1 the model reaches 1.326× speed-up at 78.6% (a 0.8 percentage point drop). On Math500, LLaDA with no exit reaches 33.4% and with γ = 0.3 reaches 1.158× at 31.9%. With no reasoning template, LLaDA reaches 76.6% on GSM8k and 33.8% on Math500.

• Early exits compound with multi-token decoding. For LLaDA on GSM8k, block parallel confidence decoding (bpcd) alone gives 3.282× at 79.0%; adding early exits (γ = 0.1) gives 4× total at 78.9%, described as a further 1.2× over bpcd alone. For Dream-7B-Instruct, entropy decoding with no exit reaches 79.8%, and γ = 0.7 gives 1.1378× at 76.7%; Dream requires higher exit thresholds for speed-ups, which the authors suggest may be due to Dream's adaptation from an NTP model.

• Any-order decoding has limited benefit for reasoning. Naive or block any-order decoding either samples a large portion of tokens in left-to-right order (nearly 50% for GSM8k) or underperforms left-to-right sampling. The authors also describe a "pathological" behavior where MDLMs decode an incorrect answer first and follow it with a post-hoc rationale.

• Posterior traces are high quality even when the model fails. On the GSM8k training set, LLaDA-8B-Instruct (with multi-token entropy decoding and any-order decoding) answers 1419 of 7473 problems incorrectly. On those 1419 question-answer pairs, both judges — Qwen2.5-Math-PRM and GPT-4o — rate roughly 40% of the instruct-tuned model's posterior reasoning chains as correct (scores of 0.38 and 0.43), and even the base model's posterior chains are rated above 31% correct (0.31 and 0.36). Posterior conditioning also outperforms STaR-style prefix-hint prompting on the same checkpoints.

• Fine-tuning on posterior data matches human traces. LLaDA-8B-Base with no template scores 13.9% on the GSM8K test set; reasoning-as-infilling alone raises it to 51.2%. LoRA fine-tuning on GSM8k human-written data (n = 7473) reaches 64.6% (+13.4), while LoRA fine-tuning on GSM8k posterior data (n = 7473) reaches 66.1% (+14.9). For reference, LLaDA-8B-Instruct with miscellaneous instruction data (n = 4.5 million) reaches 75.96%.

• Fine-tuning also improves early exiting. The posterior-finetuned model holds 66.2% at γ = 0.3 (1.2×) and 64.7% at γ = 0.4 (1.5×); the human-trace-finetuned model holds 58.9% at γ = 0.3 (1.4×) and 54.1% at γ = 0.4 (1.7×). The paper summarizes this as early exits yielding 1.5× acceleration with less than a 1.4 percentage point drop in accuracy. Early exits have limited effect on the original base model even at high thresholds.

• Intermediate MDLM scores beat a dedicated PRM. Using LLaDA-8B-Instruct on the GSM8k test set with greedy left-to-right one-token sampling, gold-answer log probabilities from the answer block correlate more strongly with final correctness at intermediate steps than scores from Qwen2.5-Math-PRM, a 7B process reward model. Even without gold labels, MDLM answer-block entropy is weakly correlated with correctness.

• Autoregressive intermediate scores are weaker. Appending "The answer is \boxed{[ANSWER]}" to a partial reasoning trace lets autoregressive models (Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct) compute answer likelihoods in one forward pass, but their intermediate scores are significantly less predictive of correctness at early and intermediate steps than both φ_iprm and the dedicated PRM; at early steps they are less predictive than label-free MDLM answer entropy.

• Comparison with Prophet on exits. On GSM8k the methods are matched at comparable NFE budgets (74.8% vs. 74.3% at 1.55×). On CommonsenseQA, entropy exits are both more accurate and faster at every evaluated threshold (76.2% at 1.86× vs. 72.2% at 1.02×). Full sweeps are reported in Appendix G.

• Works with variable-length diffusion models. DreamOn-7B, fine-tuned for variable-length decoding on a coding dataset, scores 1.4% on GSM8k with no template and no exit. With reasoning-as-infilling and no exit it reaches 60.6% (fixed length) or 66.4%; enabling deletions gives a 6% reduction in sampling time, and adding early exits (γ = 0.5) gives an additional 37% reduction in sampling time while maintaining accuracy (1.8686× at 64.7%, and 1.6732× at 66.0% in the variable-length setting). The paper notes the model achieves these results despite being fine-tuned on coding rather than math.

Methodology in Plain English

The researchers leave the MDLM's training objective untouched and change only how the model is prompted at generation time. Instead of putting all user context in a prefix, they pre-fill the output sequence with a template: a block of masked reasoning positions, a delimiter, and a block of masked answer positions. Because MDLMs predict every masked position conditioned on all unmasked ones, designating the answer block means the model's probability distribution over the answer is readable at any point during generation.

Three capabilities follow from this. For early exits, they sum the entropies of the answer-block marginals to get an upper bound on the joint answer entropy, and stop unmasking reasoning tokens once that bound drops below a threshold γ. For post-hoc reasoning, they pre-fill the correct answer into the answer block and let the model sample the reasoning block, effectively drawing from the posterior over reasoning given question and answer, then use those traces for supervised fine-tuning. For scoring, they take a partial reasoning trace plus a reference answer and sum the log probabilities the model assigns to the answer tokens, producing a per-step reward. They test these ideas on GSM8k and Math500 with LLaDA-8B-Instruct, LLaDA-8B-Base, Dream-7B-Instruct, and DreamOn-7B, using a generation length of 256, block size 32, the delimiter "The answer is boxed{..", and 10 answer tokens, and they compare samplers including entropy-bound (eb), block parallel confidence decoding (bpcd), and a proposed multi-token entropy method called med. Posterior traces are judged by GPT-4o and Qwen2.5-Math-7B PRM.

Why This Matters

The paper argues that the extra compute MDLMs spend modeling all masked positions can be repurposed: the same forward pass that generates reasoning also yields live answer uncertainty and answer-conditioned distributions that NTP models simply do not expose. This reframes MDLM reasoning as a design space rather than a curiosity, and it suggests bootstrapping reasoning data without human traces or a separate reward model.

Real-world applications:

• Cheaper reasoning inference. Skipping remaining reasoning steps once the answer is certain, combined with parallel decoding, cuts cost on math and word-problem workloads.
• Self-generated training data. Posterior traces offer a substitute for expensive human-written rationales when only question-answer pairs are available.
• In-training reward signals. Step-level answer-likelihood scores can be used for steering, reflection tokens, and filtering during post-training without training a separate process reward model.
• Adaptive reasoning budgets. Integration with variable-length diffusion and deletion operations lets systems spend fewer or more tokens per query.

Industry relevance: providers serving reasoning models can reduce inference cost per query; teams building RL or GRPO-style post-training pipelines can obtain dense intermediate rewards and synthetic traces from the model itself, lowering dependence on human annotation and auxiliary reward models, and lowering GPU time relative to full-length generation.

Future Directions

• Extending early exits and post-hoc reasoning to tasks with long or open-ended answers; the authors explicitly limit evaluation to tasks with short answers.
• Removing the dependence on question-answer pairs, since both posterior trace generation and intermediate scoring require an answer, though the authors note this assumption is standard in RL with verifiable rewards such as GRPO.
• Improving the speed-ups for models adapted from NTP backbones (such as Dream), which currently need higher exit thresholds than LLaDA.
• Turning the intermediate scores into training objectives — the paper sketches terminating low-quality chains early, steering toward correct solutions, auto-inserting reflection tokens at failures, and incorporating dense intermediate feedback into fine-tuning.

Target Audience

Researchers and engineers working on diffusion language models, efficient reasoning inference, or synthetic reasoning-data generation for post-training will get the most from this paper. It is also relevant to practitioners building process reward models or RL-based post-training pipelines, since it offers an alternative to a separately trained PRM, and to anyone tracking alternatives to next-token prediction as the default language modeling paradigm.

Authors’ 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.