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

1 Introduction . The rapid advances in deep learning have benefited a lot from large datasets like [ 1 ] . However , in the real world , data may be distributed on numerous mobile devices and the Internet of Things ( IoT ) , requiring decentralized training of deep networks . Driven by such realistic needs , federated learning [ 2 , 3 , 4 ] has become an emerging research topic where the model training is pushed to a large number of edge clients and the raw data never leave local devices . A notorious trap in federated learning is training with non-IID data . Due to diverse user behaviors , large heterogeneity may be present in different clients ’ local data , which has been found to result in unstable and slow convergence [ 5 ] and cause suboptimal or even detrimental model performance [ 6 , 7 ] . There have been a plethora of works exploring promising solutions to federated learning on non-IID data . They can be roughly divided into four categories : 1 ) client drift mitigation [ 5 , 8 , 9 , 10 ] , which modifies the local objectives of the clients , so that the local model is consistent with the global model to a certain degree ; 2 ) aggregation scheme [ 11 , 12 , 13 , 14 , 15 ] , which improves the model fusion mechanism at the server ; 3 ) data sharing [ 6 , 16 , 17 , 18 ] , which introduces public datasets or synthesized data to help construct a more balanced data distribution on the client or on the server ; ∗corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 4 ) personalized federated learning [ 19 , 20 , 21 , 22 ] , which aims to train personalized models for individual clients rather than a shared global model . However , as suggested by [ 7 ] , existing algorithms are still unable to achieve good performance on image datasets with deep learning models , and could be no better than vanilla FedAvg [ 2 ] . To identify the reasons behind this , we perform a thorough experimental investigation on each layer of a deep neural network . Specifically , we measure the Centered Kernel Alignment ( CKA ) [ 23 ] similarity between the representations from the same layer of different clients ’ local models . The observation is thought-provoking : comparing different layers learned on different clients , the classifier has the lowest feature2 similarity across different local models . Motivated by the above discovery , we dig deeper to study the variation of the weight of the classifier in federated optimization , and confirm that the classifier tends to be biased to certain classes . After identifying this devil , we conduct several empirical trials to debias the classifier via regularizing the classifier during training or calibrating classifier weights after training . We surprisingly find that post-calibration strategy is particularly useful — with only a small fraction of IID data , the classification accuracy is significantly improved . However , this approach can not be directly deployed in practice since it infringes the privacy rule in federated learning . Based on the above findings and considerations , we propose a novel and privacy-preserving approach called Classifier Calibration with Virtual Representations ( CCVR ) which rectifies the decision boundaries ( the classifier ) of the deep network after federated training . CCVR generates virtual representations based on an approximated Gaussian Mixture Model ( GMM ) in the feature space with the learned feature extractor . Experimental results show that CCVR achieves significant accuracy improvements over several popular federated learning algorithms , setting the new state-of-the-art on common federated learning benchmarks like CIFAR-10 , CIFAR-100 and CINIC-10 . To summarize , our contributions are threefold : ( 1 ) We present the first systematic study on the hidden representations of different layers of neural networks ( NN ) trained with FedAvg on non-IID data and provide a new perspective of understanding federated learning with heterogeneous data . ( 2 ) Our study reveals an intriguing fact that the primary reason for the performance degradation of NN trained on non-IID data is the classifier . ( 3 ) We propose CCVR ( Classifier Calibration with Virtual Representations ) — a simple and universal classifier calibration algorithm for federated learning . CCVR is built on top of the off-the-shelf feature extractor and requires no transmission of the representations of the original data , thus raising no additional privacy concern . Our empirical results show that CCVR brings considerable accuracy gains over vanilla federated learning approaches . 2 Related Work . Federated learning [ 2 , 3 , 4 ] is a fast-growing research field and remains many open problems to solve . In this work , we focus on addressing the non-IID quagmire [ 6 , 24 ] . Relevant works have pursued the following four directions . Client Drift Mitigation . FedAvg [ 2 ] has been the de facto optimization method in the federated setting . However , when it is applied to the heterogeneous setting , one key issue arises : when the global model is optimized with different local objectives with local optimums far away from each other , the average of the resultant client updates ( the server update ) would move away from the true global optimum [ 9 ] . The cause of this inconsistency is called ‘ client drift ’ . To alleviate it , FedAvg is compelled to use a small learning rate which may damage convergence , or reduce the number of local iterations which induces significant communication cost [ 25 ] . There have been a number of works trying to mitigate ‘ client drift ’ of FedAvg from various perspectives . FedProx [ 5 ] proposes to add a proximal term to the local objective which regularizes the euclidean distance between the local model and the global model . MOON [ 8 ] adopts the contrastive loss to maximize the agreement of the representation learned by the local model and that by the global model . SCAFFOLD [ 9 ] performs ‘ client-variance reduction ’ and corrects the drift in the local updates by introducing control variates . FedDyn [ 10 ] dynamically changes the local objectives at each communication round to ensure that the local optimum is asymptotically consistent with the stationary points of the global objective . FedIR [ 26 ] applies importance weight to the local objective , which alleviates the imbalance caused by non-identical class distributions among clients . 2We use the terms representation and feature interchangeably . Aggregation Scheme . A fruitful avenue of explorations involves improvements at the model aggregation stage . These works are motivated by three emerging concerns . First , oscillation may occur when updating the global model using gradients collected from clients with a limited subset of labels . To alleviate it , [ 11 ] proposes FedAvgM which adopts momentum update on the server-side . Second , element-wise averaging of weights may have drastic negative effects on the performance of the averaged model . [ 12 ] shows that directly averaging local models that are learned from totally distinct data distributions can not produce a global model that performs well on the global distribution . The authors further propose FedDF that leverages unlabeled data or artificial samples generated by GANs [ 27 ] to distill knowledge from the local models . [ 13 ] considers the setting where each client performs variable amounts of local works and proposes FedNova which normalizes the local updates before averaging . Third , a handful of works [ 14 , 15 ] believe that the permutation invariance of neural network parameters may cause neuron mismatching when conducting coordinate-wise averaging of model weights . So they propose to match the parameters of local models while aggregating . Data Sharing . The key motivation behind data sharing is that a client can not acquire samples from other clients during local training , thus the learned local model under-represents certain patterns or samples from the absent classes . The common practices are to share a public dataset [ 6 ] , synthesized data [ 16 , 17 ] or a condensed version of the training samples [ 18 ] to supplement training on the clients or on the server . This line of works may violate the privacy rule of federated learning since they all consider sharing raw input data of the model , either real data or artificial data . Personalized Federated Learning . Different from the above directions that aim to learn a single global model , another line of research focuses on learning personalized models . Several works aim to make the global model customized to suit the need of individual users , either by treating each client as a task in meta-learning [ 19 , 28 , 20 , 29 ] or multi-task learning [ 30 ] , or by learning both global parameters for all clients and local private parameters for individual clients [ 21 , 31 , 32 ] . There are also heuristic approaches that divide clients into different clusters based on their learning tasks ( objectives ) and perform aggregation only within the cluster [ 33 , 34 , 22 , 35 ] . In this work , we consider training a single global classification model . To the best of our knowledge , we are the first to decouple the representation and classifier in federated learning — calibrating classifier after feature learning . Strictly speaking , our proposed CCVR algorithm does not fall into any aforementioned research direction but can be readily combined with most of the existing federated learning approaches to achieve better classification performance . 3 Heterogeneity in Federated Learning : The Devil Is in Classifier . 3.1 Problem Setup . We aim to collaboratively train an image classification model in a federated learning system which consists of K clients indexed by [ K ] and a central server . Client k has a local dataset Dk , and we set D = ⋃ k∈ [ K ] Dk as the whole dataset . Suppose there are C classes in D indexed by [ C ] . Denote by ( x , y ) ∈ X × [ C ] a sample in D , where x is an image in the input space X and y is its corresponding label . Let Dkc = { ( x , y ) ∈ Dk : y = c } be the set of samples with ground-truth label c on client k. We decompose the classification model into a deep feature extractor and a linear classifier . Given a sample ( x , y ) , the feature extractor fθ : X → Z , parameterized by θ , maps the input image x into a feature vector z = fθ ( x ) ∈ Rd in the feature space Z . Then the classifier gφ : Z → RC , parameterized by φ , produces a probability distribution gφ ( z ) as the prediction for x. Denote by w = ( θ , φ ) the parameter of the classification model . Federated learning proceeds through the communication between clients and the server in a roundby-round manner . In round t of the process , the server sends the current model parameter w ( t−1 ) to a set U ( t ) of selected clients . Then each client k ∈ U ( t ) locally updates the received parameter w ( t−1 ) to w ( t ) k with the following objective : min w ( t ) k E ( x , y ) ∼Dk [ L ( w ( t ) k ; w ( t−1 ) , x , y ) ] , ( 1 ) where L is the loss function . Note that L is algorithm-dependent and could rely on the current global model parameter w ( t−1 ) as well . For instance , FedAvg [ 2 ] computes w ( t ) k by running SGD on Dk for a number of epochs using the cross-entropy loss , with initialization of the parameter set to w ( t−1 ) ; FedProx [ 5 ] uses the cross entropy loss with an L2-regularization term to constrain the distance between w ( t ) k and w ( t−1 ) ; MOON [ 8 ] introduces a contrastive loss term to address the feature drift issue . In the end of round t , the selected clients send the optimized parameter back to the server and the server updates the parameter by aggregating heterogeneous parameters as follows , w ( t ) = ∑ k∈U ( t ) pkw ( t ) k , where pk = |Dk|∑ k′∈U ( t ) |Dk ′ | .

Summary: This paper claims that current federated learning (FL) systems suffer from performance reduction due to high variance in the features learned by the final layers of the network across clients. The authors motivate their approach by training a local model per client and observing the representations learned by each layer. They show ​that the first layers generate similar representations across clients while the representations of more advanced layers are less similar. Therefore, they propose a simple post-processing step to correct the weights of the classification layer of the shared model after training with standard FL procedures (such as FedAvg). They assume the representation in the last layer follows a Gaussian distribution per class.  At the end of the training stage, they store statistics of the representations in the last layer for each client. Based on these statistics, they generate pseudo-data on the server and retrain only the last layer. The final model composed of the initial layers (that were trained regularly with standard FL approaches) and the final layer that was relearned.