cs.AIOct 9, 2026
SaveCorrection Training for Diffusion Models
Organizations: University of California, Berkeley · Reve
Abstract
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.
Figures & tables
Figure 1: ImageNet samples with correction. Without the use of guidance , our correction method can achieve high-quality class-conditional (top) and unconditional (bottom) image samples.
Figure 2: Denoising and prediction correction. During training, the denoiser and corrector predict the same clean target. Stop-gradient blocks the correction loss through the denoiser’s prediction. During sampling, each denoising prediction is refined times before the next sampling step.
Figure 3: Prediction correction improves structural coherence during sampling. We show clean-sample predictions along denoising trajectories without (a) and with (b) correction. In (a), early structural errors persist as later denoising steps add details. In (b), iterative correction revises the prediction at each , producing more coherent object structures before the next denoising step. Each triplet shows this refinement within one denoising step. Additional visualizations are shown in Appendix D.2 (Fig. 36 ).
Figure 4: Correction improves class-conditional sample quality. We compare baseline (top) and correction (bottom), starting from the same noise in each column. Correction gives more coherent and detailed images for different classes and noises compared against the baseline.
| DiT-XL/2 | FID | Prec. | Rec. |
|---|---|---|---|
| Unguided (baseline) | 9.58 | 0.668 | 0.676 |
| Corrective sampling | 2.07 | 0.811 | 0.603 |
| CFG (reference) | 2.51 | 0.818 | 0.581 |
| SiT-XL/2 | FID | Prec. | Rec. |
| Unguided (baseline) | 8.55 | 0.687 | 0.669 |
| Corrective sampling | 1.92 | 0.812 | 0.608 |
Table 1: Correction improves unguided class-conditional generation. We compare FID-50K, precision, and recall on ImageNet-256 class-conditional image generation across 4 base models. We report consistent improvement on FID and precision over unguided baseline and it is competitive with CFG. All methods are evaluated at approximately 250 NFEs. See Tab. 13 for evaluation details.
Figure 5: Efficiency of correction training. FID converges within relatively few corrector training steps.
Figure 6: Iterative correction provides a controllable refinement trajectory on ImageNet. Left: Evolvement of the same samples over more correction steps. The main content remains the same while quality improves. Right: FID and precision–recall across . Correction gives a controllable tradeoff between coverage and quality and lower minimum FID.
Figure 7: Correction improves unconditional sample quality on ImageNet. Without correction, generated samples exhibit artifacts and structural inconsistencies (top), while corrective sampling produces diverse and high-quality images (bottom). The samples in each column share the same initial noise.
Figure 10: Correction preserves more diversity and produces less oversaturation than CFG. Samples from SiT REPA -XL/2 with increasing correction depth (bottom) or CFG scale (top), at approximately matched NFE, same initial noise.
Figure 11: Correction improves object structure in text-to-image generation. Compared with no correction (top), correction (bottom) separates the violin bodies, clarifies the guitar and car geometry, and organizes cluttered interiors. All samples are without CFG; full prompts are in Appendix B.3 .
Figure 13: Correction improves the diversity of CFG samples while preserving quality. Columns compare no guidance ( ) and CFG ( ); rows compare no correction ( ) and correction ( ). All settings use the same eight seeds and number of denoising steps. Corrected samples show more varied bench shapes and viewpoints, with less oversaturation and a more natural appearance. Prompt: “a photo of a bench”.
| Method | Quality | Diversity |
|---|---|---|
| Baseline | 39.25% | 39.75% |
| Correction | 60.75% | 60.25% |
Table 2: Correction is preferred for quality and diversity. Released CFG ( ) versus CFG followed by correction ( , ).
Figure 14: Matched parameters. Repeated correction outperforms a deeper denoiser at a comparable parameter count and matched inference budget.
| Model | Params (M) | Sampling | Steps | Layer evals | FID@8K | FID@50K |
| Matched parameters | ||||||
| 675 | Unguided | 240 | 6,720 | 8.93 | – | |
| 675 | CFG ( ) | 120 | 6,720 | 6.17 | – | |
| + | 681 | 240 | 6,720 | 8.74 | – | |
| + | 681 | 120 | 6,720 | 5.66 | – | |
| + | 681 | 96 | 6,720 | 5.56 | – | |
Table 5: Iterative correction outperforms matched baselines. Although a same-layer denoiser-corrector pair achieves similar performance with a large monolithic denoiser with a single correction ( + , ), iterative correction ( + , ) outperforms the monolithic baseline. Even the smaller pair outperforms the larger monolithic baseline with . We select the best CFG scale from .
Figure 15: Iterative correction sharpens predictions. We show the effect of iterative correction on JiT-L/16. (a) The base denoiser’s predictions ( ) are often incoherent and contain artifacts. Repeated correction resolves ambiguous structures and makes faces, object contours, and other details more coherent. (b) In extreme cases, correction appears to push this coherence too far, turning ambiguous structures in the base prediction into overly regular geometric shapes. Further selected examples appear in Appendix D.2 . Analytic target is computed using Eq. 10 in Appendix.
Figure 16: Iterative correction can increase denoising error.
Figure 17: Measuring the training–inference gap with classifier two-sample tests. We train a classifier to distinguish samples from the training distribution and the inference distribution, and report AUROC and binary cross-entropy (BCE). Overall, correction makes classification harder, indicating a smaller gap between training and inference.
Figure 18: Correction beyond the supervised training depth improves generation (lower FID). Supervising more correction steps shifts the best-performing tested inference correction count toward larger values, often approximately twice the maximum supervised depth. This suggests that “extrapolating” the learned correction beyond its training inputs may play an important role in the observed gains.
Figure 19: A frozen denoiser can support high quality refinement, but iterative correction can also oversaturate samples and produce abrupt changes in the sample prediction.
Figure 20: Frozen-denoiser correction is competitive at small but deteriorates earlier.
Appendix figures & tables36 assets
Supplementary material from the paper’s appendix.
Appendix
| Variant | ||||||
|---|---|---|---|---|---|---|
| Proposal only ( ) | 24.63 | 17.96 | 14.07 | 13.35 | 14.27 | 16.06 |
| With conditioning | 24.72 | 18.22 | 14.84 | 13.79 | 14.67 | 16.58 |
Table 6: Proposal-only correction performs comparably to -conditioned correction. Unguided FID50K on JiT-B/16 at approximately 250 NFE.
Figure 21: Proposal-only correction reaches higher peak precision. SiT on ImageNet with approximately capacity-matched correction heads, sweeping – . Stars mark each arm’s lowest-FID setting.
| Corrector time input | FID | Precision | Recall |
|---|---|---|---|
| True timestep | 2.825 | 0.763 | 0.721 |
| Constant time embedding | 2.895 | 0.745 | 0.734 |
Table 7: Explicit time conditioning has little effect on FID in this setting. We compare true and fixed timestep inputs to the corrector on SiT-REPA at , using FID-50K at 250 NFE. The denoiser uses the true timestep in both cases.
| Stop- gradient | Inference-time iterations | ||||||
|---|---|---|---|---|---|---|---|
| Correcting model predictions | |||||||
| ✓ | 24.63 | 17.96 | 14.07 | 13.35 | 14.27 | 16.06 | |
| ✗ | 29.08 | 17.38 | 12.88 | 19.25 | 24.64 | 28.92 | |
| ✓ | – | 24.66 | 21.02 | 21.68 | 22.51 | 23.30 | 24.67 |
| Iterating over hidden states | |||||||
Table 8: Output-space iteration enables effective iterative correction. Component ablations on JiT-B/16, evaluated by unguided FID50K at approximately 250 NFE.
Figure 22: Output-space correction remains effective under repeated application. Qualitative comparisons of the component ablations in Tab. 8 at , using the same initial noise and class labels.
Figure 23: Full stop-gradient degrades more slowly at large correction counts in these comparisons. FID@5K versus , using EMA weights, the full correction interval, and at most 250 NFE. SG denotes stop-gradient; full SG blocks both proposal and base-combination paths.
| End | |||||
|---|---|---|---|---|---|
| Start | |||||
| 7.70 | 6.79 | 7.72 | 9.05 | 11.85 full | |
| — | 11.16 | 11.12 | 12.59 | 13.76 | |
| — | — | 14.54 | 15.09 | 17.94 | |
| — | — | — | 17.10 | 20.05 | |
| — | — | — | — | 19.96 | |
Table 9: Allocating correction to early steps yields the lowest FID at a fixed sampling budget. JiT-L/16 with , evaluated by FID8K. Rows and columns specify the start and end noise levels, decreasing from noise to data.
Figure 24: Skipping early correction refines the baseline pose. Gray boxes show the matched baseline; green, blue and orange mark early-only, skip-early and full-range correction, respectively. Restricting correction to early steps changes the generated structure. Matched-noise comparisons of full-trajectory correction and the interval selected by FID8K for each .
Figure 25: (a) Full-range correction changes fine texture and shape. (b) Full-range correction changes object layout and appearance. Gray boxes show the matched baseline; green, blue and orange mark early-only, skip-early and full-range correction, respectively.
| Setting | DiT-XL/2 | SiT-XL/2 | SiT REPA -XL/2 | JiT-L/16 | T2I (Sec. 3.2 ) |
|---|---|---|---|---|---|
| Data and initialization | |||||
| Training data | ImageNet | ImageNet | ImageNet | ImageNet | BLIP3o mixture |
| Image resolution | |||||
| Denoising space | SD-VAE latents | SD-VAE latents | SD-VAE latents | RGB pixels | FLUX.2-VAE latents |
| Denoiser init. | Released | Released | Released | Released | 202K base |
| Corrector init. | Independent copy of the pretrained denoiser | ||||
Table 10: Iterative refinement uses correctors trained with a single correction pass. Training settings for the main ImageNet and quantitative text-to-image comparisons. We follow each model’s original training recipe as closely as possible. Updates count additional training, not pretraining; and denote separate denoiser and corrector networks. Sampling settings are reported separately.
| Method | CFG | CFG interval | Correction interval | Steps | NFE | |
| DiT-XL/2; DDPM | ||||||
| Baseline | 1 | 0 | — | — | 250 | 250 |
| CFG | 1.5 | 0 | — | 125 | 250 | |
| Correction | 1 | 2 | — | 83 | 249 | |
| SiT-XL/2; Euler–Maruyama SDE | ||||||
| Baseline | 1 | 0 | — | — | 250 | 250 |
Table 11: The main quantitative comparisons use approximately 250 network evaluations per image. Sampling settings for Tables 1 and 12 . is the number of corrections per active step; intervals use noise time ( : noise, : data). NFE includes correction and both CFG branches. Guided T2I applies correction before CFG.
| Design choice | Option | Used in |
|---|---|---|
| Output formulation | Direct | Table 1 — JiT |
| Gated native | Table 1 — DiT, SiT, SiT-REPA; Tables 9 – 12 | |
| Prediction type | Full target | Tables 1 – 12 |
| Residual | Table 7 — timestep conditioning |
Table 14: Corrector parameterization choices. We distinguish the output formulation from full-target versus residual prediction, which can be used with either formulation. Here denotes the chosen output space. The equations use clean-sample input by default; input parameterization and loss space are separate choices.
| Figure column | Full prompt |
|---|---|
| 1. Green guitar | In the center of a dimly lit stage, a vintage green guitar with a glossy finish leans gently against a classic black amplifier. The polished wooden flooring of the stage reflects the warm, golden hue from a single spotlight above, which casts an inviting glow over the instrument. A microphone stand with a chrome finish stands to the left of the guitar, poised for the evening’s performance. |
| 2. Vintage car | a vintage blue Porsche 356 is captured mid-turn on a winding asphalt road, its polished surface reflecting the bright sunlight. The car’s classic design is accentuated by its rounded headlights and sleek bodywork. Alongside the road, a low stone wall partially covered in moss provides a contrast to the well-maintained vehicle. |
| 3. Sunlit kitchen | A spacious kitchen with a large wooden table at the center, cluttered with an assortment of pots, pans, and cooking utensils. The table is surrounded by matching wooden chairs, and the walls are lined with shelves filled with spices and kitchenware. Sunlight streams in through a window, casting a warm glow on the array of cookware. |
| 4. Three violins | Three exquisitely crafted violins, with their elegant F-shaped sound holes and smooth, glossy wooden surfaces, rest gently beside a grand piano. The piano’s polished ebony finish reflects the soft, gentle curves of the violins that lie on a velvet-lined case. Behind them, the piano’s open lid reveals its intricate strings and hammers, inviting a symphony of sounds to be played. |
| 5. Rabbit and seashell | On the soft, warm sand of the beach, a fluffy white rabbit with rounded ears is caught in a curious moment, gently placing its paw on the ribbed surface of a pink scallop shell. The scallop, slightly open, reveals its smooth interior contrasting with its coarse outer texture, while hues of pink and orange from the setting sun reflect off its surface. There’s a tranquil ocean backdrop with the gentle ebbing of the tide, and the fading daylight casts a golden glow over the scene, highlighting the rabbit’s soft fur and the shell’s subtle color. |
| 6. Victorian library | A voluminous library, bathed in soft, muted blue hues, stretches majestically across a spacious Victorian-style living room. Intricately carved wooden bookshelves, filled with an array of books of various sizes and colors, line the walls and a grand fireplace sits as the room’s focal point. The vray render showcases hyperrealistic textures, from the plush, patterned rug beneath the polished wooden tables, to the delicate fabric of the antique armchairs scattered throughout the space. |
Table 15: Full prompts for the text-to-image examples in Figure 12 .
Figure 26: Correction and guidance produce different visual trade-offs. Matched-noise samples from SiT REPA -XL/2 as correction depth (top) or guidance scale (bottom) increases, at approximately matched NFE.
Figure 27: Correction improves FID with a different precision–recall and saturation trade-off. DiT-XL/2 at approximately matched NFE. The dashed orange line marks the best evaluated CFG FID; the green line marks real-image saturation.
Figure 28: Correction and guidance comparisons. (a) toucan, (b) lorikeet, and (c) hummingbird samples from SiT-XL/2. For each class, the upper row varies correction depth without CFG; the lower row varies guidance scale in a separately trained denoising-only model. Both models use EMA weights after 100K updates and approximately 250 NFE. All 16 fixed seeds appear in the same row-major order in every panel. FID labels use 50K samples across ImageNet classes.
Figure 29: Correction and guidance comparisons. (a) goldfish, (b) macaw, and (c) goldfinch samples from SiT-XL/2. For each class, the upper row varies correction depth without CFG; the lower row varies guidance scale in a separately trained denoising-only model. Both models use EMA weights after 100K updates and approximately 250 NFE. All 16 fixed seeds appear in the same row-major order in every panel. FID labels use 50K samples across ImageNet classes.
Figure 30: Correction and guidance comparisons across additional classes. (a) Daisy, (b) golden retriever, and (c) volcano. SiT-XL/2 with the same checkpoints, sampling settings, and 16 fixed seeds as Figs. 28 and 29 . Top: correction without CFG. Bottom: CFG . FID labels are ImageNet-wide FID@50K.
| CFG Correction | Correction CFG | ||||
|---|---|---|---|---|---|
| Quality | Diversity | Quality | Diversity | ||
| 1 | 2 | 33.38 | 55.38 | 37.25 | 54.50 |
| 1 | 4 | 49.38 | 57.25 | 47.50 | 59.38 |
| 1 | 6 | 55.75 | 55.75 | 52.38 | 60.75 |
| 1 | 8 | 54.13 | 56.38 | 56.75 | 63.88 |
| 2 | 2 | 48.50 | 60.50 | 55.38 | 59.13 |
Table 16: Neither correction order consistently leads in both quality and diversity. Full sweep of candidate vote shares (%) against released CFG ( , ); higher is better. Bold marks the setting reported in Table 2 . Each entry contains 800 responses across 80 prompts. The two orders were rated separately against the baseline.
Figure 31: Quality and diversity are assessed separately. Actual teddy-bear stimuli and questions for Table 2 , using CFG followed by correction ( , ): single images for quality and four-image sets for diversity.
Figure 32: Correction improves generation quality beyond continued denoiser training or parameter-matched deeper denoisers.
Figure 33: Matched effective depth. Correction outperforms a deeper denoiser at matched effective depth per denoising step and inference budget, using fewer parameters.
Figure 34: Additional iterative-correction examples. The seven selected examples not included in Fig. 15 (a), using the same column layout and fixed-input protocol. Repetition strengthens object contours and details, while the predictions need not align with the original clean image or the finite-bank analytic estimate. Additional overcorrection examples. The three selected examples not included in Fig. 15 (b). Repeated correction can simplify irregular textures or amplify thin structures into unnatural patterns.
Figure 35: Shared-input noisy-state comparisons. Clean images (top) and interpolated noisy states (bottom), comparing Real , , and . Panels (a,b,g,h) use ; panels (c,d,e,f) use . Source queries (classes), in panel order, are 264649 (207), 114599 (88), 1047154 (817), 1233725 (963), 114599 (88), 265339 (207), 1047733 (817), and 265339 (207). Repeated query IDs are shown at different noise levels.
Figure 36: Correction refines intermediate predictions before denoising continues. Baseline and correction start from the same noise and use 15 denoising steps. Panels (a–c) show dog, cheeseburger, and macaw examples at selected timesteps . Each prediction row shows all eight correction passes at fixed ; the final row of each panel shows the actual final samples.
Figure 37: Additional examples with joint and frozen denoisers. Each row compares joint training (left triplet) and a frozen denoiser (right triplet) at , starting from the same noise. Selected JiT-L/16 samples use 100K checkpoints, EMA weights, no CFG, and 45 Heun steps; larger adds correction compute.
Figure 38: Correction reduces the high-frequency deficit of generated samples. Left: . Right: . Dashed lines indicate real-image references.
Figure 39: Uncurated samples: correction depth changes both structure and detail. Each strip follows one sample as increases, with the initial noise and denoising grid fixed.
Figure 40: Uncurated samples: correction refines unguided SiT samples across random classes. Within each pair of rows, baseline ( ) and correction ( ) share the same class and initial noise at approximately matched NFEs.
Figure 41: Uncurated samples: correction refines unguided SiT REPA samples across random classes. Within each pair of rows, baseline ( ) and correction ( ) share the same class and initial noise at approximately matched NFEs.
Figure 42: Uncurated samples: early correction changes structure as well as appearance in JiT. Within each pair of rows, baseline ( ) and correction ( ) share the same class and initial noise at matched NFEs.
Figure 43: High-quality unconditional generation without guidance. Selected ImageNet samples from SiT REPA -XL/2 with iterative correction.
Figure 44: Uncurated samples: unconditional generation supports iterative correction without text or CFG. Each row shares its initial noise across . The checkpoint and denoising steps are fixed; larger adds correction compute. All seeds are shown without selection.
Figure 45: Uncurated samples: unconditional generation supports iterative correction without text or CFG. Each row shares its initial noise across . The checkpoint and denoising steps are fixed; larger adds correction compute. All seeds are shown without selection.
Figure 46: Uncurated samples: unconditional generation supports iterative correction without text or CFG. Each row shares its initial noise across . The checkpoint and denoising steps are fixed; larger adds correction compute. All seeds are shown without selection.
Figure 47: Uncurated samples: unconditional generation supports iterative correction without text or CFG. Each row shares its initial noise across . The checkpoint and denoising steps are fixed; larger adds correction compute. All seeds are shown without selection.