PaperlayerView exact revision on arXiv

Source-linked AI summary

VM-UNet: Vision Mamba UNet for Medical Image Segmentation

Jiacheng Ruan, Jincheng Li, Suncheng Xiang

arXiv:2402.02491v2eess.IVcs.CV

TL;DR

CNNs have limited long-range modeling and Transformers have quadratic complexity, motivating more efficient medical image segmentation architectures. The paper proposes pure SSM-based VM-UNet with VSS blocks and an asymmetric U-shaped design, and reports competitive results on ISIC17, ISIC18, and Synapse.

  • Problem

    CNNs are limited by local receptive fields, Transformers have quadratic image-size complexity, and pure SSM-based medical segmentation remains insufficiently explored.

  • Method

    VM-UNet is a pure SSM-based U-shaped segmentation model built around VSS blocks, with asymmetric encoder-decoder stages and SS2D feature extraction.

  • Results

    VM-UNet achieves competitive performance across skin lesion and multi-organ segmentation experiments on ISIC17, ISIC18, and Synapse.

  • Takeaways & Limitations

    VM-UNet establishes a baseline for exploring pure SSM-based models in medical image segmentation.

  • Takeaways & Limitations

    The evaluation is limited to small model sizes, and skin lesion segmentation remains problematic for light-colored areas and hair-disturbed images.

Abstract

from arXiv · show

In the realm of medical image segmentation, both CNN-based and Transformer-based models have been extensively explored. However, CNNs exhibit limitations in long-range modeling capabilities, whereas Transformers are hampered by their quadratic computational complexity. Recently, State Space Models (SSMs), exemplified by Mamba, have emerged as a promising approach. They not only excel in modeling long-range interactions but also maintain a linear computational complexity. In this paper, leveraging state space models, we propose a U-shape architecture model for medical image segmentation, named Vision Mamba UNet (VM-UNet). Specifically, the Visual State Space (VSS) block is introduced as the foundation block to capture extensive contextual information, and an asymmetrical encoder-decoder structure is constructed with fewer convolution layers to save calculation cost. We conduct comprehensive experiments on the ISIC17, ISIC18, and Synapse datasets, and the results indicate that VM-UNet performs competitively in medical image segmentation tasks. To our best knowledge, this is the first medical image segmentation model constructed based on the pure SSM-based model. We aim to establish a baseline and provide valuable insights for the future development of more efficient and effective SSM-based segmentation systems. Our code is available at https://github.com/JCruan519/VM-UNet.

I. INTRODUCTION

Medical image segmentation has advanced through CNN, Transformer, and SSM-based approaches, but existing architectures trade local modeling limits against quadratic computational cost. VM-UNet introduces a pure SSM-based U-shaped model and evaluates it competitively across three datasets.

  • Motivation: CNNs struggle with long-range information because of local receptive fields, while Transformers incur quadratic complexity for dense prediction.These limitations motivate an alternative architecture for medical image segmentation.
  • Motivation: SSMs such as Mamba model long-distance dependencies with linear complexity relative to input size.SSMs have also been studied in language and general vision before their medical segmentation applications.
  • Proposed model: VM-UNet is introduced as a pure SSM-based medical image segmentation model with encoder, decoder, and skip connections.Its encoder and decoder use VSS blocks with patch merging and patch expanding operations, respectively.
  • Evaluation: Comprehensive experiments on Synapse, ISIC17, and ISIC18 indicate that VM-UNet achieves competitive segmentation performance.The evaluation covers organ segmentation and skin lesion segmentation tasks.
  • Contributions: The paper positions VM-UNet as the first pure SSM-based segmentation model and a baseline for future SSM-based methods.The authors state that the baseline is intended to provide insights for more efficient and effective systems.

II. RELATED WORKS

Medical image segmentation methods have evolved from CNN-based U-shaped models toward Transformer-based architectures that better model long-range dependencies. Related approaches include UNet modifications, CNN–Transformer hybrids, and purely Transformer-based U-shaped models.

  • Overview: Existing medical image segmentation approaches can be broadly grouped into CNN-based, Transformer-based, and SSM-based methods.These approaches share the goal of efficiently modeling contextual information.
  • CNN-based methods: UNet established a simple, scalable U-shaped foundation for many subsequent medical image segmentation methods.Later approaches frequently modify its architecture while retaining its segmentation-oriented design.
  • Transformer-based methods: Transformer-based methods were introduced to improve long-range dependency modeling in medical image segmentation.TransFuse processes local and global features through parallel CNN and ViT encoders, while Swin-UNet combines Swin Transformer with UNet.

C. SSM-based Models

SSM-based models extend sequence modeling to long-distance dependencies with linear scaling, and VM-UNet applies this framework as a pure SSM-based medical segmentation baseline. The section establishes the continuous-to-discrete SSM formulation underlying the model.

  • SSM-based models: Mamba and related SSMs provide an alternative to CNNs and Transformers for long-distance sequence modeling with linear scaling in sequence length.Mamba adds a data-dependent SSM layer to the S4 framework.
  • SSM-based models: Pure SSM-based medical segmentation remains insufficiently explored, motivating VM-UNet as a foundational model for contextual information capture.Earlier medical segmentation models U-Mamba and SegMamba were hybrid SSM-CNN approaches.
  • SSM preliminaries: Classical SSMs map a one-dimensional input sequence through implicit states to an output using a linear ordinary differential equation.The state matrix is A, while B and C are projection parameters.
  • SSM preliminaries: S4 and Mamba discretize the continuous system using a timescale parameter and typically apply zero-order hold for deep-learning computation.The resulting discrete system can be evaluated through linear recurrence or global convolution.
  • VM-UNet: VM-UNet is organized around an overall architecture and a VSS block that serves as its core component.The paper introduces the architecture before detailing the VSS block and training loss.

B. Vision Mamba UNet (VM-UNet)

VM-UNet uses an asymmetric encoder–decoder with patch-based resolution changes, VSS blocks, and parameter-free additive skip connections. Its design allocates VSS blocks across progressively changing channel dimensions to produce segmentation outputs.

  • Overall architecture: VM-UNet contains patch embedding, an encoder, a decoder, a final projection layer, and skip connections in an asymmetric structure.The architecture differs from prior symmetrical U-shaped designs.
  • VSS block: The VSS block is VM-UNet’s core module, using two normalized branches and SS2D-based feature extraction.One branch uses a linear layer and activation, while the other adds depthwise separable convolution before SS2D.
  • Encoder: Patch embedding converts the input into non-overlapping 4 × 4 patches with default channel dimension C = 96 before encoder processing.Layer normalization is applied after embedding.
  • Encoder: The encoder has four stages using VSS blocks and channel counts [C, 2C, 4C, 8C].Patch merging follows the first three stages to reduce spatial dimensions and increase channels.
  • Decoder: The decoder uses VSS blocks with channel counts [8C, 4C, 2C, C], while patch expanding restores spatial resolution.A final projection layer restores feature size to match the segmentation target.
  • Skip connections: Skip connections use straightforward addition without introducing additional parameters.This provides feature fusion without extra learned skip-connection modules.

C. VSS Block

The VSS block is VM-UNet’s core module, combining two processing branches with SS2D-based directional scanning, normalization, feature mixing, and residual output formation.

  • VSS Block: The VSS block splits normalized input into two branches, sending one through linear activation and the other through convolution, activation, and SS2D.The branches are subsequently merged through element-wise multiplication.
  • VSS Block: The block applies Layer Normalization, element-wise multiplication, linear feature mixing, and a residual connection to produce its output.SiLU is used as the default activation function.
  • VSS Block: SS2D expands an input image into four directional sequences before S6 processing and later merges the scanned features.The directions run along both diagonals in both traversal orientations.
  • Loss Function: VM-UNet uses fundamental BceDice and CeDice losses for binary and multi-class segmentation, respectively.The cited loss formulation combines cross-entropy and Dice terms with empirically default weights of 1.
  • Contribution: The paper presents VM-UNet as a baseline for exploring purely SSM-based medical image segmentation models.The authors describe it as the first medical image segmentation model constructed from a pure SSM-based model.

IV. EXPERIMENTS

The experiments evaluate VM-UNet on skin lesion and organ segmentation tasks using the ISIC17, ISIC18, and Synapse datasets.

  • Experiments: VM-UNet is evaluated on skin lesion segmentation and organ segmentation tasks.The experiments cover both lesion and multi-organ medical image segmentation settings.

A. Datasets

The study uses two skin lesion datasets and one multi-organ abdominal CT dataset, with specified training-test splits and evaluation metrics.

  • ISIC Datasets: ISIC17 and ISIC18 are publicly available skin lesion segmentation datasets containing 2,150 and 2,694 dermoscopy images, respectively.The datasets provide segmentation mask labels.
  • SS2D Implementation: The S6 block receives features shaped [B, L, D] and outputs features with the same shape after recurrent state updates.Its pseudo-code applies linear projections, state transitions, and sequence reconstruction.
  • ISIC Datasets: The ISIC17 split contains 1,500 training images and 650 test images, while ISIC18 contains 1,886 training images and 808 test images.Both datasets use a 7:3 training-test split.
  • Synapse Dataset: Synapse contains 30 abdominal CT cases with 3,779 axial images covering eight abdominal organs.The dataset uses 18 cases for training and 12 cases for testing, reporting DSC and HD95.

B. Implementation Details

Implementation compares VM-UNet with segmentation models across skin lesion and Synapse tasks, while also examining pretrained initialization and dropout settings.

  • ISIC Results: VM-UNet achieves the best mIoU, DSC, Acc, and Sen on ISIC17, and leads mIoU, DSC, and Acc on ISIC18.Against TransFuse, it improves ISIC18 mIoU by 0.72% and DSC by 0.44%.
  • Synapse Results: On Synapse, VM-UNet improves DSC by 1.95% and HD95 by 2.34mm over Swin-UNet.For the Stomach organ, VM-UNet achieves a DSC of 81.40%.

1) Initial weights.:

VM-UNet benefits from stronger pretrained initialization, while dropout and architecture choices affect performance differently across datasets and configurations.

  • Initial weights: 2.67% average mIoU and 1.65% average DSC improvements were obtained across the two ISIC datasets with VMamba-S pretrained weights.The comparison is against VM-UNet without pretrained weights.
  • Dropout value: 80.45% mIoU and 89.17% DSC were the best ISIC17 results without dropout, whereas ISIC18 peaked at 81.35% mIoU and 89.71% DSC with dropout 0.2.Increasing dropout harmed ISIC17 performance, while moderate dropout improved ISIC18 performance.
  • Architecture: The asymmetric strong-encoder and weak-decoder design reduced parameters and computational load compared with a symmetric architecture.The symmetric configuration increased parameter count by 0.1M and computational load by 0.24 GFLOPs, while performance declined.

4) Input Resolution.:

The input-resolution study examines how longer visual sequences affect Mamba-based segmentation, while visual comparisons show robustness across challenging ISIC18 cases.

  • Input resolution: Increasing image resolution increases the input-sequence length used to investigate Mamba’s visual-domain performance.The study specifically examines pure SSM-based medical image segmentation as resolution increases.
  • Input resolution: Experiments at 512×512 resolution required a single NVIDIA A100 GPU because an NVIDIA RTX A6000 encountered an out-of-memory issue.The stated resolution-specific hardware constraint applies to the 512×512 experiments.
  • Input resolution: The input-resolution analysis did not meet expectations, and the specific reasons require further investigation.This is the stated limitation of the resolution study.
  • Visualization: VM-UNet remains more stable than TransFuse and UTNetV2 on non-target misclassification and small-target segmentation in ISIC18 visualizations.The visual comparisons also describe accurate delineation of complex target boundaries.
  • Visualization: VM-UNet reduces redundant segmentation predictions and accurately delineates complex target edges in the reported ISIC18 visualizations.Figure 4 presents visual comparisons on the ISIC18 dataset, with GT denoting ground truth.

F. Discussion

The discussion identifies limitations in VM-UNet’s current SSM-based version, including sequence-length generalization, model size, resolution coverage, and skin-lesion robustness. It also frames the model as a baseline for future medical-imaging applications.

  • Limitations: VM-UNet may generalize less effectively beyond the training sequence length because structured SSMs favor continuous-time data modalities.The authors identify this as a limitation for datasets exceeding the training sequence length.
  • Limitations: Approximately 30M parameters leave room for manual streamlining or compression to improve applicability in real-world medical scenarios.The discussion explicitly proposes design or compression strategies to reduce the model size.
  • Limitations: The empirical evaluation uses small model sizes and does not systematically test pure SSM segmentation at larger resolutions.The authors state that these settings warrant further research before real-world deployment.
  • Limitations: On skin-lesion segmentation, VM-UNet segments dark areas but struggles with light-colored regions and is sensitive to hair disturbance.The authors link these failures to contextual discriminability in visual representation and note possible downstream impacts.
  • Conclusion and future work: VM-UNet is presented as a first pure SSM-based medical image segmentation baseline, with future work extending SSMs to detection, registration, and reconstruction.The conclusion describes pure SSM models as competitive and identifies additional medical-imaging tasks for future exploration.
Loading 2402.02491v2…