多轮反射掩模技术能够激发掩模扩散模型中的推理能力
Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models
摘要
虽然自回归模型中的推理通常是通过逐步思考与反思来完成的，但其对之前输出的优化仍然依赖于完全的顺序生成方式，即使只需要进行局部修改而已。相比之下，Mask Diffusion Model中的掩蔽机制能够自然地支持对之前输出进行局部修改，从而可以在不丢弃之前结果的情况下进行选择性优化，并从头开始生成新的结果。这一特性更接近于人类通过迭代式局部优化来纠正错误的方式。不过，现有的Mask Diffusion Model并不支持多轮掩蔽与去噪操作。我们提出了Reflective Masking技术，通过轻量级训练过程让这些模型具备这种内在推理能力。Reflective Masking能够在测试阶段实现灵活扩展，使得模型能够根据不断变化的上下文重新处理之前的输出。为了利用之前几轮中的信息，我们还引入了History Reference机制，这是一种无需参数的机制，可以利用中间去噪状态来进行修正。我们的方法不需要任何架构上的调整，且可以轻松应用于现有的Mask Diffusion Model中。在各种任务和领域中，包括文本生成、数独求解和图像编辑等，Reflective Masking始终优于基于传统掩蔽机制的模型，显示出强大的通用性，因此将Reflective Masking视为在Mask Diffusion Model中进行推理的基本机制。
English Abstract
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.