MambaIRv2: Attentive State Space Restoration

The Mamba-based image restoration backbones have recently demonstrated significant potential in balancing global reception and computational efficiency. However, the inherent causal modeling limitation of Mamba, where each token depends solely on its predecessors in the scanned sequence, restricts the full utilization of pixels across the image and thus presents new challenges in image restoration. In this work, we propose MambaIRv2, which equips Mamba with the non-causal modeling ability similar to ViTs to reach the attentive state space restoration model. Specifically, the proposed attentive state-space equation allows to attend beyond the scanned sequence and facilitate image unfolding with just one single scan. Moreover, we further introduce a semantic-guided neighboring mechanism to encourage interaction between distant but similar pixels. Extensive experiments show our MambaIRv2 outperforms SRFormer by even 0.35dB PSNR for lightweight SR even with 9.3% less parameters and suppresses HAT on classic SR by up to 0.29dB. Code is available at https://github.com/csguoh/MambaIR.

한 줄 요약

MambaIRv2는 Mamba의 인과성 제약을 극복한 주목 기반 상태공간 복원 모델로, 단일 스캔과 0.35dB PSNR 개선을 달성한다.

핵심 기여도

**Attentive State-space Equation (ASE)**: Mamba의 상태공간 방정식에 프롬프트 학습을 결합하여, 이미지 전체의 유사 픽셀을 탐색 가능하게 함.
**Semantic Guided Neighboring (SGN)**: 의미 기반 픽셀 재구성으로, 멀리 떨어진 유사 픽셀 간 상호작용을 촉진.
**단일 스캔 기반 복원**: 기존 4방향 스캔 대신 1회 스캔으로 43% 파라미터 감소 및 50% 계산량 절감.
**성능 개선**: Urban100에서 SRFormer 대비 0.35dB, HAT 대비 0.29dB PSNR 상승.

핵심 아이디어

기존 Mamba는 인과적 상태공간 모델링으로 인해 이미지 내 모든 픽셀을 활용하지 못하며, 멀리 떨어진 픽셀 간 상호작용이 제한된다. 이를 해결하기 위해, MambaIRv2는 **ViT의 비인과적 처리 능력을 Mamba에 통합**하는 전략을 제안한다. 구체적으로, **Attentive State-space Equation (ASE)**를 통해 Mamba의 상태공간 출력 행렬 $ \mathbf{C} $에 프롬프트를 추가하여, 이미지 내 유사 픽셀을 탐색하는 메커니즘을 구현한다. 이는 **단일 스캔으로도 전역 정보를 수집**할 수 있게 하며, 기존의 다방향 스캔이 필요 없게 만든다. 또한, **Semantic Guided Neighboring (SGN)**을 통해 의미적으로 유사한 픽셀을 1D 시퀀스 상에서 가까이 배치하여, **장거리 감쇠(long-range decay)**를 완화한다. 이러한 아이디어는 Mamba가 주목 기반 복원을 수행하도록 유도하며, **비인과적 처리와 효율성을 동시에 달성**한다.

기술적 접근법

**모델 구조**: MambaIRv2는 ASE와 SGN을 핵심 모듈로 갖는 Attentive State-space Module (ASSM)을 사용.
**ASE**: 상태공간 방정식의 $ \mathbf{C} $ 행렬에 프롬프트를 잔차 합으로 추가하여, 이미지 전체의 유사 픽셀을 탐색.
**SGN**: 픽셀에 의미 라벨을 할당하고, 의미적으로 유사한 픽셀을 1D 시퀀스 상에서 가까이 배치.
**스캔 방식**: 기존 MambaIR의 4방향 스캔 대신, **단일 스캔**으로 처리.
**하이퍼파라미터**: MambaIRv2-light는 250K 이터레이션 동안 DIV2K 데이터셋에서 훈련.

주요 결과

**Urban100 데이터셋 (2× lightweight SR)**: SRFormer 대비 0.35dB PSNR 상승, 9.3% 적은 파라미터.
**Manga109 데이터셋 (2× classic SR)**: HAT 대비 0.29dB PSNR 상승.
**성능 효율성**: MambaIR 대비 43% 파라미터 감소, 50% 계산량 절감.
**개별 모듈 실험**: ASE 추가 시 0.05~0.09dB PSNR 상승, SGN 추가 시 최종 32.97/39.24 dB PSNR 달성.

의의 및 한계

MambaIRv2는 Mamba의 인과성 제약을 극복하고, **비인과적 처리와 효율성을 결합**함으로써 이미지 복원 분야에서 새로운 가능성을 제시한다. 특히, ASE와 SGN은 Mamba의 상태공간 모델링을 주목 기반으로 확장하여, **전역 정보 활용 능력을 향상**시키는 데 기여한다. 그러나, **프롬프트 기반 메커니즘의 일반화 능력**이나, **다양한 복원 태스크에서의 성능 일관성**은 추가 실험을 통해 검증이 필요하다. 또한, SGN은 의미 라벨링에 의존하므로, **라벨링 과정의 정확성과 계산 비용**이 한계로 작용할 수 있다.

실용적 활용

MambaIRv2는 고해상도 이미지 복원, JPEG 압축 해제, 노이즈 제거 등 다양한 이미지 복원 작업에 적용 가능하다. 특히, **단일 스캔 기반의 효율성**은 실시간 처리나 모바일 환경에서 유리하며, **낮은 파라미터 수량**은 에너지 효율적인 모델 설계에 기여할 수 있다. 연구적으로는 Mamba와 ViT의 결합 가능성 탐색에 기초 자료가 될 수 있다.