Technical Summary: Patch-MoE Mamba

Problem Statement

Medical image segmentation relies heavily on Convolutional Neural Networks (CNNs) and Transformers. However, CNNs are inherently limited by local receptive fields, struggling to model long-range dependencies, while Transformers often incur quadratic computational and memory complexity. Recently, State Space Models (SSMs), particularly Mamba-based architectures, have emerged as efficient alternatives with linear sequence complexity.

Despite their promise, existing Mamba-based segmentation models face two critical limitations:

1. Disruption of Local Spatial Structure: Most models convert 2D feature maps into 1D sequences using pixel-wise raster scanning. This process can separate spatially adjacent pixels in the sequence, weakening local 2D spatial coherence. This is detrimental to dense prediction tasks requiring accurate boundary localization and local structure preservation.
2. Inflexible Directional Fusion: Current methods typically fuse outputs from multiple scanning directions via simple summation. This fixed aggregation assumes all directions are equally important at every spatial location, failing to adapt to the diverse sizes, shapes, textures, and boundary complexities of medical targets (e.g., small, low-contrast lesions vs. large, irregular structures).

Methodology

The authors propose Patch-MoE Mamba, a patch-ordered Mixture-of-Experts (MoE) state space architecture designed to address the aforementioned limitations. The architecture follows a U-Net-style design, replacing the standard Visual State Space (VSS) blocks in VM-UNetV2 with a new Patch-MoE VSS block.

1. Hierarchical Patch-Ordered Scanning

Instead of raster-scanning individual pixels, the method partitions feature maps into local patches and performs directional scanning in a patch-ordered manner.

• Mechanism: Given a feature map, the spatial grid is divided into non-overlapping patches of size $p \times p$. Within each patch, pixels are enumerated in row-major order before moving to the next patch.
• Benefit: This ensures all pixels within a patch remain consecutive in the sequence, preserving local spatial neighborhoods while still enabling long-range dependency modeling across patches.
• Hierarchical Strategy: To capture multi-scale context, the model employs a set of patch sizes $\{p^{(1)}, p^{(2)}, \dots\}$. Smaller patches capture fine-grained boundary details, while larger patches model coarser structures. Different scanning directions (forward, reverse, width-height forward, width-height reverse) can utilize distinct patch sizes to handle anisotropic structures.

2. MoE-Based Directional Fusion

To replace fixed summation, the authors introduce an adaptive fusion module treating directional outputs as "experts."

• Expert Construction: The module utilizes five experts:
  • Four direction-specific experts ($E_1$ to $E_4$): The normalized outputs of the four directional Mamba scanners.
  • One concatenation expert ($E_{concat}$): Formed by concatenating the four normalized maps along the channel dimension and projecting them back via a $1\times1$ convolution. This expert captures cross-directional correlations and provides a holistic, direction-agnostic view.
• Spatial-Aware Gating: A router network computes location-dependent weights to adaptively fuse these experts. The router is driven by a blend of local context (via depth-wise convolution) and global context (via adaptive average pooling), controlled by a learnable scalar $\alpha$.
• Residual Stabilization: A residual bypass adds the raw directional outputs back to the fused result ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). This ensures robust baseline responses and preserves strong directional signals, stabilizing training.

3. Semantics and Detail Infusion (SDI)

The encoder outputs are processed by an SDI module (adapted from U-Net v2) to enhance feature maps by infusing semantic information from higher levels and integrating finer details from lower levels via the Hadamard product.

Key Contributions

1. Patch-Ordered Scanning: A novel scanning mechanism that replaces pixel-wise raster scanning with patch-ordered traversal, effectively preserving local 2D spatial coherence while maintaining the efficiency of state space models.
2. MoE-Based Directional Fusion: An adaptive fusion module that replaces simple summation with a learnable gating network. It dynamically weights multiple directional experts and a concatenation expert based on spatial context, allowing the model to adapt to varying lesion scales and complexities.
3. Hierarchical Multi-Scale Modeling: The integration of hierarchical patch sizes allows the network to simultaneously capture fine-grained boundary details and coarse semantic structures.
4. Residual Stabilization: The inclusion of a residual path for raw directional outputs ensures training stability and prevents the loss of critical directional signals.

Experimental Results

The authors evaluated Patch-MoE Mamba on five public polyp segmentation datasets (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) and two skin lesion segmentation datasets (ISIC 2017, ISIC 2018).

• Performance: Patch-MoE Mamba achieved state-of-the-art performance across all datasets, consistently outperforming CNN-based (U-Net, U-Net v2) and Mamba-based baselines (VM-UNet, VM-UNetV2).
  • On the challenging ETIS dataset (known for low contrast and small polyps), the model improved the Dice score from 72.56% (VM-UNetV2) to 74.04%.
  • On ISIC 2018, it achieved a Dice score of 89.34%, surpassing the previous best of 88.36%.
• Ablation Studies:
  • Replacing raster scanning with patch-ordered scanning alone improved average Dice from 83.46% to 84.02%.
  • Adding the MoE fusion module further boosted the average to 84.30%.
  • Hierarchical patch sizes (coarse in early stages, fine in deeper stages) yielded the best results.
• Computational Complexity: While the patch-ordered scanning itself does not increase computational cost (token count remains unchanged), the inclusion of the concatenation expert increases parameters to ~70M and FLOPs to ~28G. The authors note this overhead is expected due to the denser channel interactions required for the MoE mechanism.

Significance and Claims

The paper claims that Patch-MoE Mamba offers a general and effective solution for medical image segmentation. By addressing the specific limitations of existing Mamba models—namely the loss of local spatial structure and the rigidity of directional fusion—the proposed architecture demonstrates superior adaptability to diverse medical targets with varying sizes, shapes, and boundary complexities. The results suggest that preserving local spatial neighborhoods through patch-ordered scanning, combined with adaptive expert fusion, significantly enhances boundary localization and segmentation accuracy, particularly in challenging scenarios involving low-contrast or small lesions.