MobileUNETR: A Lightweight End-To-End Hybrid Vision Transformer For Efficient Medical Image Segmentation

TL;DR

This work tackles the need for efficient, accurate skin lesion segmentation by introducing MobileUNETR, a lightweight end-to-end hybrid CNN-Transformer architecture. It features a two-part design: a compact encoder that blends CNN-based local feature extraction with MobileViT-based global context, and a novel lightweight hybrid decoder that fuses multi-scale local and global information during upsampling. The model achieves strong segmentation performance on ISIC 2016/2017/2018 and PH2 datasets while drastically reducing parameter count (~3M) and compute (~1.3 GFLOPs), outperforming several large CNN/Transformer baselines. These results suggest that carefully designed lightweight hybrids can rival larger models, enabling practical deployment in resource-limited clinical settings, with code to be released publicly.

Abstract

Skin cancer segmentation poses a significant challenge in medical image analysis. Numerous existing solutions, predominantly CNN-based, face issues related to a lack of global contextual understanding. Alternatively, some approaches resort to large-scale Transformer models to bridge the global contextual gaps, but at the expense of model size and computational complexity. Finally many Transformer based approaches rely primarily on CNN based decoders overlooking the benefits of Transformer based decoding models. Recognizing these limitations, we address the need efficient lightweight solutions by introducing MobileUNETR, which aims to overcome the performance constraints associated with both CNNs and Transformers while minimizing model size, presenting a promising stride towards efficient image segmentation. MobileUNETR has 3 main features. 1) MobileUNETR comprises of a lightweight hybrid CNN-Transformer encoder to help balance local and global contextual feature extraction in an efficient manner; 2) A novel hybrid decoder that simultaneously utilizes low-level and global features at different resolutions within the decoding stage for accurate mask generation; 3) surpassing large and complex architectures, MobileUNETR achieves superior performance with 3 million parameters and a computational complexity of 1.3 GFLOP resulting in 10x and 23x reduction in parameters and FLOPS, respectively. Extensive experiments have been conducted to validate the effectiveness of our proposed method on four publicly available skin lesion segmentation datasets, including ISIC 2016, ISIC 2017, ISIC 2018, and PH2 datasets. The code will be publicly available at: https://github.com/OSUPCVLab/MobileUNETR.git

Figures (7)

• Figure 1: Examples of typical skin cancer instances showcasing the typical noise and complicates in dermoscopic images.
• Figure 2: MobileUNETR Architecture: (a) Main MobileUNETR architecture showcases a hierarchical hybrid encoder-decoder architecture to extract and combine coarse and fine-grained features in an end-to-end framework. (b) Lightweight convolution stem for low-level feature extraction and spatial downsampling. (c) Hybrid Encoder Block efficiently extracts local and global features at each stage. (d) MobileViT Block for global feature extraction and long-range context understanding. (e) Novel Decoder Block for efficient upscaling and combining local/global features while allowing the model to dynamically adapt features during the decoder stage.
• Figure 3: Parameter count and GFLOP distributions (smaller is better) spanning SOTA models ranging from CNN, Transformer and Hybrid Architectures. We demonstrate a significant reduction in both model size and computational complexity compared to current SOTA architectures while achieving superior performance.
• Figure :
• Figure :