Knowledge Distillation in Lightweight U-Net Transformer Architectures for Brain Tumor Segmentation

Authors

  • Toat Tuloh Universitas Harapan Bangsa
  • Purwono Purwono
  • Iis Setiawan Mangkunegara Universitas Harapan Bangsa

Keywords:

U-Net 3D; knowledge distillation; Dice score; IoU; GFLOPs.

Abstract

Brain tumor segmentation from Magnetic Resonance Imaging was an essential step for therapy planning, prognosis evaluation, and treatment monitoring in glioma patients. Manual delineation required substantial time and was prone to inter-observer variability. Although deep learning models achieved high segmentation accuracy, performance improvements were often accompanied by increased computational complexity, limiting their applicability in resource-constrained clinical environments. To address this issue, an efficiency-oriented segmentation framework was developed based on a lightweight three-dimensional U-Net enhanced with a shallow Transformer module and guided by knowledge distillation. The main contribution of this study was the integration of logit-level and feature-level distillation to improve segmentation capability while maintaining low inference complexity. The framework emphasized a balanced trade-off between segmentation accuracy and computational efficiency rather than benchmark maximization. Experiments were conducted using the BraTS 2020 and BraTS 2021 datasets. The official training sets were internally split into eighty percent for training and twenty percent for validation. The student network was trained using a hybrid segmentation loss combined with temperature-scaled logit distillation and bottleneck feature alignment from a higher-capacity teacher model. Model performance was evaluated using Dice score, Intersection over Union, and computational complexity measured in floating-point operations. On the BraTS 2021 dataset, the proposed model achieved Dice scores of 0.6538 for Whole Tumor, 0.5382 for Tumor Core, and 0.5304 for Enhancing Tumor. Per-class Dice values were 0.9706 for background, 0.3028 for necrotic or non-enhancing tumor core, 0.4938 for edema, and 0.4936 for enhancing tumor. The corresponding Intersection over Union values followed similar trends. The model maintained an inference complexity of approximately 54.38 gigafloating-point operations for input patches of size 64 × 64 × 64. These findings indicated that the proposed framework achieved a stable balance between segmentation performance and computational efficiency, supporting practical deployment under limited computational resources

References

Y. R. Park et al., “Expedited safety reporting through an alert system for clinical trial management at an academic medical center: Retrospective design study,” JMIR Medical Informatics, vol. 8, no. 2, 2020, doi: 10.2196/14379.

X. Feng, H. Bai, D. Kim, G. Maragkos, J. Machaj, and R. Kellogg, “Brain Tumor Segmentation with Patch-Based 3D Attention UNet from Multi-parametric MRI,” in Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries, vol. 12963, A. Crimi and S. Bakas, Eds., in Lecture Notes in Computer Science, vol. 12963. , Cham: Springer International Publishing, 2022, pp. 90–96. doi: 10.1007/978-3-031-09002-8_8.

S. Baig et al., “Segmentation of pre- and posttreatment diffuse glioma tissue subregions including resection cavities,” Neuro-Oncology Advances, vol. 6, no. 1, p. vdae140, Jan. 2024, doi: 10.1093/noajnl/vdae140.

F. Yepes-Calderon and J. Gordon McComb, “Manual Segmentation Errors in Medical Imaging. Proposing a Reliable Gold Standard,” in Applied Informatics, vol. 1051, H. Florez, M. Leon, J. M. Diaz-Nafria, and S. Belli, Eds., in Communications in Computer and Information Science, vol. 1051. , Cham: Springer International Publishing, 2019, pp. 230–241. doi: 10.1007/978-3-030-32475-9_17.

K. A. Wahid et al., “Large scale crowdsourced radiotherapy segmentations across a variety of cancer anatomic sites,” Sci Data, vol. 10, no. 1, p. 161, Mar. 2023, doi: 10.1038/s41597-023-02062-w.

T. Stein et al., “Efficient Web-Based Review for Automatic Segmentation of Volumetric DICOM Images,” in Bildverarbeitung für die Medizin 2019, H. Handels, T. M. Deserno, A. Maier, K. H. Maier-Hein, C. Palm, and T. Tolxdorff, Eds., in Informatik aktuell. , Wiesbaden: Springer Fachmedien Wiesbaden, 2019, pp. 158–163. doi: 10.1007/978-3-658-25326-4_33.

N. Rasool and J. Iqbal Bhat, “Unveiling the Complexity of Medical Imaging through Deep Learning Approaches,” Chaos Theory and Applications, vol. 5, no. 4, pp. 267–280, Dec. 2023, doi: 10.51537/chaos.1326790.

B. Rajesh, M. Satyanarayana, and P. S. Vasarao, “Automated Image Segmentation for Complex Scenes Using U-Net Architecture,” IJRIAS, vol. X, no. VII, pp. 77–83, 2025, doi: 10.51584/IJRIAS.2025.100700006.

A. A. Novikov, D. Major, M. Wimmer, D. Lenis, and K. Buhler, “Deep Sequential Segmentation of Organs in Volumetric Medical Scans,” IEEE Trans. Med. Imaging, vol. 38, no. 5, pp. 1207–1215, May 2019, doi: 10.1109/TMI.2018.2881678.

L. Zhang, M. Li, P. Zhang, and P. Liu, “EfficientSegNet: Lightweight Semantic Segmentation with Multi-Scale Feature Fusion and Boundary Enhancement,” Sensors, vol. 25, no. 19, p. 5934, Sep. 2025, doi: 10.3390/s25195934.

M. Poddar, J. S. Marwaha, W. Yuan, S. Romero-Brufau, and G. A. Brat, “An operational guide to translational clinical machine learning in academic medical centers,” npj Digit. Med., vol. 7, no. 1, p. 129, May 2024, doi: 10.1038/s41746-024-01094-9.

E. Y. Zhu, C. Zhao, H. Yang, J. Li, Y. Wu, and R. Ding, “A Comprehensive Review of Knowledge Distillation- Methods, Applications, and Future Directions,” IJIRCST, vol. 12, no. 3, pp. 106–112, May 2024, doi: 10.55524/ijircst.2024.12.3.17.

X. Qi, M. Hou, and P. Gao, “Dual-View Brain Tumor MRI Segmentation with Embedding Coordinate Attention Mechanism,” in 2023 IEEE 5th International Conference on Power, Intelligent Computing and Systems (ICPICS), Shenyang, China: IEEE, Jul. 2023, pp. 140–146. doi: 10.1109/ICPICS58376.2023.10235355.

J. Wei, J. Chen, Y. Wang, H. Luo, and W. Li, “Improved deep learning image classification algorithm based on Swin Transformer V2,” PeerJ Computer Science, vol. 9, p. e1665, Oct. 2023, doi: 10.7717/peerj-cs.1665.

H. Tabani, A. Balasubramaniam, S. Marzban, E. Arani, and B. Zonooz, “Improving the Efficiency of Transformers for Resource-Constrained Devices,” Jun. 30, 2021, arXiv: arXiv:2106.16006. doi: 10.48550/arXiv.2106.16006.

L. Zhou, J. Zhao, and R. Shi, “Classifier - Centric Knowledge Distillation Based on Class Activation Weight Sharing,” in 2025 6th International Conference on Computer Vision, Image and Deep Learning (CVIDL), Ningbo, China: IEEE, May 2025, pp. 1205–1212. doi: 10.1109/CVIDL65390.2025.11085768.

Z. Liu, C. Zhang, H. Peng, Q. Xu, and Y. Gao, “Drug distribution management system based on IoT,” KSII Transactions on Internet and Information Systems, vol. 16, no. 2, pp. 424–444, 2022, doi: 10.3837/tiis.2022.02.004.

A. Galkin, Yu. Davidich, and H. Samchuk, “Development Of Mathematical Models For Evaluating Demand Parameter Elasticity In Retail Networks Under Consumer-Driven Logistics,” MEC, vol. 4, no. 185, pp. 262–266, Sep. 2024, doi: 10.33042/2522-1809-2024-4-185-262-266.

K. Alrfou and T. Zhao, “GCtx-UNet: Efficient Network for Medical Image Segmentation,” 2024, arXiv. doi: 10.48550/ARXIV.2406.05891.

L. Zhao, X. Qian, Y. Guo, J. Song, J. Hou, and J. Gong, “MSKD: Structured knowledge distillation for efficient medical image segmentation,” Computers in Biology and Medicine, vol. 164, p. 107284, Sep. 2023, doi: 10.1016/j.compbiomed.2023.107284.

M. M. Danesh Pajouh, “Efficient Brain Tumor Segmentation Using a Dual-Decoder 3D U-Net with Attention Gates (DDUNet),” Apr. 14, 2025, Open Science Framework. doi: 10.31219/osf.io/ws5xf_v1.

Anonymous, “Data augmentation.” Zenodo, Jul. 07, 2023. doi: 10.5281/ZENODO.8123405.

J. Panic, A. Defeudis, G. Balestra, V. Giannini, and S. Rosati, “Normalization Strategies in Multi-Center Radiomics Abdominal MRI: Systematic Review and Meta-Analyses,” IEEE Open J. Eng. Med. Biol., vol. 4, pp. 67–76, 2023, doi: 10.1109/OJEMB.2023.3271455.

S. Ali, N. Ali, F. Mohamed, T. Kamal, and M. Salih, “Region segmentation for lung cancer CT image using 3D U- Net model,” Turkish Journal of Internal Medicine, vol. 7, no. 3, pp. 98–108, Jul. 2025, doi: 10.46310/tjim.1580929.

Y. Cao, J. Xu, S. Lin, F. Wei, and H. Hu, “Global Context Networks,” Dec. 24, 2020, arXiv: arXiv:2012.13375. doi: 10.48550/arXiv.2012.13375.

B. Choudhary, D. Singh, and D. S. Sisodia, “Hierarchical Feature and Attention-Based Distillation to Improve Student Model Performance,” in 2025 IEEE International Conference on Computer, Electronics, Electrical Engineering & their Applications (IC2E3), Srinagar Garhwal, India: IEEE, May 2025, pp. 1–6. doi: 10.1109/IC2E365635.2025.11167035.

D. Chen, J.-P. Mei, C. Wang, Y. Feng, and C. Chen, “Online Knowledge Distillation with Diverse Peers,” AAAI, vol. 34, no. 04, pp. 3430–3437, Apr. 2020, doi: 10.1609/aaai.v34i04.5746.

D. Liu, M. Kan, S. Shan, and X. Chen, “Function-Consistent Feature Distillation,” 2023, arXiv. doi: 10.48550/ARXIV.2304.11832.

A. Kheterpal and K. Singh Gill, “The EfficientNetB5 Solution to Multi-Class Bone Marrow Classification Challenges,” in 2024 12th International Conference on Internet of Everything, Microwave, Embedded, Communication and Networks (IEMECON), Jaipur, India: IEEE, Oct. 2024, pp. 1–5. doi: 10.1109/IEMECON62401.2024.10846120.

A. Tragakis et al., “GLFNET: Global-Local (frequency) Filter Networks for efficient medical image segmentation,” 2024, doi: 10.48550/ARXIV.2403.00396.

U. Baid et al., “The RSNA-ASNR-MICCAI BraTS 2021 Benchmark on Brain Tumor Segmentation and Radiogenomic Classification,” Sep. 12, 2021, arXiv: arXiv:2107.02314. doi: 10.48550/arXiv.2107.02314.

A. Burrello, M. Risso, B. A. Motetti, E. Macii, L. Benini, and D. J. Pagliari, “Enhancing Neural Architecture Search With Multiple Hardware Constraints for Deep Learning Model Deployment on Tiny IoT Devices,” IEEE Trans. Emerg. Topics Comput., vol. 12, no. 3, pp. 780–794, Jul. 2024, doi: 10.1109/TETC.2023.3322033.

N. E. Varghese, A. John, U. D. A. C, and M. J. Pillai, “Transformer-augmented lightweight U-Net (UAAC-Net) for accurate MRI brain tumor segmentation,” Neurological Research, vol. 47, no. 12, pp. 1212–1227, Dec. 2025, doi: 10.1080/01616412.2025.2517312.

S. Banerjee and S. Mitra, “Novel Volumetric Sub-region Segmentation in Brain Tumors,” Front. Comput. Neurosci., vol. 14, p. 3, Jan. 2020, doi: 10.3389/fncom.2020.00003.

X. Gu, R. Jin, and M. Li, “Progressively Relaxed Knowledge Distillation,” in 2024 International Joint Conference on Neural Networks (IJCNN), Yokohama, Japan: IEEE, Jun. 2024, pp. 1–8. doi: 10.1109/IJCNN60899.2024.10650544.

S. Li et al., “Distilling a Powerful Student Model via Online Knowledge Distillation,” IEEE Trans. Neural Netw. Learning Syst., vol. 34, no. 11, pp. 8743–8752, Nov. 2023, doi: 10.1109/TNNLS.2022.3152732.

J. Zhu, Z. Tang, P. Ma, Z. Liang, and C. Wang, “DLKUNET : A Lightweight and Efficient Network With Depthwise Large Kernel for Medical Image Segmentation,” Int J Imaging Syst Tech, vol. 35, no. 1, p. e70035, Jan. 2025, doi: 10.1002/ima.70035.

D.-H. Le, “Integrating multiple microRNA functional similarity networks for improved disease–microRNA association prediction,” Biology Methods and Protocols, vol. 10, no. 1, p. bpaf065, Jan. 2025, doi: 10.1093/biomethods/bpaf065.

E. Radiya-Dixit, D. Zhu, and A. H. Beck, “Automated Classification of Benign and Malignant Proliferative Breast Lesions,” Sci Rep, vol. 7, no. 1, p. 9900, Aug. 2017, doi: 10.1038/s41598-017-10324-y.

Prathipati Silpa Chaitanya and Susanta Kumar Satpathy, “Advancing Brain Tumour Detection and Classification: Knowledge Distilled ResNeXt Model for Multi-Class MRI Analysis,” IJCESEN, vol. 10, no. 4, Dec. 2024, doi: 10.22399/ijcesen.730.

Downloads

Published

2026-07-31

How to Cite

Toat Tuloh, Purwono, P., & Iis Setiawan Mangkunegara. (2026). Knowledge Distillation in Lightweight U-Net Transformer Architectures for Brain Tumor Segmentation. Journal of Advanced Health Informatics Research, 4(1), 24–39. Retrieved from https://ejournal.ptti.web.id/index.php/jahir/article/view/369

Issue

Section

Articles