[논문 리뷰] MambaIR: A Simple Baseline for Image Restoration with State-Space Model

Hang Guo, Jinmin Li외 4명|arXiv (Cornell University)|2024. 02. 23.

Image and Signal Denoising Methods인용 수 18

한 줄 요약

MambaIR은 선택적 상태-공간 모델(Mamba)을 기반으로 한 간단하고 강력한 이미지 복원 백본을 도입하고, 복원 priors를 추가하여 Transformer 기반 기준선과의 비교 가능하거나 우수한 결과를 선형 복잡도로 달성한다.

ABSTRACT

Recent years have seen significant advancements in image restoration, largely attributed to the development of modern deep neural networks, such as CNNs and Transformers. However, existing restoration backbones often face the dilemma between global receptive fields and efficient computation, hindering their application in practice. Recently, the Selective Structured State Space Model, especially the improved version Mamba, has shown great potential for long-range dependency modeling with linear complexity, which offers a way to resolve the above dilemma. However, the standard Mamba still faces certain challenges in low-level vision such as local pixel forgetting and channel redundancy. In this work, we introduce a simple but effective baseline, named MambaIR, which introduces both local enhancement and channel attention to improve the vanilla Mamba. In this way, our MambaIR takes advantage of the local pixel similarity and reduces the channel redundancy. Extensive experiments demonstrate the superiority of our method, for example, MambaIR outperforms SwinIR by up to 0.45dB on image SR, using similar computational cost but with a global receptive field. Code is available at \url{https://github.com/csguoh/MambaIR}.

연구 동기 및 목표

• CNN과 Transformer를 넘어서는 저수준 이미지 복원에 상태-공간 모델 사용의 필요성을 제안한다.
• 복원 특화 RSSB를 제안하여 로컬 공간 선험 및 채널 상호작용으로 Mamba를 강화한다.
• 다양한 복원 작업에서 유사한 계산 비용으로 Transformer 기반 기준선보다 더 나은 성능을 낼 수 있음을 보인다.

제안 방법

• Mamba를 얕은 특징 추출, 깊은 특징 추출, 고품질 재구성의 세 단계 이미지 복원 파이프라인으로 적응시킨다.
• 복원 작업을 위해 Vision State-Space Module (VSSM)과 합성곱 로컬 프라이어와 채널 어텐션을 결합한 Residual State-Space Block (RSSB)을 도입한다.
• 4 방향 탐색으로 2D 공간 의존성을 포착하는 2D Selective Scan Module (2D-SSM)을 활용한다.
• 병렬 분기로 Hadamard 곱 결합을 갖춘 Vision State-Space Module을 구현하여 장거리 의존성과 채널 상호작용을 제공한다.
• Mamba의 인과 처리(causal processing)를 2D 이미지 데이터로 확장하기 위해 2D Selective Scan Module을 활용한다.
• 전역 수용 영역을 보존하는 비주목적(attentional) 선형 시간 복잡도 접근법을 강조한다.

Figure 1 : The Effective Receptive Field (ERF) visualization [ 38 ] for EDSR [ 35 ] , RCAN [ 68 ] , SwinIR [ 34 ] , HAT [ 7 ] , and the proposed MambaIR. A larger ERF is indicated by a more extensively distributed dark area. Only the proposed MambaIR achieves a significant global effective receptive

실험 결과

연구 질문

• RQ1 상태-공간 모델을 이미지 복원 과제의 기본 백본으로 어떻게 활용할 수 있는가?
• RQ2 로컬 패치 재발현(local patch recurrence) 및 채널 상호작용과 같은 어떤 복원 선험이 저수준 비전에서 Mamba의 효과를 향상시키는가?
• RQ3 간단한 RSSB 기반 설계가 유사한 계산 비용에서 Transformer 기반 복원 모델에 비견될 성능을 달성할 수 있는가?

주요 결과

• MambaIR은 유사한 계산 비용에서 SwinIR과 같은 Transformer 기반 기준선보다 우수한 여러 작업에서 경쟁력 있는 복원 성능을 달성한다.
• RSSB 설계는 Mamba 백본에 로컬 패치 재발현과 채널 상호작용을 도입하여 성능을 향상시킨다.
• Vision State-Space Module과 2D-Selective Scan은 선형 복잡도로 장거리 모델링을 가능하게 하며 글로벌 수용 필드를 유지한다.
• 이 방법은 고전적 SR, 경량 SR 및 다수 데이터셋에서 실제 이미지 노이즈 제거에 강력한 결과를 보인다.
• 공략 분석(Ablation studies)은 RSSB 내부의 로컬 컨볼루션과 채널 어텐션의 중요성이 복원 품질에 중요한 역할을 함을 보여준다.

Figure 2 : The overall network architecture of our MambaIR, as well as the (a) Residual State-Space Block (RSSB), the (b) Vision State-Space Module (VSSM), and the (c) 2D Selective Scan Module (2D-SSM).