Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation

Duy D. Nguyen, Phat T. Tran-Truong

Abstract

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.

SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation

Abstract

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.
Paper Structure (10 sections, 3 equations, 2 figures, 4 tables)

This paper contains 10 sections, 3 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Embedding:
  4. Encoder:
  5. Decoder:
  6. Experimental Result
  7. Brain Tumor Segmentation (BraTs)
  8. Multi-Organ CT Segmentation (Synapse)
  9. Automated Cardiac Diagnosis (ACDC)
  10. Limitation and Conclusion

Figures (2)

  • Figure 1: Architecture Overview: The model input is a 3D volume $\mathbb{R}^{C \times D \times H \times 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.
  • Figure 2: 3D Mix Vision Mamba Block with minimal reliance on convolutional layers.