불확실성 인식 확산 사전 기반 이미지 복원을 위한 효율적 변분 추론

초록

U‑DA​VI는 사전 학습된 확산 모델을 이용한 변분 추론을 한 단계로 수행하면서, 훈련 시에 시간적 불일치 기반의 경량 불확실성 추정과 공간적 적응 잡음을 결합해 불확실한 영역을 집중 학습한다. 이를 통해 블러 제거와 4배 초해상도 복원에서 기존 DA‑VI보다 PSNR는 비슷하거나 약간 상승하고, FID는 현저히 낮아 시각적 품질이 크게 향상된다.

상세 분석

본 논문은 확산 기반 사전(prior)을 활용한 이미지 복원 문제에서 두 가지 주요 한계를 동시에 해결하고자 한다. 첫 번째는 기존 확산‑기반 변분 추론(DA‑VI)이 높은 연산 비용 없이 단일 전방 패스로 후방 분포 샘플을 생성할 수는 있지만, 불확실성이 큰 영역—예를 들어 경계, 미세 텍스처, 눈동자 주변—에서 세밀한 복원이 부족하다는 점이다. 두 번째는 불확실성 추정을 위한 전통적 베이지안 방법(딥 앙상블, MC‑Dropout 등)이 훈련 단계마다 다중 순전파를 요구해 실용성이 떨어진다.

U‑DA​VI는 이러한 문제를 해결하기 위해 ‘시간적 불일치’를 이용한 경량 불확실성 지도(u)를 도입한다. 각 훈련 샘플마다 영구적인 재구성 메모리(EMA) (\bar{x})를 유지하고, 현재 생성된 이미지 (\hat{x})와의 L1 차이를 정규화해 픽셀별 불확실성을 계산한다. EMA 업데이트는 (\eta = 2/(N+1)) 로 설정해 메모리 윈도우 N을 효과적으로 제어한다. 이 과정은 기존 전방 패스에 거의 비용을 추가하지 않으며, 불확실성 지도는 매 iteration마다 갱신돼 학습이 진행될수록 안정된 영역은 낮은 값, 지속적으로 변동하는 영역은 높은 값을 유지한다.

불확실성 지도는 ‘Perturbed Posterior Bridge(PPB)’에 직접 반영된다. 기존 PPB는 측정 y와 원본 x₀ 사이를 가우시안 잡음으로 연결했지만, U‑DA​VI는 잡음에 ((1+\lambda u)) 를 곱해 불확실한 픽셀에 더 큰 노이즈를 주입한다. 이렇게 하면 확산 과정에서 해당 픽셀의 중간 상태 (\hat{x}t)가 더 큰 변동성을 보이며, 교사‑학생 스코어 매칭 손실 (|s\psi - s_\theta|^2) 에서 큰 그래디언트가 발생한다. 결과적으로 학습이 불확실한 영역에 집중돼 세밀한 디테일 복원이 강화된다.

또한, 이 불확실성‑가이드가 훈련 단계에만 적용되므로 추론 시에는 기존 DA‑VI와 동일하게 단일 전방 패스 (\hat{x}=I_\phi(y+hz)) 로 수행된다. 따라서 NFE(함수 평가 횟수)는 1에 머물러 실시간 응용이 가능하다.

실험에서는 FFHQ(256×256) 데이터셋을 사용해 Gaussian 블러(61×61, σ=3.0)와 4배 초해상도 두 작업을 평가했다. U‑DA​VI는 DA‑VI 대비 PSNR를 0.10.2 dB 정도 상승시키고, FID는 0.81.2 포인트 낮추어 시각적 품질이 크게 개선되었다. 특히, zero‑shot 일반화 테스트에서 CelebA‑HQ에 대해 동일한 트렌드가 유지돼 불확실성‑가이드가 데이터셋 간 전이에도 강인함을 보여준다. λ 하이퍼파라미터에 대한 민감도 분석에서는 λ=1.0(블러)·0.5(초해상도)가 가장 좋은 PSNR‑FID 균형을 제공한다는 점을 확인했다.

요약하면, U‑DA​VI는 (1) 비용 효율적인 불확실성 추정, (2) 불확실성 기반 적응 잡음 주입, (3) 기존 DA‑VI와 동일한 단일‑스텝 추론 구조를 유지한다는 세 가지 혁신을 통해 확산‑사전 기반 이미지 복원에서 정량·정성 모두를 향상시킨다.