[論文レビュー] MambaU-Lite: A Lightweight Model based on Mamba and Integrated Channel-Spatial Attention for Skin Lesion Segmentation

Thi‐Nhu‐Quynh Nguyen, Quang Tri Ho+4 more|arXiv (Cornell University)|Dec 2, 2024

Cutaneous Melanoma Detection and Management被引用数 4

ひとこと要約

MambaU-Lite は、Mamba の効率的なシーケンスモデリングと CNN、および統合されたチャネル・スパティアルアテンションを備えた新規 P-Mamba ブロックを統合した軽量な医用画像セグメンテーションモデルである。ISIC2018 および PH2 データセットにおいて、わずか 420K パラメータと 1.25G FLOPs で最先端の性能を達成し、計算コストとメモリコストを最小限に抑えながら、Dice スコアと IoU スコアで既存の軽量モデルを上回っている。

ABSTRACT

Early detection of skin abnormalities plays a crucial role in diagnosing and treating skin cancer. Segmentation of affected skin regions using AI-powered devices is relatively common and supports the diagnostic process. However, achieving high performance remains a significant challenge due to the need for high-resolution images and the often unclear boundaries of individual lesions. At the same time, medical devices require segmentation models to have a small memory foot-print and low computational cost. Based on these requirements, we introduce a novel lightweight model called MambaU-Lite, which combines the strengths of Mamba and CNN architectures, featuring just over 400K parameters and a computational cost of more than 1G flops. To enhance both global context and local feature extraction, we propose the P-Mamba block, a novel component that incorporates VSS blocks along-side multiple pooling layers, enabling the model to effectively learn multiscale features and enhance segmentation performance. We evaluate the model's performance on two skin datasets, ISIC2018 and PH2, yielding promising results. Our source code will be made publicly available at: https://github.com/nqnguyen812/MambaU-Lite.

研究の動機と目的

• 計算コストとメモリコストを最小限に抑えつつ、高い精度を維持する軽量なディープラーニングモデルを皮膚腫瘍セグメンテーション用に開発すること。
• 曖昧な境界と高解像度の入力画像を有する皮膚腫瘍のセグメンテーションに、効率的なアーキテクチャを用いて挑戦すること。
• Mamba の長距離モデリング能力を CNN と統合し、グローバルな文脈とローカル特徴抽出の両方を向上させること。
• リソース制限のある医療デバイスへのデプロイを想定し、モデルサイズと推論コストを低減すること。
• 新規なハイブリッドアーキテクチャを用いて、ISIC2018 や PH2 のようなベンチマークデータセットにおけるセグメンテーション性能を向上させること。

提案手法

• Mamba と CNN のコンポONENTを組み合わせたハイブリッド U-Net ライクなアーキテクチャを提案し、コアモジュールとして P-Mamba ブロックを採用する。
• P-Mamba ブロックを導入し、VSS（ベクトル選択スキャン）ブロックと複数のプーリング層を統合することで、マルチスケール特徴を捉える。
• スキップ接続に CBAM（畳み込みブロックアテンションモジュール）を採用し、空間的およびチャネルワイドな特徴表現を強化する。
• セグメンテーション性能を最適化するために、Dice 損失と Tversky 損失を組み合わせたデュアル損失関数を採用し、ハイパーパramータ γ₁=0.7 と γ₂=0.3 を使用する。
• U-Lite をインspired した軽量設計により、パラメータ数を 0.42M に削減し、メモリ使用量を 1.67MB に抑える。
• Adam 最適化手法を用い、学習率スケジューリングにコサインアニーリングを採用し、データオーグメンテーションを適用してモデルを訓練する。

Figure 1 : The architecture of the proposed MambaU-Lite model

実験結果

リサーチクエスチョン

• RQ1 ハイブリッド Mamba-CNN アーキテクチャは、パラメータ数と FLOPs を最小限に抑えながら、皮膚腫瘍分析で高いセグメンテーション精度を達成できるか？
• RQ2 提案された P-Mamba ブロックは、標準の Mamba や CNN ブロックと比較して、特徴学習をどのように向上させるか？
• RQ3 スキップ接続にチャネル・スパティアルアテンションを統合することで、挑戦的な皮膚腫瘍データセットにおけるセグメンテーション性能が向上するか？
• RQ4 MambaU-Lite は、既存の軽量モデルと比較して、精度、パラメータ数、推論コストの観点で優れているか？
• RQ5 MambaU-Lite は、アーキテクチャの複雑さを最小限に抑えながら、ISIC2018 や PH2 のような多様な皮膚腫瘍データセットに一般化できるか？

主な発見

• ISIC2018 データセットでは、MambaU-Lite は DSC 0.9057、IoU 0.8361 を達成し、U-Net、Attention U-Net、UNeXt、DCSAU-Net、U-Lite を上回った。
• PH2 データセットでは、MambaU-Lite は DSC 0.9572、IoU 0.9189 を達成し、軽量モデルの中で新たな最先端性能を樹立した。
• モデルはたった 0.42 百万パラメータ、1.67MB のメモリ使用量にとどまり、U-Lite（0.88M パラメータ、3.51MB）や他のモデルと比べて顕著に小型化されている。
• U-Lite よりも FLOP 数が高い（1.25G 対 0.69G）にもかかわらず、MambaU-Lite は優れた精度を達成しており、より優れた効率-精度トレードオフを示している。
• P-Mamba ブロックはマルチスケール特徴を効果的に捉えており、境界が不明瞭な病変のセグメンテーション向上に寄与している。
• スキップ接続に CBAM を統合することで、特徴の精錬が促進され、より正確なマスク予測が得られた。

Figure 2 : The main components’ architectures of the proposed MambaU-Lite model