🤖 machine learning

MambaX-Net: Dual-Input Mamba-Enhanced Cross-Attention Network for Longitudinal Prostate MRI Segmentation

MambaX-Net is a novel semi-supervised, dual-input 3D segmentation architecture that leverages a Mamba-enhanced Cross-Attention module and a Shape Extractor to effectively segment prostate zones in longitudinal Active Surveillance MRI by utilizing previous time-point data and pseudo-labels, thereby outperforming existing models in scenarios with limited expert annotations.

Original authors: Yovin Yahathugoda, Davide Prezzi, Patricia A. Gutierrez, Piyalitt Ittichaiwong, Vicky Goh, Sebastien Ourselin, Michela Antonelli

Published 2026-08-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yovin Yahathugoda, Davide Prezzi, Patricia A. Gutierrez, Piyalitt Ittichaiwong, Vicky Goh, Sebastien Ourselin, Michela Antonelli

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Prostate cancer is a common disease among men, and for many cases that are not immediately life-threatening, doctors choose a strategy called active surveillance. Instead of rushing into surgery or radiation, which can cause significant side effects, patients are monitored closely over time. This involves regular blood tests and repeated magnetic resonance imaging, or MRI, scans to watch for any signs that the cancer is growing or changing. The goal is to catch progression early while avoiding unnecessary treatment for those who do not need it. However, reading these scans is a demanding task for radiologists. They must carefully outline the entire prostate gland and its internal sections to measure volume and spot subtle changes. Doing this manually for every scan in a long-term monitoring program is slow and prone to human error, creating a need for artificial intelligence to help automate the process.

The challenge for computers is that prostate anatomy is not static. Over months or years, the gland can change shape, shrink, or shift slightly, and the quality of the MRI images can vary. Most existing artificial intelligence models are trained to look at a single scan in isolation, like a snapshot. They do not understand the story of the patient's history. If a computer sees a new scan without knowing what the prostate looked like a year ago, it might struggle to distinguish between a normal change and a sign of disease. Furthermore, training these systems usually requires thousands of images that have been painstakingly drawn by expert doctors, a resource that is scarce and expensive to produce.

A team of researchers at King's College London has developed a new approach to solve these problems, described in a recent study. They created a system named MambaX-Net, which is designed to look at two scans from the same patient at once: the current image and the one from the previous visit. Instead of treating the new scan as a standalone picture, the system uses the old scan and the previous medical drawing of the prostate as a guide. This allows the computer to understand how the anatomy has evolved over time, rather than just guessing based on a single moment.

The core of this new system relies on two specific innovations. First, it uses a component that acts like a memory bank, taking the old medical drawing and converting it into a compact map of the prostate's shape and size. This map serves as an anatomical anchor, reminding the system what the gland should look like for that specific person. Second, the system employs a specialized processing unit that compares the current scan with the previous one, searching for connections across the entire image. This unit is designed to be efficient, allowing it to track long-range relationships within the 3D volume of the prostate without getting bogged down by the massive amount of data. By combining the memory of the old shape with a deep comparison of the new image, the system can refine its boundaries, especially in tricky areas where the prostate meets other tissues or where the shape is irregular.

To test this idea, the researchers did not rely on a perfect set of expert drawings, which are rare for longitudinal studies. Instead, they used a clever training method. They started with a powerful, pre-existing artificial intelligence model to generate rough, automated drawings for a large set of patient scans. These automated drawings, while not perfect, served as a starting point. The new system then learned to correct these rough drafts by using the history of the patient's scans. The researchers tested their system on a real-world dataset of patients undergoing active surveillance, a group that included difficult cases such as men who had undergone prostate surgery or those with poor image quality due to metal implants or other factors.

The results showed that this new method significantly outperformed existing state-of-the-art models. Even when trained on a relatively small number of patients with imperfect, automated labels, the system produced more accurate outlines of the prostate and its internal zones than models that had been trained on much larger datasets. It was particularly successful at maintaining smooth, continuous boundaries, avoiding the jagged or fragmented errors that often plague other systems. The system proved robust even when the automated labels it learned from contained mistakes, suggesting that the use of historical context helps the computer ignore noise and focus on the true anatomy.

The study also explored how well the system could handle different types of errors in the training data. They found that while the system was resilient to some level of inaccuracy, it was most sensitive to errors in the outer layer of the prostate, a thin and irregular region that is difficult to define. This aligns with the known difficulty radiologists face in marking this specific area. Despite this sensitivity, the system consistently delivered better results than its competitors, demonstrating that leveraging the timeline of a patient's care is a powerful way to improve medical imaging analysis.

This work suggests a path forward for clinical practice where expert time is limited. By using the patient's own medical history as a guide, artificial intelligence can learn to segment complex organs with high precision, even when the training data is imperfect. The researchers note that while their system is a significant step forward, it is designed to work with pairs of scans, looking at one visit at a time in sequence. They acknowledge that for very long-term monitoring, the system would need to be carefully managed to ensure that small errors do not accumulate over many years. Nevertheless, the ability to accurately track the prostate over time without requiring a massive library of hand-drawn maps offers a promising tool for making active surveillance more efficient and reliable for patients.

Technical Summary: MambaX-Net for Longitudinal Prostate MRI Segmentation

Problem Statement
Active Surveillance (AS) is a standard treatment for low-to-intermediate-risk prostate cancer, requiring regular monitoring via serial MRI to track disease progression. Accurate segmentation of the whole prostate (WP) and its sub-regions (peripheral zone PZ, transition zone TZ) is critical for volume estimation and automated lesion detection. However, existing deep learning models face two primary limitations in this context:

Longitudinal Modeling Gap: Most state-of-the-art (SOTA) models are trained on single-time-point data. They treat each scan in isolation, failing to leverage the rich temporal information available in AS protocols where patients undergo multiple scans over time.
Annotation Scarcity: Longitudinal datasets are scarce and expensive to annotate manually. Existing semi-supervised approaches often struggle with the domain shift between public pre-training datasets and private clinical AS data, or fail to handle the noise inherent in pseudo-labels generated for unlabeled longitudinal scans.
Architectural Limitations: While Transformers capture long-range context, they often require large datasets and struggle with fine-grained details. Conversely, CNNs (like nnU-Net) are robust but lack explicit mechanisms to model temporal evolution without rigid image registration, which is unreliable for prostate MRI due to non-rigid deformations.

Methodology
The authors propose MambaX-Net, a novel semi-supervised, dual-input 3D segmentation architecture designed specifically for longitudinal data. The model extends the 3D nnU-Net framework to process a pair of scans: the current scan at time $t$ ( $I_t$ ) and the previous scan at time $t-1$ ( $I_{t-1}$ ), along with the segmentation mask from the previous time point ( $M_{t-1}$ ).

Key architectural components include:

Dual-Encoder Structure: Two encoders with shared weights process $I_t$ and $I_{t-1}$ . The encoder for $I_{t-1}$ is truncated to extract only the top three feature levels to balance computational efficiency with the preservation of high-resolution spatial details.
Shape Extractor Module (SEM): A sequential block of 3D convolutional layers that encodes the previous segmentation mask ( $M_{t-1}$ ) into a latent anatomical representation ( $f_{SEM}^{t-1}$ ). This serves as an explicit anatomical prior, capturing patient-specific shape and size information without requiring explicit image registration.
Mamba-Enhanced Cross-Attention Module (M-CAM): This is the core innovation. It integrates Mamba blocks (State Space Models) into a cross-attention mechanism.
- The current scan features ( $f_t$ ) act as queries.
- The previous scan features ( $f_{t-1}$ ) combined with the SEM output ( $f_{SEM}^{t-1}$ ) act as keys and values.
- Mamba blocks process these embeddings to efficiently capture long-range spatial and temporal dependencies with linear scaling, avoiding the quadratic complexity of Transformers.
- The module performs implicit feature alignment in the latent space, fusing the current scan's image data with the structural guidance of the previous mask.
Self-Training Framework: To address the lack of expert annotations in longitudinal datasets, the authors employ a semi-supervised strategy. A pre-trained nnU-Net on the public PI-CAI dataset generates pseudo-labels for the unlabeled AS training data. These pseudo-labels are used to fine-tune MambaX-Net, with the previous time point's pseudo-label serving as input and the current time point's pseudo-label used for loss calculation.
Loss Function: A modified Focal Tversky Dice Cross-Entropy (FT_DCE) loss is used to handle class imbalance and focus on hard negative examples, particularly at the irregular prostate apex and base.

Key Contributions

First Mamba-based Longitudinal Segmentation: MambaX-Net is the first architecture to jointly model longitudinal MRI evolution and explicit prior-mask anatomical context using Mamba blocks.
Novel Module Integration: The introduction of the M-CAM and SEM modules allows the model to ground global attention in patient-specific geometry, refining zone boundaries without relying on rigid registration.
Data Efficiency: The framework demonstrates that high-quality segmentation can be achieved with limited expert annotations by leveraging noisy pseudo-labels and temporal context, outperforming models trained on larger datasets.

Experimental Results
The model was evaluated on a private longitudinal AS dataset (317 patients) and compared against SOTA baselines (nnU-Net-V2, SwinUNETR-V2, SegMamba) and a previous dual-scan model (DSM).

Performance: MambaX-Net significantly outperformed all baselines in Dice Similarity Coefficient (DSC) and 95th percentile Hausdorff Distance (HD95) across the WP, PZ, and TZ, particularly in low-data regimes ( $n \le 50$ patients for fine-tuning).
Robustness: The model maintained superior performance even when trained on noisy pseudo-labels, whereas single-scan models (especially SwinUNETR) suffered severe degradation.
Qualitative Analysis: Visualizations showed that MambaX-Net produced smooth, contiguous 3D volumes with accurate boundary delineation, even in challenging cases involving post-surgical anatomy (TURP) or image artifacts, where other models exhibited boundary leakage or disjointed predictions.
Efficiency: Despite having a large parameter count (426.4M), MambaX-Net requires less peak GPU memory (14.4 GB) and is faster (276.3 ms) than SwinUNETR (25.1 GB, 649.5 ms) due to the linear scaling of Mamba blocks, avoiding the quadratic memory bottlenecks of Transformers.

Significance and Claims
The paper claims that MambaX-Net offers a feasible path toward clinical translation for longitudinal prostate MRI analysis without the need for extensive manual annotation or explicit image registration. By effectively combining the efficiency of Mamba with the structural guidance of a shape extractor, the model overcomes the limitations of both CNNs and Transformers in handling temporal data. The authors emphasize that the model's ability to leverage pseudo-labels while maintaining high accuracy in data-scarce, noisy environments addresses a critical bottleneck in deploying AI for Active Surveillance.

The authors modestly acknowledge limitations, noting that the study is based on a single-institution dataset and that the dual-input design currently processes pairwise time points rather than modeling three or more simultaneously. They also note that performance can degrade if the pseudo-labels become highly corrupted, particularly affecting the PZ, suggesting future work in uncertainty-aware self-training.

Technical Summary: MambaX-Net for Longitudinal Prostate MRI Segmentation

More like this