Correction Training for Diffusion Models

Summary

The paper identifies exposure bias as a source of error accumulation in iterative diffusion sampling: denoisers train on real images mixed with noise, but inference feeds them inputs produced by their own imperfect earlier predictions. It introduces a corrector that receives the base denoiser's clean-sample prediction and is trained toward the same ground-truth target. During sampling, the corrector is applied repeatedly at intermediate denoising steps to reduce artifacts before they persist in the final sample. Although training uses only base-denoiser predictions, the corrector can refine its own outputs at inference. Experiments across DiT, SiT, and JiT report substantially better previously reported FID scores for unguided and unconditional generation, while largely preserving diversity. The method performs competitively with classifier-free guidance (CFG); combining the two further improves diversity and reduces CFG's typical oversaturation. The authors also distill the refined predictions into the base diffusion model, retaining the quality gains without additional inference-time correction overhead.