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

Zong Fan, Zhimin Wang, Chaojie Zhang, Muzaffer Özbey, Umberto Villa, Yao Hao, Zhongwei Zhang, Xiaowei Wang, Hua Li

Research output: Chapter in Book/Report/Conference proceedingConference contributionpeer-review


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.

Original languageEnglish
Title of host publicationMedical Imaging 2024
Subtitle of host publicationImage Processing
EditorsOlivier Colliot, Jhimli Mitra
ISBN (Electronic)9781510671560
StatePublished - 2024
EventMedical Imaging 2024: Image Processing - San Diego, United States
Duration: Feb 19 2024Feb 22 2024

Publication series

NameProgress in Biomedical Optics and Imaging - Proceedings of SPIE
ISSN (Print)1605-7422


ConferenceMedical Imaging 2024: Image Processing
Country/TerritoryUnited States
CitySan Diego


Dive into the research topics of 'Self-supervised learning based on StyleGAN for medical image classification on small labeled dataset'. Together they form a unique fingerprint.

Cite this