MobileUNETR: A Lightweight End-To-End Hybrid Vision Transformer For Efficient Medical Image Segmentation
Shehan Perera Yunus Erzurumlu Deepak Gulati Alper Yilmaz








MobileUNETR






Perera et al.




Photogrammetric Computer Vision Lab, The Ohio State University 

Wexner Medical Center, The Ohio State University 

{perera.27, yilmaz.15, erzurumlu.1}@osu.edu

deepakkumar.gulati@osumc.edu




MobileUNETR: A Lightweight End-To-End Hybrid Vision Transformer For Efficient Medical Image Segmentation
    Shehan Perera10009-0005-3831-0404 Yunus Erzurumlu10009-0006-5798-5842 Deepak Gulati20000-0003-3374-5992 Alper Yilmaz10000-0003-0755-2628
    September 9, 2024
============================================================================================================================================



§ 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.githttps://github.com/OSUPCVLab/MobileUNETR.git.






§ INTRODUCTION





Skin cancer, among the most prevalent and rapidly increasing forms of cancer worldwide, poses a significant global health challenge <cit.>. Given the various forms of skin cancer that appear between patients and different levels of severity, accurately identifying and categorizing skin lesions becomes a complex task. One of the primary difficulties in diagnosing this form of cancer lies in visual inspection of the lesions. The subjective nature of the visual process, influenced by factors such as lighting conditions, individual expertise, and the inherent variability in the way skin cancer presents itself in different patients, make visual categorization a difficult task. To improve diagnostic precision, dermatologists use dermoscopy, a non-invasive technique for skin surface microscopy. Dermoscopy provides physicians with high-resolution images of the affected skin, allowing a closer examination of the characteristics of the lesion <cit.>. Although this advancement has undoubtedly improved the accuracy of human visual analysis, it has not completely eliminated the challenges associated with human subjectivity. Dermatologists, even with the help of dermoscopic images, may still differ in their interpretation of skin lesions. This lack of consistency in diagnosis among medical professionals emphasizes the need for additional tools that can offer objective and standardized assessments. Recognizing these challenges, there has been a growing effort to integrate Computer-Aided Diagnostic (CAD) systems to support physicians in the diagnosis of skin lesions. 

Early iterations of CAD systems designed for skin cancer segmentation were often approached through complicated multi-step image processing pipelines <cit.>, <cit.>. Techniques employed in these early iterations include color-space transformations, principal component analysis, and the use of hand-crafted features, to name a few. Despite their progress in medical diagnosis, these approaches struggled to accurately delineate affected skin regions. Rule-based and hand-crafted systems often oversimplified complex, variable skin lesions, including artifacts and noise from body hair. 

The development of deep learning and its adoption represents a crucial step towards enhancing the efficiency and accuracy of CAD systems. These systems employ advanced neural networks to delineate the boundaries of the lesions, allowing a more precise assessment of their characteristics. Deep learning algorithms, with their ability to automatically learn intricate patterns and features directly from data, have demonstrated superior performance in segmenting skin lesion <cit.>]. These algorithms can discern subtle variations in color, texture, and shape, adapting dynamically to the diverse manifestations of skin cancer between different individuals.  

Central to the success of deep learning for medical image segmentation is the introduction of the encoder-decoder architecture. Encoder-Decoder architectures implemented via Fully Convolutional Neural Networks (FCNNs) have particularly excelled in this domain and have become the State-Of-The-Art (SOTA) for many segmentation tasks <cit.>. Although highly successful, one of the major drawbacks of FCNN/CNN based approaches is their lack of long-range contextual understanding. Although CNNs excel at capturing local features within an image, they inherently struggle to gather broader context information or a global relationship between different elements. In particular, in the case of skin cancer, where lesions can vary significantly from patient to patient, a global understanding becomes crucial to help the model overcome ambiguities. To overcome context limitations within CNNs, researchers have resorted to larger and deeper models to help improve the overall receptive field through pure convolutions <cit.>. However, this solution comes with its own set of challenges. Larger models require more computational resources, making them computationally expensive and slower to train and deploy. Additionally, the pursuit of a larger receptive field through sheer model size may lead to diminishing returns, emphasizing the need for a more efficient and effective approaches. The integration of self-attention modules introduced in the Transformer <cit.> architecture with convolutional layers has been suggested as a means to enhance the non-local modeling capability <cit.> and offer promising long-range contextual understanding benefits for many downstream tasks. 

Originally developed for Natural Language Processing (NLP), the Transformer architecture has seen significant adoption to many computer vision tasks. With the initial Vision Transformer <cit.> that allowed Transformers to perform image classification, researchers were provided with an architecture that is capable of modeling long-range dependencies and gathering global context clues at every stage of the model. However, as a trade-off the self-attention mechanism, central to the Transformer architecture, proves computationally expensive, especially at large spatial dimensions. Additionally, ViTs produce single-scale features, in contrast to multi-scale features typically generated by CNN models <cit.>. This trade-off between global awareness and computational efficiency presents a significant challenge when employing transformer architectures in resource-limited real-world applications.



To overcome current bottlenecks in widely adopted CNN and Transformer architectures we introduce MobileUNETR, a novel end-to-end transformer based encoder decoder architecture for efficient image segmentation. At a high level, challenging and complex image segmentation tasks often benefit from feature extraction capabilities that consider local and global contextual information within the feature encoding stage. However, segmentation approaches typically focus on optimizing the feature extractor while overlooking the importance of developing novel decoding strategies. Common segmentation frameworks in medical imaging, utilizing complex CNN and/or Transformer structures, generally favor excluding Transformer based decoders, opting instead for pure CNNs <cit.>. This choice can be attributed to the fact that, despite being great at capturing global information, Transformers are unable to capture intricate local details which are highly useful when generating accurate segmentation masks. To overcome the over-reliance on pure CNN layers within the decoder stage, we propose a novel highly effective and light-weight decoder capable of learning and integrating local/global details to generate highly accurate segmentation masks. 

We demonstrate the advantages of MobileUNETR in terms of model size, run time complexity, and accuracy on four publicly available skin lesion segmentation datasets, including ISIC 2016 <cit.>, ISIC 2017 <cit.>, ISIC 2018 <cit.>, and PH2 <cit.> datasets. We demonstrate a significant increase in performance across all datasets and advanced architectures and training methodologies while reducing the model size and complexity by 10x and 23x, respectively.

Our main contributions can be summarized as follows. 



  * We propose a novel lightweight and efficient end-to-end Transformer based hybrid model for skin lesion segmentation, where local and global contextual features are enforced at each stage to retain global awareness of a given scene.   


  * To overcome the over-reliance on CNN based decoding strategies we introduce a novel Transformer based hybrid decoder that simultaneously utilize low-level and global features at different resolutions for highly accurate and well aligned mask generation.


  * The proposed architecture surpasses large and highly complex CNN, Transformer, and Hybrid models in segmentation with only 3 millions parameters and 1.3 GFLOP of complexity resulting in a 10x and 23x reduction in computational complexity, respectively, than the current SOTA models. 





§ RELATED WORKS


Skin lesion segmentation is critical in automated dermatological diagnosis; however, it is difficult due to lesion diversity and the presence of noise in the images. Traditional image processing methods have given way to advanced deep learning systems, particularly Convolutional Neural Networks (CNNs), and then Transformer-based methods, which have considerably improved segmentation accuracy and reliability.



 §.§ CNN Based Methods


After the increasing popularity of Deep Neural Networks (DNNs) and Convolutional Neural Networks (CNNs), they have become the go-to tools for skin lesion segmentation tasks. They have creatively solved difficulties such as feature discernment and data variability management. The field has seen notable developments, such as the introduction of a multistage fully convolutional network (FCN) to the field by <cit.>, which incorporates a parallel integration method to enhance the segmentation of skin lesions' boundaries. <cit.> have contributed similarly by creating an improved convolutional-deconvolutional network specifically optimized for dermoscopic image analysis and integrating various color spaces to better diagnose lesions. <cit.> expanded on this trend of architectural innovation with their DoubleU-Net, which combines multiple U-Net structures to improve segmentation accuracy. 

In parallel, efforts to develop automated detection systems have been prominent. <cit.> worked on the early detection of malignant skin lesions using dilated convolutions across multiple architectures such as VGG16, VGG19 by <cit.>, MobileNet by <cit.>, and InceptionV3 by <cit.>, as well as the HAM10000 dataset by <cit.>  for training and testing. The use of pre-trained networks and deep learning models is also evident in the multiple winning solutions at the ISIC 2018 Challenge <cit.>, where many built their model on the DeepLab <cit.> architecture using pre-trained weights from PASCAL VOC-2012 <cit.> and used ensemble approach's among others with models such as VGG16, U-net, DenseNet by <cit.>, and Inceptionv3, fine-tuning these with additional training iterations for state-of-art performance. 

<cit.> and <cit.> proved the flexibility of these models in varied context-aware settings by improving feature extraction in CNNs, the former through modified skip connections and the latter with multistage UNets. Furthermore, <cit.> introduced a new focal Tversky loss function to address data imbalance, a significant difficulty in medical imaging, improving the precision recall balance for small lesion structures.

The ISIC 2019 challenge also led to several new studies using CNNs for dermoscopic medical imaging. <cit.>, <cit.>, and <cit.> used a variety of CNN architectures with different data augmentation methods. These studies demonstrated CNNs ability to segment skin lesions locally, but their performance shortcomings can be attributed to their inability to extract valuable global context information.



 §.§ Transformer Based Methods


Being limited to only local features forced researchers to seek new approaches. This caused an evolution towards the usage of global feature-based tools. This evolution is distinguished by a shift from standard CNN-based techniques toward novel ways of using transformers and self-attention mechanisms. <cit.> pioneered the Dense Deconvolutional Network (DDN) in skin lesion segmentation, employing dense layers and chained residual pooling to capture long-range relationships, a significant departure from prior approaches. Furthermore, <cit.> investigated adversarial learning with SegAN, improving segmentation accuracy by adeptly capturing subtle relationships, a significant development in dermatological imaging. <cit.> and <cit.> substantially advanced skin lesion segmentation with their new methodologies. Mirikharaji and Hamarneh implemented a star-shape prior (SSP) in a fully convolutional network to improve accuracy and reliability by penalizing non-star-shaped regions while preserving global structures. The use of shape priors to segment complex skin lesion patterns was demonstrated in this study. <cit.> supplemented this with CPFNet, which uses pyramidal modules to collect global context in feature maps, successfully managing skin lesion variability and enhancing delineation accuracy in intricate lesion patterns. 

Using transformers in neural networks, pioneered by <cit.>, was a significant turning point. <cit.> and <cit.> introduced transformers to computer vision. <cit.> demonstrated the effectiveness of self-attention mechanisms in image recognition models, which is helpful for complex skin lesion patterns. Additionally, <cit.> created TransUNet, which combines Transformers and U-Net to improve medical picture segmentation. The strength of TransUNet is in effectively encoding picture patches from CNN feature maps, which is essential for capturing detailed global context in segmentation tasks. <cit.> demonstrated TransUNet's performance in skin lesion segmentation, stressing its superior accuracy and dice coefficient over standard models, emphasizing the benefits of merging CNNs with transformers in medical imaging. Moreover, <cit.> developed the boundary-aware Transformer (BAT) for segmentation of skin lesions. 




BAT incorporates a boundary-wise attention gate in its transformer structure to address unclear lesion boundaries, efficiently collecting global and local information in skin lesion imaging. FAT-Net, a feature-adaptive transformer network for segmentation of skin lesion, was introduced by <cit.>. FAT-Net adeptly maintains long-range dependencies and contextual nuances by incorporating an extra transformer branch into the standard encoder-decoder structure, precisely addressing the variability and irregularity in skin lesions and improving melanoma analysis. 





§ METHODOLOGY


In this section, we introduce MobileUNETR, our high-performance, efficient, and lightweight architecture for skin lesion segmentation. As shown in Figure <ref>, the core MobileUNETR architecture consists of two main modules: (1) First, a lightweight hybrid encoder that efficiently generates coarse high-level and fine-grained low-level features; and (2) A novel lightweight hybrid decoder that effectively combines multilevel features while factoring in local and global context clues to generate high-accuracy semantic segmentation masks.



 §.§ Model Complexity


The overarching goal of the medical imaging community is the pursuit of performance over complexity on a particular task such as skin lesion segmentation. One of the main contributions of the proposed MobileUNETR architecture is to demonstrate that well-constructed lightweight and efficient models can offer much better performance compared to large computationally expensive architectures. As seen in Figure <ref>, the proposed architecture is 10X smaller and 23X more computationally efficient against SOTA architectures in skin lesion segmentation while generating better results. Simplifying the model not only enhances training and performance on small datasets but also facilitates deployment in resource-limited environments.



 §.§ Encoder


Two major groups of deep learning architectures exist in medical vision research, CNNs and Transformers, each with their own advantages and disadvantages. CNNs have been the defacto approach for many medical vision applications due to its efficiency, natural inductive biases and its ability to hierarchically encode features. However, despite its success, pure CNN based feature encoders are unable to effectively gain global contextual understanding of a given scene. Many hand-crafted approaches have been proposed to help CNNs obtain a larger receptive field such as dilated convolutions <cit.> and deeper models, however, image size and computational complexity constraints further research is required to help improve overall performance. Unlike CNNs, Transformers are designed to achieve a true global understanding of a scene. However, their computational constraints at large spatial resolutions hinder their adoption for efficient deep learning applications. 

By exploiting the natural advantages and disadvantages of CNNs and Transformer architectures, our proposed encoder maximizes feature representation capabilities while significantly minimizing computational complexity and parameter count.  At a high level, the feature extraction modules can be broken down into two stages: 1) CNN based local feature extraction and downsampling, 2) Hybrid Transformer/CNN based local and global representation learning.

CNN based local feature extraction: End-to-end transformer models for computer vision, such as ViT and its derivatives <cit.> result in large computationally complex models due to large sequence lengths generated for each input image. By combining the sequence length bottleneck in Transformers and the natural tendency of ViTs to learn low-level features in early layers <cit.>, simple CNN based early feature extraction substitution can be added to significantly reduce the computational complexity of the architecture. Specifically, MobileNet <cit.> downsampling blocks are used within the proposed architecture to minimize the computational complexity of the low-level feature extraction stage without compromising the learned feature representations. Additionally, CNN based features allow the model to better incorporate spatial information compared to pure ViT based approaches while effectively reducing the spatial dimensions of the input data, allowing downstream transformer layers to efficiently learn global feature representations. 








Hybrid Transformer/CNN blocks: Once efficient down-sampling is performed via CNNs to mitigate the computational complexity associated with large spatial resolutions, the MobileViT block is used to simultaneously extract local and global representations. The MobileViT block allows us to incorporate the long-range contextual benefits of Transformers while maintaining spatial ordering and local inductive biases. The operation can be broken down into two main components, as seen in Figure <ref>. First, CNN-based depth-wise separable convolution <cit.> is applied to encode spatial information and project features into high-dimensional space. Finally, to model long-range dependencies, the tensor is unfolded into non-overlapping flattened patches, and self attention layers are applied to capture interpatch relationships. This combination allows each feature map to have local and global understanding of the scene at each stage, improving its contextual understanding of the scene. 



 §.§ Decoder


Most segmentation models emphasize the importance of the encoder stage for great segmentation performance. Here, local and global understanding is favored to ensure that relevant features that contain both are learned, compressed, and passed forward to the next stage. Most encoder-decoder approaches that use CNNs, Transformers, or CNN/Transformer dual encoders for feature extraction heavily rely on pure convolution to map extracted features to the final segmentation mask. A drawback of this approach is that, by using pure CNN layers within the decoder, we force the model to use information extracted at the bottleneck to learn features that ensure local continuity without providing it the capability to recalibrate itself using global contextual information. Additionally, naively stacking CNN layers can lead to large decoder modules, adding to the overall computational complexity of the encoder-decoder architecture. Our novel hybrid decoder architecture is a fast, computationally efficient, and lightweight approach that allows the model to hierarchically construct the final segmentation mask while ensuring that local and global context features are used at every stage of the decoding process. The proposed decoder module at  1.5 million parameters efficiently combine the benefits of CNN and Transformer architectures, into an alternative to CNN based decoding methods. 

Simple Hybrid Decoder: Typical CNN decoder modules in medical imaging <cit.> extract and refine the features of the encoder with a combination of transpose and standard convolutions. Using this structure allows the model to hierarchically increase the spatial resolution while refining features at each level with the help of features provided via skip connections. Despite their success, decoders that rely solely on CNNs face challenges in dynamically adapting their own features to ensure that the features learned at each stage are globally aligned. The proposed lightweight decoder performs three operations at each stage to ensure that the features extracted at each decoding stage are locally and globally aligned Figure <ref>e. First, the feature map from the previous stage is upsampled via transpose convolutions. Next, we perform a local refinement of the upsampled features by combining information with the respective skip connections. Finally, the Transformer/CNN hybrid layers are used to allow the model to dynamically adjust itself based on long-range global contexts. By combining local refinement with global refinement stages, we allow the decoder to generate features that improve segmentation results by improving both local and global boundaries that are well aligned at each stage. 







§ EXPERIMENTAL RESULTS


To showcase MobileUNETR's effectiveness as a highly competitive segmentation architecture, we perform multiple experiments across widely popular skin lesion segmentation datasets as well as compare and contrast its performance of the proposed model against high performing segmentation models. 





 §.§ Dataset






To evaluate the performance of our efficient and lightweight model, MobileUNETR, we used four publicly available datasets for segmenting skin lesions. The International Skin Imaging Collaboration (ISIC) has developed and released three widely used datasets ISIC 2016, ISIC 2017 and ISIC 2018 for the task of skin lesion segmentation. Additionally, we evaluated our model performance on the PH2 dataset made available by Dermatology Service of Hospital Pedro Hispano, Portugal. The dataset breakdowns are provided below. 



 §.§ Implementation Details


Our proposed MobileUNETR and accompanying experiments are trained and evaluated using PyTorch on a server equipped with a CPU and RTX 3090 GPU. All models follow a simple training procedue with AdamW <cit.> optimizer with parameters B1 = 0.9 and B2 = 0.999, employing a batch size of 8. The experimental setup incorporates a linear warming-up stage spanning 40 epochs, during which the learning rate gradually increases from 0.0004/40 to 0.0004. Subsequently, a cosine annealing scheduler is employed to decay the learning rate over 400 epochs. Adhering to established practices, we employ straightforward data preparation and augmentation techniques available in PyTorch, ensuring the accessibility and reproducibility of our results.




 §.§ Results on ISIC 2016





The ISIC 2016 dataset represents one of the first standardized skin lesion segmentation task comprising of 900 training images and 300 testing images. Our proposed MobileUNETR is bench marked against nine different architectures, encompassing FCNN, Attention-augmented FCNNs, Generative Adversarial Network (GAN)-based  methods, and Transformer-based methods. Performance results across seven metrics are consolidated in Table 1 with a 2.17% and 1.21% increate in IoU and Dice metrics respectively. 



 §.§ Results on ISIC 2017


The ISIC 2017 improves the scope of skin lesion segmentation by expanding the data corpus size. This dataset comprises 2500 training images with 600 testing images. Our proposed MobileUNETR is systematically bench marked against 12 diverse architectures. We showcase that our model consistently demonstrates improvements in IOU, Dice, and accuracy metrics in all architectures while maintaining a lightweight and efficient design. Results are presented in Table 2 where we boast a 2.47% and 1.84% increase in IoU and Dice metrics respectively.




 §.§ Results on ISIC 2018


The ISIC 2018 dataset stands out as the most comprehensive among commonly utilized skin lesion segmentation datasets. The dataset comprises of 2694 training images together with 1000 testing images. Similar to previous experiments, proposed MobileUNETR is bench marked against a diverse set of 10 architectures, covering a wide range of architectures. Performance outcomes across seven metrics for ISIC 2018 are consolidated in Table 3. Our results consistently reveal enhancements ranging from 2.54%to 1.71% in IOU, and Dice metrics across all architectures, while maintaining a lightweight and efficient design.




 §.§ Results on ISIC PH2




Finally, we present the evaluation of MobileUNETR's performance using the PH2 dataset. Unlike the earlier ISIC datasets, PH2 represents a relatively compact dataset, providing an opportunity to highlight the generalization capabilities of our hybrid architecture in handling smaller datasets. Aligning with our previous experiments, we benchmarked the proposed architecture against nine diverse architectures, and performance results are presented in Table 4. Our results consistently reveal improvements ranging from 2.68% and 1.3% in IOU and Dice metrics, respectively. Successful experiments on PH2 demonstrate the adaptability of our proposed model for applications involving sparse datasets.



 §.§ Comparison to Advanced Training Techniques


As an alternative to designing lightweight deep learning architectures a class of advanced training techniques called Parameter Efficient Fine Tuning (PEFT) <cit.> has been prevalent in recent research. To demonstrates that despite the compact size of the architecture, our model achieves results that rival those of larger architectures employing advanced training techniques we compare our method with recent solutions employing these advanced techniques. Table 5 showcases the effectiveness of well-designed lightweight architectures, proving they can be as effective as large complex models and emphasizing that over-parameterization is not the future of modern deep learning.



§ CONCLUSION


Encoder-decoder architectures provide researchers with a strong architectural paradigm for medical image segmentation. Although it has been used successfully to push the boundaries of medical image segmentation, larger and more complex versions of the encoder decoder paradigm may not be the solution for modern deep learning architectures. This paper introduces MobileUNETR, an innovative and efficient hierarchical hybrid Transformer architecture with tailored for image segmentation. Unlike existing methods, MobileUNETR efficiently integrates local and global information in both the encoder and decoder stages, leveraging the benefits of convolutions and transformers. This integration allows the encoder to extract local and global features during the encoding stage, while allowing the decoder to reconstruct these features, ensuring both local and global alignment in the final segmentation mask. By incorporating local and global features at each level, MobileUNETR avoids the need for large, complex, and over-parameterized models. This not only enhances performance, but also significantly reduces model size and complexity. Extensive experiments were carried out that compared and contrasted our proposed medical image segmentation method with four widely used public datasets (ISIC 2016, ISIC 2017, ISIC 2018, and PH2 dataset). Comparative analyses with state-of-the-art methods demonstrate the effectiveness of our MobileUNETR architecture, showcasing superior accuracy performance and excellent efficiency in model training and inference. Across all datasets MobileUNETR demonstrates a 1.3% to 2.68% increase in Dice and IoU metrics with a 10x and 23x reduction in parameters and computational complexity, compared to current SOTA models. We hope that our method can serve as a strong foundation for medical imaging research, since the application of MobileUNETR in image segmentation is endless. Additionally, we hope that our work here has opened the door to motivate further research in efficient architectures in medical imaging research. 


splncs04

