Correction Training for Diffusion Models

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

한 줄 요약

확산 모델의 예측 오류 누적 문제를 해결하기 위해 Corrector 모듈을 도입하여 샘플 품질을 향상시킨다.

핵심 기여도

Corrector 모듈을 도입하여 기존 Denoiser의 예측을 반복적으로 보정함으로써 샘플 품질을 개선.
DiT, SiT, JiT 아키텍처에서 FID 점수를 기존 보고치보다 크게 개선.
Corrector를 CFG와 결합하면 다변화를 향상시키고 과포화 현상을 감소.
Corrector의 예측을 Base 모델로 Distillation하여 추론 비용 증가 없이 품질 개선 효과 유지.

핵심 아이디어

기존 확산 모델에서 Denoiser는 노이즈가 섞인 실제 이미지로 훈련되지만, 추론 시에는 이전 단계의 불완전한 예측에 기반한 입력을 받는다. 이는 Training-Inference Gap을 유발하며, 오류가 누적되어 최종 샘플에 예측 불가능한 아티팩트를 남기는 원인이다. 이를 해결하기 위해, 본 연구는 Denoiser의 예측을 입력으로 받아 동일한 Ground-Truth를 대상으로 훈련하는 Corrector 모듈을 도입한다. Corrector는 단일 단계에서 훈련되지만, 추론 시에는 반복적으로 자기 자신의 출력을 보정할 수 있어 샘플 품질을 향상시킨다. 이는 Denoiser의 오류 누적을 보완하고, Training-Inference Gap을 줄이는 핵심 아이디어이다.

기술적 접근법

**Corrector 모듈**: Denoiser의 예측을 입력으로 받아, 동일한 Ground-Truth를 대상으로 훈련.
**Iterative Correction**: Denoising 단계에서 Denoiser의 예측을 반복적으로 보정.
**Distillation**: Corrector의 반복 예측을 Base 모델로 학습시켜 추론 비용 없이 품질 유지.
**모델 아키텍처**: DiT, SiT, JiT에서 실험.
**Metric**: FID, Recall, NFEs (Number of Function Evaluations) 사용.

주요 결과

**DiT, SiT, JiT 모델에서 FID 점수 개선**: Unconditional generation에서 기존 보고치 대비 FID 점수 크게 개선.
**CFG 대비 성능**: Corrector는 CFG와 유사한 품질을 제공하면서도 더 높은 다양성을 유지.
**CFG + Corrector 결합**: CFG의 과포화 문제를 감소시키며, 샘플의 시각적 현실감을 향상.
**Distillation 효과**: Corrector의 예측을 Base 모델로 학습시켜 추론 비용 없이 품질 개선 유지.

의의 및 한계

**의의**: Training-Inference Gap 문제를 해결하는 새로운 학습 프레임워크 제시. 기존 CFG 대비 샘플 다양성과 품질을 동시에 개선.
**한계**: Corrector의 반복적 보정 효과에 대한 이론적 해석은 미비.
**추론 효율성**: Distillation을 통해 Corrector의 추가 비용을 제거할 수 있으나, 초기 훈련 비용은 증가.

실용적 활용

**Unconditional Image Generation**: 고품질 샘플 생성이 필요한 시각 콘텐츠 생성 분야에 적용 가능.
**Text-to-Image Generation**: 텍스트 기반 이미지 생성에서 샘플 다양성과 현실감을 동시에 향상.
**Efficient Inference**: Distillation을 통해 품질 향상 효과를 유지하면서도 추론 비용 증가 없이 활용 가능.