원문 캡처 · sunovivid.github.io
원문 캡처 · sunovivid.github.io

확산 모델의 예측을 같은 깨끗한 이미지 목표에 맞춰 반복 보정하는 학습법이 제안됐다. 연구진은 이 방법이 네 가지 모델의 무가이드 생성 품질을 높였고, 텍스트-이미지 생성에서는 분류기 없는 안내(classifier-free guidance, CFG)와 함께 사용할 때 다양성과 사실감을 개선했다고 보고했다.

확산 모델은 학습할 때 실제 이미지에 노이즈를 섞은 입력으로부터 깨끗한 이미지를 예측한다. 하지만 생성 과정에서는 모델이 앞선 단계에서 만든 예측을 다음 입력으로 사용한다. 예측 오차가 누적되면 학습 때 보지 못한 상태로 입력이 벗어날 수 있고, 초기에 생긴 구조적 오류가 후속 단계에서도 남을 수 있다는 설명이다.

제안한 방법은 기본 노이즈 제거 모델의 예측을 보정하는 모델을 별도로 학습한다. 보정 모델은 기본 모델과 같은 깨끗한 이미지를 목표로 예측을 다듬는다. 생성할 때는 노이즈 제거 단계마다 같은 노이즈 수준에서 예측을 여러 번 보정한 뒤 다음 단계로 넘어간다. 예시 설정에서는 보정을 8회 적용했다.

연구진은 DiT, SiT, SiT-REPA, JiT를 대상으로 무가이드 기준 모델보다 FID를 개선했다고 밝혔다. ImageNet 256×256 클래스 조건부 평가에서는 5만 개 샘플을 사용했고, 보정 방식의 네트워크 평가 횟수는 CFG와 같거나 더 적게 맞췄다. 다만 모델별 설정은 논문을 참조해야 한다.

CFG와 비교한 결과에서 연구진은 보정 방식이 생성물의 다양성을 더 잘 유지하고 과도한 채도를 줄이는 사례를 제시했다. 보정과 CFG를 함께 적용하는 것도 가능하다. 텍스트-이미지 비교에서는 같은 8개 시드와 50회의 노이즈 제거 단계를 사용했으며, 보정 적용 시 벤치의 기본 CFG 설정만 사용한 기준보다 더 다양한 시점과 사실적인 이미지가 나왔다고 보고했다.

사람 평가에는 GenEval의 단일 객체 프롬프트 80개를 사용했다. 각 기준에 프롬프트당 10회씩, 총 1,600건의 판단을 Rapidata를 통해 수집했다. CFG와 보정을 함께 적용한 결과는 품질 평가에서 60.75%, 다양성 평가에서 60.25%를 얻었고, CFG만 적용한 결과는 각각 39.25%, 39.75%였다.

연구진은 반복 보정의 효과가 단순히 매개변수나 계산량을 늘린 결과가 아니라고 주장한다. LightningDiT를 ImageNet 256×256에서 처음부터 8만 회 업데이트해 학습한 비교에서는 28층 노이즈 제거 모델의 FID가 8.93이었다. 같은 총 층 수를 14층 노이즈 제거 모델과 14층 보정 모델로 나누고 단계마다 한 번 보정했을 때는 FID 8.74로 비슷했다. 보정을 세 번 반복하고 노이즈 제거 단계를 절반인 120회로 줄여 총 계산량을 같게 했을 때는 FID가 5.66으로 낮아졌다. 이 평가는 8천 개 샘플, 무가이드 조건에서 이뤄졌다.

연구진은 보정이 학습 입력과 생성 중 입력 사이의 차이를 줄여 성능을 높일 수 있다고 가정한다. 실제 이미지에 노이즈를 더한 학습 상태와 모델의 이전 예측으로 만든 생성 상태를 구분하도록 분류기를 학습한 실험에서, 보정을 반복할수록 두 상태를 구분하기 어려워졌다. 한편 학습 입력에서 측정한 제곱 오차는 보정을 세 번 이상 적용하면 기본 모델보다 높았다. 따라서 이 설명은 생성 중 상태에 대한 분석에 바탕을 둔 연구진의 가설이며, 학습 입력에서의 오차 감소를 뜻하지 않는다.

텍스트-이미지 실험에는 ImageNet과 텍스트-이미지 생성을 지원하는 구현인 NanoGen을 사용했다. 기반 모델은 BLIP-3o의 Short-Caption, Long-Caption, JourneyDB 하위 데이터셋으로 사전 학습됐으며, 해당 데이터셋의 원래 공개본에는 이미지-텍스트 쌍이 약 3,600만 개 포함됐다.

부록: 원문 예시

text
@misc{ahn2026correction,   title  = {Correction Training for Diffusion Models},   author = {Ahn, Donghoon and Wang, Runqian and Efros, Alexei A. and Park, Taesung and Belhe, Yash},   year   = {2026} }