기술 요약: TDiR (Transformer 기반 확산 모델을 이용한 이미지 복원)

문제 정의

이미지 복원은 노이즈, 블러(blur), 색상 왜곡(color cast), 빛의 산란 등으로 인해 저하된 버전으로부터 깨끗하고 원래의 이미지를 복구하는 것을 목표로 하는 컴퓨터 비전의 핵심 과제입니다. 이러한 저하 현상은 객체 탐지, 매핑, 분류와 같은 다운스트림 작업에서의 이미지 활용도를 크게 떨어뜨립니다. 비생성적(non-generative) 방법론들이 고수준 비전 작업에서 성공을 거두기도 했으나, 수중 재구성, 디노이징(denoising), 디레인(deraining)과 같은 저수준 비전 작업에서 발견되는 복잡하고 다봉형인(multimodal) 분포를 처리하는 데는 어려움을 겪는 경우가 많습니다. 특히 수중 이미지는 환경 조건에 따라 변하는 빛의 산란, 흡수 및 색상 변화라는 독특한 과제를 안고 있으며, 이는 단일 저하 입력에서 단일 깨끗한 출력으로 매핑하는 것보다 여러 잠재적 저하 상태를 탐색할 수 있는 능력을 갖춘 모델을 필요로 합니다.

방법론

본 논문은 통합된 프레임워크 내에서 다양한 이미지 복원 작업을 해결하기 위해 설계된 트랜스포머 기반 조건부 확산 모델인 TDiR을 제안합니다.

• 아키텍처: TDiR의 핵심은 PromptIR 백본을 변형한 트랜스포머 기반 U-Net 아키텍처입니다. PromptIR은 별도의 작업별 모델 없이도 다양한 수준의 저하에 적응하기 위해 프롬프트 모듈을 활용합니다.
• 확산 과정 (Diffusion Process): 모델은 순방향 확산 과정(가우시안 노이즈를 점진적으로 추가)과 역방향 디노이징 단계로 구성된 확산 프레임워크 내에서 작동합니다.
  • 순방향 과정 (Forward Process): 입력 이미지 $y_0$에 $T$번의 반복에 걸쳐 가우시안 노이즈를 추가하여 순수 노이즈와 유사하게 만듭니다.
  • 역방향 과정 (Backward Process): 디노이징 네트워크는 이 과정을 역전시켜 가우시안 분포로부터 데이터셋의 경험적 분포로 이동하는 법을 학습합니다.
• 핵심 아키텍처 수정 사항: 저자들은 저하된 입력을 확산 타임스텝(diffusion timestep)의 임베딩과 결합하는 전용 디코더 경로를 도입하는 경량화된 수정을 제안합니다. 이를 통해 네트워크는 순방향 과정 중에 도입된 특정 노이즈 레벨에 따라 복원 과정을 조절할 수 있습니다.
• 학습 전략: 트랜스포머의 데이터 집약적인 특성을 해결하기 위해, 저자들은 사전 학습된 인코더(frozen)를 백본으로 사용합니다. 학습 중에는 디코더 파라미터만 최적화됩니다. 모델은 수중 향상, 디노이징, 디레인을 포함하는 결합된 데이터셋을 통해 학습됩니다. 목적 함수는 예측된 노이즈와 실제 이미지에 추가된 노이즈 사이의 $L_1$ 손실을 최소화합니다.
• 추론 (Inference): 복원 과정은 반복적인 디노이징 절차를 포함하며, 이미지당 여러 번의 순방향 패스(forward pass)를 요구합니다.

주요 기여

1. 통합 조건부 확산 모델: 다양한 복원 작업(수중, 디노이징, 디레인)에 대해 작업별 모델 교체 없이 적용 가능한, 사전 학습된 PromptIR 백본을 노이즈 조건부 디노이저로 변환하는 TDiR을 제안했습니다.
2. 경량 타임스텝 컨디셔닝: 저하된 입력과 확산 타임스텝 임베딩을 결합하는 디코더 경로를 도입하여, 모델이 노이즈 레벨에 따라 복원 전략을 동적으로 적응할 수 있도록 했습니다.
3. 포괄적인 평가: 5개의 벤치마크 데이터셋(UIEB, Test-60, BSD68, Rain100L 등)과 4개의 평가 지표(PSNR, SSIM, UCIQE, UIQM)를 사용하여 TDiR을 세 가지 작업에 대해 엄격하게 평가했습니다. 본 연구는 TDiR을 18개의 최첨단(SOTA) 기술과 비교했습니다.
4. 정성적 분석: 순수 판별적(discriminative) 네트워크와 비교하여 확산 기반 복원 과정의 장점을 조사하였으며, 잔차 패치 기반 아티팩트(residual patch-based artifacts)와 같은 제약 사항을 확인하면서도 다봉형 저하를 처리하는 데 있어의 이점을 강조했습니다.

실험 결과

논문은 T-DiR이 평가된 모든 작업에서 현재의 최첨단 방법들을 능가한다고 보고합니다:

• 수중 이미지 향상 (Underwater Image Enhancement): UIEB 데이터셋에서 TDiR은 22.90 dB PSRE 및 0.8724 SSIM이라는 가장 높은 참조 기반 점수를 달성했으며, 이는 가장 가까운 경쟁 모델인 MetaUE를 각각 0.89 dB와 0.011 SSIM 차이로 앞선 수치입니다. 무참조(no-reference) 방식인 Test-60 서브셋에서는 UIQM 점수 2.73을 기록하여 다른 방법들을 크게 상회했습니다. 정성적 결과는 PromptIR 및 기타 베이스라인 모델에 비해 우수한 색상 충실도와 녹색 색조 억제력을 보여줍니다.
• 이미지 디노이징 (Image Denoising): BSD68 데이터셋에서 TDiR은 노이즈 레벨 $\sigma=15, 25, 50$에 대해 모든 벤치마크 방법(특화된 단일 작업 모델 포함)을 능가했습니다. 특히 $\sigma=50$에서 TDiR은 전문 모델인 PromptIR을 넘어섰습니다. 높은 노이즈 레벨에서 패치 아티팩트로 인해 SSIM은 약간 감소했으나, 모델은 균일한 영역에서 눈에 띄는 입자 아티팩트가 적고 지각적으로 더 깨끗한 결과를 생성했습니다.
• 이미지 디레인 (Image Deraining): Rain100L 데이터셋에서 TDiR은 PSNR 37.43을 달성하여, 재학습된 PromptIR* 베이스라인보다 0.40 dB 높았으며 AirNet 및 MPRNet과 같은 다른 방법들을 크게 앞질렀습니다. 정성적 분석 결과, 균일한 영역(하늘, 모래)에서 빗줄기가 완전히 제거되었으며 경쟁 모델들에 비해 미세한 질감이 더 잘 보존됨을 확인했습니다.

의의 및 주장

본 논문은 확산 모델과 트랜스포머의 결-합이 특히 복잡하고 다봉형인 저하가 발생하는 시나리오에서 이미지 복원을 위한 강력한 솔루션을 제공한다고 주장합니다. 저자들은 확산 모델이 저하된 이미지의 다양한 분포를 탐색하는 데 탁월하며, 이를 통해 실제 세계의 가변성을 반영하는 더 미세한 복원이 가능하다고 강조합니다. 이는 조명과 투명도가 크게 변하는 수중 이미징에 특히 유익합니다.

그러나 논문은 다음과 같은 한계점에 대해 겸허한 태도를 유지합니다:

• 계산 비용: 확산 모델의 반복적인 특성으로 인해 TDiR은 단일 패스 판별 모델보다 추론 비용이 훨씬 높으며, 이는 실시간 응용 분야에 제한이 될 수 있습니다.
• 아티팩트: 트랜스포머 기반 디노이저의 패치 단위 처리로 인해 스티치 경계에서 블로킹 아티팩트(blocking artifacts)가 발생할 수 있으며, 이는 지각적 품질 대비 PSNR/SSIM 점수를 미세하게 낮출 수 있습니다.
• 일반화 트레이드오프: 통합 모델은 다재다능하지만, 특정 시나리오(예: 평탄한 배경을 가진 저심도 수중 장면)에서는 전문화된 단일 작업 모델보다 일관되게 낮은 성능을 보이며, 이는 범용성과 최고 정확도 사이의 트레이드오프를 시사합니다.
• 깊이 의존성: 모델은 효과적인 수중 복상을 위해 깊이 관련 단서에 의존하며, 이는 매우 도전적인 장면에서는 부재할 수 있습니다.

저자들은 TDiR이 여러 작업에서 이미지 품질의 유의미한 향상을 보여주었지만, 향후 연구를 통해 중첩 패치 추론(overlapping-patch inference), 더 적은 확산 단계로의 증류(distillation), 또는 작업 조건부 미세 조정(task-conditioned fine-tuning) 등의 기술을 통해 이러한 한계를 해결할 수 있다고 결론지었습니다.