Correction Training Targets Error Accumulation in Diffusion Sampling

Diffusion models generate images through a sequence of denoising steps. During training, a denoiser usually receives an interpolation between a real image and noise. At inference time, however, the input at a later step is influenced by the model’s own earlier prediction. The model is therefore asked to process inputs that do not exactly match the distribution seen during training. Small mistakes made early in the trajectory can persist or compound, which may help explain why unguided diffusion generation remains difficult.

The paper “Correction Training for Diffusion Models” addresses this issue by adding a separate corrector rather than redesigning the entire diffusion backbone. Its framework can be summarized as follows:

• The base denoiser first predicts the clean sample at the current timestep.
• A corrector receives that prediction as its input.
• The corrector is trained toward the same ground-truth clean target used by the base model.
• At sampling time, the corrector refines intermediate predictions, and its operation can be repeated on its own outputs.

This setup turns prediction refinement into an explicit learning problem. Instead of relying only on a sampling schedule or external guidance signal, the corrector is trained using the kinds of outputs produced by the base denoiser. The intended effect is to narrow the gap between the model’s training inputs and the states encountered along its actual sampling trajectory. Correcting predictions before the next denoising transition may also prevent early artifacts from being carried through to the final image.

The authors evaluate the method with DiT, SiT, and JiT architectures. They report substantial improvements over previously reported FID results for unguided generation, including unconditional settings where high-quality synthesis is particularly challenging. Classifier-free guidance can improve visual quality, but it may also affect diversity and lead to overly strong or saturated outputs. Correction training pursues a different trade-off: improve the model’s own predictions while largely preserving the variety of generated samples.

The approach is notable because it is modular. The original denoiser remains responsible for the main generation process, while the additional corrector focuses on errors exposed during iterative sampling. This suggests a broader design direction for diffusion systems: rather than depending exclusively on stronger guidance, larger backbones, or more elaborate samplers, a model can be trained to recognize and repair its intermediate mistakes.

The available material does not establish how the method behaves across all datasets, resolutions, or sampling budgets, so those questions require examination of the full paper and further studies. Still, the work highlights a practical source of quality loss in diffusion inference and offers a simple mechanism for addressing it.

Source: Hugging Face Daily Papers