mlsanityData to 7 Oct 2026

Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models

arXiv 2605.26436 ↗First posted 6 Apr 2026cs.CL cs.AIcited by 0

On the map: Diffusion models → Discrete and text diffusion

Abstract (arXiv)

Discrete masked diffusion language models such as LLaDA generate text through iterative denoising, where mask tokens are progressively replaced with predicted tokens. LLaDA2.1 introduced a Token-to-Token (T2T) editing mechanism that accelerates generation by directly replacing committed tokens suspected of being incorrect. However, we identify fundamental limitations of T2T editing: it couples error detection with replacement, pollutes the generation context with potentially incorrect tokens, and introduces a train-inference noise mismatch where systematic model-generated errors differ from the random perturbations seen during training. We propose Token-to-Mask (T2M) remasking, a training-free, drop-in replacement for T2T editing that resets suspected erroneous tokens back to the mask state, allowing the diffusion process to re-predict them under cleaner context. We design and empirically validate three complementary error detection strategies -- probability-based, trigger-mirrored, and temporal-difference-based -- and provide a unified theoretical analysis showing that T2M remasking purifies the generation context, converts systematic inference errors back to the model's native mask noise type, and enables delayed commitment for joint multi-position optimization. Comprehensive experiments across 12 benchmarks spanning knowledge, reasoning, mathematics, coding, and instruction following show that T2M generally improves performance on tasks requiring precise token-level output, with the largest gain on mathematics (+5.92% on CMATH). Error analysis on CMATH reveals that the dominant failure mode is last-mile token corruption -- where correct reasoning produces a corrupted final answer -- and that T2M repairs 59.4% of such cases.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Discrete and text diffusion, one of 14 directions in Diffusion models. Builds on 27 of its 28 references that are in the map. First posted on arXiv 6 Apr 2026.

Builds on

27 of 28 references.

Closest ancestors · 11

Same direction, newest first

  1. 2026Learn from Your Mistakes: Self-Correcting Masked Diffusion Models12 Feb
  2. 2026LLaDA2.1: Speeding Up Text Diffusion via Token Editing9 Feb
  3. 2026CORE: Context-Robust Remasking for Diffusion Language Models4 Feb
  4. 2025Fine-Tuning Masked Diffusion for Provable Self-Correction1 Oct
  5. 2025Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models28 Sep
  6. 2025Dream 7B: Diffusion Large Language Models21 Aug
  7. 2025MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models18 Aug
  8. 2025Remasking Discrete Diffusion Models with Inference-Time Scaling1 Mar
  9. 2024Simple and Effective Masked Diffusion Language Models11 Jun
  10. 2024Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data6 Jun
  11. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution25 Oct

Other references · 16

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Denoising Diffusion Probabilistic Modelscited by 20,851
  3. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 8,964
  4. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  5. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  6. 2017TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehensioncited by 2,895
  7. 2023GPQA: A Graduate-Level Google-Proof Q&A Benchmarkcited by 2,748
  8. 2022Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Themcited by 1,767
  9. 2024MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmarkcited by 1,758
  10. 2021Structured Denoising Diffusion Models in Discrete State-Spacescited by 1,665
  11. 2023Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generationcited by 1,332
  12. 2023Instruction-Following Evaluation for Large Language Modelscited by 1,143
4 more
  1. 2019DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphscited by 883
  2. 2024GSM-Plus: A Comprehensive Benchmark for Evaluating the Robustness of LLMs as Mathematical Problem Solverscited by 155
  3. 2023CMATH: Can Your Language Model Pass Chinese Elementary School Math Test?cited by 89
  4. 2025ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMscited by 42

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Remask, Don't Replace: Token-to-Mask Refinement in Diffusion Language Models
  2. 2026Self-Generated Error Training for Token Editing in Diffusion Language Models
  3. 2026NAVIRA: Decoupled Stochastic Remasking for Masked Diffusion Language Models
  4. 2026Beyond Masks: Efficient, Flexible Diffusion Language Models via Deletion-Insertion Processes
  5. 2026Know When to Hold 'em: Correct-Token Retention in Uniform-State Diffusion Language Models
  6. 2026Edit-Based Refinement for Parallel Masked Diffusion Language Models
  7. 2026Beyond Token Revision: Investigating Mask-and-Replace Diffusion for Zero-Shot Text-to-Speech
  8. 2026Re-evaluating Confidence Remasking in Masked Diffusion Language Models