Mamba Goes HoME: Hierarchical Soft Mixture-of-Experts for 3D Medical Image Segmentation

Qualitative segmentation results. From left input slice, ground truth, and Mamba-HoME prediction.

TL;DR

We introduce Hierarchical Soft Mixture-of-Experts (HoME), a two-level token-routing layer for efficient long-context modeling, specifically designed for 3D medical image segmentation. Built on the Mamba Selective State Space Model (SSM) backbone, HoME enhances sequential modeling through adaptive expert routing. In the first level, a Soft Mixture-of-Experts (SMoE) layer partitions input sequences into local groups, routing tokens to specialized per-group experts for localized feature extraction. The second level aggregates these outputs through a global SMoE layer, enabling cross-group information fusion and global context refinement. This hierarchical design, combining local expert routing with global expert refinement, enhances generalizability and segmentation performance, surpassing state-of-the-art results across datasets from the three most widely used 3D medical imaging modalities and varying data qualities.

Paper

Mamba Goes HoME: Hierarchical Soft Mixture-of-Experts for 3D Medical Image Segmentation
Szymon Płotka^1,2*, Gizem Mert³, Maciej Chrabaszcz^4,5, Ewa Szczurek^1,3, Arkadiusz Sitek^6,7
1University of Warsaw, ²Jagiellonian University, ³Helmholtz Munich, ⁴Warsaw University of Technology, ⁵NASK, ⁶Massachusetts General Hospital, ⁷Harvard Medical School
Advances in Neural Information Processing Systems (NeurIPS) 2025

Datasets

For pre-training, training, and evaluation of Mamba-HoME, we use the following datasets:

Dataset	Source	Modality
AbdomenAtlas 1.1	Download	CT
TotalSegmentator MRI	Download	MRI
PANORAMA	Download	CT
AMOS	Download	CT / MRI
FeTA 2022	Download	MRI
MVSeg	Download	3D US

Installation

Step 1. Install Anaconda on Linux

[Optional] Install Anaconda on Linux

wget https://repo.anaconda.com/archive/Anaconda3-2024.06-1-Linux-x86_64.sh
bash Anaconda3-2024.06-1-Linux-x86_64.sh -b -p ./anaconda3
./anaconda3/bin/conda init
source ~/.bashrc

Step 2. Create a new virtual environment:

cd MambaHoME
conda create -n mambaHoME python=3.11
conda activate mambaHoME

Important

Before installing all dependencies, please install the Mamba packages first as provided in the official repository.

Step 3. Install the Mamba SSM package and its dependencies using the following steps:

pip install "causal-conv1d>=1.4.0"

# Install the core Mamba SSM package only
pip install mamba-ssm

# Install Mamba SSM along with causal Conv1d
pip install "mamba-ssm[causal-conv1d]"

# Install Mamba SSM with developer dependencies
pip install "mamba-ssm[dev]"

# (Optional) Troubleshooting: Install with no build isolation
pip install . --no-build-isolation

Step 4. Install all dependencies

pip install -r requirements.txt

Usage

Note

We trained our proposed Mamba-HoME and other state-of-the-art methods using 32-bit floating-point (float32) precision.

Step 1. Data pre-processing

The first step in our pipeline is resampling all datasets to a standardized voxel spacing.

python convert_spacing.py \
    --input_path /path/to/raw_dataset \
    --output_path /path/to/preprocessed_dataset \
    --spacing 0.8 0.8 3.0

Step 2. Training

The next step is to train the MambaHoME model using the prepared dataset.

python train.py \
    --data_train ../dataset/train --data_val ../dataset/val \
    --batch_size 2 --classes 3 --epochs 800 --lr 1e-4 --weight_decay 1e-4 \
    --optimizer AdamW --scheduler CALR --patch_size 128 128 128 --feature_size 48 \
    --use_checkpoint False --num_workers 12 --pin_memory True --use_pretrained False \
    --load_checkpoint False --checkpoint_name "MambaHoME" --model MambaHoME \
    --parallel True --num_devices 4 --strategy ddp

Step 3. Inference

After training, the best-performing checkpoint can be used for inference on test datasets.

python inference.py \
    --data_path ../dataset/test \
    --model_path ../checkpoints/MambaHoME_best.pth \
    --output_dir ../inference_results \
    --roi_size 192 192 48 \
    --sw_batch_size 4 \
    --num_workers 4 \
    --overlap 0.5 \
    --gpu

Requirements

• Linux
• CUDA 11.6+
• Python 3.11+
• PyTorch 2.4+
• MONAI 1.3.0+

Citation

If you use the code or methods in this repository, please cite:

@article{plotka2025mamba,
  title={Mamba Goes HoME: Hierarchical Soft Mixture-of-Experts for 3D Medical Image Segmentation},
  author={Płotka, Szymon and Mert, Gizem and Chrabaszcz, Maciej and Szczurek, Ewa and Sitek, Arkadiusz},
  journal={arXiv preprint arXiv:2507.06363},
  year={2025}
}

Related work

Our method builds upon the following works and their official implementations:

• Mamba: Linear-Time Sequence Modeling with Selective SSMs (Gu & Dao, 2024)
• From Sparse to Soft Mixtures of Experts (Puigcerver et al., 2024)
• SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation (Xing et al., 2024)

License

This project is released under the MIT License.

Acknowledgments

We acknowledge the use of the HPC cluster at Helmholtz Munich for the computational resources used in this study.