No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data
Authors: Mi Luo
Fei Chen
Dapeng Hu
Yifan Zhang
Jian Liang
Jiashi Feng
Venue: NeurIPS 2021 Poster (2021)
Keywords: federated learning

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

No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data
Mi Luo', Fei Chen', Dapeng Hu', Yifan Zhang', Jian Liang" Jiashi Feng"
National University of Singapore
Huaawi Noah's 's Ark Lab
Institute of Automation, Chinese Academy of Sciences (CAS)
promyluo7, liangjian92, jshfeng) 8gma i ccm chen. fwwuawei. com, (dapeng. hu, yifan. zhang) @u nus. edu
Abstract
A central challenge in training classification models in the real-world federated system is learning with non-IlD data. To cope with this, most of the existing works involve enforeing regularization in local optimization model aggregation scheme at the server. Other works also share public datasets synthesized samples supplement the training of under-represented classes introduce a certain level of personalization. Though effective, they lack a deep understanding of how the data heterogeneity affects each layer of deep classi- fication model. In this paper, we bridge this gap by performing an experimental analysis of the representations learned by different layers. Our observations are surprising: (1) there exists greater bias in the classifier than other layers, and (2) the classification performance can be significantly improved by post-calibrating the classifier after federated training. Motivated by the above findings, we propose novel and simple algorithm called Classifier Calibration with Virtual Representa- tions (CCVR), which adjusts the classifier using virtual representations sampled from an approximated gaussian mixture model. Experimental results demonstrate that CCVR achieves state-of-the the-tarrtt performance benchmarks including CIFAR- 10, CIFAR 100, and CINIC-10. We hope that our simple yet effective method can shed some light on the future research of federated learning with non-IID data.
improving the
popular federated learning
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 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 large number of edge clients and the raw data never leave local devices.
A notorious trap in federated learning is training with non-lID data. Due to diverse user behaviors, large heterogeneity may be present 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 certain degree; 2) aggregation scheme [11, 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 more balanced data distribution on the client or on the server;
"corresponding author.
35th Conference on Neural Information Processing Systems (NcurIPS 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 toought-pooooing: comparing different layers learned on different clients, the classifier has the lowest feature' 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 debias the classifier via regularizing the classifier during training or calibrating classifier weights after training. We surprisingly find that post-calibration strategy particularly useful with only small fraction of IID data, the classification accuracy is significantly improved. However, this approach cannot 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 preneervireeeenngggggg 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-an art on common federated learning benchmarks like 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 a simple and universal classifier calibration algorithm for federated learning. CCVR is built 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 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.
We use the terms representation and feature interchangeably.

---

Layer 1
Layer 4
Layer 7 (Classifier)
0.8
Client ID
Client
Cilent
0.8
0.4
0.2
0.0
9. 6
0 66666110000
Client ID
Client ID
Client ID
Figure 1: CKA similarities three different layers of different 'client model-client model' pairs.
CKA Similarity et Round
CKA Similarity at Round 100
Smilanby
CKA Smilannty
0.6
04
0.4
FedAug-IlD
0.2
0.2
FedAvg-NIID
0.0
0.0
Layer
Layer
Figure means of CKA similarities of different layers in different local models.
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 limited subset of labels. 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 cannot produce 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 local works and proposes FedNova which normalizes the local updates before averaging. Third, 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 cannot acquire samples from other clients during local training, thus the Icarned local model under-represents certain patterns or samples from the absent classes. The common practices are to share public dataset [6]. synthesized data [16, 17] 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 artificial data.
Personalized Federated Learning. Different from the above directions that aim to learn a single global model, another line of research focuses 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 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 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 & has a local dataset D and we set
3

---

Label Distribution
NIID-Round
NIID-Round 100
IID-Round 100
1.50
6 a
4000
9 0 2
1.2
Client ID
Client ID
Client ID
Client ID
1.00
3000
0.75
2000
0.50
1000
O 6N456600 56666666
9 6
Classes
Classes
Classes
Classes
Figure 3: Label distribution of CIFAR-II 1- across clients (the first graph) and the classifier weight norm distribution across clients in different rounds and data partitions (the three graphs the right).
D= UkEjK] D as the whole dataset. Suppose there are classes 'D indexed by [C]. Denote by (x. E X' [C] sample in D, where is image in the input space X' and y is its corresponding label. Let D = {(.. v) E D y c} be the set of samples with ground ground-truth label on client k. We decompose the classification model into a deep feature extractor and a linear classifier. Given sample (z. y), the feature extractor fo X' + Z, parameterized by 0, maps the input image a into a feature vector 2 (z) E IR in the feature space Z. Then the classifier ge Z - RC, parameterized by sp, produces probability distribution Pip (p(2) as the prediction for æ. Denote by w (0, .) the parameter of the classification model.
Federated learning proceeds through the communication between clients and the server in round- by-round manner. In round / of the process, the server sends the current model parameter w to set (+t) of selected clients. Then each client & E ((() locally updates the received parameter (f-11) (t) w with the following objective:
min E (a) [C(w (1), (t-1) xr x),y w
where L is the loss function. Note that L is model parameter (t-1) as well. For instance, FedAvg [2] computes w (t) by running SGD on D* for number of epochs using the cross-entropy loss, with initialization of the parameter set to (t-1)); FedProx [5] uses the cross entropy loss with Læ-regularization term to constrain the distance between w and (t-1), MOON [8] introduces contrastive loss term 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 follows,
and could rely on the current global
|Dx
(t)
Σ kEU(()
g (t) where pkc Exeeo [DK|'
3.2 A Closer Look at Classification Model: Classifier Bias
To vividly understand how non-lID data affect the classification model in federated learning, we perform an experimental study on heterogeneous local models. For the sake of simplicity, we choose CIFAR-10 with 10 clients which is a standard federated learning benchmark, and a convolutional neural network with 7 layers used in [8]. As for the non-IlI experiments, we partition the data according to the Dirichlet distribution with the concentration parameter set as 0.1. More details are covered in the Appendix. To be specific, for each layer in the model, we leverage the recently proposed Centered Kernel Alignment (CKA) [23] to measure the similarity of the output features between two local models, given the same input testing samples. CKA outputs a similarity score between (not similar at all) and (identical). We train the model with FedAvg for 100 communication rounds and each client optimizes for 10 local epochs at each round.
We first selectively show the pairwise CKA features similarity of three different layers across local models Figure 1. Three compared layers here are the first layer, the middle layer (Layer 4), and the last layer (the classifier), respectively. Interestingly, we find that features outputted by the deeper layer show lower CKA similarity. indicates that, for federated models trained on non-lID data, the deeper layers have heavier heterogeneity across different clients. By averaging the pairwise CKA

---

Table 1:
Accuracy @ (%) on CIFAR-IO with different degrees of heterogeneity.
Method
a = 0.5
Q 0.1
0.05 a
FedAvg
68.62-4:0.7
58.55
52.333+:.43
FedAvg + clsnorm
69.65:10.35 (t 1.03)
58.9420.08 (1 0.39)
51.74:1.002 (.50.59
FedAvg + clsprox
68. 68.82:10.75 0.20)
59.04 (t 0.49)
52.38:0.78 (1 0.05)
FedAvg clsnorm clsprox
68.75:10.75
0.13)
58.80=0.30 (1 0.25)
52.39:10.24 0.06)
FedAvg + calibration (whole data)
72.51+0..3 3.89)
04.90=00.4 6.15)
57.53+ (1.5.20
FodAvg
FodProx
before
Test Accuracy
70
calbration-whhle data calibration data
N.D.0.5
NID-0011
NID-0.05
N.D.0.5
NID--0.1 0.1 Data Distribution
NID-0.05
Dota Distributian
Figure 4: The effect classifier calibration using different amounts of data.
features similarity Figure 1, we can obtain a single value to approximately represent the similarity of the feature outputs by each layer across different clients. We illustrate the approximated layer-wise features similarity in Figure 2. The results show that the models trained with non-IlI data have consistently lower feature similarity across clients for all layers, compared with those trained on IID data. The primary finding is that, for non-lID training, the classifier shows the lowest features similarities, among all the layers. The low CKA similarities of the classifiers imply that the local classifiers change greatly to fit the local data distribution.
To perform a deeper analysis on the classifier trained on non-lID data, inspired by [36], we illustrate the L2 norm of the local classifier weight vectors in Figure We observe that the classifier weight norms would be biased to the class with more training samples at the initial training stage. At the end of the training, models trained on non-III data suffer from much heavier biased classifier than the models trained on IID data.
Based on the above observations about the classifier, we hypothesize that: because the classifier is the closest layer to the local label distribution, it can be easily biased to the heterogeneous local data, reflected by the low features similarity among different local classifiers and the biased weight norms. Furthermore, we believe that debiasing the classifier is promising to directly improve the classification performance.
3.3 Classifier Regularization and Calibration
To effectively debias the classifier, we consider the following regularization and calibration methods.
Classifier Weight 12-normalization. To eliminate the bias in classifier weight norms, we normalize the classifier weight vectors during the training and the inference stage. We abbreviate it to 'clsnorm', In particular, the classifier is linear transformation with weight oي [p),,..,,PP] Pcc], followed by normalization and softmax. Given a feature z, the output of the classifier is
(2(2); C
Vi [C].
EV=t 2110011111
Classifier Quadratic Regularization. Beyond restricting the weight norms of classifier, we also consider adding a proximal term similar to [5] only to restrict the classifier weights to be close to the received global classifier weight vectors from the server. We write it as "elsprox" for short. The loss function in Eq. (1) can be specified as
((w (t), (I-1), x, 3) = E( (fog (x)), y) + "110!"" (1-1))
where l is the cross-entropy loss and is the regularization factor.
5

---

end Classifier Post- Post-calibration with IID Samples. In addition to regularizing the classifier during federated training, we also consider a post-processing technique to adjust the learned classifier. After the federated training, we fix the feature extractor and calibrate the classifier by SGD optimization with a cross-entropy loss on III samples. Note that this calibration strategy requires IID raw features collected from heterogeneous clients. Therefore, it can only serve as an experimental study use but cannot be applied to the real federated learning system.
We conduct experiments to compare the above three methods on CIFAR-I0 with three different degrees of data heterogeneity and present the results in Table We observe that regularizing the L2-norm of classifier weight (clsnorm) is effective for light data heterogeneity but would have less help or even lead to damages along with the increase of the heterogeneity. Regularizing the classifier parameters (clsprox) is consistently effective but with especially minor improvements. Surprisingly. we find that calibrating the classifier of the FedAvg model with all training samples brings significant performance improvement for all degrees of data heterogeneity.
To further understand the classifier calibration technique. we additionally perform calibrations with different numbers of data samples and different off-the-shelf federated models trained by FedAvg and FedProx. The results are shown in Figure 4 and observe that data-based classifier calibration performs consistently well, even with 1 1050 training data samples for calibration use. These significant performance improvements after adjusting the classifier strongly verify our aforementioned hypothesis, i.e., the devil is in the classifier.
Classifier Calibration with Virtual Representations
Motivated by the above observations, we pro- pose Classifier Calibration with Virtual Repre- sentations (CCVR) that runs on the server after federated training the global model. CCVR uses virtual features drawn from an estimated Gaussian Mixture Model (GMM), without ac- cessing any real images. Suppose Sa and dp the feature extractor and classifier of the global model, respectively, where w = (0, g) the parameter trained by certain federated learning algorithm, e.g. FedAvg. We shall use fg to extract features and estimate the corre- sponding feature distribution, and re- re-train using generated virtual representations.
Algorithm Virtual Representation Generation Input:
Feature extractor fa of the global model, number M, virtual features for class e
# Server executes:
2 Send fg to clients.
3 # Clients execute:
4 foreach client l E [K] do
foreach class E [C] do
Produce Zr,k.j = fg (Z.,, for j-th sample in D. for E [Ne,a].
Compute Me,k and Ee,k using Eq. (2).
Send {MM,, Ee,e) : C [C]} to server. 10 end
Feature Distribution Estimation. For se- mantics related tasks such as classification, the features learned by deep neural networks can be approximated with mixture of Gaussian distribution. Theoretically, any continuous dis- tribution can be approximated by using a fi- nite number of mixture of gaussian distribu- tions [37]. In CCVR, we assume that fea- tures of each class in DD follow a Gaussian distribution. The server estimates this distribu- tion by computing the mean ft, and the covariance Σ, for each class e of 'D using gathered local statistics from clients, without accessing true data samples their features. In particular, the server first sends the feature extractor fo of the trained global model to clients. Let No.k = |D be the number of samples of class e on client k, and set N. Σk= Ne.k- Client & produces features (Ze.k.1..... , Ne,,,N., } for class c, where 2e,k.j = (o(kk,,.))) is the feature of the j-th sample in D and computes local mean Me.k and covariance Ee,t of D as:
# Server executes:
12 foreach class E [C] do
13
Compute Me and Σ. using Eq. (3) and (4). Draw a set Ge of M. features from A(eee Σg) with ground truth label c.
14
end
Output:
Set of virtual representations Uoe|o Go
Ne. Zek..j Ea,
New (Ze,kij - Mc,k) (Zekk. - pc.k)
He,k = Na.k j=1
(2)
Na, - j=1

---

Then client & uploads ((Hc, Ee,)) : c E [C]] to server. For the server to compute the global statisties of D. is sufficient represent the global mean He, and covariance E. using Mo, 's and De,k 8 for each class c. The global mean can be straightforwardly written as
Ne.k
Ne. Nc
kckkjj
Hec..
(3)
No
k=
For the covariance, note that by definition we have
No.
(Ne,k E,,k
T ZackjZekk
No pecklee T
whenever Ne.k 2 1. Then the global covariance can be written as
K No.k
No Mullé
Σ
N
Ne
k=
j=1
K
N c.
K Na. N₂
Ne
N Melle
T.
(4)
N 1
k=1
Virtual Representations Generation. After obtaining pe's and E, 's, the server generates set Ge of virtual features with ground truth label e from the Gaussian distribution ((pe 2.). The number
Mc = Gel of virtual features for each class e could be determined by the fraction N inter-class distribution. See Algorithm 1.
D to reflect the
The last step of our CCVR method classifier re-training using virtual representations. We take out the classifier from the global model, initialize its parameter as 6 and re-train the parameter & for the objective
Classifier Re-Training,
min E (xz,))-Uccee G. (=)
where l is the cross-entropy loss. We then obtain the final classification model gG fă consisting of the pre-trained feature extractor and the calibrated classifier.
Privacy Protection. CCVR protects privacy at the basic level because each client only uploads
their local Gaussian statistics rather than the raw representations. Note that CCVR is just post-hoc method, so can be easily combined with some privacy protection techniques [38] to further secure privacy. In the Appendix, we provide empirical analysis on the privacy-preserve aspect.
5
Experiment
5.1 Experiment Setup
Federated Simulation. We consider image classification task and adopt three datasets from the popular FedML benchmark [39], i.e., CIFAR-10 [40], CIFAR- 100 [40] and CINIC-10 [41]. Note that CINIC-10 is constructed from ImageNet [42] and CIFAR-IO. 10, whose samples are very similar but not drawn from identical distributions. Therefore, it naturally introduces distribution shifts which is suited to the heterogeneous nature of federated learning. To simulate federated learning scenario, we randomly split the training set of each dataset into K batches, and assign one training batch to each client. Namely, each client owns its local training evaluation of the classification performance of the global model. For hyperparameter tuning, we first take out a 15% subset of training set for validation. After selecting the best hyperparameter, we return the validation set to the training set and retrain the model. We are interested in the NIID partitions of three datasets, where class proportions and number data points of each client are unbalanced. Following [14, 15], we sample p; Dirx (ax) and assign a Pi.a proportion of the samples from class i to client ks. We set a as 0.5 unless otherwise specified. For fair comparison, we apply the same data augmentation techniques for all methods.
We hold out the testing set at the server for
7

---

Table 2: Accuracy (%) on CIFAR-10 with different degrees of heterogeneity (a
{0.5, 0.0, 0.05)), CIFAR 100 and CINIC-10.
Method
a 0.5
a 0.1
a = 0.05
CIFAR-100
CINIC-10
FodAvg
68 08.6220..77
58.55+0.98
52.33++0.43
66.2540.54
60.20=22.04
FedProx
69.07=1.07
58.93-0064
53.00:10.32
66.311:0.39
60.52::2.07
No Calibration
FedAwgM
69.00
59.22.1:111
51.98.10.91
66.43.10.23
60.46.10.73
MOON
70.48+00.36
57.36-H0.85
49.91+0.38
67.02-40.33
65.67++2.10
FodAvg
71:03=0,40 <t 2.11)
62.68:00554
(1:5.15
54.95:0.61 (1 (.22)
66.60220.65 (1 0.355
69.9920.54 ( 9.79)
CCVR (Ours.)
FedAvgM
FedPrrx
71.49:0088 :1 2.49)
70.99:11.21
192))
62.60.0.433
62.644+1.07
(t 5-42)
3.67)
55.7911.07
54.57-00.58
(1.279)
(12559)
66.61:10.48 (1030)
66.71+:0.16
0.2%)
70.87--0.61 (1 10.41)
70.05:10.66 (.53)
MOON
71.29:00.11
(1-081)
62.22:40.70
(1))
55.60+0.63
(15669)
67.17=00.77
(10115)
69:42==0.65 (?3775)
FodAvg
72.31=0.53 <1 3.89)
64.70=0.94
(1165
57.534
(1
66.84220.0
(1559) 73,47=0030 (1 13:27)
FodProx
72.26=:1.22 (.19)
64.63::0.93
5.70)
57.33==0,72 (1:433)
66.68:20.43
(10.37) 73.10:10.57 12.50)
Oracle
FodAvgM
73.30=00.19: <1 4.30
64.24 1.32 (1 5.2)
57. L08 (135
66.94::0.22 (1 (1055)
72.88-+0033 3,7 (i 12,42)
MOON
72.05=0.16
64.94::0.58
58.14::0.47 (1
67.56::0.44 (1054) 03.38==0233 (17771)
Baselines and Implementation. We consider comparing the test accuracies of the representative federated learning algorithms FedAvg [2], FedProx [5], FedAvgM [11, 26] and the state-of-ther arrt method MOON [8] before and after applying our CCVR. For FedProx and MOON, we carefully tune the coefficient of local regularization term l and report their best results. For FedAvgM, the server momentum is set to be 0.1. We use a simple 4-layer CNN network with a 2-layer MLP projection head described in [8] for CIFAR-10. For CIFAR-100 and CINIC-10, we adopt MobileNetV2 [43]. For CCVR, to make the virtual representations more Gaussian-like, we apply ReLU and Tukey's transformation before classifier re-training. For Tukey's transformation, the parameter is set to be 0.5. For each dataset, all methods are evaluated with the same model for fair comparison. The proposed CCVR algorithm only has one important hyperparameter, the number of feature samples Me to generate. Unless otherwise stated, M, is set to 100, 500 and 1000 for CIFAR-10, CIFAR- 100 and CINIC-10 respectively. All experiments run with PyTorch 1.7.1. More details about the implementation and datasets are summarized in the Appendix.
5.2 Can classifier calibration improve performance of federated learning?
In Table 2, we present the test accuracy on all datasets before and after applying our CCVR. We also report the results under an ideal setting where the whole data are available for classifier calibration (Oracle). These results indicate the upper bound of classifier calibration.
CCVR consistently improves all baseline methods. First, it can be observed that applying classifier calibration increases accuracies for all baseline methods, even with the accuracy gain up 10.41 % on CINIC-10. This is particularly inspiring because CCVR requires no modification to the original federated training process. One can easily get considerable accuracy profits by simply post-procest post-processing the trained global model. Comparing the accuracy gains of different methods after applying CCVR and whole data calibration, we find that the accuracies of FedAvg and MOON get the greatest increase. On CINIC-10, the oracle results of FedAvg even outstrip those of all other baselines, implying that FedAvg focuses more on learning high-quality features but ignores learning a fair classifier. It further confirms the necessity of classifier calibration.
5.3 In what situation does CCVR work best?
We observe that though there is improvement on CIFAR-100 by applying CCVR, it seems subtle compared with that of other two datasets. This is not surprising, since the final accuracy achieved by classifier calibration is not only dependent on the degree to which the classifier is debaised, but also closely correlated with the quality of pre-trained representations. In CIFAR-100, each class only has 500 training images, so the classification task itself is very difficult and the model may learn representations with low separability. It is shown that the accuracy obtained with CCVR on CIFAR-100 is very close to the upper bound, indicating that CCVR does a good job of correcting the classifier, even if it is provided with a poor feature extractor.
We also note that CCVR achieves huge improvements on CINIC-10. To further analyze the reason of this success and the characteristics of CCVR, we now show the t-SNE visualization [44] of the features learned by FedAvg on CINIC-I0 dataset in Figure 5. From the first and second sub-graphs, we can observe that some classes dominate the classification results, while certain classes are rarely
8

---

ground truth
caliaration
calibratonn
t-SNE visualization of the features learned by FedAvg on CINIC-10. 10. The features colored by the ground truth and the predictions of the classifier hefore and after applying CCVR. Best viewed in color.
Figure 5:
predicted correctly. For instance, the classifier makes wrong prediction for most the samples belonging to the grey class. Another evidence showing there exists a great bias in the classifier is that, from the upper right corner of the ground truth sub-graph, we can see that the features colored green and those colored purple can be casily separated. However, due to biases in the classifier, nearly all purple features are wrongly classified as the green class. Observing the third sub-graph, we find that by applying CCVR, these misclassifications are alleviated. We also find that, with CCVR. mistakes are basically made when identifying easily-confused features that are close to the decision boundary rather than majority of features that belong certain classes. This suggests that the classifier weight has been adjusted to be more fair to each class. In summary, CCVR may be more effective when applied to the models with good representations but serious classifier biases.
5.4 How to forecast the performance of classifier calibration?
We resort to Sliced Wasserstein Distance [45], which is a popular metric to measure the distances between distribu- tions, to quantify the separability of GMM. The experiments conducted on CIFAR-10 10 with 0.1. We first com- pute the Wasserstein distances between any two mixtures, then we average all the distances to get mean distance. The farther the distance, the better the separability of GMM. We visualize the relationship between the accuracy gains and the separability of GMM in Figure 6. It is observed that the mean Wasserstein distance of GMM is positively correlated with the accuracy upper bound of classifier cal- ibration. It verifies our claim in Section 5.3: CCVR may be more effective when applied to the models with good (sep- arable) representations. In practice, one can use the mean Wasserstein distance of GMM to evaluate the quality of the simulated representations, as well as to forecast the potential performance of classifier calibration.
GMM Separability v.s. Accuracy Gain
CCVR
Whole Data
Accuracy Gain (%)
6-
Senbuuhh
NOCK
Poakorg
60
70
B0
90
100
Mean Wasserstein Distance
Figure 6: GMM's separability.
5.5 How many virtual features to generate?
One important hyperparameter in our CCVR is the number of virtual features M. for each class c to generate. We study the effect of M,, by tuning it from (0, 50, 100, 500, 1000, 2000] on three different partitions of CIFAR-10 (ar E (0.05, 0.1, 0.5)) when applying CCVR to FedAvg. The results are provided in Figure 7. In general, even sampling only a few features can significantly increase the classification accuracy. Additionally, it is observed that on the two more heterogeneous distributions (the left two sub-gra sub-graphs), more samples produces higher accuracy. Although results on NIID-0.5 give similar hint in general, an accuracy decline when using medium number of virtual samples is observed. This suggests that M, is more sensitive when faced with a more balanced dataset. This can be explained by the nature of CCVR: utilizing virtual feature distribution to mimic the original feature distribution. As result, if the number of virtual samples is limited, the simulated distribution may deviates from the true feature distribution. The results on NIID-0.5 implies that this trap could be easier to trigger when CCVR dealing with a more balanced original distribution. To conclude, though CCVR can provide free lunch for federated classification, one should still be very careful when tuning Mc to achieve higher accuracy. Generally speaking, larger value of M, is better.

---

72
Test Accuracy
55
62-
60-
50
68
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Virtual Features
Virtual Features
a Virtual Features
Figure 7: Accuracy@) @1 (%) CCVR CIFAR-10 with different numbers of virtual samples.
5.6 Does differrnt levels of heterogeneity affeet CCVR's performance?
We study the effect of heterogeneity on CIFAR- by generating various non-III partitions from Dirichlet distribution with different concentration parameters a. Note that partition with smaller more imbalanced.
can be seen from Table that CCVR steadily improves aceuracy for all the methods on all partitions. Typically, the improvements is greater when dealing with more heterogeneous data, implying that the amount of bias existing the classifier positively linked with the imbalanceness of training data. Another interesting discovery is that vanilla MOON performs worse than FedAvg and FedProx when equals to 0.1 or 0.05, but the oracle results after classifier calibration is higher than those of FedAvg and FedProx. It indicates that MOON's regularization on the representation brings severe negative effects on the classifier. As a consequence, MOON learns good representations but poor classifier. In that case, applying CCVR observably improves the original results, making the performance of MOON on par with FedAvg and FedProx.
Limitations
In this work, mainly focus the characteristic of the classifier in federated learning, because it found to change the most during local training. However, our experimental results show that in highly heterogeneous setting, only calibrating the classifier still cannot achieve comparable accuracies to that obtained on IID data. This is because the performance of classifier calibration highly relies on the quality learned representations. Thus, it's more important learn good feature space. Our experiments reveal that there may exist a trade-off in the quality of representation and classifier in federated learning on non-IID data. Namely, the methods that gain the greatest benefits from classifier calibration typically learn high-qualit representations but poor classifier. We believe this finding is intriguing for future research and there is still long way to tackling the non-lID quagmire.
Moreover, we mainly focus on the image classification task in this work. Our experiments validate that the Gaussian assumption works well for visual model like CNN. However, this conclusion may not hold for language tasks or for other architectures like LSTM [46] and Transformer [47]. We believe the extensions of this work to other tasks and architectures worth exploring.
7 Conclusion
In this work, we provide a new perspective to understand why the performance of deep learning- based classification model degrades when trained with non-III data in federated learning. We first anatomize the neural networks and study the similarity of different layers of the models on different clients through recent representation analysis techniques. We observe that the classifiers different local models are less similar than any other layer, and there is a significant bias among the classifier. then propose a novel method called Classifier Calibration with Virtual Representations (CCVR), which samples virtual features from approximated Gaussian Mixture Model (GMM) for classifier calibration to avoid uploading raw features to the server. Experimental results on three image datasets show that CCVR steadily improves over several popular federated learning algorithms.
10

---

Acknowledgement
We would like to thank the anonymous reviewers for their insightful comments and suggestions.
References
[1] Deng, J., Dong, R. Socher, et al. Imagenet: A large-scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, pages 248 255. lcce, 2009.
-afiiiiatttotfiiiienttt [2] McMahan, B., E. Moore, D. Ramage, et al. Communication Efficient Learning of Deep Networks from Decentralized Data. In AISTATS, pages 1273-1282 1282. 2017.
[3] Kairouz, P.. H. B. McMahan, B. Avent, et al. Advances and open problems in federated learning. arXiv preprint arXiv: (eti:::992..077.,,. 2019.
[4] Li. A. K. Sahu, A. Talwalkar, al. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50 60, 2020.
Li, A. K. Sahu, M. Zaheer, al. Federated optimization in hetcrogencous networks. arXiv preprint arXiv:1812.00 181.06677. 2018.
[6] Zhao, Y., M. Li, L. Lai, et al. Federated learning with non-iid data, arXiv preprint arXiv:1806.00582 2018.
[7] Li, Q. Y. Diao, Q. Chen, et al. Federated learning on non-iid data silos: An experimental study. arXiv preprint Stff::2122.0000 2021.
[8] Li, Q. B. He, D. Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition. 2021.
Karimireddy, P., S. Kale, M. Mohri, al. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132-5143. PMLR, 2020.
[10] Acar, A. E., Y. Zhao, R. Navarro, et Federated learning based on dynamic regularization. In International Conference on Learning Representations. 2021.
[11] Hsu, T.-M. H., H. Qi, M. Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv: arXiv:1909.06:: 1909006339, 2019.
[12] Lin, T., L. Kong, S. U. Stich, et al. Ensemble distillation for robust model fusion in federated learning. arXiv preprint arXiv:2006.0775 07242, 2020.
[13] Wang. J., Liu, H. Liang, al. Tackling the objective inconsistency problem in heterogeneous federated optimization, arXiv preprint arXiv:2007.0748) 2020.
[14] Yurochkin, M., M. Agarwal, S. Ghosh, Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learning, pages 7252 7261. PMLR, 2019.
[15] Wang, H., M. Yurochkin, Y. Sun, et al. Federated learning with matched averaging. In International Conference on Learning Representations. 2020.
[16] Hao. W., M. El-Khamy, J. Lee, et al. Towards fair federated learning with zero-shot data augmentation. arXiv preprint arXiv:2104.1 2021.,2021.
[17] Jeong. S. Oh, H. Kim, al. Communicatoon-effiiiee on-device machine learning: Federated distillation and augmentation under non-iid private data. arXiv preprint 1rX.w.181.11149, 2018
[18] Goetz, A. Tewari, Federated learning via synthetic data, arXiv preprint arXiv:2008, 200.:0000.04489,, 2020.
[19] Fallah, A., A. Mokhtari, A. Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.07945 2020.
[20] Jiang, Y., Konečny, K. Rush, et al. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv: 12488, 2019.
[21] Bui, D., K. Malik, J. Goetz, et al. Federated user representation learning, arXiv preprint 1909:1900.123 2019.

---

[22] Sattler, F., K.-R. Müller, W. Samek. Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE Transactions Neural Networks and Learning Systems, 2020.
[23] Komblith, S., M. Norouzi, Lee, al. Similarity of neural network representations revisited. In International Conference on Machine Learning, pages 5519-3529 PMLR, 2019.
[24] Hsich, K., A. Phanishayee, O. Mutlu, et al. The non-iid data quagmire of decentralized machine learning. arXiv preprint arXiv: 1910.0018, 2019.
[25] Li. X., K. Huang, W. Yang, et al. On the convergence of fedavg on non-iid data. arXiv preprint arXiv: 1 1907.02189 (89,002199900 2019.
[26] Hsu, T.MM. H., H. Qi, M. Brown. Federated visual classification with real-world data distribution. In Computer Vision- ECCV2200, 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part 16, pages 76-92. Springer, 2020.
[27] Goodfellow, I. J., J. Pouget-Abadic, M. Mirza, al. Generative adversarial networks. arXiv preprint arXiv: arXiv:1406.22 (ee6.100.6666 2601, 2014.
[28] Chen, F., M. Luo, Z. Dong, et al. Federated meta-learning with fast convergence and efficient communica- tion. arXiv preprint 1x22::/802.0776 2018.
[29] Khodak, M., M.-F F. Balcan, A. S. Talwalkar. Adaptive gradient-based meta-leaming methods. In Advances in Neural Information Processing Systems, pages 5917-5928. 2019.
[30] Smith, V., C.-K. Chiang. M. Sanjabi, et al. Federated multi-task learning. In Advances in Neural Information Procexsing Systems, pages 4424 2017.
[31] Liang, P. P., Liu, Ziyin, Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.0/4 011222,,..,,,,,.... 2020.
[32] Arivazhagan, M. G., V. Aggarwal, A. K. Singh, et al. Federated learning with personalization layers. arXiv preprint arxiv: 1912.0081,, 2019.
[33] Ghosh, A., J. Chung. D. Yin, et al. An efficient framework for clustered federated learning. arXiv preprint WiXiv:2006.04083 2020.
[34] Ghosh, A., J. Hong, D. Yin, et al. Robust federated learning in a heterogeneous environment. arXiv preprint arXiv: 0027/996606622 2019.
[35] Xie, M., G. Long, T. Shen, al. Multi-center federated learning. arXiv preprint arXiv:2005.01022 2020.
[36] Kang, B., S. Xie, M. Rohrbach, et al. Decoupling representation and classifier for long-tailed recognition. In International Conference on Learning Representations, 2020.
[37] Lindsay, B. Mixture models: theory, geometry and applications. NSF-CBMS regional conference series in probability and statistics, pages i-163. JSTOR, 1995.
[38] Vepakomma, P., T. Swedish, R. Raskar, et al. No peek: A survey of private distributed deep learning. arXiv preprint arXiv:1812.0322 2018.
[39] He, C., J. So, et al. Fedml: A research library and benchmark for federated machine learning. arXiv preprint 20X.v:2007.13578, 2020.
[40] Krizhevsky, A., G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[41] Darlow, L. N., E. J. Crowley, A. Antoniou, al. Cinic-10 not imagenet or cifar-10. arXiv preprint wrXiv:110.003505. 2018.
[42] Russakovsky, O., J. Deng, Su, et al. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (ICCV), 115(3):211-22) 11-222,,, 2015.
[43] Sandler, M., A. Howard, M. Zhu, et al. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510-4520, 2018.
[44] Van der Maaten, L., G. Hinton. Visualizing data using t-sne. Journal of machine learning research, 911). 2008.
12

---

[45] Kolouri, S., K. Nadjahi, U. Simsekli, et al. Generalized sliced wasserstein distances. In H. Wallach, H.. Larochelle, A. Beygelzimer, E. d'Alché-Bue, E. Fox, R. Garnett, eds., Advances in Neural Information Processing Systems, vol. 32. Curran Associates, Inc., 2019.
[46] Hochreiter, S., J. Schmidhuber. Long short-term memory. Neural Comput. 98)) 1997.
[47] Vaswani, A., N. Shazeer, N. Parmar, et al. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, NIPS' 17, page 6000-6010. Curran Associates Inc., Red Hook, NY, USA, 2017.
13


Reviews:
[{"review_id": "lN97tsAyq_V", "reviewer": "Reviewer_B69V", "summary": "- This work proposes a classifier calibration technique to handle the data heterogeneity commonly encountered in federated learning setups.\n- This work motivates its proposal by assessing the similarities of representations learned in different layers of the network, and finds that the final layer (classifier layer) is the one most dissimilar across clients.\n- They propose a post-learning classifier calibration technique through representations sampled from a gaussian mixture model estimated from representations from client data.\n- Experiments on CIFAR-10, CIFAR-100, and CINIC-10 with FedAvg, FedProx, and MOON show improvements over models learned with no calibration.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n\n- The paper is very well written with a clear flow of ideas and adequate references.\n- The experiments are well motivated with the prior experiment clearly leading to the next.\n- The insight that the earlier layers in a network tend to learn similar features while the last layer has the most dissimilarity is very interesting and could motivate future work in this area.\n\n\n**Improvements**\n\n- Table 1: While `clsnorm` and `clsprox` are understandable on closer study, it might help to map these terms to relevant sections in Sec 3.3\n- Figure 6: It will be helpful to overlay all three images on one another to enable the reader to understand the difference in heterogeneity on a common scale. Currently, with difference Y-scales for the three images, it is not easy to compare them with one another\n- Figures 4 and 6: It will also be helpful to have error bars on these figures.\n- There are minor typos in some places which the authors should go through and fix in the updated manuscript", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- The authors discuss the limitations of their method in Section 5.3", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "iredwdfH8Q", "reviewer": "Reviewer_445Q", "summary": "The authors proposed CCVR to post-calibrate federated learning classifiers models. Insights were first drawn from preliminary experiments, and CCVR is then tested atop many popular FL methods to demonstrate its efficacy.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of using simple statistics to calibrate federated learning models is great. Especially when the authors were able to show prominent gains with a post-hoc method that does not require modifications to existing training methods. The field of FL generally lacks this type of insightful yet powerful research.\n\nWith that said, I am providing some comments below:\n\n- This work is missing a critical reference: [1] addresses client distribution drift on large scale FL datasets. It falls under the first category in L28.\n\n- It seems a bit indirect having to aggregate per-client Gaussian mixtures into a global Gaussian mixture. One might argue that this step removes per-client statistics in exchange for better privacy, but in practice the clients might still need to send the full matrices over to the server for the aggregation. Hence it is reasonable that a server utilizes these statistics for learning/fine-tuning. I wish to either see the authors address this privacy comment, or provide some insights on how per-client CCVR (sampling from per-client Gaussians) compares with global CCVR (sampling from global Gaussian)\n\n- In the case of extremely non-identical client distribution such as when alpha -> 0, some clients might hold very sparse classes, and the mean/covariance would be mostly represented by a few samples. The statistics may expose a lot about the client’s single examples. \n\n- Although this work investigates a general phenomenon about FL, the datasets used (CIFAR/CINIC) are still very basic datasets that were not intended for realistic non-identical setups. Sure the datasets can be synthesized to have non-identical clients, by the conventional Dirichlet assignment neglects the fact that client data pool sizes can be drastically different. See [1]. An additional verification on larger datasets would be very helpful.\n\n- Isn’t section 3.3 simply “data sharing for FL model improvement”? If it is, I am receiving contradicting messages from the Introduction and the analysis in this section.\n\n- It’ll be good to provide a brief analysis on the sensitivity against µ, the regularization coefficient.\n\n- In table 1 clsnorm / clsprox is not clearly described. Consider including the names in section 3.3\n\n- Figure 5 does not convey the message effectively. The embeddings are basically identical across rows and the differences are exhibited in the colors. Consider using confusion matrices.\n\n[1] Federated visual classification with real-world data distribution.\n\nI am happy to increase my rating should the comments be addressed properly.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I recommend discussing more limitations should the space allow.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "RTLWMcm3FEo", "reviewer": "Reviewer_BVNc", "summary": "This paper proposes to bridge the gap between data heterogeneity and layer-wise sensitivity in deep models in Fed learning. The contributions of this paper are as follows: (1) the authors claim to be the first systematic study on hidden representation sensitivity to IID noise, the result of which reveals that the classifier layer is the most sensitive. (2) the authors then propose a calibration procedure that preserves the privacy of inputs while dealing with mitigating data drift. ", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: \n- Reasonable summary of related work, broad synthesis of categories related to heterogeneity mitigation is useful.\n- The proposed method is very simple and easy to understand. \n- Great reproducibility; supplemental material provides code and data for all experiments. \n\nWeaknesses:\n- One of the main selling points is that the authors fit a GMM on the last representation layer before the classifier from which to draw samples for calibration. The authors suggest that the alternative approach of directly using hidden representations from actual samples violates privacy. However, fitting a GMM to the last layer essentially leaks information about the distribution (i.e., the sufficient statistics) from which the client data is generated. In this sense, it is unclear whether the proposed approach actually respects the privacy constraint. \n- Weak experimental baselines. Only FedAvg, FedProx, and MOON are used for comparison. No baselines from personalized federated learning or aggregation schemes are included. In particular, a more fair comparison should consider personalized federated approaches such as meta-learning, which is in principle similar to the proposed GMM + calibration approach. \n- Limited technical contribution. Why use GMM in favor of other (more powerful) generative models? Section 5.3 mentions the separability of learned representations. If a GMM is used to generate synthetic samples, perhaps a more rigorous analysis can be provided to formalize 5.3 findings (i.e., performance guarantees or limitations based on GMM parameters, etc).\n\nOverall the novelty is not very high and experiments do not make a strong case for a NIPS paper.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "70ywncsgU6y", "reviewer": "Reviewer_JNfg", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall it is a nice paper; however, I don't think that it is original or significant enough to be accepted to NeurIPS. I am willing to change my mind in light of new information from the authors and by addressing the below comments.  \n\nThis paper presents the following merits:  \n- The approach is simple and doesn't impose a large overhead over current learning procedures.  \n- The approach is general in the sense that it can be integrated with many FL learning procedures.  \n- In CINIC-10 the improvements are significant over compared baseline methods.  \n- The experimental setup is clear. Code was provided which is great.  \n- The analysis of the method is good and exhaustive.  \n\nSome points that I found problematic:  \n- I think that the comparison made in the paper to baseline methods is severely lacking. Therefore, it is hard to accurately quantify how significant the improvements over current methods really are. Many studies were done in the past few years showing that under heterogeneous distribution FedAvg does not perform well (e.g., [1]). A comparison to more recent baselines (even ones that can not be integrated with this method) is needed.  \n- All the results in the paper are with 10 clients. Real-life federate learning systems have much more than that. I believe that the performance should be tested under a varying amount of clients.  \n- The results on CIFAR-100 indicate that this method will no perform well when the data per client is limited. This point was also stated by the authors and I appreciate that. However, federated learning is of most use under this scenario exactly, when there is enough data per client one may do local training only. I think that this point weakens the paper's claim.  \n- Although the analysis made in the paper showing that the representations of later layers are less similar across clients is nice, I think this is not surprising, especially given that there isn't a domain shift across clients. This is in line with the current understanding of NNs that early layers learn general (and transferable) representations while the representation of later layers is more specific [2].  \n\nSome minor comments:\n- The authors cannot expect the average user to be familiar with Centered Kernel Alignment (CKA) and should provide a detailed description about it in the appendix.  \n- I think that the colors in the t-SNE figures should be associated with the actual classes (and refer in the text to points associated with that class by the class name, e.g., L297).\n- Some references should be altered from arXiv to the publication venue (e.g., [3], [4]).\n\n\n[1] Arivazhagan, M. G., Aggarwal, V., Singh, A. K., & Choudhary, S. (2019). Federated learning with personalization layers. arXiv preprint arXiv:1912.00818.  \n[2] Yosinski, J., Clune, J., Bengio, Y., & Lipson, H. (2014, December). How transferable are features in deep neural networks?. In Proceedings of the 27th International Conference on Neural Information Processing Systems-Volume 2 (pp. 3320-3328).  \n[3] Li, T., Sahu, A. K., Talwalkar, A., & Smith, V. (2020). Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3), 50-60.  \n[4] Wang, H., Yurochkin, M., Sun, Y., Papailiopoulos, D., & Khazaeni, Y. (2020, April). Federated Learning with Matched Averaging. In International Conference on Learning Representations.  ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors addressed the potential limitations and negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers appreciate the simple yet effective strategy in this paper that handles heterogeneous data in federated learning. The experimental finding that the last layer exhibits the most dissimilarity could inspire future research. Great reproducibility and good writing also make the paper stand out. Therefore, I recommend acceptance. In addition, please incorporate the new results in the final version.
Decision: Accept (Poster)