[論文レビュー] MedMamba: Vision Mamba for Medical Image Classification

Yubiao Yue, Zhenzhang Li|arXiv (Cornell University)|Mar 6, 2024

COVID-19 diagnosis using AI被引用数 65

ひとこと要約

MedMamba は Vision Mamba ベースのアーキテクチャを導入し、SS-Conv-SSM ブロックで局所特徴と長距離特徴を医用画像分類のために捉える。複数のモダリティで評価し、競争力のある結果。

ABSTRACT

Since the era of deep learning, convolutional neural networks (CNNs) and vision transformers (ViTs) have been extensively studied and widely used in medical image classification tasks. Unfortunately, CNN's limitations in modeling long-range dependencies result in poor classification performances. In contrast, ViTs are hampered by the quadratic computational complexity of their self-attention mechanism, making them difficult to deploy in real-world settings with limited computational resources. Recent studies have shown that state space models (SSMs) represented by Mamba can effectively model long-range dependencies while maintaining linear computational complexity. Inspired by it, we proposed MedMamba, the first Vision Mamba for generalized medical image classification. Concretely, we introduced a novel hybrid basic block named SS-Conv-SSM, which purely integrates the convolutional layers for extracting local features with the abilities of SSM to capture long-range dependencies, aiming to model medical images from different image modalities efficiently. By employing the grouped convolution strategy and channel-shuffle operation, MedMamba successfully provides fewer model parameters and a lower computational burden for efficient applications without sacrificing accuracy. We thoroughly evaluated MedMamba using 16 datasets containing ten imaging modalities and 411,007 images. Experimental results show that MedMamba demonstrates competitive performance on most tasks compared with the state-of-the-art methods. This work aims to explore the potential of Vision Mamba and establish a new baseline for medical image classification, thereby providing valuable insights for developing more powerful Mamba-based artificial intelligence algorithms and applications in medicine. The source codes and all pre-trained weights of MedMamba are available at https://github.com/YubiaoYue/MedMamba.

研究の動機と目的

• CNN の医用画像における長距離依存性の捉え方の制限を解決する。
• 局所モデリングと長距離モデリングを組み合わせた軽量で効率的なアーキテクチャを提案する。
• 多様な医用画像モダリティで MedMamba を評価し、SSM ベースの医用画像分類の新しいベースラインを確立する。

提案手法

• 畳み込み局所特徴と状態空間モデルの長距離依存性を融合する SS-Conv-SSM モジュールを導入する。
• SS2D（2D 選択スキャン）を開発し、複数方向の特徴系列を走査して統合する。
• パッチ埋め込み層（4×4 パッチ）を使用し、チャネル拡張を伴う SS-Conv-SSM ブロックを4段階適用する。
• レイヤー正規化を適用し、チャネルシャッフルで相互作用をとる双方向ブランチ構造（Conv ブランチと SSM ブランチ）。
• データ拡張や事前学習なしで訓練し、アーキテクチャ推進の利得を示す。
• 画像データのエンドツーエンドのモデリングを目的として、離散化された状態空間モデリング（Abar, Bbar）と畳み込みカーネルアセンブリ（Kbar）による理論的根拠を提供する。

実験結果

リサーチクエスチョン

• RQ1 状態空間モデルと畳み込み特徴を統合することは、モダリティを跨いだ医用画像分類にどのような影響を与えるか？
• RQ2 事前学習やデータ拡張なしで、MedMamba は多様な医用データセットで競争力のある性能を達成できるか？
• RQ3 SS2D 選択走査機構が特徴抽出と長距離依存性モデリングに与える影響は何か？

主な発見

• MedMamba は CNN ベースおよび Transformer ベースのベースラインと比較して、競争力があり、いくつかのデータセットでは優れた性能を示す。
• このアーキテクチャは皮膚病変、超音波、胸部X線、内視鏡、MedMNIST由来のタスクを含む複数のモダリティで良好に機能する。
• SS-Conv-SSM は局所的な特徴（畳み込み）とグローバルな特徴（SSM）を効果的に組み合わせ、医用画像の微細情報と粗大情報を捉える。
• 5つの公開/非公開医療データセットでの広範な実験は、MedMamba が SSM ベースの医用画像分類の新しいベースラインを確立できることを示唆する。
• 再現とさらなる研究のためのソースコードが公開されている。