Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations (1707.03237v3)

Abstract: 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.

• The paper proposes the Generalised Dice Loss to effectively balance class disparities in medical image segmentation.
• It compares multiple loss functions across 2D and 3D networks, showing that Generalised Dice Loss provides stable performance under aggressive training conditions.
• The study highlights the practical benefit of volume-based weighting to enhance detection of small, rare pathological regions.

Generalised Dice Overlap as a Deep Learning Loss Function for Highly Unbalanced Segmentations

The paper "Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations" by Carole H. Sudre et al. addresses the fundamental challenge in medical image segmentation of handling class imbalance, particularly when segmenting rare pathological regions such as brain tumors or white matter lesions.

The paper explores the efficacy of several loss functions for deep learning-based image segmentation under conditions of severe class imbalance. Specifically, it evaluates the behavior and robustness of the Weighted Cross-Entropy (WCE), Dice Loss (DL), Sensitivity-Specificity (SS), and a proposed Generalized Dice Loss (GDL). These loss functions are tested in both 2D and 3D contexts using four neural network architectures: UNet, TwoPathCNN, DeepMedic, and HighResNet.

Methods and Experiments

Loss Functions for Unbalanced Data

1. Weighted Cross-Entropy (WCE): This loss function introduces weights to the standard cross-entropy formulation to account for class imbalance. The weight for the foreground class is inversely proportional to the number of foreground pixels.
2. Dice Loss (DL): An overlap-based loss function that directly optimizes the Dice Score Coefficient, which is a common evaluation metric for segmentation performance.
3. Sensitivity-Specificity (SS): This function combines sensitivity and specificity, aiming to balance the trade-off between these two metrics.
4. Generalized Dice Loss (GDL): The paper proposes using the Generalized Dice Score, a metric originally designed for multi-class segmentation, as a loss function. This metric incorporates volume-based weighting to address the imbalance between different label sets effectively.

Network Architectures

The experiments utilized four established neural network architectures in both 2D and 3D for segmentation tasks:

• UNet and TwoPathCNN for 2D images.
• DeepMedic and HighResNet for 3D volumes.

Datasets and Experimental Setup

Two segmentation tasks were performed:

1. Tumor Segmentation: Applied using the BRATS dataset.
2. White Matter Hyperintensity Segmentation: Utilized an in-house dataset of 524 subjects.

Various learning rates and patch sampling strategies were tested to evaluate the robustness and performance of the loss functions. The experiments were designed to observe the sensitivity of these functions to training hyperparameters, particularly under highly unbalanced conditions.

Results and Findings

2D Results

Across different training conditions, the Generalized Dice Loss (GDL) exhibited superior robustness to learning rates and sample sizes compared to WCE, DL, and SS. Specifically, GDL maintained stable performance even with aggressive learning rates and larger patch sizes, which typically exacerbate imbalance issues.

3D Results

In the 3D context, the imbalance was more pronounced. GDL again demonstrated the best overall performance, where other loss functions, such as WCE and SS, either failed to train effectively or yielded significantly poorer segmentation accuracy. Notably, GDL provided consistent Dice scores across a wide range of learning rates and sampling strategies, indicating its utility in highly unbalanced conditions typical of 3D medical images.

Practical and Theoretical Implications

Practical Implications:

• The robustness of GDL makes it a compelling choice for medical image analysis tasks where imbalanced class distributions are prevalent.
• The integration of volume-based weighting in GDL enhances its performance in detecting small pathological regions, as evidenced by its success in capturing punctate lesions.

Theoretical Implications:

• The paper reinforces the importance of choosing an appropriate loss function tailored to the specific characteristics of the dataset and the type of imbalance encountered.
• The results suggest that overlap-based loss functions, particularly those incorporating class rebalancing mechanisms like GDL, provide a more versatile approach for segmentation tasks across different imaging modalities and dimensions.

Future Developments

Future research could delve into even more extreme cases of class imbalance, such as the detection of microscopic anomalies in medical images. Additionally, exploring the combination of GDL with other advanced training techniques, such as data augmentation and semi-supervised learning, could further enhance segmentation performance.

In conclusion, the proposed Generalized Dice Loss demonstrates significant promise for addressing class imbalance in deep learning-based medical image segmentation, providing a robust method that can be effectively applied across various segmentation tasks and network architectures.