확산 확률 모델의 SNR-t 편향 규명

Diffusion 모델이 추론 시 훈련 때와 달리 신호 대 잡음비(SNR)와 타임스텝 간의 불일치를 겪는 SNR-t 편향 문제를 최초로 규명했다. 추가 학습 없이도 생성 품질을 획기적으로 개선할 수 있는 플러그앤플레이 방식의 보정 기법을 제시하여 실무적 가치가 높다.

Diffusion 모델은 훈련 시 특정 타임스텝 t에 정확히 대응하는 노이즈 수준(SNR)을 학습한다. 하지만 실제 생성(추론) 단계에서는 이전 단계의 예측 오류와 수치적 오차가 누적되면서, 현재 타임스텝 t라고 가정된 샘플의 실제 SNR이 훈련 때 보았던 수준보다 낮아지는 현상이 발생한다. 이는 모델이 훈련 데이터 분포에서 벗어난 샘플을 처리하게 만들어 노이즈를 과도하게 예측하는 악순환을 초래한다.

이 논문은 모델이 생성 초기에는 이미지의 전체적인 윤곽(저주파)을 잡고 나중에 세부 묘사(고주파)에 집중한다는 특성에 주목한다. 이미지를 주파수 영역으로 변환한 뒤, 현재 예측된 샘플과 모델이 재구성한 샘플 사이의 차이 신호를 활용해 샘플을 이상적인 궤적으로 밀어주는 가이드를 제공한다. 특히 생성 단계에 따라 저주파와 고주파 보정 강도를 다르게 조절함으로써 이미지의 구조적 일관성과 세부 디테일을 동시에 확보한다.

전체적인 접근 방식은 추론 단계에서 발생하는 SNR-t 편향을 줄이기 위해 예측 샘플을 이상적인 SNR 분포로 이동시키는 차분 보정(Differential Correction)을 수행한다. 별도의 재학습 없이 추론 루프 내에서 모델의 자체 출력값을 활용해 보정 신호를 생성한다.

핵심 메커니즘인 DCW는 이산 웨이브렛 변환(DWT)을 사용하여 샘플을 네 개의 주파수 서브밴드(LL, LH, HL, HH)로 분해한다. 각 성분 f에 대해 보정 계수 λ를 곱한 차분 신호를 더해주는 연산 [x_hat_f + λ * (x_hat_f - x_theta_f) → 보정된 x_f]를 수행한다. 여기서 x_theta_f는 모델이 예측한 원본 이미지의 주파수 성분이며, 이 연산은 샘플을 더 정확한 SNR 궤적으로 유도하는 역할을 한다.

구현 상세 측면에서 보정 계수 λ는 타임스텝에 따라 동적으로 변한다. 저주파 성분(LL)에는 생성 초기에 큰 가중치를 주어 구조를 잡고, 고주파 성분(LH, HL, HH)에는 생성 후기에 가중치를 높여 디테일을 복구한다. 최종적으로 역 이산 웨이브렛 변환(iDWT)을 통해 픽셀 공간으로 복원하며, 이 과정은 추가적인 신경망 연산(NFE)을 요구하지 않아 효율적이다.

IDDPM 모델을 사용한 CIFAR-10 실험에서 20단계 샘플링 기준 FID 점수를 13.19에서 7.57로 약 42.6% 개선했다. ADM 모델을 이용한 ImageNet 128x128 실험에서도 FID를 12.28에서 10.34로 낮추며 생성 품질 향상을 입증했다.

최신 모델인 FLUX에 적용했을 때 10단계의 적은 샘플링 횟수에서도 기존 모델 대비 과도한 평활화(over-smoothing)나 노출 부족 문제를 해결하고 훨씬 선명한 이미지를 생성함을 확인했다. 또한 NVIDIA A6000 GPU 기준 시간 오버헤드는 0.08%~0.47% 수준으로 매우 미미하여 실시간 생성 시스템에 즉시 적용 가능하다.

SNR-t 편향은 훈련 시 SNR(t) = α_bar_t / (1 - α_bar_t)로 고정된 관계가 추론 시의 누적 오차로 인해 붕괴되면서 발생한다. 저자들은 Tweedie 공식을 활용하여 모델의 재구성 출력 x_theta_0가 실제 데이터 x0의 평균 추정치임을 보이고, 추론 시 샘플의 L2 노름이 항상 훈련 시보다 크게 나타나 결과적으로 SNR이 낮아짐을 수학적으로 증명했다.

제안된 DCW는 주파수 도메인에서의 보정이 픽셀 도메인보다 노이즈 간섭에 강인하다는 점을 이용한다. 보정 신호는 현재 단계의 예측값과 모델이 추정한 원본 이미지 사이의 차이에서 유도되며, 이는 샘플을 이상적인 분포로 밀어주는 그래디언트 가이드 역할을 한다. 특히 역과정 분산(σ_t)을 보정 계수의 지표로 활용하여 별도의 하이퍼파라미터 튜닝 없이도 안정적인 성능을 낸다.

본 논문은 Diffusion 모델의 내부적인 SNR-t 편향에 집중하고 있으며, 텍스트-이미지 정렬(Text-Image Alignment)과 같은 외부적인 조건부 제어 오차에 대해서는 직접적인 해결책을 제시하지 않는다.