反思掩码(RM)激发掩码扩散模型的推理能力
阅读原文· arxiv.org自回归模型推理依赖链式思维和反思,但局部修改仍需完全顺序生成。掩码扩散模型(MDMs)的掩码机制天然支持选择性局部编辑,但现有MDMs不支持多轮掩码与去噪。本文提出反思掩码(RM),通过轻量后训练激发MDMs的多轮掩码能力,迭代修改先前输出。RM还引入免参数的历史参考机制,利用中间去噪状态提升修订效果。无需架构改动,在文本生成、数独、图像编辑等任务上一致优于标准掩码基线。
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.