日常阅读论文——U-Mamba

1、摘要

很抽象。卷积神经网络（CNNs）和变换器（Transformers）是生物医学图像分割中最常用的两种结构，但由于其固有的局部性和计算复杂性，这两种结构在处理长范围相关性方面的能力都很有限。为了解决这一难题，我们引入了U-Mamba，一种用于生物医学图像分割的通用网络。受状态空间序列模型（SSM）的启发，我们设计了一个混合CNN-SSM模块，该模块将卷积层的局部特征提取能力与SSM捕捉长相关性的能力相结合。此外，U-Mamba还拥有一个自配置机制，使其能够自动适应各种数据集，而无需人工干预。我们对四种不同的任务进行了广泛的实验，包括CT和MR图像中的3D腹部器官分割、内窥镜图像中的器械分割和显微镜图像中的细胞分割。实验结果表明，U-Mamba在所有任务中的性能均优于最新的基于CNN和基于Transformer的分割网络。这为生物医学图像分析中的有效长范围相关性建模开辟了新的途径。代码、模型和数据可在https://wanglab.ai/u-mamba.html上公开获取。

2、引言

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

两种流行的网络架构，CNN [27]和Transformers [46]，在医学图像分割中已经获得了突出地位。CNN，如U-Net [41，22]和DeepLab [7]，在提取分层图像特征方面是有效的，并且比幼稚的全连接网络更有参数效率。它们的共享权重架构使它们擅长捕捉平移不变性和识别局部特征。Transformer最初是为自然语言处理而设计的，但已经成功地适用于图像处理，例如用于图像识别的Vision Transformer（ViT）[9]和用于各种视觉任务的通用骨干的Swin Transformer [28]。与CNN相比，Transformers本身并不处理空间图像层次结构，而是将图像视为一系列补丁。它具有更好的捕获全局信息的能力。由于这种互补性，许多研究探索了通过混合网络架构将Transformer并入CNN，例如TransUNet [5，6]，UNETR [17]，nnFormer [49]和SwinUNETR [16]。

虽然变形金刚提高了长距离依赖性的能力，但它们通常在计算上非常昂贵。这是因为自注意机制与输入大小成二次比例，使得它们是资源密集型的，尤其是对于通常具有高分辨率的生物医学图像。因此，如何有效地增强神经网络的长程依赖性仍然是一个悬而未决的问题。最近，状态空间序列模型（SSM）[11，15]，特别是结构化状态空间序列模型（S4）[14]，已经作为用于构建深度网络的高效且有效的构建块（或层）出现，其在连续长序列数据分析中获得了前沿性能[10，14]。Mamba [12]使用选择机制进一步改进了S4，允许模型以输入依赖的方式选择相关信息。通过与硬件感知实现相结合，Mamba在语言和基因组学等密集模态上超过了Transformers。另一方面，状态空间模型在视觉任务上也显示出有希望的结果，例如图像[40]和视频[24]分类。由于图像块和图像特征可以被投射为序列[9，28]，SSM的这些吸引人的特征促使我们探索使用Mamba块来增强CNNs中的远程建模能力的潜力。

在本文中，我们提出了U-MAMBA，一个通用的网络，用于三维和二维生物医学图像分割。基于创新的混合CNNSSM体系结构，U-Mamba能够同时捕获本地化的细粒度特征和图像中的长距离依赖性。该网络通过在特征大小上提供线性扩展，而不是通常与Transformer相关的二次复杂性，从流行的基于Transformer的架构中脱颖而出。此外，U-Mamba的自配置功能可实现对各种数据集的无缝适应，从而增强其在一系列生物医学分割任务中的可扩展性和灵活性。四个不同数据集的定量和定性结果表明，U-Mamba实现了上级性能，大大超过了基于Transformer的网络。这为网络设计的未来发展铺平了道路，这些网络设计能够高效地对生物医学成像中的长距离相关性进行建模。

3、方法

U-Mamba遵循编码器-解码器网络结构，以有效的方式捕获本地特征和远程上下文。图1显示了U-Mamba块和完整网络结构的概述。接下来，我们首先介绍Mamba块，然后说明U-Mamba的细节。

（一个U-Mamba构建模块包含两个连续的残差模块，后面是基于SSM的Mamba模块，用于增强远程依赖性建模。U-Mamba继承了nnU-Net的自配置功能，网络块的数量在数据集中自动确定。）

4、Mamba：带有扫描的选择性结构化状态空间序列模型（S6）

其中状态矩阵A ∈ RN×N，B、C ∈ RN是其参数，x（t）∈ RN表示隐式潜在状态。SSM提供了几个所需的属性，例如每个时间步长的线性计算复杂度和用于有效训练的并行计算。然而，简单的SSM通常比等效的CNN需要更多的内存，并且在训练过程中经常遇到消失的梯度，这阻碍了它们在一般序列建模中的更广泛应用。

结构化状态空间序列模型（S4）[14]通过在状态矩阵A上施加结构化形式并引入有效算法，显着改进了简单的SSM。具体来说，状态矩阵是用高阶多项式投影算子（HIPPO）[13]制作和初始化的，允许构建具有丰富能力和高效远程推理能力的深度序列模型。S4作为一种新型的网络架构，在具有挑战性的长距离竞技场基准测试[44]上已经超过了变压器[46]。最近，Mamba [12]进一步推进了离散数据建模中的SSM（例如，文本和基因组）通过两个关键的改进。首先，Mamba采用了一种依赖于输入的选择机制，与传统的时间和输入不变的SSM不同，它允许从输入中进行有效的信息过滤。这通过基于输入数据对SSM参数进行参数化来实现。其次，开发了一种硬件感知算法，在序列长度上线性缩放，通过扫描递归地计算模型，使Mamba比现代硬件上的以前的方法更快。此外，Mamba架构将SSM块与线性层合并，非常简单，并且在包括语言和基因组学在内的各种长序列领域中表现出最先进的性能，在训练和推理过程中展示了显着的计算效率。

5、U-Mamba（与U-Net相结合）

Mamba已经在各种离散数据上展示了令人印象深刻的结果，但它在建模图像数据，特别是在生物医学成像方面的潜力仍然有待开发。鉴于图像本质上是从连续信号中采样的离散数据，并且在展平时可以被视为长序列，我们建议利用Mamba的线性缩放优势来增强CNN的长距离依赖建模。Transformer 已经成功地应用于处理图像，例如ViT [9]和SwinTransformer [28]，但是它们通常遭受大图像的高计算负担，因为自我注意力具有二次复杂性这些因素促使我们利用Mamba中的线性缩放特性来增强CNNs的长程依赖性。

U-Net [41]及其变体[22]是医学图像分割中广泛使用的网络架构。它们通常采用对称的编解码器结构，通过卷积运算提取多尺度图像特征。然而，由于卷积核本质上是局部的，因此这种体系结构设计对图像中的长范围依赖性建模的能力有限。每个卷积层从有限的感受野捕获特征。虽然U-Net中的跳跃连接有助于将低级细节与高级特征相结合，但它们主要是加强局部特征的集成，而不是扩展网络对长期依赖性建模的能力。

U-Mamba是一种融合了U-Net和Mamba的全局上下文理解技术的医学图像分割算法。如图1a所示，每个构建块包含两个连续的残差块[18]，后面是Mamba块[12]。残差块包含纯卷积层，其后是实例归一化（IN）[45]和泄漏ReLU [36]。然后，具有形状（B，C，H，W，D）的图像特征被平坦化并转置为（B，L，C），其中L = H ×W×D。在通过层标准化[2]之后，特征进入包含两个平行分支的Mamba块。在第一个分支中，特征通过线性层扩展到（B，2L，C），然后是1D卷积层，SiLU激活函数[20]，并与SSM层一起。在第二分支中，特征也被扩展到（B，2L，C），其中线性层之后是SiLU激活函数。然后，将两个分支的特征与Hadamard乘积合并。最后，特征被投影回原始形状（B、L、C），然后重新成形并转置为（B、C、H、W、D）。

图1b显示了完整的U-Mamba网络架构，其中编码器是用上述模块构建的，以捕获本地特征和远程依赖关系。该解码器由残差块和转置卷积组成，主要关注细节局部信息和分辨率恢复。此外，我们继承了U-Net中的skip连接，将编码器的层次特征连接到解码器。最终的解码器特征被传递到1 × 1 × 1卷积层和Softmax层，以产生最终的分割概率图。此外，我们还实现了一个U-Mamba变体，其中只有瓶颈使用U-Mamba块，其他都是普通的残差块。我们使用“U-Mamba_Bot”和“U-Mamba_Enc”来区分这两种网络变体，它们分别在瓶颈和所有编码器块中使用U-Mamba块。

5、实验和结果

我们使用了四个医学图像数据集来评估U-Mamba在不同图像尺寸、分割目标和模式下的性能和可扩展性。

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

我们将U-Mamba与两种基于CNN的分割网络（nnUNet [22]和SegResNet [39]）和两种基于变换器的网络（UNETR [17]和SwinUNETR [16]）进行了比较，这两种网络广泛用于医学图像分割竞赛。为了公平比较，我们还实现了SegResNet，UNETR，和SwinUNETR到nnU-Net框架中，并使用他们建议的优化器（例如，Adam [26]和AdamW [29]）进行模型训练。我们使用了nnU-Net [22]中的默认图像预处理。所有网络都是在同一批大小的NVIDIA A100 GPU上从头开始训练1000个epoch的（表2）。根据Dice Similarity Coefficient（DSC）和Normalized Surface Distance（NSD）的建议，我们使用Dice Similarity Coefficient（DSC）和Normalized Surface Distance（NSD）进行了三项语义分割任务，包括CT和MRI扫描中的器官分割以及内窥镜图像中的器械分割。F1评分用于评估细胞分割质量，因为它是一个实例分割任务。

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

表4显示了定量2D分割结果。MRI扫描通常比CT扫描具有更不均匀的间距，并且2D分割网络也是实践中的常见选择[4]。因此，我们还比较了腹部器官分割的2D设置下的所有网络，其中每个3D MRI扫描都转换为多个2D切片。U-Mamba再次证明了其在三个2D分割任务上优于现有方法的优势，分别为器官，仪器和细胞分割实现了0.7625，0.6504和0.5607的最佳平均DSC和F1分数。图2中的定性结果表明，nnU-Net、UNETR和SwinUNETR由于外观不均匀而混淆了脾脏。UNETR和SwinUNETR还在显微镜图像中生成许多细胞离群值，并且nnU-Net未能成功识别内窥镜图像中的大针驱动器。相比之下，U-Mamba在这些场景中表现出明显的优势，表明它具有更好的捕获全局上下文的能力。

6、讨论与总结

本文引入U-Mamba来解决由于CNNs的固有局部性和Transformers的计算复杂性而导致的在建模长范围依赖性方面的挑战。综合实验结果表明，U-Mamaba算法在多种模态和分割目标上的性能均优于现有的基于CNN和Transformer的分割网络。特别是，U-Mamba在处理具有异构外观的对象方面表现出了上级的能力，从而减少了分割离群值。性能的提升很大程度上归功于U-Mamba的架构设计，它可以同时提取多尺度局部特征和捕获长范围依赖性。

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

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