Imbalance-aware loss functions improve medical image classification

Deep learning models offer unprecedented opportunities for diagnosis, prognosis, and treatment planning. However, conventional deep learning pipelines often encounter challenges in learning unbiased classifiers within imbalanced data settings, frequently exhibiting bias towards minority classes. In this study, we aim to improve medical image classification by effectively addressing class imbalance. To this end, we employ differentiable loss functions derived from classification metrics commonly used in imbalanced data settings: Matthews correlation coefficient (MCC) and the F1 score. We explore the efficacy of these loss functions both independently and in combination with cross-entropy loss and various batch sampling strategies on diverse medical datasets of 2D fundoscopy and 3D magnetic resonance images. Our findings demonstrate that, compared to conventional loss functions, we achieve notable improvements in overall classification performance, with increases of up to +12% in balanced accuracy and up to +51% in class-wise F1 score for minority classes when utilizing cross-entropy coupled with metrics-derived loss. Additionally, we conduct feature visualization to gain insights into the behavior of these features during training with imbalance-aware loss functions. Our visualization reveals a more pronounced clustering of minority classes in the feature space, consistent with our classification results. Our results underscore the effectiveness of combining cross-entropy loss with class-imbalance-aware loss functions in training more accurate classifiers, particularly for minority classes.