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

TL;DR

This paper introduces MambaX-Net, a dual-scan, semi-supervised 3D segmentation framework designed for longitudinal prostate MRI analysis in Active Surveillance. By fusing current and previous-timepoint images through a Mamba-enhanced cross-attention module (M-CAM) and seeding anatomical priors via a Shape Extractor Module (SEM), the model effectively leverages temporal information to improve prostate zone delineation under limited expert labels. A specialized loss (FT_DCE) and a self-training pipeline using pseudo-labels enable robust training without extensive annotations, with experiments showing superior performance over state-of-the-art baselines, especially in data-scarce scenarios. The work highlights data efficiency, robustness to label quality, and potential clinical impact for longitudinal monitoring of prostate cancer.

Abstract

Active Surveillance (AS) is a treatment option for managing low and intermediate-risk prostate cancer (PCa), aiming to avoid overtreatment while monitoring disease progression through serial MRI and clinical follow-up. Accurate prostate segmentation is an important preliminary step for automating this process, enabling automated detection and diagnosis of PCa. However, existing deep-learning segmentation models are often trained on single-time-point and expertly annotated datasets, making them unsuitable for longitudinal AS analysis, where multiple time points and a scarcity of expert labels hinder their effective fine-tuning. To address these challenges, we propose MambaX-Net, a novel semi-supervised, dual-scan 3D segmentation architecture that computes the segmentation for time point t by leveraging the MRI and the corresponding segmentation mask from the previous time point. We introduce two new components: (i) a Mamba-enhanced Cross-Attention Module, which integrates the Mamba block into cross attention to efficiently capture temporal evolution and long-range spatial dependencies, and (ii) a Shape Extractor Module that encodes the previous segmentation mask into a latent anatomical representation for refined zone delination. Moreover, we introduce a semi-supervised self-training strategy that leverages pseudo-labels generated from a pre-trained nnU-Net, enabling effective learning without expert annotations. MambaX-Net was evaluated on a longitudinal AS dataset, and results showed that it significantly outperforms state-of-the-art U-Net and Transformer-based models, achieving superior prostate zone segmentation even when trained on limited and noisy data.

Figures (5)

• Figure 1: High-level architecture of MambaX-Net showing a U-Net-like dual-encoder structure $(Enc_t, Enc_{t-1})$ with shared weights for processing scans at two time points ($I_t$ and $I_{t-1}$), and a single decoder $(Dec_t)$. The M-CAM block fuses latent features from $I_t$ and $I_{t-1}$ with the SEM output, producing enhanced representations for prostate zone segmentation.
• Figure 2: The Mamba-enhanced Cross-Attention Module (M-CAM) structure.
• Figure 3: Distribution of HD95 computed on $AS_{test}$ across WP, PZ, and TZ for the models trained on $AS_{train}$ with $n=50$.
• Figure 4: Qualitative comparison of each architecture's best-performing model on two $AS_{test}$ patients, comparing models trained with n=291 (nnU-NetAS, SwinUNETRAS and SegMambaAS) against those trained with n=50 (MambaX-Net and DSM).
• Figure 5: Effect of pseudo-label noise on the model's performance. Noise % represents the percentage of pseudo-labels replaced with Gaussian noise used for training, the y-axis shows the mean DSC computed on $AS_{test}$.