Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

VM-UNet: Vision Mamba UNet for Medical Image Segmentation

Jiacheng Ruan, Jincheng Li, Suncheng Xiang

TL;DR

VM-UNet introduces a pure State Space Model-based U-Net variant for medical image segmentation, leveraging Vision Mamba's SS2D-based VSS blocks to capture long-range context with linear complexity. The architecture employs an asymmetric encoder–decoder and a simple additive skip connection, serving as a baseline for pure SSM-based segmentation. Across ISIC skin lesion datasets and the Synapse multi-organ dataset, VM-UNet achieves competitive to state-of-the-art results, highlighting the potential of SSMs as an efficient alternative to CNNs and Transformers. The work provides a reproducible baseline and paves the way for future development of more efficient SSM-based segmentation systems in medical imaging.

Abstract

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.

VM-UNet: Vision Mamba UNet for Medical Image Segmentation

TL;DR

VM-UNet introduces a pure State Space Model-based U-Net variant for medical image segmentation, leveraging Vision Mamba's SS2D-based VSS blocks to capture long-range context with linear complexity. The architecture employs an asymmetric encoder–decoder and a simple additive skip connection, serving as a baseline for pure SSM-based segmentation. Across ISIC skin lesion datasets and the Synapse multi-organ dataset, VM-UNet achieves competitive to state-of-the-art results, highlighting the potential of SSMs as an efficient alternative to CNNs and Transformers. The work provides a reproducible baseline and paves the way for future development of more efficient SSM-based segmentation systems in medical imaging.

Abstract

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.
Paper Structure (24 sections, 7 equations, 5 figures, 6 tables, 1 algorithm)

This paper contains 24 sections, 7 equations, 5 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. CNN-based Models
  4. Transformer-based Models
  5. SSM-based Models
  6. Methods
  7. Preliminaries
  8. Vision Mamba UNet (VM-UNet)
  9. VSS Block
  10. Loss Function
  11. Experiments
  12. Datasets
  13. Implementation Details
  14. Comparison to the State-of-the-art Models
  15. ISIC Datasets.
  16. ...and 9 more sections

Figures (5)

  • Figure 1: A brief introduction to the development of U-shape architecture in the field of medical image segmentation.
  • Figure 2: (a) The overall architecture of VM-UNet. (b) VSS block is the main construction block of VM-UNet, and SS2D is the core operation in VSS block.
  • Figure 3: (a) The scan expanding operation in SS2D. (b) The scan merging operation in SS2D.
  • Figure 4: Visual Comparisons on the ISIC18 Dataset. GT stands for ground truth.
  • Figure 5: Some failure examples of our VM-UNnet on ISIC18 Dataset.