기술 요약: 결정론적 요소와 확률적 요소의 상호 학습을 통한 확산 모델 기반 이미지 복원

문제 정의

이미지 복원(Image Restoration, IR)은 열화된 관측값으로부터 잠재적인 시각 정보를 재구성하는 것을 목표로 하며, 이는 서로 다른 유형의 열화(예: 노이즈, 비, 그림자, 저조도) 간의 특징 공간 내 상당한 분포 차이로 인해 복잡한 작업입니다. CNN, 트랜스포머, GAN을 포함한 기존의 딥러닝 방법들은 태스크별 재설계 없이는 다양한 IR 태스크에 걸쳐 일반화하는 데 어려움을 겪는 경우가 많습니다. 확산 모델(Diffusion Models)은 높은 충실도와 해석 가능성을 제공하지만, 현재의 접근 방식들은 두 가지 주요 한계점에 직면해 있습니다:

1. IR에서의 비효율성: 표준 확산 모델은 순수 노이즈로부터 역과정을 시작하는데, 이는 이미 구조적 정보를 포함하고 있는 IR 태스크에서 비효율적입니다.
2. 결정론적 요소와 확률적 요소의 정적 모델링: 결정론적(잔차 기반) 및 확률적(노이즈 기반) 확산을 통합하려는 이전의 시도들(예: RDDM)은 각 경로를 위해 독립적인 U-Net을 사용합니다. 이러한 정적 분리는 서로 다른 열화 유형의 고유한 확률 분포를 포착하지 못하며, 정밀한 재구성을 방해하는 과도한 다양성을 유발할 수 있습니다.

방법론

저자들은 이러한 한계를 해결하기 위해 **결정론적 요소와 확률적 요소의 상호 학습(Interactive Learning of Determinism and Stochasticity, ILDS)**을 채택한 새로운 확산 프레임워크를 제안합니다. 핵심 아키텍처는 두 개의 단순한 U-Net을 기반으로 하는 이중 경로 네트워크를 활용합니다.

1. 순방향 확산 과정 구축

노이즈만을 추가하는 표준 확산과 달리, 제안된 순방향 과정은 잔차(residuals)(결정론적 요소)와 노이즈(확률적 요소)를 모두 포함합니다. 전이 과정은 다음과 같이 정의됩니다:
$$x_t = x_{t-1} + \alpha_t x_{res} + \beta_t \epsilon_{t-1}$$
여기서 $x_{res} = x_{in} - x_0$는 열화 잔차를 나타내며, $\epsilon$은 가우시안 노이즈입니다. 이를 통해 순방향 궤적은 열화 관련 구조를 보존할 수 있으며, 역과정이 순수 노이즈보다 목표 상태에 더 가까운 상태에서 시작할 수 있게 하여 필요한 샘플링 단계를 줄여줍니다.

2. 상호 학습 메커니즘

독립적인 모델링의 한계를 극려하기 위해, 저자들은 U-Net 아키텍처 내에 상호 학습 전략을 도입했습니다:

• 이중 경로 예측: 두 개의 U-Net이 각각 노이즈($\epsilon_\theta$)와 잔차($x_{res}^\theta$)를 예측합니다.
• 특징 융합(Feature Fusion): 두 경로가 고립되어 작동하는 대신, 두 U-Net의 인코더는 특징을 공유합니다. 구체적으로, 첫 번째 다운샘플링 블록 이후, 양쪽 경로의 특징이 결합되어 $1 \times 1$ 컨볼루션($\tilde{h}_{res}$ 및 $\tilde{h}_\epsilon$)을 통해 처리됩니다.
• 동적 시간 의존적 모델링: 두 경로의 타임스텝 입력은 시간적 축적($\bar{\alpha}_t$)과 노이즈 강도($\bar{\beta}_t$)를 제어하는 파라미터에 의해 동적으로 조절됩니다. 이를 통해 모델은 열화 유형의 특정 특징 확률 분포에 따라 결정론적 성분과 확률적 성분의 기여도를 적응적으로 조정할 수 있습니다.
• 공유 컨텍스트: 상호작용된 특징들은 디코딩 단계 동안 교차 경로 간의 공유된 컨텍스트 정보를 제공하여, 결정론적 가이드와 확률적 다양성이 서로 영향을 미치도록 보장합니다.

3. 훈련 및 생성

모델은 실제 사후 분포와 학습된 역전이 사이의 쿨백-라이블러(KL) 발산을 최소화하도록 훈련됩니다. 최종 손실 함수는 노이즈 및 잔차 예측에 대한 제곱 오차의 합입니다:
$$L = L_{res}(\theta) + L_{\epsilon}(\theta)$$
특히, 상호작용 메커니즘이 훈련 중 자기 적응적 평형을 촉진하므로 이 구성 요소들은 가중치가 부여되지 않습니다. 추론 시, 모델은 조건부 분포로부터 반복적으로 샘플링하여 고품질 이미지를 생성하며, 노이즈 스케일 인자 $\eta = 0$일 때 과정은 완전히 결정론적으로 변합니다.

주요 기여

1. 새로운 확산 프레임워크: 이미지 열화의 결정론적 측면과 확률적 측면을 상호작용적으로 모델링하여, 모델이 높은 충실도(정밀도) 또는 다양성을 요구하는 태스크에 적응할 수 있도록 하는 프레임워크를 설계했습니다.
2. 동적 시간 의전적 모델링: 이중 확산 경로를 위한 타임스텝 입력을 동적으로 조절하기 위해 시간적 축적 및 노이즈 강도 파라미터를 도입하여, 다양한 IR 시나리오에 대한 일반화 능력을 향상시켰습니다.
3. 효율적인 이중 경로 아키텍처: 공유 인코더 특징을 가진 두 개의 단순한 U-Net만을 사용하여 이중 경로 확산을 구현했습니다. 이 설계는 최소한의 파라미터 오버헤드로 경쟁력 있는 성능을 달성하며 높은 확장성을 가집니다.

실험 결과

제안된 방법은 네 가지 대표적인 IR 태스크인 **이미지 디노이징(Denoising), 그림자 제거(Shadow Removal), 저조도 개선(Low-Light Enhancement), 이미지 디레인(Deraining)**에 대해 평가되었습니다.

• 성능: 본 방법은 모든 태스크에서 우수하거나 경쟁력 있는 성능을 달old했습니다.
  • 디노이징 (SIDD): 최고의 PSNR (41.67 dB) 및 SSIM (0.9385)을 달성하여 RDDM 및 기타 베이스라인을 능가했습니다.
  • 그림자 제거 (ISTD): 최고의 PSNR (29.29 dB) 및 MAE (8.776)를 확보했으며, SSIM은 두 번째로 높은 수치를 기록했습니다.
  • 디레인 (Raindrop): 최고의 PSNR (31.30 dB) 및 SSIM (0.9460)을 달성하였으며, 다른 방법들이 아티팩트를 남기는 곳에서도 비 줄기를 성공적으로 제거했습니다.
  • 저조도 개선 (LOL): 두 번째로 우수한 성능 (PSPSNR 24.04, SSIM 0.9118)을 기록했으며, 지표상으로는 LLFormer에 약간 뒤처졌으나 강력한 구조 보존 능력을 유지했습니다.
• 효율성: 상호 학습은 샘플링 효율성을 크게 향-상시켰습니다. 모델은 고품질 결과를 얻기 위해 단 1~4회의 샘플링 단계만을 필요로 했으며, 이는 다른 확산 기반 방법들(예: DiffusionVI의 경우 1000단계 필요)과 대조적입니다.
• 절제 연구(Ablation Studies):
  • 상호작용 모듈을 제거했을 때 모든 태스크에서 성능 저하가 발생하여 특징 융합의 이점을 확인했습니다.
  • 동적 시간 의존적 모델링($\bar{\alpha}_t, \bar{\beta}_t$)을 사용하는 것이 정적 타임스텝 모델링($t$)보다 우수한 성능을 보였습니다.
  • 최적의 샘플링 단계는 태스크마다 달랐습니다 (저조도 1, 디노이징 2, 그림자 제거 3, 디레인 4).

의의 및 주장

본 논문은 제안된 프레임워크가 재구성 정확도와 다양성의 균형을 맞추는 멀티 태스크 이미지 복원을 위한 통합 솔루션을 제공한다고 주장합니다. 상호 학습을 통해 결정론적 모델링과 확률적 모델링을 통합함으로써, 이 접근 방식은 GAN 기반 방법보다 더 강력한 해석 가능성을 제공하고, 태스크 특화된 CNN이나 트랜스포머보다 더 나은 일반화 능력을 제공합니다.

저자들은 특히 밝기 변화가 큰 저조도 개선 태스크에서 밝기 복원이 전문화된 베이스라인에 비해 불충분할 수 있다는 점을 포함하여 한계를 인정합니다. 또한 현재 지도 학습을 위해 쌍을 이룬 깨끗한 이미지-열화 데이터에 의존하고 있음을 언급했습니다. 향후 연구는 조명에 민감한 태스크를 위한 확산 방정식 최적화와 데이터 의존도를 줄이기 위한 비지도 학습 전략 개발로 향할 것입니다.

전반적으로, 본 연구는 단순하고 상호작용적인 이중 경로 확산 아키텍처가 이미지 열화 과정의 이질성을 효과적으로 다룰 수 있음을 보여주며, 다양한 이미지 복원 과제에 대해 확장 가능하고 효율적인 대안을 제시합니다.