UNETR++ with Voxel-Focused Attention: Efficient 3D Medical Image Segmentation with Linear-Complexity Transformers
Source abstract
Abstract
There have been significant breakthroughs in developing models for segmenting 3D medical images, with many promising results attributed to the incorporation of Vision Transformers (ViT). However, the fundamental mechanism of transformers, known as self-attention, has quadratic complexity, which significantly increases computational requirements, especially in the case of 3D medical images. In this paper, we investigate the UNETR++ model and propose a voxel-focused attention mechanism inspired by TransNeXt pixel-focused attention. The core component of UNETR++ is the Efficient Paired Attention (EPA) block, which learns from two interdependent branches: spatial and channel attention. For spatial attention, we incorporated the voxel-focused attention mechanism, which has linear complexity with respect to input sequence length, rather than projecting the keys and values into lower dimensions. The deficiency of UNETR++ lies in its reliance on dimensionality reduction for spatial attention, which reduces efficiency but risks information loss. Our contribution is to replace this with a voxel-focused attention design that achieves linear complexity without low-dimensional projection, thereby reducing parameters while preserving representational power. This effectively reduces the model’s parameter count while maintaining competitive performance and inference speed. On the Synapse dataset, the enhanced UNETR++ model contains 21.42 M parameters, a 50% reduction from the original 42.96 M, while achieving a competitive Dice score of 86.72%.
Source: crossref · Quality: baseline
Research context
OpenAlex topics: Medical Image Segmentation Techniques · Advanced Neural Network Applications · Advanced Image and Video Retrieval Techniques
SciVX discovery tags are terminology-based. OpenAlex topics are source-provided classifications.
Affiliations
Computer Science Discipline, University of KwaZulu-Natal, Durban 4000, South Africa
References (43)
Show source-provided references
1. National Health Service (2022). Magnetic Resonance Imaging (MRI) Scan.
2. National Health Service (2023). CT Scan.
3. Padhi, S., Rup, S., Saxena, S., and Mohanty, F. (2019, January 28–29). Mammogram Segmentation Methods: A Brief Review. Proceedings of the 2019 2nd International Conference on Intelligent Communication and Computational Techniques (ICCT), Jaipur, India.10.1109/ICCT46177.2019.8968781 ↗
4. Computer-aided diagnosis in medical imaging: Historical review, current status and future potential10.1016/j.compmedimag.2007.02.002 ↗
5. Chen, Y., Zhuang, Z., and Chen, C. (2023, January 26–28). Object Detection Method Based on PVTv2. Proceedings of the 2023 IEEE 3rd International Conference on Electronic Technology, Communication and Information (ICETCI), Changchun, China.10.1109/ICETCI57876.2023.10176423 ↗
6. PVT v2: Improved baselines with Pyramid Vision Transformer10.1007/s41095-022-0274-8 ↗
7. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. arXiv.10.1109/ICCV48922.2021.00986 ↗
8. Focal Modulation Networks
9. Dong, X., Bao, J., Chen, D., Zhang, W., Yu, N., Yuan, L., Chen, D., and Guo, B. (2022). CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows. arXiv.10.1109/CVPR52688.2022.01181 ↗
10. Wang, W., Dai, J., Chen, Z., Huang, Z., Li, Z., Zhu, X., Hu, X., Lu, T., Lu, L., and Li, H. (2023). InternImage: Exploring Large-Scale Vision Foundation Models with Deformable Convolutions. arXiv.10.1109/CVPR52729.2023.01385 ↗
11. Shi, D. (2024). TransNeXt: Robust Foveal Visual Perception for Vision Transformers. arXiv.10.1109/CVPR52733.2024.01683 ↗
12. UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation10.1109/TMI.2024.3398728 ↗
13. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., and Polosukhin, I. (2023). Attention Is All You Need. arXiv.
14. Deep Learning Techniques for Automatic MRI Cardiac Multi-Structures Segmentation and Diagnosis: Is the Problem Solved?10.1109/TMI.2018.2837502 ↗
15. The Multimodal Brain Tumor Image Segmentation Benchmark (BRATS)10.1109/TMI.2014.2377694 ↗
16. Landman, B., Xu, Z., Igelsias, J., Styner, M., Langerak, T., and Klein, A. (2015, January 5). Miccai multi-atlas labeling beyond the cranial vault–workshop and challenge. Proceedings of the MICCAI Multi-Atlas Labeling Beyond Cranial Vault—Workshop Challenge, Munich, Germany.
17. Navab, N., Hornegger, J., Wells, W.M., and Frangi, A.F. (2015, January 5–9). U-Net: Convolutional Networks for Biomedical Image Segmentation. Proceedings of the Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015, Munich, Germany.10.1007/978-3-319-24553-9 ↗
18. He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep Residual Learning for Image Recognition. arXiv.10.1109/CVPR.2016.90 ↗
19. Huang, H., Lin, L., Tong, R., Hu, H., Zhang, Q., Iwamoto, Y., Han, X., Chen, Y.W., and Wu, J. (2020). UNet 3+: A Full-Scale Connected UNet for Medical Image Segmentation. arXiv.10.1109/ICASSP40776.2020.9053405 ↗
20. Ourselin, S., Joskowicz, L., Sabuncu, M.R., Unal, G., and Wells, W. (2016, January 17–21). 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation. Proceedings of the Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016, Athens, Greece.10.1007/978-3-319-46720-7 ↗
21. Zhou, Z., Siddiquee, M.M.R., Tajbakhsh, N., and Liang, J. (2018). UNet++: A Nested U-Net Architecture for Medical Image Segmentation. arXiv.10.1007/978-3-030-00889-5_1 ↗
22. Angermann, C., and Haltmeier, M. (2019, January 13). Random 2.5 d u-net for fully 3d segmentation. Proceedings of the Machine Learning and Medical Engineering for Cardiovascular Health and Intravascular Imaging and Computer Assisted Stenting: First International Workshop, MLMECH 2019, and 8th Joint International Workshop, CVII-STENT 2019, Held in Conjunction with MICCAI 2019, Shenzhen, China. Proceedings 1.10.1007/978-3-030-33327-0_19 ↗
23. Isensee, F., Petersen, J., Klein, A., Zimmerer, D., Jaeger, P.F., Kohl, S., Wasserthal, J., Koehler, G., Norajitra, T., and Wirkert, S. (2018). nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation. arXiv.10.1007/978-3-658-25326-4_7 ↗
24. Liu, Z., Mao, H., Wu, C.Y., Feichtenhofer, C., Darrell, T., and Xie, S. (2022). A ConvNet for the 2020s. arXiv.10.1109/CVPR52688.2022.01167 ↗
25. Roy, S., Koehler, G., Ulrich, C., Baumgartner, M., Petersen, J., Isensee, F., Jaeger, P.F., and Maier-Hein, K.H. (2023, January 23–27). Mednext: Transformer-driven scaling of convnets for medical image segmentation. Proceedings of the International Conference on Medical Image Computing and Computer-Assisted Intervention, Daejeon, Republic of Korea.10.1007/978-3-031-43901-8_39 ↗
26. Yang, G., Aviles-Rivero, A., Roberts, M., and Schönlieb, C.B. LKAU-Net: 3D Large-Kernel Attention-Based U-Net for Automatic MRI Brain Tumor Segmentation. Proceedings of the Medical Image Understanding and Analysis.
27. Beyond self-attention: External attention using two linear layers for visual tasks
28. Huang, T., Huang, L., You, S., Wang, F., Qian, C., and Xu, C. (2022). Lightvit: Towards light-weight convolution-free vision transformers. arXiv.
29. de Bruijne, M., Cattin, P.C., Cotin, S., Padoy, N., Speidel, S., Zheng, Y., and Essert, C. (October, January 27). Convolution-Free Medical Image Segmentation Using Transformers. Proceedings of the Medical Image Computing and Computer Assisted Intervention–MICCAI 2021, Strasbourg, France.10.1007/978-3-030-87193-2 ↗
30. Chen, J., Lu, Y., Yu, Q., Luo, X., Adeli, E., Wang, Y., Lu, L., Yuille, A.L., and Zhou, Y. (2021). TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation. arXiv.
31. Hatamizadeh, A., Tang, Y., Nath, V., Yang, D., Myronenko, A., Landman, B., Roth, H.R., and Xu, D. (2022, January 3–8). UNETR: Transformers for 3D Medical Image Segmentation. Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), Waikoloa, HI, USA.10.1109/WACV51458.2022.00181 ↗
32. Zhou, H.Y., Guo, J., Zhang, Y., Yu, L., Wang, L., and Yu, Y. (2021). nnformer: Interleaved transformer for volumetric segmentation. arXiv.
33. Cao, H., Wang, Y., Chen, J., Jiang, D., Zhang, X., Tian, Q., and Wang, M. Swin-unet: Unet-like pure transformer for medical image segmentation. Proceedings of the European Conference on Computer Vision.
34. SCANeXt: Enhancing 3D medical image segmentation with dual attention network and depth-wise convolution10.1016/j.heliyon.2024.e26775 ↗
35. Huang, X., Deng, Z., Li, D., and Yuan, X. (2021). Missformer: An effective medical image segmentation transformer. arXiv.
36. Hatamizadeh, A., Nath, V., Tang, Y., Yang, D., Roth, H.R., and Xu, D. (2021). Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images. Proceedings of the International MICCAI Brainlesion Workshop, Springer.10.1007/978-3-031-08999-2_22 ↗
37. Rahman, M.M., and Marculescu, R. (2023, January 2–7). Medical Image Segmentation via Cascaded Attention Decoding. Proceedings of the 2023 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), Waikoloa, HI, USA.10.1109/WACV56688.2023.00616 ↗
38. Niknejad, M., Firouzbakht, M., and Amirmazlaghani, M. (2024, January 14–16). Enhancing Precision in Dermoscopic Imaging using TransUNet and CASCADE. Proceedings of the 2024 32nd International Conference on Electrical Engineering (ICEE), Tehran, Iran.10.1109/ICEE63041.2024.10667782 ↗
39. Azad, R., Niggemeier, L., Huttemann, M., Kazerouni, A., Aghdam, E.K., Velichko, Y., Bagci, U., and Merhof, D. (2023). Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation. arXiv.10.1109/WACV57701.2024.00132 ↗
40. Rahman, M.M., and Marculescu, R. (2023). Multi-scale Hierarchical Vision Transformer with Cascaded Attention Decoding for Medical Image Segmentation. arXiv.10.1109/WACV56688.2023.00616 ↗
41. Rahman, M.M., Munir, M., and Marculescu, R. (2024). EMCAD: Efficient Multi-scale Convolutional Attention Decoding for Medical Image Segmentation. arXiv.10.1109/CVPR52733.2024.01118 ↗
42. Perera, S., Navard, P., and Yilmaz, A. (2024). SegFormer3D: An Efficient Transformer for 3D Medical Image Segmentation. arXiv.10.1109/CVPRW63382.2024.00503 ↗
43. Tsai, T.Y., Yu, A., Maadugundu, M.S., Mohima, I.J., Barsha, U.H., Chen, M.H.F., Prabhakaran, B., and Chang, M.C. (2025). IntelliCardiac: An Intelligent Platform for Cardiac Image Segmentation and Classification. arXiv.