多輪反思遮罩在遮罩擴散模型中引發推理

摘要

在自迴歸（AR）模型的推理過程中，常透過思維鏈推理與反思來進行，但對先前輸出的修改仍依賴完全序列化的生成，即便僅需局部編輯時亦然。相較之下，遮蔽擴散模型（MDM）中的遮蔽機制自然支援對先前輸出進行明確的局部編輯，使其能夠在無須拋棄先前的答案並從頭生成的條件下，進行選擇性的修正。雖然這項特性更貼近人類透過反覆局部修正來改正錯誤的方式，但現有的MDM並不支援多輪次的遮蔽與去噪。我們提出反思性遮蔽（RM），透過輕量級後訓練來激發MDM中這項內在的推理能力。RM提供一種原生的測試階段擴展機制，使MDM能夠根據不斷變化的上下文，反覆回顧並修正其先前的輸出。為進一步利用類似AR推理中來自先前回合的見解，我們引入歷史參考（History Reference）——一種無需參數的機制，能在修正過程中利用中間的去噪狀態。我們的方法無需改變架構，且易於套用至現有的MDM。在包含文字生成、數獨與影像編輯等多樣任務與模態中，反思性遮蔽一致地優於標準的遮蔽式基準方法，並展現出強大的泛化能力，使RM成為MDM上進行推理的基礎原語。

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.