Zhantao Cao, Yuanbing Shi, Shuli Zhang, Huanan Chen, Weide Liu, Guanghui Yue, Huazhen Lin
BACKGROUND: Recently, deep convolutional neural networks (CNNs) have shown great potential in medical image classification tasks. However, the practical usage of the methods is constrained by two challenges: 1) the challenge of using nonindependent and identically distributed (non-IID) datasets from various medical institutions while ensuring privacy, and 2) the data imbalance problem due to the frequency of different diseases.
PURPOSE: The objective of this paper is to present a novel approach for addressing these challenges through a decentralized learning method using a prototypical contrastive network to achieve precise medical image classification while mitigating the non-IID problem across different clients.
METHODS: We propose a prototype contrastive network that minimizes disparities among heterogeneous clients. This network utilizes an approximate global prototype to alleviate the non-IID dataset problem for each local client by projecting data onto a balanced prototype space. To validate the effectiveness of our algorithm, we employed three distinct datasets of color fundus photographs for diabetic retinopathy: the EyePACS, APTOS, and IDRiD datasets. During training, we incorporated 35k images from EyePACS, 3662 from APTOS, and 516 from IDRiD. For testing, we used 53k images from EyePACS. Additionally, we included the COVIDx dataset of chest X-rays for comparative analysis, comprising 29 986 training images and 400 test samples.
RESULTS: In this study, we conducted comprehensive comparisons with existing works using four medical image datasets. Specifically, on the EyePACS dataset under the balanced IID setting, our method outperformed the FedAvg baseline by 3.7% in accuracy. In the Dirichlet non-IID setting, which presents an extremely unbalanced distribution, our method showed a notable 6.6% enhancement in accuracy over FedAvg. Similarly, on the APTOS dataset, our method achieved a 3.7% improvement in accuracy over FedAvg under the balanced IID setting and a 5.0% improvement under the Dirichlet non-IID setting. Notably, on the DCC non-IID and COVID-19 datasets, our method established a new state-of-the-art across all evaluation metrics, including WAccuracy, WPrecision, WRecall, and WF-score.
CONCLUSIONS: Our proposed prototypical contrastive loss guides the local client's data distribution to align with the global distribution. Additionally, our method uses an approximate global prototype to address unbalanced dataset distribution across local clients by projecting all data onto a new balanced prototype space. Our model achieves state-of-the-art performance on the EyePACS, APTOS, IDRiD, and COVIDx datasets.
