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

• URL: http://arxiv.org/abs/2605.26436v2
• Date: Sun, 31 May 2026 04:33:39 GMT
• Title: Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models
• Abstract summary: 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.
• Score: 0.6916773850242582
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Escaping Confidence Trap: Evolutionary Decoding for Mathematical Reasoning in Diffusion LLMs [69.52853771994451]
  We analyze the decoding trajectories of LLaDA 2.0 and identify a recurring diffusion confidence trap.<n>We propose Evolutionary Decoding, a training-free test-time scaling framework that views diffusion decoding as an evolutionary process.
  arXiv  Detail & Related papers  (2026-08-01T11:48:25Z)
• 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)
• Follow the Latent Roadmap: Navigating Revocable Decoding for Diffusion LLMs with Anchor Tokens [22.687236570529844]
  We propose ASRD (Anchor Supervised Revocable Decoding), a training-free framework that operates within the embedding space.<n>We show that ASRD outperforms recent remasking baselines, achieving accuracy improvements of up to 6.4% while accelerating inference throughput by up to 7.2times$.
  arXiv  Detail & Related papers  (2026-06-15T15:23:47Z)
• 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)
• SimSD: Simple Speculative Decoding in Diffusion Language Models [61.33773959352141]
  Diffusion large language models (dLLMs) offer faster inference through parallel or blockwise decoding.<n>Their masked language modeling formulation remains incompatible with standard token-level speculative decoding.<n>We propose a speculative decoding algorithm for dLLMs, named SimSD, that equips dLLMs with temporally valid token-level contexts.<n>Our method achieves up to 7.46x higher decoding throughput while maintaining and even improving average generation quality.
  arXiv  Detail & Related papers  (2026-06-01T17:46:46Z)
• Remask, Don't Replace: Token-to-Mask Refinement in Masked Diffusion Language Models [0.6916773850242582]
  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.
  arXiv  Detail & Related papers  (2026-04-20T18:43:28Z)
• Learn from Your Mistakes: Self-Correcting Masked Diffusion Models [31.536464269884103]
  Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models.<n>We propose a framework that trains a model to perform both unmasking and correction.<n>We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence.
  arXiv  Detail & Related papers  (2026-02-12T05:17:31Z)
• 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)
• Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models [51.12873073612084]
  Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance depends on the inference time order of unmasking.<n>We propose Lookahead Unmasking (LookUM), which addresses these concerns by reformulating sampling as path selection over all possible unmasking orders.<n>LookUM requires only two to three paths to achieve peak performance, demonstrating remarkably efficient path selection.
  arXiv  Detail & Related papers  (2025-11-04T02:37:37Z)
• From Denoising to Refining: A Corrective Framework for Vision-Language Diffusion Model [72.73512218682187]
  We introduce ReDiff, a refining-enhanced diffusion framework that teaches the model to identify and correct its own errors.<n>Our approach features a two-stage training process: first, we instill a foundational revision capability by training the model to revise synthetic errors; second, we implement a novel online self-correction loop.<n>This mistake-driven learning endows the model with the crucial ability to revisit and refine its already generated output, effectively breaking the error cascade.
  arXiv  Detail & Related papers  (2025-10-22T06:58:55Z)
• Finish First, Perfect Later: Test-Time Token-Level Cross-Validation for Diffusion Large Language Models [47.5976588836299]
  Diffusion large language models (dLLMs) offer advantages such as accelerated parallel decoding and bidirectional context modeling.<n>The vanilla decoding strategy in discrete dLLMs suffers from a critical limitation: once a token is accepted, it can no longer be revised in subsequent steps.<n>We propose Tolerator, a training-free decoding strategy that leverages cross-validation among predicted tokens.
  arXiv  Detail & Related papers  (2025-10-06T17:56:46Z)
• 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)
• Masking as an Efficient Alternative to Finetuning for Pretrained Language Models [49.64561153284428]
  We learn selective binary masks for pretrained weights in lieu of modifying them through finetuning. In intrinsic evaluations, we show that representations computed by masked language models encode information necessary for solving downstream tasks.
  arXiv  Detail & Related papers  (2020-04-26T15:03:47Z)

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.