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

2025-07-08

AI Summary

To address weak cross-modal generalization, high data heterogeneity, and inefficient long-range dependency modeling in 3D medical image segmentation, this paper proposes HoME—a Hierarchical Mixture of soft Experts. HoME introduces a novel two-stage hierarchical expert mechanism: bottom-level local experts enable modality-adaptive sparse computation via token-wise routing, while top-level experts integrate global contextual information to enhance generalization. Crucially, HoME replaces self-attention with the Mamba state-space model for efficient long-range spatial modeling, and incorporates Soft Mixture of Experts (SMoE) with adaptive token routing to improve computational efficiency. Evaluated on six mainstream 3D datasets across CT, MRI (T1/T2), and PET modalities, HoME achieves significant improvements over state-of-the-art methods, yielding an average Dice score gain of 2.3%. Moreover, it demonstrates superior robustness to annotation noise and missing modalities.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Multi-modal Vision
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

In recent years, artificial intelligence has significantly advanced medical image segmentation. However, challenges remain, including efficient 3D medical image processing across diverse modalities and handling data variability. In this work, 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 state-space model (SSM) backbone, HoME enhances sequential modeling through sparse, adaptive expert routing. The first stage employs a Soft Mixture-of-Experts (SMoE) layer to partition input sequences into local groups, routing tokens to specialized per-group experts for localized feature extraction. The second stage aggregates these outputs via a global SMoE layer, enabling cross-group information fusion and global context refinement. This hierarchical design, combining local expert routing with global expert refinement improves generalizability and segmentation performance, surpassing state-of-the-art results across datasets from the three most commonly used 3D medical imaging modalities and data quality.

Problem

• Efficient 3D medical image segmentation across diverse modalities
• Handling data variability in medical imaging
• Improving generalizability and segmentation performance

Innovation

• Hierarchical Soft Mixture-of-Experts (HoME) design
• Two-level token-routing for efficient modeling
• Mamba SSM backbone with adaptive expert routing