Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models

What is the native form of reasoning for Mask Diffusion Models (LLaDa, DiffusionGemma, etc.)?

For autoregressive language models, reasoning has largely been framed as continuation: generate a chain of thought, reflect on it, and append more tokens.

But Mask Diffusion Models are not left-to-right. They generate and refine a full canvas with bidirectional context. So perhaps their natural reasoning mechanism should not be continuation, but revision.
In our recent work, we introduce Reflective Masking, a lightweight post-training method that elicits multi-turn self-revision in existing Mask Diffusion Models.

The core idea is simple:
Instead of committing to every generated token, the model learns a token-level revision policy:
✅ Keep reliable tokens
🟨 Re-mask uncertain or wrong tokens
✨ Reveal better replacements
This turns generation into an iterative process of sparse, local self-correction.

A key challenge in this setting is that revision can loop: a model may re-mask an incorrect token and later generate the same mistake again. To address this, we introduce History Reference, a parameter-free memory mechanism that exposes previous denoising states to the model. This helps the model remember what it has already tried and avoid repeated errors.

We find this direction exciting because it suggests a different view of reasoning for diffusion language models:
Autoregressive reflection thinks by continuing.
Reflective Masking thinks by revising.

Recent systems such as DiffusionGemma and Gemini Diffusion highlight the growing promise of diffusion-based text generation, including bidirectional context and iterative refinement. Our work studies a complementary question: how can existing MDMs be post-trained to explicitly perform sparse, multi-turn reasoning through self-revision?

We evaluate Reflective Masking across several settings:
🧩 Sudoku revision
🧠 mathematical reasoning
💻 code generation
🎨 image editing

Across these tasks, Reflective Masking improves performance by allowing the model to selectively revisit previous outputs rather than regenerate from scratch or commit too early.
The broader takeaway is that diffusion language models may not simply be faster parallel generators. They may enable a different reasoning paradigm: reasoning as iterative state refinement.