Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yufeng Jiang, Zongxi Li, Xiangyan Chen, Haoran Xie, Jing Cai

TL;DR

This work proposes 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.

Abstract

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.

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

TL;DR

This work proposes 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.

Abstract

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.

Paper Structure

This paper contains 27 sections, 23 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Visual Transformers and Self-Attention Mechanisms
  4. Mamba and Linear Attention in Vision
  5. Mamba Models for 2D Medical Image Segmentation
  6. Methodology
  7. MLLA-UNet architecture
  8. Stem
  9. Feature Compression
  10. Feature Expansion
  11. Mamba-Like Linear Attention (MLLA) Block
  12. Position Encoding Strategy for Mamba-like Mechanism
  13. EDSM and EUSM
  14. Predicting head
  15. Loss function
  16. ...and 12 more sections

Figures (3)

  • Figure 1: Architecture Overview: The model consists of a stem for initial processing, followed by multiple MLLA Blocks for feature extraction at various scales. EDSM layers reduce spatial dimensions, while skip connections preserve information across different layers. EUSM layers reconstruct the output, aided by fully connected layers for classification. Feature concatenation merges information from different paths, and final expanding adjusts the output to the desired dimensions and class count.
  • Figure 2: Illustration of EDSM and EUSM.
  • Figure 3: Visualization of segmentation results from various methods. The first three rows depict WORD, BTCV and AMOS with CT images, the forth row showcases ATLAS2023 with CE-MRI images. The six columns from left to right correspond to the original image, the ground truth (GT), the nnUNetv2 results, the SwinUNetR results, the STU-Net-B results, and our MLLA-UNet results. Red and white boxes highlight challenging regions where our method demonstrates superior performance in preserving anatomical details and boundary accuracy compared to other methods.