Style-KD: Class-imbalanced medical image classification via style knowledge distillation

Abstract

Problem: In the medical domain, obtaining training images on a large scale is difficult due to privacy and cost issues, and as further disease incidence varies widely, class imbalance has been a critical problem. Importantly, securing samples for rare diseases will always be difficult, but a diagnosis model should be able to predict the diseased samples accurately. When a model is trained on imbalanced data, it is expected that (1) the model only exhibits biased results for the majority classes, and (2) the recognition performance for minority classes will decline rapidly, resulting in low specificity. Objective: We aim to improve the prediction performance of such few sample classes using our auxiliary data from an efficient generation model and style knowledge distillation. In this paper, we introduce a style knowledge distillation method to address the class imbalance problem. Methods: The proposed method generates auxiliary images in which the style of a training sample is transferred. We train our model with balanced training samples by distilling the styles of auxiliary images. Conclusion: We demonstrate that it considerably increases the accuracy of classes with few samples, notably without adversely impacting classes with large samples, using the APTOS2019, and ISIC2018. This approach is straightforward but very effective for medical image datasets with skewed class distributions.