Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models

2026-06-15

AI Summary

Existing masked diffusion models lack the capacity for multi-round iterative reasoning, making it difficult to emulate human-like local refinement processes. This work proposes a lightweight post-training approach that introduces a reflective masking mechanism, enabling multi-step iterative denoising at test time. It further incorporates a history-aware reference mechanism that leverages intermediate denoising states without requiring additional parameters or architectural modifications, thereby enhancing inference performance. As the first method to realize multi-round reflective reasoning within masked diffusion models, it significantly outperforms standard baselines across diverse cross-modal tasks—including text generation, Sudoku solving, and image editing—demonstrating strong generalization and potential as a universal reasoning primitive.

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.

Problem

• Mask Diffusion Models
• multi-turn reasoning
• local refinement
• iterative revision
• reasoning capability

Innovation

• Reflective Masking
• Mask Diffusion Models
• Iterative Refinement
• History Reference
• Local Editing