让掩码扩散模型像人一样逐步修正答案,无需重来。
Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models

- 通过轻量级后训练引入多轮掩码修正机制
- 在文本、数独、图像编辑任务中显著优于基线方法
- 无需修改模型结构,适用于现有掩码扩散模型
尽管自回归模型常通过思维链和反思来推理,但其对先前输出的优化仍依赖全序列生成,即使只需局部修改。相比之下,掩码扩散模型(MDMs)的掩码机制天然支持对已有输出的显式局部修改,可在不丢弃先前结果的情况下实现选择性修正。然而,现有MDMs不支持多轮掩码与去噪。本文提出反射掩码(Reflective Masking, RM),通过轻量级后训练激发MDMs内在的迭代修正能力。RM提供原生的测试时扩展能力,使模型基于动态上下文多次回溯并修正先前输出。为进一步利用前轮推理信息,引入无参数的历史参考机制,利用修正过程中的中间去噪状态。该方法无需架构改动,可直接应用于现有MDMs。在文本生成、数独求解、图像编辑等多任务多模态场景中,RM持续优于标准掩码基线,展现出强泛化能力,有望成为MDMs推理的基础范式。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。