BackPlay: Head-Only Look-Back Self-Correction for Diffusion Language Models

arXiv:2601.06428v3 Announce Type: replace Abstract: Diffusion Language Models (DLMs) decode multiple tokens in parallel, but aggressive multi-token decoding amplifies cross-token dependency errors and can sharply degrade generation quality. We propose BackPlay, a frozen-backbone self-correction framework that trains only a lightweight correction head on a finetuned DLM without updating any backbone or adapter parameters. Because the head is trained on errors produced by the same frozen generator used at inference time, its training distribution aligns with the error patterns of the deployed model. We further introduce Look-back Correction, a training mechanism that injects predictions from earlier, more corrupted denoising states into later, richer contexts, enabling the head to leverage later context to detect mistakes made in earlier generation steps. During inference, BackPlay periodically revisits previously generated tokens through selective remasking and regeneration to limit error accumulation. Across mathematical reasoning and code generation benchmarks, BackPlay improves the speed--quality trade-off of the underlying DLM under multi-token decoding.

Related

• Parallelism and Generation Order in Masked Diffusion Language Models: Limits Today, Potential Tomorrow
• Diffusion Language Models Know the Answer Before Decoding
• LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models
• NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization
• Improving the Throughput of Diffusion-based Large Language Models via a Training-Free Confidence-Aware Calibration

Source: arXiv cs.LG | 2026-04-24