Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation (2102.04525v4)

Abstract: Automatic segmentation methods are an important advancement in medical image analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most medical image segmentation tasks. Issues with class imbalance pose a significant challenge in medical datasets, with lesions often occupying a considerably smaller volume relative to the background. Loss functions used in the training of deep learning algorithms differ in their robustness to class imbalance, with direct consequences for model convergence. The most commonly used loss functions for segmentation are based on either the cross entropy loss, Dice loss or a combination of the two. We propose the Unified Focal loss, a new hierarchical framework that generalises Dice and cross entropy-based losses for handling class imbalance. We evaluate our proposed loss function on five publicly available, class imbalanced medical imaging datasets: CVC-ClinicDB, Digital Retinal Images for Vessel Extraction (DRIVE), Breast Ultrasound 2017 (BUS2017), Brain Tumour Segmentation 2020 (BraTS20) and Kidney Tumour Segmentation 2019 (KiTS19). We compare our loss function performance against six Dice or cross entropy-based loss functions, across 2D binary, 3D binary and 3D multiclass segmentation tasks, demonstrating that our proposed loss function is robust to class imbalance and consistently outperforms the other loss functions. Source code is available at: https://github.com/mlyg/unified-focal-loss

• The paper presents a novel hierarchical framework that unifies Dice and cross entropy-based losses to tackle class imbalance in medical segmentation.
• It introduces the Unified Focal loss which emphasizes hard-to-classify examples, significantly improving metrics such as DSC and IoU.
• Empirical evaluations across five diverse datasets demonstrate marked improvements in minority class segmentation, bolstering clinical utility.

An Overview of "Unified Focal Loss: Generalising Dice and Cross Entropy-Based Losses to Handle Class Imbalanced Medical Image Segmentation"

The paper "Unified Focal Loss: Generalising Dice and Cross Entropy-Based Losses to Handle Class Imbalanced Medical Image Segmentation" presents a novel approach to loss function design aimed at addressing the challenge of class imbalance in medical image segmentation. Given the prevalence of machine learning techniques, especially deep neural networks, in automatic segmentation tasks, the choice of loss function remains pivotal for achieving convergence and optimizing performance in the face of class imbalance. The authors propose the Unified Focal loss, a new loss function that amalgamates and extends the properties of Dice and cross entropy-based loss functions to better handle the issue of class imbalance in medical image datasets.

Key Contributions

The authors present several key contributions in this work:

1. Unified Hierarchical Framework: The paper introduces a hierarchical framework that unifies various distribution-based and region-based loss functions by leveraging a generalized formulation. This framework highlights the procedural relationship among different loss functions, offering a pathway to derive specialized loss models, such as the proposed Unified Focal loss.
2. Proposed Unified Focal Loss: The Unified Focal loss integrates the capabilities of the Dice and cross entropy-based losses, incorporating mechanisms like focal adjustments to handle both input and output imbalances. It effectively improves the segmentation performance across several challenging medical datasets by emphasizing hard-to-classify examples and allowing class imbalance adaptation.
3. Comprehensive Evaluation: The paper provides empirically rigorous evaluation across five diverse and publicly available medical imaging datasets—CVC-ClinicDB, DRIVE, BUS2017, BraTS20, and KiTS19—encompassing 2D binary, 3D binary, and 3D multiclass segmentation tasks. By doing so, they validate the robustness and general superiority of the Unified Focal loss over existing loss functions in handling datasets with varying degrees of class imbalance.

Implications and Performance Metrics

The evaluation results underscore the Unified Focal loss's capability to consistently outperform other mainstream loss functions, including the standard cross entropy, Dice, Tversky, and their respective focal variants. Metrics such as DSC and IoU are used to objectively quantify performance, demonstrating that the Unified Focal loss achieves higher validation scores across datasets. Notably, on the highly imbalanced KiTS19 dataset, Unified Focal loss shows marked improvements in segmenting minority classes like tumors.

The implications of these results are particularly relevant for clinical applications where precision and recall are crucial. Improved handling of class imbalances not only enhances the reliability of automated segmentation tools but also bolsters their potential for integration into clinical workflows, ultimately aiding in better diagnostic and therapeutic decision-making. From a theoretical perspective, the Unified Focal loss enriches the existing landscape by providing a more generalized and flexible loss function formulation, potentially catalyzing further advancements in model training paradigms.

Future Directions

Given the framework's adaptability and the promising results obtained, future research could explore extending this loss function to other types of segmentation challenges, such as those involving non-medical or natural scenes. Another avenue for exploration is the potential integration of the Unified Focal loss with state-of-the-art segmentation architectures beyond U-Net variants, thereby enhancing performance across more sophisticated systems. Automating the tuning of hyperparameters within this framework can also streamline its adoption and improve performance reproducibility.

In conclusion, the paper provides a substantial contribution to the domain of medical image segmentation by addressing the persistent challenge of class imbalance via a well-crafted and generalized loss function—Unified Focal loss. This approach not only simplifies the landscape of loss function selection but also significantly enhances segmentation outcomes, promising a notable impact on both current research efforts and future applications in medical imaging.