Author: Weiguo Lu, Hua-Chieh Shao, Guoping Xu, You Zhang 👨🔬
Affiliation: Medical Artificial Intelligence and Automation (MAIA) Lab, Department of Radiation Oncology, UT Southwestern Medical Center, Medical Artificial Intelligence and Automation (MAIA) Laboratory, Department of Radiation Oncology, UT Southwestern Medical Center 🌍
Purpose:
Neural network-based lesion segmentation remains a significant challenge due to the low contrast between lesions and surrounding tissues (high ambiguity) and the variability of lesion shapes. We propose a Margin-based Focal distance (MbFd) loss specifically for lesion segmentation performance enhancement, by controlling the logit map distance between the lesion and the background.
Methods:
MbFd tackles lesion segmentation by integrating two complementary components: a bidirectional margin-based logit distance regulator and an intra-class distance regularizer. The bidirectional regulator manages lesion and background classes by computing their absolute logit map distances before the sigmoid or softmax activation layers. A predefined margin is applied to penalize both large deviations (to reduce over-confident predictions and facilitate network optimization) and small deviations (to mitigate ambiguities), enabling controlled refinement of the logit distribution. Meanwhile, the intra-class distance regularizer aims to increase the group-wise logit map distance between lesion and background pixels to enhance inter-class separation and accelerate training. It penalizes misclassified pixels and incorporates a distance-based weighting scheme to emphasize class-specific segmentation for lesion and background logit maps. Both components of MbFd re-calibrate the logit distribution, improving segmentation accuracy and robustness.
Results:
MbFd was evaluated on two public lesion datasets: the 2D BUSI ultrasound dataset for breast lesion segmentation and the 3D BraTS MRI dataset for brain tumor segmentation. MbFd demonstrated consistent improvements over the standard cross-entropy loss. For BUSI, the Dice-Similarity-Coefficient (DSC) improved by approximately 1-3.5% for segmentation backbones including nnUNet, SegResNet, and SwinUNETR. For BraTS, MbFd enhanced DSC by about 0.2-1% for the same three segmentation frameworks. In addition, MbFd reduces the expected-calibration-error (ECE), improving the probability distribution matching between the predicted output and the ‘ground truth’.
Conclusion:
MbFd effectively re-calibrates the logit map distance distribution between lesion and background regions, which enhances the accuracy of segmentation networks particularly in challenging lesion segmentation scenarios.