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 K times before the next sampling step.
Figure 3
Figure 3: Correction uses NFEs more effectively than denoising steps. We show clean-sample predictions x^ along denoising trajectories without (a) and with (b) correction. At the same NFE budget, correction (b) produces sharper, more coherent samples than the baseline (a) with fewer steps. Each triplet in (b) shows the refined clean-sample prediction at fixed t . Additional visualizations with matched denoising step counts are shown in Appendix D.2 (Fig. 35 ).
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 K . 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
Figure 10: Correction preserves more diversity and produces less oversaturation than CFG. Samples from SiT-XL/2 with increasing correction depth K (top) or CFG scale w (bottom), 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 ( w=1 ) and CFG ( w=6 ); rows compare no correction ( K=0 ) and correction ( K=6 ). 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 ( w=6 ) versus CFG followed by correction ( w=4 , K=8 ).
Table 15
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
D28
675
Unguided
240
6,720
8.93
–
D28
675
CFG ( w=5 )
120
6,720
6.17
–
D14 + C14
681
K=1
240
6,720
8.74
–
D14 + C14
681
K=3
120
6,720
5.66
–
D14 + C14
681
K=4
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 ( D14 + C14 , K=1 ), iterative correction ( D14 + C14 , K=3 ) outperforms the monolithic baseline. Even the smaller D7+C7 pair outperforms the larger monolithic baseline with K=3 . We select the best CFG scale from w∈{1,…,10} .
Figure 15: Iterative correction sharpens predictions. We show the effect of iterative correction on JiT-L/16. (a) The base denoiser’s predictions ( K=0 ) 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. 9 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 K 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 K but deteriorates earlier.
Appendix figures & tables35 assets
Supplementary material from the paper’s appendix.
Appendix
Variant
K=0
K=1
K=2
K=3
K=4
K=5
Proposal only ( x^ )
24.63
17.96
14.07
13.35
14.27
16.06
With zt conditioning
24.72
18.22
14.84
13.79
14.67
16.58
Appendix
Table 6: Proposal-only correction performs comparably to zt -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 K=1 – 8 . 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
Appendix
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 K=1 , using FID-50K at 250 NFE. The denoiser uses the true timestep in both cases.
Stop- gradient
Ktrain
Inference-time iterations K
0
1
2
3
4
5
Correcting model predictions x^
✓
1
24.63
17.96
14.07
13.35
14.27
16.06
✗
1
29.08
17.38
12.88
19.25
24.64
28.92
✓
1 – 5
24.66
21.02
21.68
22.51
23.30
24.67
Iterating over hidden states h
Appendix
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 K=5 , 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 K , 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 rend
Start rstart
0.8
0.6
0.4
0.2
0.0
1.0
7.70
6.79
7.72
9.05
11.85 full
0.8
—
11.16
11.12
12.59
13.76
0.6
—
—
14.54
15.09
17.94
0.4
—
—
—
17.10
20.05
0.2
—
—
—
—
19.96
Appendix
Table 9: Allocating correction to early steps yields the lowest FID at a fixed sampling budget. JiT-L/16 with K=5 , 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 K=0 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 K .
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 K=0 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
256×256
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
Appendix
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; D and C denote separate denoiser and corrector networks. Sampling settings are reported separately.
Method
CFG w
K
CFG interval
Correction interval
Steps
NFE
DiT-XL/2; DDPM
Baseline
1
0
—
—
250
250
CFG
1.5
0
1→0
—
125
250
Correction
1
2
—
1→0
83
249
SiT-XL/2; Euler–Maruyama SDE
Baseline
1
0
—
—
250
250
Appendix
Table 11: The main quantitative comparisons use approximately 250 network evaluations per image. Sampling settings for Tables 1 and 12 . K is the number of corrections per active step; intervals use noise time t ( 1 : noise, 0 : data). NFE includes correction and both CFG branches. Guided T2I applies correction before CFG.
Design choice
Option
Used in
Output formulation
Direct x
Table 1 — JiT
Gated native q
Table 1 — DiT, SiT, SiT-REPA; Tables 9 – 12
Prediction type
Full target
Tables 1 – 12
Residual
Table 7 — timestep conditioning
Appendix
Table 14: Corrector parameterization choices. We distinguish the output formulation from full-target versus residual prediction, which can be used with either formulation. Here q 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.
Appendix
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 K (top) or guidance scale w (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 K=0,2,4,6 without CFG; the lower row varies guidance scale w=1,3,5,7 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 K=0,2,4,6 without CFG; the lower row varies guidance scale w=1,3,5,7 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.
CFG → Correction
Correction → CFG
w
K
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
Appendix
Table 16: Neither correction order consistently leads in both quality and diversity. Full sweep of candidate vote shares (%) against released CFG ( w=6 , K=0 ); 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 30: Quality and diversity are assessed separately. Actual teddy-bear stimuli and questions for Table 2 , using CFG followed by correction ( w=4 , K=8 ): single images for quality and four-image sets for diversity.
Figure 31: Correction improves generation quality beyond continued denoiser training or parameter-matched deeper denoisers.
Figure 32: Matched effective depth. Correction outperforms a deeper denoiser at matched effective depth per denoising step and inference budget, using fewer parameters.
Figure 33: 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 34: Shared-input noisy-state comparisons. Clean images (top) and interpolated noisy states (bottom), comparing Real , K=0 , and K=9 . Panels (a,b,g,h) use t=0.95→s=0.55 ; panels (c,d,e,f) use t=0.9→s=0.5 . 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 35: 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 t∈{0.93,0.80,0.60,0.07} . Each prediction row shows all eight correction passes at fixed t ; the final row of each panel shows the actual final samples.
Figure 36: Additional examples with joint and frozen denoisers. Each row compares joint training (left triplet) and a frozen denoiser (right triplet) at K=0,5,9 , starting from the same noise. Selected JiT-L/16 samples use 100K checkpoints, EMA weights, no CFG, and 45 Heun steps; larger K adds correction compute.
Figure 38: Uncurated samples: correction depth changes both structure and detail. Each strip follows one sample as K increases, with the initial noise and denoising grid fixed.
Figure 39: Uncurated samples: correction refines unguided SiT samples across random classes. Within each pair of rows, baseline ( K=0 ) and correction ( K=6 ) share the same class and initial noise at approximately matched NFEs.
Figure 40: Uncurated samples: correction refines unguided SiT REPA samples across random classes. Within each pair of rows, baseline ( K=0 ) and correction ( K=6 ) share the same class and initial noise at approximately matched NFEs.
Figure 41: Uncurated samples: early correction changes structure as well as appearance in JiT. Within each pair of rows, baseline ( K=0 ) and correction ( K=6 ) share the same class and initial noise at matched NFEs.
Figure 42: High-quality unconditional generation without guidance. Selected ImageNet samples from SiT REPA -XL/2 with iterative correction.
Figure 43: Uncurated samples: unconditional generation supports iterative correction without text or CFG. Each row shares its initial noise across K=0,…,5 . The checkpoint and denoising steps are fixed; larger K adds correction compute. All seeds are shown without selection.
Figure 44: Uncurated samples: unconditional generation supports iterative correction without text or CFG. Each row shares its initial noise across K=0,…,5 . The checkpoint and denoising steps are fixed; larger K 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 K=0,…,5 . The checkpoint and denoising steps are fixed; larger K 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 K=0,…,5 . The checkpoint and denoising steps are fixed; larger K adds correction compute. All seeds are shown without selection.
Steering diffusion models toward conditions unseen during training typically requires either retraining with conditional inputs or per-step gradient computations, both of which incur substantial computational overhead. We present Noise-Aligned RFM Steering (NA-RFM), a general recipe for efficiently steering diffusion models without gradient guidance during inference, enabling fast controllable generation. The method combines two offline-computed signals: noise alignment, a high-noise correction from PCA statistics of the target examples and the full data, and Recursive Feature Machine (RFM) activation steering, which learns a target-discriminative direction from labeled forward-process activations. During sampling, noise alignment provides coarse control at high noise, while the RFM direction is reused over intermediate/late timesteps through lightweight activation edits. Experiments on CIFAR-10, ImageNet, CelebA, and fine-grained bird species show improved target accuracy over gradient-based post-hoc guidance baselines, improved FID on the class-guidance benchmarks, and substantial inference speedups. Code: https://github.com/isotrivial/na-rfm.
Qingsong Wang, Mikhail Belkin, Yusu Wang
Halıcıoğlu Data Science Institute, University of California San Diego, La Jolla, CA, USA
The practical success of conditional image generation hinges on fine-grained differences in condition alignment and visual fidelity. Classifier-free guidance (CFG) is central to this success, but its lack of an explicit criterion makes it difficult to assess whether the guided trajectory is progressing as intended. To address this gap, we show that spectral alignment provides a principled criterion for understanding guidance behavior and improving guided diffusion sampling through adaptive correction. Our analysis identifies the spectra of intermediate states as an indicator of consistency with the expected spectral evolution of the forward process. Based on this observation, we introduce Spectral Correction Guidance, a method that corrects deviations from an analytic reference spectrum during sampling. The proposed method is training-free and applicable across diffusion backbones and conditional generation tasks without modifying the underlying model. Experiments demonstrate consistent gains in preference-based metrics over baseline guidance methods in text-to-image generation and improved generation quality over CFG on ImageNet. These improvements persist across a range of guidance scales and with fewer denoising steps. Our analyses and ablations provide insight into guidance behavior and how the proposed method affects generation quality.
Gihoon Kim, Taesup Kim
Graduate School of Data Science, Seoul National University
The post-training pipeline for diffusion models currently has two stages: supervised fine-tuning (SFT) on curated data and reinforcement learning (RL) with reward models. A fundamental gap separates them. SFT optimizes the denoiser only on ground-truth states sampled from the forward noising process; once inference deviates from these ideal states, subsequent denoising relies on out-of-distribution generalization rather than learned correction, exhibiting the same exposure bias that afflicts autoregressive models, but accumulated along the denoising trajectory instead of the token sequence. RL can in principle address this mismatch, yet its terminal reward signal is sparse, suffers from credit-assignment difficulty, and risks reward hacking. We propose SOAR (Self-Correction for Optimal Alignment and Refinement), a bias-correction post-training method that fills this gap. Starting from a real sample, SOAR performs a single stop-gradient rollout with the current model, re-noises the resulting off-trajectory state, and supervises the model to steer back toward the original clean target. The method is on-policy, reward-free, and provides dense per-timestep supervision with no credit-assignment problem. On SD3.5-Medium, SOAR improves GenEval from 0.70 to 0.78 and OCR from 0.64 to 0.67 over SFT, while simultaneously raising all model-based preference scores. In controlled reward-specific experiments, SOAR surpasses Flow-GRPO in final metric value on both aesthetic and text-image alignment tasks, despite having no access to a reward model. Since SOAR's base loss subsumes the standard SFT objective, it can directly replace SFT as a stronger first post-training stage after pretraining, while remaining fully compatible with subsequent RL alignment.