Source-linked AI summary
SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation
Zhaohu Xing, Tian Ye, Yijun Yang, Guang Liu, Lei Zhu
TL;DR
High-dimensional 3D medical images make global relationship modeling computationally challenging, while CNNs struggle with global relationships. SegMamba uses Mamba-based whole-volume modeling and achieves strong segmentation performance across BraTS2023 and AIIB2023 while maintaining inference efficiency.
Problem
CNN-based 3D segmentation methods struggle to model global relationships because convolution is inherently local.
Method
SegMamba combines a U-shaped architecture with tri-orientated Mamba blocks, gated spatial convolution, and uncertainty-enhanced skip connections for multi-scale volume modeling.
Results
SegMamba achieves the highest reported segmentation metrics on both BraTS2023 and AIIB2023 while maintaining better training memory and inference time than comparison approaches.
Takeaways & Limitations
SegMamba demonstrates effective long-range dependency modeling for volumetric medical image segmentation with outstanding inference efficiency.
Abstract
from arXiv · showhide
The Transformer architecture has shown a remarkable ability in modeling global relationships. However, it poses a significant computational challenge when processing high-dimensional medical images. This hinders its development and widespread adoption in this task. Mamba, as a State Space Model (SSM), recently emerged as a notable manner for long-range dependencies in sequential modeling, excelling in natural language processing filed with its remarkable memory efficiency and computational speed. Inspired by its success, we introduce SegMamba, a novel 3D medical image \textbf{Seg}mentation \textbf{Mamba} model, designed to effectively capture long-range dependencies within whole volume features at every scale. Our SegMamba, in contrast to Transformer-based methods, excels in whole volume feature modeling from a state space model standpoint, maintaining superior processing speed, even with volume features at a resolution of {$64\times 64\times 64$}. Comprehensive experiments on the BraTS2023 dataset demonstrate the effectiveness and efficiency of our SegMamba. The code for SegMamba is available at: https://github.com/ge-xing/SegMamba
1 Introduction
The introduction motivates broader receptive fields and global feature modeling for 3D medical image segmentation, while noting the efficiency challenge of Transformer-based approaches. It presents SegMamba, combining a U-shaped architecture with Mamba and introducing ToM and GSC modules for whole-volume, multi-scale modeling.
- Motivation: Extending the receptive field within 3D space is identified as critical for improving medical image segmentation performance.Large-kernel convolution is cited as one approach for modeling a broader range of features.
- Related work: Transformers use self-attention to extract global information, with UNETR learning single-scale features and SwinUNETR extracting multi-scale features.The introduction notes that these transformer-based methods improve segmentation performance.
- Related work: Mamba uses selection mechanisms and hardware-aware algorithms to model long-range dependencies while improving training and inference efficiency.The introduction also mentions U-Mamba and Vision Mamba as medical or vision applications of Mamba-based modeling.
- Contributions: SegMamba combines a U-shaped structure with Mamba to model whole-volume global features at multiple scales for 3D medical image segmentation.The paper describes it as the first method specifically using Mamba for 3D medical image segmentation.
- Contributions: The proposed tri-orientated Mamba module enhances whole-volume sequential modeling of 3D features, while gated spatial convolution enhances spatial feature representation.These are introduced as the ToM and GSC modules, respectively.
2 Method
SegMamba combines a multi-scale 3D encoder with TSMamba blocks, a convolutional decoder, and uncertainty-aware skip connections. Its TSMamba blocks preserve spatial structure with gated convolutions and model global dependencies across three orientations.
- Architecture: SegMamba comprises a 3D encoder with multiple TSMamba blocks, a convolution-based 3D decoder, and FUE-enhanced skip connections.The encoder models global information at different scales, while the decoder predicts segmentation results.
- Encoder: The encoder begins with a 7×7×7 depth-wise convolution using padding 3×3×3 and stride 2×2×2, producing 48-channel features at half input resolution.For input volume I ∈ R^C×D×H×W, the stem extracts z0 ∈ R^48×D/2×H/2×W/2.
- TSMamba Blocks: Each TSMamba block combines gated spatial convolution and tri-orientated Mamba with layer normalization and MLP-based feature enrichment.The blocks are applied across encoder stages with corresponding down-sampling layers.
- Gated Spatial Convolution: GSC uses 3×3×3 and 1×1×1 convolution blocks, pixel-wise multiplication for gating, feature fusion, and a residual connection.This module captures spatial relationships before 3D features are flattened into a sequence for Mamba processing.
- Tri-orientated Mamba: ToM flattens 3D features into forward, reverse, and inter-slice sequences to model global dependencies from three directions.The three directions are denoted f, r, and s, respectively.
- Feature-Level Uncertainty Estimation: FUE estimates uncertainty for each multi-scale skip feature and enhances lower-uncertainty features using ˜zi = zi + zi · (1 − ui).It averages across channels and applies a sigmoid function before feature reuse.
3 Experiments
Experiments evaluate SegMamba against six CNN- and Transformer-based methods on BraTS2023, AIIB2023, and CRC-500, showing strong segmentation performance across tumor, airway, and colorectal cancer tasks. Ablations further assess the contributions of GSC and ToM and examine SegMamba’s efficiency for global modeling.
- BraTS2023: 93.61%, 92.65%, and 87.71% are SegMamba’s highest Dice scores on BraTS2023 WT, TC, and ET, with corresponding HD95s of 3.37, 3.85, and 3.48.UX-Net has the best comparison average Dice of 89.69% and average HD95 of 4.81.
- AIIB2023: 88.59%, 70.21%, and 61.33% are SegMamba’s highest AIIB2023 scores for IoU, DLR, and DBR, respectively.The airway-tree task contains many tiny branches, making robust and continuous segmentation challenging.
- CRC-500: 48.46% Dice and 28.52 HD95 are SegMamba’s best CRC-500 results despite the dataset’s typically small cancer regions.The results indicate accurate detection of colorectal cancer regions.
- The Effectiveness of Proposed Modules: 2.88% Dice and 13.95% HD95 are the improvements from adding GSC in M2 over the basic Mamba-only method M1.The passage attributes this gain to improved spatial representation before the ToM module; M3 then models global information from three directions and reports 47.22% Dice and 33.32 HD95.
- The High Efficiency of SegMamba: SegMamba’s efficiency is evaluated against UX-Net, SwinUNETR, and self-attention-based global modeling in an ablation study focused on long-range and global relationships.Table 5 reports training memory, inference memory, inference time, and out-of-memory status, although the supplied passage does not provide their numerical values.
4 Conclusion
The paper introduces SegMamba, a general 3D medical image segmentation method based on Mamba, with modules for tri-oriented sequential modeling, gated spatial convolution, and feature-level uncertainty estimation. It also presents CRC-500 and reports effective, universal, and efficient volumetric segmentation with long-range dependency modeling.
- 4 Conclusion: SegMamba is presented as the first general 3D medical image segmentation method based on Mamba.The method includes tri-orientated Mamba, gated spatial convolution, and feature-level uncertainty estimation modules.
- 4 Conclusion: The gated spatial convolution module models spatial relationships before tri-orientated Mamba, while feature-level uncertainty estimation enhances multi-scale skip-connection features.
- 4 Conclusion: CRC-500 is introduced as a large-scale dataset for 3D colorectal cancer segmentation to support related research.
- 4 Conclusion: SegMamba models long-range dependencies within volumetric data while maintaining outstanding inference efficiency, with experiments demonstrating effectiveness and universality.