Source-linked AI summary
Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations
Carole H Sudre, Wenqi Li, Tom Vercauteren, Sébastien Ourselin, M. Jorge Cardoso
TL;DR
Severely imbalanced medical-image segmentation can make loss-function training unstable, and existing strategies vary in their sensitivity to learning-rate choices. The paper evaluates these losses across 2D and 3D tasks and proposes Generalized Dice Loss (GDL), which showed the strongest reliability across setups, including a 0.02% foreground-background ratio.
Problem
Severely unbalanced segmentation problems can destabilize segmentation frameworks, motivating evaluation of loss functions for rare medical-image structures.
Method
The paper compares four loss functions across learning rates, sampling strategies, and 2D and 3D networks, while proposing Generalized Dice Loss for deep-network training.
Results
Across two tasks and networks, overlap-based losses were more robust under greater imbalance, with GDLv showing the strongest reliability across setups.
Takeaways & Limitations
Loss-function choice is crucial for deep-learning segmentation of highly unbalanced problems, including the study's 0.02% foreground-background 3D case.
Abstract
from arXiv · showhide
Deep-learning has proved in recent years to be a powerful tool for image analysis and is now widely used to segment both 2D and 3D medical images. Deep-learning segmentation frameworks rely not only on the choice of network architecture but also on the choice of loss function. When the segmentation process targets rare observations, a severe class imbalance is likely to occur between candidate labels, thus resulting in sub-optimal performance. In order to mitigate this issue, strategies such as the weighted cross-entropy function, the sensitivity function or the Dice loss function, have been proposed. In this work, we investigate the behavior of these loss functions and their sensitivity to learning rate tuning in the presence of different rates of label imbalance across 2D and 3D segmentation tasks. We also propose to use the class re-balancing properties of the Generalized Dice overlap, a known metric for segmentation assessment, as a robust and accurate deep-learning loss function for unbalanced tasks.
1 Introduction
Medical-image segmentation often targets pathological regions occupying a very small fraction of an image, creating severe class imbalance. This work evaluates loss functions across 2D and 3D multi-class tasks and varying learning rates and sampling strategies.
- Rare pathological regions can make segmentation frameworks unstable because they occupy a very small fraction of medical images.
- Existing imbalance strategies either modify training-sample selection or use more robust loss functions.
- The study compares three published loss functions across multi-class 2D and 3D segmentation problems.
- The study assesses loss-function robustness to learning rate and sample rate.
2 Methods
The methods compare imbalance-aware loss functions under binary foreground-versus-background formulations and across representative 2D and 3D networks. The proposed Generalized Dice Loss uses class rebalancing through label-dependent weighting.
- All compared loss functions are analyzed using binary foreground-versus-background segmentation to quantify class imbalance.The authors note that a one-class formulation would mitigate imbalance but generalize less easily to multiple classes.
- Weighted cross-entropy assigns a weight to the foreground class and extends to more than two classes.
- Dice loss uses the Dice overlap score as a training objective, with epsilon preventing division-by-zero when reference and prediction sets are empty.
- Sensitivity-Specificity loss balances sensitivity and specificity with λ set to 0.05, while epsilon handles empty-set divisions.
- Generalized Dice Loss introduces a segmentation-overlap metric into discriminative deep-network training and weights labels to provide invariance to label-set properties.
- The experiments reimplement four representative networks: UNet and TwoPathCNN for 2D, plus DeepMedic and HighResNet for 3D.
3 Experiments and Results
The experiments evaluate imbalance-sensitive losses across two 2D and two 3D networks, using tumor and white-matter-lesion segmentation under varied patch sizes and learning rates. GDL showed greater robustness across hyperparameters and experiments, while intermediate learning rates and smaller patches with larger batches generally performed better.
- Tasks and datasets: The study tests tumor segmentation on BRATS and age-related white matter hyperintensity segmentation, covering pathology with varying locations, shapes, and sizes.
- 2D results: In 2D, WCE and DL2 were less able to cope with the fast learning rate 10^-3 in TwoPathCNN, whereas SS efficiency was more network dependent.
- 2D results: 10^-4 generally provided the best training, while smaller patches with larger batches produced stronger performance across sampling strategies.
- 3D results: In 3D, WCE was unable to train under higher imbalance, SS dropped relative to GDLv, and DL2 failed at higher learning rates.Learning rate 10^-5 did not provide a loss plateau after 3000 iterations, while smaller patches performed better under LR=10^-4.
- 3D results: Across the 3D test experiments, GDLv was more robust than the other losses, with only small relative-performance variations for less unbalanced samples.
4 Discussion
Across tasks and networks, overlap-based losses became more robust as class imbalance increased, with GDLv showing the strongest reliability across setups. The study emphasizes that loss-function choice is crucial for highly unbalanced segmentation.
- When imbalance increased, losses based on overlap measures appeared more robust than other strategies.
- In the 3D example, GDLv increased the ability to capture punctuate lesions.
- GDLv showed the strongest reliability across the tested setups.
- The most unbalanced study case had a foreground-background ratio of 0.02% in the 3D white matter-lesion experiment.
- Future work will examine more extreme imbalance, including lacunes and perivascular spaces with a reported ratio of 1/100000.