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

Michael Yeung; Evis Sala; Carola-Bibiane Schönlieb; Leonardo Rundo

Source-linked AI summary

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

Michael Yeung, Evis Sala, Carola-Bibiane Schönlieb, Leonardo Rundo

arXiv:2102.04525v4Feb 8, 2021eess.IVcs.CVcs.LG

TL;DR

The paper examines how relationships among numerous loss functions can guide segmentation heuristics. It proposes a hierarchical framework and reports that Unified Focal loss achieved the highest DSC and IoU scores across five datasets, with study comparisons restricted to a subset of the framework.

Problem
The study focuses on the large number of available loss functions and the need to understand their relationships for segmentation heuristics.
Method
The paper proposes a new hierarchical framework encompassing segmentation loss functions, with optimisation simplified to tuning a single γ.
Results
Unified Focal loss consistently achieved the highest DSC and IoU scores across the five datasets.
Takeaways & Limitations
Unified Focal loss provides consistently strong segmentation performance across the evaluated datasets.
Takeaways & Limitations
The framework and comparisons were restricted to only a subset of the available options.

Abstract

from arXiv · show

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

1. Introduction

Medical image segmentation must address severe foreground–background imbalance, especially for organs and tumours, while loss-function choice remains insufficiently informed. The paper proposes a unified framework and evaluates it across diverse imbalanced datasets and segmentation settings.

Motivation: Class imbalance is a dominant medical segmentation challenge because foreground structures may be much smaller than the surrounding background.The problem is especially prevalent for tumours, which are often smaller than their organs of origin.
Motivation: False positives and false negatives both matter clinically because they can respectively increase treatment margins or cause inadequate treatment and incomplete resection.The cited consequences arise in image-guided interventions.
Prior work: Existing comparisons disagree about which region-based or compound loss is best, with class-imbalance severity identified as a major confounding factor.Prior studies reported different preferred losses across datasets with varying imbalance.
Motivation: Loss-function selection is consequential but not widespread in practice, motivating large-scale comparisons and heuristics for imbalanced settings.The paper links loss functions to model optimisation and seeks to inform their choice.
Contributions: The study extends prior comparisons using five class-imbalanced datasets spanning multiple imaging modalities and proposes a hierarchical classification of Dice- and cross-entropy-based losses.The Unified Focal loss is derived from this classification to generalise both loss families.
Contributions: The proposed loss consistently improves segmentation quality over six related losses, with a better recall–precision balance and robustness to class imbalance.These are stated as the paper’s principal contributions and reported outcome.

2. Background

The paper frames semantic segmentation as pixel-level classification trained through loss-driven optimisation. It introduces distribution-based, region-based, and compound loss families, then details cross entropy as a pixel-wise objective whose imbalance can disadvantage smaller objects.

Segmentation and losses: Semantic segmentation performs pixel-level classification directly, unlike instance segmentation, which requires an additional object-detection stage.The paper focuses on semantic segmentation.
Cross entropy loss: Categorical cross entropy uses one-hot ground-truth labels and predicted class probabilities summed over classes and pixels.The paper extends cross entropy to multiclass segmentation using class and pixel indices.
Cross entropy loss: Pixel-wise cross entropy can over-represent larger objects under class imbalance, producing poorer segmentation of smaller objects.This is the stated imbalance-related limitation of the standard objective.

2.2. Focal loss

Focal loss modifies binary cross entropy to address class imbalance by reducing the contribution of easy examples. Its parameters control class weighting and the strength of this down-weighting, with γ = 0 recovering binary cross entropy.

Focal loss: Focal loss addresses standard cross-entropy imbalance by down-weighting easy examples so harder examples contribute more to learning.It is presented as a variant of binary cross entropy.
Focal loss: The Focal loss adds a modulating factor to binary cross entropy and is parameterised by α and γ.α controls class weights, while γ controls the degree of down-weighting of easy-to-classify pixels.
Focal loss: When γ = 0, the Focal loss simplifies to binary cross entropy.This provides a direct special-case relationship between the two losses.
Multiclass extension: For multiclass segmentation, categorical Focal loss uses a vector of class weights, ground-truth class probabilities, and categorical cross entropy.The multiclass formulation is defined using class and pixel-wise quantities.

2.3. Dice loss

Dice loss is derived from the Dice similarity coefficient and is naturally suited to segmentation overlap. Although it helps with class imbalance, its gradient can become unstable when denominators are small.

Dice loss: The Dice similarity coefficient is the segmentation-overlap metric from which Dice loss is defined.The coefficient is expressed using voxel-level true positives, false positives, and false negatives.
Dice variants: Generalised Dice corrects class weights by inverse volume, while Generalised Wasserstein Dice combines the Wasserstein metric with Dice loss.These are presented as variants of the basic Dice loss.
Dice loss: Dice loss is adapted to class imbalance even in its simplest formulation.Its region-based formulation directly relates predictions to segmentation overlap.
Dice loss: Dice-gradient calculations can be unstable for highly imbalanced data because they involve small denominators.This is identified as an inherent limitation of the Dice loss gradient.

2.4. Tversky loss

The Tversky loss generalises Dice-based similarity by separately weighting false positives and false negatives, supporting class-imbalanced segmentation. Its parameters can be selected to emphasize recall and balance precision with recall.

The formulation represents foreground and background probabilities and ground-truth indicators for each pixel.
The Tversky loss is defined using the Tversky index for C classes.
When Dice loss is applied to class-imbalanced problems, resulting segmentations often have high precision but low recall.
Assigning greater weight to false negatives improves recall and balances precision and recall.
β is often set higher than α, with β = 0.7 and α = 0.3 commonly used.
The asymmetric similarity loss becomes functionally equivalent to Tversky loss when α and β sum to 1.

2.5. Focal Tversky loss

Focal Tversky loss adapts Tversky loss by focusing on harder examples. It reduces to Tversky loss at γ = 1, while the reported optimum γ = 4/3 instead enhances easy-example loss.

Focal Tversky loss applies a focal adaptation to the Tversky index.
γ < 1 increases the degree of focusing on harder examples.
At γ = 1, Focal Tversky loss simplifies to Tversky loss.
γ = 4/3 was reported as optimal and enhances rather than suppresses the loss of easy examples.
Near the end of training, confident classifications make the Tversky index approach 1, so the loss enhancement affects this regime.

2.6. Combo loss

Combo loss is a compound loss that combines Dice similarity with modified cross entropy. Its parameters control the relative component weighting and the penalties assigned to false-positive and false-negative predictions.

The paper uses the definition combining cross entropy with DSC, consistent with Combo loss and the KiTS19 state-of-the-art loss.
Combo loss minimises multiple loss functions together as a weighted sum of Dice similarity and modified cross entropy.
α ∈ [0, 1] controls the relative contributions of the Dice and cross entropy terms.
β controls the relative weights assigned to false positives and false negatives.
β > 0.5 penalises false-negative predictions more heavily.
The terminology “Dice and cross entropy loss” has been used for multiple related compound-loss definitions.

2.7. Hybrid Focal loss

Hybrid Focal loss extends compound Dice–cross-entropy approaches with tunable mechanisms for output and input imbalance. It replaces the Dice and cross-entropy components with Focal Tversky and Focal loss, respectively.

Combo loss and DiceFocal loss inherit benefits from Dice and cross entropy but do not exploit their full benefits for class imbalance.
Their cross-entropy components are partially robust to output imbalance, but their Dice components lack equivalent positive-negative weighting.
The Dice components are not adapted to handle input imbalance.
Hybrid Focal loss incorporates tunable parameters for output imbalance and focal parameters for input imbalance.
It replaces Dice loss with Focal Tversky loss and cross entropy loss with Focal loss.
λ ∈ [0, 1] determines the relative weighting of the two component losses.

2.8. Unified Focal loss

The Unified Focal loss unifies Dice- and cross entropy-based losses while addressing focal-loss suppression and hyperparameter complexity in class-imbalanced segmentation.

Motivation: Hybrid Focal loss has six hyperparameters to tune, creating a larger search space that can affect convergence near the end of training.The parameters include α and γ from Focal loss, α/β and γ from Focal Tversky loss, and λ for component weighting.
Unified framework: The Unified Focal loss groups functionally equivalent hyperparameters, reducing the Hybrid Focal loss from six parameters to three.δ controls positive-versus-negative weighting, γ controls background suppression and rare-class enhancement, and λ weights the component losses.
Asymmetric variants: Asymmetric focal weighting suppresses background loss while preserving or enhancing the rare-class contribution.The modified Focal loss removes focal weighting from rare-class terms, whereas modified Focal Tversky loss removes it from background terms.
Component interaction: Complementary pairing of Focal and Focal Tversky losses enables simultaneous background suppression and foreground enhancement.The authors describe this effect as analogous to increasing the signal-to-noise ratio.
Generalisation: All described Dice- and cross entropy-based losses are special cases of the Unified Focal loss, including Dice and cross entropy at γ = 0 and δ = 0.5.Setting λ = 0 recovers Dice loss, while λ = 1 recovers cross entropy loss.

3. Materials and methods

The experiments evaluate Unified Focal loss variants on five class-imbalanced medical imaging datasets spanning 2D binary, 3D binary, and 3D multiclass segmentation.

Datasets: Five datasets—CVC-ClinicDB, DRIVE, BUS2017, BraTS20, and KiTS19—were selected to assess segmentation under class imbalance.Foreground proportions were calculated per image and averaged across each dataset.
Evaluation: Accuracy was assessed using DSC, IoU, recall, and precision, while models used 2D or 3D U-Net architectures according to task dimensionality.The evaluated losses included cross entropy, Focal, Dice, Tversky, Focal Tversky, Combo, and Unified Focal variants.
Segmentation tasks: The study covers 2D binary segmentation with CVC-ClinicDB, DRIVE, and BUS2017, 3D binary segmentation with BraTS20, and 3D multiclass segmentation with KiTS19.The 3D tasks use patch-wise analysis, while the listed 2D tasks use full-image analysis.

4. Experimental results

Experiments span 2D binary, 3D binary, and 3D multiclass segmentation across five class-imbalanced medical imaging datasets. The Unified Focal loss generally achieved the strongest performance, with robust behavior across γ values and improved qualitative segmentations.

The experiments cover 2D binary segmentation on CVC-ClinicDB, DRIVE, and BUS2017, followed by 3D binary BraTS20 and 3D multiclass KiTS19 evaluations.
2D binary segmentation: Across the three 2D datasets, the Unified Focal loss variants showed no statistically significant performance difference from each other.
2D binary segmentation: Cross entropy-based losses generally performed worst, while the asymmetric Unified Focal loss significantly outperformed other losses across CVC-ClinicDB, DRIVE, and BUS2017.
2D binary segmentation: The Unified Focal losses maintained strong DSC performance across γ ∈ [0.1, 0.9], with γ = 0.5 providing robust performance and γ = 0.9 performing worst.
2D binary segmentation: Compound losses produced the highest-quality example segmentations, with the Unified Focal loss showing the strongest visual results and cross entropy producing more false negatives.
3D multiclass segmentation: Kidney tumour segmentations varied considerably by loss function, with the Unified Focal loss producing the visibly most accurate tumour contour.

5. Discussion and conclusions

The study proposes the Unified Focal loss as a hierarchical framework generalising Dice- and cross entropy-based losses, and evaluates it across five imbalanced medical segmentation datasets and three task settings. It consistently achieved the highest DSC and IoU scores, while exposing limitations in the compared loss-function scope and multiclass γ optimisation.

Contribution: The Unified Focal loss unifies Dice-based and cross entropy-based loss functions within a hierarchical framework for class-imbalanced segmentation.
Evaluation: The evaluation compared Unified Focal loss with six other losses across five class-imbalanced datasets involving 2D binary, 3D binary, and 3D multiclass segmentation.
Results: The Unified Focal loss consistently achieved the highest DSC and IoU scores across the five datasets, with slightly better performance from its asymmetric variant.
Results: The γ hyperparameter can be the single tuning parameter for Unified Focal loss, and its observed stability made optimisation easy in these experiments.
Results: Loss-function choice substantially affected performance, including poor kidney-tumour segmentation by distribution-based losses on the highly imbalanced KiTS19 dataset.
Discussion: Compound losses performed consistently well, combining cross entropy-based gradient stability with Dice-based robustness to class imbalance; Unified Focal segmentations also had the highest qualitative quality.
Discussion: No difference in training time was observed between the evaluated loss functions.
Limitations: The study restricted comparisons to selected popular Dice- and cross entropy-based losses, excluding boundary-based losses and other variants.