[論文レビュー] MambaIR: A Simple Baseline for Image Restoration with State-Space Model

Hang Guo, Jinmin Li+4 more|arXiv (Cornell University)|Feb 23, 2024

Image and Signal Denoising Methods被引用数 18

ひとこと要約

MambaIR は、選択的状態空間モデル（Mamba）に基づく単純で強力な画像復元バックボーンを、復元事前知識を強化して導入し、線形計算量で 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を越えた低レベルの画像復元において、状態空間モデルの採用を動機づける。
• Mamba を局所的な空間事前情報とチャネル相互作用で強化する復元特化の RSSB を提案する。
• 複数の復元タスクで、類似の計算コストで MambaIR が 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 の因果処理を2D画像データに拡張するため、2D Selective Scan Module を採用する。
• グローバル受容野を維持する、非注意機構の線形時間アプローチを強調する。

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 局所パッチ再発とチャネル相互作用といった復元事前知識は、低レベル視覚タスクにおける Mamba の有効性をどのように向上させるか？
• RQ3 単純な RSSB ベースの設計は、類似の計算コストで Transformer ベースの復元モデルと競合する性能を達成できるか？

主な発見

• MambaIR は競争力のある復元性能を達成し、SwinIR などの Transformer ベースのベースラインを、類似の計算コストでいくつかのタスクで上回る。
• RSSB の設計は、局所パッチ再発とチャネル相互作用を Mamba バックボーンに組み込むことで性能を向上させる。
• Vision State-Space Module と 2D-Selective Scan は、線形計算量で長距離モデリングを実現し、グローバルな受容野を維持する。
• 本手法は、古典的な SR、軽量 SR、実画像ノイズ除去の複数データセットで強い結果を示す。
• アブレーション研究は、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).