Remask, Don't Replace: Token-to-Mask Refinement in Masked Diffusion Language Models

• URL: http://arxiv.org/abs/2604.18738v1
• Date: Mon, 20 Apr 2026 18:43:28 GMT
• Title: Remask, Don't Replace: Token-to-Mask Refinement in Masked Diffusion Language Models
• Abstract summary: Masked diffusion language models such as LLaDA2.1 rely on Token-to-Token editing to correct their own generation errors.<n>We propose Token-to-Mask (T2M) remasking, which resets the position to the mask state so that the next denoising step re-predicts it from an in-distribution context.<n>Across 8 benchmarks, T2M improves accuracy on tasks that require exact token-level output.
• Score: 0.6916773850242582
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Masked diffusion language models such as LLaDA2.1 rely on Token-to-Token (T2T) editing to correct their own generation errors: whenever a different token crosses a confidence threshold, the committed token is overwritten. We identify three structural failure modes of this rule. The trigger cannot fire when no single alternative is confident enough; the replacement is computed under a context that may itself contain errors; and the uniform perturbations used to train the T2T stream do not resemble the coherent, semantically plausible mistakes that the model actually makes at inference. As an alternative, we propose Token-to-Mask (T2M) remasking. Rather than overwriting a suspect token with a new guess, T2M resets the position to the mask state, so that the next denoising step re-predicts it from an in-distribution context. The method is training-free, modifies only the editing rule, and introduces no new parameters. We pair it with three detection heuristics and give a short theoretical account of why a mask is a better conditioning signal than an erroneous token. Across 8 benchmarks, T2M improves accuracy on tasks that require exact token-level output. Its largest gain is +5.92 points on CMATH, where we attribute 79.9% of baseline errors to last-mile corruption (correct reasoning followed by a garbled final answer); T2M repairs 41.3% of these cases.

Related papers

• Self-Generated Error Training for Token Editing in Diffusion Language Models [0.6916773850242582]
  Token-to-token (T2T) editing lets LLaDA2.1 revise committed tokens during block-diffusion decoding.<n>We study this training-inference mismatch and propose self-generated T2T.
  arXiv  Detail & Related papers  (2026-06-15T18:13:54Z)
• NAVIRA: Decoupled Stochastic Remasking for Masked Diffusion Language Models [0.0]
  PRISM addresses this by learning token-level quality scores and remasking unreliable tokens.<n>We propose NAVIRA, an inference-time decoding policy that separates these two operations.<n>Temperature-controlled remasking reduces repeated correction of the same positions.
  arXiv  Detail & Related papers  (2026-06-04T11:24:47Z)
• DMax: Aggressive Parallel Decoding for dLLMs [77.24184219948337]
  We present DMax, a new paradigm for efficient diffusion language models (dLLMs)<n>It mitigates error accumulation in parallel decoding, enabling aggressive decoding while preserving generation quality.<n>At the core of our approach is On-Policy Uniform Training, a novel training strategy that efficiently unifies masked and uniform dLLMs.
  arXiv  Detail & Related papers  (2026-04-09T14:35:42Z)
• Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models [0.6916773850242582]
  LLaDA2.1 introduced a Token-to-Token (T2T) editing mechanism that accelerates generation.<n>We propose Token-to-Mask (T2M) remasking, a training-free, drop-in replacement for T2T editing.<n>T2M resets suspected erroneous tokens back to the mask state, allowing the diffusion process to re-predict them under cleaner context.
  arXiv  Detail & Related papers  (2026-04-06T17:39:11Z)
• LogicDiff: Logic-Guided Denoising Improves Reasoning in Masked Diffusion Language Models [0.0]
  We introduce LogicDiff, an inference-time method that replaces confidence-based unmasking with logic-role-guided unmasking.<n>A dependency-ordered scheduler then unmasks tokens in logical dependency order.<n> LogicDiff improves LLaDA-8B-Instruct accuracy from 22.0% to 60.7% on GSM8K and from 23.6% to 23.6% on MATH-500, with less than 6% speed overhead.
  arXiv  Detail & Related papers  (2026-03-24T13:08:10Z)
• Self-Aware Markov Models for Discrete Reasoning [8.161697757509701]
  We introduce a method based on learning a Markov transition kernel that is trained on its own outputs.<n>This design enables tokens to be remasked, allowing the model to correct its previous mistakes.<n>On the Sudoku-Extreme dataset we clearly outperform other flow based methods with a validity of 95%.
  arXiv  Detail & Related papers  (2026-03-17T15:30:06Z)
• CORE: Context-Robust Remasking for Diffusion Language Models [51.59514489363897]
  We propose Context-Robust Remasking (CORE), a training-free framework for inference-time revision.<n>Rather than trusting static token probabilities, CORE identifies context-brittle tokens by probing their sensitivity to targeted masked-context perturbations.<n>On LLaDA-8B-Base, CORE delivers consistent improvements across reasoning and code benchmarks, outperforming compute-matched baselines and improving MBPP by up to 9.2 percentage points.
  arXiv  Detail & Related papers  (2026-02-04T00:12:30Z)
• Residual Context Diffusion Language Models [90.07635240595926]
  Residual Context Diffusion (RCD) is a module that converts discarded token representations into contextual residuals and injects them back for the next denoising step.<n>RCD consistently improves frontier dLLMs by 5-10 points in accuracy with minimal extra computation overhead.
  arXiv  Detail & Related papers  (2026-01-30T13:16:32Z)
• Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models [22.525318796588568]
  Think-at-Hard (TaH) is a dynamic latent thinking method that iterates deeper only at hard tokens.<n>TaH boosts reasoning performance across five challenging benchmarks.
  arXiv  Detail & Related papers  (2025-11-11T18:57:02Z)
• REAR: Rethinking Visual Autoregressive Models via Generator-Tokenizer Consistency Regularization [130.46612643194973]
  reAR is a simple training strategy introducing a token-wise regularization objective.<n>On ImageNet, it reduces gFID from 3.02 to 1.86 and improves IS to 316.9 using a standardization-based tokenizer.<n>When applied to advanced tokenizers, it achieves a gFID of 1.42 with only 177M parameters, matching the performance with larger state-of-the-art diffusion models (675M)
  arXiv  Detail & Related papers  (2025-10-06T02:48:13Z)
• Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models [40.902681492117786]
  RemeDi is a mask-based DLM that predicts token distributions and per-token confidence scores at each step.<n>We train a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model to detect and remask incorrect tokens.<n>Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.
  arXiv  Detail & Related papers  (2025-09-28T05:39:49Z)
• Mask the Correct Tokens: An Embarrassingly Simple Approach for Error Correction [38.463639262607174]
  Previous error correction methods usually take the source (incorrect) sentence as encoder input and generate the target (correct) sentence through the decoder. We propose a simple yet effective masking strategy to achieve this goal.
  arXiv  Detail & Related papers  (2022-11-23T19:05:48Z)
• Mask Transfiner for High-Quality Instance Segmentation [95.74244714914052]
  We present Mask Transfiner for high-quality and efficient instance segmentation. Our approach only processes detected error-prone tree nodes and self-corrects their errors in parallel. Our code and trained models will be available at http://vis.xyz/pub/transfiner.
  arXiv  Detail & Related papers  (2021-11-26T18:58:22Z)
• Composable Sparse Fine-Tuning for Cross-Lingual Transfer [56.86192078426372]
  Fine-tuning all parameters of a pre-trained model has become the mainstream approach for transfer learning. We introduce a new fine-tuning method with both these desirable properties. It outperforms adapters in zero-shot cross-lingual transfer by a large margin.
  arXiv  Detail & Related papers  (2021-10-14T17:27:29Z)
• ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators [108.3381301768299]
  Masked language modeling (MLM) pre-training methods such as BERT corrupt the input by replacing some tokens with [MASK] and then train a model to reconstruct the original tokens. We propose a more sample-efficient pre-training task called replaced token detection.
  arXiv  Detail & Related papers  (2020-03-23T21:17:42Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.