Prostate Cancer Detection in Bi-Parametric MRI using Zonal Anatomy-Guided U-Mamba with Multi-Task Learning

Michael S. Larsen, Syed Farhan Abbas, Gabriel Kiss, Mattijs Elschot, Tone F. Bathen, Frank Lindseth

Keywords: Deep Learning, Medical Image Analysis, Prostate Cancer, Mamba, Multi-Task Learning

Abstract: Prostate cancer (PCa) remains a leading cause of cancer-related morbidity, emphasizing the need for accurate and non-invasive diagnostic tools. While deep learning models have advanced PCa detection in magnetic resonance imaging (MRI), they often fail to integrate anatomical knowledge. This study evaluates U-Mamba, a deep learning architecture designed to enhance long-range dependency modeling with linear time complexity, for PCa detection. Furthermore, a multi-task learning (MTL) extension, U-Mamba MTL, is introduced to incorporate prostate zonal anatomy, aligning with clinical diagnostic workflows. The models were assessed using diverse datasets, including the PI-CAI hidden tuning cohort (N=100) and an in-house collected out-of-distribution cohort (N=200). Results demonstrate that U-Mamba achieves state-of-the-art detection performance, while U-Mamba MTL further improves PCa detection through the auxiliary zonal segmentation task. These findings highlight the potential of integrating U-Mamba with anatomical context to improve PCa detection. The code and model weights are available at https://github.com/mokkalokka/U-MambaMTL.

Primary Subject Area: Detection and Diagnosis

Secondary Subject Area: Segmentation

Paper Type: Both

Registration Requirement: Yes

Reproducibility: https://github.com/mokkalokka/U-MambaMTL

Submission Number: 196