MambaIR: A Simple Baseline for Image Restoration with State-Space Model

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}.

한 줄 요약

MambaIR은 이미지 복원에서 Mamba 기반 모델의 국부적 퍼픽스 유사성과 채널 상호작용을 강화하여 기존 CNN 및 Transformer 대비 뛰어난 성능을 보인다.

핵심 기여도

MambaIR은 Mamba에 Residual State Space Block (RSSB)를 도입하여, 이미지 복원에 특화된 국부적 패치 재귀와 채널 상호작용을 강화함.
RSSB는 기존 Mamba의 토큰별 동일 처리 문제를 해결하기 위해 추가적인 컨볼루션과 채널 어텐션을 포함함.
MambaIR은 SwinIR 대비 0.36dB 높은 성능을 보이며, 유사한 계산 비용으로 전역 수용 필드를 유지함.
MambaIR은 CNN과 Transformer의 한계를 극복하며, 이미지 복원 분야에서 새로운 베이스라인을 제시함.

핵심 아이디어

기존 이미지 복원 모델은 CNN의 국부적 수용 필드 제약 또는 Transformer의 제곱 복잡도로 인해 한계를 겪는다. Mamba는 선형 복잡도로 장거리 의존성을 효과적으로 모델링하지만, 저수준 비전 작업에서의 국부적 퍼픽스 유사성과 채널 중복 문제를 해결하지 못한다. 본 연구는 Mamba에 Residual State Space Block (RSSB)을 도입하여, 컨볼루션과 채널 어텐션을 통해 국부적 패치 재귀와 채널 상호작용을 강화함으로써, Mamba의 장점을 유지하면서도 이미지 복원에 필요한 특성들을 보완한다. RSSB는 기존 Mamba의 토큰별 동일 처리 문제를 해결하고, 이미지 복원에 특화된 특징을 도입함으로써 성능을 향상시킨다.

기술적 접근법

**모델 구조**: MambaIR은 3단계로 구성됨 — 얕은 특징 추출, 깊은 특징 추출, 고해상도 이미지 재구성.
**RSSB (Residual State Space Block)**: MambaIR의 핵심 구성 요소로, 컨볼루션과 채널 어텐션을 추가하여 Mamba의 성능을 강화함.
**국부적 패치 재귀**: RSSB는 Mamba가 토큰을 동일하게 처리하는 문제를 해결하기 위해, 추가 컨볼루션을 통해 국부적 픽셀 유사성을 강조함.
**채널 어텐션**: 채널 간 상호작용을 강화하여 이미지 복원에 유리한 정보를 보존함.
**스킵 연결**: 각 블록 내부에 스킵 연결을 사용하여 특징 집약을 용이하게 함.
**Mamba의 선형 복잡도**: Mamba는 입력 크기와 선형 비례하는 계산 복잡도를 가지며, 이는 Transformer의 제곱 복잡도 문제를 해결함.

주요 결과

MambaIR은 SwinIR 대비 0.36dB 높은 성능을 보이며, 유사한 계산 비용으로 전역 수용 필드를 유지함.
RSSB는 Mamba에 국부적 패치 재귀와 채널 상호작용을 도입함으로써, 기존 Mamba의 토큰별 동일 처리 문제를 해결함.
실험 결과, RSSB 내 컨볼루션과 채널 어텐션의 결합이 MLP 기반의 변형보다 우수한 성능을 보임.
MambaIR은 다양한 이미지 복원 작업에서 기존 CNN 및 Transformer 기반 모델을 상회함.

의의 및 한계

MambaIR은 CNN과 Transformer의 한계를 극복하고, 이미지 복원 분야에서 새로운 베이스라인을 제시함. Mamba의 선형 복잡도와 전역 수용 필드를 유지하면서, RSSB를 통해 이미지 복원에 특화된 국부적 패치 재귀와 채널 상호작용을 강화함으로써, 기존 모델보다 뛰어난 성능을 보인다. 그러나 MambaIR은 주로 이미지 슈퍼해상도 작업에서 평가되었으며, 다른 저수준 비전 작업에서의 일반화 가능성은 추가 연구가 필요하다. 또한, MambaIR은 기존 Mamba 기반 모델과 비교했을 때, 복잡도가 약간 증가할 수 있으나, 여전히 선형 복잡도를 유지함.

실용적 활용

MambaIR은 이미지 슈퍼해상도, 노이즈 제거, 저해상도 이미지 복원 등 다양한 저수준 비전 작업에 적용 가능하다. 특히, 전역 수용 필드와 선형 복잡도를 유지하면서 높은 성능을 보이는 점에서, 실시간 이미지 처리나 모바일 기기에서의 활용이 기대된다. 또한, MambaIR은 기존 CNN 및 Transformer 기반 모델의 대안으로, 복잡도와 성능의 균형을 필요로 하는 연구 및 산업 분야에서 유용하게 사용될 수 있다.