[论文解读] SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation

Zhaohu Xing, Ye Tian+3 more|arXiv (Cornell University)|Jan 24, 2024

Medical Image Segmentation Techniques被引用 9

一句话总结

SegMamba 在 U 形架构中使用 Mamba 状态空间块来建模 3D 医学影像中的长程依赖，与 CNN 和 transformer 基线相比，在推理效率方面实现了更优的 BraTS2023 分割表现。

ABSTRACT

The Transformer architecture has shown a remarkable ability in modeling global relationships. However, it poses a significant computational challenge when processing high-dimensional medical images. This hinders its development and widespread adoption in this task. Mamba, as a State Space Model (SSM), recently emerged as a notable manner for long-range dependencies in sequential modeling, excelling in natural language processing filed with its remarkable memory efficiency and computational speed. Inspired by its success, we introduce SegMamba, a novel 3D medical image extbf{Seg}mentation extbf{Mamba} model, designed to effectively capture long-range dependencies within whole volume features at every scale. Our SegMamba, in contrast to Transformer-based methods, excels in whole volume feature modeling from a state space model standpoint, maintaining superior processing speed, even with volume features at a resolution of {$64 imes 64 imes 64$}. Comprehensive experiments on the BraTS2023 dataset demonstrate the effectiveness and efficiency of our SegMamba. The code for SegMamba is available at: https://github.com/ge-xing/SegMamba

研究动机与目标

• 在高分辨率 3D 医学影像中推动全局、长程依赖的建模。
• 提出基于 Mamba 的编码器，以在多尺度上捕获整体体积特征。
• 结合带跳跃连接的 CNN 解码器实现精确分割。
• 在对大体积 3D 图像进行训练和推理时保持高效性。

提出的方法

• 提出三部分的 SegMamba 架构：一个包含多个模块的基于 Mamba 的编码器、一个基于 CNN 的解码器，以及在 U 形结构中的跳跃连接。
• 在 Mamba 块之前将 3D 特征展平为 1D 长序列以实现高效的序列建模，然后通过 σ 运算恢复到 3D。
• 在 Section 2.1 所定义的方程中，使用层归一化、Mamba 块和 MLP 构成的残差式块结构。
• 通过 7x7x7 的逐通道卷积对干 stem 进行下采样以产生多尺度特征 z^0，并在 Mamba 块中逐步处理。
• 使用交叉熵损失、SGD 优化器、多项式学习率调度和标准数据增强进行训练；在推理阶段应用测试时数据增强。

实验结果

研究问题

• RQ1 SegMamba 是否能够相较于基于 Transformer 的方法，使用 Mamba 高效地建模 3D 医学体积的长程相关性？
• RQ2 SegMamba 的编码器–解码器结构在 BraTS2023 上是否提高分割精度，同时在高分辨率（如 64x64x64 特征图）下保持推理速度？

主要发现

• SegMamba 在 BraTS2023 的 WT、TC、ET 三个分割指标上达到 state-of-the-art 的 Dice 分数（WT: 93.61，TC: 92.65，ET: 87.71），并且 HD95 分数为 WT: 3.37，TC: 3.85，ET: 3.48。
• 在 BraTS2023 的平均 Dice 为 91.32%，平均 HD95 为 4.01，在所报告的结果中超过基于 CNN 和 Transformer 的基线。
• 根据表 1，相较于 UX-Net 与 SwinUNETR-V2，平均 Dice 分别高出 1.63% 与 1.93% 。
• 展示基于 Mamba 的长程建模在 3D 医学图像分割中的高效性，在高分辨率（64x64x64）下保持竞争力的推理速度。