⚡ electrical engineering

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

This paper introduces HoME, a hierarchical soft mixture-of-experts framework built on the Mamba backbone that utilizes a two-level token-routing mechanism to enhance long-context modeling and achieve state-of-the-art performance in 3D medical image segmentation across diverse modalities and data qualities.

Original authors: Szymon Płotka, Gizem Mert, Maciej Chrabaszcz, Ewa Szczurek, Arkadiusz Sitek

Published 2026-02-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Szymon Płotka, Gizem Mert, Maciej Chrabaszcz, Ewa Szczurek, Arkadiusz Sitek

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive library containing millions of books (which represent a 3D medical scan of a human body). Your goal is to find specific, tiny details—like a small tumor or a specific organ boundary—while also understanding the big picture of how the whole library is arranged.

This paper introduces a new AI system called Mamba-HoME to solve this problem. Here is how it works, explained through simple analogies:

The Problem: Too Big for One Brain

Previous AI models tried to look at the whole library at once.

Old Method A (CNNs): Like a librarian who only looks at one bookshelf at a time. They are fast and good at seeing local details, but they can't see how the shelves connect across the whole room. They miss the "big picture."
Old Method B (Transformers): Like a librarian who tries to read every single book in the library simultaneously to understand the connections. This is great for seeing the big picture, but it takes forever and requires a super-computer to handle the memory. It gets too slow for high-resolution 3D scans.

The Solution: The "Mamba-HoME" Team

The authors created a new team of AI workers that combines the speed of Method A with the big-picture vision of Method B, but with a clever twist called HoME (Hierarchical Soft Mixture-of-Experts).

Think of HoME as a two-level management system for a massive construction crew:

Level 1: The Local Foremen (The "Local" Experts)

When the team receives a huge 3D scan, they don't send every single pixel to the main boss. Instead, they break the scan into small, manageable neighborhoods (groups).

In each neighborhood, a Local Foreman (an "expert") is assigned to handle the specific details of that area.
If a neighborhood has a lot of bone texture, a "bone expert" handles it. If it has fluid, a "fluid expert" takes over.
This is efficient because the experts only talk to the people in their own neighborhood first. They don't waste time shouting across the whole city.

Level 2: The City Planners (The "Global" Experts)

Once the Local Foremen have done their work, they send a summary up to the City Planners (the "global" experts).

These planners look at the summaries from all the neighborhoods.
They figure out how the neighborhoods fit together. For example, "The liver (neighborhood A) is right next to the stomach (neighborhood B)."
They refine the final map, ensuring the local details make sense in the global context.

Why "Mamba"?

The paper uses a backbone technology called Mamba. Think of Mamba as a highly efficient conveyor belt.

Unlike older systems that had to stop and look at every single item on the belt to understand the whole line (which is slow), Mamba moves along the belt linearly. It remembers what it saw earlier and uses that memory to understand what it's seeing now, all without getting bogged down.
This allows the system to handle massive 3D scans (like a whole human body) much faster and with less computer memory than previous methods.

The Results: Better Maps, Less Memory

The authors tested this system on three different types of medical "cameras":

CT scans (like X-rays that show bones and organs in 3D).
MRI scans (which show soft tissues like the brain and liver in great detail).
Ultrasound (which is often noisy and grainy).

What they found:

Accuracy: Mamba-HoME drew the boundaries of organs and tumors more accurately than any other current method. It was better at finding the "edges" of things.
Efficiency: Even though it is very smart, it didn't need a massive amount of computer memory (RAM) to run, which is a huge problem for 3D medical images.
Versatility: The system learned from CT and MRI scans and was so good at understanding the "shape" of organs that it could also work well on Ultrasound scans, even though it wasn't explicitly trained on them first.

In a Nutshell

The paper claims that Mamba-HoME is a new way to teach computers to "see" inside the human body. It does this by using a two-tier team of specialists: local experts who handle small details quickly, and global experts who ensure those details fit together perfectly. This approach is faster, uses less computer power, and creates more accurate medical maps than the current best tools.

Note: The paper focuses entirely on the technical performance of the AI model in segmentation tasks (drawing lines around organs). It does not claim the model is currently being used in hospitals for patient diagnosis or treatment planning, nor does it discuss future clinical applications beyond the technical benchmarks presented.

Technical Summary: Mamba Goes HoME

Problem Statement

Three-dimensional (3D) medical image segmentation is critical for diagnosis and treatment planning across modalities like Computed Tomography (CT), Magnetic Resonance Imaging (MRI), and Ultrasound (US). However, existing architectures face significant trade-offs:

Convolutional Neural Networks (CNNs) offer linear computational complexity but suffer from limited receptive fields, hindering the capture of global spatial patterns essential for understanding the hierarchical relationship between local lesions and global anatomical structures.
Vision Transformers (ViTs) utilize global attention to model long-range dependencies but incur quadratic computational complexity relative to the number of tokens. This makes them computationally prohibitive for high-resolution 3D volumetric data.
Selective State-Space Models (SSMs), such as Mamba, provide linear complexity for long-sequence modeling, offering an efficient alternative to ViTs. However, standard SSMs are not inherently designed to adaptively handle diverse local patterns found in complex, multi-scale medical data.
Mixture-of-Experts (MoE) frameworks can dynamically route features to specialized subnetworks to handle local patterns, but combining the global efficiency of SSMs with the localized adaptability of MoE in 3D medical imaging remains largely unexplored.

Methodology: Mamba-HoME

The authors propose Mamba-HoME, a novel architecture that integrates the Mamba Selective State-Space Model (SSM) backbone with a Hierarchical Soft Mixture-of-Experts (HoME) layer. The design aims to efficiently model local-to-global pattern hierarchies with linear computational complexity.

1. Hierarchical Soft Mixture-of-Experts (HoME)

The core innovation is a two-level token-routing layer designed to balance local specialization with global context integration:

Level 1: Local Grouping and Routing. Input sequences are partitioned into local groups. A Soft Mixture-of-Experts (SMoE) layer routes tokens within these groups to specialized local experts. This enables localized feature extraction while preserving spatial locality and reducing memory overhead compared to global routing.
Level 2: Global Aggregation. The outputs from the first level are aggregated and passed through a second-level global SMoE layer. This facilitates cross-group information fusion and global context refinement.
Scalability: By performing routing within groups before global aggregation, the method avoids the prohibitive memory costs of global routing on high-resolution 3D data (where token counts can reach millions), maintaining a computational complexity of $O(Nd)$.

2. Mamba-HoME Block

The architecture embeds the HoME layer into a U-shaped encoder-decoder network (inspired by SegMamba):

Backbone: Utilizes Mamba layers for efficient long-range dependency modeling with linear time complexity.
Gated Spatial Convolution (GSC): Integrated to capture local spatial priors before flattening data into 1D sequences for Mamba processing.
Dynamic Tanh (DyT): Replaces standard Layer Normalization with a point-wise nonlinearity ( $f(x) = w \cdot \tanh(\alpha \cdot x) + b$ ). This stabilizes gradients and reduces computational overhead by avoiding mean and variance calculations, accelerating training and inference.
Hierarchical Progression: The number of experts increases monotonically with depth ( $E_1 < E_2 < \dots$ ), while group sizes decrease ( $K_1 > K_2 > \dots$ ), allowing for progressively finer-grained processing as the network deepens.

Key Contributions

HoME Layer: Introduction of a two-level token-routing MoE layer that efficiently captures local-to-global pattern hierarchies by grouping tokens for local expert routing followed by global expert refinement.
Unified Architecture: Design of the Mamba-HoME Block, which integrates Mamba's SSMs with HoME, combining memory-efficient long-sequence processing with hierarchical expert routing.
Multi-Stage Segmentation Network: Embedding these components into a U-shaped architecture specifically tailored for 3D medical image segmentation, enhancing contextual representation across spatial and depth dimensions.

Experimental Results

The authors evaluated Mamba-HoME on four public datasets (PANORAMA, AMOS, FeTA 2022, MVSeg) and one in-house CT dataset, covering CT, MRI, and US modalities.

Performance: Mamba-HoME achieved state-of-the-art (SOTA) results in segmentation accuracy (Dice Similarity Coefficient - DSC) and boundary precision (Hausdorff Distance 95% - HD95) across all tested modalities.
- On the PANORAMA dataset, it achieved a DSC of 56.7% (pre-trained) and 54.8% (from scratch), outperforming baselines like SegMamba and Swin UNETR.
- On the FeTA 2022 fetal MRI dataset, it reached a DSC of 87.7% (pre-trained).
- On the MVSeg 3D US dataset, it achieved a DSC of 85.0% (pre-trained).
Efficiency: Despite having a larger parameter count (170.1M) compared to some competitors, the model demonstrated low GPU memory usage during inference, a critical advantage for high-resolution 3D data. While inference was approximately 30% slower than some baselines, the authors note this is a compelling trade-off for the significant gains in accuracy.
Generalizability: The model demonstrated strong cross-modal generalization. A model pre-trained on CT and MRI data successfully adapted to 3D Ultrasound tasks with minimal fine-tuning, outperforming SOTA methods in both DSC and HD95 metrics.
Ablation Studies:
- Optimal performance was found with specific expert counts ( $E_1 \in \{4, 8, 12, 16\}$ , $E_2 \in \{8, 16, 24, 32\}$ ) and group sizes ( $K \in \{2048, 1024, 512, 256\}$ ).
- Dynamic Tanh (DyT) normalization improved training and inference speed by ~6% compared to Layer Normalization without sacrificing accuracy.
- Supervised pre-training on large-scale CT and MRI datasets significantly enhanced performance and robustness across diverse anatomical regions and modalities.

Significance and Claims

The paper claims that Mamba-HoME represents a significant advancement in 3D medical image segmentation by effectively bridging the gap between the efficiency of SSMs and the adaptability of MoE frameworks. The authors assert that their hierarchical design allows models to capture the inherent local-to-global spatial hierarchies of medical data more effectively than previous CNN or Transformer-based approaches.

The work highlights the potential of modality-agnostic feature representations, demonstrating that a model pre-trained on CT and MRI can generalize effectively to Ultrasound, addressing the challenge of data variability across different imaging modalities. The authors position Mamba-HoME as a scalable solution for high-resolution 3D data that maintains competitive memory and compute efficiency while delivering superior segmentation accuracy.

Limitations Acknowledged: The authors note that scalability to extremely large-scale datasets (e.g., >10,000 scans) and the behavior under large-scale self-supervised learning (>200,000 scans) remain unexplored, identifying these as directions for future work.