SegResMamba: An Efficient Architecture for 3D Medical Image Segmentation

TL;DR

Transformer-based approaches for 3D medical segmentation suffer from high memory and environmental costs. SegResMamba integrates Tri-oriented Mamba ToM blocks with convolutional layers in an encoder–decoder to capture global context efficiently, formalized as $ToM(z) = Mamba(z_f) + Mamba(z_r) + Mamba(z_s)$. It achieves competitive mean Dice on BTCV, BraTS2021, and Spleen datasets with substantially lower MACs and training memory, and displays reduced $CO_2$ emissions compared with other SOTA methods, making it suitable for resource-constrained deployment. This work demonstrates that SSM-based hybrids can deliver scalable, eco-friendly 3D medical image segmentation without sacrificing accuracy.

Abstract

The Transformer architecture has opened a new paradigm in the domain of deep learning with its ability to model long-range dependencies and capture global context and has outpaced the traditional Convolution Neural Networks (CNNs) in many aspects. However, applying Transformer models to 3D medical image datasets presents significant challenges due to their high training time, and memory requirements, which not only hinder scalability but also contribute to elevated CO$_2$ footprint. This has led to an exploration of alternative models that can maintain or even improve performance while being more efficient and environmentally sustainable. Recent advancements in Structured State Space Models (SSMs) effectively address some of the inherent limitations of Transformers, particularly their high memory and computational demands. Inspired by these advancements, we propose an efficient 3D segmentation model for medical imaging called SegResMamba, designed to reduce computation complexity, memory usage, training time, and environmental impact while maintaining high performance. Our model uses less than half the memory during training compared to other state-of-the-art (SOTA) architectures, achieving comparable performance with significantly reduced resource demands.

Figures (3)

• Figure 1: a) Overview of SegResMamba architecture, b) Convolution Mamba mixed block, and c) Tri-oriented Mamba
• Figure 2: Average Dice Scores for BTCV, Spleen, and BraTS2021 datasets plotted against training memory (in GB) for different models using image size $128\times128\times128$ for BTCV and BRATS dataset and $96\times96\times96$ for Spleen dataset with batch size 1.
• Figure 3: Mean dice score of BraTS dataset against CO$_2$ emission with 5-fold cross-validation settings for different models.