SegFormer3D: an Efficient Transformer for 3D Medical Image Segmentation

TL;DR

SegFormer3D addresses the resource-intensiveness of ViT-based 3D medical image segmentation by introducing a memory-efficient, 4-stage hierarchical Transformer that processes multiscale volumetric features and uses an all-MLP decoder. The model achieves substantial reductions in parameters (~34x) and GFLOPs (~13x) while maintaining competitive accuracy on BraTS, Synapse, and ACDC without pretraining. This approach demonstrates that high performance can be attained with lightweight architectures, improving accessibility for data-limited medical imaging settings and enabling deployment on constrained hardware. The work provides practical design principles—overlapped patch embedding, efficient self-attention, and all-MLP decoding—that can guide future memory-efficient Transformer models in 3D medical imaging.

Abstract

The adoption of Vision Transformers (ViTs) based architectures represents a significant advancement in 3D Medical Image (MI) segmentation, surpassing traditional Convolutional Neural Network (CNN) models by enhancing global contextual understanding. While this paradigm shift has significantly enhanced 3D segmentation performance, state-of-the-art architectures require extremely large and complex architectures with large scale computing resources for training and deployment. Furthermore, in the context of limited datasets, often encountered in medical imaging, larger models can present hurdles in both model generalization and convergence. In response to these challenges and to demonstrate that lightweight models are a valuable area of research in 3D medical imaging, we present SegFormer3D, a hierarchical Transformer that calculates attention across multiscale volumetric features. Additionally, SegFormer3D avoids complex decoders and uses an all-MLP decoder to aggregate local and global attention features to produce highly accurate segmentation masks. The proposed memory efficient Transformer preserves the performance characteristics of a significantly larger model in a compact design. SegFormer3D democratizes deep learning for 3D medical image segmentation by offering a model with 33x less parameters and a 13x reduction in GFLOPS compared to the current state-of-the-art (SOTA). We benchmark SegFormer3D against the current SOTA models on three widely used datasets Synapse, BRaTs, and ACDC, achieving competitive results. Code: https://github.com/OSUPCVLab/SegFormer3D.git

Figures (5)

• Figure 1: Parameter Count vs. Performance on BraTs We compare Segformer3D to existing 3D volumetric image segmentation architectures evaluating model performance with respect to parameter count. The green bars represent model parameters while the purple plot shows the mean dice performance for each architecture. We demonstrate that at 4.5 million parameters Segformer3D is a highly competitive lightweight architecture for 3D medical image segmentation.
• Figure 2: Segformer3D Overview: The model input is a 3D volume $\mathbb{R}^{D \times C \times H \times W}$. We extract multiscale volumetric features using a 4 stage hierarchical Transformer. An all-MLP decoder then upsamples and aggregate local and global attention features from the encoding stage to generate the final segmentation mask.
• Figure :
• Figure :
• Figure :