Generative model based federated learning on non-IID data

Abstract

Federated Learning (FL) is a training strategy of a deep learning model for distributed data when the data cannot be centralized for privacy reasons. To share the knowledge that can be learned from local datasets, a server regularly aggregates parameters of models which are independently trained on each edge device (i.e. client). Among the parameter aggregation methods, FedAvg uses weighted summation, where weights are the ratio of the number of training datasets. FedAvg has been shown to work for classification with independently and identically distributed (IID) data, but not for non-IID data. In this paper, we propose an approach using the generative model to address the FL performance degradation in non-IID data. In particular, we first train the Conditional Generative Adversarial Network (cGAN) to estimate the mixture distribution of non-IID data. For better conditioning of the cGAN, we propose a new regularization term based on a contrastive loss. After training of cGAN, each client augments the local dataset using the trained cGAN and performs FedAvg on augmented local datasets. Here, we propose the data augmentation strategy based on the Bernoulli sampling. Experiments with the extreme and realistic non-IID settings show that augmented datasets can improve the performance of FL on non-IID data. Our method achieves an accuracy (99.1%) comparable to that of federated learning for the IID MNIST dataset, even though each client contains samples from only one label.