While reasoning on autoregressive (AR) models is often performed by chain-of-thought reasoning and reflection, their refinement of previous outputs still relies on fully sequential generation, even when only local edits are needed. In contrast, the masking mechanism in Mask Diffusion Models (MDMs) naturally supports explicit local edits on previous outputs, allowing selective refinement without discarding previous answers and generating another from scratch. While this property more closely aligns with how humans correct mistakes by iterative local refinement, existing MDMs do not support multi-turn masking and denoising. We propose Reflective Masking (RM), which elicits such an intrinsic reasoning capability in MDMs via lightweight post-training. RM provides a native test-time scaling, where an MDM iteratively revisits and revises its prior outputs based on evolving context. To exploit insights from previous turns like AR reasoning, we further introduce History Reference, a parameter-free mechanism that leverages intermediate denoising states during revision. Our approach requires no architectural changes and is easily applicable to existing MDMs. Across diverse tasks and modalities, including text generation, Sudoku, and image editing, Reflective Masking consistently outperforms standard masking-based baselines and demonstrates strong generality, positioning RM as a fundamental primitive for reasoning on MDMs.
翻译:尽管自回归(AR)模型中的推理通常通过思维链推理与反思来执行,但其对先前输出的改进仍依赖于完全顺序生成,即使仅需局部编辑时也是如此。相比之下,掩码扩散模型(MDMs)中的掩码机制天然支持对先前输出进行显式局部编辑,允许在不丢弃之前答案的情况下进行选择性优化,而无需从头生成。尽管这一特性更贴近人类通过迭代局部修正来纠正错误的方式,现有MDMs并不支持多轮掩码与去噪。我们提出反思式掩码(Reflective Masking,RM),通过轻量级后训练激发MDMs中这种内在的推理能力。RM提供了一种原生的测试时扩展能力,使MDM能够基于动态演化的上下文迭代性地重新审视并修正其先前输出。为利用类似AR推理中来自先前轮次的洞察,我们进一步引入历史参考(History Reference),一种无需参数即可在修订过程中利用中间去噪状态的机制。该方法无需修改架构,并易于应用于现有MDMs。在包含文本生成、数独和图像编辑在内的多种任务与模态中,反思式掩码持续优于标准掩码基线方法,展现出显著的通用性,将RM定位为MDMs上推理的基本原语。