xLSTM崛起！魔芯科技、中科大等新研究揭示其在医学图像分割中超越Mamba

公众号新闻

2024-07-09 06:07

近日，魔芯科技、浙江大学、中科大等多家科研机构的研究人员联合推出了一项突破性的医学图像分割技术——xLSTM-UNet。这一创新技术在 2D 和 3D 医学图像分割领域取得了显著成果。

项目页面：

http://tianrun-chen.github.io/xLSTM-Unet/

论文链接：

https://arxiv.org/abs/2407.01530

开源代码：

https://github.com/tianrun-chen/xLSTM-UNet-PyTorch

“简单替换Mamba为xLSTM，即可显著提升性能”

在医学图像处理领域，精准的图像分割技术对于疾病的诊断、治疗规划和研究至关重要。然而，随着高分辨率和高维成像模式的发展，处理长距离依赖关系的挑战变得愈发严峻。近年来，一些研究提出了整合具有长距离依赖性的计算模块，如 Mamba，并成功地将其集成到传统的 UNet 架构中，这些变体（如 UMamba、VM-Unet、Mamba-Unet、Swin-UMamba 和 SegMamba 等）展示了显著的性能提升。

然而，最近的一项研究表明，利用扩展长短期记忆网络（xLSTM）作为核心组件，可以显著提升图像分割性能。由魔芯科技、浙江大学和中科大等单位的研究团队开发的 xLSTM-UNet，通过将改进后 UNet 架构中的 Mamba 模块替换为 xLSTM，展现了令人惊讶的性能提升。

实验结果显示，仅通过这一简单的替换操作，xLSTM-UNet 在多个 2D 和 3D 医学图像分割数据集上均超越了现有的基于 CNN、基于 Transformer 以及基于 Mamba 的分割网络。这些数据集包括 CT 和 MRI 等多种医学影像，覆盖了不同的解剖结构和病变区域。通过广泛的实验验证，研究人员发现 xLSTM-UNet 不仅在分割精度上有所提升，还在计算效率和内存使用上表现出色。

具体来说，xLSTM-UNet 结合了卷积层的局部特征提取能力和 xLSTM 的长距离依赖捕捉能力。xLSTM 能够有效捕捉全局上下文信息，克服了传统 CNNs 在处理长距离依赖关系时的局限性，同时避免了 ViTs 在处理高分辨率图像时的巨大计算开销。研究表明，xLSTM 在序列建模方面的优势，使其在医学图像分割中也能发挥出色的性能。

这一发现不仅验证了 xLSTM 在图像分割领域的潜力，还为未来的研究提供了新的方向，即在传统深度学习架构中引入更强大的序列建模组件，以应对复杂的图像分析任务。xLSTM-UNet 的成功，展示了一种高效且强大的解决方案，为医学图像分割领域带来了新的希望。

1. xLSTM 的引入：xLSTM（扩展长短期记忆网络）是 LSTM 的继任者，能够有效处理长距离依赖关系，并且在计算和内存复杂度方面表现出线性增长。相比传统的 LSTM，xLSTM 在神经语言处理（NLP）和图像分类任务中表现出色，展示了其在处理序列建模任务中的强大能力。

2. 结合 UNet 架构：xLSTM-UNet 采用了经典的 UNet 架构，结合了卷积层和 xLSTM 的优势。具体而言，xLSTM-UNet 在编码器部分引入了多个 xLSTM 层，用于捕捉图像的长距离依赖关系和全局上下文信息。同时，保留了 UNet 架构中的跳跃连接（skip connections），将编码器中的特征直接传递到解码器，以保留图像的细节信息。

3. 多层次特征提取：xLSTM-UNet 在多个层次上进行特征提取，既包括局部特征，也包括全局特征。通过在编码器的多个层次上引入 xLSTM 块，xLSTM-UNet 能够在不同分辨率和感受野上提取丰富的特征信息，从而提高分割精度。

实验结果显著，性能全面提升！

在多个代表性的医学图像分割数据集上，xLSTM-UNet 展现了卓越的性能，显著超越了基于 CNN、Transformer 和 Mamba 的分割网络。具体而言，xLSTM-UNet 在腹部 MRI、内窥镜图像和显微镜图像等数据集上的表现均显著优于现有方法，展示了其卓越的性能和鲁棒性。

1. 腹部 MRI 数据集：在该数据集上，xLSTM-UNet 在 Dice 相似系数（DSC）和归一化表面距离（NSD）两个关键指标上均取得了最高分，显著优于先前的最先进模型 U-Mamba。这表明 xLSTM-UNet 在腹部器官分割任务中具有更高的精度和可靠性。

2. 内窥镜图像和显微镜图像细胞分割数据集：在这些数据集上，xLSTM-UNet 同样在 DSC 和 NSD 指标上取得了最佳成绩，证明了其在不同医学图像分割任务中的鲁棒性和可靠性。具体而言，xLSTM-UNet 在内窥镜图像中的外科手术器械分割和显微镜图像中的细胞分割任务中均表现出色，进一步验证了其广泛的适用性。

3. 3D 医学分割任务：在 BraTS2023 数据集上，xLSTM-UNet 的表现超越了其他基线方法。无论是 Dice 相似系数还是 HD95（Hausdorff 距离 95%），xLSTM-UNet 在所有评估指标上均显示出其在精确分割脑肿瘤区域方面的有效性。这表明 xLSTM-UNet 在处理复杂的 3D 医学图像分割任务时具有显著的优势。

通过在多个数据集上的全面实验，xLSTM-UNet 不仅展示了其在 2D 和 3D 医学图像分割任务中的卓越性能，还证明了其在不同应用场景中的广泛适用性和鲁棒性。具体而言，xLSTM-UNet 在腹部 MRI、内窥镜图像和显微镜图像等多种医学图像分割任务中均表现出色，显著超越了现有的基于 CNN、Transformer 和 Mamba 的分割网络。

其在 Dice 相似系数（DSC）、归一化表面距离（NSD）以及 Hausdorff 距离 95%（HD95）等关键评估指标上均取得了领先成绩，进一步验证了其在处理复杂医学图像分割任务时的高效性和可靠性。这些实验结果充分表明，xLSTM-UNet 不仅在特定任务中表现优异，而且具有广泛的应用潜力，能够适应多种医学图像分割需求。

xLSTM启发未来图像分割研究

xLSTM-UNet 在多个医学图像分割数据集上表现卓越，显著优于现有的基于 CNN、Transformer 和 Mamba 的网络。这种卓越表现得益于 xLSTM 的引入。xLSTM 在处理长距离依赖性方面表现出色，并且在其他领域已经取得了显著成就。通过将 xLSTM 与 UNet 架构相结合，研究人员成功地提升了模型的全局特征捕捉能力和计算效率。

xLSTM-UNet 解决了医学图像分割中处理长距离依赖性和优化计算资源的双重挑战，为该领域注入了新的活力。这项研究不仅展示了 xLSTM-UNet 的卓越性能，还为未来的研究和应用开辟了新的可能性。展望未来，随着 xLSTM 的进一步开发和优化，我们有理由相信，它将在图像分割乃至更广泛的领域中取得与 Mamba 和 Transformer 相媲美的成就，推动这一领域的持续进步。研究相关代码和所使用数据均已经开源。