魔芯科技、中科大提出仅将 UMamba 的 Mamba 替换为 xLSTM，就能在各种2D和3D图像分割中取得SOTA表现！

极市平台

2024-07-07

↑ 点击蓝字关注极市平台

作者丨AI视界引擎

来源丨AI视界引擎

编辑丨极市平台

极市导读

利用Vision-LSTM（xLSTM）作为其骨干网络的UNet结构深度学习神经网络。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

卷积神经网络（CNNs）和视觉变换器（ViT）在生物医学图像分割中起到了关键作用。然而，由于其固有的局部性和计算开销，这些方法在处理长距离依赖关系方面的能力仍然受限。

为了解决这些问题，在本技术报告中，我们首先提出了xLSTM-UNet，这是一种利用Vision-LSTM（xLSTM）作为其骨干网络的UNet结构深度学习神经网络，专用于医学图像分割。xLSTM最近被提出作为长短期记忆（LSTM）网络的继任者，并且在神经语言处理（NLP）和图像分类中（如Vision-LSTM或ViL实现）表现出优异的性能，优于变换器和状态空间模型（SSMs）如Mamba。在此，我们设计的xLSTM-UNet旨在将这些成功经验扩展到生物医学图像分割领域。

通过结合卷积层的局部特征提取优势和xLSTM捕获长距离依赖关系的能力，xLSTM-UNet为综合图像分析提供了一个强大的解决方案。

我们通过实验验证了xLSTM-UNet的有效性。我们的研究结果表明，xLSTM-UNet在多个生物医学分割数据集上（包括腹部MRI中的器官、内窥镜图像中的仪器以及显微镜图像中的细胞），其表现持续优于领先的基于CNN、变换器和Mamba的分割网络。通过全面的实验，本技术报告突出了基于xLSTM的架构在推进2D和3D生物医学图像分析方面的潜力，代码、模型和数据集均已开源。

项目页面：http://tianrun-chen.github.io/xLSTM-Unet/

论文链接：https://arxiv.org/abs/2407.01530

开源代码：https://github.com/tianrun-chen/xLSTM-UNet-PyTorch。

1 Introduction

生物医学图像分割是医学成像中的一项关键任务，它能够精确地描绘解剖结构和异常情况，对诊断、治疗计划和研究至关重要。近年来，深度学习方法在肿瘤分割和3D计算机断层扫描（CT）中的器官分割以及显微镜图像中的细胞分割方面取得了显著成功。这些进展突显了深度学习对生物医学图像分割领域的变革性影响，为更准确和高效的诊断和治疗规划工具铺平了道路。传统上，卷积神经网络（CNNs）一直是这一领域深度学习方法的基础，利用其强大的局部特征提取能力。最近，视觉变换器（ViTs）因其通过自注意力机制捕获全局上下文的能力，成为一种强有力的替代方案。尽管它们取得了成功，但CNNs和ViTs都面临固有的限制。由于局部感受野，CNNs在处理长距离依赖关系方面表现不足，而ViTs在处理高分辨率图像或高维成像模态（如3D图像或超光谱成像，如受激拉曼散射（SRS）成像或中红外（IR）光谱成像）时会遇到巨大的计算开销。

为了解决这些挑战，最近的工作提出了整合长距离依赖关系且在序列长度方面具有线性计算和内存复杂度的计算模块。在这些计算模块中，状态空间模型（SSMs），如Mamba，显示出巨大的成功。SSMs在处理长距离依赖关系方面表现优异，并已成功整合到传统的UNet架构中。UMamba，VM-Unet，Mamba-Unet，Swin-UMamba和SegMamba等变体已展示出其显著成功。

与此同时，扩展长短期记忆（xLSTM）最近作为长短期记忆（LSTM）网络的强大继任者出现，在序列建模中挑战变换器。像SSMs一样，xLSTM可以处理长距离依赖关系并保持线性计算和内存复杂度。然而，xLSTM在神经语言处理（NLP）和图像分类（在其Vision-LSTM（ViL）实现中）中表现出优异的性能。这种成功自然引出了一个问题：xLSTM，或ViL，能否在图像分割，特别是在医学图像分割领域中表现出色？

答案是肯定的！在本技术报告中，作者们介绍了xLSTM-UNet，这是首个能够执行2D和3D医学图像分割任务并达到最新技术水平（SOTA）结果的xLSTM支持的U-Net图像分割网络。作者们在各种2D和3D医学分割场景中进行了全面的实验，包括腹部MRI中的器官，内窥镜中的仪器，显微镜中的细胞，以及3D脑MRI体积中的癌症分割。结果显示，xLSTM-UNet在现有的基于CNN和变换器的分割方法以及其基于Mamba的对应方法中表现出色。这些发现突显了基于xLSTM的架构在医学图像分割领域设立新基准的潜力，提供了在广泛应用中的改进精度和效率。为了进一步推进该领域的研究，作者们将在http://tianrun-chen.github.io/xLSTM-UNet/ 公开模型和代码，以便在自动病理检测、隐蔽图像分割、精确农业、环境监测、卫星图像分析和工业检测等各个领域进行未来的探索。

2 Method

图1展示了xLSTM-UNet网络架构。xLSTM-UNet遵循传统的UNet结构。输入信息首先经过一个卷积层进行初始下采样。然后，使用前述xLSTM构建块构建的多个后续层捕捉局部特征和长距离依赖关系，形成编码器的主要部分。需要注意的是，xLSTM-UNet的设计目标是利用U-Net和xLSTM的最佳特性，以改进医学图像理解中的全局理解。因此，作者们在编码器中的多个层中使用了xLSTM，而不仅仅是在下采样完成后的压缩潜在空间中应用xLSTM。每层包含两个连续的残差块，一个普通卷积层和一个实例归一化（IN），然后是一个xLSTM块。具体来说，通过残差块的图像特征形状为(B, C, H, W, D)，首先被展平并转置为(B, H, W, D, C)$，然后进行层归一化，再输入到ViL块中。这种在多个层中引入xLSTM的做法有助于在多种分辨率/感知域中进行特征提取，而这些xLSTM块提取的信息被重塑为(B, C, H, W, D)并在解码步骤中与各层连接，以促进分割掩码的生成。

在编码之后，由残差块和转置卷积组成的解码器专注于细致地恢复局部信息。此外，作者们继承了U-Net架构中的跳跃连接，将编码器的层级特征连接到解码器。最终的解码器特征被输入到一个1 x 1卷积层，结合Softmax层，生成最终的分割概率图。此外，按照之前的工作，作者们还实现了一个变体，其中U-xLSTM块仅在瓶颈中使用，称为“ours_bot”，而“ours_enc”则表示在所有编码器块中应用xLSTM块的网络。

3 Experiments

3.1 Datasets

作者们使用了与UMamba相同的基准测试用于2D医学图像分割，并使用了与SegMamba相同的基准测试用于3D医学图像分割。

为了验证所提出方法（xLSTM-UNet）的有效性，作者们使用了几个具有代表性的医学图像分割数据集。这些数据集涵盖了器官分割、仪器分割和细胞分割，并跨越了各种分辨率和图像模态。通过使用这些多样化的数据集，作者们能够全面评估xLSTM-UNet在不同场景中的性能和适用性，从而证明其在医学图像分割中的有效性和优越性。

腹部MRI：作者们使用了来自MICCAI 2022 AMOS挑战赛的腹部MRI数据集，该数据集专门用于腹部器官的分割。数据集由放射科医生使用MedSAM和ITK-SNAP精心标注。按照U-Mamba的设置，作者们使用60个标记的MRI扫描进行训练，并使用50个MRI扫描进行测试，训练集中包含5615张切片，测试集中包含3357张切片。该数据集包含13个不同的类别标签：肝脏、脾脏、胰腺、右肾、左肾、胃、胆囊、食管、主动脉、下腔静脉、右肾上腺、左肾上腺和十二指肠。在2D分割任务中，图像被裁剪成320 x 320像素的补丁。在3D分割任务中，补丁的大小设置为48 x 160 x 224像素。

内窥镜图像：内窥镜图像数据集来自MICCAI 2017 EndoVis挑战赛，专注于内窥镜图像中七种手术器械的分割。这些器械包括大针驱动器、程序夹持钳、单极弯曲剪刀、卡迪埃钳、双极钳、血管密封器和一个额外的超声探头。作者们遵循官方数据集划分，训练集包含1800个图像帧，测试集包含1200个图像帧。训练图像从八个视频中提取，而测试集包含两个新视频中的未见图像。图像被裁剪为384 x 640像素的大小，以适应nnU-Net框架进行训练和测试。

显微镜图像：显微镜图像数据集来自NeurIPS 2022细胞分割挑战赛，专注于各种显微镜图像中的细胞分割。作者们使用了1000张图像进行训练，101张图像进行评估。原始任务是实例分割任务。作者们在实验中将实例分割转换为语义分割任务，按照之前的工作中的数据处理方法进行处理。图像被裁剪为512 x 512像素的大小，以适应nnU-Net框架进行训练和测试。

BraTS2023：BraTS2023数据集包括1251个3D脑MRI体积。每个体积包含四种成像模态（T1, T1Gd, T2和T2-FLAIR）和三个分割目标：全肿瘤（WT）、增强肿瘤（ET）和肿瘤核心（TC）。在训练过程中，作者们使用128×128×128的随机裁剪大小处理3D数据。

3.2 Implemetation details

该网络基于UMamba实现。损失函数使用Dice损失和交叉熵损失的总和。作者们采用AdamW优化器，权重衰减为0.05。学习率设定为：腹部MRI数据集为0.005，内窥镜数据集为0.01，训练Microscopy数据集上的xLSTM-UNet_Bot为0.007，训练Microscopy数据集上的xLSTM-UNet_Enc为0.0015，BraTS2023数据集为0.01。批量大小设定为：3D腹部MRI数据集为2，2D腹部MRI数据集为30，内窥镜数据集为2，Microscopy数据集为12，BraTS2023数据集为4。所有网络均从头开始训练，使用单个NVIDIA A100 GPU进行1000个周期。有关更多实现细节，请参阅作者们的代码库。

3.3 Baselines

在2D医学分割中，为了确保公平的比较，作者们遵循UMamba中的评估协议。选择了两个基于CNN的分割网络（nnU-Net和SegResNet）和两个基于变换器的网络（UNETR和SwinUNETR），以及UMamba本身，该模型有两种变体：U-Mamba_Bot和U-Mamba_Enc。类似于作者们的配置，U-Mamba_Bot仅在瓶颈处应用，而U-Mamba_Enc在每个编码器中使用。作者们使用Dice相似系数（DSC）和归一化表面距离（NSD）作为评估指标来评估腹部MRI和内窥镜数据集上的语义分割任务。对于Microscopy数据集上的细胞分割，作者们采用了F1分数。

在3D医学分割中，对于3D腹部MRI数据集，基线方法和任务与2D医学分割中使用的保持一致。对于BraTS2023数据集，为了确保公平的比较，作者们遵循SegMamba中的评估协议。使用相同的基线方法，包括三种基于CNN的方法（SegresNet、UX-Net、MedNeXt），三种基于变换器的方法（UNETR、SwinUNETR、SwinUNETR V2），以及Mamba基方法SegMamba本身。使用Dice和HD95作为评估指标。

3.4 Quantitative and Qualitative Results for 2D Segmentation

表1展示了各种方法在腹部MRI 2D、内窥镜和显微镜数据集上的分割性能。作者们提出的xLSTM-UNet在所有基线方法中表现最佳，并实现了最新技术水平（SOTA）。

值得注意的是，xLSTM-UNet的两个变体在所有数据集上都表现出优异的性能。具体来说，xLSTM-UNet_Enc在腹部MRI 2D数据集上展示了最高的性能，DSC为0.7747，NSD为0.8374，显著超过了之前的SOTA模型U-Mamba。此外，xLSTM-UNet_Bot分别获得了DSC和NSD分数0.7636和0.8322，超越了结构类似的U-Mamba_Bot。同样，在内窥镜数据集上，xLSTM-UNet_Bot和xLSTM-UNet_Enc都获得了最佳的DSC和NSD分数，分别为0.6843和0.7001。对于显微镜数据集，xLSTM-UNet_Enc和xLSTM-UNet_Bot分别获得了F1分数0.6036和0.5818，均超越了之前的SOTA结果，表明它们在细胞分割任务中的稳健性。

2D医学图像的可视化分割示例进一步展示了xLSTM-UNet的有效性。如图2所示，与其他模型相比，xLSTM-UNet对异质外观更具鲁棒性，分割错误更少。这些可视化证据支持了定量结果，突显了xLSTM-UNet在多样医学图像分割任务中的卓越性能和可靠性。

3.5 Quantitative and Qualitative Results for 3D Segmentation

3D医学图像分割通常比2D分割更具挑战性，因为它涉及处理大量信息。随着维度增加，计算复杂性急剧增加，分辨率的提高导致计算量呈立方增长。准确的空间关系建模对于实现满意的分割结果也至关重要。这些因素使得具有计算效率的xLSTM构建块非常适合这一任务。

作者们在BraTS2023和腹部MRI 3D数据集上进行了评估。表2展示了在BraTS2023数据集上的性能比较，包括全肿瘤（WT）、肿瘤核心（TC）和增强肿瘤（ET）区域的指标。所提出的方法在所有评估指标上表现优越，包括Dice和HD95，超过了其他基线方法，如SegresNet、UX-Net、MedNeXt、UNETR、SwinUNETR、SwinUNETR-V2和SegMamba。具体而言，所提出的方法获得了最高的平均Dice得分91.80，突显了其在准确分割脑肿瘤区域方面的有效性。表3展示了腹部MRI 3D数据集上器官分割的性能比较。所提出的方法xLSTM-UNet_Bot获得了最高的Dice得分0.8483和最佳的NSD得分0.9153，超过了其他方法如nnU-Net、SegResNet、UNETR、SwinUNETR和U-Mamba_Bot。这表明了作者们的方法在分割MRI图像中的腹部器官时的稳健性和准确性。

xLSTM在复杂空间域中有效建模语义信息的能力是成功的关键因素。优越的实验结果也表明，基于xLSTM的构建块非常适合解决复杂成像应用中的语义分割挑战。