Self-Generated Error Training for Token Editing in Diffusion Language Models

2026-06-15

AI Summary

This work addresses the misalignment between training and inference objectives in diffusion-based language models—where training relies on random perturbations while inference requires correcting the model’s own high-confidence errors—by proposing a self-generated Text-to-Text (T2T) approach. The method leverages gradient-free draft generation and masked-fill prediction, using the model’s self-generated errors as supervision signals for recovery training. This is the first approach to incorporate self-generated errors into the training process, effectively aligning the training and inference distributions. By reducing required edit intensity, it improves correction accuracy at critical positions, such as the last digit of an answer. Built upon LLaDA2.1-mini with short-cycle LoRA continued pretraining, and combined with block-wise diffusion decoding and Q-Mode T2T inference, the model significantly mitigates issues like over-self-correction and end-digit transcription errors across multiple benchmarks.

Abstract

Token-to-token (T2T) editing lets LLaDA2.1 revise committed tokens during block-diffusion decoding. The released recipe trains this editor on random vocabulary corruptions, but at inference the editor sees the model's own fluent, high-confidence draft errors instead. We study this training-inference mismatch and propose self-generated T2T, which performs a no-gradient draft pass, fills masked positions with predicted tokens, and supervises recovery in a second pass under these self-generated corruptions. We implement the update as a short LoRA continued-pretraining pass on LLaDA2.1-mini and evaluate on several benchmarks under the official Q-Mode T2T procedure with unchanged inference parameters. The method generally improves accuracy while reducing T2T edit intensity, mitigating failure modes such as final-digit transcription errors after otherwise correct reasoning and excessive self-correction before short factual answers.

Problem

• training-inference mismatch
• token editing
• diffusion language models
• self-generated errors
• T2T editing

Innovation

• self-generated error training
• token-to-token editing
• diffusion language models
• LoRA continued-pretraining
• training-inference mismatch