[論文レビュー] VmambaIR: Visual State Space Model for Image Restoration

Yuan Shi, Bin Xia+6 more|arXiv (Cornell University)|Mar 18, 2024

Image and Signal Denoising Methods被引用数 8

ひとこと要約

VmambaIRはUnetフレームワーク内の新規 Omni Selective Scan を活用し、状態空間モデルを用いた画像復元をモデル化することで、deraining、single-image super-resolution、real-world super-resolutionの分野で、計算量を抑えつつパラメータを減らした最先端の結果を提供します。

ABSTRACT

Image restoration is a critical task in low-level computer vision, aiming to restore high-quality images from degraded inputs. Various models, such as convolutional neural networks (CNNs), generative adversarial networks (GANs), transformers, and diffusion models (DMs), have been employed to address this problem with significant impact. However, CNNs have limitations in capturing long-range dependencies. DMs require large prior models and computationally intensive denoising steps. Transformers have powerful modeling capabilities but face challenges due to quadratic complexity with input image size. To address these challenges, we propose VmambaIR, which introduces State Space Models (SSMs) with linear complexity into comprehensive image restoration tasks. We utilize a Unet architecture to stack our proposed Omni Selective Scan (OSS) blocks, consisting of an OSS module and an Efficient Feed-Forward Network (EFFN). Our proposed omni selective scan mechanism overcomes the unidirectional modeling limitation of SSMs by efficiently modeling image information flows in all six directions. Furthermore, we conducted a comprehensive evaluation of our VmambaIR across multiple image restoration tasks, including image deraining, single image super-resolution, and real-world image super-resolution. Extensive experimental results demonstrate that our proposed VmambaIR achieves state-of-the-art (SOTA) performance with much fewer computational resources and parameters. Our research highlights the potential of state space models as promising alternatives to the transformer and CNN architectures in serving as foundational frameworks for next-generation low-level visual tasks.

研究の動機と目的

• CNN、Transformer、および拡散モデルの長距離依存性と効率性への対応における制約を解消し、画像復元の改善を促す。
• 2D画像データに対して線形計算量を持つ状態空間モデルベースのアーキテクチャを開発する。
• 情報の流れを六方向で捉えるためのOSS（Omni Selective Scan）ブロックを備えたマルチスケールのUNetを設計する。
• 低リソースでderaining、SR、real-world SRタスクにおけるVmambaIRの有効性を示す。

提案手法

• 提案されたOSSブロックを積み上げたUNet風アーキテクチャを採用する。
• 入力を2つのフローで処理し、特徴次元をマッピングするCNNを用いるOSSモジュールを導入する。
• 階層的情報の流れを調整するEfficient Feed-Forward Network(EFFN)を組み込む。
• 高周波成分のモデリングのため、六方向（双方向スキャンを含む三次元）で情報の流れをモデル化するOmni Selective Scan(OSS)を実装し、Mambaブロックを用いる。
• 画像特徴の効率的な系列モデリングのため、ZOHベースの離散化を用いた離散化済みState Space Model(SSM)を用いる。

実験結果

リサーチクエスチョン

• RQ1 線形複雑度の状態空間モデリングアプローチは、Transformer/CNNベースの画像復元性能と同等またはそれを超えることができるか？
• RQ2 Omni Selective Scanは、単方向のMambaブロックを超えて、画像における総合的で多方向の情報フローのモデリングを可能にするか？
• RQ3 SR、real-world SR、derainingタスクでの復元精度と効率性に対するOSS、双方向チャネルスキャン、およびEFFNの寄与は何か。

主な発見

• VmambaIRは、画像復元タスク全般で最先端の性能を達成しており、画像 deraining、single-image super-resolution、real-world image super-resolutionを含む。
• 実世界の4×超解像では、VmambaIRはベースライン手法の計算コストの約26%を使用しつつ、より高い再構成精度を実現する。
• アブレーション研究は、OSSが一方向スキャンより性能を大幅に改善し、双方向チャネルスキャンが精度を高め、EFFNが情報フローと効率を向上させることを示している。
• VmambaIRは、いくつかのベンチマークで既存のSOTA手法と比較して、より少ないパラメータとFLOPsで高周波の詳細をより良く提供します。
• 定性的な結果は、顔の目・鼻などのディテールの保存が向上し、欠陥が少ないことを示す。