aiqbench analysis
This research introduces Token-to-Mask (T2M), a training-free inference-time correction method for diffusion language models that resolves internal inconsistencies by remasking low-confidence positions rather than replacing them. The study demonstrates that this approach effectively converts additional compute into performance gains while maintaining the sequential-step advantages over traditional autoregressive decoding across various tasks.
Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Language Models
arXiv:2604.18738v3 Announce Type: replace Abstract: Diffusion language models (dLLMs) generate text through iterative denoising, filling multiple masked positions at each step. Positions filled in the same step are predicted without conditioning on one another's newly filled values and can therefore be mutually inconsistent; once retained, these inconsistencies become context for later predictions. We introduce \emph{Token-to-Mask} (T2M), a training-free inference-time correction method that identifies low-confidence positions using the model's probability of the current token, remasks them, and reconstructs them in later denoising steps. On dLLMs equipped with correction mechanisms, a single T2M configuration transfers across tasks and models without retuning and broadly improves task metrics over each model's native correction mechanism. In controlled experiments, we decompose correction methods into a detector that identifies suspicious tokens and an action that determines how to revise them. Holding the detector fixed, remasking yields higher task metrics than replacement; across the tested detector--action combinations, current-token-probability detection paired with remasking performs best. Compared with direct editing, T2M converts additional inference compute into performance gains more effectively and, on most tasks, retains a sequential-step advantage over autoregressive token-by-token decoding.
Abstract:Diffusion language models (dLLMs) generate text through iterative denoising, filling multiple masked positions at each step. Positions filled in the same step are predicted without conditioning on one another's newly filled values and can therefore be mutually inconsistent; once retained, these inconsistencies become context for later predictions. We introduce \emph{Token-to-Mask} (T2M), a training-free inference-time correction method that identifies low-confidence positions using the model's probability of the current token, remasks them, and reconstructs them in later denoising steps. On dLLMs equipped with correction mechanisms, a single T2M configuration transfers across tasks and models without retuning and broadly improves task metrics over each model's native correction mechanism. In controlled experiments, we decompose correction methods into a detector that identifies suspicious tokens and an action that determines how to revise them. Holding the detector fixed, remasking yields higher task metrics than replacement; across the tested detector--action combinations, current-token-probability detection paired with remasking performs best. Compared with direct editing, T2M converts additional inference compute into performance gains more effectively and, on most tasks, retains a sequential-step advantage over autoregressive token-by-token decoding.
Computation and Language (cs.CL)
arXiv:2604.18738 [cs.CL]
(or arXiv:2604.18738v3 [cs.CL] for this version)
https://doi.org/10.48550/arXiv.2604.18738 Focus to learn more arXiv-issued DOI via DataCite
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 270 words
Related coverage
NAVIRA: Decoupled Stochastic Remasking for Masked Diffusion Language Models
arXiv · Jun 5, 2026
CoDR: Training-Free Confidence-Drift Remasking for Diffusion Language Models
arXiv · Oct 8, 2026
Google AI Releases DiffusionGemma, a 26B MoE Open Model Using Text Diffusion for Up to 4x Faster Generation
MarkTechPost · Jun 10, 2026
Low-Confidence Remasking Traps Flexibility: Realizing Arbitrary-Order Potential for Diverse Rollouts in Diffusion LLMs
arXiv · Sep 30, 2026