arXiv:2510.19990cs.LG2025-10被引 5

用掩码扩散模型实现更高效、可解释的推理,支持提前退出和事后推理。

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

  • 将推理过程视为填空任务,显式划分推理与答案区域。
  • 在GSM8k上通过后验采样提升准确率14.9%,接近人工标注效果。
  • 可实时评估中间步骤正确性,适合强化学习与高效推理场景。

推理范式(语言模型先推理再作答)已在数学问题求解等任务中取得突破。现有工具基于逐词预测训练的模型,而近期工作提出新选择:掩码扩散语言模型(MDLMs),其训练目标是填充随机掩码序列中的位置。本文提出“推理即填空”提示方法,通过预先填充标记推理与答案区域,解锁MDLM统一推理能力。由于答案位置明确,生成过程中可直接获取答案条件分布,支持模型在确信时提前退出。同一框架支持事后推理:给定问答对,MDLM可从其后验分布中采样高质量推理路径,而自回归模型难以实现此操作。在GSM8k数据集上,对LLaDA-8B-Base微调后验轨迹,准确率提升+14.9%,接近人工标注轨迹的+13.4%增益。此外,给定参考答案,答案区域分布可用于评分中间推理步骤,提供与正确性相关性更强的中间奖励。相比自回归模型,MDLM在中间步骤的答案似然得分对正确性的预测力显著更强。

原文摘要 · Abstract (English)

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.

推理机制扩散模型强化学习早期退出

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。