Self-supervised learning based on StyleGAN for medical image classification on small labeled dataset

Medical image classification plays a vital role in disease diagnosis, tumor staging, and various clinical applications. Deep learning (DL) methods have become increasingly popular for medical image classification. However, medical images have unique characteristics that pose challenges for training DL-based models, including limited annotated data, imbalanced distribution of classes, and large variations in lesion structures. Self-supervised learning (SSL) methods have emerged as a promising solution to alleviate these issues through directly learning useful representations from large-scale unlabeled data. In this study, a new generative self-supervised learning method based on the StyleGAN generator is proposed for medical image classification. The style generator, pre-trained on large-scale unlabeled data, is integrated into the classification framework to effectively extract style features that encapsulate essential semantic information from input images through image reconstruction. The extracted style feature serves as an auxiliary regularization term to leverage knowledge learned from unlabeled data to support the training of the classification network and enhance model performance. To enable efficient feature fusion, a self-attention module is designed for this integration of the style generator and classification framework, dynamically focusing on important feature elements related to classification performance. Additionally, a sequential training strategy is designed to train the classification model on a limited number of labeled images while leveraging large-scale unlabeled data to improve classification performance. The experimental results on a chest X-ray image dataset demonstrate superior classification performance and robustness compared to traditional DL-based methods. The effectiveness and potential of the model were discussed as well.