Source-linked AI summary
Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models
Yanming Zhang, Yihan Bian, Jingyuan Qi, Yuguang Yao, Lifu Huang, Tianyi Zhou
TL;DR
Existing Mask Diffusion Models cannot actively revisit committed predictions, limiting multi-turn self-correction. Reflective Masking enables iterative local revision, with History Reference preserving intermediate states; across image editing, Sudoku, and text generation, it consistently improves performance over standard masking-based baselines.
Problem
Existing Mask Diffusion Models passively fix confidently denoised tokens and cannot actively revisit and correct committed predictions during generation.
Method
Reflective Masking uses lightweight post-training to enable uncertainty-driven selective revision, while History Reference preserves intermediate denoising states without architectural changes.
Results
Across image editing, Sudoku, and text generation, Reflective Masking consistently improves performance, with History Reference particularly effective for autonomous exploration.
Takeaways & Limitations
Reflective Masking establishes explicit in-place revision as a native form of test-time scaling and a promising reasoning primitive for Mask Diffusion Models.
Takeaways & Limitations
Evaluation covers simpler tasks than challenging long-horizon reasoning, and computational constraints leave transfer to significantly larger-scale training regimes unexamined.
Abstract
from arXiv · showhide
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.
1 Introduction
The introduction identifies a limitation of autoregressive and existing mask-diffusion reasoning: models cannot reliably revise prior errors during multi-turn generation. It presents Reflective Masking, lightweight post-training, and History Reference as mechanisms for iterative, context-aware revision without architectural changes, with improvements demonstrated across image editing, Sudoku, and text generation.
- Motivation: Autoregressive models struggle in multi-turn and long-horizon reasoning because they propagate prior errors and must regenerate entire sequences for local corrections.This makes local revision computationally inefficient and can preserve incorrect intermediate assumptions.
- Motivation: Mask Diffusion Models naturally support localized revision by keeping context fixed while resampling uncertain tokens through iterative masked updates.Their native remasking mechanism avoids regenerating the entire sequence, suggesting a potential advantage for reasoning.
- Problem: Existing MDM decoding fixes confidently denoised tokens, preventing active reconsideration and correction of previously committed predictions.The paper argues that reasoning-level self-correction requires models to identify unreliable predictions and revise them during generation.
- Reflective Masking: Reflective Masking activates self-initiated, context-aware revision through lightweight post-training without architectural modifications.The training paradigm uses a data-generation strategy designed to produce stable signals aligned with the model’s native output distribution.
- History Reference: History Reference preserves intermediate decoding states, giving MDMs a stateful view of denoising trajectories and access to insights from historical generations.The mechanism is parameter-free and addresses a capability that multi-turn reflective masking otherwise lacks relative to autoregressive reasoning.
- Evaluation: Reflective Masking demonstrates consistent improvements and strong generality across image editing, Sudoku, and text generation tasks.These tasks span rich instructional guidance, structured error correction, and minimal supervision without direct final-answer hints.
2 Related Work
Prior work establishes that mask diffusion models naturally permit token revision, but existing methods largely use one-shot denoising or externally imposed re-masking rather than iterative, intrinsic revision. Autoregressive editing methods are also indirect because forward-only generation generally requires regenerating prior outputs.
- Mask diffusion models: Mask diffusion models naturally allow tokens to be revisited through masking, yet existing approaches primarily focus on one-shot denoising instead of iterative revision.This limitation motivates treating masking as a mechanism for revising prior outputs across turns.
- Editing and revision in autoregressive models: Autoregressive editing methods include insertion-based generation, edit-based modeling, and controllable text editing frameworks.These approaches add editing capabilities through special tokens or span-level regeneration.
- Editing and revision in autoregressive models: Because autoregressive generation is forward-only, revising prior outputs typically requires regenerating entire sequences through additional decoding passes.This makes editing indirect and often inefficient because earlier decisions cannot be modified in place.
- RemeDi and re-masking approaches: RemeDi uses self-reflective re-masking with a dual-stream architecture, but requires architectural modifications and auxiliary training objectives.Its design increases both training and inference complexity while enabling revision of intermediate predictions.
- RemeDi and re-masking approaches: Other re-masking methods use mixed noise schedules, predictor–corrector strategies, or per-step resampling, treating masking as externally driven rather than intrinsic model behavior.The paper instead frames masking as a native capability that can be activated for internal revision.
3 Reflective Masking
Reflective Masking lets Mask Diffusion Models revise outputs through per-position keep, re-mask, and reveal decisions, while History Reference supplies parameter-free conditioning on prior denoising states. A matching oracle-label training paradigm teaches these actions without architectural changes.
- Inference decision rule: RM allows each position to keep its token, re-mask it for re-prediction, or reveal a token when masked.The inference rule uses the model’s per-position probabilities: re-mask when MASK is more probable than the current token, otherwise keep; masked positions reveal the most likely vocabulary token.
- Inference decision rule: Unlike standard absorbing MDM inference, RM revisits and revises past decisions, enabling test-time scaling.The rule operates iteratively on the current state, allowing local refinement rather than requiring a fresh sequence generation.
- History Reference: History Reference adds a parameter-free accumulated embedding of prior states with O(1) per-step updates.Past states are expressed in the current reference frame using lag-dependent rotations and decay, helping distinguish identical current states and avoid recurring errors.
- Training paradigm: The model is trained with per-position oracle labels that map masked tokens to reveal, wrong tokens to re-mask, and correct tokens to keep.Synthetic trajectories combine mask and wrong-token corruption with position-wise transitions, and training minimizes cross-entropy over reveal, mask, and keep actions.
- Training objective: Conditioning on the accumulated history cannot increase the optimal training risk relative to conditioning only on the current state.With a sufficiently expressive model, the richer input can improve or preserve the best achievable expected risk while recovering the intended inference rule.
4 Experiments
The experiments test Reflective Masking and History Reference across image editing, Sudoku revision, and autonomous text reasoning, spanning increasing demands for exploration and guidance. Across these settings, the method enables selective iterative correction and consistently improves over masking-based baselines.
- Experimental Tasks: Experiments span image editing with strong external guidance, Sudoku with constrained error correction, and text generation requiring fully autonomous reasoning.The tasks systematically vary supervision and exploration demands.
- Image Editing: In image editing, RM localizes and masks target regions, regenerating only those areas while preserving unmasked content.This produces higher-quality edits and avoids the consistency problems caused by globally distributed editing noise.
- Sudoku Revision: In Sudoku revision, History Reference greatly reduces repeated mistakes and constraint conflicts compared with the variant without HR.A decay factor alone improves over the no-history baseline but performs worse than the HR-only variant.
- Text Reasoning: On text reasoning tasks, RM iteratively re-masks erroneous tokens in chain-of-thought regions and corrects initially incorrect answers using evolving context.The evaluated tasks include mathematical problem solving, code generation, and ARC-Challenge.
- Text Reasoning: RM consistently improves over both LLaDA and Vanilla SFT across math, code, and ARC-Challenge benchmarks, with gains on nearly all Minerva MATH subject categories.The performance gain on MBPP is larger than that on MATH500.
5 Limitations and future work
The evaluation covers image editing, Sudoku revision, and text generation, but these tasks are simpler than challenging long-horizon reasoning problems. Computational constraints also leave transfer under significantly larger-scale training regimes unexplored.
- Limitations: RM is evaluated on image editing, Sudoku revision, and text generation, but these tasks remain simpler than challenging long-horizon reasoning problems studied in AR models.The authors attribute this gap to the limited reasoning capability of current base mask diffusion models.
- Future work: Computational resources constrain the experiments and prevent investigation of RM transfer under significantly larger-scale training regimes.Whether RM transfers effectively at larger training scales remains an open question.
6 Conclusion · A Theoretical Analysis
Reflective Masking is a lightweight post-training framework that gives Mask Diffusion Models an intrinsic ability to revise predictions iteratively. It also introduces History Reference, a parameter-free mechanism that uses intermediate denoising states to support multi-turn refinement.
- 6 Conclusion: Reflective Masking elicits reflective revision as an intrinsic capability of Mask Diffusion Models.The framework is introduced as a lightweight post-training approach.
- 6 Conclusion: RM replaces one-way denoising with iterative revisiting, re-masking, and refinement of previous predictions.Refinement is conditioned on the evolving context.
- 6 Conclusion: History Reference exposes intermediate denoising states to the model during reflective revision.It is described as a parameter-free mechanism.
- 6 Conclusion: The proposed framework supports multi-turn masking and denoising in Mask Diffusion Models.This capability follows from repeatedly revisiting and refining previous predictions.
- 6 Conclusion: RM treats generation as an evolving process rather than a one-way denoising sequence.The model can re-mask earlier predictions as context changes.
- 6 Conclusion: History Reference helps stabilize the multi-turn refinement process by exposing intermediate denoising states.The passage identifies stabilization as a benefit of the mechanism.
A.1 Bayes-Consistent Revision-Policy Learning
The revision-policy objective is Bayes-consistent: under a fixed, θ-independent corruption proposal, population cross-entropy recovers the conditional distribution of optimal revision actions. The resulting categorical policy represents calibrated correctness beliefs at visible tokens and gold-token distributions at MASK positions, while History Reference cannot increase Bayes risk.
- Population consistency: Under a fixed, θ-independent corruption proposal, cross-entropy training recovers the conditional distribution of the optimal revision action at every edit position.Theorem 1 establishes this population result under a rich-family idealization.
- Population consistency: History Reference conditioning cannot increase Bayes risk, and the analyst-side population minimum lower-bounds the Bayes risk available to the implementation.The proof conditions on X = (z, ϕ(H, c), c, E), whereas implementation uses (a(t), c).
- State-specific policy: At visible-token positions, the minimizer is a binary calibrated belief assigning α_i(X) to the visible token and 1 − α_i(X) to MASK.The Bayes-optimal inference rule re-masks when pθ(MASK | X) exceeds pθ(z_i | X).
- State-specific policy: At MASK positions, the minimizer is a conditional language model over the gold token, recovering the standard mask-prediction sampling step.This is the other state-specific behavior of the same per-position categorical policy.
- Policy interpretation: RM therefore trains one categorical policy with calibrated visible-token correctness and MASK-position gold-token beliefs, read through re-masking and candidate-reveal actions.The policy uses different atoms of the learned distribution according to the current state type.
A.2 Plug-in Excess-Risk Theorem · A.3 History Reference Information Bound · A.4 Top-k Variant: Distribution-Shift TV Bound
The appendices establish excess-risk guarantees for plug-in revision, characterize History Reference as an information hierarchy, and bound distribution shift from practical top-k corruption. Together, they show when Bayes risk is irreducible, history cannot hurt optimal estimation, and practical corruption is the relevant inference-time support.
- A.2 Plug-in Excess-Risk Theorem: The Bayes risk contains irreducible contributions from both MASK and visible-token positions, vanishing only when the conditional label distribution is degenerate.These contributions are 1 − max_y P*_i(y | X) at MASK positions and min(α_i(X), 1 − α_i(X)) at visible-token positions.
- A.2 Plug-in Excess-Risk Theorem: At the population minimum pθ = P*, the CE-loss minimizer is Bayes-optimal under 0–1 revision risk, while further CE optimization cannot reduce the non-zero Bayes risk.The result follows by combining Theorems 1 and 7.
- A.3 History Reference Information Bound: History Reference enlarges conditioning from (z, c) to (z, ϕ(H, c), c), permitting at least as low an optimal 0–1 estimation risk as the no-history baseline.The full-history conditioning uses (z, H, c), and the infimum over its larger measurable estimator family cannot increase.
- A.3 History Reference Information Bound: The accumulated history feature has Bayes risk between the no-history baseline and the full trajectory, but whether it strictly improves risk is empirical.Its information content is compressed and generally many-to-one over distinct trajectories.
- A.4 Top-k Variant: Distribution-Shift TV Bound: The practical corruption distribution is non-uniform and θ-independent, using pretrained top-k text proposals, fixed wrong digits for Sudoku, and source-image VQ tokens for image editing.Image positions already matching the target follow the keep branch rather than corruption.
- A.4 Top-k Variant: Distribution-Shift TV Bound: Under bounded per-token log-probabilities with ε = 10^-8 and B ≈18, the loss gap from uniform to practical corruption is controlled through joint and per-position total variation.The joint bound scales as 2B · |E| · TV(P, Q).
- A.4 Top-k Variant: Distribution-Shift TV Bound: Only wrong-rule positions contribute to the distribution-shift bound; mask-rule and non-edit branches contribute zero.Theorem 1 therefore applies on the practical proposal support, where the population minimizer recovers labels under ν(prac), the more relevant inference-time region.
B Engineering Implementation
The implementation uses a relative-lag recurrence for strict O(1) per-step updates, maintaining one O(N × d) buffer instead of the full token-level trajectory. It also specifies numerical-stability properties and uses the relative-lag formulation in all reported experiments.
- Recurrence: The relative-lag formulation updates the accumulated embedding in strict O(1) time per step by rotating the previous state and adding the current-step contribution.The recurrence follows from applying γR−1 to the previous closed form and using the rotation composition rule, with the current contribution using R0 = I.
- Memory layout: The implementation maintains a single a ∈ R^N×d buffer, reducing memory from O(T × N × d) for the full trajectory to O(N × d).Each update consists of one N × d block-diagonal rotation followed by an element-wise addition.
- Numerical stability: Orthogonal rotations preserve norms, and with γ ∈ (0, 1], the running sum’s norm grows at most linearly in t.This provides the stated numerical-stability guarantee for the running recurrence.
- Relative-lag formulation: Unlike the absolute-step alternative, the relative-lag formulation anchors phase at the model’s most recent input and is used in all reported experiments.The absolute-step formulation anchors the trajectory at the start of denoising while admitting the same per-step recurrence.
- Loss reweighting: Training supports per-position CE-loss weights, including constant math-launcher weights and edit-position emphasis in the Sudoku launcher, without changing the population minimizer.The Sudoku weighting is based on an X-measurable edit-vs-non-edit distinction.
C Training Algorithm
The training algorithm samples per-position corruption rules and boundaries, constructs a deterministic multi-step trajectory, and trains on oracle actions at a uniformly sampled step. Image editing additionally applies independent CFG dropout and auxiliary regularization.
- Sampling: The procedure uses trajectory length T = 6 and samples a target sequence with its task condition from the data distribution.The target sequence and task condition are denoted (x∗, c) ∼ D.
- Sampling: Each editable position independently samples a wrong or mask rule, while non-edit positions deterministically keep the target token throughout the trajectory.Wrong-rule positions also sample a source token; mask-rule positions enter directly at MASK.
- Trajectory construction: Rule-specific boundaries define trajectories that transition wrong tokens or MASK states into target tokens without post-processing.Wrong positions use β_i and µ_i boundaries, whereas mask positions sample only µ_i.
- Training objective: The algorithm samples t uniformly from {0, . . . , T −1}, constructs oracle per-position actions from the current state, and minimizes cross-entropy against those labels.The accumulated embedding is computed from the trajectory history and conditioned on c, while z is retained for label construction.
- Image-editing specialization: For image editing, independent text and image CFG dropout each occur with probability 0.10; the generic text-generation and Sudoku algorithm omits this step.Image-content embeddings are overwritten with wte(MASK) after RMSNorm when image dropout fires.
- Image-editing specialization: Image editing supplements the primary cross-entropy with unlikelihood and stage-2 ordering losses, which are heuristic regularizers outside Theorem 1.Their detailed forms are deferred to a separate technical companion.
D Results gallery on image editing
Figure 5 presents image-editing examples across four editing categories and compares the proposed edits with Lumina and Lumina-SFT baselines. The proposed edits concentrate pixel changes within predicted-mask regions, whereas baselines show changes outside those regions.
- Image editing examples: Figure 5 spans object replacement, attribute modification, object insertion, and localized scene editing.The examples cover four image-editing types.
- Image editing examples: Each row displays the source image, predicted mask, proposed edit, and Lumina / Lumina-SFT baselines.Pixel-wise difference heat maps indicate larger pixel changes with brighter regions relative to the source.
- Mask-localized editing: The proposed edits concentrate pixel changes within predicted-mask regions, while the baselines exhibit changes outside those regions.This comparison is shown using pixel-wise difference heat maps relative to the source image.
E More results on text reasoning task
Additional text-revision examples show the model making targeted corrections by re-masking logically inconsistent tokens and revising dependent positions to preserve consistency. The examples also include format-level corrections through selective re-masking.
- Text revisions: The model re-masks tokens that are logically inconsistent with their surrounding context and re-predicts them.This behavior is illustrated in Cases 1 and 4.
- Text revisions: A single corrected token can trigger downstream revisions, with dependent positions re-masked and re-predicted to match the corrected context.This cascading behavior is illustrated in Cases 2 and 5.
- Text revisions: The model also performs format-level correction by re-masking tokens whose surface form violates the intended format.This behavior is illustrated in Case 3.