Correction Training for Diffusion Models

Correction Training for Diffusion Models. By Donghoon Ahn and four others, from the University of California, Berkeley and Reve. Published on the fourteenth of October. Its code repository is worked out: about fifty-five source files, with a readme and a licence.

A diffusion model samples iteratively, and iteration lets errors pile up. The cause named here is exposure bias. During training, the denoiser only ever sees real images interpolated with noise. At inference, its input is built from its own imperfect predictions at earlier steps, which is something it never encountered in training. The model is asked to work on material unlike anything it practised on, and small mistakes early on get carried forward. The authors suggest this may be one reason diffusion models often struggle without extra help, such as classifier-free guidance, or CFG.

The response is to add a correction step aimed squarely at that training-inference gap. A corrector takes the base model's prediction as its input and is trained against the same ground-truth target the base model was. That turns correction into an explicit learning task, and the training inputs for it come from the base model itself. In other words, the corrector is trained on exactly the kind of imperfect prediction the denoiser produces, which is the kind of input the denoiser was never shown. The base model is left as it is. Nothing about its training changes; a second learner is added whose whole job is to start from where the first one actually ends up, and move that result closer to the truth.

At sampling time, at every denoising step, the denoiser's prediction of the clean sample is refined iteratively by the corrector. Because the fix happens at intermediate steps rather than only at the end, an artifact that appears early can be caught before it persists into the final sample. One property stands out. The corrector is trained only on predictions from the base denoiser, yet at inference it can be applied repeatedly to its own outputs, and doing so gives better results.

There is also a way to pay for it. Correction adds inference cost, but the authors show that cost can be amortized by distilling the refined predictions back into the base diffusion model. The quality gains are retained, and the extra inference-time overhead goes away.

The method is evaluated across three model families: DiT, SiT and JiT. It substantially improves on previously reported FID scores for unguided generation, and notably for unconditional generation, where producing high-quality samples without guidance has been particularly hard. Sample quality improves substantially while diversity is largely preserved. Against classifier-free guidance, correction alone is competitive. Used together with CFG, it goes further, improving diversity and reducing the oversaturation that CFG typically causes.

The unconditional result is the one the authors single out. Without any guidance and without any condition to steer by, a diffusion model has nothing to lean on but its own predictions, which is exactly the situation where errors from earlier steps should matter most. That is also where correction shows its most notable improvement.

Picture a student who only ever practised on clean exercises, and is then asked to keep building on their own rough working. The corrector amounts to practice on one's own rough working, made into a training task. The framework asks for little: one extra learned step, trained on the base model's outputs, reaching quality competitive with guidance, after which distillation can fold even that step back into the original model.