MLLA-UNet: Mamba-like Linear Attention in an Efficient U-Shape Model for Medical Image Segmentation

📄 arXiv: 2410.23738v1 📥 PDF

作者: Yufeng Jiang, Zongxi Li, Xiangyan Chen, Haoran Xie, Jing Cai

分类: eess.IV, cs.CV

发布日期: 2024-10-31

🔗 代码/项目: GITHUB

💡 一句话要点

提出MLLA-UNet，结合线性注意力与Mamba机制，高效解决医学图像分割难题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 医学图像分割 U-Net 线性注意力 Mamba架构 深度学习 自适应选择 多尺度特征融合

📋 核心要点

  1. 医学图像分割面临高变异性、低对比度等挑战，传统方法难以有效处理。
  2. MLLA-UNet结合线性注意力和Mamba机制，降低计算复杂度的同时保持分割精度。
  3. 实验结果表明，MLLA-UNet在多个数据集上达到SOTA，平均DSC达88.32%。

📝 摘要（中文）

医学成像技术的进步带来了更复杂多样的图像，同时也带来了诸如高解剖变异性、模糊的组织边界、低器官对比度和噪声等挑战。传统的分割方法难以应对这些挑战，因此深度学习方法，特别是U型架构，变得越来越重要。然而，标准自注意力的二次复杂度使得Transformer在高分辨率图像上计算成本过高。为了解决这些挑战，我们提出了MLLA-UNet（类Mamba线性注意力UNet），这是一种新颖的架构，通过其线性注意力和Mamba启发式自适应机制的创新组合，以及用于增强特征处理的有效对称采样结构，实现了线性计算复杂度，同时保持了高分割精度。我们的架构有效地保留了重要的空间特征，同时以降低的计算复杂度捕获了长距离依赖关系。此外，我们还引入了一种用于多尺度特征融合的新型采样策略。实验表明，MLLA-UNet在六个具有24个不同分割任务的具有挑战性的数据集（包括但不限于FLARE22、AMOS CT和ACDC）上实现了最先进的性能，平均DSC为88.32％。这些结果突出了MLLA-UNet优于现有方法。我们的贡献包括新颖的2D分割架构及其经验验证。

🔬 方法详解

问题定义：医学图像分割任务面临着高分辨率图像带来的计算负担，以及医学图像本身固有的挑战，例如器官形状和大小的变异性、组织边界模糊、器官对比度低以及噪声干扰。传统的基于卷积神经网络的分割方法在捕捉长距离依赖关系方面存在局限性，而标准的Transformer架构由于自注意力的二次复杂度，难以应用于高分辨率医学图像。

核心思路：MLLA-UNet的核心思路是利用线性注意力机制和Mamba架构的自适应选择能力，在降低计算复杂度的同时，有效地捕捉医学图像中的长距离依赖关系和关键特征。通过Mamba-like的机制，模型能够动态地调整对不同特征的关注程度，从而更好地适应医学图像的复杂性和多样性。

技术框架：MLLA-UNet采用经典的U-Net架构，包含编码器和解码器两个主要部分。编码器负责提取图像的特征表示，解码器则利用这些特征逐步恢复分割结果。关键模块是MLLA（Mamba-Like Linear Attention）模块，它替代了传统U-Net中的卷积层或标准自注意力层。此外，论文还提出了一种新的采样策略，用于多尺度特征融合，以进一步提高分割精度。

关键创新：MLLA-UNet的关键创新在于MLLA模块，它结合了线性注意力和Mamba架构的优点。线性注意力降低了计算复杂度，使其能够处理高分辨率图像，而Mamba架构的自适应选择能力则增强了模型对关键特征的关注。这种结合使得MLLA-UNet能够在计算效率和分割精度之间取得良好的平衡。与现有方法相比，MLLA-UNet在捕捉长距离依赖关系和处理高分辨率图像方面具有优势。

关键设计：MLLA模块的具体实现细节包括线性注意力的计算方式、Mamba架构的状态空间模型参数设置、以及如何将两者有效结合。此外，新的采样策略也需要仔细设计，以确保多尺度特征融合的有效性。损失函数方面，可能采用Dice Loss或Cross-Entropy Loss等常用的分割损失函数，并根据具体数据集进行调整。网络深度和通道数等超参数的选择也需要通过实验进行优化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MLLA-UNet在六个具有挑战性的医学图像分割数据集上取得了最先进的性能，平均DSC达到了88.32%。与现有方法相比，MLLA-UNet在FLARE22、AMOS CT和ACDC等数据集上均取得了显著的提升，证明了其在医学图像分割任务中的优越性。代码已开源。

🎯 应用场景

MLLA-UNet在医学图像分割领域具有广泛的应用前景，可用于辅助医生进行疾病诊断、治疗计划制定和手术导航。例如，它可以应用于CT、MRI等多种医学影像的器官分割、病灶检测和肿瘤分割等任务。该研究的成果有助于提高医学图像分析的效率和准确性，从而改善患者的治疗效果。

📄 摘要（原文）

Recent advancements in medical imaging have resulted in more complex and diverse images, with challenges such as high anatomical variability, blurred tissue boundaries, low organ contrast, and noise. Traditional segmentation methods struggle to address these challenges, making deep learning approaches, particularly U-shaped architectures, increasingly prominent. However, the quadratic complexity of standard self-attention makes Transformers computationally prohibitive for high-resolution images. To address these challenges, we propose MLLA-UNet (Mamba-Like Linear Attention UNet), a novel architecture that achieves linear computational complexity while maintaining high segmentation accuracy through its innovative combination of linear attention and Mamba-inspired adaptive mechanisms, complemented by an efficient symmetric sampling structure for enhanced feature processing. Our architecture effectively preserves essential spatial features while capturing long-range dependencies at reduced computational complexity. Additionally, we introduce a novel sampling strategy for multi-scale feature fusion. Experiments demonstrate that MLLA-UNet achieves state-of-the-art performance on six challenging datasets with 24 different segmentation tasks, including but not limited to FLARE22, AMOS CT, and ACDC, with an average DSC of 88.32%. These results underscore the superiority of MLLA-UNet over existing methods. Our contributions include the novel 2D segmentation architecture and its empirical validation. The code is available via https://github.com/csyfjiang/MLLA-UNet.