[论文解读] 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-selective-scan）用于扫描并合并多方向特征序列。
• 使用分块嵌入层（4x4 公分）后跟随四个阶段的 SS-Conv-SSM 模块并进行通道扩展。
• 应用层归一化、双分支结构（卷积分支和 SSM 分支）并通过通道洗牌实现交互。
• 在不使用数据增强或预训练的情况下进行训练，以证明架构驱动的增益。],
• research_questions":["将状态空间模型与卷积特征相结合对不同模态的医学影像分类有何影响？","在多样化的医学数据集上，MedMamba 能否在不进行预训练或数据增强的情况下实现具有竞争力的性能？","SS2D 选择性扫描机制对特征提取和长程依赖建模有何影响？"]
• key_findings":["MedMamba 在与 CNN 和 Transformer 基线相比时，表现具有竞争力，且在若干数据集上甚至更优。","该架构在多模态下均表现良好，包括皮肤病变、超声、胸部 X 光、内镜，以及来自 MedMNIST 的任务。","SS-Conv-SSM 有效地将局部（卷积）与全局（SSM）特征结合，捕捉医学影像中的细粒度与粗粒度信息。","在五个公开/私有医学数据集上的大量实验表明，MedMamba 可以为基于 SSM 的医学影像分类建立新的基线。","源代码公开，便于复现与进一步研究。

实验结果

研究问题

• RQ1 将状态空间模型与卷积特征相结合对不同模态的医学影像分类有何影响？
• RQ2 在多样化的医学数据集上，MedMamba 能否在不进行预训练或数据增强的情况下实现具有竞争力的性能？
• RQ3 SS2D 选择性扫描机制对特征提取和长程依赖建模有何影响？

主要发现

• MedMamba 在与 CNN 和 Transformer 基线相比时，表现具有竞争力，且在若干数据集上甚至更优。
• 该架构在多模态下均表现良好，包括皮肤病变、超声、胸部 X 光、内镜，以及来自 MedMNIST 的任务。
• SS-Conv-SSM 有效地将局部（卷积）与全局（SSM）特征结合，捕捉医学影像中的细粒度与粗粒度信息。
• 在五个公开/私有医学数据集上的大量实验表明，MedMamba 可以为基于 SSM 的医学影像分类建立新的基线。
• 源代码公开，便于复现与进一步研究。