Correction Training for Diffusion Models

ListenThe Corrector: Beating CFG Without the AI Look10:13

Show notes

A look at **Correction Training for Diffusion Models** by Donghoon Ahn, which trains a separate "corrector" network to iteratively refine a denoiser's predictions at inference time. The method matches or beats classifier-free guidance on FID across four ImageNet backbones while preserving diversity and avoiding CFG's oversaturated "AI look" — and the authors tie the gain to narrowing the training–inference exposure-bias gap. We dig into the parameter-matched ablation, the two-sample classifier result, and what it means for image and video model work. **Sources** - Correction Training for Diffusion Models — Donghoon Ahn (blog)

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

Correction Training trains a separate corrector network that takes the base diffusion denoiser’s prediction as input and regresses toward the same clean target. At inference, the corrector is applied K times per denoising step, iteratively refining predictions before the sampler advances. On ImageNet-256 (DiT, SiT, SiT-REPA, JiT) and on NanoGen text-to-image, correction improves FID over the unguided baseline across all four backbones at an inference budget matched to CFG, and matches or beats CFG on FID while preserving diversity and avoiding CFG’s oversaturated “AI look.” The authors frame the mechanism as a reduction of the exposure-bias gap between training noisy states (built from real images) and inference noisy states (built from the model’s own predictions).

Key claims

Correction significantly improves FID over the unguided baseline on DiT, SiT, SiT-REPA, and JiT at ~250 NFEs per image, matching or beating CFG with the same or smaller inference budget [FID table, ImageNet-256 class-conditional].
The gain is not explained by added parameters or compute: a 14-layer denoiser + 14-layer corrector with K=1 correction per step matches a 28-layer monolithic denoiser (FID 8.74 vs 8.93), while K=3 corrections at half the denoising steps (same compute) drops FID to 5.66 [LightningDiT ablation, 80K updates, ImageNet-256].
Correction preserves diversity and avoids oversaturation. Increasing CFG scale on SiT-REPA-XL/2 pushes samples toward similar close-up portraits with saturated colors; increasing K with no CFG keeps more varied poses, framings, and backgrounds [Qualitative side-by-side at matched NFE].
The method extends to text-to-image. On NanoGen (pretrained on BLIP-3o Short-Caption + Long-Caption + JourneyDB subsets, ~36M pairs), correction fixes artifacts without CFG and reduces the “overcooked” look [Qualitative T2I samples].
Human preference favors CFG + correction over CFG only on both quality (60.75% vs 39.25%) and diversity (60.25% vs 39.75%) across 80 GenEval single-object prompts × 10 raters per criterion (1,600 total judgments) via Rapidata [Human study table].
Iterated correction can increase the denoising MSE on training inputs x_t while improving sample quality — the error metric is measured in the wrong distribution [Discussion, Fig. on denoising error vs K].
A two-sample classifier trained to distinguish training noisy states from sampling noisy states on SiT achieves high AUROC without correction, especially at low noise levels; one correction lowers the classifier’s AUROC, four corrections lower it further — i.e., correction brings the inference-time state distribution closer to the training distribution [Classifier two-sample test, held-out AUROC per noise level].
Training the corrector with more training-time iterations (K_train > 1) does not improve performance and roughly scales the optimal inference-time K by n — the corrector’s effective “step size” shrinks proportionally [K_train ablation].
Iterated correction can be distilled back into the denoiser (one denoiser call + K corrector calls → one denoiser call) [Distillation extension; qualitative only on the project page].

Method

The setup is a standard pretrained diffusion / flow-matching denoiser. Correction Training adds a second network — the corrector — of comparable size. During training, for each noisy state x_t the denoiser produces a prediction x̂_0; the corrector is trained to map (x̂_0, x_t, t) back toward the clean target x_0 using the same denoising loss. The denoiser can be kept frozen or jointly updated; the project page frames both as options.

At inference, a single denoising step becomes: (1) one base-denoiser forward pass to produce x̂_0; (2) K corrector forward passes, each taking the current refined prediction as input and producing a new refined prediction, with K a sampling-time hyperparameter. The sampler then advances to the next noise level using the K-th corrected prediction. The total NFE budget is held equal to CFG’s (which also doubles per-step compute) in all headline comparisons.

The authors analyze the mechanism via a two-sample classifier trained to distinguish training noisy states (produced by forward-noising real images) from sampling noisy states (produced by running the denoiser on initial noise). On a standard SiT model the classifier easily distinguishes the two, especially at low noise levels where errors have accumulated; one correction lowers AUROC, four corrections lower it further. The corrector’s job, on this reading, is to pull the model’s trajectory back toward the training manifold even if that increases per-step MSE on training inputs.

A follow-up ablation varies K_train (the number of correction iterations supervised during training). Larger K_train does not improve quality and roughly scales the optimal inference-time K by a factor of n — suggesting the training-time iteration merely spreads the same correction mapping across more small steps.

Results

All ImageNet-256 numbers are class-conditional FID on 50K samples at ~250 total network evaluations per image.

DiT, SiT, SiT-REPA, JiT: correction significantly improves FID over the unguided baseline on all four backbones at matched or smaller NFE than CFG. The exact per-backbone numbers are in the project-page FID table; a representative point is SiT-REPA-XL/2 going from the unguided baseline to a sharply lower FID at K=8 with no CFG [Project-page FID table].
Parameter-matched LightningDiT ablation (80K updates, 675M vs 681M parameters): 28-layer denoiser FID 8.93; 14+14 denoiser + corrector at K=1 FID 8.74; 14+14 at K=3 with half the denoising steps (compute-matched) FID 5.66 [LightningDiT ablation, FID on 8K samples, no guidance].
DiT headline curves: FID drops monotonically with K for ~5 corrections and precision/recall both improve, without the diversity collapse seen when increasing CFG scale [DiT FID + precision-recall figure].
Text-to-image on NanoGen: cleaner object structures and fewer artifacts at K>0 vs K=0. Combining correction with CFG reduces the AI-look and increases diversity relative to CFG alone — same seeds, same 50 denoising steps, same prompt, four cells (w ∈ {1, 6}) × (K ∈ {0, 6}) [T2I qualitative panel].
Human study (Rapidata, 80 GenEval single-object prompts, 10 raters per prompt per criterion, 1,600 judgments total): CFG + correction wins quality 60.75% / diversity 60.25% vs CFG only [Human study table].

Why it’s interesting

Correction Training sits at the intersection of two threads on the wiki. First, within Classifier-Free Guidance variants, it adds a tenth axis: an inference-time refiner module that produces CFG-like quality gains without the CFG two-pass, without the diversity collapse, and without the oversaturation. It differs from Random noise augmentation during diffusion training eliminates the need for CFG (which fixes noise miscalibration at training time with per-channel augmentation) and from Continuous Adversarial Flow Models (which fixes generalization miscalibration via adversarial post-training) by instead adding a separate network and applying it iteratively — a direction nothing else in the cluster currently occupies. Second, as an exposure-bias fix, it is close in spirit to HY-SOAR: Self-Correction for Optimal Alignment and Refinement in Diffusion Models (HY-SOAR), which also trains on model-rolled-out off-trajectory states to close the training–inference gap, but uses a self-correction loss on the base denoiser rather than a separate corrector module. The two-sample classifier result — correction reduces train/inference distinguishability even as it increases training-input MSE — is the first direct empirical validation of the exposure-bias framing that both papers invoke. The distillation extension (correction → single-pass denoiser) is briefly mentioned but not quantified; this is the obvious composition with Diffusion Distillation work.