[論文レビュー] UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation

Renkai Wu, Yinghao Liu+2 more|arXiv (Cornell University)|Mar 29, 2024

Cutaneous Melanoma Detection and Management被引用数 28

ひとこと要約

Parallel Vision Mamba (PVM Layer) に基づく高パラメータ効率の皮膚病変セグメンテーションモデル UltraLight VM-UNet を提案し、0.049M パラメータで公衆データセット上に競争力のある性能を達成。

ABSTRACT

Traditionally for improving the segmentation performance of models, most approaches prefer to use adding more complex modules. And this is not suitable for the medical field, especially for mobile medical devices, where computationally loaded models are not suitable for real clinical environments due to computational resource constraints. Recently, state-space models (SSMs), represented by Mamba, have become a strong competitor to traditional CNNs and Transformers. In this paper, we deeply explore the key elements of parameter influence in Mamba and propose an UltraLight Vision Mamba UNet (UltraLight VM-UNet) based on this. Specifically, we propose a method for processing features in parallel Vision Mamba, named PVM Layer, which achieves excellent performance with the lowest computational load while keeping the overall number of processing channels constant. We conducted comparisons and ablation experiments with several state-of-the-art lightweight models on three skin lesion public datasets and demonstrated that the UltraLight VM-UNet exhibits the same strong performance competitiveness with parameters of only 0.049M and GFLOPs of 0.060. In addition, this study deeply explores the key elements of parameter influence in Mamba, which will lay a theoretical foundation for Mamba to possibly become a new mainstream module for lightweighting in the future. The code is available from https://github.com/wurenkai/UltraLight-VM-UNet .

研究の動機と目的

• 計算リソースが限られるモバイル/臨床利用のために、軽量な医用画像分割を促進する。
• Visionタスクにおける Mamba のパラメータ削減が性能に与える影響を調査する。
• 精度を維持しつつパラメータの成長を抑える並列処理戦略（PVM Layer）を開発する。
• ISIC 2017、ISIC 2018、PH2 データセットで UltraLight VM-UNet の有効性を示す。

提案手法

• 入力チャネルを4等分し、VSS Blocks によって並列処理する Parallel Vision Mamba Layer (PVM Layer) を提案する。
• パラメータを最小化するよう、チャネル数を慎重に制御した Vision Mamba 系コア（SS2D、S4D および関連する射影）を使用する。
• マルチスケール特徴融合のためのChannel and Spatial Attention Bridgesを用いたSkip接続を備えるU-Net風のエンコーダ-デコーダを統合する。
• チャネル数と並列VSS Blockがパラメータと性能に与える影響を分析するアブレーションを実施する。
• 標準的なデータ拡張と BCE-Dice 損失を用いて3つの公的な皮膚病変データセットを評価し、DSC、SE、SP、ACC を報告する。

実験結果

リサーチクエスチョン

• RQ1 Mamba コンポーネントの入力チャネル数を減らすことは、パラメータ数と性能にどう影響するか？
• RQ2 並列処理方式（PVM Layer）は、パラメータを大幅に削減しつつ、セグメンテーション性能を維持または向上させることができるか？
• RQ3 ISIC2017、ISIC2018、PH2 データセットにおけるパラメータ削減とセグメンテーション精度のトレードオフはどのようになるか？
• RQ4 スキップ接続融合モジュール（CAB/SAB）は超軽量設定で性能に有意に寄与するか？

主な発見

• UltraLight VM-UNet は 0.049M パラメータと 0.060 GFLOPs を達成し、3つの皮膚病変データセットで競争力を維持する。
• 提案された PVM Layer は、VSS Block 経路で最大 93.1% のパラメータ削減を実現し、4つの並列ブロックに処理を分散して入力を4分の1チャネルにする。
• UltraLight VM-UNet は ISIC2017、ISIC2018、PH2 で DSC が約 0.909–0.926、高い ACC および堅牢な SE/SP 指標を達成する。
• アブレーションにより PVM Layer を標準畳み込みに置換するとパラメータが増え、性能が低下することが示され、並列 Vision Mamba 設計の重要性を裏付ける。
• VM-UNet および LightM-UNet と比較して、UltraLight VM-UNet はそれぞれパラメータを 99.82% および 87.84% 削減しつつ、競争力のある性能を維持する。