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

원문 초록

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.