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

Technical Summary: MHMamba for 3D Brain Tumor Segmentation

Problem Statement
Brain tumor segmentation in 3D MRI faces significant challenges due to the high heterogeneity of tumor morphology and multimodal contrast. While Convolutional Neural Networks (CNNs) excel at modeling local textures and boundaries, they struggle with long-range dependencies due to limited receptive fields. Conversely, Transformer-based models can capture global context but suffer from quadratic computational and memory complexity ($O(N^2)$), making them difficult to apply to high-resolution 3D volumes without aggressive downsampling or windowing, which can lead to context incoherence and boundary inconsistencies. Furthermore, existing hybrid approaches often suffer from statistical mismatches between convolutional and attention branches (e.g., Batch/Instance Normalization vs. Layer Normalization), leading to unstable training and feature fusion issues.

Methodology
The paper proposes MHMamba (Multi-Head Mamba), a U-shaped encoder-decoder architecture designed to model global features with linear complexity while preserving high-resolution details. The framework consists of three primary components:

1. Multi-Head Mamba (MHM) Module:

   • Replaces standard blocks in the encoder with a multi-head State Space Model (SSM).
   • The channel dimension is split into parallel SSM heads ($N_h$), allowing different heads to learn complementary long-range dependency patterns from multimodal inputs without fixed modality assignments.
   • Each head serializes 3D spatial features into a sequence, applies selective state updates (Mamba), and fuses the outputs via concatenation and projection.
   • This design maintains linear complexity while enhancing global modeling capabilities and training stability under multimodal conditions.

2. Channel-Spatial Calibration Attention (CSCA) Module:

   • Designed to align feature statistics and enhance lesion responses after the multi-head Mamba processing.
   • It employs a dual-path mechanism:
     • Channel Path: Uses global average and max pooling followed by an MLP to generate channel weights.
     • Spatial Path: Innovatively integrates four spatial statistics (mean, standard deviation, maximum, and minimum) to generate a spatial weight map, capturing heterogeneity and salient peaks.
   • A gating mechanism adaptively fuses these branches, followed by a residual connection to preserve original feature responses.

3. Adaptive Gated Fusion (AGF) Module:

   • Located at the skip connections between the encoder and decoder.
   • Addresses the conflict between high-level semantics and low-level details by dynamically fusing multi-scale features.
   • Input features are split into subgroups along the channel dimension. For each subgroup, the module learns independent fusion weights to balance encoder and decoder features, preventing feature conflicts that often lead to inconsistent boundary predictions.

Key Contributions

• Multi-Head Mamba Architecture: The paper introduces a multi-head SSM approach for 3D MRI that decomposes channels into parallel heads. This enhances long-range modeling and multimodal adaptability while retaining the linear complexity advantage of Mamba.
• Feature Calibration and Fusion: The integration of the CSCA module to unify feature scales and statistics, and the AGF module to dynamically fuse multi-scale features at skip connections, significantly improves boundary consistency and the detection of small-volume lesions.
• Systematic Validation: The method was rigorously evaluated on the BraTS2021 and BraTS2023 datasets, demonstrating stable improvements over strong baselines (including nnU-Net, SegMamba, and Transformer-based models) in terms of Dice scores and Hausdorff Distance (HD95).

Experimental Results
Experiments conducted on BraTS2021 and BraTS2023 datasets show that MHMamba achieves state-of-the-art or near-state-of-the-art performance:

• BraTS2021: Achieved a mean Dice score of 91.02% and an average HD95 of 3.38, outperforming baselines like SegMamba (90.32% Dice, 4.18 HD95) and nnU-Net (89.40% Dice, 5.58 HD95).
• BraTS2023: Achieved a mean Dice score of 90.23% and an average HD95 of 3.67.
• Ablation Studies:
  • The addition of CSCA, AGF, and the Multi-Head Mamba module individually contributed to performance gains across all tumor subregions (Whole Tumor, Tumor Core, and Enhancing Tumor).
  • The optimal configuration was found to be 4 heads, which balanced representational capacity and computational efficiency; increasing to 8 heads led to performance degradation due to over-partitioning and optimization noise in 3D training with small batch sizes.
  • Qualitative results indicated that MHMamba produces less fragmented predictions with tighter boundaries, particularly in heterogeneous tumor regions and small enhancing areas.

Significance and Claims
The paper claims that MHMamba effectively addresses the trade-off between global context modeling and local detail preservation in 3D brain tumor segmentation. By leveraging the linear complexity of State Space Models through a multi-head architecture, the method avoids the quadratic memory overhead of Transformers while overcoming the local receptive field limitations of CNNs. The proposed modules (CSCA and AGF) specifically target the statistical mismatches and context incoherence often found in multimodal 3D segmentation, resulting in a robust, efficient, and accurate framework suitable for clinical applications such as radiotherapy planning and preoperative assessment. The authors emphasize that the method maintains the efficiency advantages of Mamba-based modeling while significantly improving sensitivity to tumor cores and small-volume lesions.