Taming Mambas for Voxel Level 3D Medical Image Segmentation

TL;DR

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.

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.

Figures (4)

• Figure 1: From left to right: (a) The unidirectional Mamba layer, which processes input sequences only in the forward direction. The layers within the gray square collectively form the Mamba Block. (b) The bidirectional Mamba layer, consisting of two unidirectional Mamba layers: the left branch processes the forward sequences, while the right branch processes the reversed sequences.
• Figure 2: U-Net architecture integrating our proposed Mamba Layers. By properly selecting the Mamba Layers (turquoise arrows), SegMamba, BiSegMamba, and MultiSegMamba are obtained. To obtain SegMambaSkip the currently displayed Mamba Layers (turquoise arrows) must be replaced by the standard U-Net convolution (blue arrow) and corresponding Mamba Layers must be placed within the skip connections (gray arrows).
• Figure 3: To achieve multi-directionality, four Bidirectional 3D Mamba Layers are employed, corresponding to four out of the six possible permutations of the triplet (H, W, D). The outputs of each layer is stacked, and the mean per token is computed.
• Figure 4: Visualization of segmentation results for four sample cases from the Synapse Abdomen evaluation set. Annotation errors are marked with red dashed boxes. The figure is best viewed in color and zoomed in. From left to right: Input, Ground Truth (GT), SegMamba, BiSegMamba, SegMambaSkip, and MultiSegMamba.