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

TL;DR

SegMaFormer is a lightweight hybrid model combining Mamba and Transformer modules, achieving high segmentation accuracy with significantly fewer parameters and computational resources on 3D medical imaging benchmarks.

Contribution

This work introduces SegMaFormer, a novel hybrid architecture that efficiently combines Mamba and Transformer modules for 3D medical image segmentation, reducing complexity while maintaining performance.

Findings

Reduces parameters by up to 75x compared to state-of-the-art models.

Achieves competitive Dice scores on multiple benchmarks.

Substantially decreases FLOPs while maintaining accuracy.

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.