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

arXiv: https://arxiv.org/abs/2605.26436 · HTML: https://mlsanity.com/paper/2605.26436/

## Key facts

- No citing papers in the map yet (data through 14 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.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs](https://mlsanity.com/paper/1903.00161/) (2019)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)
- [GSM-Plus: A Comprehensive Benchmark for Evaluating the Robustness of LLMs as Mathematical Problem Solvers](https://mlsanity.com/paper/2402.19255/) (2024)
- [CMATH: Can Your Language Model Pass Chinese Elementary School Math Test?](https://mlsanity.com/paper/2306.16636/) (2023)
- [LLaDA2.1: Speeding Up Text Diffusion via Token Editing](https://mlsanity.com/paper/2602.08676/) (2026)
- [ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs](https://mlsanity.com/paper/2510.04767/) (2025)
- [Fine-Tuning Masked Diffusion for Provable Self-Correction](https://mlsanity.com/paper/2510.01384/) (2025)
- [Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models](https://mlsanity.com/paper/2509.23653/) (2025)
- [MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models](https://mlsanity.com/paper/2508.13148/) (2025)
- [Learn from Your Mistakes: Self-Correcting Masked Diffusion Models](https://mlsanity.com/paper/2602.11590/) (2026)
- [CORE: Context-Robust Remasking for Diffusion Language Models](https://mlsanity.com/paper/2602.04096/) (2026)

## 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.