[論文レビュー] MambaUIE&SR: Unraveling the Ocean's Secrets with Only 2.8 GFLOPs

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

Oceanographic and Atmospheric ProcessesEarth and Planetary Sciences被引用数 3

ひとこと要約

本稿では、状態空間モデル（SSM）アーキテクチャであるMambaを基盤とする、MambaUIE&SRと呼ばれる新しい水中画像増強モデルを提案する。わずか2.8 GFLOPsの計算コストで最先端の性能を達成している。視覚的状態空間（VSS）ブロック、動的相互作用ブロック（DIB）、空間ゲート付きフィードフォワードネットワーク（SGFN）を統合することで、グローバルな文脈とローカルな細粒度特徴を効率的に捉え、最先端の手法と比較して計算コストを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で性能を発揮できない理由を解明し、ローカルな細粒度特徴をより効果的に保持する手法を提案すること。
• グローバルな文脈モデリングとローカルな詳細回復の両立を図る、軽量でパラメータ効率の良いアーキテクチャを設計すること。
• FLOPsを著しく削減しながらUIE分野で最先端の性能を達成し、効率性と正確性のトレードオフを打ち破ること。

提案手法

• 長距離依存性をモデル化し、グローバルな文脈を捉えるために、視覚的状態空間（VSS）ブロックを統合したカスタムMambaベースのアーキテクチャ、MambaUIEを提案する。
• ブロック内での特徴の適応的相互作用を可能にするために、動的相互作用ブロック（DIB）を導入する。これにより、ブロック内での特徴集約が強化される。
• ローカル特徴学習を強化し、細粒度のディテールを保持するために、空間ゲート付きフィードフォワードネットワーク（SGFN）を採用する。
• VSSブロックとDIB、SGFNを階層的なエンコーダ-デコーダ構造に統合し、グローバル情報とローカル情報の処理のバランスを取る。
• MambaのSSM機構の線形計算複雑度を活用することで、低FLOP動作を最適化する。
• UIEBデータセット上でエンドツーエンドに学習させ、統合的な増強およびスーパーレゾリューション（SR）能力を学習する。

実験結果

リサーチクエスチョン

• RQ1 Mamba状態空間モデルは、高い効率性を実現するために水中画像増強に効果的に適応可能か？
• RQ2 Mambaの直接的適用がUIEで失敗する理由は何か？また、ローカルな細粒度特徴をより効果的に保持するにはどうすればよいか？
• RQ3 SSMとローカル特徴学習メカニズムを組み合わせたハイブリッドアーキテクチャは、低FLOPコストで最先端の性能を達成できるか？
• RQ4 提案された動的相互作用ブロック（DIB）は、標準的なアテンションやフィードフォワード機構と比較して、特徴集約をどのように改善するか？
• RQ5 SSMベースのモデルを活用することで、UIEにおける正確性-FLOPトレードオフを打ち破ることは可能か？

主な発見

• MambaUIE&SRはわずか2.8 GFLOPsのFLOP数を達成し、最先端の手法と比較して計算コストを67.4％削減した。
• モデルはUIEBデータセット上で定量的指標（例：PSNR、SSIM）および定性的な視覚的結果の両面で、既存の最先端手法を上回った。
• DIBとSGFNの統合により、特徴集約が顕著に向上し、テクスチャーや色のディテールがよりよく保持された。
• 提案されたアーキテクチャは、少数のパラメータで高い正確性を達成しており、優れたパラメータ効率性を示した。
• これは、SSMに基づく最初のUIEモデルであり、水中画像増強分野で伝統的なFLOPsと正確性のトレードオフを成功裏に打ち破った。
• 本手法は、多様な水中画像劣化パターンにわたり、優れた汎化性とロバストネスを示した。