[論文レビュー] Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining

Jiarun Liu, Hao Yang+9 more|arXiv (Cornell University)|Feb 5, 2024

Vehicle License Plate Recognition被引用数 6

ひとこと要約

Swin-UMamba は 2D 医用画像分割のための Mamba ベースの UNet で、ImageNet 事前学習済みの Mamba ブロックを活用して、複数のデータセットで精度と効率を向上させる。軽量な Mamba ベースのデコーダを持つ Swin-UMamba dagger というバリアント。

ABSTRACT

Accurate medical image segmentation demands the integration of multi-scale information, spanning from local features to global dependencies. However, it is challenging for existing methods to model long-range global information, where convolutional neural networks (CNNs) are constrained by their local receptive fields, and vision transformers (ViTs) suffer from high quadratic complexity of their attention mechanism. Recently, Mamba-based models have gained great attention for their impressive ability in long sequence modeling. Several studies have demonstrated that these models can outperform popular vision models in various tasks, offering higher accuracy, lower memory consumption, and less computational burden. However, existing Mamba-based models are mostly trained from scratch and do not explore the power of pretraining, which has been proven to be quite effective for data-efficient medical image analysis. This paper introduces a novel Mamba-based model, Swin-UMamba, designed specifically for medical image segmentation tasks, leveraging the advantages of ImageNet-based pretraining. Our experimental results reveal the vital role of ImageNet-based training in enhancing the performance of Mamba-based models. Swin-UMamba demonstrates superior performance with a large margin compared to CNNs, ViTs, and latest Mamba-based models. Notably, on AbdomenMRI, Encoscopy, and Microscopy datasets, Swin-UMamba outperforms its closest counterpart U-Mamba_Enc by an average score of 2.72%.

研究の動機と目的

• 長距離依存性をモデリングする必要性を、局所的な CNN の受容野や高コストな ViT を超えて医用画像分割に対して動機づける。
• ImageNet ベースの事前学習に適合した Mamba ベースのエンコーダを導入する。
• 正確な分割のためのスキップ接続と深い監督を強化した U-Net 風デコーダを設計する。
• 効率性のために軽量な Swin-UMamba dagger バリアントを Mamba ベースのデコーダで提案する。
• Mamba ブロックを用いたデータ効率的な分割を可能にするための事前学習の重要性を示す。

提案手法

• 2D 医用画像の特徴を抽出する ImageNet で事前訓練された Mamba ベースのエンコーダを採用する。
• 長距離依存性を扱う 2D selective-scan VSS (SS2D) ブロックを使用する。
• エンコーダ初期化に ImageNet の事前学習を活用するため VMamba-Tiny と重みを共有する。
• スキップ接続と深い監督を備えた U 字型アーキテクチャの Swin-UMamba デコーダを構築する。
• パッチ展開とパラメータ/ FLOPs を削減した軽量デコーダ変種 Swin-UMamba dagger を提供する。
• Dice + クロスエントロピー損失、深い監督、AdamW、コサイン減衰、事前学習重みの段階的凍結で訓練する。

実験結果

リサーチクエスチョン

• RQ1 ImageNet ベースの事前学習は医用画像分割における Mamba ベースモデルの性能を向上させるか？
• RQ2 Swin-UMamba は CNN-, ViT-, および他の Mamba ベースの分割モデルと比較して多様なデータセットでどのように比較されるか？
• RQ3 Mamba ベースのデコーダ（Swin-UMamba dagger）はより少ないパラメータと FLOPs で競争力のある結果を達成できるか？
• RQ4 事前学習は Mamba ブロックを用いた医用分割の収束安定性とデータ効率にどのような影響を与えるか？
• RQ5 長距離モデリング能力が 2D 医用画像タスクにどれだけ翻訳されるか？

主な発見

• Swin-UMamba および Swin-UMamba dagger は AbdomenMRI、Endoscopy、Microscopy データセットで CNN-, ViT-, および従来の Mamba ベースモデルを上回る。
• ImageNet ベースの事前学習は Swin-UMamba の DSC（AbdomenMRI で約 3.04 ポイント）と NSD（約 4.19 ポイント）の顕著な向上をもたらす。
• 事前学習は収束の速さと訓練の安定性を可能にし、Swin-UMamba は AbdomenMRI でベースラインよりはるかに少ない反復で訓練できる。
• Swin-UMamba dagger はパラメータと FLOPs が大幅に少ない（27M パラメータ、15.0G FLOPs）一方、Swin-UMamba（40M、58.4G）および U-Mamba 系より競争力のある結果を達成。
• Endoscopy では、事前学習済み Swin-UMamba dagger が非事前学習版に対して顕著な改善をもたらし、事前学習のデータ効率の利点を強調している。