SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation
This work addresses efficiency issues in 3D medical image segmentation for researchers and practitioners, offering an incremental improvement over existing Transformer and Mamba-based models.
The paper tackled the high computational complexity and parameter counts in 3D medical image segmentation models by introducing SegMaFormer, a lightweight hybrid architecture that reduces parameters by up to 75x and substantially decreases FLOPs while achieving competitive performance on benchmarks like Synapse, BraTS, and ACDC.
The advent of Transformer and Mamba-based architectures has significantly advanced 3D medical image segmentation by enabling global contextual modeling, a capability traditionally limited in Convolutional Neural Networks (CNNs). However, state-of-the-art Transformer models often entail substantial computational complexity and parameter counts, which is particularly prohibitive for volumetric data and further exacerbated by the limited availability of annotated medical imaging datasets. To address these limitations, this work introduces SegMaFormer, a lightweight hybrid architecture that synergizes Mamba and Transformer modules within a hierarchical volumetric encoder for efficient long-range dependency modeling. The model strategically employs Mamba-based layers in early, high-resolution stages to reduce computational overhead while capturing essential spatial context, and reserves self-attention mechanisms for later, lower-resolution stages to refine feature representation. This design is augmented with generalized rotary position embeddings to enhance spatial awareness. Despite its compact structure, SegMaFormer achieves competitive performance on three public benchmarks (Synapse, BraTS, and ACDC), matching the Dice coefficient of significantly larger models. Empirically, our approach reduces parameters by up to 75x and substantially decreases FLOPs compared to current state-of-the-art models, establishing an efficient and high-performing solution for 3D medical image segmentation.