xLSTM成功杀入医学图像分割领域，取代Mamba成为新SOTA！魔芯科技联合多单位推出

CV君

发布于 2024-07-12 16:49:14

1680

发布于 2024-07-12 16:49:14

近年来，Mamba作为一种处理长距离依赖关系的计算模块，在医学图像分割领域取得了显著进展。而近期，来自浙江大学，魔芯科技，中科大等单位的研究人员将Mamba替换为xLSTM，提出 xLSTM-UNet(xLSTM-UNet can be an Effective 2D \& 3D Medical Image Segmentation Backbone with Vision-LSTM (ViL) better than its Mamba Counterpart)，通过引入扩展长短期记忆网络（xLSTM），在多个医学图像分割任务中表现出色，超越了基于Mamba的模型！

论文链接：https://arxiv.org/abs/2407.01530
开源代码：https://github.com/tianrun-chen/xLSTM-UNet-PyTorch

什么是xLSTM-UNet？

xLSTM-UNet是一种结合了扩展长短期记忆网络（xLSTM）和经典UNet架构的深度学习模型。传统的卷积神经网络（CNNs）在处理长距离依赖关系时存在局限，而视觉Transformer（ViTs）在处理高分辨率图像时计算开销巨大。xLSTM-UNet通过将xLSTM引入UNet架构，成功地克服了这些挑战。

为什么xLSTM-UNet如此特别？

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

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

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

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

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

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

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

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