Search papers, labs, and topics across Lattice.
This paper identifies failure modes in Token-to-Token (T2T) editing, a common error correction technique in masked diffusion language models, where tokens are directly replaced based on confidence thresholds. They propose Token-to-Mask (T2M) remasking, which resets a suspect token to a mask state, allowing the model to re-predict it from a cleaner context in the next denoising step. T2M, a training-free method, significantly improves accuracy on tasks requiring exact token-level output, achieving a +5.92 point gain on CMATH by repairing 41.3% of last-mile corruption errors.
Masking, not replacement, unlocks better error correction in diffusion language models by providing a cleaner context for subsequent denoising steps.
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.
