MHMamba: Multi-Head Mamba for 3D Brain Tumor Segmentation

TL;DR

MHMamba introduces a novel multi-head state-space model within a U-shaped architecture for efficient, stable, and accurate 3D brain tumor segmentation, addressing limitations of CNNs and Transformers.

Contribution

The paper proposes MHMamba, combining multi-head state-space modeling with a U-shaped network, enhancing long-range dependency modeling and boundary accuracy in 3D MRI segmentation.

Findings

Achieved significant improvements in accuracy, boundary smoothness, and sensitivity on BraTS datasets.

Maintained linear complexity while enhancing long-range dependencies and stability.

Validated effectiveness and versatility through extensive experiments and ablations.

Abstract

Brain tumors exhibit high heterogeneity in morphology and multimodal contrast, making manual slice-by-slice de lineation time-consuming and experience-dependent, thus necessitating efficient and stable automated segmentation methods. To address the limitations of CNNs in modeling long-range dependencies, and the heavy computational and memory overhead and inter-block contextual in coherence of Transformers in 3D MRI, this paper proposes Multi-Head Mamba (MHMamba). This method combines a U-shaped architecture with a multi-head state-space model (Mamba), splitting the channel dimension into parallel SSM heads and aggregating them with residuals. This enhances long-range representation and improves the stability of multimodal training while maintaining linear complexity. To further align statistics and enhance lesion response, we designed a channel-space calibration module for multi-head outputs and introduced an adaptive fusion mechanism at skip connections to dynamically connect global semantics with local details, thereby improving boundary consistency and the detection of small-volume lesions. We conducted experiments and ablations on BraTS2021 and BraTS2023. The results showed that MHMamba achieved stable and significant improvements in overall accuracy, boundary smoothness, and sensitivity to tumor core and small-volume enhancement areas, while preserving the linear-complexity advantage of Mamba-based modeling, thus verifying the effectiveness and versatility of the method.