Variational Deep Unfolding with Mamba-Based Nonlocal Modeling for Underwater Image Enhancement

문제 정의

수중 이미징은 해양 공학 및 해양 연구에 매우 중요하지만, 캡처된 데이터는 빈번하게 낮은 가시성, 불균일한 조명, 색상 감쇄 및 노이즈 문제를 겪습니다. 이러한 열화 현상은 후속 컴퓨터 비전 작업의 성능을 크게 저해합니다.

• 고전적 방법: 전통적인 접근 방식은 모델 프리(model-free) 기법(물리적 가정을 배제하고 픽셀 강도를 조작하는 방식)과 물리 모델 기반(Jaffe-McGlamery 모델과 같은 수중 이미지 형성 모델을 기반으로 역문제를 해결하는 방식) 기법으로 나뉩니다. 변분법(variational methods)은 해석 가능성을 제공하지만, 그 효과는 적절한 사전 확률(prior)을 수동으로 선택하는 것에 크게 의존합니다.
• 딥 러닝 방법: 데이터 구동형 딥 러닝 접근 방식은 데이터로부터 직접 사전 확률을 학습하지만, 종종 복잡한 아키텍처와 대규모 데이터셋을 필요로 합니다. 순수 데이터 구동 모델은 해석 가능성이 부족할 수 있으며, 표준 딥 러닝 아키텍처(예: Transformer)는 장거리 의존성을 포착하기 위해 높은 계산 비용을 초래할 수 있습니다.
• 공백: 모델 기반 최적화의 해석 가능성과 딥 러닝의 유연성 사이의 간극을 메울 필요가 있으며, 특히 수중 장면의 장거리 자기 유사성(long-range self-similarities)을 효율적으로 포착하는 과제를 해결해야 합니다.

방법론

저자들은 변분 모델링을 학습 가능한 아키텍처에 통합한 **모델 기반 딥 언폴딩 네트워크(model-based deep unfolding network)**를 제안합니다. 이 접근 방식은 변분 베이스라인, 딥 언폴딩 아키텍처, 그리고 특화된 학습 전략의 세 가지 주요 구성 요소로 이루어져 있습니다.

1. 변분 베이스라인 (Variational Baseline)

이 방법은 데헤이징(dehazing) 분해에 기초한 변분 정식화에 근거합니다. 저자들은 다음과 같이 수정된 이미지 형성 모델을 제 propose합니다:
$$I = (J + N)t + A(1 - t)$$
여기서 $I$는 열화된 이미지, $J$는 복원된 이미지, $N$은 잔여 아티팩트를 흡수하기 위한 곱셈적 잔차 성분, $t$는 투과 맵(transmission map), $A$는 전역적 후방 산란광입니다.

최소화해야 할 에너지 범함수(energy functional)는 다음을 포함합니다:

• 데이터 충실도(Data Fidelity): 위에서 언급한 분해 모델에 기반합니다.
• 비로컬 정규화 ($\alpha$-항): 미세한 디테일과 질감을 보존하기 위한 비로컬 그래디언트 제약을 적용합니다.
• 전변량(Total Variation, $\beta$-항): 깊이 불연속성을 보존하면서 공간적 매끄러움을 보장하기 위해 투과 맵에 적용됩니다.
• 잔차 제약 ($\lambda$-항): 잔차 성분의 흡수를 제한합니다.
• 비로컬 그래디언트 충실도 ($\mu$-항): 복원된 이미지의 그래디언트와 특정 벡터 필드 $V$(열화된 입력의 증폭된 그래디언트에서 유도됨) 사이의 거리에 페널티를 부여합니다. 이 항은 저대비 수중 이미지에 숨겨진 구조적 정보를 강화하도록 설계되었습니다.
• 투과 충실도 ($\rho$-항): 추정된 투과 맵이 초기 추정치(Dark Channel Prior로부터 유도됨)와 가깝게 유지되도록 보장합니다.

저자들은 정의된 허용 공간 내에서 이 변분 문제의 **해의 존재성(existence of a minimizer)**을 확립하는 이론적 증명을 제공합니다.

2. 딥 언폴딩 네트워크 (Deep Unfolding Network)

Chambolle–Pock 기반의 프라이멀-듀얼(primal-dual) 전략을 사용하는 최적화 알고리즘이 $K$개의 스테이지를 가진 딥 뉴럴 네트워크로 "언폴딩(unfolded)"됩니다. 이 프레임워크에서 최적화 변수($J, t, N$)의 반복 업데이트는 네트워크 레이어에 대응합니다.

• 학습 가능한 사전 확률 (Learnable Priors): 이미지 사전 확률($J$) 및 투과 맵($t$)과 관련된 근사 연산자(proximal operators)는 학습 가능한 뉴럴 네트워크로 대체됩니다.
• Mamba 통합: 자기 주의 집중(self-attention)의 높은 계산 비용 없이 장거리 의존성과 자기 유사성을 효율적으로 포착하기 위해, 저자들은 Mamba 레이어(상태 공간 모델, State Space Models)를 통합합니다.
  • MambaResNet: 이미지 $J$를 위한 근사 연산자를 대체합니다. 입력 이미지, 화이트 밸런스가 향상된 버전, 히스토그램 평활화된 버전의 세 가지 브랜치를 처리하는 멀티 브랜치 설계를 사용합니다. 비로컬 상호작용을 모델링하기 위해 2D 선택적 스캔(SS2D) 블록을 채택합니다.
  • ProxNet: 투과 맵 $t$를 위한 경량 잔차 네트워크입니다.
  • MambaNet: 그래디언트 영역에서 작동하며 비로컬 그래디언트 충실도 항을 모델링하는 모듈입니다.
• 초기화: 네트워크는 $J$의 경우 열화된 이미지로, $N$은 0으로, $t$는 Dark Channel Prior 추정치로 초기화됩니다.

3. 근사 궤적 손실 (Proximal Trajectory Loss)

언폴딩 네트워크가 기저의 최적화 로직을 준수하도록 보장하기 위해, 저자들은 근사 궤적 손실을 도입합니다. 최종 출력만을 감독하는 대신, 이 손실은 네트워크의 중간 단계와 "이상적인" 복원 정규화 도구(지면 참(ground truth)을 포함한 단순화된 최적화 문제로부터 유도됨)의 반복 사이의 일관성을 강제합니다. 이는 중간 단계가 예상되는 최적화 경로에서 벗어나는 것을 방지합니다.

주요 기여

1. 변분 정식화: 아티팩트를 처리하고 에지를 보존하기 위해 곱셈적 잔차 성분과 비로컬 그래디언트 유형의 제약을 포함하는 새로운 수중 향상 에너지 범함수를 제안했습니다.
2. 이론적 분석: 제안된 비볼록(non-convex) 변분 문제에 대한 해의 존재성에 대한 증명을 제공했습니다.
3. Mamba 기반 언폴딩: 근사 연산자를 Mamba 기반 모듈로 대체하여 수중 장면의 장거리 의전성을 효율적으로 모델링할 수 있는 딥 언폴딩 네트워크를 구축했습니다.
4. 근사 궤적 손실: 네트워크의 언폴딩 단계와 이상적인 최적화 알고리즘의 반복 사이의 일관성을 강제하여 해석 가능성과 성능을 개선하는 새로운 손실 함수를 설계했습니다.

실험 결과

이 방법은 UIEB 및 SUIM-E 데이터셋에서 평가되었으며, 고전적 방법(예: CBAF, BayesianRetinex), 순수 딥 러닝 방법(예: UColor, Phaseformer, Mamba-UIE), 그리고 다른 언폴딩 방법(UnfoldNet)과 비교되었습니다.

• 정량적 성능: 제안된 언폴딩 방법은 두 데이터셋 모두에서 모든 지표(PSNR, SSIM, LPIPS)에 대해 최첨단(state-of-the-art) 결과를 달성했습니다. 예를 들어, UIEB에서 PSNR 25.16, SSIM 0.9140을 달성하여 두 번째로 우수한 방법(Mamba-UIE)을 상당한 차이로 앞질렀습니다.
• 시각적 품질: 시각적 비교 결과, 제안된 방법은 경쟁 모델들에 비해 색상 정확도(다른 방법들에서 흔히 나타나는 붉은색/녹색 톤 감소)를 더 잘 회복하고 미세한 디테일(예: 물고기 비늘, 해저 질감)을 더 잘 보존함을 보여주었습니다.
• 어블레이션 연구 (Ablation Studies):
  • 데헤이징 vs. Retinex: 데헤이징 기반 정식화가 Retinex 기반 대안보다 우수한 성능을 보였습니다.
  • 그래디언트 제약: 비로컬 그래디언트 제약을 제거하면 에지 대비와 지각적 품질이 저하되었습니다.
  • 궤적 손실: 근사 궤적 손실을 통해 중간 단계를 감독하는 것이 최종 출력에 대한 MSE만을 사용하는 것보다 최종 지표를 크게 향상시켰습니다.
  • 효율성: 다중 헤드 주의 집중(MHA)과의 비교에서, Mamba 기반 아키텍처는 높은 해상도에서도 더 높은 처리량(throughput)을 유지하면서 더 낮은 FLOPs와 GPU 메모리 사용량을 달성했습니다.

의의 및 주장

본 논문은 제안된 프레임워크가 모델 기반 최적화와 딥 러닝 사이의 간극을 성공적으로 메운다고 주장합니다. 물리적으로 동기화된 변분 모델을 상태 공간 모델(Mamba)의 효율성과 결합함으로써, 저자들은 다음을 달성했습니다:

• 개선된 복원: 최근의 최첨단 방법들과 비교하여 우수한 시각적 품질과 정량적 성능을 보여줍니다.
• 해석 가능성: 언폴딩 구조는 해의 존재성에 대한 이론적 보증과 함께 기저 최적화 문제의 로직을 보존합니다.
• 효율성: Mamba 레이어를 사용함으로써, 수중 이미징의 계산적 병목 현상인 Transformer 기반 접근 방식의 문제를 해결하면서 선형 복잡도로 전역적 문맥(비로컬 모델링)을 효율적으로 포착할 수 있습니다.

저자들은 결과가 유망하지만, 향-후 연구에서는 (현재 Dark Channel Prior에 의존하는) 초기 투과 추정치를 개선하고, 모델을 비디오 향상 및 시간적 일관성(temporal consistency)으로 확장하는 데 집중할 수 있다고 언급했습니다.