PaperlayerView exact revision on arXiv

Source-linked AI summary

Correction Training for Diffusion Models

Donghoon Ahn, Runqian Wang, Alexei A. Efros, Taesung Park, Yash Belhe

arXiv:2610.14136v2cs.AI

TL;DR

Diffusion sampling can create a training–inference gap because later inputs depend on the model’s earlier predictions. The paper trains a corrector to refine those predictions repeatedly, reporting improved quality across generation settings and a distillation route that avoids extra inference cost.

  • Problem

    At sampling time, inputs depend on imperfect earlier predictions and can differ from the real-image interpolations used for training, creating a training–inference gap.

    Source: §1 · p. 2, §2.1 · p. 3

  • Method

    A corrector takes the base denoiser’s prediction as input, learns to predict the same clean target, and is applied repeatedly during sampling.

    Source: §1 · p. 2

  • Results

    Correction improves sample quality across class-conditional, unconditional, and text-to-image generation, achieves quality competitive with guidance while better preserving diversity, and can be distilled to remove inference overhead.

    Source: §1 · p. 2, §6 · p. 14

  • Takeaways & Limitations

    Distillation transfers corrected predictions into the base model, retaining much of the quality improvement without additional correction passes at inference.

    Source: §1 · p. 2

    A theoretical understanding of correction’s effectiveness and iterative behavior remains an open question.

    Source: §6 · p. 14

Abstract

from arXiv · show

Iterative sampling in diffusion models can lead to error accumulation due to exposure bias. The denoiser is trained on real images interpolated with noise, but at inference its inputs are constructed from its own imperfect predictions from the earlier steps, something it has never seen at training time. This might be one reason diffusion models often struggle without additional help, such as CFG. We introduce a correction step for diffusion models which attempts to reduce this training-inference gap. The corrector takes the base model's prediction as input and is trained against the same ground-truth target. This makes correction an explicit learning task, with training inputs provided by the base model. During sampling, at each denoising step, we iteratively refine the denoiser's clean-sample prediction. This allows us to correct predictions at intermediate denoising steps, preventing early artifacts from persisting into the final sample. Although the corrector is trained only on predictions from the base denoiser, it can be applied repeatedly to its own outputs at inference for better results. This simple framework leads to substantial improvement in sample quality while largely preserving diversity. We evaluate our method across DiT, SiT, and JiT. Our method substantially improves upon previously reported FID scores for unguided generation and, notably, unconditional generation, where producing high-quality samples without guidance has been particularly challenging. Our method achieves performance competitive with classifier-free guidance (CFG), and combining correction with CFG further improves diversity and reduces oversaturation typical of CFG. Finally, we show that the additional inference cost of correction can be amortized by distilling the refined predictions back into the base diffusion model, retaining the quality gains without additional inference-time overhead.

1 INTRODUCTION

The introduction frames exposure bias as a mismatch between training and sampling inputs, then proposes iterative correction of the denoiser’s predictions. It reports quality improvements across generation settings and describes distillation as a way to avoid extra inference cost.

1 INTRODUCTION

  • Exposure bias describes the training–inference mismatch: sampling inputs depend on imperfect earlier predictions, and errors can accumulate into artifacts and incoherent structures.The introduction contrasts this with CFG, which improves sample quality at the cost of reduced diversity and oversaturation.Source: §1 · p. 2
  • The corrector learns from base-model predictions against the same clean target, then repeatedly refines clean-sample predictions during sampling.The introduction hypothesizes that repeated correction helps compensate for errors accumulated over earlier denoising steps.Source: §1 · p. 2, Figure 2 · p. 2
  • Correction improves unguided quality across class-conditional, unconditional, and text-to-image generation, with better quality at matched diversity than CFG and less oversaturation.The introduction also reports that adding correction to CFG reduces the stereotypical “AI-generated” look of its samples.Source: §1 · p. 2
  • Distillation transfers corrected predictions back into the base model, retaining much of the quality improvement with one forward pass per sampling step.The distilled model preserves the base architecture, parameter count, and inference cost.Source: §1 · p. 2

2 LEARNING TO CORRECT DENOISING PREDICTIONS

The section frames correction as a response to sampling inputs that differ from training inputs, then develops learned refinement of clean-sample predictions during sampling. It reports gains in unguided generation and more effective use of inference compute at matched cost.

2.1 BACKGROUND: DENOISING AND SAMPLING

  • Training interpolates clean images with noise, whereas sampling inputs depend on earlier model predictions, creating a training–inference gap.At each sampling step, the predicted clean image is used to construct the next noisy state.Source: §2.1 · p. 3

2.2 PREDICTION ERROR DURING SAMPLING

  • Early structural errors can persist as later denoising adds details, so the paper makes correction an explicit task applied to the denoiser’s predictions.Correction refines the prediction at a fixed noise level before it updates the sampler state.Source: §2.2 · p. 3

2.3 LEARNING TO CORRECT MODEL PREDICTIONS

  • The corrector learns to map the denoiser’s initial prediction to the same clean target, with denoising and correction objectives trained jointly on online-generated inputs.The corrector can build on structure already captured by the initial prediction and focus on remaining inaccuracies.Source: §2.3 · p. 4

2.4 ITERATIVE CORRECTION AT INFERENCE

  • Training supervises one correction pass, but inference can repeatedly apply the corrector to its own clean-sample estimates.Each iteration remains an estimate of the same clean image.Source: §2.4 · p. 4
  • At matched total inference cost, correction is more effective than simply increasing the number of denoising steps.Successive corrections refine intermediate predictions throughout the sampling trajectory.Source: §2.4 · p. 4
  • Correction consistently improves FID and precision over unguided baselines and is competitive with CFG on ImageNet-256 class-conditional generation across four models at approximately 250 NFEs.The table compares FID-50K, precision, and recall.Source: Table 1 · p. 5

3 RESULTS

Correction improves unguided diffusion generation across ImageNet and text-to-image settings, and can be paired with CFG or distilled to reduce inference cost. Across these evaluations, it improves quality while retaining diversity, with gains also reported at matched inference budgets.

3.1 IMAGENET GENERATION

  • At matched NFE, correction lowers FID from 9.58 to 2.07 on DiT, 8.55 to 1.92 on SiT, and 5.53 to 1.31 on SiTREPA.The reported trade-off is higher precision at the cost of some recall; SiTREPA uses 50K correction updates, 1.25% of its 4M pretraining updates.Source: §3.1 · p. 6
  • For unconditional ImageNet generation, correction improves quality without guidance and at K = 4 substantially improves FID over both the baseline and best evaluated autoguidance scale.Larger K trades recall for higher precision.Source: §3.1 · p. 6, Table 2 · p. 7
  • Correction preserves more varied poses, framings, and backgrounds than higher CFG scales, which produce less diverse, increasingly saturated samples.The comparisons vary correction depth K and CFG scale w.Source: §3.1 · pp. 6–7, Figure 9 · p. 7

3.2 TEXT-TO-IMAGE GENERATION

  • At approximately matched NFE, correction raises GenEval and DPG-Bench scores both without CFG and when applied before CFG.The comparisons use the same-checkpoint baseline.Source: §3.2 · p. 7, Table 3 · p. 8
  • With CFG, correction alleviates reduced diversity while retaining quality refinement; raters preferred CFG plus correction over the default-CFG baseline in more than 60% of responses.The human preference study used 80 GenEval single-object prompts.Source: §3.2 · pp. 8–9, Figure 11 · p. 8, Table 4 · p. 9
  • On unconditional web-image generation, correction improves realism while preserving diversity and the overall structure of uncorrected samples.Source: §3.2 · p. 9

3.3 MAKING CORRECTION EFFICIENT

  • A shared backbone achieves performance comparable to separate networks with roughly half the total parameters, and distillation yields comparable FID while reducing NFE by approximately 67%.The student matches predictions after k* correction iterations and requires one forward pass per sampling step.Source: §3.3 · p. 9

4 DISCUSSION

The discussion finds that repeated correction can outperform parameter-matched deeper denoisers, even when denoising error rises. It examines links to the training–inference gap and how correction depth and denoiser updates affect results.

4.1 PARAMETER-MATCHED BASELINES

  • From scratch, D7 + C7 with K = 3 achieves FID 4.21 versus 4.88 for D28, using the same denoising steps and roughly half as many parameters.For D14 + C14 versus D28, a single correction (K = 1) gives similar FID: 8.74 versus 8.93.Source: §4.1 · p. 10

4.2 HOW ITERATIVE CORRECTION CHANGES PREDICTIONS AND SAMPLING

  • Repeated correction makes predictions more coherent, but can also produce overly regular structures and move predictions farther from the optimal denoising prediction.The denoising error increases after correction for k ≥ 3, even as visual coherence improves.Source: §4.2 · pp. 11–12, Figure 13 · p. 11, Figure 14 · p. 11
  • Classifier tests find that correction makes training and inference states harder to distinguish, with iterative correction generally reducing the gap further.The authors suggest this may contribute to sample quality even when reconstruction error on training inputs increases.Source: §4.2 · p. 12, Figure 15 · p. 12

4.3 HOW TRAINING AFFECTS ITERATIVE CORRECTION

  • The best-performing inference correction count exceeds the supervised training depth, while denoising error remains above the base denoiser for correction counts beyond that depth.Increasing Ktrain shifts the best-performing inference K higher while the best achieved FID remains similar.Source: §4.3 · p. 12, Figure 16 · p. 12
  • Freezing the denoiser gives similar FID at K = 5, but further correction increases FID, unlike joint training, which continues to lower it.The authors suggest a jointly updated denoiser exposes the corrector to a broader range of errors and may help it generalize across correction counts.Source: §4.3 · p. 13, Figure 17 · p. 13

5 RELATED WORK

The paper relates its explicit prediction corrector to prior iterative and self-conditioning approaches while distinguishing its training objective and use during diffusion sampling.

5 RELATED WORK

  • Unlike self-conditioning, which supplies previous clean-data predictions as denoiser inputs, this method trains an explicit corrector on denoiser predictions and reuses it across iterations.Source: §5 · p. 13
  • The method retains the denoising objective and applies correction at each noise level before the sampling update, without requiring the original noisy input.Source: §5 · p. 13
  • The closest related work identified is Schiff et al. (2026), which repeatedly applies a corrector trained on generated samples in discrete-diffusion language modeling.Source: §5 · p. 13

6 CONCLUSION

Correction training refines diffusion predictions using the original clean-data targets, and the corrector can be applied repeatedly at inference. The conclusion reports sample-quality gains across several generation settings, while the appendix describes practical design choices for implementing correction training.

6 CONCLUSION

  • The corrector is trained to correct diffusion-model predictions against the original clean-data targets, despite being trained for a single correction pass.The conclusion states that it supports repeated refinement at inference.Source: §6 · p. 13
  • Correction training substantially improves sample quality in class-conditional and unconditional ImageNet and text-to-image generation.It can achieve quality competitive with guidance while better preserving diversity, and can be distilled to remove inference overhead.Source: §6 · pp. 13–14
  • A theoretical understanding of correction’s effectiveness and iterative behavior remains an open question.Source: §6 · p. 14

A DESIGN CHOICES AND ALTERNATIVE REALIZATIONS

  • The appendix’s default setup uses a separate corrector initialized from the base denoiser, omits the noisy input, retains timestep conditioning, and applies stop-gradient to protect denoiser updates.The corrector loss does not update the denoiser under this setup; design choices are varied one at a time.Source: A DESIGN CHOICES AND ALTERNATIVE… · p. 18

A.1 CORRECTOR INTERFACE

The interface, training, and inference ablations examine which correction inputs and configurations work well. They report trade-offs for noisy-state conditioning and timestep inputs, differences between output-space and hidden-state refinement, and sensitivity to stop-gradient, correction count, and correction interval.

A.1 CORRECTOR INTERFACE

  • Proposal-only correction achieves higher precision at slightly lower recall than conditioning on both the proposal and noisy input.The SiT comparison is shown in Figure 19; on JiT, proposal-only correction achieves comparable FID to z_t-conditioned correction.Source: §A.1 · p. 18, Figure 19 · p. 18, Table 8 · p. 18
  • True and fixed timestep conditioning yield similar final FID in the tested residual-correction setting, but this may not generalize to other settings.The pretrained timestep pathway encodes the original noisy input’s noise level, not the errors in the clean-sample proposal.Source: §A.1 · p. 18, Table 9 · p. 19
  • Output-space correction benefits from repeated application, whereas hidden-state refinement does not achieve comparable gains; increasing correction iterations during training does not improve sample quality.The comparison is reported in Table 10 and illustrated qualitatively at K = 5 in Figure 20.Source: §A.1 · p. 19, Table 10 · p. 19, Figure 20 · p. 19

A.2 TRAINING CONFIGURATION

  • Full stop-gradient yields similar best FID to alternatives in additional pretrained-model comparisons and degrades less at larger correction counts.In an earlier comparison, removing stop-gradient gave a lower best FID, but samples became visually simpler under repeated correction.Source: §A.1 · p. 19, §A.2 · p. 20, Figure 21 · p. 20

A.3 INFERENCE CONFIGURATION

  • Increasing K improves FID up to a setting-dependent optimum, after which further correction can degrade it.Figure 21 compares FID@5K across K using EMA weights, the full correction interval, and at most 250 NFE.Source: §A.3 · p. 20, Figure 21 · p. 20
  • For JiT, restricting correction to early steps and reallocating saved calls to denoising improves FID at a fixed sampling budget.Table 11 identifies start and end noise levels; the figures compare early-only, skip-early, and full-range correction, with matched-noise examples.Source: §A.3 · p. 20, Table 11 · p. 20, Figure 22 · p. 21

B.1 TRAINING CONFIGURATIONS

The appendix documents backbone and sampling configurations, then discusses alternative corrector parameterizations and gradient paths. It presents clean-sample input as a practical starting point, while describing options for preserving pretrained output representations and combining corrector and base predictions.

B.1 TRAINING CONFIGURATIONS

  • The experiments retain each pretrained model’s backbone and conditioning, use a separate corrector copy, and preserve each model’s VAE or direct-RGB setup.The training and quantitative sampling settings are summarized separately; the main quantitative comparisons use approximately 250 network evaluations per image.Source: §B.1 · p. 23, Table 12 · p. 23, Table 13 · p. 24

B.2 CORRECTOR PARAMETERIZATION

  • Small-scale single-correction experiments found that clean-sample-space correction performed best among tested same-space input–output choices, but do not establish a universally optimal parameterization.The tests used K = 1 without additional noisy-state conditioning; direct noise- or velocity-space correction remains an option.Source: §B.2 · p. 25
  • For pretrained denoisers, clean-sample input with gated native-space output preserves the pretrained output representation without supplying z_t to the corrector.The corrector output is combined with the retained base prediction through a learned gate; the gate is initialized to zero, preserving the base prediction at initialization.Source: §B.2 · p. 25
  • The fixed-base-component formulation uses v = ϵ − x to connect input and output parameterizations, but was tested only in small-scale single-correction experiments.Repeated correction and larger-scale settings were not extensively investigated.Source: §B.2 · p. 26
  • The default full stop-gradient detaches both the corrector input proposal and retained base prediction, so the separate denoiser is updated only by its denoising loss.Full stop-gradient achieves competitive FID and degrades less at large correction counts in the reported comparisons.Source: §B.2 · p. 26
  • Residual prediction can be used with gating, but the paper reports no dedicated comparison; full-target prediction is retained in the main formulation for simplicity.The residual output projection is zero-initialized to preserve the base prediction at initialization.Source: §B.2 · p. 26

B.3 FIGURE AND TABLE DETAILS

The appendix specifies settings for unconditional generation, text-to-image examples, guidance comparisons, parameter sharing, distillation, and power-spectrum analyses. These details describe model checkpoints, sampling budgets, and evaluation procedures rather than reporting comparative outcomes.

B.3 FIGURE AND TABLE DETAILS

  • Unconditional ImageNet evaluation uses SiTREPA-XL/2 with the null class, 50,000 images, and sweeps over correction count and autoguidance settings.Correction uses detached proposal and base predictions; the K sweep adjusts denoising steps to keep the budget near 250 NFE.Source: §B.3 · p. 27
  • Text-to-image examples compare K = 0 and K = 6 without CFG, using 50 Euler steps and matched prompts and rank seeds.The paired settings use 50 and 350 NFE, respectively.Source: §B.3 · p. 27, Table 17 · p. 27
  • Guidance comparisons vary CFG scale and correction count at 50 Euler steps, applying CFG before correction.The selected prompt is a photo of a bench, with all eight seeds retained per setting.Source: §B.3 · p. 27
  • The parameter-sharing comparison uses plain SiT-XL/2 without CFG and evaluates FID50K after 100K updates.Separate and shared correctors use different correction counts and denoising steps; shared parameters receive gradients from both losses.Source: §B.3 · pp. 27–28
  • Distillation trains a student to match a frozen EMA teacher’s K = 2 predictions on class-conditional ImageNet, without a ground-truth loss.The student is initialized from the teacher’s denoiser and evaluated with FID@50K using the native SiT Euler–Maruyama sampler.Source: §B.3 · p. 28
  • The power-spectrum analysis averages radial power across 100 ImageNet-256 samples for K = 0 through 10 at approximately 250 NFE.It plots log10 power and log10(Psample/Preal) over normalized frequencies 0.1–0.5.Source: §B.3 · p. 28

C.1 COMPARING CORRECTION AND CFG

Correction and CFG show different visual trade-offs, and correction reaches a lower minimum FID than the best evaluated CFG setting at approximately matched sampling compute. When combined, the two methods’ order trades off quality and diversity preference rather than consistently leading on both.

C.1 COMPARING CORRECTION AND CFG

  • As guidance scale increases, displayed samples show more similar close-up compositions and stronger colors, while correction retains more varied poses and backgrounds.These selected examples do not establish distributional mode collapse.Source: §C.1 · p. 28, Figure 24 · p. 28
  • Correction reaches a lower minimum FID than the best evaluated CFG setting at approximately matched sampling compute.Figure 25 also reports precision–recall and saturation trade-offs; its dashed orange line marks the best evaluated CFG FID and its green line marks real-image saturation.Source: §C.1 · p. 28, Figure 25 · p. 29
  • Stronger CFG raises mean HSV saturation above the real-image reference, whereas correction remains near or below the unguided level.This ImageNet measurement complements text-to-image examples that illustrate refinement at weaker guidance, rather than measuring a saturation reduction.Source: §C.1 · p. 28
  • Additional SiT-XL/2 comparisons show fixed-noise samples across nine selected classes, using separately trained models at approximately 250 network evaluations.The class examples are illustrative, and their ImageNet-wide FID labels are not estimates for the displayed classes.Source: §C.1 · pp. 28–29, Figure 26 · p. 30, Figure 27 · p. 31

C.2 COMBINING CFG AND CORRECTION

  • At matching guidance and correction settings, correction before CFG more often yields higher quality preference, while correction after CFG more often yields higher diversity preference.Neither order consistently leads on both criteria; Table 18 reports the full vote-share sweep against released CFG.Source: §C.2 · p. 33, Table 18 · p. 33

C.3 HUMAN EVALUATION

The human evaluation compares CFG-then-correction against a released baseline using separate quality and diversity judgments. The supplied passages describe the evaluation design and stimuli but do not report preference outcomes.

C.3 HUMAN EVALUATION

  • The evaluation collects human judgments through Rapidata, with annotator qualification and reliability filtering, plus English-language and adult-age filters.The passage describes the service as supporting high-throughput collection of diverse judgments.Source: §C.3 · p. 33
  • Table 4 uses 800 judgments per metric across 80 GenEval single-object prompts, with one image for quality and four-image sets for diversity.Both methods use 50 Euler steps, and the evaluated method applies CFG before correction using the conditional text.Source: §C.3 · p. 33
  • The full sweep compares both correction orders across w ∈ {1, 2, 4, 6} and K ∈ {2, 4, 6, 8}, separately against the same released baseline.At w = 1 the orders use identical images, so differences reflect repeated human ratings; Figure 29 shows the main-table setting’s stimuli and questions.Source: §C.3 · p. 33, Figure 29 · p. 34

D.1 TRAINING AND CAPACITY CONTROLS

Correction improves generation quality beyond continued denoiser training and deeper-denoiser controls, and iterative correction also outperforms a deeper model at matched effective depth. The visual analyses show how correction refines intermediate predictions and changes frequency content, with further iterations capable of overshooting the real-image spectrum.

D.1 TRAINING AND CAPACITY CONTROLS

  • Over the 50K-update training window, single-pass correction improves FID substantially faster than continued training or either deeper-denoiser control.The correction model has the same total sequential transformer depth as the 56-layer denoisers in this comparison.Source: §D.1 · p. 35, Figure 30 · p. 35
  • At matched effective depth and inference budget, D7 + C7 achieves FID@50K 4.21 versus 4.88 for D28, using 346M versus 675M parameters.Both configurations use 240 denoising steps and 6,720 transformer-layer evaluations per sample; Figure 31 presents the comparison.Source: §D.1 · p. 35, Figure 31 · p. 35

D.2 ITERATIVE CORRECTION VISUALIZATIONS

  • Additional iterative-correction visualizations include selected examples at fixed noisy inputs, shared-input noisy-state comparisons, and joint-versus-frozen-denoiser samples.The examples are qualitative; the analytic comparison uses a uniform empirical prior over a finite same-class image bank excluding the original query cohort, and corrected predictions need not reconstruct the original image.Source: §D.2 · p. 36, Figure 32 · p. 37, Figure 33 · p. 38
  • At fixed timestep, correction refines the clean-sample prediction before the sampler updates its noisy state, sharpening structure and suppressing visible artifacts in examples.In the dog example, corrections at t = 0.73–0.60 reveal a coherent shape that carries through to the final image, while the baseline retains artifacts.Source: §D.2 · p. 36, Figure 34 · p. 39

D.3 POWER SPECTRUM ANALYSIS

  • Increasing correction iterations reduces the baseline’s high-frequency spectral deficit, with K = 6 approximately matching the real-image reference.Further iterations can raise both low- and high-frequency power above the reference.Source: §D.3 · p. 41, Figure 36 · p. 41

E ADDITIONAL SAMPLES

The additional galleries show correction-depth progressions and paired baseline-versus-correction samples across SiT, SiTREPA, JiT, and unconditional-generation settings. Their captions specify shared classes or initial noise for comparisons, while the text-to-image-model galleries show unconditional samples without text or CFG.

E.1 EFFECT OF CORRECTION DEPTH

  • The correction-depth gallery follows the same class and initial noise as K increases, then presents paired galleries over randomly drawn ImageNet classes.The examples use class conditioning without CFG.Source: §E.1 · p. 41, Figure 37 · p. 42

E.2 SIT: RANDOM-CLASS GENERATION

  • SiT random-class samples pair baseline K = 0 and correction K = 6 at approximately matched NFEs, sharing class and initial noise.The galleries retain uncurated samples.Source: Figure 38 · p. 43

E.3 SITREPA: RANDOM-CLASS GENERATION

  • SiTREPA random-class samples pair baseline K = 0 and correction K = 6 at approximately matched NFEs, sharing class and initial noise.The galleries retain uncurated samples.Source: Figure 39 · p. 44

E.4 JIT: RANDOM-CLASS GENERATION

  • JiT random-class samples pair baseline K = 0 and correction K = 6 at matched NFEs, sharing class and initial noise.The gallery notes that early correction changes structure as well as appearance.Source: Figure 40 · p. 45

E.5 UNCONDITIONAL GENERATION

  • Selected SiTREPA-XL/2 ImageNet samples illustrate unconditional generation with iterative correction and without guidance.The caption identifies these as selected samples.Source: Figure 41 · p. 46

E.6 UNCONDITIONAL GENERATION FROM A TEXT-TO-IMAGE MODEL

  • Figures 42–45 show uncurated unconditional-generation samples without text or CFG; within each figure, rows share initial noise across K = 0 through 5.The checkpoint and denoising steps are fixed; larger K adds correction compute. All seeds are shown without selection.Source: Figure 42 · p. 47, Figure 43 · p. 48, Figure 44 · p. 49

Related papers

Loading 2610.14136v2…