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

Abstract

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