Adaptive Voxel-Weighted Loss Using L1 Norms in Deep Neural Networks for Detection and Segmentation of Prostate Cancer Lesions in PET/CT Images

TL;DR

This paper introduces L1-weighted Dice Focal Loss, a novel loss function that improves prostate cancer lesion detection in PET/CT images by balancing gradients and enhancing model performance across various architectures.

Contribution

We propose L1DFL, a new loss function that adaptively weights voxels based on difficulty, leading to better detection accuracy and balanced false positive/negative rates in prostate cancer segmentation.

Findings

L1DFL outperforms traditional loss functions with at least 4% higher Dice score.

F1 scores improved by 6% and 26% over Dice Loss and Dice Focal Loss.

L1DFL maintains robustness across lesion sizes and reduces false positives.

Abstract

Accurate automated detection of recurrent prostate cancer in PSMA PET/CT scans is challenging due to heterogeneous lesion size, activity, anatomical location, and intra- and inter-class imbalances. Conventional deep learning loss functions often produce suboptimal optimization, as gradients are dominated by easy background voxels or extreme outliers. To address this, we propose L1-weighted Dice Focal Loss (L1DFL), which harmonizes gradient magnitudes across voxels using L1 norms to adaptively weight samples based on classification difficulty, resulting in well-calibrated predictions with a bimodal separation between correct and incorrect predictions. We trained three 3D convolutional networks (Attention U-Net, SegResNet, U-Net) and a transformer-based UNETR model on 380 PSMA PET/CT scans. PET and CT volumes were concatenated as input to the models. We also fine-tuned SAM-Med3D foundation model with the different loss functions and evaluated their performance. Across architectures, L1DFL consistently outperformed Dice Loss (DL) and Dice Focal Loss (DFL), achieving at least a 4% improvement in Dice Similarity Coefficient. F1 scores were higher by 6% and 26% compared to DL and DFL, respectively. While DFL produced more false positives and DL struggled with larger lesions, L1DFL achieved balanced detection, minimizing false detections while maintaining high true positive rates. The gradient harmonization mechanism ensured robustness across varying lesion sizes, volumes, and spread. The code is publicly available at: https://github.com/ObedDzik/pca_segment.git.