Taming Mambas for Voxel Level 3D Medical Image Segmentation

TL;DR

This paper introduces Mamba, a recurrent neural network based on state space models, which effectively handles long-context 3D medical image segmentation tasks with linear complexity, outperforming traditional CNNs and transformers.

Contribution

The paper proposes Mamba, a novel RNN architecture utilizing state space models for efficient, high-performance 3D medical image segmentation at the voxel level.

Findings

Mamba achieves superior segmentation accuracy compared to CNNs and transformers.

Mamba operates with linear complexity, enabling processing of large 3D volumes.

Mamba outperforms existing models on benchmark datasets.

Abstract

Recently, the field of 3D medical segmentation has been dominated by deep learning models employing Convolutional Neural Networks (CNNs) and Transformer-based architectures, each with their distinctive strengths and limitations. CNNs are constrained by a local receptive field, whereas transformers are hindered by their substantial memory requirements as well as they data hungriness, making them not ideal for processing 3D medical volumes at a fine-grained level. For these reasons, fully convolutional neural networks, as nnUNet, still dominate the scene when segmenting medical structures in 3D large medical volumes. Despite numerous advancements towards developing transformer variants with subquadratic time and memory complexity, these models still fall short in content-based reasoning. A recent breakthrough is Mamba, a Recurrent Neural Network (RNN) based on State Space Models (SSMs) outperforming Transformers in many long-context tasks (million-length sequences) on famous natural language processing and genomic benchmarks while keeping a linear complexity.