GMD-Mamba: Lightweight 3D Medical Image Segmentation Using a Grouped Multi-Directional Mamba

Rahim Ullah, Hasan Abbas, Hazrat Ali, Syed Waqar Nabi

Mamba has the capability to capture global context in linear time for 3D medical image segmentation. However, existing methods employ Mamba with full-width directional scans at every resolution, keeping the model large (64-138M parameters and high FLOPs). We propose Grouped Multi-Directional (GMD) Mamba, a lightweight network based on the GMD scan that divides input channels $C$ into groups $G$ and flattens each group in different directions in one pass. This gives multi-directional context at $O(C^2/G)$ instead of $O(K \cdot C^2)$, where $K$ is the number of scan directions. Additionally, we employ Mamba only at the deepest encoder stages, resulting in 19.2M parameters. On BraTS2023, it achieves the best average Dice of 92.0% and the lowest HD95 of 2.7 at $\sim$1/7 the parameters of SegMamba-V2. For model validation, we also ran experiments on the MSD Spleen dataset, achieving competitive performance. The implementation details and code are publicly available at https://github.com/2802990u/GMD-Mamba.