扩散模型的纠正训练
摘要
论文提出一种针对扩散模型的“修正训练”框架,通过训练一个以基模型自身预测为输入的修正器来缩小训练-推理差距,显著提升无引导/无条件生成的样本质量,并可通过蒸馏将收益无损保留到基模型中。
查看缓存全文
缓存时间: 2026/10/15 06:45
# Correction Training for Diffusion Models
Source: [https://arxiv.org/html/2610.14136](https://arxiv.org/html/2610.14136)
Runqian WangAffiliation:University of California, BerkeleyAlexei A\. EfrosAffiliation:University of California, BerkeleyTaesung ParkAffiliation:ReveYash BelheAffiliation:Reve
###### Abstract
Iterative sampling in diffusion models during inference can lead to error accumulation due to exposure bias\. The denoiser is trained on real training images interpolated with noise, but at inference its inputs areconstructed fromits 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\.In this paper, we introduce a correction step for diffusion models which attempts to reduce this training\-inference gap\. The corrector takes the base model’s own prediction as input and is trained against the same ground\-truth target as the base model\. 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\. We find that applying correction beyond the single step supervised during training is key to the significant quality gains\. This simple framework leads to substantial improvement in sample quality while largely preserving diversity\. We extensively evaluate our method across widely used architectures, including DiT, SiT, and JiT\. Our method substantially improves upon previously reported FID scores forunguided generationandunconditional generation\. 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\.
Figure 1:ImageNet samples with correction\.Without the use of guidance, our correction method can achieve high\-quality class\-conditional \(top\) andunconditional\(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 refinedKKtimes before the next sampling step\.## 1Introduction
Diffusion models generate images through a sequence of denoising predictions\([Ho et al\., 2020](https://arxiv.org/html/2610.14136#bib.bib9);[Song et al\., 2020](https://arxiv.org/html/2610.14136#bib.bib33)\)\. At each sampling step, a neural network estimates clean data from a noisy input, and this estimate is used to update the sample\. But one curious property of diffusion models is the training\-inference gap arising from what is known as exposure bias\([Ning et al\., 2023](https://arxiv.org/html/2610.14136#bib.bib23);[Ning et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib24);[Li et al\., 2024a](https://arxiv.org/html/2610.14136#bib.bib16);[Yu & Zhan, 2025](https://arxiv.org/html/2610.14136#bib.bib39)\)\. The denoiser is trained on real training images interpolated with noise\. At inference time, however, it receives noisy inputs that depend onits own imperfect predictionsfrom earlier steps\. These inputs can differ from the noisy training inputs, creating a mismatch between training and inference\. As noise decreases along the denoising trajectory, prediction errors can accumulate, often producing visual artifacts and incoherent structures in the generated images\. A common way to improve sample quality is classifier\-free guidance \(CFG\)\([Ho & Salimans, 2022](https://arxiv.org/html/2610.14136#bib.bib8)\), which generates better quality samples, but at the cost of reduced diversity and oversaturation\.
In this paper, we take a more direct approach: correcting themodel’s own prediction\. We introduce acorrectorthat takes the base model’s prediction as input and is trained against the same ground\-truth target, as illustrated in Fig\.[2](https://arxiv.org/html/2610.14136#S0.F2)\. This makes correction an explicit learning task, with training inputs provided by the base model\. During sampling, at each denoising step, we therefore iteratively correct the denoiser’s clean\-sample prediction\. Although the corrector is trained only on predictions from the base denoiser, we find that it can be applied repeatedly to its own outputs at inference time without additional training\. This iterative correction further improves sample quality, producing sharper images with more coherent structures\.We hypothesize that repeated correction helps compensate for errors accumulated over earlier denoising steps, reducing the denoiser’s training–inference gap and contributing to the observed improvement \(see Sec\.[4\.2](https://arxiv.org/html/2610.14136#S4.SS2)for details\)\.
We evaluate our method across class\-conditional, unconditional, and text\-to\-image generation, using established models including DiT\([Peebles & Xie, 2023](https://arxiv.org/html/2610.14136#bib.bib25)\), SiT\([Ma et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib22)\), REPA\([Yu et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib40)\), and JiT\([Li & He, 2026](https://arxiv.org/html/2610.14136#bib.bib17)\)\. Correction improves unguided sample quality across these settings with a relatively small amount of additional training\. The gain remains at matched numbers of function evaluations \(NFEs\), suggesting that introducing correction steps can make more effective use of inference compute than denoising alone\. In unconditional generation, correction also provides substantial improvements without any guidance signal\. Comparisons with CFG on DiT and SiT show better quality at matched diversity, as measured by FID\([Heusel et al\., 2017](https://arxiv.org/html/2610.14136#bib.bib7)\)and recall\([Kynkäänniemi et al\., 2019](https://arxiv.org/html/2610.14136#bib.bib13)\)\. Qualitatively, correction shows less oversaturation and greater visual realism than CFG\. Moreover, adding correction to CFG reduces the stereotypical “AI\-generated” look of its samples\. We present conditional and unconditional samples using correction in Fig\.[1](https://arxiv.org/html/2610.14136#S0.F1)\.
Finally, we show that the corrector can be made as efficient as standard diffusion models while retaining its quality advantage\. We distill corrected predictions back into the base diffusion model\. The student learns to match predictions obtained through repeated correction and requires only one forward pass per sampling step\. The resulting model retains much of the quality improvement while preserving the base model’s architecture, parameter count, and inference cost\. This allows the benefits of correction to be incorporated into the base model without additional correction passes during generation\. We hope that the effectiveness, simplicity, and efficiency of corrector can help to bridge the long\-standing gap between unconditional and conditional generation\.
## 2Learning to Correct Denoising Predictions
### 2\.1Background: Denoising and Sampling
Diffusion models learn to recover data from noisy observations\([Sohl\-Dickstein et al\., 2015](https://arxiv.org/html/2610.14136#bib.bib30);[Song & Ermon, 2019](https://arxiv.org/html/2610.14136#bib.bib32);[Ho et al\., 2020](https://arxiv.org/html/2610.14136#bib.bib9);[Song et al\., 2020](https://arxiv.org/html/2610.14136#bib.bib33)\)\. Flow matching\([Lipman et al\., 2022](https://arxiv.org/html/2610.14136#bib.bib19);[Liu et al\., 2022](https://arxiv.org/html/2610.14136#bib.bib20);[Albergo et al\., 2025](https://arxiv.org/html/2610.14136#bib.bib2)\)extends the diffusion paradigm by using a linear interpolation between noise and image\. Specifically, given a clean sample𝒙∼pdata\{\\bm\{x\}\}\\sim p\_\{\\mathrm\{data\}\}and independent noiseϵ∼𝒩\(0,I\)\\bm\{\\epsilon\}\\sim\\mathcal\{N\}\(0,I\), the training input is
𝒛t=\(1−t\)𝒙\+tϵ,t∈\[0,1\],\\color\[rgb\]\{0,0,0\}\{\\bm\{z\}\_\{t\}\}=\(1\-t\)\{\\bm\{x\}\}\+t\\bm\{\\epsilon\},\\qquad t\\in\[0,1\],\(1\)wheret=0t=0denotes data andt=1t=1denotes noise\. An𝒙\\bm\{x\}\-prediction networkDθD\_\{\\theta\}estimates the clean sample from\(𝒛t,t\)\(\{\\bm\{z\}\_\{t\}\},t\)\. Its denoising objective is
ℒdenoise=𝔼𝒙,ϵ,t\[∥Dθ\(𝒛t,t\)−𝒙∥22\]\.\\color\[rgb\]\{0,0,0\}\\mathcal\{L\}\_\{\\mathrm\{denoise\}\}=\\mathbb\{E\}\_\{\{\\bm\{x\}\},\\bm\{\\epsilon\},t\}\\\!\\left\[\\lVert D\_\{\\theta\}\(\{\\bm\{z\}\_\{t\}\},t\)\-\{\\bm\{x\}\}\\rVert\_\{2\}^\{2\}\\right\]\.\(2\)At each sampling step, the denoiser predicts the clean image, and the sampler uses this prediction to construct the next noisy state\. For example, an Euler step fromtttos<ts<tgives
𝒛s=st𝒛t\+\(1−st\)𝒙^\(0\),𝒙^\(0\)=Dθ\(𝒛t,t\)\.\\bm\{z\}\_\{s\}=\\frac\{s\}\{t\}\\bm\{z\}\_\{t\}\+\\left\(1\-\\frac\{s\}\{t\}\\right\)\\hat\{\\bm\{x\}\}^\{\(0\)\},\\qquad\\hat\{\\bm\{x\}\}^\{\(0\)\}=D\_\{\\theta\}\(\\bm\{z\}\_\{t\},t\)\.\(3\)The updated state𝒛s\\bm\{z\}\_\{s\}then becomes the denoiser’s next input\. During training, however, inputs are constructed by interpolating real images with independent noise \(Eq\.[1](https://arxiv.org/html/2610.14136#S2.E1)\)\. Sampling inputs therefore depend on the model’s own earlier predictions and may differ from the inputs encountered during training, creating a training–inference gap\([Ning et al\., 2023](https://arxiv.org/html/2610.14136#bib.bib23);[Ning et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib24);[Li et al\., 2024a](https://arxiv.org/html/2610.14136#bib.bib16)\)\.
Algorithm 1Training loss\# D: denoiser; C: corrector\# weight: correction loss weightt=sample\_time\(x\)\# per imagee=randn\_like\(x\)z=\(1\-t\)\*x\+t\*epred=D\(z,t\)refined=C\(stopgrad\(pred\),t\)loss=mse\(pred,x\)loss\+=weight\*mse\(refined,x\)
Algorithm 2Sampling with correction\# D: denoiser; C: corrector\# K: corrections per stepz=randn\(shape\)fort,sinzip\(times,times\[1:\]\):pred=D\(z,t\)\# correction stepsfor\_inrange\(K\):pred=C\(pred,t\)v=\(z\-pred\)/tz=z\+\(s\-t\)\*vsample=z
### 2\.2Prediction Error During Sampling
Fig\.[3](https://arxiv.org/html/2610.14136#S2.F3)\(a\) illustrates how structural errors can persist along the sampling trajectory\. Early predictions contain incoherent object structures\. Later denoising steps add texture and facial details such as teeth, eyes, and a tongue while largely preserving the incorrect object structure\. The clean targets used for training, however, are real images with coherent object structures\. The denoiser learns to recover these structures from noise, but is not explicitly trained to further refine its own predictions at the same noise level\.
One widely used way to obtain sharper, more coherent samples is classifier\-free guidance \(CFG\)\([Ho & Salimans, 2022](https://arxiv.org/html/2610.14136#bib.bib8)\)\. CFG strengthens conditioning by amplifying the difference between conditional and unconditional predictions\. However, it can reduce sample diversity and produce oversaturated images\.
In this paper, we make prediction correction an explicit learning task: we train a correctorCϕC\_\{\\phi\}that takes the denoiser’s own prediction as input\. During sampling, we apply this learned mapping repeatedly to improve sample quality\. As we show in Sec\.[4](https://arxiv.org/html/2610.14136#S4), repeated correction can improve generation even when it increases reconstruction error on training inputs\. We investigate this behavior and its relationship to the training–inference gap\.
### 2\.3Learning to Correct Model Predictions
Standard denoising training learns the mapping𝒛t→𝒙\\bm\{z\}\_\{t\}\\rightarrow\\bm\{x\}by reconstructing images corrupted by a prescribed noise process\. Given a noisy observation𝒛t\\bm\{z\}\_\{t\}, the denoiser produces an initial prediction𝒙^\(0\)\\hat\{\\bm\{x\}\}^\{\(0\)\}, which may still contain inaccuracies\. We introduce*prediction correction*as a reconstruction task that uses this prediction as its input\. The corrector takes𝒙^\(0\)\\hat\{\\bm\{x\}\}^\{\(0\)\}and predicts the same clean target𝒙\\bm\{x\}, learning the mapping𝒙^\(0\)→𝒙\\hat\{\\bm\{x\}\}^\{\(0\)\}\\rightarrow\\bm\{x\}\(Fig\.[2](https://arxiv.org/html/2610.14136#S0.F2)\)\.
When𝒙^\(0\)\\hat\{\\bm\{x\}\}^\{\(0\)\}already captures much of the structure in𝒙\\bm\{x\}, the initial denoising pass has already performed much of the reconstruction\. The corrector can then build on this estimate and focus on the remaining inaccuracies\. We refine𝒙^\(0\)\\hat\{\\bm\{x\}\}^\{\(0\)\}by applying correctorCϕC\_\{\\phi\}again at the original timett:𝒙^\(1\)=Cϕ\(𝒙^\(0\),t\)\.\\hat\{\\bm\{x\}\}^\{\(1\)\}=C\_\{\\phi\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\},t\)\.The corrected prediction directly estimates the clean sample𝒙\\bm\{x\}\. We train it with
ℒcorr=𝔼𝒙,ϵ,t\[∥Cϕ\(sg\(𝒙^\(0\)\),t\)−𝒙∥22\],\\mathcal\{L\}\_\{\\mathrm\{corr\}\}=\\mathbb\{E\}\_\{\\bm\{x\},\\bm\{\\epsilon\},t\}\\\!\\left\[\\lVert C\_\{\\phi\}\(\\operatorname\{sg\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\}\),t\)\-\\bm\{x\}\\rVert\_\{2\}^\{2\}\\right\],\(4\)wheresg\\operatorname\{sg\}is the stop\-gradient operation\.
Starting from a pretrained denoiser, we jointly optimizeℒtotal=ℒdenoise\+λcorrℒcorr,\\mathcal\{L\}\_\{\\mathrm\{total\}\}=\\mathcal\{L\}\_\{\\mathrm\{denoise\}\}\+\\lambda\_\{\\mathrm\{corr\}\}\\mathcal\{L\}\_\{\\mathrm\{corr\}\},whereλcorr\\lambda\_\{\\mathrm\{corr\}\}weights the correction task\. The denoising term supervises the initial prediction, while the correction term supervises its refinement\. The two tasks are trained jointly, with corrector inputs generated online by the current model\. Algorithm[1](https://arxiv.org/html/2610.14136#alg1)summarizes the loss computation in Python\-style pseudocode\.
Figure 3:Prediction correction improves structural coherence during sampling\.We show clean\-sample predictions𝒙^\\hat\{\\bm\{x\}\}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 eachtt, 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](https://arxiv.org/html/2610.14136#A4.SS2.SSS0.Px4)\(Fig\.[36](https://arxiv.org/html/2610.14136#A4.F36)\)\.
### 2\.4Iterative Correction at Inference
Training supervises only a single correction step, from the initial denoising prediction𝒙^\(0\)\\hat\{\\bm\{x\}\}^\{\(0\)\}to the refined prediction𝒙^\(1\)\\hat\{\\bm\{x\}\}^\{\(1\)\}\. At inference, however, the model can repeatedly refine its own outputs, further improving sample quality\. The output𝒙^\(k\)\\hat\{\\bm\{x\}\}^\{\(k\)\}at thekk\-th correction iteration is another estimate of the same clean image\. This estimate may still contain inaccuracies, motivating further applications of the learned correction\.
Starting from the denoising prediction, we applyKKcorrection passes:
𝒙^\(0\)\\displaystyle\\hat\{\\bm\{x\}\}^\{\(0\)\}=Dθ\(𝒛t,t\),\\displaystyle=D\_\{\\theta\}\(\\bm\{z\}\_\{t\},t\),\(5\)𝒙^\(k\+1\)\\displaystyle\\hat\{\\bm\{x\}\}^\{\(k\+1\)\}=Cϕ\(𝒙^\(k\),t\),k=0,…,K−1\.\\displaystyle=C\_\{\\phi\}\(\\hat\{\\bm\{x\}\}^\{\(k\)\},t\),\\qquad k=0,\\ldots,K\-1\.Fig\.[3](https://arxiv.org/html/2610.14136#S2.F3)illustrates how successive corrections refine intermediate predictions throughout the sampling trajectory\. This provides a way to spend additional inference compute on improving the predictions at each time\. At matched total inference cost, we find this more effective than simply increasing the number of denoising steps\.
## 3Results
We evaluate correction on class\-conditional and unconditional generation \(Sec\.[3\.1](https://arxiv.org/html/2610.14136#S3.SS1)\) and text\-to\-image generation \(Sec\.[3\.2](https://arxiv.org/html/2610.14136#S3.SS2)\)\. We also demonstrate that the base and correction models can share most parameters and that additional correction passes can be eliminated with distillation \(Sec\.[4](https://arxiv.org/html/2610.14136#S3.T4)\)\. We detail our experimental settings in Appendix[B](https://arxiv.org/html/2610.14136#A2)\.
### 3\.1ImageNet Generation
We evaluate our method on widely used transformer\-based architectures, including DiT\([Peebles & Xie, 2023](https://arxiv.org/html/2610.14136#bib.bib25)\), SiT\([Ma et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib22)\), and JiT\([Li & He, 2026](https://arxiv.org/html/2610.14136#bib.bib17)\)\. We initialize both the base denoiser and a separate corrector from publicly released checkpoints and jointly fine\-tune them following Algorithm[1](https://arxiv.org/html/2610.14136#alg1)\. Unless otherwise specified, we follow each baseline’s training recipe and evaluation protocol\. Number of function evaluations \(NFE\) under correction measures both denoiser steps and correction steps\. We approximately match the total NFE by reducing the number of denoising steps as the number of correction iterationsKKincreases\.
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\.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](https://arxiv.org/html/2610.14136#A2.T13)for evaluation details\.
Figure 5:Efficiency of correction training\.FID converges within relatively few corrector training steps\.Class\-conditional generation\.We present class conditional image samples in Fig\.[4](https://arxiv.org/html/2610.14136#S3.F4)\. Correction produces sharp, realistic, and diverse images, allwithout applying any guidance\. The resulting samples exhibit realism competitive with that achieved by guidance\-based methods such as classifier\-free guidance\([Ho & Salimans, 2022](https://arxiv.org/html/2610.14136#bib.bib8)\), perturbation guidance\([Ahn et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib1)\)and autoguidance\([Karras et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib11)\)\. See Appendix[E\.1](https://arxiv.org/html/2610.14136#A5.SS1)for more samples\.
In Tab\.[1](https://arxiv.org/html/2610.14136#S3.T1), we evaluate generation using FID, precision, and recall\([Heusel et al\., 2017](https://arxiv.org/html/2610.14136#bib.bib7);[Kynkäänniemi et al\., 2019](https://arxiv.org/html/2610.14136#bib.bib13)\)\. Correctionsignificantly reduces FID at matched NFE\(DiT: 9\.58→\\to2\.07, SiT: 8\.55→\\to1\.92, and SiTREPA: 5\.53→\\to1\.31\.\)\. This suggests that correction is a better allocation of inference compute than simply increasing denoising steps\. These gains require little additional training: for SiTREPA, 50K updates amount to just 1\.25% of its 4M pretraining updates\. Fig\.[5](https://arxiv.org/html/2610.14136#S3.F5)compares FID convergence with correction versus continued denoising\. Fig\.[6](https://arxiv.org/html/2610.14136#S3.F6)shows how the number of correction iterationsKKaffects the final samples and their fidelity–diversity trade\-off\. Correction effectively suppresses artifacts and improves object coherence, yielding lower FID and higher precision at the cost of some recall\. We further analyze the power spectra of the generated samples in Appendix[D\.3](https://arxiv.org/html/2610.14136#A4.SS3)\.
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 acrossKK\. Correction gives a controllable tradeoff between coverage and quality and lower minimum FID\.##### Unconditional generation\.
Unconditional generation remains a challenging task since no class or text conditioning is available to constrain the outputs\([Li et al\., 2024b](https://arxiv.org/html/2610.14136#bib.bib18)\)\. We show that for unconditional generation, our prediction correction significantly improves image quality and coherence, all without requiring guidance\. We present qualitative comparisons of corrected and uncorrected samples in Fig\.[7](https://arxiv.org/html/2610.14136#S3.F7)\. Correction removes structural inconsistencies and artifacts from the baseline samples while maintaining the main subject\.
Quantitatively, we compare correction and autoguidance\([Karras et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib11)\)in Tab\.[9](https://arxiv.org/html/2610.14136#S3.F9)on ImageNet unconditional generation\. Correction withK=4K=4substantially improves FID over the baseline and the best evaluated autoguidance scale \(w∈\{1,2,…,10\}w\\in\\\{1,2,\\dots,10\\\}, with the released pretrained model as the guide\), while largerKKtrades recall for higher precision \(Fig\.[9](https://arxiv.org/html/2610.14136#S3.F9)\)\.
Figure 7:Correction improvesunconditionalsample 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\.![[Uncaptioned image]](https://arxiv.org/html/2610.14136v2/unconditional_matched50k_preview.png)
Figure 8:Correction lowers unconditional FID without guidance\.50K training updates, 250 NFE, autoguidancew=7w=7\.Figure 9:Correction reaches lower FID and higher peak precision than autoguidance\.Dashed: best autoguidance FID\. Dotted: each method’s observed peak precision\.
##### Diversity preservation\.
Although comparison with CFG is not our main goal, our method addresses two key limitations of CFG: reduced diversity and oversaturation\. Fig\.[10](https://arxiv.org/html/2610.14136#S3.F10)shows samples from the same ImageNet class as we vary the CFG scalewwand the number of correction iterationsKK\. At higher CFG scales, the samples become less diverse, converging on canonical close\-up portraits, while their colors become increasingly saturated\. With correction, the dogs appear in more varied poses, framings, and backgrounds while retaining sharp details and realism\. Additional class\-wise comparisons are shown in Figs\.[28](https://arxiv.org/html/2610.14136#A3.F28)and[29](https://arxiv.org/html/2610.14136#A3.F29)in the appendix\.
Figure 10:Correction preserves more diversity and produces less oversaturation than CFG\.Samples from SiTREPA\-XL/2 with increasing correction depthKK\(bottom\) or CFG scaleww\(top\), at approximately matched NFE, same initial noise\.
### 3\.2Text\-to\-Image Generation
To assess our method on a more complex, practically relevant task, we evaluate text\-to\-image generation using NanoGen\([Leng et al\., 2026](https://arxiv.org/html/2610.14136#bib.bib15)\), a unified training framework for diffusion transformers\. We use a pretrained DiT\([Peebles & Xie, 2023](https://arxiv.org/html/2610.14136#bib.bib25)\)with a DDT head\([Wang et al\., 2026](https://arxiv.org/html/2610.14136#bib.bib35)\), operating in the latent space of the FLUX\.2 VAE\([Labs, 2025](https://arxiv.org/html/2610.14136#bib.bib14)\)\. We use the JourneyDB\([Sun et al\., 2023](https://arxiv.org/html/2610.14136#bib.bib34)\), Short\-Caption, and Long\-Caption subsets of the BLIP\-3o pretraining mixture\([Chen et al\., 2025](https://arxiv.org/html/2610.14136#bib.bib3)\), with aspect\-preserving preprocessing to256×256256\\times 256\. For the quantitative comparison, we jointly fine\-tune the denoiser and corrector for 50K steps, adding a relatively small training cost compared with pretraining the base model\.

Figure 11:Correction improves object structure intext\-to\-imagegeneration\.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](https://arxiv.org/html/2610.14136#A2.SS3.SSS0.Px2)\.Without CFG
With CFG
Figure 12:Correction improves text\-to\-image alignment with and without CFG\.We observe consistent improvements on GenEval and DPG\-Bench by applying correction over the baseline, both with and without CFG, evaluated at approximately matched NFE\. See Table[13](https://arxiv.org/html/2610.14136#A2.T13)for detailed settings\.Fig\.[12](https://arxiv.org/html/2610.14136#S3.F12)illustrates changes in object structure without guidance: correction separates the violin bodies, clarifies the guitar and car geometry, and organizes cluttered interiors\. Quantitatively, Tab\.[12](https://arxiv.org/html/2610.14136#S3.F12)shows higher GenEval\([Ghosh et al\., 2023](https://arxiv.org/html/2610.14136#bib.bib5)\)and DPG\-Bench\([Hu et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib10)\)scores with correction at approximately matched NFE, first without CFG and then with correction applied before CFG\. Both comparisons use the same\-checkpoint baseline and DiffusionBench’s native evaluation pipeline; sampling protocols are provided in Table[13](https://arxiv.org/html/2610.14136#A2.T13), and additional guidance scales in Appendix[B\.3](https://arxiv.org/html/2610.14136#A2.SS3)\.
Fig\.[13](https://arxiv.org/html/2610.14136#S3.F13)compares different combinations of CFG and correction\. Stronger CFG alone can improve conditioning adherence and sample quality, but often at the cost of reduced diversity and oversaturated or stylized appearances\. Lowering the guidance scale alleviates these effects, but may leave samples under\-refined\. We find thatapplying correction jointly with CFG alleviates the diversity issuewhile preserving the quality refinement\. We observe that CFG alone produces visually similar samples, while combining CFG and correction leads to higher diversity \(Fig\.[13](https://arxiv.org/html/2610.14136#S3.F13)\)\.

Figure 13:Correction improves the diversity of CFG samples while preserving quality\.Columns compare no guidance \(w=1w=1\) and CFG \(w=6w=6\); rows compare no correction \(K=0K=0\) and correction \(K=6K=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”\.##### Unconditional generation on web images\.
Using the same model, we also presentunconditionalgeneration results in Appendix[E\.6](https://arxiv.org/html/2610.14136#A5.SS6)\. Surprisingly, correction improves realism while preserving both the diversity of web images and the overall structure of the corresponding uncorrected samples\.
Table 2:Correction is preferred for quality and diversity\.Released CFG \(w=6w=6\) versus CFG followed by correction \(w=4w=4,K=8K=8\)\.Human evaluation\.To assess sample quality and diversity through human judgments, we conduct a preference study \(Tab\.[2](https://arxiv.org/html/2610.14136#S3.T2)\) using 80 GenEval single\-object prompts on the Rapidata platform111[https://www\.rapidata\.ai/](https://www.rapidata.ai/)\. Raters answer “Which image has better visual quality?” for individual images and “Which group is more diverse? Ignore image quality\.” for four\-image groups\. Raters prefer CFG\+\+Correction over the baseline using DiffusionBench’s default CFG scale\([Leng et al\., 2026](https://arxiv.org/html/2610.14136#bib.bib15)\)in more than 60% of responses\. Evaluation details are provided in Appendix[C\.3](https://arxiv.org/html/2610.14136#A3.SS3)\.
### 3\.3Making correction efficient
##### Parameter\-efficient correction through weight sharing\.
The denoiser and corrector can share a backbone while retaining separate lightweight input embedders, which contribute only a small fraction of the total parameters\. On SiT\-XL/2, this configuration achieves performance comparable to separate networks with roughly half the total parameters \(Tab\.[4](https://arxiv.org/html/2610.14136#S3.T4)\)\. In the shared configuration, the backbone receives gradients from both the denoising and correction losses\.
We also compare against a larger single denoiser with roughly matched parameter count, and find that iterative correction yields substantially larger gains than simply doubling the denoiser depth \(Sec\.[4\.1](https://arxiv.org/html/2610.14136#S4.SS1), Tab\.[5](https://arxiv.org/html/2610.14136#S4.T5)\)\.
Table 3:A shared backbone retains similar correction quality\.Using shared layers from SiT\-XL/2 on ImageNet with 100K steps leads to similar FID, suggesting that correction’s advantage does not primarily come from using more parameters\.
Table 4:Distillation achieves similar FID under much smaller NFE\.Distilling a SiT\-XL/2 correction model into the base model leads to a similar FID performance with 3 times less NFE on ImageNet\.
##### Distilling correction into the base model\.
Weight sharing reduces parameter overhead, but each correction iteration still requires an additional forward pass\. We therefore distill corrected predictions into the original denoiser\. The student is initialized from the teacher’s base denoiser and trained to match its predictions afterk∗k^\{\*\}correction iterations:
ℒdistill=wdistill\(t\)∥Dθ′\(𝒛t,t\)−x^\(k∗\)∥22,\\mathcal\{L\}\_\{\\mathrm\{distill\}\}=w\_\{\\mathrm\{distill\}\}\(t\)\\lVert D\_\{\\theta^\{\\prime\}\}\(\\bm\{z\}\_\{t\},t\)\-\\hat\{x\}^\{\(k^\{\*\}\)\}\\rVert\_\{2\}^\{2\},\(6\)wherek∗k^\{\*\}controls the number of teacher correction iterations\. Tab\.[4](https://arxiv.org/html/2610.14136#S3.T4)reports results for class\-conditional ImageNet\-256 generation using SiT\-XL/2 withk∗=2k^\{\*\}=2\. The student achieves comparable FID while reducing NFE by approximately 67%\. We usewdistill\(t\)=1/t2w\_\{\\mathrm\{distill\}\}\(t\)=1/t^\{2\}fort\>0t\>0, corresponding to velocity\-space MSE\.
## 4Discussion
In this section, we investigate where the quality gains from correction might come from\. We first compare a denoiser–corrector pair with a single deeper denoiser containing the same total number of transformer layers \(Sec\.[4\.1](https://arxiv.org/html/2610.14136#S4.SS1)\)\. We then empirically examine the behavior of iterative correction and its relationship to the training–inference gap \(Sec\.[4\.2](https://arxiv.org/html/2610.14136#S4.SS2)\)\. Finally, we examine how the number of corrections seen during training and denoiser updates affect the quality gains from iterative correction \(Sec\.[4\.3](https://arxiv.org/html/2610.14136#S4.SS3)\)\.
### 4\.1Parameter\-matched baselines
##### A single correction achieves similar FID to a monolithic denoiser\.
Figure 14:Matched parameters\.Repeated correction outperforms a deeper denoiser at a comparable parameter count and matched inference budget\.A natural question is whether a denoiser–corrector pair with a single correction \(K=1K=1\) can outperform a monolithic denoiser with the same total number of transformer layers \(Fig\.[14](https://arxiv.org/html/2610.14136#S4.F14), first and second rows\)\. We also use the same number of denoising steps for both configurations\. More specifically, we compare a 28\-layer denoiser \(D28D\_\{28\}\) with a 14\-layer denoiser and a 14\-layer corrector \(D14\+C14D\_\{14\}\+C\_\{14\}\), both trained from scratch for 80K updates using LightningDiT\([Yao et al\., 2025](https://arxiv.org/html/2610.14136#bib.bib37)\)on ImageNet 256×\\times256\([Russakovsky et al\., 2015](https://arxiv.org/html/2610.14136#bib.bib26)\)\. The two configurations have approximately the same parameter count \(675M and 681M, respectively\)\. With a single correction at each denoising step \(K=1K=1\), the denoiser–corrector pair achieves similar FID to the monolithic denoiser \(8\.74 versus 8\.93\)\. Thus, splitting the model into a denoiser and a corrector gives little quality gain when trained from scratch\.
##### Iterative correction outperforms a monolithic denoiser\.
We then reuse the same corrector forK=3K=3iterations at each denoising step \(Fig\.[14](https://arxiv.org/html/2610.14136#S4.F14), third row\)\. This lowers FID to 5\.66, while halving the number of denoising steps keeps the total transformer\-layer evaluation budget unchanged \(Tab\.[5](https://arxiv.org/html/2610.14136#S4.T5)\)\. These results suggest thatiterativecorrection makes effective use of inference compute\.
Even when we match the effective depth within each denoising step, the denoiser–corrector pair still performs better\. The shallowD7\+C7D\_\{7\}\+C\_\{7\}pair withK=3K=3achieves lower FID thanD28D\_\{28\}\(4\.21 versus 4\.88\), using the same number of denoising steps and roughly half as many parameters \(Tab\.[5](https://arxiv.org/html/2610.14136#S4.T5)\. See Fig\.[33](https://arxiv.org/html/2610.14136#A4.F33)in Appendix\)\. These results are from training from scratch\. When starting from a pretrained model, as in our main experiments in Sec\.[3](https://arxiv.org/html/2610.14136#S3), a single correction can outperform a deeper monolithic denoiser and converges much faster \(Fig\.[32](https://arxiv.org/html/2610.14136#A4.F32)in the Appendix\)\.
ModelParams \(M\)SamplingStepsLayer evalsFID@8K↓\\downarrowFID@50K↓\\downarrowMatched parametersD28D\_\{28\}675Unguided2406,7208\.93–D28D\_\{28\}675CFG \(w=5w=5\)1206,7206\.17–D14D\_\{14\}\+C14C\_\{14\}681K=1K=12406,7208\.74–D14D\_\{14\}\+C14C\_\{14\}681K=3K=31206,7205\.66–D14D\_\{14\}\+C14C\_\{14\}681K=4K=4966,7205\.56–Matched effective depthD28D\_\{28\}675Unguided2406,7208\.194\.88D7D\_\{7\}\+C7C\_\{7\}346K=3K=32406,7207\.514\.21Table 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 \(D14D\_\{14\}\+C14C\_\{14\},K=1K=1\), iterative correction \(D14D\_\{14\}\+C14C\_\{14\},K=3K=3\) outperforms the monolithic baseline\. Even the smallerD7\+C7D\_\{7\}\+C\_\{7\}pair outperforms the larger monolithic baseline withK=3K=3\. We select the best CFG scale fromw∈\{1,…,10\}w\\in\\\{1,\\ldots,10\\\}\.
### 4\.2How iterative correction changes predictions and sampling
Figure 15:Iterative correction sharpens predictions\.We show the effect of iterative correction on JiT\-L/16\.\(a\)The base denoiser’s predictions \(K=0K=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](https://arxiv.org/html/2610.14136#A4.SS2)\. Analytic target is computed using Eq\.[10](https://arxiv.org/html/2610.14136#A4.E10)in Appendix\.##### Iterative correction sharpens the denoiser’s predictions\.
The corrector is trained to reconstruct the clean target from the denoiser’s prediction\. Its output is another clean\-image prediction, allowing the same mapping to be applied repeatedly at a fixed noise level\. We examine how this repetition changes the predictions\. Fig\.[15](https://arxiv.org/html/2610.14136#S4.F15)shows this visually\. We use noisy data𝒛t=\(1−t\)𝒙\+tϵ\\bm\{z\}\_\{t\}=\(1\-t\)\\bm\{x\}\+t\\bm\{\\epsilon\}as input to JiT\-L/16 and show its predictions before and after correction\. As the correction countKKincreases, object structures become increasingly coherent, producing realistic\-looking, globally coherent images atK=9K=9\. For example, in \(a\), the blurry, “average\-like” predictions of dogs, pandas, and cars become distinct, well\-formed objects\. In \(b\), the corrector sometimes makes overly coherent structures, such as disk\-like pizza toppings or strawberry seeds arranged in a regular grid\.
In our from\-scratch experiments, a single correction yields similar FID to a deeper, parameter\-matched denoiser \(Sec\.[4\.1](https://arxiv.org/html/2610.14136#S4.SS1)\)\. However, repeated application of the same corrector substantially improves sample quality at a matched inference budget\. Our intuition is that splitting the computation into denoising and correction exposes a refinement operation that can be reused\. The denoiser first maps the noisy input into a clean\-image prediction, and the corrector learns to refine this prediction toward the original clean target\. A single correction may provide refinement similar to that achieved by additional denoiser layers\. Repeated correction, however, reapplies this learned mapping to its own outputs, beyond the single iteration used during training\. This may reinforce coherent structures, sometimes excessively, producing overly regular patterns such as the disks and grids\.
Figure 16:Iterative correction can increase denoising error\.
##### Does iterative correction reduce denoising error?
Iterative correction sharpens predictions, but the corrected outputs can differ substantially from the original clean image \(Fig\.[15](https://arxiv.org/html/2610.14136#S4.F15)\)\. Thus, we analyze whether these visual improvements also reduce denoising error\. We denote the denoising error after correction by
ℒdenoise\(k\)=𝔼𝒙,ϵ,t\[∥𝒙^\(k\)−𝒙∥22\],\\mathcal\{L\}\_\{\\mathrm\{denoise\}\}^\{\(k\)\}=\\mathbb\{E\}\_\{\\bm\{x\},\\bm\{\\epsilon\},t\}\\\!\\left\[\\lVert\\hat\{\\bm\{x\}\}^\{\(k\)\}\-\\bm\{x\}\\rVert\_\{2\}^\{2\}\\right\],\(7\)wherekkis the number of correction iterations\.
We compare this error with the base denoiser’s error \(Eq\.[2](https://arxiv.org/html/2610.14136#S2.E2)\) by measuringℒdenoise\(k\)−ℒdenoise\\mathcal\{L\}\_\{\\mathrm\{denoise\}\}^\{\(k\)\}\-\\mathcal\{L\}\_\{\\mathrm\{denoise\}\}\. Fig\.[16](https://arxiv.org/html/2610.14136#S4.F16)shows positive differences fork≥3k\\geq 3, meaning that the denoising errorincreasesafter correction\. Thus, even as predictions become sharper and more coherent, iterative correction can move them farther on average from the optimal denoising prediction\.
##### Iterative correction reduces the training–inference gap\.
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\.How can iterative correction improve sample quality even when it increases denoising error? Denoising error is measured on noisy inputs drawn from thetraining distribution, whereas intermediate states encountered during sampling can follow a different distribution\([Song et al\., 2025](https://arxiv.org/html/2610.14136#bib.bib31)\)\. Each corrected prediction is used to construct the next noisy state, which becomes the denoiser’s input\. Thus, we hypothesize that correction may improve generation by bringing these states closer to the training distribution\. We measure the training–inference gap using a classifier two\-sample test\([Lopez\-Paz & Oquab, 2016](https://arxiv.org/html/2610.14136#bib.bib21)\)\. We train classifiers to distinguish intermediate states𝒛t\\bm\{z\}\_\{t\}from the training distribution,𝒛t=\(1−t\)𝒙\+tϵ\\bm\{z\}\_\{t\}=\(1\-t\)\\bm\{x\}\+t\\bm\{\\epsilon\}, from those along denoising trajectories at inference\. Specifically, we experiment with SiT\([Ma et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib22)\)on 100K intermediate states, and report binary cross\-entropy \(BCE\) and AUROC on held\-out datasets\. A smaller training\-inference gap should make the two distributions harder to distinguish, bringing AUROC closer to0\.50\.5\(random chance\)\.
As shown by the BCE in Fig\.[17](https://arxiv.org/html/2610.14136#S4.F17), the classifier can distinguish training and inference states, with classification becoming easier as the noise level decreases\. This trend is consistent with errors accumulating during sampling due to exposure bias\([Ning et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib24);[Ning et al\., 2023](https://arxiv.org/html/2610.14136#bib.bib23)\)\. Notably, a single correction \(gray lines\) yields higher BCE and lower AUROC, suggesting that correctionreduces the training–inference gap\. Iterative correction generally reduces the gap further \(blue lines\), which may contribute to better sample quality even when reconstruction error on training inputs increases\.
### 4\.3How training affects iterative correction
##### Extrapolation is key to substantial quality gains\.
Figure 18:Correction beyond the supervised training depth improves generation \(lower FID\)\.Supervising more correction steps shifts the best\-performing tested inference correction countKKtoward 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\.We denote the number of correction iterations applied during training byKtrainK\_\{\\mathrm\{train\}\}\. In our default setting \(Ktrain=1K\_\{\\mathrm\{train\}\}=1\), the corrector sees only the denoiser’s initial prediction during training\. We also train the corrector to predict the target after multiple correction iterations \(Ktrain\>1K\_\{\\mathrm\{train\}\}\>1\)\. IncreasingKtrainK\_\{\\mathrm\{train\}\}shifts the best\-performing inference correction countKKtoward larger values, while the best achieved FID remains similar \(Fig\.[18](https://arxiv.org/html/2610.14136#S4.F18), left\)\. Notably,the best performance comes from applying correction more times than supervised during training\. The right panel shows that increasingKtrainK\_\{\\mathrm\{train\}\}lowers denoising error below the base denoiser’s only for correction iterations included in training\. ForK\>KtrainK\>K\_\{\\mathrm\{train\}\}, the error remains above this baseline, even at the correction count with the best FID\.
Our intuition is that the corrector learns a mapping toward the optimal denoising prediction\. Supervising more correction iterations may simply spread this correction over more steps, so more inference iterations are needed to obtain similar gains fromextrapolatingthe learned mapping beyond the iterations used during training\.
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\.
##### Denoiser updates affect iterative correction\.
Figure 20:Frozen\-denoiser correction is competitive at smallKKbut deteriorates earlier\.To examine how denoiser updates affect iterative correction, we compare \(a\) jointly training the denoiser and corrector, our default setting in Algorithm[1](https://arxiv.org/html/2610.14136#alg1), with \(b\) freezing the denoiser and training only the corrector under the same correction framework\. Interestingly, as shown in Fig\.[19](https://arxiv.org/html/2610.14136#S4.F19), samples in \(b\) improve with correction but change more abruptly in structure and appearance, with slight oversaturation typical of CFG\. Fig\.[20](https://arxiv.org/html/2610.14136#S4.F20)shows similar FID for both settings atK=5K=5\. With further correction, FID continues to decrease with joint training but increases again with a frozen denoiser despite visually improved samples, similar to the trend at higher CFG scales\. One possible interpretation is that, in \(b\), the corrector specializes in the errors of a particular denoiser checkpoint, making it less robust to iterative application\. In \(a\), the denoiser continues to evolve during training, exposing the corrector to a broader range of errors\. This may help the corrector generalize to its own outputs and remain effective across different correction countsKK\.
## 5Related work
Stacked generalization\([Wolpert, 1992](https://arxiv.org/html/2610.14136#bib.bib36)\)trains a higher\-level predictor on held\-out predictions of base models to improve generalization\. In diffusion models, self\-conditioning\([Chen et al\., 2022](https://arxiv.org/html/2610.14136#bib.bib4)\)uses previous clean\-data predictions as additional denoiser inputs, and recent concurrent works extend this idea to iterative refinement in language modeling and structured reasoning\([Yoo et al\., 2026](https://arxiv.org/html/2610.14136#bib.bib38);[Helbling et al\., 2026](https://arxiv.org/html/2610.14136#bib.bib6)\)\. We train an explicit corrector on denoiser predictions and reuse it across correction iterations, without requiring the original noisy input\. Idempotent Generative Networks\([Shocher et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib29)\)also learn iterative mappings through an explicit idempotency objective\. Our method retains the denoising objective and applies correction at each noise level to refine predictions before the sampling update\. The most relevant work to ours is[Schiff et al\. \(2026\)](https://arxiv.org/html/2610.14136#bib.bib28), which trains a corrector on generated samples and applies it repeatedly during sampling in discrete diffusion for language modeling\.
## 6Conclusion
We introduce correction training, which learns to correct a diffusion model’s predictions using the original clean\-data targets\. Despite being trained for a single correction, the corrector supports repeated refinement at inference\. Experiments on class\-conditional and unconditional ImageNet generation and text\-to\-image generation demonstrate substantial improvements in sample quality\. Correction can achieve competitive quality with guidance while better preserving diversity\. Correction can also be distilled to remove inference overhead\. Our analysis offers empirical insights into why correction works, but a theoretical understanding of its effectiveness and iterative behavior remains an open question for future research\.
#### Acknowledgments
We would like to thank Konpat Preechakul for valuable insights and feedback, as well as Michael Gharbi for insightful discussions, and Nicolas Dufour for providing detailed feedback on the manuscript\. We are grateful to A\. Sophia Koepke, Anne Harrington, Amil Dravid, Jorge Diaz Chao, Isabella Yu, Yutong Bai, Yong\-Hyun Park, Sangyoon Lee, Sihyun Yu, Susung Hong, Joonghyuk Shin, Gyuseong Lee, and Sanghyun Lee for helpful discussions\. This work was supported, in part, by NSF IIS\-2403305 and ONR MURI grants\.
## References
- Ahn et al\. \(2024\)Donghoon Ahn, Hyoungwon Cho, Jaewon Min, Wooseok Jang, Jungwoo Kim, SeonHwa Kim, Hyun Hee Park, Kyong Hwan Jin, and Seungryong Kim\.Self\-rectifying diffusion sampling with perturbed\-attention guidance\.In*European Conference on Computer Vision*, pp\. 1–17\. Springer, 2024\.
- Albergo et al\. \(2025\)Michael Albergo, Nicholas M Boffi, and Eric Vanden\-Eijnden\.Stochastic interpolants: A unifying framework for flows and diffusions\.*Journal of Machine Learning Research*, 26\(209\):1–80, 2025\.
- Chen et al\. \(2025\)Jiuhai Chen, Zhiyang Xu, Xichen Pan, Yushi Hu, Can Qin, Tom Goldstein, Lifu Huang, Tianyi Zhou, Saining Xie, Silvio Savarese, et al\.Blip3\-o: A family of fully open unified multimodal models\-architecture, training and dataset\.*arXiv preprint arXiv:2505\.09568*, 2025\.
- Chen et al\. \(2022\)Ting Chen, Ruixiang Zhang, and Geoffrey Hinton\.Analog bits: Generating discrete data using diffusion models with self\-conditioning\.*arXiv preprint arXiv:2208\.04202*, 2022\.
- Ghosh et al\. \(2023\)Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt\.Geneval: An object\-focused framework for evaluating text\-to\-image alignment\.*Advances in Neural Information Processing Systems*, 36:52132–52152, 2023\.
- Helbling et al\. \(2026\)Alec Helbling, Andrey Bryutkin, Mauro Martino, Nima Dehmamy, and Hendrik Strobelt\.Flow reasoning models: Scaling reasoning through iterative self\-refinement\.*arXiv preprint arXiv:2606\.29150*, 2026\.
- Heusel et al\. \(2017\)Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter\.Gans trained by a two time\-scale update rule converge to a local nash equilibrium\.*Advances in neural information processing systems*, 30, 2017\.
- Ho & Salimans \(2022\)Jonathan Ho and Tim Salimans\.Classifier\-free diffusion guidance\.*arXiv preprint arXiv:2207\.12598*, 2022\.
- Ho et al\. \(2020\)Jonathan Ho, Ajay Jain, and Pieter Abbeel\.Denoising diffusion probabilistic models\.*Advances in neural information processing systems*, 33:6840–6851, 2020\.
- Hu et al\. \(2024\)Xiwei Hu, Rui Wang, Yixiao Fang, Bin Fu, Pei Cheng, and Gang Yu\.Ella: Equip diffusion models with llm for enhanced semantic alignment\.*arXiv preprint arXiv:2403\.05135*, 2024\.
- Karras et al\. \(2024\)Tero Karras, Miika Aittala, Tuomas Kynkäänniemi, Jaakko Lehtinen, Timo Aila, and Samuli Laine\.Guiding a diffusion model with a bad version of itself\.*Advances in Neural Information Processing Systems*, 37:52996–53021, 2024\.
- Kobayashi et al\. \(2026\)Yuya Kobayashi, Masato Ishii, Yuhta Takida, Takashi Shibuya, and Yuki Mitsufuji\.Spectral prior for reducing exposure bias in diffusion models\.*arXiv preprint arXiv:2607\.22091*, 2026\.
- Kynkäänniemi et al\. \(2019\)Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila\.Improved precision and recall metric for assessing generative models\.*Advances in neural information processing systems*, 32, 2019\.
- Labs \(2025\)Black Forest Labs\.FLUX\.2: Frontier Visual Intelligence\.[https://bfl\.ai/blog/flux\-2](https://bfl.ai/blog/flux-2), 2025\.
- Leng et al\. \(2026\)Xingjian Leng, Jaskirat Singh, Zhanhao Liang, Ethan Smith, Martin Bell, Aninda Saha, Yuhui Yuan, and Liang Zheng\.Diffusionbench: On holistic evaluation of diffusion transformers\.*arXiv preprint arXiv:2606\.24888*, 2026\.
- Li et al\. \(2024a\)Mingxiao Li, Tingyu Qu, Ruicong Yao, Wei Sun, and Marie\-Francine Moens\.Alleviating exposure bias in diffusion models through sampling with shifted time steps\.In*International Conference on Learning Representations*, volume 2024, pp\. 16816–16838, 2024a\.
- Li & He \(2026\)Tianhong Li and Kaiming He\.Back to basics: Let denoising generative models denoise\.In*Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp\. 36115–36125, 2026\.
- Li et al\. \(2024b\)Tianhong Li, Dina Katabi, and Kaiming He\.Return of unconditional generation: A self\-supervised representation generation method\.*Advances in Neural Information Processing Systems*, 37:125441–125468, 2024b\.
- Lipman et al\. \(2022\)Yaron Lipman, Ricky TQ Chen, Heli Ben\-Hamu, Maximilian Nickel, and Matt Le\.Flow matching for generative modeling\.*arXiv preprint arXiv:2210\.02747*, 2022\.
- Liu et al\. \(2022\)Xingchao Liu, Chengyue Gong, and Qiang Liu\.Flow straight and fast: Learning to generate and transfer data with rectified flow\.*arXiv preprint arXiv:2209\.03003*, 2022\.
- Lopez\-Paz & Oquab \(2016\)David Lopez\-Paz and Maxime Oquab\.Revisiting classifier two\-sample tests\.*arXiv preprint arXiv:1610\.06545*, 2016\.
- Ma et al\. \(2024\)Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden\-Eijnden, and Saining Xie\.Sit: Exploring flow and diffusion\-based generative models with scalable interpolant transformers\.In*European Conference on Computer Vision*, pp\. 23–40\. Springer, 2024\.
- Ning et al\. \(2023\)Mang Ning, Enver Sangineto, Angelo Porrello, Simone Calderara, and Rita Cucchiara\.Input perturbation reduces exposure bias in diffusion models\.*arXiv preprint arXiv:2301\.11706*, 2023\.
- Ning et al\. \(2024\)Mang Ning, Mingxiao Li, Jianlin Su, Albert Ali Salah, and Itir Onal Ertugrul\.Elucidating the exposure bias in diffusion models\.In*International Conference on Learning Representations*, volume 2024, pp\. 15167–15189, 2024\.
- Peebles & Xie \(2023\)William Peebles and Saining Xie\.Scalable diffusion models with transformers\.In*2023 IEEE/CVF International Conference on Computer Vision \(ICCV\)*, pp\. 4172–4182\. IEEE, 2023\.
- Russakovsky et al\. \(2015\)Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C\. Berg, and Li Fei\-Fei\.ImageNet Large Scale Visual Recognition Challenge\.*International Journal of Computer Vision \(IJCV\)*, 115\(3\):211–252, 2015\.doi:10\.1007/s11263\-015\-0816\-y\.
- Sadat et al\. \(2024\)Seyedmorteza Sadat, Otmar Hilliges, and Romann M Weber\.Eliminating oversaturation and artifacts of high guidance scales in diffusion models\.In*The Thirteenth International Conference on Learning Representations*, 2024\.
- Schiff et al\. \(2026\)Yair Schiff, Omer Belhasin, Roy Uziel, Guanghan Wang, Marianne Arriola, Gilad Turok, Ran Zilberstein, Michael Elad, and Volodymyr Kuleshov\.Learn from your mistakes: Self\-correcting masked diffusion models\.*arXiv preprint arXiv:2602\.11590*, 2026\.
- Shocher et al\. \(2024\)Assaf Shocher, Amil Dravid, Yossi Gandelsman, Inbar Mosseri, Michael Rubinstein, and Alexei Efros\.Idempotent generative network\.In*International Conference on Learning Representations*, volume 2024, pp\. 35440–35460, 2024\.
- Sohl\-Dickstein et al\. \(2015\)Jascha Sohl\-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli\.Deep unsupervised learning using nonequilibrium thermodynamics\.In*International conference on machine learning*, pp\. 2256–2265\. pmlr, 2015\.
- Song et al\. \(2025\)Kiwhan Song, Jaeyeon Kim, Sitan Chen, Yilun Du, Sham Kakade, and Vincent Sitzmann\.Selective underfitting in diffusion models\.*arXiv preprint arXiv:2510\.01378*, 2025\.
- Song & Ermon \(2019\)Yang Song and Stefano Ermon\.Generative modeling by estimating gradients of the data distribution\.*Advances in neural information processing systems*, 32, 2019\.
- Song et al\. \(2020\)Yang Song, Jascha Sohl\-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole\.Score\-based generative modeling through stochastic differential equations\.*arXiv preprint arXiv:2011\.13456*, 2020\.
- Sun et al\. \(2023\)Keqiang Sun, Junting Pan, Yuying Ge, Hao Li, Haodong Duan, Xiaoshi Wu, Renrui Zhang, Aojun Zhou, Zipeng Qin, Yi Wang, et al\.Journeydb: A benchmark for generative image understanding\.*Advances in neural information processing systems*, 36:49659–49678, 2023\.
- Wang et al\. \(2026\)Shuai Wang, Zhi Tian, Weilin Huang, and Limin Wang\.Ddt: Decoupled diffusion transformer\.In*Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp\. 40633–40642, 2026\.
- Wolpert \(1992\)David H Wolpert\.Stacked generalization\.*Neural networks*, 5\(2\):241–259, 1992\.
- Yao et al\. \(2025\)Jingfeng Yao, Bin Yang, and Xinggang Wang\.Reconstruction vs\. generation: Taming optimization dilemma in latent diffusion models\.In*2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\)*, pp\. 15703–15712\. IEEE, 2025\.
- Yoo et al\. \(2026\)Jaehoon Yoo, Wonjung Kim, Floor Eijkelboom, Chanhyuk Lee, Nicholas M Boffi, Seunghoon Hong, and Jinwoo Kim\.Self\-conditioned flow map language models via fixed\-point flows\.*arXiv preprint arXiv:2607\.00714*, 2026\.
- Yu & Zhan \(2025\)Meng Yu and Kun Zhan\.Frequency regulation for exposure bias mitigation in diffusion models\.In*Proceedings of the 33rd ACM International Conference on Multimedia*, pp\. 10370–10378, 2025\.
- Yu et al\. \(2024\)Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie\.Representation alignment for generation: Training diffusion transformers is easier than you think\.*arXiv preprint arXiv:2410\.06940*, 2024\.
## Appendix Contents
## Appendix ADesign choices and alternative realizations
##### A practical realization of correction training\.
In this section, we examine the design choices in our framework\. The default configuration used in the main paper is a simple implementation compatible with pretrained models: we use a separate corrector initialized from the base denoiser, omit the noisy input𝒛t\\bm\{z\}\_\{t\}, retain the timestep\-conditioning pathway to preserve the original architecture, and apply stop\-gradient to prevent the correction loss from updating the denoiser\. Many other combinations of these design choices are possible within our correction training framework\. Instead of exhaustively evaluating all combinations, we vary one component at a time from the default configuration to understand how each choice affects performance\.
### A\.1Corrector interface
##### Access to the noisy state𝒛t\\bm\{z\}\_\{t\}from corrector\.
Prior approaches\([Chen et al\., 2022](https://arxiv.org/html/2610.14136#bib.bib4)\)to learning from model predictions often also provide the original input\. In our diffusion setting, this corresponds to giving the corrector access to𝒛t\\bm\{z\}\_\{t\}in addition to the denoiser’s prediction\. Tab\.[6](https://arxiv.org/html/2610.14136#A1.T6)shows that, on JiT, correction without𝒛t\\bm\{z\}\_\{t\}achieves comparable FID to correction with𝒛t\\bm\{z\}\_\{t\}, while requiring no architectural changes to the denoiser\.
The SiT comparison in Fig\.[21](https://arxiv.org/html/2610.14136#A1.F21)further shows how this choice affects the precision–recall trade\-off\. Here,pp\-only denotes conditioning on the proposal𝒙^\\hat\{\\bm\{x\}\}alone, whereasp\+𝒛tp\+\\bm\{z\}\_\{t\}provides both the proposal and the noisy input\. Proposal\-only correction achieves higher precision at slightly lower recall\. One possible explanation is that, without access to the original noisy input, the corrector is less constrained by it and can modify the proposal more aggressively\.
Table 6:Proposal\-only correction performs comparably tozt\\bm\{z\}\_\{t\}\-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, sweepingK=1K=1–88\. Stars mark each arm’s lowest\-FID setting\.
##### Timestep conditioning\.
By default, we provide the true timestep to reuse the pretrained model’s timestep\-conditioning pathway\. However, Tab\.[7](https://arxiv.org/html/2610.14136#A1.T7)shows that using the true timestep or a constant time embedding yields similar final FID in our residual\-correction setting\. This finding may not generalize to other settings, including correctors trained from scratch: the pretrained pathway associatesttwith the noise level of the original noisy input, which does not directly describe the errors in the clean\-sample prediction received by the corrector\.
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 atK=1K=1, using FID\-50K at 250 NFE\. The denoiser uses the true timestep in both cases\.
##### Prediction\-space versus hidden\-state refinement\.
We compare correction in output space \(pixels or latents\) with refinement of the network’s hidden states\. Output\-space correction benefits from repeated application, whereas the hidden\-state variants do not achieve comparable gains \(Tab\.[8](https://arxiv.org/html/2610.14136#A1.T8)\)\. Increasing the number of correction iterations during training does not improve sample quality in either setting\. Fig\.[22](https://arxiv.org/html/2610.14136#A1.F22)presents qualitative comparisons atK=5K=5\.
On the other hand, removing stop\-gradient from the base denoiser’s prediction achieves a lower best FID, but retaining it yields more consistent gains at largerKK\. Without stop\-gradient, samples become visually simpler under repeated correction, as shown in Fig\.[22](https://arxiv.org/html/2610.14136#A1.F22)\(b\)\.
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](https://arxiv.org/html/2610.14136#A1.T8)atK=5K=5, using the same initial noise and class labels\.
### A\.2Training configuration
##### Stop\-gradient versus end\-to\-end optimization\.
We further evaluate the effect of stop\-gradient on additional pretrained models \(Fig\.[23](https://arxiv.org/html/2610.14136#A1.F23)\)\. Unlike the earlier comparison in Table[8](https://arxiv.org/html/2610.14136#A1.T8), the variants achieve similar best FID here\. Full stop\-gradient \(our default\) nevertheless leads to less degradation at larger correction countsKK\.
Figure 23:Full stop\-gradient degrades more slowly at large correction counts in these comparisons\.FID@5K versusKK, 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\.
### A\.3Inference configuration
##### Correction countKK\.
IncreasingKKimproves FID up to a setting\-dependent optimum, after which further correction can degrade it \(Tab\.[8](https://arxiv.org/html/2610.14136#A1.T8); Fig\.[23](https://arxiv.org/html/2610.14136#A1.F23)\)\.
##### Correction interval\.
Restricting JiT correction to early steps and reallocating the saved calls to denoising improves FID at a fixed sampling budget \(Tab\.[9](https://arxiv.org/html/2610.14136#A1.T9)\)\.
fullFull trajectory \(1→01\\to 0\)\.
Table 9:Allocating correction to early steps yields the lowest FID at a fixed sampling budget\.JiT\-L/16 withK=5K=5, evaluated by FID8K\. Rows and columns specify the start and end noise levels, decreasing from noise to data\.\(a\)

\(b\)
Figure 24:Skipping early correction refines the baseline pose\.Gray boxes show the matchedK=0K=0baseline; 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 eachKK\.\(a\)

\(b\)

Figure 25:\(a\) Full\-range correction changes fine texture and shape\. \(b\) Full\-range correction changes object layout and appearance\.Gray boxes show the matchedK=0K=0baseline; green, blue and orange mark early\-only, skip\-early and full\-range correction, respectively\.
## Appendix BImplementation Details
### B\.1Training configurations
##### Common backbone settings\.
We retain each pretrained model’s backbone and original conditioning\. The separate corrector uses an independent copy of the denoiser backbone; latent\-space models retain their original VAE, while JiT operates directly on RGB images\. Training and quantitative sampling settings are summarized in Tables[10](https://arxiv.org/html/2610.14136#A2.T10)and[13](https://arxiv.org/html/2610.14136#A2.T13)\.
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;DDandCCdenote separate denoiser and corrector networks\. Sampling settings are reported separately\.†\\daggerSiTREPA\\text\{SiT\}\_\{\\mathrm\{REPA\}\}is initialized from the[publicly released SiT\-XL/2 checkpoint](https://github.com/sihyun-yu/REPA)trained with the REPA loss for 4M iterations; the REPA loss is not used during additional training\.
‡\\ddaggerT2I denotes text\-to\-image generation \(Sec\.[3\.2](https://arxiv.org/html/2610.14136#S3.SS2)\)\. We use the[DiffusionBench codebase](https://github.com/End2End-Diffusion/diffusion-bench)and its[T2I checkpoint](https://huggingface.co/diffusion-bench/diffusion-bench/blob/main/t2i/ckpts/t2i-ddt-en28d1152hd72-dn2d2048hd128-flux2-vpred-t4-v0-200k.pt), trained in the FLUX\.2 VAE latent space using velocity prediction\.
Table 11:The main quantitative comparisons use approximately 250 network evaluations per image\.Sampling settings for Tables[1](https://arxiv.org/html/2610.14136#S3.T1)and[12](https://arxiv.org/html/2610.14136#S3.F12)\.KKis the number of corrections per active step; intervals use noise timett\(11: noise,00: data\)\. NFE includes correction and both CFG branches\. Guided T2I applies correction before CFG\.MethodCFGwwKKCFG intervalCorrection intervalStepsNFEDiT\-XL/2; DDPMBaseline10——250250CFG1\.501→01\\to 0—125250Correction12—1→01\\to 083249SiT\-XL/2; Euler–Maruyama SDEBaseline10——250250CFG1\.501→01\\to 0—125250Correction12—1→01\\to 083249SiTREPA\-XL/2; Euler–Maruyama SDEBaseline10——250250CFG1\.800\.7→00\.7\\to 0—148250Correction15—0\.7→0\.50\.7\\to 0\.5123248JiT\-L/16; Heun ODE \(final Euler step\)Baseline10——125249CFG2\.400<t<0\.90<t<0\.9—63250Correction19—1→0\.81\\to 0\.845242T2I \(DiT \+ DDT\); Euler ODEBaseline10——250250Correction16—1→01\\to 035245CFG401→01\\to 0—125250CFG \+ Correction421→01\\to 01→01\\to 041246
Table 12:Unconditional ImageNet settings\.Sampling settings for Table[9](https://arxiv.org/html/2610.14136#S3.F9), using SiTREPA\-XL/2\.With CFG\(w=2w=2\)
With CFG\(w=6w=6\)
Table 13:Correction’s alignment gains depend on guidance strength and iteration count\.Additional guidance scales use the same checkpoint as Table[12](https://arxiv.org/html/2610.14136#S3.F12), but fix 50 Euler steps rather than NFE\. Correction is applied before CFG\. Bold marks the best score within each block\.
### B\.2Corrector parameterization
Algorithm[1](https://arxiv.org/html/2610.14136#alg1)usesxx\-prediction\. For anϵ\\epsilon\- orvv\-predicting denoiser, we recommend the following two practical options\.
##### Direct clean\-sample prediction\.
The simplest starting point is to convert the denoiser’s output into a clean\-sample proposal and train the corrector to predict a refined clean sample:
𝒙^\(1\)=Cϕ\(sg\(𝒙^\(0\)\),t\)\.\\hat\{\\bm\{x\}\}^\{\(1\)\}=C\_\{\\bm\{\\phi\}\}\\\!\\left\(\\operatorname\{sg\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\}\),t\\right\)\.The result can then be converted into the representation required by the sampler using𝒛t\\bm\{z\}\_\{t\},tt, and the noise schedule\. These conversions are performed outside the corrector network;𝒛t\\bm\{z\}\_\{t\}need not be supplied as an additional network input\. The training loss can still be computed in velocity or noise space\.
Why convert the proposal toxxrather than retain the denoiser’s original output space? In our small\-scale experiments with a single correction \(K=1K=1\) and without additional noisy\-state conditioning,xx\-space correction performs best among the tested same\-space input–output choices, including withϵ\\epsilon\- andvv\-predicting denoisers\. This makes clean\-sample input and output a useful starting point, especially when training the corrector from scratch\. These experiments do not establish a universally optimal parameterization since it’s small\-scale run onK=1K=1\. Directly correcting noise or velocity remains an option\.
##### Reusing pretrained output representations\.
When initializing the corrector from a pretrained denoiser, another consideration is how much of its learned representation to retain\. For example, SiT\([Ma et al\., 2024](https://arxiv.org/html/2610.14136#bib.bib22)\)is trained to predict velocity\. Changing its output toxxrequires adapting the pretrained output representation\. Directxx\-prediction can still work and converged quickly in our experiments, but retaining the original output parameterization offers another way to preserve the representation of pretrained models\. This motivated our use of clean\-sample input with nativeϵ\\epsilon\- orvv\-output in the main pretrained\-model experiments\.
Retaining the native output also raises a practical question: how should it be combined with a clean\-sample input? An exact conversion betweenxx,ϵ\\epsilon, andvvrequires𝒛t\\bm\{z\}\_\{t\}or another component in addition to the quantity being converted\. Providing𝒛t\\bm\{z\}\_\{t\}as an additional condition is possible, but its empirical benefits were mixed \(App\.[A\.1](https://arxiv.org/html/2610.14136#A1.SS1.SSS0.Px1)\), so we omit it by default\. Instead, we retain the base prediction in its native output space and linearly combine it with the corrector’s native\-space output through a learned gate\. This allows us to use a clean\-sample input while retaining the pretrained output space, without supplying𝒛t\\bm\{z\}\_\{t\}to the network or requiring an explicit conversion inside it\.
Let𝒒\\bm\{q\}denote the target in the chosen output space and𝒒^\(0\)\\hat\{\\bm\{q\}\}^\{\(0\)\}the base prediction\. Our full stop\-gradient formulation is:
𝒒^\(1\)\\displaystyle\\hat\{\\bm\{q\}\}^\{\(1\)\}=sg\(𝒒^\(0\)\)\+𝜸⊙\[Cϕq\(sg\(𝒙^\(0\)\),t\)−sg\(𝒒^\(0\)\)\],\\displaystyle=\\operatorname\{sg\}\(\\hat\{\\bm\{q\}\}^\{\(0\)\}\)\+\\bm\{\\gamma\}\\odot\\left\[C\_\{\\bm\{\\phi\}\}^\{q\}\\\!\\left\(\\operatorname\{sg\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\}\),t\\right\)\-\\operatorname\{sg\}\(\\hat\{\\bm\{q\}\}^\{\(0\)\}\)\\right\],\(8\)ℒcorr\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{corr\}\}=𝔼𝒙,ϵ,t\[λq\(t\)‖𝒒^\(1\)−𝒒‖22\],\\displaystyle=\\mathbb\{E\}\_\{\\bm\{x\},\\bm\{\\epsilon\},t\}\\\!\\left\[\\lambda\_\{q\}\(t\)\\left\\\|\\hat\{\\bm\{q\}\}^\{\(1\)\}\-\\bm\{q\}\\right\\\|\_\{2\}^\{2\}\\right\],\(9\)where𝒒\\bm\{q\}is the target in the chosen output space, such asvvfor SiT orϵ\\epsilonfor DiT, andλq\(t\)\\lambda\_\{q\}\(t\)preserves the model’s loss weighting\. The unconstrained gate is initialized to zero, preserving the base prediction at initialization\. It is learned per channel, except for DiT and plain SiT, which use a scalar shared across channels\. For better performance, one can extend it to timestep\-dependent gating or 2D tensor gates, but we keep the \(per\-channel\) scalar for simplicity\.
##### Correction with a fixed base component\.
Gating is not the only way to combine different input and output parameterizations\. We also tested a formulation that combines the corrector’s output with a fixed component of the base prediction\. Under our conventionv=ϵ−xv=\\epsilon\-x, anxx\-space prediction is combined with the base noise prediction, while anϵ\\epsilon\-space prediction is combined with the base clean\-sample prediction:
𝒙^\(1\)=\{𝒙^,ϵ^−sg\(𝒗^\(0\)\),sg\(ϵ^\(0\)\)−𝒗^,ϵ^\(1\)=\{𝒙^\+sg\(𝒗^\(0\)\),ϵ^,sg\(𝒙^\(0\)\)\+𝒗^,𝒗^\(1\)=\{sg\(ϵ^\(0\)\)−𝒙^,ϵ^−sg\(𝒙^\(0\)\),𝒗^\.\\hat\{\\bm\{x\}\}^\{\(1\)\}=\\begin\{cases\}\\hat\{\\bm\{x\}\},\\\\ \\hat\{\\bm\{\\epsilon\}\}\-\\operatorname\{sg\}\(\\hat\{\\bm\{v\}\}^\{\(0\)\}\),\\\\ \\operatorname\{sg\}\(\\hat\{\\bm\{\\epsilon\}\}^\{\(0\)\}\)\-\\hat\{\\bm\{v\}\},\\end\{cases\}\\quad\\hat\{\\bm\{\\epsilon\}\}^\{\(1\)\}=\\begin\{cases\}\\hat\{\\bm\{x\}\}\+\\operatorname\{sg\}\(\\hat\{\\bm\{v\}\}^\{\(0\)\}\),\\\\ \\hat\{\\bm\{\\epsilon\}\},\\\\ \\operatorname\{sg\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\}\)\+\\hat\{\\bm\{v\}\},\\end\{cases\}\\quad\\hat\{\\bm\{v\}\}^\{\(1\)\}=\\begin\{cases\}\\operatorname\{sg\}\(\\hat\{\\bm\{\\epsilon\}\}^\{\(0\)\}\)\-\\hat\{\\bm\{x\}\},\\\\ \\hat\{\\bm\{\\epsilon\}\}\-\\operatorname\{sg\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\}\),\\\\ \\hat\{\\bm\{v\}\}\.\\end\{cases\}These expressions illustrate alternative fixed\-base\-component formulations for each output space\. Our small\-scale experiments used the velocity\-space formulation on the right and minimized velocity\-space MSE\.
Here𝒙^\\hat\{\\bm\{x\}\},ϵ^\\hat\{\\bm\{\\epsilon\}\}, and𝒗^\\hat\{\\bm\{v\}\}denote the corrector’s predictions\. The input can independently be chosen from the base clean, noise, or velocity predictions\. All combinations minimize velocity\-space MSE without a learned gate\. Forxx\- andϵ\\epsilon\-space predictions, this objective also requires compensating for errors in the retained base component, so these predictions are not directly supervised against the clean sample or noise alone\.
By explicitly using the decompositionv=ϵ−xv=\\epsilon\-x, this formulation may offer a more principled way to connect different input and output parameterizations than a learned blend\. Our exploration was limited by time constraints, however: we tested it only in small\-scale experiments with a single correction and did not extensively investigate repeated correction or larger\-scale settings\.
##### Gradient paths\.
The correction loss can reach the denoiser through the corrector’s input𝒙^\(0\)\\hat\{\\bm\{x\}\}^\{\(0\)\}and, when used, the retained base prediction𝒒^\(0\)\\hat\{\\bm\{q\}\}^\{\(0\)\}\. Our default,Full stop\-gradient \(SG\), detaches both \(Eq\.[8](https://arxiv.org/html/2610.14136#A2.E8)\)\. For separate networks, the denoiser is then updated only by its denoising loss; it is not frozen\. We also compareProposal\-only SG, which detaches only the corrector’s input, andNo SG, which detaches neither\. Full SG achieves competitive FID and degrades less at large correction counts in our comparisons \(App\.[A\.2](https://arxiv.org/html/2610.14136#A1.SS2.SSS0.Px1)\)\.
##### Full\-target versus residual prediction\.
Independently of the input and output spaces, the network can predict either the full target or a residual relative to the base prediction\. This choice applies both to direct clean\-sample prediction and to native\-space correction with gating\. Without gating, residual prediction takes the form:
𝒒^\(1\)=sg\(𝒒^\(0\)\)\+Rϕq\(sg\(𝒙^\(0\)\),t\)\.\\hat\{\\bm\{q\}\}^\{\(1\)\}=\\operatorname\{sg\}\(\\hat\{\\bm\{q\}\}^\{\(0\)\}\)\+R\_\{\\bm\{\\phi\}\}^\{q\}\\\!\\left\(\\operatorname\{sg\}\(\\hat\{\\bm\{x\}\}^\{\(0\)\}\),t\\right\)\.With gating, the gate can instead scale the predicted residual\. Thus, residual prediction is not an alternative to gating; the two can be used together\. The correction objective remains unchanged\. Zero\-initializing the residual output projection preserves the base prediction at initialization\. Although we do not present a dedicated comparison in this paper, residual prediction also worked well in our experiments, sometimes slightly better than full\-target prediction\. We use full\-target prediction in the main formulation for simplicity, but residual prediction is another practical option within correction training\.
##### Choosing a formulation\.
For a first implementation, particularly when training from scratch, we recommend clean\-sample input and direct clean\-sample output\. When reusing a pretrained denoiser, clean\-sample input with gated native\-space prediction is another practical starting point\. We retained this form in the main pretrained\-model experiments for compatibility with the pretrained output representation and its slightly better empirical performance\. Neither choice is required by correction training: input space, output space, full\-target versus residual prediction, and loss space are distinct design choices\. Other combinations can also be used, provided that the information needed for any conversion is available\. The network can also take noise or velocity proposals as input; we use clean\-sample input in the equations to illustrate our default\. Full\-target versus residual prediction can be chosen independently\.
Table 14:Corrector parameterization choices\.We distinguish the output formulation from full\-target versus residual prediction, which can be used with either formulation\. Hereqqdenotes the chosen output space\. The equations use clean\-sample input by default; input parameterization and loss space are separate choices\.
### B\.3Figure and Table Details
##### Unconditional generation \(Table[9](https://arxiv.org/html/2610.14136#S3.F9); Figure[9](https://arxiv.org/html/2610.14136#S3.F9)\)\.
We evaluate SiTREPA\-XL/2 with the null class after 50K updates, using EMA weights, 50,000 images, and the native Euler–Maruyama sampler\. Correction detaches both the proposal and the base prediction\. We sweepK=0,…,10K=0,\\ldots,10using⌊250/\(K\+1\)⌋\\lfloor 250/\(K\+1\)\\rfloordenoising steps, and autoguidancew=1,…,10w=1,\\ldots,10using the released pretrained guide at 250 NFE\. The table reports the lowest\-FID settings from these sweeps, in addition toK=1K=1\. Sampling settings are summarized in Table[13](https://arxiv.org/html/2610.14136#A2.T13)\.
##### Text\-to\-image examples \(Figure[12](https://arxiv.org/html/2610.14136#S3.F12)\)\.
We compareK=0K=0andK=6K=6using the same 50K\-update EMA checkpoint, without CFG, and 50 Euler steps \(50 and 350 NFE\)\. Paired images share their prompt and recorded rank seed\. Full prompts are listed in Table[15](https://arxiv.org/html/2610.14136#A2.T15)\.
Table 15:Full prompts for the text\-to\-image examples in Figure[12](https://arxiv.org/html/2610.14136#S3.F12)\.
##### Guidance and correction \(Figure[13](https://arxiv.org/html/2610.14136#S3.F13)\)\.
We comparew∈\{1,6\}w\\in\\\{1,6\\\}andK∈\{0,6\}K\\in\\\{0,6\\\}at 50 Euler steps\. The selected prompt is*a photo of a bench*, with all eight seeds retained per setting\. CFG is applied before correction\.
##### Parameter sharing \(Table[4](https://arxiv.org/html/2610.14136#S3.T4)\)\.
We compare plain SiT\-XL/2 models after 100K updates, without CFG, using FID50K\. The shared model uses separate input embedders and shares the transformer blocks and output head\. Both separate and shared correctors detach the proposal and base prediction, but shared parameters receive gradients from both losses\. The baseline is a separate denoising\-only continuation\. Baseline, separate\-corrector, and shared\-corrector settings use\(K,steps\)=\(0,250\),\(2,83\),\(3,62\)\(K,\\text\{steps\}\)=\(0,250\),\(2,83\),\(3,62\), respectively\.
##### Distillation \(Table[4](https://arxiv.org/html/2610.14136#S3.T4)\)\.
We freeze the separate SiT\-XL/2 teacher at EMA100K and initialize the student from its denoiser\. The student minimizes velocity MSE against the teacher’sK=2K=2prediction, without a ground\-truth loss, for 100K updates on class\-conditional ImageNet\. We sample Gaussian noise and uniform time, using batch size 256, AdamW with learning rate10−410^\{\-4\}, 100 warmup updates, zero weight decay, and EMA decay 0\.9999\. We evaluate the final EMA with FID@50K using the native SiT Euler–Maruyama sampler and ADM evaluator\.
##### Power spectra \(Figure[38](https://arxiv.org/html/2610.14136#A4.F38)\)\.
We average radial power over 100 ImageNet\-256 samples per setting forK=0,…,10K=0,\\ldots,10at approximately 250 NFE\. The panels showlog10\\log\_\{10\}power andlog10\(Psample/Preal\)\\log\_\{10\}\(P\_\{\\mathrm\{sample\}\}/P\_\{\\mathrm\{real\}\}\)over normalized frequencies 0\.1–0\.5\.
## Appendix CAdditional Results
### C\.1Comparing Correction and CFG
Fig\.[26](https://arxiv.org/html/2610.14136#A3.F26)compares correction and CFG on ImageNet using the same checkpoint and initial noise\. In these selected examples, increasing guidance produces more similar close\-up compositions and stronger colors, whereas correction retains more varied poses and backgrounds\. These observations concern the displayed samples; they do not by themselves establish distributional mode collapse\.
Figure 26:Correction and guidance produce different visual trade\-offs\.Matched\-noise samples from SiTREPA\-XL/2 as correction depthKK\(top\) or guidance scaleww\(bottom\) increases, at approximately matched NFE\.Fig\.[27](https://arxiv.org/html/2610.14136#A3.F27)quantifies the trade\-offs on DiT\. At approximately matched sampling compute, correction reaches a lower minimum FID than the best evaluated CFG setting\.
Following[Sadat et al\. \(2024\)](https://arxiv.org/html/2610.14136#bib.bib27), we also measure mean HSV saturation against a real\-image reference\. Stronger CFG raises saturation above the reference, whereas correction remains near or below the unguided level\. This ImageNet measurement complements the text\-to\-image examples; the latter illustrate refinement at weaker guidance rather than measuring a reduction in saturation caused by correction\.
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\.##### Additional Correction and Guidance Comparisons
Figures[28](https://arxiv.org/html/2610.14136#A3.F28),[29](https://arxiv.org/html/2610.14136#A3.F29), and[30](https://arxiv.org/html/2610.14136#A3.F30)retain all 16 fixed noise seeds \(0–15\) for nine classes selected from an exploratory 24\-class comparison\. No samples are omitted within these classes\. The correction and guidance rows use separately trained SiT\-XL/2 models at the same 100K\-update budget, with EMA weights\. Correction uses no CFG\. All settings use approximately 250 network evaluations \(249–252\)\. Panels share the same initial Gaussian noise; their denoising grids and subsequent stochastic trajectories differ\. FID labels are existing ImageNet\-wide FID@50K results, rather than FID estimates for the displayed class\. These examples illustrate visual trade\-offs and do not estimate the number of modes covered by either method\.
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 depthK=0,2,4,6K=0,2,4,6without CFG; the lower row varies guidance scalew=1,3,5,7w=1,3,5,7in 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 depthK=0,2,4,6K=0,2,4,6without CFG; the lower row varies guidance scalew=1,3,5,7w=1,3,5,7in 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](https://arxiv.org/html/2610.14136#A3.F28)and[29](https://arxiv.org/html/2610.14136#A3.F29)\. Top: correctionK=0,2,4,6K=0,2,4,6without CFG\. Bottom: CFGw=1,3,5,7w=1,3,5,7\. FID labels are ImageNet\-wide FID@50K\.
### C\.2Combining CFG and Correction
##### Combination order\.
CFG and correction can be applied in either order\. Applying CFG first combines the conditional and null predictions before the corrector refines the guided prediction using the conditional text\. Applying correction first refines both branches independently and then combines them with CFG\. We use CFG→\\rightarrowcorrection for the guided qualitative examples and the main human\-preference comparison because we prefer its visual appearance, while correction→\\rightarrowCFG is used for the guided alignment results in Tab\.[12](https://arxiv.org/html/2610.14136#S3.F12)\.
Table[16](https://arxiv.org/html/2610.14136#A3.T16)compares quality and diversity preferences at matchingw,Kw,Kacross the two orders\. Among the evaluated guided settings, correction before CFG more often yields higher quality preference, while correction after CFG more often yields higher diversity preference\. Thus neither order consistently leads on both criteria\.
Table 16:Neither correction order consistently leads in both quality and diversity\.Full sweep of candidate vote shares \(%\) against released CFG \(w=6w=6,K=0K=0\); higher is better\. Bold marks the setting reported in Table[2](https://arxiv.org/html/2610.14136#S3.T2)\. Each entry contains 800 responses across 80 prompts\. The two orders were rated separately against the baseline\.
### C\.3Human evaluation
We evaluate our method with human raters, using[Rapidata](https://www.rapidata.ai/)to collect human\-feedback\. It enables high\-throughput collection of diverse human judgments, roughly takes few minutes on thousands of human responses\. Annotator qualification and reliability filtering are applied to control response quality\. We additionally used language filter “English” and age filter “adult”\.
Table[2](https://arxiv.org/html/2610.14136#S3.T2)reports unweighted preference votes on all 80 GenEval single\-object prompts, with 10 judgments per prompt \(800 per metric\)\. Quality uses one image per method and diversity uses four\-image sets\. We first apply CFG and then correct the guided prediction using the conditional text\. Both methods use 50 Euler steps\.
Table[16](https://arxiv.org/html/2610.14136#A3.T16)provides the full results for both correction orders overw∈\{1,2,4,6\}w\\in\\\{1,2,4,6\\\}andK∈\{2,4,6,8\}K\\in\\\{2,4,6,8\\\}, highlighting the setting reported in the main text\. Each setting is compared separately against the same released baseline\. Atw=1w=1, the two orders use identical images, so differences reflect repeated human ratings\. Figure[31](https://arxiv.org/html/2610.14136#A3.F31)illustrates the questions and stimuli for the main table’s CFG\-then\-correction setting \(w=4w=4,K=8K=8\)\.
Figure 31:Quality and diversity are assessed separately\.Actual teddy\-bear stimuli and questions for Table[2](https://arxiv.org/html/2610.14136#S3.T2), using CFG followed by correction \(w=4w=4,K=8K=8\): single images for quality and four\-image sets for diversity\.
## Appendix DDiscussion
### D\.1Training and capacity controls
##### Additional\-training and capacity controls\.
We examine how quickly the corrector converges by comparing correction training with three baselines: continued training of the base denoiser and two deeper denoisers with different initialization schemes\. We use the pretrainedSiTREPA\\text\{SiT\}\_\{\\mathrm\{REPA\}\}\-XL/2 denoiser with 28 transformer layers from Tab\.[1](https://arxiv.org/html/2610.14136#S3.T1)\. For the deeper denoisers, we append 28 transformer layers before the final unembedding layer, matching the total number of transformer layers in the denoiser–corrector pair\. The added layers are either randomly initialized or copied from the pretrained denoiser, with the latter matching our corrector initialization\. We train all configurations for 50K steps and report their progress in Fig\.[32](https://arxiv.org/html/2610.14136#A4.F32)\.We use a single correction pass \(K=1K=1\) to isolate its benefit from the additional gains of iterative refinement\.
Figure 32:Correction improves generation quality beyond continued denoiser training or parameter\-matched deeper denoisers\.As can be seen in Fig\.[32](https://arxiv.org/html/2610.14136#A4.F32)\(a\) and \(b\), correction improves FID substantially faster over the 50K\-update training window we used, whereas continued training and both deeper\-denoiser controls improve much less\. Importantly, this gain appears with only asingle correction pass: despite using the same total sequential transformer depth as the 56\-layer denoisers, theD28\+C28D\_\{28\}\+C\_\{28\}correction model achieves substantially better generation quality\.These results suggest thatadditional training or model depth alone cannot explain the gains, pointing to an advantage of learning to refine an already structured prediction over learning directly from noisy observations\.
##### Iterative correction improves quality at matched effective depth\.
Figure 33:Matched effective depth\.Correction outperforms a deeper denoiser at matched effective depth per denoising step and inference budget, using fewer parameters\.One might argue that the gain from repeating correction simply comes from increasing the effective depth within each denoising step\. To test whether this is the case, we compare the 28\-layer monolithic denoiser \(D28D\_\{28\}\) with a 7\-layer denoiser and a 7\-layer corrector \(D7\+C7D\_\{7\}\+C\_\{7\}\), usingK=3K=3corrections to match the effective depth of the monolithic denoiser \(7\+3×7=287\+3\\times 7=28; Fig\.[33](https://arxiv.org/html/2610.14136#A4.F33)\)\. For this comparison, we train both configurations from scratch for 80K updates\. Both use exactly 240 denoising steps and 6,720 transformer\-layer evaluations per sample\. Surprisingly, despite using roughly half as many parameters \(346M versus 675M\), the shallow denoiser–corrector pair achieves a lower FID@50K \(4\.21 versus 4\.88; Tab\.[5](https://arxiv.org/html/2610.14136#S4.T5)\)\. This highlights that iterative correction does more than improve prediction accuracy through additional denoiser layers\.
### D\.2Iterative correction visualizations
##### Analytic target\.
Letℬy=\{𝒙i\}i=1Ny\\mathcal\{B\}\_\{y\}=\\\{\\bm\{x\}\_\{i\}\\\}\_\{i=1\}^\{N\_\{y\}\}be a bank of clean images from classyy, with a uniform empirical prior\. For𝒛t=\(1−t\)𝒙\+tϵ\\bm\{z\}\_\{t\}=\(1\-t\)\\bm\{x\}\+t\\bm\{\\epsilon\}, whereϵ∼𝒩\(𝟎,𝑰\)\\bm\{\\epsilon\}\\sim\\mathcal\{N\}\(\\bm\{0\},\\bm\{I\}\)and0<t≤10<t\\leq 1, the posterior mean under this prior is
𝒙analytic\(𝒛t,t,y\)\\displaystyle\\bm\{x\}\_\{\\mathrm\{analytic\}\}\(\\bm\{z\}\_\{t\},t,y\)=∑i=1Nywi\(𝒛t,t\)𝒙i∑i=1Nywi\(𝒛t,t\),\\displaystyle=\\frac\{\\sum\_\{i=1\}^\{N\_\{y\}\}w\_\{i\}\(\\bm\{z\}\_\{t\},t\)\\,\\bm\{x\}\_\{i\}\}\{\\sum\_\{i=1\}^\{N\_\{y\}\}w\_\{i\}\(\\bm\{z\}\_\{t\},t\)\},\(10\)wi\(𝒛t,t\)\\displaystyle w\_\{i\}\(\\bm\{z\}\_\{t\},t\)=exp\(−‖𝒛t−\(1−t\)𝒙i‖222t2\)\.\\displaystyle=\\exp\\\!\\left\(\-\\frac\{\\\|\\bm\{z\}\_\{t\}\-\(1\-t\)\\bm\{x\}\_\{i\}\\\|\_\{2\}^\{2\}\}\{2t^\{2\}\}\\right\)\.
##### Additional iterative\-correction examples\.
Figures[34](https://arxiv.org/html/2610.14136#A4.F34)and[34](https://arxiv.org/html/2610.14136#A4.F34)show the remaining selected examples from the visualization in Fig\.[15](https://arxiv.org/html/2610.14136#S4.F15), without repeating its six rows\. All use the same JiT\-L/16 checkpoint after 100K training updates and its second EMA weights, with class conditioning and CFG scale11\. For each row, the denoiser receives a forward input𝒛t=\(1−t\)𝒙\+tϵ\\bm\{z\}\_\{t\}=\(1\-t\)\\bm\{x\}\+t\\bm\{\\epsilon\}, and the corrector is repeatedly applied to the clean prediction at that fixed timestep\. The selected examples use noise levelst∈\{0\.8,0\.9,0\.95,0\.975\}t\\in\\\{0\.8,0\.9,0\.95,0\.975\\\}, where11denotes noise and00denotes data\. The Analytic column uses the same formula \(Eq\.[10](https://arxiv.org/html/2610.14136#A4.E10)\), with a uniform empirical prior over a finite same\-class image bank excluding the original query cohort\. The examples were selected for qualitative illustration and are not a random\-sample quality estimate\.


Figure 34:Additional iterative\-correction examples\.The seven selected examples not included in Fig\.[15](https://arxiv.org/html/2610.14136#S4.F15)\(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](https://arxiv.org/html/2610.14136#S4.F15)\(b\)\. Repeated correction can simplify irregular textures or amplify thin structures into unnatural patterns\.
##### Interpolation of corrected predictions toward noise\.
Figure[35](https://arxiv.org/html/2610.14136#A4.F35)shows all eight selected noisy\-state comparisons\. They uset=0\.9→s=0\.5t=0\.9\\to s=0\.5ort=0\.95→s=0\.55t=0\.95\\to s=0\.55, with the same intervalt−s=0\.4t\-s=0\.4throughout\. For each example and everyKK, we hold the original forward input𝒛t\\bm\{z\}\_\{t\}fixed and construct𝒛s\(K\)=\(s/t\)𝒛t\+\(1−s/t\)𝒙^\(K\)\\bm\{z\}\_\{s\}^\{\(K\)\}=\(s/t\)\\bm\{z\}\_\{t\}\+\(1\-s/t\)\\hat\{\\bm\{x\}\}^\{\(K\)\}\. The real reference is\(1−s\)𝒙\+sϵ\(1\-s\)\\bm\{x\}\+s\\bm\{\\epsilon\}with the same original Gaussian\. Interpolation is applied to unclipped predictions; all cells use the same display range, without contrast normalization\. This is an illustrative interpolation to a lower noise level rather than a single step from the sampling grid\. The corrected predictions need not reconstruct the original image: the comparison concerns the plausibility of visible structure after interpolation, rather than paired reconstruction accuracy\.
Figure 35:Shared\-input noisy\-state comparisons\.Clean images \(top\) and interpolated noisy states \(bottom\), comparing*Real*,K=0K=0, andK=9K=9\. Panels\(a,b,g,h\)uset=0\.95→s=0\.55t=0\.95\\to s=0\.55; panels\(c,d,e,f\)uset=0\.9→s=0\.5t=0\.9\\to 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\.
##### Correction within a denoising step\.
Fig\.[3](https://arxiv.org/html/2610.14136#S2.F3)provides a qualitative view of how correction changes intermediate clean\-sample predictions during sampling\. Without correction, structural errors in an intermediate prediction can persist into subsequent denoising steps because the prediction is immediately used to update the noisy state\. Correction instead refines the prediction at fixedttbefore this update\. In the examples in Fig\.[3](https://arxiv.org/html/2610.14136#S2.F3), successive corrections sharpen object structure and suppress visible artifacts within a single denoising step, so the subsequent sampler state is constructed from an improved prediction\.
To illustrate what correction does at each step, Fig\.[36](https://arxiv.org/html/2610.14136#A4.F36)shows the individual correction passes for three examples fromSiTREPA\\text\{SiT\}\_\{\\mathrm\{REPA\}\}\-XL/2\. For each example, the baseline and correction trajectories start from the same initial noise and use the same 15 Euler denoising steps, withK=8K=8for correction\. In the dog example \(Fig\.[36](https://arxiv.org/html/2610.14136#A4.F36)\), already at early\-to\-middle stages of sampling \(t=0\.73t=0\.73–0\.600\.60\), successive corrections reduce structural artifacts and reveal a coherent dog shape\. These improvements carry through to the final image, while the baseline retains structural artifacts even after completing all denoising steps\.
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 timestepst∈\{0\.93,0\.80,0\.60,0\.07\}t\\in\\\{0\.93,0\.80,0\.60,0\.07\\\}\. Each prediction row shows all eight correction passes at fixedtt; 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\) atK=0,5,9K=0,5,9, starting from the same noise\. Selected JiT\-L/16 samples use 100K checkpoints, EMA weights, no CFG, and 45 Heun steps; largerKKadds correction compute\.
### D\.3Power spectrum analysis
##### Frequency\-domain effects of correction\.
To further characterize the changes induced by correction, we examine the frequency content of generated samples\. Prior work\([Yu & Zhan, 2025](https://arxiv.org/html/2610.14136#bib.bib39);[Kobayashi et al\., 2026](https://arxiv.org/html/2610.14136#bib.bib12)\)has reported lower subband energy in reverse\-process samples than in forward\-corrupted real data\. This motivates examining whether correction recovers frequency components that are underrepresented in baseline samples\.
We compare the power spectra of samples generated withK=0,…,10K=0,\\ldots,10correction iterations against those of real images\. As shown in Fig\.[38](https://arxiv.org/html/2610.14136#A4.F38), the baseline samples exhibit a spectral deficit, particularly at high frequencies\. Increasing the number of correction iterations reduces this deficit, withK=6K=6yielding a spectrum that approximately matches the real\-image reference\. Further iterations can raise both low\- and high\-frequency power above the reference, indicating that the effect depends on the correction count\.
Figure 38:Correction reduces the high\-frequency deficit of generated samples\.Left:log10P\\log\_\{10\}P\. Right:log10\(P/Preal\)\\log\_\{10\}\(P/P\_\{\\mathrm\{real\}\}\)\. Dashed lines indicate real\-image references\.
## Appendix EAdditional Samples
### E\.1Effect of Correction Depth
We show the same class and initial noise as correction depth increases, followed by paired galleries over randomly drawn ImageNet classes\. All examples use class conditioning without CFG\.
Figure 39:Uncuratedsamples: correction depth changes both structure and detail\.Each strip follows one sample asKKincreases, with the initial noise and denoising grid fixed\.
### E\.2SiT: Random\-Class Generation
Figure 40:Uncuratedsamples: correction refines unguided SiT samples across random classes\.Within each pair of rows, baseline \(K=0K=0\) and correction \(K=6K=6\) share the same class and initial noise at approximately matched NFEs\.
### E\.3SiTREPA: Random\-Class Generation
Figure 41:Uncuratedsamples: correction refines unguided SiTREPAsamples across random classes\.Within each pair of rows, baseline \(K=0K=0\) and correction \(K=6K=6\) share the same class and initial noise at approximately matched NFEs\.
### E\.4JiT: Random\-Class Generation
Figure 42:Uncuratedsamples: early correction changes structure as well as appearance in JiT\.Within each pair of rows, baseline \(K=0K=0\) and correction \(K=6K=6\) share the same class and initial noise at matched NFEs\.
### E\.5Unconditional Generation
Figure 43:High\-quality unconditional generation without guidance\.Selected ImageNet samples from SiTREPA\-XL/2 with iterative correction\.
### E\.6Unconditional Generation from a Text\-to\-Image Model
Figure 44:Uncuratedsamples: unconditional generation supports iterative correction without text or CFG\.Each row shares its initial noise acrossK=0,…,5K=0,\\ldots,5\. The checkpoint and denoising steps are fixed; largerKKadds correction compute\. All seeds are shown without selection\.Figure 45:Uncuratedsamples: unconditional generation supports iterative correction without text or CFG\.Each row shares its initial noise acrossK=0,…,5K=0,\\ldots,5\. The checkpoint and denoising steps are fixed; largerKKadds correction compute\. All seeds are shown without selection\.Figure 46:Uncuratedsamples: unconditional generation supports iterative correction without text or CFG\.Each row shares its initial noise acrossK=0,…,5K=0,\\ldots,5\. The checkpoint and denoising steps are fixed; largerKKadds correction compute\. All seeds are shown without selection\.Figure 47:Uncuratedsamples: unconditional generation supports iterative correction without text or CFG\.Each row shares its initial noise acrossK=0,…,5K=0,\\ldots,5\. The checkpoint and denoising steps are fixed; largerKKadds correction compute\. All seeds are shown without selection\.相似文章
扩散模型的校正训练
本文提出了一种用于扩散模型的校正训练方法,通过训练校正模块来细化去噪预测、缩小训练与推理之间的差距,从而在无引导条件下显著提升样本质量,并可通过蒸馏将改进效果整合回基础模型,避免额外的推理开销。
SALD:扩散模型的自参照优势学习
SALD 提出了一种自参考训练框架,通过在两个噪声水平下用同一模型评估图像-文本对来获取训练信号,无需外部教师网络或额外可训练参数,包含优势引导扩散(AGD)、时间优势记忆(TAM)和谱优势分解(SAD)三个组件,在多个架构和数据集上持续提升了生成质量。
CoDR:面向扩散语言模型的免训练置信度漂移重遮蔽方法
论文提出 CoDR(Confidence-Drift Remasking),一种免训练、与采样器无关的后处理精炼方法,通过检测已提交 token 的置信度漂移并对其进行重遮蔽与重新生成,从而在掩码扩散语言模型的推理与编码任务上以较少的前向传播次数提升准确率。
DiffusionART:高效、稳定且可扩展的扩散模型强化方法
本文提出DiffusionART,一种通过聚合奖励方向进行扩散模型强化学习微调的新方法,消除了一致性力并用信任区域掩码实现免EMA调参的稳定训练,在HunyuanVideo-13B上以约12倍更少的GPU小时数取得显著更高的训练奖励。
通过在线上下文蒸馏将智能体经验内化到扩散模型权重中
论文提出 D-OPCD(Diffusion On-Policy Context Distillation),将文本生成图像智能体框架(harness)通过记忆、技能和迭代优化获得的增益蒸馏进扩散模型权重中,使模型在仅使用原始查询时也能保留部分收益,平均直接生成分数从 60.52 提升至 65.09,并通过 AutoSkillEvolver 实现框架与模型的持续协同进化。