U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

增强生物医学图像分割中的长距离依赖性

The code, models,and data are publicly available at ：https://wanglab.ai/u-mamba.html

   这篇论文介绍了一种名为U-Mamba的新型网络架构，旨在提高生物医学图像分割任务中的长距离依赖建模能力。U-Mamba结合了卷积神经网络（CNNs）的局部特征提取能力和状态空间序列模型（SSMs）处理长序列数据的优势。

摘要 卷积神经网络（CNNs）和Transformer是生物医学图像分割中最受欢迎的架构，但由于其固有的局部性或计算复杂性，它们在处理长距离依赖性方面存在局限性。为了应对这一挑战，我们引入了U-Mamba，一个用于生物医学图像分割的通用网络。受状态空间序列模型（SSMs）的启发，SSMs是一种新型深度序列模型，以其处理长序列的强大能力而闻名，我们设计了一个混合CNN-SSM块，该块结合了卷积层的局部特征提取能力与SSMs捕捉长距离依赖性的能力。此外，U-Mamba具有自我配置机制，能够自动适应各种数据集而无需人工干预。我们在四个不同的任务上进行了广泛的实验，包括CT和MR图像中的3D腹部器官分割、内窥镜图像中的器械分割以及显微镜图像中的细胞分割。结果表明，U-Mamba在所有任务中的表现均优于最新的基于CNN和基于Transformer的分割网络。这为生物医学图像分析中的高效长距离依赖性建模开辟了新途径。

关键字： 图像分割 、骨干网络、 State Space Models（SSM状态空间序列模型）、 Transformer

1. 介绍

  分割是生物医学图像分析中的一个基础过程，它将图像划分为对医生和生物学家有意义的解剖结构或感兴趣区域。这项任务在多种生物医学应用中至关重要，如疾病诊断、癌症微环境量化、治疗策略规划和疾病进展追踪。虽然手动分割是完成这项任务的一种直接方式，但它极其耗时，且需要广泛的领域知识，这在实践中是不切实际的。因此，对自动分割方法有着巨大的需求。在过去的十年里，分割方法经历了从基于数学模型或手工特征的方法到基于深度学习的范式的转变，这极大地提高了各种生物医学任务（如3D计算机断层扫描(CT)中的肿瘤分割=和器官分割=，以及显微镜图像中的细胞分割）的分割准确性和效率。

  在医学图像分割中，两种流行的网络架构——卷积神经网络（CNNs）和Transformer——已经占据了重要地位。卷积神经网络，如U-Net和DeepLab，在提取层次化图像特征方面非常高效，并且比简单的全连接网络具有更高的参数效率。它们共享的权重架构使它们在捕获平移不变性和识别局部特征方面表现出色。Transformer最初是为自然语言处理设计的，但已被成功应用于图像处理，例如用于图像识别的Vision Transformer（ViT）和用于各种视觉任务中通用骨干的SwinTransformer。与CNNs不同，Transformer本质上不处理空间图像层次结构，而是将图像视为一系列图像块（patches）的序列，这使其具有更好的捕获全局信息的能力。由于这一互补特性，许多研究已探索通过混合网络架构将Transformer集成到CNNs中，如TransUNet、UNETR、nnFormer和SwinUNETR。

  虽然Transformer提高了处理长距离依赖的能力，但它们通常计算成本非常高。这是因为自注意力机制的计算量与输入大小呈二次方关系，使得它们资源密集，特别是对于通常具有高分辨率的生物医学图像而言。因此，如何有效地增强CNN中的长距离依赖仍然是一个悬而未决的问题。最近，状态空间序列模型（SSMs），特别是结构化状态空间序列模型（S4），已成为构建深度网络的一个高效且有效的构建块（或层），在连续长序列数据分析中取得了前沿性能。Mamba通过引入选择机制进一步改进了S4，允许模型以输入依赖的方式选择相关信息。通过与硬件感知实现相结合，Mamba在如语言和基因组学等密集模态上超越了Transformer。另一方面，状态空间模型在视觉任务上也表现出了有前景的结果，如图像和视频分类。由于图像块和图像特征可以被视为序列，SSM的这些吸引人的特性促使我们探索使用Mamba块来增强CNN中长距离建模能力的潜力。

  在本文中，我们提出了一种名为U-Mamba的通用网络，用于3D和2D生物医学图像分割。基于创新的混合CNN-SSM（卷积神经网络-状态空间模型）架构，U-Mamba能够捕获图像中的局部细粒度特征和长距离依赖性。与常见的基于Transformer的架构相比，该网络通过提供特征大小的线性缩放能力而脱颖而出，这与Transformer通常相关的二次复杂性形成对比。此外，U-Mamba的自配置能力使其能够无缝适应各种数据集，从而在各种生物医学分割任务中提高了其可扩展性和灵活性。在四个不同的数据集上进行的定量和定性结果表明，U-Mamba取得了卓越的性能，大幅超越了基于Transformer的网络。这为未来网络设计的发展铺平了道路，这些设计将有效且高效地模拟生物医学成像中的长距离依赖性。

2. 方法

  U-Mamba遵循编码器-解码器网络结构，以有效的方式捕获局部特征和长距离上下文。图1展示了U-Mamba块和完整网络结构的概览。接下来，我们首先介绍Mamba块，然后详细说明U-Mamba。

图1：U-Mamba（编码器）架构概览：

a. U-Mamba构建块包含两个连续的残差块（Residual blocks），随后是基于SSM（状态空间模型）的Mamba块，以增强长距离依赖性的建模。

b. U-Mamba采用编码器-解码器框架，其中编码器部分包含U-Mamba块，解码器部分包含残差块，并通过跳跃连接（skip connections）连接两者。注意：此图示仅作为概念性表示。U-Mamba从nnU-Net继承了自配置特性，网络块的数量会根据数据集自动确定。详细的网络配置如表2所示。

U-Mamba构建块描述：
   U-Mamba旨在结合U-Net和Mamba的优点，以在医学图像分割中实现全局上下文理解。如图1a所示，每个构建块包含两个连续的残差块，后跟一个Mamba块。残差块包含普通的卷积层，后跟实例归一化（IN）和Leaky ReLU激活函]。形状为(B, C, H, W, D)的图像特征随后被展平并转置为(B, L, C)，其中L = H × W × D。在通过层归一化后，特征进入Mamba块，该块包含两个并行分支。在第一分支中，特征首先通过一个线性层扩展到(B, 2L, C)，然后经过一个一维卷积层、SiLU激活函数，并与SSM层一起处理。在第二分支中，特征也通过一个线性层扩展到(B, 2L, C)，后跟SiLU激活函数。之后，来自两个分支的特征通过哈达玛积（Hadamard product）合并。最后，特征被投影回原始形状(B, L, C)，然后重新塑形并转置回(B, C, H, W, D)。

U-Mamba网络架构描述：
  图1b展示了完整的U-Mamba网络架构，其中编码器由上述构建块构建，以捕获局部特征和长距离依赖。解码器由残差块和转置卷积组成，专注于详细的局部信息和分辨率恢复。此外，我们继承了U-Net中的跳跃连接，以连接来自编码器的层次化特征到解码器。最终的解码器特征被传递到一个1×1×1卷积层，与Softmax层一起产生最终的分割概率图。此外，我们还实现了一个U-Mamba变体，其中只有瓶颈使用U-Mamba块，其他部分是普通的残差块。我们使用“U-Mamba_Bot”和“U-Mamba_Enc”来区分这两种网络变体，分别在瓶颈和所有编码器块中使用U-Mamba块。

2.1 Mamba：具有扫描（S6）功能的选择性结构化状态空间序列模型

   状态空间序列模型（SSMs）是一种系统，它将一维函数或序列 u(t) 映射到 y(t) ∈ R，并可以表示为以下线性常微分方程（ODE）：
                  x′(t) = Ax(t) + Bu(t),
                   y(t) = Cx(t),
  其中，状态矩阵 A ∈ R^N×N 和 B, C ∈ R^N 是其参数，x(t) ∈ R^N 表示隐含的潜在状态。状态空间序列模型（SSMs）提供了几种所需的属性，例如每个时间步的线性计算复杂度和用于有效训练的并行计算能力。然而，简单的SSMs通常比等效的卷积神经网络（CNNs）需要更多的内存，并且在训练过程中经常遇到梯度消失的问题，这阻碍了它们在一般序列建模中的更广泛应用。

  结构化状态空间序列模型Structured State Space Sequence Models（S4） 通过在状态矩阵 A 上施加结构化形式并引入一种有效的算法，显著改进了简单的SSMs。具体而言，该状态矩阵是通过高阶多项式投影算子（HIPPO） 精心设计和初始化的，从而能够构建具有丰富功能和高效长距离推理能力的深度序列模型。作为一种新型网络架构，S4 在极具挑战性的长距离竞技场基准（Long Range Arena Benchmark）上以显著优势超越了Transformer。

  最近，Mamba 通过两项关键改进进一步推进了离散数据建模（例如，文本和基因组）的顺序状态机（SSM）。首先，Mamba采用了一种依赖于输入的选择机制，与传统的不随时间和输入变化的SSM不同，它允许从输入中高效地过滤信息。这是通过根据输入数据参数化SSM参数来实现的。其次，开发了一种硬件感知算法，其在序列长度上的扩展性呈线性，用于通过扫描递归计算模型，使得Mamba在现代硬件上比以前的方法更快。此外，Mamba架构将SSM块与线性层合并，不仅显著更简单，而且在包括语言和基因组学在内的各种长序列领域展示了最先进的性能，展示了在训练和推理过程中的重大计算效率。*

2.2 U-Mamba: Marry Mamba with U-Net

  Mamba在各种离散数据上展示了令人印象深刻的结果，但它在建模图像数据，特别是在生物医学成像方面的潜力仍未被充分探索。鉴于图像本质上是从连续信号中采样的离散数据，并且在展平时可以被视为长序列，我们提出利用Mamba的线性扩展优势来增强CNNs的长距离依赖建模。Transformer已经被成功应用于处理图像，如ViT 和SwinTransformer ，但它们在处理大图像时常常因为自注意力具有二次复杂度而遭受高计算负担。这些因素激励我们利用Mamba中的线性扩展特性来增强CNNs的长距离依赖。

  U-Net 及其变体在医学图像分割中被广泛使用的网络架构。它们通常具有对称的编码器-解码器结构，通过卷积操作提取多尺度图像特征。然而，这种架构设计在建模图像中的长距离依赖方面的能力有限，因为卷积核本质上是局部的。每个卷积层只能从有限的感受野中捕获特征。虽然U-Net中的跳跃连接有助于结合低级细节和高级特征，但它们主要加强了局部特征整合，而不是扩展网络建模长距离依赖的能力。

3. 实验和结果

3.1 Datasets

  我们使用了四个医学图像数据集来评估U-Mamba在不同图像尺寸、分割目标和模态下的性能和可扩展性。表1提供了数据集的概览。所有数据集都是公开可用的，并且允许用于研究目的。

腹部CT：这个数据集来自2022年MICCAI FLARE挑战赛，专注于分割13个腹部器官，包括肝脏、脾脏、胰腺、右肾、左肾、胃、胆囊、食管、主动脉、下腔静脉、右肾上腺、左肾上腺和十二指肠。训练集包含了50个CT扫描，这些扫描来自MSD Pancreas数据集，注释来自AbdomnenCT-1K。另外50个案例来自不同的医疗中心，用于评估，注释由挑战赛组织者提供。

腹部MRI：这个数据集来自2022年MICCAI AMOS挑战赛，同样专注于腹部器官分割。原始数据集包含40个用于训练的MRI扫描和20个用于验证的MRI扫描。由于20个案例不足以得出统计学上有意义的结果，在我们的实验中，我们使用了原始的60个标记的MRI扫描进行模型训练，并另外注释了50个MRI扫描作为测试集。为了能够在模态上比较腹部器官，我们也专注于与腹部CT数据集相同的13个器官。注释是由放射科医生在MedSAM和ITK-SNAP的辅助下生成的。我们已经将这个新注释的数据集发布给社区，以促进MRI中腹部器官分割的发展。

内窥镜图像：这个数据集来自2017年MICCAI EndoVis挑战赛，专注于从内窥镜图像中分割出七种器械，包括大号针持器、Prograps钳、单极弯剪刀、Cadiere钳、双极钳、血管封闭器，以及额外的插入式超声波探头。我们遵循了官方的数据集分割，训练集分别包含了1800和1200个图像帧。训练图像来自八个视频，而测试集则包含了另外两个新视频的未见过的图像。

显微镜图像：这个数据集来自2022年NeurIPS细胞分割挑战赛，专注于各种显微镜图像中的细胞分割。我们分别使用了1000张和101张图像进行训练和评估。与上述三个任务不同，这是一个实例分割任务，其中算法需要为每个细胞实例分配一个唯一的标签。在我们的实验中，我们将实例分割转换为语义分割任务，通过预测细胞边界和内部区域，因为它们可以通过“skimage.measure.label”函数轻松转换为实例掩码。需要注意的是，我们的目的是基准测试网络架构的性能，而不是追求这项任务的最先进性能。U-Mamba也可以作为最先进实例分割框架的骨干网络。我们将这种扩展作为未来的工作。

3.2 实现和训练协议

  我们实现了U-Mamba，它是基于流行的nnU-Net框架构建的，这一选择受到两个期望特性的驱动。首先，nnU-Net的模块化设计与我们专注于引入一种新的网络架构的目标完美契合。这种设计使我们能够专注于网络的实现，同时控制其他变量，如图像预处理和数据增强。这样的设置使得在统一条件下，U-Mamba与各种方法的公平比较成为可能，网络架构是唯一的不同因素。其次，nnU-Net的显著自配置框架，能够自动为不同的分割数据集配置超参数，这是我们选择的另一个令人信服的原因。我们也保留了这一特性，允许U-Mamba轻松适应广泛的分割任务。在训练过程中，补丁大小、批量大小和网络配置（例如，分辨率状态的数量和不同轴向上的下采样操作次数）与nnU-Net保持一致（表2）。U-Mamba还通过随机梯度下降进行优化，损失函数是Dice损失和交叉熵的未加权和，因为复合损失已证明在不同任务中具有鲁棒性。在推理过程中，为了更简化和高效的评估过程，所有实验中都禁用了测试时间增强（TTA），因为TTA会分别使2D和3D数据集的计算负担增加4倍和8倍。

3.3 Benchmarking(基准测试)

  我们将U-Mamba与两种基于CNN的分割网络（nnU-Net 和 SegResNet ）以及两种基于Transformer的网络（UNETR 和 SwinUNETR ）进行了比较，这些网络在医学图像分割竞赛中被广泛使用。为了进行公平比较，我们还将SegResNet、UNETR和SwinUNETR实现到nnU-Net框架中，并使用了他们建议的优化器（例如，Adam 和 AdamW ）进行模型训练。我们使用了nnU-Net 中的默认图像预处理。所有网络都在一个NVIDIA A100 GPU上从头开始训练了1000个周期，使用相同的批量大小（表2）。

  遵循《Metrics Reloaded》中的建议，我们使用了Dice相似系数（DSC）和归一化表面距离（NSD）来评估三个语义分割任务，包括CT和MRI扫描中的器官分割，以及内窥镜图像中的器械分割。由于这是一个实例分割任务，我们使用F1分数来评估细胞分割质量。

  表3：3D腹部CT和MRI数据集器官分割的结果总结。
  U-Mamba_Bot：仅在瓶颈处使用U-Mamba块。
  U-Mamba_Enc：所有编码器块都是U-Mamba块。


  图2：可视化的腹部器官分割示例，包括CT扫描（第1行和第2行）和MRI扫描（第3行和第4行）。U-Mamba在区分腹部复杂软组织方面具有更好的能力。

3.4 定量和定性的分割结果

  表3显示了CT和MRI扫描中腹部器官的3D分割的定量结果。U-Mamba超越了基于CNN和基于Transformer的分割网络，分别在腹部CT和MR数据集上实现了平均DSC得分为0.8683（U-Mamba_Bot）和0.8501（U-Mamba_Enc）。nnU-Net的表现具有竞争力，远优于SegResNet、UNETR和SwinUNETR，但U-Mamba仍然显示出优势，如图2所示。基于六个网络的可视化分割结果，U-Mamba在分割结果中的异常值较少。例如，U-Mamba可以生成更准确的肝脏和胃的分割掩码，而其他方法在腹部CT扫描中对肝脏掩码有过度分割错误，对胃掩码有遗漏区域。同样地，在MRI扫描中的胆囊分割中，U-Mamba成功地描绘了其边界，而其他方法产生了各种分割错误。这些观察结果也反映在器官级别的DSC和NSD得分中（附录表5和表6）。

  表4. 2 D分割任务的结果总结：腹部MRI扫描中的器官分割、内窥镜图像中的器械分割以及显微镜图像中的细胞分割。

  表4展示了2D分割的定量结果。MRI扫描通常比CT扫描具有更不均匀的间距，2D分割网络在实践中也是一个常见的选择。因此，我们还比较了所有网络在2D设置下的腹部器官分割，其中每个3D MRI扫描被转换成多个2D切片。U-Mamba在这三个2D分割任务上再次展示了其优势，分别在器官、器械和细胞分割中取得了最好的平均DSC得分0.7625、0.6504和F1得分0.5607。图2中的定性结果表明，由于外观的异质性，nnU-Net、UNETR和SwinUNETR在区分脾脏时存在混淆。UNETR和SwinUNETR在显微镜图像中也产生了许多细胞异常值，而nnU-Net未能成功识别内窥镜图像中的大号针持器。相比之下，U-Mamba在这些场景中显示出明显的优势，表明其在捕捉全局上下文方面具有更好的能力。
图3. 可视化的分割示例：腹部器官分割的MRI扫描（第1行）、显微镜图像中的细胞分割（第2行）和内窥镜图像中的器械分割（第3行）。U-Mamba对异质性外观更为稳健，并且分割异常值更少。

4. 结论

  在本文中，我们介绍了U-Mamba，以应对由于CNN的固有局部性和Transformer的计算复杂性而导致的长距离依赖建模的挑战。全面的实验结果揭示了U-Mamba在不同模态和分割目标上超越了现有的基于CNN和Transformer的分割网络。特别是，U-Mamba在处理外观异质的对象方面展现出了卓越的能力，从而减少了分割异常值。性能的提升主要归功于U-Mamba的架构设计，它能够同时提取多尺度局部特征并捕获长距离依赖。

  Transformer在许多自然图像分割任务上已经展现出了更强的性能。然而，它们并没有改变医学图像分割的格局。这一点从医学图像分割竞赛中的获胜解决方案仍然主要依赖于CNNs，如nnU-Net和SegResNet，就可以看出。有趣的是，我们还观察到，在所有我们测试的场景中，基于Transformer的网络相对于基于CNN的网络表现不佳，尽管我们使用了它们原始的网络实现，并用推荐的优化器对它们进行了训练。这可能是因为需要用多个GPU而不是一个GPU来训练这些网络。然而，这将导致不公平的比较，因为所有其他网络都是只用一个GPU训练的。另一个原因可能是Transformer应该在大规模预训练和微调范式中使用。为了在我们的评估中保持一致性和公平性，我们从头开始训练了所有网络，并在所有实验中采用了相同的预处理和数据增强技术。我们相信社区合作和透明度的价值，因此，我们已经将比较方法的实现公开，供研究社区进一步检查和使用。

  此外，nnU-Net在大多数任务上都展现出了非常有竞争力的结果，但U-Mamba总体上取得了更好的分数。我们注意到nnU-Net和U-Mamba在CT和MRI扫描中对某些器官的分割性能具有互补性，例如附录表5-7中所示的主动脉、肾脏和下腔静脉。这表明通过模型集成进一步整合它们的潜力，这是一种在医学图像分割竞赛中常用的策略，用于提高性能。

  虽然这项工作的主要重点是新的分割架构，但U-Mamba还有许多途径可以进一步增强和扩展。一个直接的方向是利用大规模数据集来训练U-Mamba，旨在创建随时可部署的分割工具或为数据有限的任务提供预训练模型权重，正如TotalSegmentator和STU-Net等类似倡议所示范的。此外，U-Mamba的设计天然支持与先进技术的集成，例如对小数据集的强大数据增强、针对高度不平衡目标的损失函数，以及针对嵌套对象的基于区域的训练。此外，U-Mamba块为在分类和检测网络中应用提供了新的机会，以实现更好的长距离依赖建模。我们将这些方向留作近期未来的工作。

  总之，本文提出了一种新的架构U-Mamba，用于通用的生物医学图像分割，它整合了CNNs的局部模式识别优势和Mamba的全局上下文理解优势。U-Mamba可以自动为不同数据集配置自己，使其成为生物医学成像中多样化分割任务的多功能和灵活工具。结果表明，U-Mamba是一个有前途的候选者，可以作为下一代生物医学图像分割网络的骨干。

  表5. 腹部CT数据集中3D模型的器官分割结果。每个器官和每个指标的最佳和次佳分数分别以红色和蓝色突出显示。（IVC：下腔静脉，RAG：右肾上腺，LAG：左肾上腺）

  表6. 腹部MRI数据集中3D模型的器官分割结果。每个器官和每个指标的最佳和次佳分数分别以红色和蓝色突出显示。（IVC：下腔静脉，RAG：右肾上腺，LAG：左肾上腺）

  表7. 腹部MRI数据集中2D模型的器官分割结果。每个器官和每个指标的最佳和次佳分数分别以红色和蓝色突出显示。（IVC：下腔静脉，RAG：右肾上腺，LAG：左肾上腺）