Correction Training for Diffusion Models

Computer Science > Artificial Intelligence arXiv:2610.14136 (cs) [Submitted on 9 Oct 2026 (v1), last revised 14 Oct 2026 (this version, v2)] Title: Correction Training for Diffusion Models Authors: Donghoon Ahn, Runqian Wang, Alexei A. Efros, Taesung Park, Yash Belhe.

计算机科学 > 人工智能 arXiv:2610.14136 (cs) [提交于 2026 年 10 月 9 日 (v1)，最后修订于 2026 年 10 月 14 日 (此版本，v2)] 标题：扩散模型的校正训练 作者：Donghoon Ahn, Runqian Wang, Alexei A. Efros, Taesung Park, Yash Belhe。

Abstract: Iterative sampling in diffusion models can lead to error accumulation due to exposure bias. The denoiser is trained on real images interpolated with noise, but at inference its inputs are constructed from its own imperfect predictions from the earlier steps, something it has never seen at training time. This might be one reason diffusion models often struggle without additional help, such as CFG.

摘要：扩散模型中的迭代采样可能会因暴露偏差导致误差累积。去噪器是在插入噪声的真实图像上进行训练的，但在推理时，其输入是由之前步骤中自身不完美的预测构建而成的，这是它在训练时从未见过的情况。这可能是扩散模型在没有额外帮助（如 CFG）的情况下往往表现不佳的原因之一。

We introduce a correction step for diffusion models which attempts to reduce this training-inference gap. The corrector takes the base model’s prediction as input and is trained against the same ground-truth target. This makes correction an explicit learning task, with training inputs provided by the base model. During sampling, at each denoising step, we iteratively refine the denoiser’s clean-sample prediction. This allows us to correct predictions at intermediate denoising steps, preventing early artifacts from persisting into the final sample.

我们为扩散模型引入了一个校正步骤，旨在缩小这种训练与推理之间的差距。校正器以基础模型的预测作为输入，并针对相同的真实目标进行训练。这使得校正成为一项明确的学习任务，其训练输入由基础模型提供。在采样过程中，我们在每个去噪步骤迭代优化去噪器对清晰样本的预测。这使我们能够在中间去噪步骤校正预测，防止早期的伪影持续存在于最终样本中。

Although the corrector is trained only on predictions from the base denoiser, it can be applied repeatedly to its own outputs at inference for better results. This simple framework leads to substantial improvement in sample quality while largely preserving diversity. We evaluate our method across DiT, SiT, and JiT. Our method substantially improves upon previously reported FID scores for unguided generation and, notably, unconditional generation, where producing high-quality samples without guidance has been particularly challenging.

尽管校正器仅在基础去噪器的预测上进行训练，但它可以在推理时重复应用于其自身的输出，以获得更好的结果。这个简单的框架在很大程度上保持多样性的同时，显著提高了样本质量。我们在 DiT、SiT 和 JiT 上评估了我们的方法。我们的方法显著改善了先前报道的无引导生成以及无条件生成的 FID 分数，值得注意的是，在无引导的情况下生成高质量样本一直是一个特别具有挑战性的问题。

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.

我们的方法实现了与无分类器引导 (CFG) 相媲美的性能，并且将校正与 CFG 相结合进一步提高了多样性，并减少了 CFG 常见的过度饱和现象。最后，我们证明了通过将优化后的预测蒸馏回基础扩散模型，可以分摊校正带来的额外推理成本，从而在没有额外推理时间开销的情况下保持质量提升。