[논문 리뷰] MambaUIE&SR: Unraveling the Ocean's Secrets with Only 2.8 GFLOPs

Zhihao Chen, Yiyuan Ge|arXiv (Cornell University)|2024. 04. 22.

Oceanographic and Atmospheric ProcessesEarth and Planetary Sciences인용 수 3

한 줄 요약

이 논문은 상태공간모델(SSM) 아키텍처인 Mamba를 기반으로 한 새로운 수중 이미지 향상 모델인 MambaUIE&SR를 제안하며, 단지 2.8 GFLOPs의 계산량으로 최신 기술 수준의 성능을 달성한다. 시각적 상태공간(VSS) 블록, 동적 상호작용 블록(DIB), 그리고 공간 게이팅 피드포워드 네트워크(SGFP)를 통합함으로써, 전반적인 맥락과 국소적인 세밀한 특징을 효율적으로 포착하여 기존 최신 기술 수준의 방법들 대비 67.4%의 계산 비용 절감을 이룬다. 동시에 높은 정확도를 유지한다.

ABSTRACT

Underwater Image Enhancement (UIE) techniques aim to address the problem of underwater image degradation due to light absorption and scattering. In recent years, both Convolution Neural Network (CNN)-based and Transformer-based methods have been widely explored. In addition, combining CNN and Transformer can effectively combine global and local information for enhancement. However, this approach is still affected by the secondary complexity of the Transformer and cannot maximize the performance. Recently, the state-space model (SSM) based architecture Mamba has been proposed, which excels in modeling long distances while maintaining linear complexity. This paper explores the potential of this SSM-based model for UIE from both efficiency and effectiveness perspectives. However, the performance of directly applying Mamba is poor because local fine-grained features, which are crucial for image enhancement, cannot be fully utilized. Specifically, we customize the MambaUIE architecture for efficient UIE. Specifically, we introduce visual state space (VSS) blocks to capture global contextual information at the macro level while mining local information at the micro level. Also, for these two kinds of information, we propose a Dynamic Interaction Block (DIB) and Spatial feed-forward Network (SGFN) for intra-block feature aggregation. MambaUIE is able to efficiently synthesize global and local information and maintains a very small number of parameters with high accuracy. Experiments on UIEB datasets show that our method reduces GFLOPs by 67.4% (2.715G) relative to the SOTA method. To the best of our knowledge, this is the first UIE model constructed based on SSM that breaks the limitation of FLOPs on accuracy in UIE. The official repository of MambaUIE at https://github.com/1024AILab/MambaUIE.

연구 동기 및 목표

• 기존의 CNN 및 트랜스포머 기반 수중 이미지 향상(UIE) 방법들이 겪는 높은 계산 비용과 제한된 성능 문제를 해결하기 위해.
• 특히 Mamba를 포함한 상태공간모델(SSM)이 효율적이고 효과적인 UIE에 어떻게 활용될 수 있는지 탐색하기 위해.
• 직접 Mamba를 UIE에 적용했을 때의 열악한 성능을 해결하기 위해 국소적인 세밀한 특징을 효과적으로 유지하기 위해.
• 전반적인 맥락 모델링과 국소 세부 정보 복구 사이의 균형을 이루는 경량이며 파rameter 효율적인 아키텍처를 설계하기 위해.
• 정확도와 효율성 사이의 트레이드오프를 깨고, 극적으로 감소된 FLOPs로 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• 장기적 의존성과 전반적 맥락을 포착하기 위해 시각적 상태공간(VSS) 블록을 통합한 커스터마이징된 Mamba 기반 아키텍처인 MambaUIE를 제안한다.
• 블록 내에서의 적응형 특징 상호작용을 가능하게 하여 블록 내 특징 집합을 향상시키기 위해 동적 상호작용 블록(DIB)을 도입한다.
• 국소적 특징 학습을 강화하고 세밀한 디테일을 유지하기 위해 공간 게이팅 피드포워드 네트워크(SGFP)를 활용한다.
• 계층적 인코더-디코더 구조에서 VSS 블록, DIB, SGFP를 조합하여 전역적 정보 처리와 국소적 정보 처리의 균형을 맞춘다.
• Mamba의 SSM 메커니즘이 선형 복잡도를 가지므로, 저 FLOP 동작을 최적화하기 위해 이를 활용한다.
• UIEB 데이터셋에서 종합적으로 훈련하여 공동으로 향상 및 초해상도(SR) 능력을 학습한다.

실험 결과

연구 질문

• RQ1 Mamba 상태공간모델은 수중 이미지 향상에 효과적으로 적응되어 고효율을 달성할 수 있는가?
• RQ2 왜 Mamba의 직접적 적용이 UIE에서 실패하는가? 국소적인 세밀한 특징을 더 잘 유지하기 위해서는 어떻게 해야 하는가?
• RQ3 SSM과 국소적 특징 학습 메커니즘을 조합한 하이브리드 아키텍처가 저 FLOP 비용으로 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ4 제안된 동적 상호작용 블록(DIB)은 표준 어텐션 또는 피드포워드 메커니즘에 비해 특징 집합을 향상시키는 데 기여하는가?
• RQ5 SSM 기반 모델을 활용하여 수중 이미지 향상에서 정확도와 FLOPs 사이의 트레이드오프를 깰 수 있는가?

주요 결과

• MambaUIE&SR는 단지 2.8 GFLOPs의 FLOP 수치를 기록하여 최신 기술 수준의 방법 대비 계산 비용을 67.4% 절감했다.
• UIEB 데이터셋에서 정량적 지표(예: PSNR, SSIM)와 정성적 시각적 결과 모두에서 기존 최신 기술 수준의 방법들을 능가하는 성능을 보였다.
• DIB와 SGFP의 통합은 특징 집합을 크게 향상시켜 질감과 색상 디테일의 보존에 기여했다.
• 높은 정확도를 달성하면서도 매우 적은 수의 파라미터를 유지하여 강력한 파라미터 효율성을 입증했다.
• 이 연구는 수중 이미지 향상에서 전통적인 FLOPs와 정확도 사이의 트레이드오프를 성공적으로 깨뜨린 최초의 SSM 기반 UIE 모델이다.
• 다양한 수중 이미지 열화 패턴에 걸쳐 뛰어난 일반화 능력과 강건성을 보였다.