Batch-balanced focal loss: a hybrid solution to class imbalance in deep learning

Abstract

Purpose: To validate the effectiveness of an approach called batch-balanced focal loss (BBFL) in enhancing convolutional neural network (CNN) classification performance on imbalanced datasets.

Materials and methods: BBFL combines two strategies to tackle class imbalance: (1) batch-balancing to equalize model learning of class samples and (2) focal loss to add hard-sample importance to the learning gradient. BBFL was validated on two imbalanced fundus image datasets: a binary retinal nerve fiber layer defect (RNFLD) dataset ($n=\mathrm{7,258}$) and a multiclass glaucoma dataset ($n=\mathrm{7,873}$). BBFL was compared to several imbalanced learning techniques, including random oversampling (ROS), cost-sensitive learning, and thresholding, based on three state-of-the-art CNNs. Accuracy, F1-score, and the area under the receiver operator characteristic curve (AUC) were used as the performance metrics for binary classification. Mean accuracy and mean F1-score were used for multiclass classification. Confusion matrices, t-distributed neighbor embedding plots, and GradCAM were used for the visual assessment of performance.

Results: In binary classification of RNFLD, BBFL with InceptionV3 (93.0% accuracy, 84.7% F1, 0.971 AUC) outperformed ROS (92.6% accuracy, 83.7% F1, 0.964 AUC), cost-sensitive learning (92.5% accuracy, 83.8% F1, 0.962 AUC), and thresholding (91.9% accuracy, 83.0% F1, 0.962 AUC) and others. In multiclass classification of glaucoma, BBFL with MobileNetV2 (79.7% accuracy, 69.6% average F1 score) outperformed ROS (76.8% accuracy, 64.7% F1), cost-sensitive learning (78.3% accuracy, 67.8.8% F1), and random undersampling (76.5% accuracy, 66.5% F1).

Conclusion: The BBFL-based learning method can improve the performance of a CNN model in both binary and multiclass disease classification when the data are imbalanced.

Keywords: class imbalance; convolutional neural networks; deep learning; retina fundus.

Fig. 1

Workflow of the BBFL algorithm.

Fig. 2

Focal loss has a modulating factor $(1-{\hat{p}}_i{)}^{\gamma }$ that decreases the loss contribution of well-classified samples. Increasing the $\gamma$ parameter focuses less on samples that were predicted correctly with high probability. Focal loss can be applied to both binary and categorical crosss entropy loss functions.

Fig. 3

ROC plots (left) and PRC plots (right) comparing the class imbalance techniques for InceptionV3 (top), MobileNetV2 (middle), and ResNet50 (bottom) to discriminate positive and negative RNFLD cases. Both plot types were generated using the raw RNFLD predictions from the highest performing model from the five experiments.

Fig. 4

GradCAM visualizations of four positive RNFLD samples. In each $2\times 2$ grid, the top two CFP images were selected as “easy” and the bottom two as “difficult” in RNFLD interpretability. ROS and BBFL algorithms visually produce the most accurate heatmaps with the least heat variance. ROS, random oversampling.

Fig. 5

t-SNE plots using output from final GAP layer of InceptionV3 for RNFLD detection: (a) BBFL, (b) cost-sensitive, (c) random oversampling, (d) focal loss, (e) balanced-batches, and (f) baseline t-SNE.

Fig. 6

Confusion matrix comparing imbalance learning techniques to BBFL using the ResNet50 network ($n=788$). Results from multiclass glaucoma classification with fundus images graded as normal (0), mild (1), moderate (2), and severe (3).

Fig. 7

Confusion matrix comparing imbalance learning techniques to BBFL using the MobileNetV2 network ($n=788$). Results from multiclass glaucoma classification with fundus images graded as normal (0), mild (1), moderate (2), and severe (3).

Fig. 8

Confusion matrix comparing imbalance learning techniques to BBFL using the InceptionV3 network ($n=788$). Results from multiclass glaucoma classification with fundus images graded as normal (0), mild (1), moderate (2), and severe (3).