All-in-One Image Restoration via Causal-Deconfounding Wavelet-Disentangled Prompt Network

이 논문은 이미지 복원 모델에서 발생하는 위상 상관관계와 편향된 손상 패턴 추정을 해결하기 위해 웨이블릿 기반의 인과적 교란 및 해리 메커니즘을 도입한 CWP-Net 을 제안하여 다양한 손상 유형을 통합적으로 복원하는 성능을 입증합니다.

핵심

이 논문은 **"온갖 더러운 사진을 한 번에 깨끗하게 만드는 똑똑한 AI"**에 대한 이야기입니다.

기존의 사진 복원 기술은 "비 제거용", "흐림 제거용", "어두운 사진용"처럼 용도별로 AI 를 따로 만들어야 했습니다. 마치 비가 올 때는 우산을, 눈이 올 때는 모자를 따로 챙겨야 하는 것처럼 불편하고 저장 공간도 많이 차지했습니다.

최근에는 "모든 것을 한 번에 해결하는 (All-in-One)" AI 가 등장했지만, 이 논문은 **"아직도 이 AI 들이 속고 있습니다"**라고 지적합니다.

이 논문의 핵심 내용을 쉬운 비유로 설명해 드릴게요.

---

1. 왜 기존 AI 는 실패할까요? (두 가지 치명적인 실수)

이 논문은 기존 AI 가 두 가지 큰 실수를 저지르고 있다고 말합니다.

① "우연한 착각" (Spurious Correlation)

• 상황: AI 가 학습할 때, '비'가 내린 사진은 대부분 '동물'이 나오는 장면이고, '흐림'이 낀 사진은 '건물'이 나오는 장면이었다고 가정해 봅시다.
• 실수: AI 는 비를 제거하는 법을 배우는 대신, **"동물이 나오면 비가 온 거야!"**라고 착각하게 됩니다.
• 결과: 나중에 '동물이 없는 비'나 '동물이 있는 흐림' 사진을 보면 AI 는 당황해서 엉뚱한 작업을 하거나 실패합니다. 즉, 사진의 내용 (동물, 건물) 과 결함 (비, 흐림) 을 혼동하는 것입니다.

② "편향된 추측" (Biased Estimation)

• 상황: AI 가 "이 사진이 비인지, 안개인지, 흐림인지"를 판단할 때, 학습 데이터가 불균형하면 (예: 비 사진은 많고 안개 사진은 적음) 자주 보는 것만 잘 맞추고, 드문 것은 엉터리로 맞춥니다.
• 결과: 정확한 판단을 못 하면, 그 판단을 바탕으로 사진을 고치는 과정도 엉망이 됩니다.

---

2. 이 논문이 제안한 해결책: "CWP-Net" (현명한 사진 수리공)

이 논문은 **파동 (Wavelet)**이라는 수학적 도구를 이용해 위 두 가지 문제를 해결하는 새로운 AI, CWP-Net을 제안합니다.

비유: "사진을 주파수 (소리의 높낮이) 로 나누어 보는 안경"

일반적인 AI 는 사진을 '그림'으로만 봅니다. 하지만 CWP-Net 은 사진을 **'주파수' (소리의 높낮이)**로 나누어 봅니다.

• 저주파 (LL): 사진의 전체적인 모양, 배경, 색감 (예: 하늘, 벽).
• 고주파 (LH, HL, HH): 사진의 디테일, 가장자리, 결함 (예: 비줄기, 안개, 노이즈).

이 안경을 끼고 보면, 비줄기나 안개는 주로 '고주파' 영역에 집중되어 있고, 배경 (동물, 건물) 은 '저주파' 영역에 있다는 것을 명확히 구분할 수 있습니다.

이 AI 의 두 가지 핵심 기능

1. "결함과 내용을 분리하는 안경" (Wavelet Attention Module)

• 기능: AI 가 사진을 볼 때, 비나 안개 같은 '결함' 부분에만 집중하고, '동물'이나 '건물' 같은 '내용'은 무시하도록 훈련시킵니다.
• 효과: 앞서 말한 "동물이 나오면 비가 온 거야"라는 착각을 막아줍니다. 결함만 보고 결함만 고치기 때문에, 어떤 배경이든 상관없이 똑같이 잘 고칩니다.

2. "상황에 맞춰 변신하는 지시봉" (Wavelet Prompt Block)

• 기능: AI 가 "이게 비인지, 안개인지"를 100% 정확히 알 수 없을 때 (편향된 추측), **가상의 대안 (Alternative Variable)**을 만들어냅니다.
• 작동 원리: "아마 비일 수도 있고, 안개일 수도 있는데, 두 가지 경우를 모두 고려해서 고쳐보자"라고 확률적으로 접근합니다. 마치 요리사가 "소금이 부족할 수도 있고, 설탕이 부족할 수도 있으니, 두 가지를 모두 조금씩 넣어보자"라고 유연하게 대처하는 것과 같습니다.
• 효과: 결함의 종류를 정확히 몰라도, 가장 적합한 방식으로 사진을 복원할 수 있게 됩니다.

---

3. 요약: 이 기술이 왜 대단한가요?

• 한 번에 해결: 비, 안개, 흐림, 노이즈, 어둠 등 모든 결함을 하나의 모델로 해결합니다.
• 착각 없음: 사진의 내용 (동물, 사람) 과 결함 (비, 안개) 을 명확히 구분해서, 어떤 상황에서도 일관된 성능을 냅니다.
• 정교한 복원: 단순히 흐릿하게 지우는 게 아니라, **세부적인 질감 (Texture)**까지 살려서 원래 사진처럼 자연스럽습니다.

한 줄 요약:

"기존 AI 가 사진의 '내용'과 '결함'을 헷갈려서 엉뚱한 일을 했다면, 이 새로운 AI 는 결함만 골라내는 현명한 수리공이 되어, 어떤 상황에서도 완벽한 사진을 만들어냅니다."

이 기술은 자율주행차, 감시 카메라, 스마트폰 카메라 등 실제 생활에서 다양한 환경의 사진을 즉시 깨끗하게 만들어야 하는 모든 분야에 큰 도움을 줄 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 이미지 복원 (Image Restoration) 은 노이즈, 흐림, 안개, 저조도 등 다양한 열화 (degradation) 를 제거하여 선명한 이미지를 복원하는 기술입니다. 기존의 표준 접근법은 특정 작업 (예: 제노이징, 디블러링 등) 에 맞춰 별도의 모델을 학습시키거나, 공통 아키텍처를 공유하되 작업별 가중치를 분리하여 사용합니다.
• 문제점:
  1. 높은 저장 비용: 각 열화 유형마다 별도의 모델을 저장해야 하므로 메모리 비용이 큽니다.
  2. 열화 패턴의 사전 지식 필요: 실제 환경에서는 열화의 종류와 정도를 알기 어렵습니다.
  3. All-in-One Image Restoration (AiOIR) 의 한계: 하나의 모델로 여러 열화를 처리하려는 AiOIR 방법이 등장했으나, 여전히 두 가지 치명적인 결함이 존재합니다.
     • 허위 상관관계 (Spurious Correlation): 열화 패턴 (Degradation Pattern) 과 비열화 의미론적 특징 (Semantic Features, 예: 배경, 사물) 사이에 인과관계가 아닌 우연한 상관관계가 학습됩니다. (예: '비' 데이터셋에는 동물이 자주 등장하지만 '흐림' 데이터셋에는 없음). 이로 인해 모델이 열화 대신 의미론적 특징을 기반으로 잘못된 추론을 합니다.
     • 편향된 열화 패턴 추정 (Biased Estimation): 불균형한 학습 데이터로 인해 열화 분류기 (Degradation Classifier) 가 열화 패턴을 정확히 추정하지 못하며, 이는 모델의 일반화 성능을 저하시킵니다.

2. 제안 방법론: CWP-Net (Methodology)

저자들은 구조적 인과 모델 (Structural Causal Model, SCM) 을 활용하여 AiOIR 의 인과 메커니즘을 분석하고, 이를 해결하기 위해 **Causal-deconfounding Wavelet-disentangled Prompt Network (CWP-Net)**을 제안했습니다.

A. 인과적 분석 (Causal Analysis)

• SCM 프레임워크: 열화된 이미지 ($X$) 와 복원된 이미지 ($Y$) 사이의 진정한 인과관계 ($P(Y|do(X))$) 를 학습하는 것이 목표입니다.
• 허위 상관관계 해결: 의미론적 특징 ($C$) 과 열화 패턴 ($T$) 사이의 허위 상관관계 ($C \sim T$) 를 차단해야 합니다.
• 편향된 추정 해결: 관측 불가능한 열화 패턴 ($T$) 대신, 인과적 교란을 제거하기 위한 **대체 변수 (Alternative Variable, $P$)**를 도입하여 백도어 조정 (Backdoor Adjustment) 을 수행합니다.

B. 핵심 모듈 구성

1. 웨이브릿 어텐션 모듈 (Wavelet Attention Modules):

   • 목적: 허위 상관관계를 해결하기 위해 열화 특징 ($D$) 과 의미론적 특징 ($C$) 을 명시적으로 분리 (Disentangle) 합니다.
   • 구조:
     • WAE (Encoder): 인코더에 삽입되어 이산 웨이브릿 변환 (DWT) 을 통해 저주파 및 고주파 계수를 추출합니다. 저주파 (Low-frequency) 부분의 어텐션 맵을 열화 표현으로 사용하여, 열화된 영역에 높은 가중치를 부여하고 의미론적 영역은 무시하도록 설계되었습니다.
     • WAD (Decoder): 디코더에 삽입되어 복원된 특징을 정제합니다.
   • 효과: 모델이 열화 패턴을 의미론적 배경과 분리하여 학습하도록 유도합니다.

2. 웨이브릿 프롬프트 블록 (Wavelet Prompt Block, WPB):

   • 목적: 편향된 열화 패턴 추정을 해결하기 위해 대체 변수 $P$를 생성하고 백도어 조정을 수행합니다.
   • 구성:
     • Degradation-based Weight Estimator (DWE): WAE 에서 추출된 열화 표현을 기반으로 K-Means 클러스터링을 수행하여 현재 이미지의 열화 유형을 파악하고, 각 웨이브릿 서브밴드에 적용할 가중치 ($\omega$) 를 추정합니다.
     • Prompt-guided Weighted Spatial Feature Transform (PWSFT): 학습 가능한 프롬프트 (Prompt) 를 생성하여 웨이브릿 특징과 상호작용합니다. 공간적 특징 변환 (SFT) 을 통해 픽셀 단위의 스케일링 및 시프트 파라미터를 학습하여 특징을 조절합니다.
   • 작동 원리: 열화 패턴이 웨이브릿 서브밴드의 왜곡 정도에 영향을 미친다는 관측 ($T \to P$) 을 바탕으로, $P$를 조정 변수로 사용하여 $P(Y|do(X))$를 추정합니다.

3. 전체 아키텍처:

   • U-Net 구조를 기반으로 하며, 인코더와 디코더의 각 스케일에 WAE 와 WAD 를 대칭적으로 배치하고, 스킵 연결 (Skip-connection) 에 WPB 를 삽입하여 다중 스케일 특징을 활용합니다.

3. 주요 기여 (Key Contributions)

1. 인과적 통찰: AiOIR 성능을 저하시키는 근본 원인인 '허위 상관관계'와 '편향된 열화 추정'을 인과적 관점에서 규명하고, 이를 해결하기 위한 이론적 프레임워크 (SCM) 를 제시했습니다.
2. CWP-Net 제안:
   • 웨이브릿 어텐션 모듈: 의미론적 특징과 열화 특징을 분리하여 허위 상관관계를 제거합니다.
   • 웨이브릿 프롬프트 블록: 편향된 추정을 보완하기 위한 대체 변수를 생성하고 백도어 조정을 구현합니다.
3. 성능 향상: 기존 최첨단 (SOTA) AiOIR 방법론들을 압도하는 성능을 달성했습니다.
4. 일반화 능력: 불균형한 데이터셋뿐만 아니라, 의미론적 특징과 열화 패턴 간의 상관관계가 제거된 균형 잡힌 테스트셋 (Balanced Test Set) 에서도 뛰어난 일반화 성능을 입증했습니다.

4. 실험 결과 (Results)

• 평가 설정: 5 가지 열화 패턴 (5-pattern) 과 7 가지 열화 패턴 (7-pattern) 설정에서 평가 수행.
• 정량적 성능:
  • 5-pattern 설정: Lin et al. [54] 대비 평균 PSNR 에서 0.59dB 향상. 특히 안개 제거 (Dehazing) 에서 1.58dB, 비 제거 (Deraining) 에서 1.10dB 향상.
  • 7-pattern 설정: 기존 2 위 방법 (IDR) 대비 평균 PSNR 2.22dB 향상. AirNet 대비 5.07dB 향상.
  • 균형 잡힌 테스트셋 (Generalization): 도메인 편향이 제거된 환경에서 AirNet 및 PromptIR 대비 월등히 높은 성능을 보이며, 의미론적 특징에 의존하지 않고 진정한 열화 정보를 복원함을 입증했습니다.
• 정성적 성능: 시각적 비교에서 기존 방법들이 제거하지 못한 잔여 열화 (비 줄무늬, 노이즈 등) 를 효과적으로 제거하고, 선명한 텍스처와 자연스러운 색상을 복원했습니다.
• 효율성: 파라미터 수와 FLOPs 면에서 Lin et al. (CLIP 기반) 나 PromptIR 보다 가볍고 추론 속도가 빠르며, 경량 모델들 (BRDNet 등) 보다 훨씬 높은 복원 품질을 제공합니다.

5. 의의 및 결론 (Significance)

이 논문은 이미지 복원 분야에서 **인과적 추론 (Causal Inference)**을 적용하여 모델의 일반화 능력을 획기적으로 개선한 선구적인 연구입니다.

• 이론적 기여: 단순히 데이터 양을 늘리는 것을 넘어, 데이터 내의 인과 구조를 분석하고 허위 상관관계를 제거하는 방법론을 제시했습니다.
• 실용적 가치: 하나의 모델로 다양한 열화 유형을 처리하면서도 높은 정확도를 유지하여, 실제 동적 환경 (자율주행, 감시 카메라 등) 에서의 적용 가능성을 높였습니다.
• 기술적 혁신: 웨이브릿 변환을 프롬프트 학습 및 인과적 교란 제거에 성공적으로 접목하여, 주파수 영역 분석과 딥러닝 기반 이미지 복원의 새로운 방향성을 제시했습니다.

결론적으로, CWP-Net 은 기존 AiOIR 방법론의 한계를 인과적 관점에서 해결하고, 웨이브릿 기반의 디커플링 기법을 통해 뛰어난 복원 성능과 일반화 능력을 동시에 달성한 획기적인 모델입니다.