Investigation of Federated Learning Algorithms for Retinal Optical Coherence Tomography Image Classification with Statistical Heterogeneity

TL;DR

Despite the effectiveness of federated learning in the utilization of private data across multiple medical institutions, the large number of clients and heterogeneous distribution of labels deteriorate the performance of both algorithms.

Abstract

Purpose: We apply federated learning to train an OCT image classifier simulating a realistic scenario with multiple clients and statistical heterogeneous data distribution where data in the clients lack samples of some categories entirely. Methods: We investigate the effectiveness of FedAvg and FedProx to train an OCT image classification model in a decentralized fashion, addressing privacy concerns associated with centralizing data. We partitioned a publicly available OCT dataset across multiple clients under IID and Non-IID settings and conducted local training on the subsets for each client. We evaluated two federated learning methods, FedAvg and FedProx for these settings. Results: Our experiments on the dataset suggest that under IID settings, both methods perform on par with training on a central data pool. However, the performance of both algorithms declines as we increase the statistical heterogeneity across the client data, while FedProx consistently performs better than FedAvg in the increased heterogeneity settings. Conclusion: Despite the effectiveness of federated learning in the utilization of private data across multiple medical institutions, the large number of clients and heterogeneous distribution of labels deteriorate the performance of both algorithms. Notably, FedProx appears to be more robust to the increased heterogeneity.

Figures (2)

• Figure 1: Illustration of FedAvg and FedProx training process. In each communication round, a subset of clients trained locally on their respective data splits with objective $L_{method}$ are selected for aggregation and used to update the global model. The updated global model parameters, $w_g$ are copied back to all the clients.
• Figure 2: (a) Label distribution of 4 randomly selected clients for IID and 2 SPC Non-IID settings. For the IID setting, the distribution of labels across clients is similar, while for 2 SPC, every client has multiple missing labels and distinct distributions (b) Image samples from the dataset