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

TL;DR

BackPlay is a self-correction framework for diffusion language models that improves multi-token decoding quality by training a lightweight correction head and revisiting previous tokens during inference.

Contribution

It introduces a novel self-correction method with look-back correction and token revisiting, enhancing decoding quality without updating the backbone model.

Findings

Improves speed-quality trade-off in mathematical reasoning tasks.

Enhances code generation accuracy with minimal additional training.

Demonstrates effectiveness across multiple benchmarks.

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.