No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data

A central challenge in training classification models in the real-world system is learning with non-IID data. To cope with this, most of the works involve enforcing regularization in local optimization or the model aggregation scheme at the server. Other works also share datasets or synthesized samples to supplement the training of-represented classes or introduce a certain level of personalization. effective, they lack a deep understanding of how the data heterogeneity each layer of a deep classification model. In this paper, we bridge gap by performing an experimental analysis of the representations learned different layers. Our observations are surprising: (1) there exists a bias in the classifier than other layers, and (2) the classification can be significantly improved by post-calibrating the classifier federated training. Motivated by the above findings, we propose a novel simple algorithm called Classifier Calibration with Virtual Representations (CCVR), which adjusts the classifier using virtual representations sampled from approximated gaussian mixture model. Experimental results demonstrate that achieves state-of-the-art performance on popular federated learning including CIFAR-10, CIFAR-100, and CINIC-10. We hope that our simple effective method can shed some light on the future research of federated with non-IID data.