Source-linked AI summary
VM-UNET-V2 Rethinking Vision Mamba UNet for Medical Image Segmentation
Mingya Zhang, Yue Yu, Limei Gu, Tingsheng Lin, Xianping Tao
TL;DR
Medical image segmentation requires long-range modeling but CNNs have limited global context and Transformers have quadratic computational complexity. VM-UNetV2 addresses this with VSS blocks and SDI-based feature infusion, achieving competitive segmentation performance and efficient complexity measures across evaluated datasets.
Problem
CNNs have limited long-range modeling, while Transformer self-attention has quadratic complexity, motivating efficient architectures for medical image segmentation.
Method
VM-UNetV2 uses VSS blocks and SDI to process encoder features and skip connections, with VMamba pretraining and deep supervision.
Results
VM-UNetV2 achieves competitive segmentation results across skin-lesion and polyp datasets, outperforming UNetV2 by up to 7% in mIoU on polyp-related datasets.
Takeaways & Limitations
VM-UNetV2 combines competitive segmentation performance with favorable FLOPs, parameter count, and inference speed.
Abstract
from arXiv · showhide
In the field of medical image segmentation, models based on both CNN and Transformer have been thoroughly investigated. However, CNNs have limited modeling capabilities for long-range dependencies, making it challenging to exploit the semantic information within images fully. On the other hand, the quadratic computational complexity poses a challenge for Transformers. Recently, State Space Models (SSMs), such as Mamba, have been recognized as a promising method. They not only demonstrate superior performance in modeling long-range interactions, but also preserve a linear computational complexity. Inspired by the Mamba architecture, We proposed Vison Mamba-UNetV2, the Visual State Space (VSS) Block is introduced to capture extensive contextual information, the Semantics and Detail Infusion (SDI) is introduced to augment the infusion of low-level and high-level features. We conduct comprehensive experiments on the ISIC17, ISIC18, CVC-300, CVC-ClinicDB, Kvasir, CVC-ColonDB and ETIS-LaribPolypDB public datasets. The results indicate that VM-UNetV2 exhibits competitive performance in medical image segmentation tasks. Our code is available at https://github.com/nobodyplayer1/VM-UNetV2.
1 Introduction
Medical image segmentation needs architectures that capture long-range information without Transformers’ quadratic cost. VM-UNetV2 combines SSM-based modeling with feature infusion and shows competitive results across skin-lesion and polyp datasets.
- Background: Encoder-Decoder networks with skip connections combine hierarchical encoder features with decoder-based pixelwise mask reconstruction.U-Net studies commonly modify backbones, skip connections, attention mechanisms, sampling, or feature fusion.
- Motivation: CNNs struggle with long-range information, whereas Transformer self-attention incurs quadratic computational cost in dense medical image segmentation.These limitations motivate architectures that preserve efficient computation while modeling global context.
- Contribution: VM-UNetV2 re-integrates low-level and high-level features, infusing semantic information into low-level features and detail into high-level features.The design is influenced by VMamba, VM-Unet, and UNetV2.
- Experiments: Experiments cover seven public datasets spanning skin-lesion and gastroenterology semantic segmentation, with outcomes suggesting competitive performance.The tested public datasets include ISIC17, ISIC18, CVC-300, CVC-ClinicDB, Kvasir, CVC-ColonDB, and ETIS-LaribPolypDB.
- Contribution: The study contributes an SSM-based UNet variant and explores combining SSMs with UNet architectures for more efficient and effective segmentation algorithms.The paper also identifies exhaustive multi-dataset experimentation as a primary contribution.
2 Methods
The method combines a discretized state-space formulation with a Vision Mamba U-Net architecture comprising Encoder, SDI, and Decoder modules. VSS blocks model image features, while SDI enhances and aligns hierarchical features before decoding.
- Preliminaries: SSM-based models discretize continuous state-space systems with a timescale parameter and can compute outputs by linear recurrence or global convolution.The structured convolutional kernel K has length L matching the input sequence length.
- VM-UNetV2 Architecture: Vision Mamba UNetV2 contains Encoder, SDI, and Decoder modules, with the Encoder generating multilevel features that are forwarded to SDI.The Encoder produces features at M levels, and SDI outputs correspond to the respective stages.
- VM-UNetV2 Architecture: The Encoder uses VSS blocks across four stages with channel counts [C, 2C, 4C, 8C] and configurable N3 values [2,9,27].C is set to 96, while N1 and N2 are each 2; the configurations support Tiny and Small VMamba backbones for ablations.
- VSS Block: Each VSS block embeds inputs, splits them into two streams, processes one with depth-wise convolution and SiLU before SS2D, then normalizes and combines the streams.The SS2D output passes through layer normalization before combination with the other stream.
- Training Objective: The segmentation tasks use Cross-Entropy and Dice loss because the dataset masks contain a target class and background.The loss weights λ1 and λ2 are often set to (1, 1).
3 Experiments and results
Experiments evaluate VM-UNetV2 on skin-lesion and gastrointestinal-polyp datasets using segmentation, efficiency, and ablation analyses. The model outperforms comparison models on reported accuracy measures and shows favorable computational results, while moderate encoder depth generally performs better.
- Segmentation results: VM-UNetV2 outperforms other models on ISIC datasets for mIoU, DSC, and Acc, and surpasses UNetV2 across polyp-dataset metrics.The reported mIoU improvement over UNetV2 reaches up to 7% on polyp-related datasets.
- Segmentation results: The comparative results are presented in Table 1 for ISIC17 and ISIC18 and Table 2 for Kvasir-SEG, ClinicDB, ColonDB, ETIS, and CVC-300.Both tables mark the best results in bold.
- Computational complexity: VM-UNetV2 has superior FLOPs and FPS to the other evaluated models in tests conducted on an NVIDIA V100 GPU.The complexity evaluation also considers parameter count and uses input images of size (3, 256, 256).
- Ablation studies: Encoder-depth and Deep Supervision ablations find relatively better segmentation metrics with encoder depth, while Deep Supervision is usually beneficial but not decisive.Whether to adopt Deep Supervision should be determined separately for different datasets.
4 Conclusions
The paper concludes that VM-UNetV2 is an SSM-based UNet medical image segmentation model using VSS blocks, SDI, VMamba initialization, and Deep Supervision. Tests on skin-disease and polyp datasets report competitive segmentation and computational efficiency.
- Conclusions: VM-UNetV2 uses VSS blocks in the Encoder and SDI in skip connections to process features in an SSM-based UNet architecture.The design is intended to use SSM-based model capabilities for medical image segmentation.
- Conclusions: Pre-trained VMamba weights initialize the Encoder, while Deep Supervision supervises multiple output features.These components are part of the proposed VM-UNetV2 configuration.
- Conclusions: Extensive tests on skin-disease and polyp datasets demonstrate that VM-UNetV2 is highly competitive in segmentation tasks.The conclusion summarizes evaluations across the two application areas.
- Conclusions: Complexity analysis indicates that VM-UNetV2 is efficient in FLOPs, Params, and FPS.The conclusion reports efficiency across all three listed computational measures.