A mixed Mamba U-net for prostate segmentation in MR images

Sci Rep 2024 Deep Learning 7 Explanations View Original
Original Paper (PDF)

Unable to display PDF. Download it here or view on PMC.

Plain-English Explanations
Page 1
Why Accurate Prostate Segmentation on MRI Matters

Prostate cancer is the most prevalent male cancer in the United States, accounting for 29% of all new male cancer diagnoses. Early detection hinges on the ability to precisely identify and outline (segment) the prostate gland in magnetic resonance imaging (MRI) scans.

Currently, radiologists must trace prostate boundaries manually in MRI images -- a task that is both time-consuming and prone to inconsistency between different readers. This inter-observer variability means two radiologists examining the same scan may draw different boundaries, leading to differing measurements and diagnoses.

Automated segmentation using deep learning has emerged as a solution, but existing methods have shortcomings. Convolutional neural networks (CNNs) are efficient but struggle to capture relationships between distant parts of an image. Transformer-based models handle long-range context but require enormous computing resources for 3D medical images.

This paper introduces MM-UNet, a new deep learning architecture that combines the strengths of CNNs with a recently developed model called Mamba, which can capture global context across entire images at much lower computational cost than Transformers.

TL;DR: Accurate automated prostate segmentation from MRI is essential for early cancer detection, and this paper proposes a new AI model that overcomes limitations of both traditional and Transformer-based approaches.
Pages 1-2
What Is Mamba and Why It Matters for Medical Imaging

Mamba is a new type of AI architecture based on the Structured State Space Sequence (S4) model. Unlike Transformers, which must compare every part of an image to every other part (quadratic complexity), Mamba processes sequences with linear complexity -- meaning it scales efficiently even for very large, high-resolution 3D medical images.

For medical image segmentation, capturing long-range dependencies (understanding how distant parts of an image relate to each other) is critical. For example, understanding the full 3D shape of the prostate requires integrating information from many slices simultaneously -- something CNNs with their fixed, local receptive fields struggle with.

Mamba addresses this gap through a selection mechanism that allows the model to dynamically decide which information to retain or discard as it processes sequential data. This makes it highly effective for long, complex sequences -- whether those sequences are words in a sentence, nucleotides in a genome, or pixel values across an MRI scan.

MM-UNet integrates Mamba into the well-established U-Net encoder-decoder architecture, which has been the gold standard for medical image segmentation. The result is a hybrid model that combines CNN efficiency, Mamba's global awareness, and U-Net's proven structural design.

TL;DR: Mamba is a new AI modeling approach that efficiently captures global context across long sequences, making it ideal for augmenting traditional CNN-based medical image segmentation.
Pages 2-5
The Three Novel Modules Inside MM-UNet

MM-UNet introduces three custom components. The first is the Adaptive Feature Fusion Module (AFFM), which intelligently merges information from different network depths. The encoder extracts coarse, high-level semantic information (what is the prostate?) while shallow layers capture fine-grained spatial details (where exactly is the boundary?). AFFM uses channel attention to weight which features matter most and suppresses irrelevant background noise.

The second component is the Global Context-Aware Module (GCM), placed at the network's bottleneck layer. It uses a Mamba block to model long-range relationships across the entire 3D volume simultaneously. This allows the network to understand the prostate's global shape and spatial context -- something a standard CNN cannot do with its localized filters.

The third component is the Multi-Scale Anisotropic Convolution Module (MACM). Prostate MRI scans typically have anisotropic resolution -- meaning the spacing between image slices in one direction is different from the pixel spacing within each slice. MACM uses parallel convolutions of different sizes (kernels of 3, 7, 13, and 15) to capture features at multiple scales, while decomposing large 3D convolutions into smaller operations to reduce computational cost.

These three modules replace the standard skip connections in a U-Net -- the pathways that directly connect encoder and decoder layers. By upgrading these connections with intelligent feature fusion and global context, MM-UNet produces smoother, more accurate prostate boundaries.

TL;DR: MM-UNet's three novel modules address feature merging, global context capture, and multi-scale resolution handling -- together producing more accurate prostate boundary delineation.
Pages 5-6
Datasets and How Performance Was Measured

The model was evaluated on two publicly available benchmark datasets. The PROMISE12 dataset contains 80 T2-weighted prostate MRI scans from four different medical institutions, covering both benign conditions and prostate cancer. The ASPS13 (NCI-ISBI 2013) dataset contains 60 MRI scans from prostate cancer patients, with expert-annotated prostate boundaries.

Three standard metrics were used to evaluate performance. The Dice Similarity Coefficient (DSC) measures how much the predicted segmentation overlaps with the expert-drawn ground truth, expressed as a percentage (higher is better). The 95th Percentile Hausdorff Distance (95HD) measures boundary accuracy in millimeters (lower is better). Average Symmetric Distance (ASD) also measures boundary error in millimeters (lower is better).

A five-fold cross-validation approach was used, meaning the data was split into five groups and the model was trained and tested on each split in turn. This reduces the chance that results are specific to one particular data split and provides a more reliable estimate of real-world performance.

TL;DR: Two established public prostate MRI datasets were used for evaluation, with performance measured by overlap accuracy (DSC) and boundary precision (95HD and ASD).
Pages 6-7
Performance Results: How MM-UNet Compared to Competing Models

MM-UNet outperformed all nine comparison methods on both datasets. On the PROMISE12 dataset, it achieved a DSC of 92.39%, 95HD of 3.43 mm, and ASD of 1.42 mm -- the best results in all three metrics. The nearest competitor (MSD-Net) achieved a DSC of 91.88%, meaning MM-UNet's segmentation overlap was meaningfully more accurate.

On the ASPS13 dataset, MM-UNet again achieved the best performance with a DSC of 92.17%, 95HD of 3.61 mm, and ASD of 1.67 mm. Compared to the previous second-best method (CCT-Unet), it improved DSC by 0.94% and ASD by 0.28 mm -- modest but clinically meaningful improvements in boundary precision.

Compared to the classic V-Net baseline, MM-UNet improved DSC by 3.35%, 95HD by 0.78 mm, and ASD by 0.61 mm, illustrating the substantial advantage of the new architecture over older deep learning approaches. Visual comparisons showed that MM-UNet produced fewer incorrect over-segmentations and under-segmentations at prostate boundaries.

TL;DR: MM-UNet achieved state-of-the-art prostate segmentation on both benchmark datasets, outperforming all nine competing models in overlap accuracy and boundary precision.
Pages 7-9
Ablation Study: Confirming Each Module's Contribution

The researchers conducted an ablation study -- systematically removing each module to measure how much it contributes to overall performance. Starting from the baseline 3D U-Net (DSC 89.90%), adding each component progressively improved results.

Adding the Mamba-based GCM alone improved DSC by 0.81% over the Res2Net encoder baseline, confirming that global context awareness helps the network better identify the prostate's full extent. Adding the AFFM improved DSC by 0.98%, showing that better feature fusion between encoder and decoder layers is especially impactful.

The complete MM-UNet with all four components achieved a 2.49% DSC improvement and 1.13 mm improvement in 95HD over the baseline -- confirming that the modules work synergistically rather than just additively. Ablation tests on MACM also showed that including all four convolution kernel sizes (3, 7, 13, 15) together with anisotropic convolutions yielded the best results.

TL;DR: Ablation tests confirmed that all three novel modules contribute meaningfully to performance, with the full model improving DSC by 2.49% over the baseline through synergistic effects.
Page 10
Limitations and Future Directions

Despite its strong performance, MM-UNet has limitations. The primary constraint is data availability: manual annotation of prostate boundaries by expert radiologists is expensive and time-consuming, meaning training datasets remain relatively small. This limits how much the model can learn and may affect performance on unusual anatomy cases.

The model was designed specifically for prostate MRI segmentation and has not been tested on other organs or imaging modalities. Its transferability to tasks like kidney, liver, or brain segmentation, or to CT and ultrasound images, remains uncertain.

Future work will focus on using self-supervised learning -- a technique for training AI on large amounts of unlabeled data -- to improve robustness and reduce dependence on expert-annotated training sets. The authors also plan to test MM-UNet on additional segmentation tasks across different anatomical structures and imaging modalities, moving toward a more generalizable AI tool for medical imaging.

TL;DR: MM-UNet shows strong performance but is currently limited to prostate MRI; future work will use self-supervised learning to expand its capabilities to other organs and imaging types.
Citation: Open Access, . Available at: PMC11358272.