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.

    Source: §I · p. 1, C. SSM-based Models · p. 3

  • 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.

    Source: C. VSS Block · p. 3

  • Results

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

    Source: 1) ISIC Datasets · p. 5, 2) Synapse Dataset · p. 5

  • Takeaways & Limitations

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

    Source: §I · p. 2, §V · p. 8

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

    Source: F. Discussion · p. 8

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.Source: §I · p. 1
  • 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.Source: §I · p. 1

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.Source: §I · p. 2

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.Source: §I · p. 2

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.Source: §I · p. 2

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.Source: §II · p. 2

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.Source: A. CNN-based Models · p. 2

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.Source: B. Transformer-based Models · p. 2

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.Source: C. SSM-based Models · p. 2
  • 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.Source: C. SSM-based Models · p. 3

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.Source: A. Preliminaries · p. 3
  • 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.Source: A. Preliminaries · p. 3

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.Source: A. Preliminaries · p. 3

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.Source: B. Vision Mamba UNet (VM-UNet) · p. 3

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.Source: Figure 2 · p. 3

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.Source: B. Vision Mamba UNet (VM-UNet) · p. 3
  • 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.Source: B. Vision Mamba UNet (VM-UNet) · p. 3

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.Source: B. Vision Mamba UNet (VM-UNet) · p. 3

Skip connections

  • Skip connections use straightforward addition without introducing additional parameters.This provides feature fusion without extra learned skip-connection modules.Source: B. Vision Mamba UNet (VM-UNet) · p. 3

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.Source: C. VSS Block · pp. 3–4
  • 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.Source: C. VSS Block · pp. 3–4
  • 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.Source: C. VSS Block · p. 4, Figure 3 · p. 4

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.Source: D. Loss Function · p. 4

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.Source: D. Loss Function · p. 4

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.Source: §IV · p. 4

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.Source: A. Datasets · pp. 4–5

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.Source: A. Datasets · p. 5

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.Source: A. Datasets · p. 5

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.Source: A. Datasets · p. 5

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%.Source: 1) ISIC Datasets · p. 5

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%.Source: 2) Synapse Dataset · p. 5

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.Source: 1) Initial weights · p. 5

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.Source: 2) Dropout Value · p. 6

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.Source: 3) Encoder-Decoder Architectures · p. 6

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.Source: 4) Input Resolution · p. 6
  • 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.Source: Table VI · p. 7
  • The input-resolution analysis did not meet expectations, and the specific reasons require further investigation.This is the stated limitation of the resolution study.Source: 4) Input Resolution · p. 7

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.Source: E. Visualization · pp. 7–8
  • 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.Source: Figure 4 · p. 7, E. Visualization · p. 8

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.Source: F. Discussion · p. 8
  • 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.Source: F. Discussion · p. 8
  • 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.Source: F. Discussion · p. 8
  • 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.Source: F. Discussion · p. 8, Figure 5 · p. 8

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.Source: §V · p. 8

Related papers

Loading 2402.02491v2…