PaperlayerView exact revision on arXiv

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

arXiv:2401.13560v4cs.CV

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.

    Source: §1 · p. 1

  • 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.

    Source: §1 · p. 2, §2 · p. 2

  • Results

    SegMamba achieves the highest reported segmentation metrics on both BraTS2023 and AIIB2023 while maintaining better training memory and inference time than comparison approaches.

    Source: §3.3 · p. 7, §3.4 · p. 8

  • Takeaways & Limitations

    SegMamba demonstrates effective long-range dependency modeling for volumetric medical image segmentation with outstanding inference efficiency.

    Source: §4 · p. 9

Abstract

from arXiv · show

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.Source: §1 · p. 1

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.Source: §1 · p. 2
  • 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.Source: §1 · p. 2

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.Source: §1 · p. 2
  • 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.Source: §1 · p. 2

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.Source: §2 · p. 2

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.Source: §2.1 · p. 3

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.Source: §2.1 · p. 3

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.Source: §2.1 · pp. 3–4

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.Source: §2.1 · p. 4

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.Source: §2.2 · p. 4

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.Source: §3.3 · p. 7

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.Source: §3.3 · p. 7

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.Source: §3.3 · p. 8

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.Source: §3.4 · p. 8

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.Source: Table 5 · p. 7, §3.4 · p. 8

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.Source: §4 · p. 8
  • The gated spatial convolution module models spatial relationships before tri-orientated Mamba, while feature-level uncertainty estimation enhances multi-scale skip-connection features.Source: §4 · pp. 8–9
  • CRC-500 is introduced as a large-scale dataset for 3D colorectal cancer segmentation to support related research.Source: §4 · p. 9
  • SegMamba models long-range dependencies within volumetric data while maintaining outstanding inference efficiency, with experiments demonstrating effectiveness and universality.Source: §4 · p. 9

Related papers

Loading 2401.13560v4…