PaperlayerView exact revision on arXiv

Source-linked AI summary

UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation

Renkai Wu, Yinghao Liu, Pengchen Liang, Qing Chang

arXiv:2403.20035v3eess.IVcs.CV

TL;DR

Medical image segmentation for mobile and clinical settings needs models that balance performance with constrained computation and memory. The paper analyzes Mamba’s parameter drivers and introduces parallel Vision Mamba through the PVM Layer in UltraLight VM-UNet. The model achieves competitive skin lesion segmentation with 0.049M parameters and 0.060 GFLOPs.

  • Problem

    Clinical and mobile medical settings require segmentation models with low parameters and minimal computational memory footprint, but improving predictive power often increases model parameters.

  • Method

    The paper analyzes Mamba’s parameter influences and proposes the PVM Layer, which processes deep features with parallel Vision Mamba while keeping total processing channels constant.

  • Results

    0.049M parameters and 0.060 GFLOPs yield strong performance competitiveness across three public skin lesion segmentation datasets.

  • Takeaways & Limitations

    Parallel Vision Mamba remains competitive while substantially reducing parameters and computational complexity, supporting UltraLight VM-UNet as a lightweight Vision Mamba model.

Abstract

from arXiv · show

Traditionally for improving the segmentation performance of models, most approaches prefer to use adding more complex modules. And this is not suitable for the medical field, especially for mobile medical devices, where computationally loaded models are not suitable for real clinical environments due to computational resource constraints. Recently, state-space models (SSMs), represented by Mamba, have become a strong competitor to traditional CNNs and Transformers. In this paper, we deeply explore the key elements of parameter influence in Mamba and propose an UltraLight Vision Mamba UNet (UltraLight VM-UNet) based on this. Specifically, we propose a method for processing features in parallel Vision Mamba, named PVM Layer, which achieves excellent performance with the lowest computational load while keeping the overall number of processing channels constant. We conducted comparisons and ablation experiments with several state-of-the-art lightweight models on three skin lesion public datasets and demonstrated that the UltraLight VM-UNet exhibits the same strong performance competitiveness with parameters of only 0.049M and GFLOPs of 0.060. In addition, this study deeply explores the key elements of parameter influence in Mamba, which will lay a theoretical foundation for Mamba to possibly become a new mainstream module for lightweighting in the future. The code is available from https://github.com/wurenkai/UltraLight-VM-UNet .

1 Introduction

Medical segmentation models must balance accuracy with the parameter and memory constraints of clinical and mobile settings. This paper proposes UltraLight VM-UNet and parallel Vision Mamba to reduce complexity while retaining competitive performance.

  • Convolution captures local features but struggles with remote information, while Transformer self-attention adds image-size-dependent quadratic complexity.
  • Clinical and mobile medical applications require low-parameter models with minimal computational memory footprint and good performance.
  • SSMs offer linear input-size complexity and remote-dependency modeling, motivating their use as lightweight foundations for vision models.
  • UltraLight VM-UNet uses a PVM Layer that processes deep features with parallel Vision Mamba while keeping the overall processing channel count constant.
  • Parallel Vision Mamba remains performance-competitive while reducing parameters and computational complexity, addressing the previously unknown effect of parallel connections in Mamba.
  • 0.049M parameters and 0.060 GFLOPs make UltraLight VM-UNet highly lightweight while maintaining competitive performance across three skin lesion segmentation datasets.

2 Related Work

Skin lesion segmentation builds on U-Net and related medical image segmentation advances, while this paper addresses the field’s model-size challenge with a lightweight Mamba-based alternative.

  • U-Net’s skip connections merge high-level and low-level features, supporting its importance in medical image segmentation.
  • Medical image segmentation commonly targets multi-scale variation and feature-refinement learning, with skin lesions presenting varied features and high melanoma-related lethality.
  • Prior skin lesion methods include attention-enhanced Swin U-Net, high-order spatial interaction, and adaptive higher-order interaction U-Net models.
  • The proposed UltraLight VM-UNet targets large model parameters while retaining strong competitive performance on three public skin lesion datasets.

3 Method

UltraLight VM-UNet combines a six-layer U-shaped architecture with parallel Vision Mamba processing for deep features. Its design reduces Mamba’s parameter burden by controlling channel-dependent components while preserving total processing channels and using attention bridges for multiscale skip fusion.

  • Architecture Overview: The six-layer U-shaped UltraLight VM-UNet uses channels [8, 16, 24, 32, 48, 64], convolution blocks for shallow features, and PVM Layers for deeper features.The decoder mirrors the encoder, while skip connections use attention bridges for feature fusion.
  • Mamba Parameter Impact Analysis: Mamba parameters depend on channels, SSM state size, convolutional kernel size, projection expansion, and step-size rank, with channel number having the strongest effect.The internal projection width d_inner is determined by expand × d_model, so it increases with input channels.
  • Mamba Parameter Impact Analysis: The input and output projections, intermediate SSM projections, internal convolution, A_logs, and D all contribute parameters controlled substantially by d_model and d_inner.The internal convolution contributes 4 × d_inner^2 + d_inner when d_conv is fixed to 4.
  • Mamba Parameter Impact Analysis: Reducing channels from 1024 to 256 lowers total parameters from 23,435,264 to 1,484,288, reducing the parameter explosion by 93.7%.This calculation keeps other parameters unchanged and demonstrates the dominant influence of channel number.
  • Parallel Vision Mamba Layer: PVM processing keeps total channels constant while maximizing parameter reduction and maintaining strong performance competitiveness.For a channel count of 1024, Method A reduces parameters by 93.7% per Vision Mamba, while the four-operation comparison Method B achieves a 74.8% overall reduction.
  • Parallel Vision Mamba Layer: The PVM Layer splits C-channel features into four C/4-channel branches, applies Mamba in parallel, concatenates them back to C channels, and projects the result.LayerNorm, residual concatenation, an adjustment factor, and projection are used in the processing path.
  • Architecture Overview: SAB and CAB modules fuse multistage features at different scales along the skip-connection path.SAB uses pooling and shared-weight dilated convolution, while CAB uses global average pooling, concatenation, fully connected layers, and sigmoid activation.

4 Experiment

Experiments evaluate UltraLight VM-UNet against lightweight and classical models, then test parallel Vision Mamba and PVM Layer across architectures. The results show substantial reductions in parameters and computational complexity while retaining competitive segmentation performance.

  • Comparison results: 99.82% fewer parameters and 98.54% fewer GFLOPs than VM-UNet, while maintaining excellent and highly competitive performance across ISIC2017, ISIC2018, and PH2.Compared with LightM-UNet, reductions are 87.84% in parameters and 84.65% in GFLOPs.
  • Comparison results: MALUNet has lower complexity than VM-UNet and LightM-UNet, but its performance is much lower than the Mamba-based models.This comparison highlights a performance–complexity imbalance for the convolution-based lightweight model.
  • Ablation experiments: Parallelism is evaluated across different SSM variants and datasets, with Table 3 marking ISIC2017, ISIC2018, and PH2 results separately.The table uses (P) to denote replacing the original form with quadruple parallelism.
  • Ablation experiments: Four-way parallel Vision Mamba uses C/4 channels per branch and reduces the overall model parameters to 36.03% of conventional Vision Mamba, with little overall GFLOPs change.Two-way parallelism reduces parameters to 51.47% of the conventional connection.
  • Ablation experiments: Replacing convolution, Vision Transformer, Mamba, or Vision Mamba modules with PVM Layer significantly decreases parameters and GFLOPs while preserving competitive performance.The reported plug-and-play experiments also improve performance for MALUNet and EGE-UNet after key-module replacement.

5 Conclusion

The study analyzes which Mamba factors drive parameter growth and proposes PVM Layer, a parallel Vision Mamba design for deep-feature processing. The resulting UltraLight VM-UNet reaches very low complexity while retaining competitive skin-lesion segmentation performance.

  • Conclusion: PVM Layer processes features with four Vision Mamba modules in parallel, each receiving one-fourth of the initial channels.The design can be generalized to any Mamba variant as a PxM Layer.
  • Conclusion: 93.7% fewer Vision Mamba parameters result when each module processes one-fourth of the original channels, because SSM input-channel count strongly affects parameter growth.The paper attributes this reduction to the explosive effect of Mamba’s SSM input-channel dimension on parameters.
  • Conclusion: 0.049M parameters and 0.060 GFLOPs characterize the proposed UltraLight VM-UNet.The model is reported as 99.82% smaller than VM-UNet and 87.84% smaller than LightM-UNet by parameter count.

A Mamba variant (SS2D) parameter impact analysis

SS2D-based VSS Blocks contain multiple parameter-sensitive components, with input-channel count exerting the strongest influence. Reducing channels sharply lowers total parameters while preserving the overall processing-channel design.

  • VSS Block composition: SS2D is embedded in a VSS Block whose branches combine linear layers, convolution, SiLU, selective scanning, and LayerNorm.The branches merge their outputs after processing.
  • SS2D processing: SS2D expands image sequences in four directions, extracts features with an S6 block, then merges them back to the original image size.The scan directions run between opposite image corners.
  • Parameter factors: The number of input channels, S6 state dimension, convolution-kernel size, expansion multiplier, and projection rank all affect VSS Block parameters.The paper identifies input-channel count as the most influential factor.
  • Parameter factors: d_inner equals expand multiplied by d_model, so the internal expansion channel grows with the projection multiplier and input-channel count.The default expansion multiplier is 2.
  • Parameter factors: Input and output projections, normalization, S6 projections, internal convolution, A_logs, and Ds contribute parameters controlled substantially by d_model and d_inner.The cited formulas show contributions from projections, convolution, state parameters, and selective integration.
  • Channel reduction example: 45,504,512 to 2,921,984 parameters results when input channels decrease from 1024 to 256, a 93.6% reduction with other parameters unchanged.The example reduces channels to one quarter of the original.

B.1 Datasets

The study evaluates UltraLight VM-UNet on three public skin-lesion datasets, using ISIC2017 and ISIC2018 for development and PH2 for external validation.

  • Datasets: The experiments use ISIC2017, ISIC2018, and PH2, with PH2 serving as external validation after training weights on ISIC2017.All three datasets provide dermatoscopic images with segmentation masks.
  • ISIC2017: ISIC2017 contains 2,000 images split into 1,250 training, 150 validation, and 600 testing images.Images are standardized to 256×256 pixels for model input.
  • ISIC2018: ISIC2018 contains 2,594 images split into 1,815 training, 259 validation, and 520 testing images.Images are standardized to 256×256 pixels for model input.
  • PH2: PH2 contains 200 images, all used for external validation after standardization to 256×256 pixels.The original images are 768×560 pixels.

B.2 Implementation details

The experiments use a common PyTorch training setup and shared augmentation procedure, with segmentation assessed using DSC, sensitivity, specificity, and accuracy.

  • Experimental setup: Experiments use Python 3.8, PyTorch 1.13.0, and a single NVIDIA V100 GPU with 32GB memory.The same implementation environment is used across experiments.
  • Training setup: Training uses horizontal and vertical flips, random rotations, BceDice loss, AdamW, 250 epochs, batch size 8, and cosine annealing learning-rate scheduling.The shared augmentations support fairer model comparisons.
  • Evaluation metrics: The evaluation metrics are Dice similarity coefficient, sensitivity, specificity, and accuracy.DSC measures similarity between predicted and ground-truth segmentation maps.
  • Evaluation metrics: Accuracy is computed as (TP + TN) divided by (TP + TN + FP + FN).TP, TN, FP, and FN denote true positives, true negatives, false positives, and false negatives.

C Visualization results

The visualization compares UltraLight VM-UNet with lightweight and classical segmentation models. The reported visual outcome is smoother, clearer, and more accurate lesion boundaries.

  • Visualization results: Figure 7 visualizes segmentation results from comparison experiments across three publicly available skin-lesion datasets.The figure presents segmentation graphs for the dataset comparisons.
  • Visualization results: UltraLight VM-UNet produces smooth, clear, and more accurate boundaries in the visualized segmentation results.The paper reports this comparison against several state-of-the-art lightweight and classical models.
  • Ablation experiments: Table 5 reports ablation experiments examining the effects of individual UltraLight VM-UNet modules.The table is identified as an ablation study rather than a visualization result.

D.1 Impact of components in the UltraLight VM-UNet

Ablations show that the PVM Layer is central to UltraLight VM-UNet’s efficiency and performance, while SAB and CAB support skip connections. Channel-number configurations are compared with and without parallel connections to assess their effects.

  • PVM Layer ablation: Replacing PVM Layers with standard 3×3 convolutions increased parameters and GFLOPs while decreasing performance.Replacing encoder or decoder PVM Layers separately increased parameters by 63.26%; replacing both increased parameters by 151% and GFLOPs by 25%.
  • PVM Layer ablation: Removing PVM Layers caused decreases across performance measures alongside increases in parameters and GFLOPs, demonstrating their crucial role.
  • Skip connections: SAB and CAB modules are used as a skip-connection bridge to further improve segmentation performance.Their removal further reduced parameters and GFLOPs, but the ablation text reports the modules as performance-improving components.
  • Channel-number ablation: Table 6 compares four channel-number combinations under parallel-free and quadruple-parallel configurations.DP denotes the percentage of parameter reduction with the parallel approach, while 74.80% is the theoretical reduction for localized module parameters.
Loading 2403.20035v3…