Lite Swin UNETR: A Lightweight Version of Swin UNETR for Efficient 3D Medical Image Segmentation

초록

In this paper, we propose Lite Swin UNETR, a lightweight and highly efficient model for 3D medical image segmentation that addresses key limitations of existing methods. While Swin UNETR demonstrates strong performance through hierarchical self-attention mechanisms, it struggles to capture fine-grained local context and incurs substantial computational costs due to standard convolutions. To overcome these challenges, we introduce a novel Lite Module that replaces conventional convolutions with an effective three-component architecture. The Lite Module integrates three key innovations: (1) inverted residual blocks with depthwise separable convolutions for improved parameter efficiency, (2) a dual-branch structure with 7×7×7 and 3×3×3 kernels that simultaneously captures broad anatomical context and fine boundary details, and (3) Squeeze-and-Excitation blocks that adaptively emphasize important channel features for enhanced feature representation. This design specifically addresses the unique challenges of 3D medical imaging, where both global shape understanding and local texture analysis are crucial for accurate segmentation. Comprehensive experiments on MSD datasets demonstrate that Lite Swin UNETR achieves superior performance across multiple metrics with high efficiency. Compared to Swin UNETR, our model reduces parameters by 37% (from 4.38M to 2.76M) and computational complexity by 52% (from 53.07G to 25.66G FLOPs), while simultaneously improving average Dice scores from 58.76% to 60.89% on the challenging Prostate dataset. When compared against other lightweight approaches including 2D-adapted architectures, MobileUNet and EfficientUNet, as well as native 3D models such as SegFormer3D, Lite Swin UNETR consistently delivers superior segmentation accuracy, particularly excelling in anatomically complex regions such as the prostate peripheral zone where precise boundary delineation is critical.

키워드