UltraLightUNet: Rethinking U-shaped Network with Multi-kernel Lightweight Convolutions for Medical Image Segmentation

Authors: Md Mostafijur Rahman, Radu Marculescu
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Ultra Lightweight CNN, Medical Imaging, Semantic Segmentation, 3D Segmentation

Abstract:
In this paper, we introduce UltraLightUNet (2D and 3D), an ultra-lightweight, multi-kernel U-shaped network for medical image segmentation. The core of UltraLightUNet consists of a new Multi-kernel Inverted Residual (MKIR) block, which can efficiently process images through multiple kernels while capturing complex spatial relationships. Additionally, our Multi-kernel Inverted Residual Attention (MKIRA) block refines and emphasizes image salient features via sophisticated convolutional multi-focal attention mechanisms. UltraLightUNet strategically employs the MKIR block in the encoder for feature extraction and the MKIRA block in the decoder for feature refinement, thus ensuring targeted feature enhancement at each stage. With only 0.316M \#Params and 0.314G #FLOPs, UltraLightUNet offers an ultra-lightweight yet powerful segmentation solution that outperforms state-of-the-art (SOTA) methods across twelve medical imaging benchmarks. Notably, UltraLightUNet surpasses TransUNet on DICE score while using 333$\times$ fewer \#Params and 123$\times$ fewer #FLOPs. Compared to the lightweight model, UNeXt, UltraLightUNet improves DICE scores by up to 6.7% with 4.7$\times$ fewer parameters. UltraLightUNet also outperforms recent lightweight models such as MedT, CMUNeXt, EGE-UNet, Rolling-UNet, and UltraLight_VM_UNet, while using significantly fewer #Params and #FLOPs. Furthermore, our 3D version, UltraLightUNet3D-M (1.42M #Params and 7.1G #FLOPs), outperforms SwinUNETR (62.19M #Params, 328.6G #FLOPs) and nn-UNet (31.2M #Params, 110.4G #FLOPs) on the FETA, MSD Brain Tumor, Prostate, and Lung Cancer segmentation benchmarks. This remarkable performance, combined with substantial computational gains, makes UltraLightUNet an ideal solution for real-time and point-of-care services in resource-constrained environments. We will make the code publicly available upon paper acceptance.

Decision: reject
Reviewer scores: 1, 1, 6, 3

Reviews and discussion:
<thread>
creation date:
1729151135005

thread_type:
official_review

<thread_object>
creator: reviewer_swmn

<content>
summary:
This manuscript proposes a novel U-shaped network incorporating various modules for medical image segmentation, with a focus on reducing computational costs. Key contributions include the Multi-Kernel Inverted Residual Block, Multi-Kernel Inverted Residual Attention, and Grouped Attention Gate. As a result, the proposed model achieves remarkable computational efficiency and delivers superior segmentation quality compared to state-of-the-art models. The manuscript is well-written and well-organized; however, more detailed technical insights into each module would enhance clarity and depth.

strengths:
This manuscript presents various experiments validating the model's performance. The multi-kernel structures employed effectively capture multi-scale contexts, enhancing segmentation accuracy without adding complexity. The reported results demonstrate that the proposed model’s segmentation performance surpasses that of significantly heavier architectures.

weaknesses:
Although well-organized, the manuscript could benefit from a deeper focus on explaining how each module reduces computational costs while maintaining high performance. The method section lacks clear evidence or mathematical proof to support the model’s design, which may present a scientific limitation. Additionally, there are no ablation studies evaluating the individual contributions of each module in terms of both segmentation performance and parameter efficiency.

questions:
1. Can you clarify the effects of the proposed modules in the model with clear evidence or mathematical insight? The current equations in the manuscript seem somewhat redundant, serving primarily as mathematical restatements. It would be helpful to provide more rigorous insights or proofs demonstrating how each module contributes to performance improvement and computational efficiency.

2. Which proposed module is the most critical? Please provide ablation studies to highlight the individual impact of each module on both segmentation performance and computational cost. This would help identify the key components driving the model's success.

3. General multi-scale approaches are not novel. For instance, the "Spatial Feature Conservation Networks (SFCNs) for Dilated Convolutions to Improve Breast Cancer Segmentation from DCE-MRI" (International Workshop on Applications of Medical AI, 2022) employs a multi-scale strategy. What distinguishes your model’s approach to multi-scale feature extraction? Please elaborate on the unique contributions and insights of your method in this context.

4. What are the differences among UltraLightUNet- T, S, and L in terms of model? Layer difference? Please explain in details in the manuscript.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729959651987

thread_type:
official_review

<thread_object>
creator: reviewer_jk4q

<content>
summary:
The paper introduces UltraLightUNet, a novel ultra-lightweight, multi-kernel U-shaped network designed to improve medical image segmentation. Leveraging a new Multi-kernel Inverted Residual (MKIR) block for efficient multi-scale feature extraction and a Multi-kernel Inverted Residual Attention (MKIRA) block for refined feature enhancement, UltraLightUNet achieves high segmentation accuracy with minimal parameters and computational load. With only few parameters and FLOPs, the 2D version of UltraLightUNet outperforms existing lightweight and transformer-based segmentation models across multiple medical imaging benchmarks, while the 3D variant, UltraLightUNet3D, achieves superior results on complex 3D medical segmentation tasks with even greater efficiency. These performance gains make UltraLightUNet a viable option for real-time applications in resource-constrained environments, such as point-of-care diagnostics.

strengths:
The paper's primary strength lies in its originality, presenting a lightweight model architecture, UltraLightUNet, that effectively combines multi-kernel convolutions and attention mechanisms, a novel approach for achieving high segmentation accuracy with low computational overhead. This innovation addresses a critical need for practical, high-performance segmentation models in resource-constrained environments, adding value to the field by bridging computational efficiency and segmentation quality. In terms of quality, the paper’s methodology appears well-supported by rigorous experimental validation across 11 datasets and comparison with SOTA models, demonstrating robust performance gains and highlighting the practical implications of the architecture's low parameter count and FLOPs. Clarity is another strength; the paper methodically explains the architecture components, from the MKIR and MKIRA blocks to grouped attention mechanisms, making it accessible to readers with a background in medical image segmentation. Finally, the significance of UltraLightUNet is considerable due to its broad applicability across a range of medical imaging tasks and its potential for real-time use in settings like point-of-care diagnostics. The model’s lightweight design, paired with high accuracy, addresses critical bottlenecks in deploying AI-driven diagnostics in clinical environments, establishing the paper as a meaningful contribution to the field. I am glad to see the author will open-source the code to promote research in this line.

weaknesses:
First, while the model is evaluated across various medical imaging datasets, these datasets are relatively straightforward, covering simple segmentation tasks and organs rather than more complex applications like CT/MRI tumor and lesion segmentation. 

Second, the masks are mostly binary segmentation tasks. The multi-class segmentation is not well explored. Adding more complicated and multi-class segmentation would better demonstrate the model's capability for broader, real-world medical imaging tasks. 

Third, while the proposed blocks—MKIRA, MKIR, MKDC, GAG, and CMFA—are illustrated in the Method section, the paper lacks sufficient theoretical or conceptual motivation for why these specific block designs should enhance segmentation performance. 

Lastly, the paper does not discuss the limitations of UltraLightUNet. A dedicated discussion on limitations would provide readers with a more balanced understanding of the model’s practical use and potential future directions for research.

questions:
Typically, there is no free lunch. Are there any limitations or tradeoffs of the method? Providing those will be helpful for readers.

How it would work on more complicated problems (e.g., 3D tumors/lesions) and multi-class segmentation problems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730086616344

thread_type:
official_review

<thread_object>
creator: reviewer_zjlj

<content>
summary:
The author proposed a 2D and 3D ultra-lightweight, multi-kernel U-shaped network for medical image segmentation, termed a UltraLightUNet. It consists of an Multi-kernel Inverted Residual (MKIR) block and an Multi-kernel Inverted Residual Attention (MKIRA) block.  MKIR was proposed to efficiently process images through multiple kernels while capturing complex spatial relationships, and MKIRA block refines and emphasizes image salient features via new sophisticated convolutional multi-focal attention mechanisms. This UltraLightUNet outperformed other methods with lower complexity.

strengths:
1. The proposed network has a low number of parameters and low computational complexity than other widely used baselines and achieved promising segmentation accuracy.
2. Methods and results are thoughtfully described.

weaknesses:
1. The overall novelty is low. The author mainly proposed two modules, CMFA and MKDC modules. MKDC just applied several depth-wise convolutional layers to extract features from different channels. This idea has already proposed in the Scale-Aware Modulation Meet Transformer [1]. MKDC module employed average and max pooling, and convolution layers, which are the channel-wise and spatial-wise attention. Many various attention-based modules have been proposed between 2018 and 2020 (like the CBAM [2]). The overall network is similar with [3] [4].
2. The experimental results are limited. First, authors reported FLOPs and Params to demonstrate that the network has a lower computational complexity than other baselines. It is achieved by mainly replacing convolutional layers with depth-wise convolutional layers. However, it will take longer time to train networks which employ depth-wise convolutional layers compared with those with standard convolution layers. Thus, training and test time are needed to be reported. Second, the comparison in Synapse, MSD prostate, and FETA is insufficient. Synapse is a popular benchmark, but only a few baseline methods were reported. Additionally, the performance reported for these baseline methods in this paper is much lower than the performance in the original paper. For example, Swin Unet reported 79.13 in their paper [5], but only 77.58 was reported for it in this manuscript. If authors run experiments for baselines on their own, please make sure the baseline networks have been fully optimized. Only seven 3D methods proposed before 2022 were compared in MSD prostate and FETA. However, 3D segmentation networks between 2023 and 2024 were not compared, and these networks usually achieve more superior performance with lower computational complexity.
3. The motivation is unclear. The author mentioned that several 3D segmentation networks, including 3D U-Net, SwinUNETR, 3D UX-Net, UNETR, nnU-net and nnFormer, have high computational demands, so they proposed their lightweight network. However, these baseline networks are proposed in 2021 and 2022, and in recent years many lightweight 3D segmentation networks have been proposed and this challenge has been tackled, such as [6][7][8][9][10]. However, authors didn't discuss and explore these lightweight networks. Additionally, modules in this manuscript were proposed based on the idea of ConvNeXt, but its lightweight version was not discussed. Some other depth-wise convolution-based lightweight networks were also not discussed [11].
4. The theoretical development is not solid. Authors reviewed Vision Transformers in the related work, but it is not related to the work in this manuscript. The network in the manuscript employs a convolutional neural network architecture and attention mechanisms. Additionally, insufficient theoretical development in the Method section, and it is unclear how this design improved the segmentation performance. 
5. No model interpretability. The interpretability of the CMFA and MKDC modules were not discussed, such as saliency maps. It is important to understand the mechanisms of attention-based modules.
6. The overall impact is low. The overall improvement in the segmentation performance is low. For example, its best DSC score in the Polyp dataset was 93.48, but other baselines achieved 93.29 and 93.18. Its best DSC score in the Synapse dataset was 78.68, but other baselines achieved 78.40. 

Minors: (1) lacking qualitative results for 3D segmentation results in Synapse, MSD Prostate, and FETA.
(2) Only overall DSC scores were reported for multi-class segmentation tasks, but organ-specific DSC scores were not reported.
(3) lack p-values and standard deviations

[1] Lin, W., Wu, Z., Chen, J., Huang, J., & Jin, L. (2023). Scale-aware modulation meet transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 6015-6026).
[2] Woo, S., Park, J., Lee, J. Y., & Kweon, I. S. (2018). Cbam: Convolutional block attention module. In Proceedings of the European conference on computer vision (ECCV) (pp. 3-19).
[3] Rahman, M. M., & Marculescu, R. (2023). Medical image segmentation via cascaded attention decoding. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 6222-6231).
[4] Rahman, M. M., Munir, M., & Marculescu, R. (2024). Emcad: Efficient multi-scale convolutional attention decoding for medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 11769-11779).
[5] Cao, H., Wang, Y., Chen, J., Jiang, D., Zhang, X., Tian, Q., & Wang, M. (2022, October). Swin-unet: Unet-like pure transformer for medical image segmentation. In European conference on computer vision (pp. 205-218). Cham: Springer Nature Switzerland.
[6] Pang, Y., Liang, J., Huang, T., Chen, H., Li, Y., Li, D., ... & Wang, Q. (2023). Slim UNETR: Scale hybrid transformers to efficient 3D medical image segmentation under limited computational resources. IEEE Transactions on Medical Imaging
[7] Tang, F., Ding, J., Quan, Q., Wang, L., Ning, C., & Zhou, S. K. (2024, May). Cmunext: An efficient medical image segmentation network based on large kernel and skip fusion. In 2024 IEEE International Symposium on Biomedical Imaging (ISBI) (pp. 1-5). IEEE.
[8] Yang, S., Zhang, X., Chen, Y., Jiang, Y., Feng, Q., Pu, L., & Sun, F. (2023). UcUNet: A lightweight and precise medical image segmentation network based on efficient large kernel U-shaped convolutional module design. Knowledge-Based Systems, 278, 110868.
[9] He, Y., Gao, Z., Li, Y., & Wang, Z. (2024). A lightweight multi-modality medical image semantic segmentation network base on the novel UNeXt and Wave-MLP. Computerized Medical Imaging and Graphics, 111, 102311.
[10] Lin, X., Yu, L., Cheng, K. T., & Yan, Z. (2023). BATFormer: Towards boundary-aware lightweight transformer for efficient medical image segmentation. IEEE Journal of Biomedical and Health Informatics, 27(7), 3501-3512.
[11] Yin, Y., Han, Z., Jian, M., Wang, G. G., Chen, L., & Wang, R. (2023). AMSUnet: A neural network using atrous multi-scale convolution for medical image segmentation. Computers in Biology and Medicine, 162, 107120.

questions:
1. Provide more detailed and necessary experimental details, including recent works in lightweight network design and other more advanced baselines, training and test time in experiments.
2. Demonstrate more solid theoretical development
3. Discuss model interpretability
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721938942

thread_type:
official_review

<thread_object>
creator: reviewer_jekk

<content>
summary:
The authors introduce UltraLightUNet, an ultra-lightweight 2D and 3D U-shaped network designed for medical image segmentation. This network features novel Multi-kernel Inverted Residual (MKIR) and Multi-kernel Inverted Residual Attention (MKIRA) blocks, aiming to effectively balance computational efficiency with high segmentation performance across multiple medical imaging benchmarks. The architecture is motivated by the need to reduce computational demands in point-of-care diagnostics, particularly in resource-constrained environments.

strengths:
- The paper is well-written and easy to follow.

- The experimental results are through, with diverse datasets and settings.

- In-depth ablation study and parameters consideration are presented

weaknesses:
- Despite the paper presents diverse experiments and ablation studies, I see a very close similarity to a CVPR 2024 paper, named EMCAD [1], which I detail in the next parts (Note that the EMCAD paper is not cited!)

- The proposed method (Figure 2) is clearly the same as last year CVPR paper, with approximately no change in any of the modules both in encoder and decoder of the network. Therefore, I see no novelty and contribution in the paper submitted.

- How is it possible for your method to just have 27000 parameters, while last year paper with the same architecture has at-least 3M parameters (its base version).

- The proposed method is not compared to other SOTA networks (such as [2]) which perform much better over some of the datasets used in the paper (such as ISIC)

[1] Rahman MM, Munir M, Marculescu R. Emcad: Efficient multi-scale convolutional attention decoding for medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition 2024 (pp. 11769-11779).


[2]Azad R, Niggemeier L, Hüttemann M, Kazerouni A, Aghdam EK, Velichko Y, Bagci U, Merhof D. Beyond self-attention: Deformable large kernel attention for medical image segmentation. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision 2024 (pp. 1287-1297).

questions:
Please refer to weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732630308712

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Global response by Authors: Part1

comment:
We sincerely thank all reviewers for their constructive feedback, which has been invaluable in improving our manuscript. We are encouraged by the recognition of **UltraLightUNet**’s contributions to lightweight, high-performance segmentation and its potential for practical applications in resource-constrained scenarios. 

We note that our contribution is along the lines of a growing direction in computer vision in general, and biomedical imaging in particular, of using ultralightweight models to perform real-time tasks with high precision. Remarkable examples of this trend are MobileNet and EfficientNet in computer vision, TinyBERT in language processing, UNeXt and EGE-UNet in biomedical imaging. We beleive that the significance of this type of work will only increase in future years.

Below, we provide a summary of the major revisions made in response to the reviewers’ comments and highlight our specific changes. We marked these changes in **blue** in the revised version of our draft.

## Key Modifications and Contributions

### 1. New Experiments on Complex Datasets

To address the reviewers’ concerns regarding the evaluation of **UltraLightUNet** on more complex segmentation tasks (**Table R5**), we conducted additional experiments on two challenging datasets:

- **MSD Task01_BrainTumour** (multi-level tumor segmentation):
  - Results show that **UltraLightUNet3D-M** achieves the **best DICE scores** on Tumor Core (TC) and Whole Tumor (WT) segmentation, while having remarkably lower #Params and #FLOPs compared to heavyweight models (**SwinUNETR**, **3D UX-Net**).
  - The base model, **UltraLightUNet3D**, also performs competitively, outperforming lightweight models like **SlimUNETR** with the lowest computational cost.

- **MSD Task06_Lung** (binary CT lesion segmentation):
  - Results demonstrate **UltraLightUNet**’s ability to handle complex binary cancer segmentation tasks efficiently, with **UltraLightUNet3D-M** achieving the **best DICE score** compared to other baselines, including lightweight (e.g., **SlimUNETR**) and heavyweight (**3D UX-Net**, **SwinUNETR**).

### 2. Theoretical Motivation, Methodological Novelty, and Interpretability

- We elaborated on the theoretical underpinnings of **UltraLightUNet**’s modules:
  - Our **multi-kernel trick** supports both **$k_1 = k_2$** (same-size kernels) and **$k_1 ≠ k_2$** (different-size kernels) for **$k_1$, $k_2$** $\in Kernels$ versus conventional multi-scale (only **$k_1 ≠ k_2$**) designs, thus allowing adaptable context extraction.
  - This conceptual distinction allows **UltraLightUNet** to adapt kernel sizes based on application-specific needs (e.g., large kernels for large objects, small kernels for small objects, or mixed for both objects segmentation).
  - **Depth-wise convolutions** across multiple kernels (**$k_1 = k_2$** or **$k_1 ≠ k_2$**) ensure lightweight computation without compromising accuracy.
- We clarified about our **methodological novelty**: 
    - End-to-end ultralightweight architecture design 
    - Novel 2D modules, 
    - New 3D modules, and
    - Both 2D and 3D versatility.
- **Activation heatmaps (Figure R1)** were added to demonstrate how MKIR and CMFA focus on critical regions in an image, thus providing interpretability evidence for the modules.

### 3. Clarifications on Related Approaches

We clarified the distinctions between **UltraLightUNet** and related methods:

- **EMCAD**: Table R1 explicitly highlights the differences, including **UltraLightUNet**’s conceptual distinction (**$k_1 = k_2$** or **$k_1 ≠ k_2$**), 3D versatility, and unified end-to-end lightweight design.
- **SAMT, CASCADE, ConvNeXt, and SFCNs**: We emphasized the advantages of our multi-kernel approach over multi-scale or single-scale strategies used in these methods.

### 4. Additional Comparisons and New Baseline Results

- **2D Segmentation**: We implemented **DeformableLKA** (*Azad et al., 2024*), a recent segmentation method, and compared its performance with **UltraLightUNet** in **Table R2**. Results confirm **UltraLightUNet**’s superior efficiency and segmentation accuracy.
- **3D Segmentation**: **SlimUNETR** (*Pang et al., 2023*), a lightweight 3D method, was implemented and compared in **Table R4**. **UltraLightUNet** outperforms **SlimUNETR** in both segmentation accuracy and computational efficiency.

### 5. Efficiency Metrics

- Reported training and inference times, parameter counts, FLOPs, and average Dice scores across six binary segmentation datasets (**Table R3**). **UltraLightUNet** demonstrated a compelling trade-off between accuracy and computational cost.

- **Table R3** will be reported in **Table 14 and Appendix A.13** in our revised manuscript.
</content>
</thread_object>

</thread>


Metareview: