Paper: Selective Hourglass Mapping for Universal Image Restoration Based on Diffusion Model (CVPR 2024)

최근 Image Restoration 분야에서는 하나의 모델로 다양한 degradation 상황에 대응하는 All-in-One Restoration (AIOR)에 대한 관심이 높아지고 있다. 실제 환경에서는 이미지가 단순히 흐릿해지거나 노이즈가 끼는 것뿐 아니라, 비, 안개, 저조도, 눈, 모션 블러 등이 복합적으로 발생하기 때문에, 단일 degradation만 처리하는 모델로는 한계가 있다.

이러한 한계를 해결하기 위해 최근 다양한 연구들이 등장하고 있으며, 그중에서도 Diffusion 기반 All-in-One 복원 모델은 점차 주목받고 있다. 대표적으로 각 열화 유형을 condition으로 넣거나 task-specific encoder를 활용하는 방식이 있지만, 여전히 각 task의 데이터 분포를 독립적으로(map separately) 학습하는 구조를 가진다. 이런 구조는 논문에서 multi-partite mapping이라 부르며, task 간의 공통 정보나 상호작용을 모델이 학습할 수 없다는 근본적인 한계를 갖는다.

따라서 본 논문에서는 다양한 열화 분포를 하나의 shared distribution으로 수렴시킨 뒤, 다시 task-specific 분포로 복원하는 Selective Hourglass Mapping 구조를 소개한다. 이 방식은 단지 각 task를 하나의 모델로 처리하는 것에 그치지 않고, inter-task interaction이나 공통 feature representation까지 모델이 적극적으로 활용할 수 있도록 해준다.

저자들이 언급한 핵심 contribution은 다음 2가지이지만, 사실상 기존 RDDM에 shared distribution term (SDT)을 추가한 게 전부다.

1. Strong condition guidance

   기존 RDDM 방식처럼 degraded image 정보를 diffusion 수식에 직접 삽입 (explicit condition)하고, 입력에도 concat (implicit condition)하여 이미지 품질을 높임.

2. Shared distribution term (SDT)

   서로 다른 degradation 분포를 SDT를 통해 점진적으로 하나의 공통 분포로 정규화시킨 후, reverse 과정에서 다시 각 task에 맞는 분포로 복원 (hourglass 구조).

RDDM 리뷰 및 수식 정리는 이전 포스팅에도 남겨두었고, 아래에 한번 더 정리해두었다.

DiffUIR 수식도 마찬가지로 아래 노트에 정리해두었다.

실제로 SDT의 유무에 따라 diffusion endpoint의 feature 분포가 달라지며, t-SNE로 시각화한 결과 (아래 Figure 5) 에서는 task별 클러스터가 뚜렷하게 분리되는 현상이 완화된 것을 확인할 수 있다.

DiffUIR의 성능은 아래 Table 1에서 확인할 수 있다.

제안된 모델은 5개 restoration task 모두에서 기존 SOTA와 비슷하거나 더 나은 성능을 달성했으며, 특히 파라미터 수가 1M도 되지 않는 Tiny 모델조차도 SOTA에 근접한 성능을 유지하는 점이 인상적이다.