Towards a Reliable Deep Learning Framework for Prostate Cancer Diagnosis using Ultrasound

Abstract

Reliable prostate cancer (PCa) detection using ultrasound is crucial for improving patient outcomes. Developing deep learning (DL) models for PCa detection is hindered by noisy labels and cancer heterogeneity. The purpose of this work is to develop a clinically applicable framework for DL-based detection of PCa from ultrasound that is robust to noise and uncertainty inherent to ultrasound images and their associated gold standard pathology labels. We begin by conducting a comprehensive evaluation of various DL baselines to assess their performance on noisy ultrasound data. We explore vision transformers for efficient feature extraction, followed by noise-resistant fine-tuning with multiple-instance learning (MIL). We then propose a novel multi-objective loss function that combines predictions from the extractor and the MIL aggregator, ensuring robust performance by maximizing agreement. Prostate ultrasound data is severely imbalanced in favor of benign cores. We propose TRUSWorthy, a unified framework that integrates self-supervised learning and MIL with an ensemble strategy that trains on diverse subsets of the majority class to enhance performance and uncertainty calibration. To our knowledge, this is the first framework to address noisy labels, data heterogeneity, and imbalance in PCa detection simultaneously. TRUSWorthy outperforms state-of-the-art methods in accuracy and uncertainty calibration. As model confidence increases, we also observe larger improvements over the baseline in both qualitative and quantitative performance. Our findings show that ensemble diversity is an effective strategy for reliable PCa detection and uncertainty estimation. They further suggest that an integrated solution has the potential to address both noise and uncertainty during clinical deployment.