확산 모델의 노이즈 제거 예측을 반복 보정해 품질을 높인다

무엇인가

이 논문은 확산 모델의 반복 샘플링에서 생기는 노출 편향(exposure bias)을 다룬다. 학습 때 노이즈 제거기는 실제 이미지에 노이즈를 선형 보간한 입력을 보지만, 추론 때는 이전 단계에서 자신이 만든 불완전한 예측으로 다음 입력이 구성된다. 이 불일치가 노이즈가 줄어드는 경로를 따라 오차로 누적되어 구조가 무너지거나 아티팩트가 남는 문제로 이어지고, 저자들은 CFG 같은 추가 장치 없이도 품질이 흔들리는 이유로 본다.

어떻게 동작하나

제안 방법의 핵심은 베이스 노이즈 제거기 Dθ의 출력 x̂(0)=Dθ(z_t,t)를 입력으로 받는 별도 보정기 Cφ를 학습하는 것이다. 보정기는 같은 정답 이미지 x를 목표로 x̂(0)→x 매핑을 배우며, 손실은 ||Cφ(sg(x̂(0)),t)-x||²이다. sg는 stop-gradient로 보정기 학습이 베이스 예측 쪽으로 역전파되지 않게 한다. 사전학습된 노이즈 제거기에서 시작해 L_total=L_denoise+λ_corr L_corr를 공동 최적화하고, 보정기 입력은 현재 모델이 온라인으로 생성한 예측을 쓴다. 즉 보정을 별도 사후처리가 아니라 명시적 학습 과제로 만든다.

무엇과 다른가

추론에서는 각 디노이징 단계마다 보정을 K번 반복한다. 먼저 x̂(0)=Dθ(z_t,t)를 얻고, x̂(k+1)=Cφ(x̂(k),t)를 k=0,…,K-1에 적용한 뒤, 최종 보정 예측으로 다음 노이즈 상태 z_s를 만든다. 학습은 K_train=1, 즉 한 번의 보정만 감독하지만, 추론에서 더 많이 반복하는 외삽이 품질 향상의 핵심이라고 저자들은 말한다. 총 NFE를 맞추기 위해 K가 커질수록 디노이징 스텝 수를 줄여 비교하며, 같은 추론 비용이면 디노이징 스텝만 늘리는 것보다 보정 반복이 더 효과적이라고 주장한다.

어떻게 쓰나

ImageNet 클래스 조건부 생성에서 DiT는 FID 9.58→2.07, SiT는 8.55→1.92, SiT REPA는 5.53→1.31로 같은 NFE에서 크게 개선됐다. SiT REPA의 추가 학습은 50K 업데이트로 사전학습 4M 업데이트의 1.25%에 불과했다. 무조건부 생성에서는 K=4 보정이 베이스라인과 autoguidance 최고 스케일(w=1…10)보다 FID를 크게 개선했다고 보고하지만, 원문에 구체 수치는 제시되지 않았다. 텍스트-이미지에서는 NanoGen 프레임워크의 DiT+DDT 헤드, FLUX.2 VAE 잠재 공간, JourneyDB·Short/Long 캡션, 256×256 설정에서 50K 스텝 공동 미세조정 후 대략 같은 NFE에서 GenEval과 DPG-Bench 점수가 올랐고, 80개 GenEval 단일 객체 프롬프트 인간 평가에서 CFG+보정이 기본 CFG 베이스라인보다 60% 이상 선호됐다. 저자들은 보정이 CFG와 경쟁할 만한 품질을 내면서 다양성을 더 보존하고, CFG와 결합하면 CFG 특유의 과포화와 스타일화를 줄인다고 주장한다.

전제와 한계

보정은 추가 순전파가 필요하지만 두 방식으로 비용을 줄인다. 먼저 SiT-XL/2에서 백본을 공유하고 입력 임베더만 분리하면, 분리 네트워크와 비슷한 성능을 약 절반 파라미터로 낸다(Tab. 4). 다음으로 보정된 예측을 베이스 모델로 증류한다. 학생은 교사의 k*번 보정 후 예측을 맞추도록 학습하며, k*=2에서 비슷한 FID를 유지하면서 NFE를 약 67% 줄였다. 손실은 w_distill(t)=1/t²를 쓰는 속도 공간 MSE다. 또한 D28(675M)과 D14+C14(681M)를 같은 80K 업데이트로 비교했을 때 K=1은 FID 8.74 대 8.93으로 비슷했지만, K=3에서는 디노이징 스텝을 절반으로 줄여 총 레이어 평가 예산을 맞춘 조건에서 FID 5.66으로 낮아졌고, D7+C7 K=3은 D28보다 낮은 FID 4.21 대 4.88을 같은 디노이징 스텝, 약 절반 파라미터로 달성했다.

저자들은 보정이 왜 효과적인지도 분석한다. 반복 보정은 예측을 더 선명하고 일관되게 만들지만, k≥3부터는 학습 분포에서 측정한 디노이징 오차가 오히려 커진다. 그럼에도 품질이 좋아지는 이유로, 보정된 예측이 다음 노이즈 상태를 학습 분포에 더 가깝게 만들어 학습-추론 간 불일치를 줄인다는 가설을 세운다. 실제로 SiT의 중간 상태 100K개로 두 표본 분류기를 학습해 BCE/AUROC를 재면, 한 번의 보정이 BCE를 높이고 AUROC를 낮춰 두 분포를 구분하기 어렵게 만들며 반복 보정은 그 격차를 더 줄인다. 또한 K_train을 늘리면 최적 추론 K가 커지고, 최고 성능은 학습 때 감독한 횟수보다 더 많이 보정할 때 나온다. 노이즈 제거기를 고정하고 보정기만 학습하면 K=5까지는 비슷하지만 더 반복하면 FID가 다시 올라가고, 공동 학습은 계속 개선된다.

개발자 입장에서는 이미 학습된 diffusion 체크포인트에 작은 추가 학습으로 보정기를 붙여 CFG 없이 품질을 올리거나, CFG와 결합해 과포화와 다양성 손실을 완화하는 선택지로 쓸 수 있다. 추론 예산을 디노이징 스텝과 보정 반복 사이에 재분배할 수 있고, 배포에서는 증류로 추가 순전파를 없앨 수 있다. 다만 저자들도 보정이 학습 분포 기준 재구성 오차를 키울 수 있고, 반복이 지나치면 피자 토핑이나 딸기 씨처럼 과도하게 규칙적인 구조를 만들 수 있다고 밝힌다. 또한 한 번의 보정만 학습한 뒤 여러 번 적용하는 것은 외삽이며, 그 효과와 반복 동작에 대한 이론적 이해는 열린 문제로 남겨져 있다.