SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation
Authors: Duy D. Nguyen, Phat T. Tran-Truong
Organizations: Department of Software Engineering, Faculty of Computer Science and Engineering, Ho Chi Minh City University of Technology (HCMUT), 268 Ly Thuong Kiet Street, Dien Hong Ward, Ho Chi Minh City, Vietnam · Vietnam National University Ho Chi Minh City, Linh Trung Ward, Ho Chi Minh City, Vietnam
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.
Figures & tables
Figure 1: Architecture Overview: The model input is a 3D volume RC×D×H×W .A four-stage hierarchical Mamba-Transformer Block is adopted to derive multiscale volumetric representations. These features are then upsampled and fused by an all-MLP decoder, integrating local and global attention cues to generate the final segmentation mask.
Architecture
Params(M)
GFLOPs
nnFormer [ 32 ]
150.5
213.4
TransUNet [ 5 ]
96.07
88.91
UNETR [ 12 ]
92.49
75.76
SwinUNETR [ 11 ]
62.83
384.2
Segformer3D [ 22 ]
4.51
17.5
SegMaFormer (ours)
2.02
15.2
Table 1: Relative comparison of model complexity in terms of parameters and GFLOPs.
Figure 2: 3D Mix Vision Mamba Block with minimal reliance on convolutional layers.
Methods
Params(M)
Avg(%)
Whole Tumor
Enhancing Tumor
Tumor Core
nnFormer [ 32 ]
150.5
86.4
91.3
81.8
86.0
Ours
2.0
83.79
91.0
76.2
84.16
SegFormer3D [ 22 ]
4.5
82.1
89.9
74.2
82.2
UNETR [ 12 ]
92.49
71.1
78.9
58.5
76.1
TransBTS [ 28 ]
–
69.6
77.9
57.4
73.5
CoTr [ 30 ]
41.9
68.3
74.6
55.7
74.8
Table 2: Comparison of segmentation performance across methods on the BraTS benchmark. Our model achieves competitive accuracy while using significantly fewer parameters and small FLOPS.
Methods
Params(M)
Avg(%)
AOR
LIV
LKID
RKID
GAL
PAN
SPL
STO
U-Mamba [ 21 ]
172.63
87.98
90.8
96.90
94.6
94.5
73.80
79.3
95.80
81.70
nnFormer [ 32 ]
150.5
86.57
92.04
96.84
86.57
86.25
70.17
83.35
90.51
86.83
Ours
2.0
83.33
89.98
96.47
90.49
90.53
57.29
70.57
93.03
78.70
SegFormer3D [ 22 ]
4.5
82.15
90.43
95.68
86.53
86.13
55.26
73.06
89.02
81.12
MISSFormer [ 13 ]
–
81.96
86.99
94.41
85.21
82.00
68.65
65.67
91.92
80.81
UNETR [ 12 ]
92.49
79.56
89.99
94.46
85.66
84.80
60.56
59.25
87.81
73.99
Table 3: Within the Synapse benchmark, models are ranked by their average Dice scores across all organ classes. The SegMaFormer performs exceptionally well, outperforming numerous established baselines and ranking just behind nnFormer and U-Mamba, even though these models use more than 75× the parameters and require substantially higher FLOPs.
Methods
Params(M)
Avg(%)
RV
Myo
LV
Primus-S [ 26 ]
23.9
92.46
-
-
-
nnFormer [ 32 ]
150.5
92.06
90.94
89.58
95.65
Ours
2.0
91.11
90.06
89.1
94.14
SegFormer3D [ 22 ]
4.5
90.96
88.5
88.86
95.53
LeViT-UNet-384 [ 31 ]
52.17
90.32
89.55
87.64
93.76
SwinUNet [ 3 ]
–
90.00
88.55
85.62
95.83
Table 4: Comparison on the ACDC benchmark. Models are ranked by their average Dice scores across RV, Myo, and LV. Our method achieves competitive performance with significantly fewer parameters.
National Institute of Technology Kurukshetra Department of Electronics and Communication Engineering · National Institute of Technology Kurukshetra Department of Computer Engineering · Indian Institute of Technology Ropar Department of Computer Science and Engineering