SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation

arxiv' 2024

paper: https://arxiv.org/pdf/2401.13560.pdf

code: https://github.com/ge-xing/SegMamba

Abstract

Transformer体系结构在建模全局关系方面显示出了非凡的能力。然而，它在处理高维医学图像时提出了重大的计算挑战。这阻碍了它在这项任务中的发展和广泛采用。Mamba作为一种状态空间模型(State Space Model, SSM)，近年来作为序列建模中一种值得关注的远程依赖关系模型，以其显著的存储效率和计算速度在自然语言处理领域表现优异。受其成功的启发，我们引入了SegMamba，一种新颖的3D医学图像分割Mamba模型，旨在有效地捕获每个尺度下整个体积特征中的远程依赖关系。与基于transformer的方法相比，从状态空间模型的角度来看，SegMamba在整体体积特征建模方面表现出色，即使体积特征的分辨率为64 × 64 × 64(序列长度约为260k)，也能保持卓越的处理速度。在BraTS2023数据集上的综合实验证明了该方法的有效性和高效性。可从https://github.com/ge-xing/SegMamba获得SegMamba的代码。

关键词:状态空间模型·曼巴·多序列建模·三维医学图像分割

1 Introduction

扩展模型的接受野是三维医学图像分割的一个关键方面。传统的卷积神经网络(cnn)在从高分辨率三维医学图像中提取全局信息方面不是很有效。因此，建议使用具有大核大小的深度卷积[13]来建模更广泛的特征。3D UX-Net[10]引入了一种新的架构，利用卷积(7 × 7 × 7)块和一个大内核来促进更大的接受域。然而，当像素之间的距离太大时，基于cnn的方法难以建立关系。

近年来，利用自关注模块提取全局信息的transformer架构被广泛应用于三维医学图像分割。TransBTS[19]结合3D-CNN提取局部空间特征，然后应用transformer对高级特征中的全局依赖关系进行建模。UNETR[5]采用Vision Transformer (ViT)[2]作为编码器来学习上下文信息，然后通过多个分辨率的跳过连接与基于cnn的解码器合并。SwinUNETR[4]利用SwinTransformer[12]作为编码器提取多尺度特征。同时设计了一个多尺度解码器，融合了每个编码器阶段的特征，在三维医学图像分割中取得了很好的效果。然而，对于以transformer为基础的方法，典型的3D医学图像的高分辨率会导致显著的计算负担和降低速度性能。

为了克服长序列建模的挑战，Mamba[3]源于状态空间模型(ssms)[8]，旨在通过选择机制和硬件感知算法对远程依赖关系进行建模，提高训练和推理的效率。许多研究探索了Mamba在计算机视觉(CV)中的应用。U-Mamba[14]将Mamba层集成到nnunet[7]的编码器中，以增强一般医学图像的分割。同时，Vision Mamba[21]提出了Vim块，该块结合了双向SSM，用于数据依赖的全局视觉上下文建模和位置嵌入，用于位置感知的视觉理解。此外，VMamba[11]设计了一个CSM模块来弥补一维阵列扫描和二维平面遍历之间的差距。在三维医学图像分割中，传统的transformer块在处理大尺寸特征时面临挑战。为了获得更强的视觉理解，有必要对高维特征中的相关性进行建模。受此启发，我们引入了SegMamba，这是一种新颖的框架，将u形结构与Mamba结合起来，在各种尺度上模拟整个体量的全局特征。据我们所知，这是利用Mamba专门用于3D医学图像分割的第一种方法。与传统的基于cnn的和基于transformer的方法相比，SegMamba在体积数据中表现出卓越的远程依赖关系建模能力，同时保持了出色的推理效率。大量的实验证明了该方法的有效性。

2 Method

SegMamba主要由三个部分组成:1)具有多个Mamba块的Mamba编码器，用于提取不同尺度的特征;2)基于卷积层的三维解码器，用于预测分割结果;3)跳过连接将多尺度特征连接到解码器，用于特征重用。图1说明了所提出的SegMamba的概述。我们在本节中进一步描述编码器和解码器的细节。

2.1 Mamba Encoder

全局特征和多尺度特征建模是三维医学图像分割的重要内容。虽然transformer体系结构可以提取全局信息，但当处理过长的特征序列时，它会产生显着的计算负担。为了减少序列长度，UNETR等基于transformer结构的方法直接对输入的三维医学图像进行下采样，分辨率为D*H*W至D/16 ×H/16 ×W/16。然而，这种方法限制了多尺度特征的建模能力，而多尺度特征对于通过解码器预测分割结果至关重要。为了克服这一限制，我们设计了Mamba块，它用更高效的Mamba层代替了transformer体系结构中的自关注模块。这使得多尺度和全局特征建模成为可能，同时在训练和推理过程中保持高效率。

如图1所示，Mamba编码器由一个stem层和多个Mamba块组成。对于stem层，我们采用深度卷积，内核大小为7 × 7 × 7，填充为3 × 3 × 3，步幅为2 × 2 × 2。给定一个三维输入体积I∈RC×D×H×W，其中C表示输入通道数，stem层提取第一个尺度特征z0∈R48× D/2 × H/2 × W/2。然后，z0通过每个Mamba块和相应的下采样层馈送。值得注意的是，SegMamba在Mamba块之前采用了平坦化操作ϕ，将3D特征重塑为1D长序列，从而实现了具有较少感应偏置的高效序列建模。在Mamba块后，通过σ运算将特征恢复到三维形状。这个过程可以定义为:

其中LN为归一化输入特征的层范数，Mamba为Mamba层，MLP为丰富特征表示的多层感知层。

2.2 Decoder

Mamba编码器提取了多尺度特征，在前人的研究基础上，我们利用基于cnn的解码器和跳跃连接形成U-shape网络来预测分割结果。

3 Experiments

3.1 Datasets

BraTS2023 dataset : BraTS2023数据集[15,1,9]共包含1,251个3D脑MRI体积。每个体积包括四种模式(T1, T1Gd, T2, T2- flair)和三个分割目标(WT:整个肿瘤，ET:增强肿瘤，TC:肿瘤核心)。所有数据都被重新采样到相同的间距(1.0,1.0,1.0)。

3.2 Evaluation Metrics

基于不同数据集的特征，我们采用各种指标来比较我们的方法与其他方法的性能。

Dice similarity coefficient (Dice score)：Dice分数是一个重叠度量，用来衡量预测和真实之间的重叠百分比。计算公式如下:

式中，X表示语义预测，Y表示GT，|·|表示基数计算操作。

95% Hausdorff distance (HD95)： HD95度量标准测量Hausdorff距离的第95个百分位数，为预测和实际情况GT之间的最大距离提供可靠的评估。由下式定义:

式中，d(X, Y)表示X与Y之间的Hausdorff Distance ，maxk95%表示第95百分位处的最大值。

3.3 Comparison Methods

为了进行全面的评估，将SegMamba与其他七种最先进的方法进行了比较，这些方法都涵盖了CNN和transformer架构。

CNN-based methods。 我们将SegMamba与SegresNet[16]、UX-Net[10]和MedNeXt[17]进行了比较。其中，UX-Net和MedNeXt是最新的3D医学图像分割方法。

Transformer-based methods。我们还比较了SegMamba与最新的基于transformer的在医学成像上的分割方法。UNETR[5]、SwinUNETR[4]、SwinUNETR- v2[6]和nnFormer[20]是最著名的基于transformer的三维医学图像分割方法。他们使用Vision Transformer[2]或SwinTransformer[12]结构作为编码器来建模全局特征。

为了确保公平的比较，我们在实验中使用了所有方法的公开代码，并保持所有设置相同。

3.4 Implementation Details

我们的模型在Pytorch 2.0.1-cuda11.7和Monai 1.2.0中实现。在训练过程中，我们对每个数据集使用128 × 128 × 128的随机裁剪大小和每个GPU bs为2的批处理。由于BraTS2023数据集中的每个体积包含4个模态，因此我们在网络的输入处将通道维度中的每个模态连接起来。我们对所有实验使用交叉熵损失和SGD优化器以及多项式学习率调度器(初始学习率为1e-2，衰减为1e-5)。我们对所有数据集运行1000个epoch，并采用以下数据增强:添加亮度，伽马，旋转，缩放，镜像和弹性变形。所有实验都是在一个云计算平台上进行的，该平台配备了4个NVIDIA A100 gpu。我们将测试时间增强(TTA)技术[18](即镜像预测和重叠滑动窗口推断[7])应用于所有三个数据集的推理，并将重叠比设置为0.5。对于每个数据集，我们随机分配70%的3D体积用于训练，10%用于验证，剩下的20%用于测试，确保每个体积在训练、验证和测试集中只出现一次。

3.5 Quantitative Comparison to Previous Methods

BraTS2023 brain tumor segmentation 。BraTS2023数据集的胶质瘤分割结果如表1所示，其中我们使用Dice评分和HD95来评价在WT、TC和ET三个分割目标上的表现。基于cnn的UX-Net方法在比较方法中表现最好，其平均Dice评分为89.69%，平均HD95为4.81。在基于transformer的方法中，swinunetrt - v2也表现出良好的性能，平均Dice得分为89.39%，平均HD95为4.51。相比之下，我们的SegMamba在WT, TC和ET上的Dice得分分别为93.61%，92.65%和87.71%，HD95得分分别为3.37,3.85,3.48，优于所有其他方法。我们的SegMamba的平均Dice得分为91.32%，比排名第二的UX-Net和排名第三的SwinUNETRv2分别提高了1.63%和1.93%。