近年来，3D医学分割领域主要由采用卷积神经网络（CNNs）和Transformer-based架构的深度学习模型主导，各有其独特的优势和局限性。
CNNs受限于局部感受野，而Transformer受制于其巨大的内存需求和数据饥饿性，使其不适合处理3D医学体积的细粒度信息。因此，在分割3D医学结构的大型医学体积时，卷积全卷积神经网络（如nnU-Net）仍然主导市场。
尽管朝着开发具有亚平方时间与内存复杂度的Transformer变体迈进，但这些模型在内容推理方面仍有所欠缺。最近，Mamba成为了一个基于状态空间模型（SSMs）的循环神经网络（RNN），在许多长文本任务（百万字符序列）上超过了Transformer，同时保持了线性复杂度。
在本文中，作者评估了Mamba-based架构与最先进的卷积和Transformer-based模型在3D医学图像分割方面的有效性，数据集包括Synapse Abdomen、MSD Brain Tumor和ACDC。
此外，作者通过提出替代的架构设计来解决现有Mamba-based架构的主要局限性，从而提高分割性能。
1 Introduction
图像分割在医学影像分析中至关重要，通常作为检查解剖结构和手术计划的前期步骤 [2]。近年来，卷积神经网络（CNN） [25] 以及特别是 U-shaped 全卷积神经网络（FCNN） [31] 在研究社区中得到了广泛应用。尽管它们具有有效性，但自视觉 Transformer （vision transformers）爆发以来，FCNN已被由卷积和多头自注意力层组成的混合架构所取代，这些混合架构旨在减轻卷积操作的特征，即局部感受野问题，并依赖 Transformer 的注意力机制 [36]。文献中已经尝试在经典的 U-Net 中集成基于 Transformer 架构。即使这些方法在性能上有所提高，但注意力机制的平方级内存占用以及与数据需求相伴随的注意力机制的计算密集性使得这些方法在大规模 3D 体积的应用中并不理想。
在这方面，最新的研究投入了大量精力来降低 Transformer 架构的计算成本，提出了线性注意机制和门控[34]。然而，在上下文建模方面，它们大多在上下文长度相当高时表现不足。最近，序列建模领域受到了一种基于状态空间模型（SSM）[16]的创新架构的影响，该架构称为Mamba[14]，它在自然语言处理（NLP）和基因组任务中展示了最先进的性能，改进了高达百万个 Token 的大上下文建模，使其成为高效处理3D卷的合适候选者，其中 Token 的数量也达到了同样的数量级。
本论文旨在通过与卷积和Transformer基础分割模型对比，研究Mamba在3D图像分割方面的有效性。此外，作者还试图通过在U-Net基础架构中提出各种策略来解决现有Mamba基础架构的主要局限性。
具体来说，作者研究了在某一或多个轴上建模方向性对结果的影响，并探索了将Mamba作为跳接连接的 selective copying 机制的应用。为了进行实验评估，作者采用了三个不同知名的 dataset，即MSD BrainTumour [1]，Synapse Multi-organ [24] 和 ACDC [4]。代码已公开以鼓励进一步的研究。
2 Related Work
卷积神经网络（CNNs）[25]多年来一直是2D和3D医学图像分割的主导解决方案。在这些模型中，U-Net[31]以其U形对称编码器-解码器结构以及 Short-Cut ，代表了后续模型一直采用的有效架构，直到今天。U-Net之后，出现了几种变体，包括Res-U-Net[10]，Dense-U-Net[6]，V-Net[29]，3D U-Net和其最新的生态系统nnU-Net[22]，每个模型都为原始框架提出了改进。尽管它们取得了进步，但由于卷积运算符的局部性，CNN在捕捉全局模式方面存在固有限制。为了应对这一问题，大量研究致力于将 Transformer 的注意力机制与U-Net基础架构相结合。这种结合旨在利用局部和全局依赖性，如MedFormer[12]，TransUNet[8]，Swin-UNet[7]，UNETR[19]，以及Swin-UNetR[18]。然而，注意力机制的二次复杂性使得必须施加约束，例如基于窗口或轴的注意力，以降低计算需求。尽管有许多研究试图降低这种复杂性[3，23，38，9]，但至今尚无方法在长序列建模中与传统注意力机制相媲美。
最近的发展引入了一种名为Mamba的新架构，该架构基于状态空间建模[16, 17]，承诺具有线性时间复杂度的长上下文内容推理能力。Mamba在需要长上下文内容推理的任务上，如自然语言处理和基因组分析等，相较于最先进的 Transformer 模型（如Pythia-6.9B[5]、GPT-J-6B[37]、OPT-6.7B[45]、Hyena[30]）表现出了优越的性能。
由于其有效性和通用性，基于Mamba的架构已迅速应用于各种领域，包括计算机视觉[26]。此外，考虑到将3D体积分割可以被视为处理由数百万个 Voxel 组成的序列，许多研究行人致力于将Mamba架构应用于2D和3D分割，取得了有前途的结果。
在这些贡献中，UMamba在领域内仍然是最重要的，因为该模型能够有效地适应新数据集而无需进行广泛的超参数调优。特别地，作者提出了一种架构，名为UMambaEnc和UMambaBot，这两种架构都继承了U-Net的核心结构，并利用了基于Mamba的层。前者将Mamba层集成到架构的编码部分，而后者将单个Mamba层集成到 Bottleneck 中。
尽管它们的有效性，作者并没有关注从使用循环网络提取具有多个空间维度的数据中衍生出的方向性问题。实际上，一旦将3D体积展开为一个序列，每个 Voxel 都被分配到一个序列中的位置。这导致模型能够通过利用前一部分的信息来分析序列的后续部分。然而，在处理序列的开 Head 分时，它缺乏上下文信息。
3 Method
在本节中，作者介绍了与 vanilla Mamba 架构（一套堆叠的 Mamba 模块）相关的所有理论概念。然后，作者详细解释了如何使用 Mamba 模块从 3D 体积中提取模式，并说明了将此类模块集成到 U-Net 架构中进行 3D 医学影像分割的方法。
Preliminaries
状态空间模型。 状态空间模型（SSM）是表示动态系统的数学表示，它将一维输入信号 在将其转换为 维潜在状态 之前映射到一维输出信号 。该系统使用 作为演化参数， 和 作为投影参数：
结合之前的方程，作者的目标是根据观察到的数据预测系统的状态。由于输入被预期为连续的，SSM的主要表示是连续时间表示。
为了在实际场景中应用等式（2），并更具体地说，将其应用到神经网络中，需要对变量进行离散化，这可以通过引入步长参数和离散化规则来实现。在本文中，这种离散化规则被称为“零阶保持”。
这导致了以下离散状态空间模型，可以采用递归方式进行计算： 由于在序列长度中梯度呈指数增长，这种基本的SSM在实际应用中表现非常糟糕。为了解决这个问题，Mamba提出了两个关键要素：对矩阵A施加结构，使用HiPPO理论[15]，并包括一个选择机制，即通过线性投影使参数B、C和Δ成为输入相关：
这种公式，再加上通过选择性扫描算法实现的流程的有效实现，使得模型能够过滤掉无关信息，构成了所谓的 S6 模型 [14]。为了构建 Mamba 块，首先使用线性投影来扩展输入 Token 的嵌入。然后，在 S6 之前应用卷积，以防止独立的输入 Token 计算，最后将 SSM 输出投影回原始维度，并采用 Shortcut 。作者使用 SiLU 非线性 [20]。这个块的视觉表示如图1 中的灰色部分所示。关于讨论的这种方法，请参阅原始出版物 [14]。
### Vision Mamba
Mamba是一种序列到序列模型，因此只能处理1D序列。为了将其应用于2D图像和3D体积，需要对像素（或 Voxel ）进行1D序列的展开。与在视觉Transformer中采用的方法不同，自注意力与像素数量的成本是平方的，这阻止了它们扩展到“实际”输入大小，并需要提取块以减少输入空间维度。Mamba使作者能够在线性时间内对输入进行序列建模，防止任何采样。块降采样是医学图像分割中的主要问题，因为需要 Voxel Level 的细节，这通常由大型医学输入数据强制执行。
Mamba的一个缺点是它不是排列不变的。与transformer自注意力机制不同，其中每个 Token 可以从序列中的每个其他 Token 处获取信息，Mamba限制每个 Token 只能从当前状态推理信息，导致对过去 Token 的近似。这意味着，当Mamba用于图像分割任务时，序列中的第一个像素（或 Voxel ）没有上下文感知。因此，作者不是直接将Mamba块纳入U-Net架构，而是受到ViT架构[11]的启发，开发了一个包装模块。包装，包括额外的LayerNorm和MLP头以及 Shortcut ，允许作者提高Mamba的稳定性。作者将这个模块称为Mamba Layer，如图1(a)所示。
随后，作者将两个Mamba Layer实例集成到一个统一模块中。这个模块名为Bidirectional 3D Mamba Layer，输入是一个具有尺寸（B, H, W, D, C）的3D体积。它折叠空间维度，通过将一个层输入序列的向后方向，双向管理序列。接下来，这个层的输出被反向到原始顺序，然后按 Token 的顺序与“直接”层的输出相加。最后，求和被归一化和 Reshape 回3D体积。这个层如图1(b)所示。
以下，作者将详细介绍将Mamba Layer集成到nnU-Net中的策略。图2展示了作者丰富Mamba层U-Net架构。
SegMamba 。最初提出的集成包括在每个池化卷积之前插入一个单向（单向）Mamba Layer，以及U-Net的瓶 Neck 分。这种战略位置旨在提高整体上下文理解，解决卷积在全局上下文中常常遇到的固有限制，同时限制额外参数的数量。
SegMambaSkip 。U-Net架构的一个普遍认可的优势是其跳过的连接[31]，这使得网络的解码部分能够访问来自编码器的细粒度细节。与此同时，Mamba被设计用来以输入相关的方式高效地选择数据，从而能够过滤掉无关信息并记住相关信息。因此，作者在将激活图与对应解码器输出连接时，通过在U-Net架构的 Shortcut 中插入一个额外的双向3D Mamba层来进行增强。
BiSegMamba 。它包括在作者的BiSegMamba之前每个下采样步骤以及U-Net Bottleneck 中。通过利用BiSegMamba中单个连续排列的两种方向，作者在计算效率和模型有效性之间取得了平衡。双向3D Mamba层使模型能够有效地考虑不同空间维度上 Token 的重要性，而无需考虑所有可能的排列。这种方法在处理远程依赖和选择性信息处理时尤为有益，可以增强在降采样和 Bottleneck 层中识别相关特征的能力。
多向Mamba 。在Mamba内部，输入 Token 的顺序很重要。因此，作者提出对所有可能的顺序进行处理，从而在3D体积的三个空间维度（H、W、D）上得到总共六个可能排列。这总共产生了12个不同的序列，分别考虑了这六个排列的前向和后向。寻求多个方向的原因在于，每个 Voxel 都需要在所有可能的方位上利用空间信息。如果作者只考虑一个序列，如（H，W，D）。flatten（），那么第一个 Token （0，0，0）和 Token （0，0，l）之间的距离将是H*W而不是1，通常H和W的值在的范围内，总共的距离为。
由于内存限制，作者只涵盖6个可能方向中的4个。2通过引入多个方向，作者保持线性复杂度，同时使每个 Token 具有优越的空间 Aware 。这种方法确保了相邻 Token 在获得的代表中确实很接近，从而增强了模型的整体空间 Aware 。为了聚合所有模块的输出序列，作者在新轴上堆叠每个序列，并计算其平均值（图3）。这个模块替代了BiSegMamba中的双向3D Mamba Layer。
### Implementation Details
有关 Patch 形状、批处理大小等流水线设置的详细信息，请参见表1。作者使用RAdam优化器，学习率0.0003，线性学习率调度器训练了所有模型300个周期。对于Mamba层参数的初始化，作者将残差层的重量初始化乘以1/√N，其中N是残差层的数量。
这与GPT-2论文中的方法相同，并已在Mamba源代码中采用。状态空间模型的内部维度（即进化参数A的大小，见第3.1节）定义为min(C,256)，其中C是输入序列中每个token的通道数。训练在运行CUDA 11.8和PyTorch 2.1.2的Nvidia A100 GPU上进行，使用CUDA和PyTorch。
4 Experiments and Results
数据集。作者使用了三个不同知名的数据集：MSD脑肿瘤[1]、Synapse多器官[24]和ACDC[4]。值得注意的是，所选实验设置与医学图像分割[7, 19, 35, 41, 43, 46]_MSD脑肿瘤._第一个是医学分割大赛（MSD脑肿瘤）的脑肿瘤分割数据集[1]。
它包含484张MRI图像，每个图像包含四个通道：FLAIR、T1w、T1gd和T2w。图像被标注为三个肿瘤亚区域：水肿（ED）、增强肿瘤（ET）和非增强肿瘤（NET）。为了与[19]中的结果保持一致，分割指标在ET、肿瘤核心（TC，即ET和NET的并集）和整个肿瘤（WT，即ED、ET和NET的并集）上进行计算。根据[19]提供的划分，作者使用95%的数据集作为训练/验证集，进行五倍交叉验证，剩余5%用于测试。
《多器官连接》第二个数据集是Synapse多器官分割数据集[24]，该数据集发布在2015年MICCAI多 Atlas 腹部标签挑战赛内。该数据集包括30个腹部CT扫描的3,779个轴位对比增强腹部CT图像，每个体积包含85至198层切片。作者采用了与[8]中相同的划分，其中18个案例用于训练，12个案例用于测试。与作者的竞争对手一样，该数据集的评价指标针对腹部器官中的8个（主动脉、胆囊、左肾、右肾、肝、胰腺、脾脏和胃）计算。
ACDC. 最后，作者使用的第三个数据集是ACDC数据集[4]。该数据集包含100个MRI扫描，每个扫描都被标注为左心室（LV）、右心室（RV）和心肌（Myo）。作者将这个数据集分为80个样本用于训练和验证，以及20个测试样本，按照[8]中的划分方法进行。
评估指标。作者采用Dice相似度系数（DSC，%）和95分位数海德距离（HD95，mm），这两个是广泛接受的指标，用于分割任务[28]。
DSC（Dice Coefficient）在实际应用中几乎与IoU（Intersection over Union，交并比）具有相同的意义。但是，在感兴趣区域远小于背景的情况下，DSC（Dice Coefficient）可能更适合，因为它更注重正确识别的区域。在这种情况下，由于赋予正确识别区域更高的权重，DSC可能比IoU更具鲁棒性和信息量。DSC指标及其与IoU之间的关系可以用以下公式表示：
其中 是模型预测， 是真实值。
另一方面，HD95计算两组点之间最大距离，考虑这些距离的95百分位数。通常，定义A和B之间边界点之间的距离的95百分位数如下：
给定预测 Mask （）中的像素集合和属于真实值的像素集合（），海森堡距离被定义为与之间的最大值和与之间的最大值在95百分位数处的值。
通过使用95百分位数，这一指标提供了一种稳健的评价方法，对数据集中的异常值或点集之间的极端差异不太敏感。
对比方法已在最近提出的医学图像分割方法上进行了性能比较。具体而言，可比较的竞争者可以分为三个主要组：卷积神经网络（CNN）-、Transformer-和Mamba-为基础的架构。
在前一组中，作者包括原始的nnU-Net [22]配置，该配置使用标准的U-Net架构（nnU-Net），以及基于U-Net（编码器中具有残差连接）的变体（nnU-Net ResEnc）。此外，基于ConvNeXt块的变体MedNeXt [32]也已被考虑在其两个变体K3和K5中。
关于基于Transformer的结构，作者将作者的 Proposal 与TransU-Net [8]、TransBTS [39]、CoTr [41]、结合卷积和Transformer模块的混合架构UNETR [19]、SwinU-Net [7]及其基于UNETR的变体Swin-UNETR [18]、LeViT-UNet-384s [43]、MISSFormer [21]以及最近发布的nn-Former [46]进行了比较。
最后，作者包括UMamba [27]在其两个变体UMambaBot和UMambaEnc。
在作者的实验中，已经采用了标准的超参数配置方案。只要有可能，就使用了自配置方法，否则，作者选择了默认配置（如果有）或者与各自数据集最接近的配置，将学习率降低到收敛为止。模型是从头开始训练的，没有任何预训练数据。始终使用了nnU-Net五倍交叉验证方案。
随着在BrainTumor数据集（表2）上的结果（如图），SegMamba，BiSegMamba和MultiSegMamba在平均指标上始终优于SegMambaSkip。在单独的许多类别中，它们也优于其他模型。在SegMamba模型中，MultiSegMamba，该模型利用更多的方向，在其他配置中表现最好，这表明建模多个方向的重要性。除了nnFormer，作者的Mamba-based架构在最佳表现的可变形器架构上获得了超过3个dice点，在nnU-Net上获得了多达1个dice点。
关于Synapse Abdomen数据集（表3），由于类别数量较多，结果显示作者的模型在肾脏和脾脏分割方面取得了显著的改进，同时平均HD95和DSC也有所提升，与最先进的架构相比。值得注意的是，包含四个独特方向使得胆囊分割取得了更明显的改进，这是最难以分割的器官。
实际上，胆囊比其他器官如肝脏明显更小，形状和位置变化更大。此外，胆囊与腹部腔内其他器官和结构之间的紧密相邻增加了在医学图像中区分它的复杂性。胆囊分割的结果表明，SegMamba达到62.21 Dice点，而其多方向版本如BiSegMamba和MultiSegMamba分别比它提高了8和10点。
最后，在ACDC数据集上的结果如图4所示。该表显示MultiSegMamba在所有模型较少方向的变种中表现优于其他所有模型，而MultiSegMamba和BiSegMamba始终优于SegMambaSkip和所有U-Mamba变体。
处理细微细节时容易出错。这证实了之前讨论的定量结果。
5 Conclusion
这篇论文旨在评估Mamba状态空间模型在3D医学图像分割方面的有效性，并与先进的卷积和Transformer结构进行比较。此外，作者还提出了一些Mamba架构的替代设计，以解决其关键限制。具体来说，作者将Mamba集成到标准U-Net框架的各个阶段，无论是 Short-Cut 还是池化操作之前，利用单向、双向和多向实现。
整个框架将卷积和状态空间模型相结合，利用前者进行精确空间信息的编码，而后者用于建模长程 Voxel Level 的交互。Mamba具有双重优势，既提供全局上下文，又实现 Voxel Level 的精确度，这在传统卷积层中由于感受野限制而缺失，而在Transformer中由于其计算复杂性而缺失。
作者的实验结果表明，与nnU-Net和不同基于transformer的网络相比，在三个知名数据集上，HD95和DSC指标有了显著提升。
作者展示了Mamba的灵活性，通过将其从原本在文本生成和大型语言模型中的使用，改编到实现在一个完全不同的任务上取得最先进结果，从而突显了Mamba在其初始设计之外的可能性。这种适应性证明了Mamba在图像编码和分割方面的有效性。
局限性与未来工作 。尽管在Mamba模型上取得了进步，可以明确指出两个关键的局限性。
首先，由于Mamba是一种因果模型，将其应用于非因果视觉数据需要进行修改。具体来说，作者通过将每个序列分别向前和向后处理来解决这个问题。然而，这会引入冗余，增加过拟合的风险。作者认为可以开发更有效的方法来解决这个问题。
[0]. Taming Mambas for Voxel Level 3D Medical Image Segmentation.