Federated Learning with Decoupled Probabilistic-Weighted Gradient Aggregation
Authors: Jian-hui Duan
Wenzhong Li
Sanglu Lu
Venue:  (2021)
Keywords: Federated Learning
Gradient Aggregation
Variational Inference

Abstract:  In the federated learning paradigm, multiple mobile clients train local models independently based on datasets generated by edge devices, and the server aggregates parameters/gradients from local models to form a global model.   However, existing model aggregation approaches suffer from high bias on both data distribution   and parameter distribution for non-IID datasets,    which result in severe accuracy drop for increasing number of heterogeneous clients.    In this paper, we proposed a novel decoupled probabilistic-weighted gradient aggregation approach called FeDEC for federated learning.  The key idea is to optimize gradient parameters and statistical parameters in a decoupled way, and aggregate the parameters from local models with probabilistic weights to deal with the heterogeneity of clients.  Since the overall dataset is unaccessible by the central server,   we introduce a variational inference method to derive the optimal probabilistic weights to minimize statistical bias.   We further prove the convergence bound of the proposed approach.   Extensive experiments using mainstream convolutional neural network models based on three federated datasets  show that FeDEC significantly outperforms the state-of-the-arts in terms of model accuracy and training efficiency.

Under review as a conference paper at ICLR 2021
FEDERATED LEARNING WITH DECOUPLED SIBEBILISTIC-WEIGHTEI GRADIENT AGGREGATION
Anonymous authors Paper under double-blind review
ABSTRACT
In the federated learning paradigm, multiple mobile clients train local models in- dependently based on datasets generated by edge devices, and the server aggre- gates parameters/gradients from local models to form global model. However, existing model aggregation approaches suffer from high bias on both data dis- tribution and parameter distribution for non-III datasets, which result in severe accuracy drop for increasing number of heterogeneous clients. In this paper, we proposed novel decoupled moobabilistic-weihhee gradient aggregation approach called FeDEC for federated learning. The key idea is to optimize gradient param- eters and statistical parameters in decoupled way, and aggregate the parameters from local models with probabilistic weights to deal with the heterogeneity clients. Since the overall dataset is unaccessible by the central server, we intro- duce a variational inference method to derive the optimal probabilistic weights to minimize statistical bias. We further prove the convergence bound of the proposed approach. Extensive experiments using mainstream convolutional neural network models based on three federated datasets show that FeDEC significantly outper- forms the state-of ofett--r--s in terms of model accuracy and training efficiency.
INTRODUCTION
Federated learning (FL) has emerged as a novel distributed machine learning paradigm that allows global machine learning model be trained by multiple mobile clients collaboratively. In such paradigm, mobile clients train local models based on datasets generated by edge devices such as sensors and smartphones, and the server is responsible to aggregate parameters/gradient from local models to form a global model without transferring data to central server. Federated learning has been drawn much attention in mobile-edge computing (Konecný ct al. (20166); Sun el all (2017) with its advantages in preserving data privacy (Zhu Jin (2020); Piang et all (2019); Keller et all (201)) and enhancing communication efficiency (Shamir et all (201 4)); Smith et all X); Zhang
at (2)T3;;; McMahan et all (2017); Wang all (2020)).
Gradient aggregation is the key technology of federated learning, which typically involves the fol- lowing three steps repeated periodically during training process: (1) the involved clients train the same type of models with their local data independently; (2) when the server sends aggregation signal to the clients, the clients transmit their parameters or gradients the server; (3) when server receives all parameters or gradients, it applies an aggregation methods to the received parameters or gradients to form the global model. The standard aggregation method FedAvg (McMahan el all (2017) and its variants such as FedProx (Lletall (2020a)). Zeno (Xie (2019)) and F-FedGGD (Li all (20206)) applied the synchronous parameter averaging method the entire model in- discriminately. Agnostic federated learning (AFL.) (Mohri et all (2019) defined an agnostic and risk-averse objective to optimize a mixture of the client distributions. FedMA (Wang et all (2020) constructed the shared global model in layer-wise manner by matching and averaging hidden elements with similar feature extraction signatures. The recurrent neural network (RNN) based aggregator elletal (201)) learned aggregation method to make it resilient to Byzantine attack.
Despite the efforts that have been made, applying the existing parameter aggregation methods for large number of heterogeneous clients in federated learning still suffers from performance issues. was reported in (Zhao et all (2018)) that the accuracy of a convolutional neural network (CNN) model trained by FedAvg reduces by up 55% for highly skewed non-III dataset. The work of

---

Under review as a conference paper at ICLR 2021
(Wang et al (2020) showed that the accuracy of FadAvg (McMahan et all (201)) and FedProx eL (2222) dropped from 61% to under 50% when the client number increases from 5 to 20 under heterogeneous data partition. A possible reason to explain the performance drops in federated learning could be the different levels of bias caused by inappropriate gradient aggregation, on which make the following observations.
Data Bias:
In the federated learning setting, local datasets are only accessible by the owner and they are typically non-IlD, Conventional approaches aggregate gradients uniformly from the clients, which could cause great bias to the real data distribution. Fig. D shows the distribution of the real dataset and the distributions of uniformly taking samples from different number of clients in the CIFAR- 10 dataset (Knihevsky (2019)). It is observed that there are great differences between the real data and the sampled distributions. The more clients involved, the more difference occurs.
GraundTruth
Central
Parcertage
Clients
0.5
Clients
Clients
10 Clients
15 Clients
15 Clients
20 Clients
0.4
20 Clents
0.2
4
1.000
0075
0.550
-0.25
0.00
0.25
0.50
1.00
Classes
Amplituse
Figure 1: The differences between real data Figure 2: The cumulative distribution func- tion of means and variances in BN layers of different FL. models (ResNetl @ICIAR--0))
and sampled datasets (CIFAR-10).
Parameter Bias: A CNN model typically contains two different types of parameters: the gradient parameters from the convolutional (Conv) layers and full connected (FC) layers; and the statis- tical parameters such as mean and variance from the batch normalization (BN) layers. Existing approaches such as FedAvg average the entire model parameters indiscriminately using distributed stochastic gradient descent (SGD), which will lead to bias on the means and variances in BN layer. Fig. D shows the means and variances in BN layer distribution of centrally-trained CNN model and that of uendAvg-trained models with different number of clients on non-lID local datasets. It observed that the more clients involved, the larger deviation between the central model and the federated learning models.
Our contributions: In the context of federated learning. the problems of data bias and parameter bias have not been carefully addressed in the literature. In this paper, we propose novel gradient aggregation approach called FeDEC. The main contribution of our work are summarized as follows. (1) We propose the key idea of optimizing gradient aggregation with decoupled probabilistic- weighted method. To the best our knowledge, we make the first attempt to aggregate gradient parameters and statistical parameters separatively, and adopt probabilistic mixture model to resolve the problem of aggregation bias for federated learning with heterogeneous clients. (2) We propose variational inference method to derive the optimal probabilistic weights for gradient aggregation, and prove the convergence bound of the proposed approach. (3) We conduct extensive experiments using five mainstream CNN models based on three federated datasets under non-III conditions. is shown that FeDEC significantly outperforms the maae--oftthe-ard in terms of model accuracy and training efficiency.
2 RELATED WORK
We summarize the related work as two categories: parameter/gradient aggregation for distributed learning and federated learning.
In distributed learning, the most famous parameter aggregation paradigm the Parameter Server Framework (Lietal (201)). In this framework, multiple servers maintain partition of the globally shared parameters and communicate with each other to replicate and migrate parameters, while the clients compute gradients locally with a portion of the training data, and communicate with the server for model update. Parameter server paradigm had motivated the development of numerous distributed optimization methods (Boyd et al, (2011;; Dean et all (2012); Dekel et all (2012); Richtärik & Takác (2016); Zhang et al. (2015).. Several works focused on
Distributed Learning:
2

---

Under review as a conference paper at ICLR 2021
improving the caaamunicatton-effiiiee for distributed learning (Shamir et al
(201 Smith et al (2018); Zhang et all (20T3)) To address the issue of model robustness. Zeno (Xie et al (2019) was proposed to make distributed machine learning tolerant to arbitrary number of faulty workers. The
RNN based aggregator (li et all (2019) adopted a meta-learning approach that utilizes a recurrent neural network (RNN) in the parameter server to learn to aggregate the gradients from the workers, and designed a coordinatewise preprocessing and postprocessing method to improve its robustness.
Federated Learning: Federated learning (Konečny al. (2015)) is an emerging edge distributed machine learning paradigm that aims to build machine-learning models based on datasets distribut- ing across multiple clients. One of the standard parameter aggregation methods is FedAvg (McMa han et all (20T2)), which combines local stochastic gradient descent (SGD) on each client with a server that performs parameter averaging. The lazily aggregated gradient (I.ag) method (Chen et all (2018) allowed clients running multiple epochs before model aggregation to reduce communication cost. For heterogeneous datasets, FedProx (L (20208) modified FedAvg by adding a hetero- geneity bound on datasets and devices to tackle heterogeneity. The FedMA (Wang et all (202I)) method, derived from AFL Mohri all (2019) and PFNM (Yurochkin et all (2019)), demonstrated that permutations of layers can affect the gradient aggregation results, and proposed layer-wiss gradient aggregation method to solve the problem. For fair resources allocation, the q-FedSGD (Li at (20206)) method encouraged more uniform accuracy distribution across devices in federated networks.
However, all the methods did not differentiate gradient parameters and statistical parameters and aggregated the entire model in a coupled manner. In this paper, we make the first attempt to decou- ple the aggregation of gradient parameters and statistical parameters with probabilistic weights to optimize the global model to achieve fast convergence and high accuracy in non-IID conditions.
3 FEDEC: A DECOUPLED GRADIENT AGGREGATION METHOD
3.1 OBJECTIVE OF FEDERATED LEARNING WITH NON-IID DATA
Consider a federated learning scenario with K clients that train their local CNN models indepen- dently based on local datasets X1, , XK and report their gradients and model parameters to central server. The objective of the server is to form an aggregate global CNN model to minimize
the loss function over the total datasets x = {x1, X,, Conventional federated learning tends to optimize the following loss function:
K
XR x
min L(W x)
La (WR, Xk),
where W is the parameters of the global model, We = 1, K) is the parameters of the k-th local model; L(-) and (.). indicate the loss functions for global model and local models accordingly. The above objective assumes training samples uniformly distributed among the clients, so that the aggregated loss can be represented by the sum of percentage-weiighte of the local losses.
As discussed in section MI conventional federated learning has two drawbacks. Firstly, local datasets are collected by mobile devices used by particular users, which are typically non-IID. Training sam- ples each client may be drawn from different distribution, therefore the data points available locally could be bias from the overall distribution. Secondly, since a neural network model is typ- ically consists of convolutional (Conv) layers and full-connected (FC) layers that formed by gradient parameters, and batch normalization (BN) layers that are formed by statistical parameters such as mean and variance, aggregating them without distinction will cause severe deviation of the global model parameters.
To address the above issues, we propose a decoupled Setaabilistic-eiighhee approach for federated learning that focuses on optimizing the following loss function:
K
min L(WW) W.
W! moanL
W! par }, x)
Toe C&
WI-LA W!-1k 111-1k xx). mean!
N
NN
where * indicates NN, mean and var; Wt N Wt miteaan and Wt var are the parameters of Conv and FC layers of the global model after t-th aggregation epoch; Wt-.k Wt-l, and Wt- are the NN mean ver
k-th local model been trained several local epoch based L. 1-th global model; क (k: 1,..., K) is
3

---

Under review as a conference paper at ICLR 2021
the probability that sample is drawn from the distribution of the k-thh client, ie., TT. 10 (k I,..., K) and Σ T
The above formulation objects to minimize the expected loss over K clients with non-llI datasets. Next, we will introduce a decoupled method called FeDEC optimize the parameters of different types of layers separatively. and derived the probability weights Tr for parameter aggregation.
3.2 DECOUPLED PROBABILISTIC- WEIGHTED GRADIENT AGGREGATION METHOD
In this section, we proposed decoupled method to derive the global model with respect to Nw (parameters of Conv and FC layers) and W mean? W ear (statistical parameters of BN layers).
3.2.1 GRADIENT AGGREGATION FOR CONV AND FC LAYERS
Since the parameters of Conv and FC layers are neural network weights which are updated by distributed gradient descent method (Nesterov (I193)), they are appropriate to be aggregated with similar approach that adapts conventional federated average for non-IlD datasets. Let
gi = W W (k = L.....K K), where N indicates NN parameters after full cal training. be the gradient of the ke-th client in the t-th training epoch. After receiving the gradients from K clients, the central server update the parameters of global model as follows.
W NN
W NN
(3)
where 33 is the learning rate for parameter update, of. (I: = 1..... K) are the probabilistic weights with > A= T 1 that derived in section 3
3.2.2 PARAMETER AGGREGATION FOR MEANS AND VARIANCES IN BN
Different from the Conv and FC layers, the BN layers mainly contain statistical parameters such as mean and variance. Conventional federated learning aggregates BN layers and other layers without distinction, which could lead to high bias of means and variances in BN layer of the global model.
Thus we propose a different way to aggregate means and variances in BN layer as follows.
In t-th training epoch, the means and variances Wt mean: Wt var in BN layer, which are updated by:
K
W!
f Wek
(4)
mean
mean.
1=
K
W var
Σ k=1
(xx
- 1)m W eart
(5)
X
where Wi.k mean and Wi.k indicate the means and variances in BN layers of the k-th client epoch vur
I o 1,... K) are probabilistic weights with Σ Tf. = 1 that are derived in section (.2.3
In the above equations, we update the mean with the weighted average of local models, and up- date the variance with the weighted pooled variance (Killeen (2(155)), which can give an unbias estimation of parameters the whole dataset under non-IID conditions (see Appendix A
3.2.3 DERIVATION OF PROBABILISTIC WEIGHTS
We adopt a mixture probabilistic model to describe non-IID datasets in federated learning. Without loss of generality, in the t-th training epoch, assume the mini-batch samples of each client fol- lows a Gaussian distribution (i(Hu, σ (k = 1,..., K)), where Ak dy are the mean and standard deviation of the distribution that vary among clients. We omit the upper script for simplicity there- after. The whole samples can be described as a Gaussian Mixture Model (GMM) with the following probability function:
p(x|A) Σ
TAP(X* oR),
where À { Uk. Ok ls 1,2, K) are the parameters of the GMM mode!
Notedd that the proposed variational inference method can be applied to other non-Gaussian distributions with slight modification.
4

---

Under review as a conference paper at ICLR 2021
In federated learning, the local data samples accessed by particular client and the central server can only observe the statistics of local dataset such as mean and standard variance. Without know- ing the overall samples, conventional expectation-mantt maximization (EM) algorithm (Dempster et al, (11771) cannot be applied to derive A. Alternatively, we introduce a variational inference method to estimate the parameters of X.
Specifically, we construct a variational Bayesian generative model to generate data that are close to the reported statistics of local models as possible, and use the generated data to estimate the GMM model parameters. The plate notions of the generative model are shown in Fig. B. The notations are explained as follows.
ΦK
z
O2
est W mean' W var. the observed statistics from the feature maps of k-th client.
z {z - ,C) is a vector of latent variables with length C, where z 10, _ 1, and C' is the number of classes for classification task. z can be viewed as a data distribu- tion that represents the probability of a sample in client & belonging to the classes.
Pk:
8
Figure 3: Bayesian generative model using plate notations.
The variational
θ {0x} are generative model parameters, and Φ {Φr} are variational parameters.
The solid línes in Fig. dennote the generative model po,, (2), )r, (st. (2]), and the dashed lines denote the variational approximation 4s. (s{( to the intractable posterior pe, (zf.s)) We approximate pa, (zi (s)) with qo. (z), ( by minimizing their divergence:
Φk. 0% = arg min divergence (z) pe, (2),s))),
02.4%
s.t. Σ
To derive the optimal value the parameters Φu and 0x, we compute the marginal likelihood of st:
log p(s)) (xo(gs. (z(1s)) pe. (z). (s)) + Eg,, (z((ett log Pa. (z). s
(8)
4 (z) (
In Eq. &, the first term is the KL-divergence (Joyce (20ID) of the approximate distribution and the posterior distribution; the second term is called the ELBO (Evidence Lower BOund) on the marginal likelihood of dataset in the k-th client.
Since log p(s jis non-negative, the minimization problem of Eq. can be converted maximize the ELBO. To solve the problem, we change the form of ELBO as:
log paz (2), st) ao. [x)
p(z))
E 9o, (zi(+))
Eg,, Pa (2) a
log qs (zf (s)))
Ege. (a (a llg poz (st (2)1)
Eneoder
Deunder
The above form is variational encoder-decoder structure: the model go. (z), (ss) can be viewed as a probabilistic encoder that given an observed statistics sf it produces a distribution over the possible values of the latent variables z!: The model pe,, (s (z can be refered as a probabilistic decoder that reconstructs the value of s/. based on the code z According to the theory of variational inference (Kingma & Welling (2)1E))) ), the problem in Eq, can be solved with stochastic gradient descent (SGD) method using fully-connected neural network to optimize the mean squared error loss function.
With the derived optimal parameters oi. 0% we can extract the latent variables z that is interpreted as the sample distribution of client-k. Therefore zíz can be used to infer the parameters (Ty, Mk σy.) of k-th component of the GMM model. Specifically, the probabilistic weights mi can be represented by
K
TKK
(10)
K
j=1
5

---

Under review as a conference paper at ICLR 2021
4 CONVERGENCE ANALYSIS
In this section, we will show that the convergence of the proposed FeDEC algorithm is theoreti- cally guaranteed. We use the following assumptions and lemmas, and the convergence guarantee is provided in Theorem 1.
Assumption (Unbiased Gradient): We assume that the stochastic gradients gf is an unbiased estimator of the true gradient f(w)). i.e., Elgi] [V(()), where f()) is any convex objective function and w is its variables.
Assumption 2 (Gradient Convex Set): We assume that gradient set G is a convex set, where all gradients g1 B2....;BK are in G, and any g= = >:g; (V); and Σi=>; = is in G.
Lemma 1 (L-Lipschitz Continuity): For function f(-) is Lipschitz continuous if there exists positive real constant such that, for all real 2ι and Z22
(f(z)) f(x2)) S L| -22/-
Lemma 2 (Jensen's Inequality): If fw) is a convex function on W, and [[f(w)] and f f(Blw]) are finite, then:
B[S(w)] 2 [(R(w))).
Definition 1 (Projection Operation): Assume W. an intermediate result of optimization, we define project operator (w.) to project w. to the domain W,, which is computed by:
II (w.)=arg arg min W
l|w- -w.ll-
wEW
Definition (Diameter of Domain): Given a function f(w), where w W, and W is f's domain of definition. The diameter of W is denoted by I:: for every W1, W2 E W: |[] w2ll S r.
Theorem (Guaranteed Convergence Rate): If convex function f(w) is L-Lipschitz continuous function, then f(w)l < L. Let I be the diameter of domain. Applying equations ((V))) for gradients aggregation, we have the following convergence rate for the proposed FeDEC algorithmm:
8
f(w")
min f(w) WEW
O
13).
(11)
23T
where w is the average result of w for total training epoch T. B is the learning rate in equation and T the total training epoch. If we let LVT. the convergence rate O VT
Proof skeleton: We provide a simple description of the proof skeleton of Theorem with the follow-
ing steps. (1) Since (.) is convex function, have / f(w") f(w) ( (w" w). (2) With assumption and 2, we have f(w")- f(w) 2 (l|w* w|| w w |3) l
where w is the intermediate result of f(w) in update time t + 1. (3) With lemma I and definition
1, by projecting w to w we have f(w*)- ) fww S 23 (llw" w || ) L
(4) Summing from /. 1 to T and with definition 1 and 2, we have Σ ((w Tf(w) <
23 1 LL*T. (5) According to lemma we have f(w"") - f(w < 257 + 5 L2 (6) Taking
D/(LTT), we can obtain the convergence rate in the theorem. The detailed proof of Theorem 1 and explanations are provided in Appendix A.
According to Theorem the FeDEC parameter aggregation algorithm is guaranteed converge, and the convergence rate can be as fast as general stochastic gradient decent which only related to the training epoch T with an associated constants. The constant related the optimization problem parameters such lipschitz constant I,, and diameter of domain r.
5 PERFORMANCE EVALUATION
In this section, we evaluate the performance of the proposed FeDEC method for federated learning.

---

Under review as a conference paper at ICLR 2021
5.1 EXPERIMENTAL SETUP
Implementation. We implement the proposed FeDEC parameter aggregation approach and the considered baselines in PyTorch (Lal (2019). We train the models in simulated federated learning environment consisting of one server and a set of mobile clients with wireless network connections. Unless explicitly specified, the default number of clients is 20, and the learning rate 3 = 0.01. We conduct experiments on a GPU-equipped personal computer (CPU: Inter Core i7- 8700 3.2GHz, GPU: Nvidia GeForee RTX 2070, Memory: 32GB DDR4 2666MHz, and OS: 64-bit Ubuntu 16.04).
Models and datasets. We conduct experiments based on 5 mainstream neural network models: 8 eHl (2016)), LeNet (Lecun et (1998), DenseNet121 121 (Huang all (2017)) Mo- b (Sandler all (2018)), and a 4-layer CNN (every CNN layer followed by a BN layer). The detailed structure of the CNN models are provided in Appendix A.
We use three real world datasets: MNIST (LeCun et all (20TT)) Fashion-MNIST (Xiao al (20T7), and CIFAR-10 (Knzhevsky (2019)) MNIST is dataset for hand written digits classi- fication with 60000 samples and each example is a 28 x 28 greyscale image. Fashion-MNT is dataset intended to replace the original MNIST for benchmarking machine learning algorithms. CIFAR-10 is a larger dataset with 10 categories. Each category has 5000 training images and 1000 validation images of size 32 32. For each dataset, we 80% of the data for training and amal- gamate the remaining data into a global test set.
We form non-lll) local datasets follows. Assume there C classes of samples in dataset.
7) [0, 1], if E classj, Each client draw samples form the dataset with probability pr(z) (V..., otherwise.
means that the client draw samples from particular class j with a fixed probability 7), and from other classes based on standard Gaussian distribution. The larger is, the more likely the client's samples concentrate on a particular class, and the more heterogeneous the local datasets are.
5.2 PERFORMANCE COMPARISON
We compare the performance of FeDEC with 5 state-of-the-art methods: FedAvg (McMahan et all (217) ) RNN based aggregator (1 ct al (2019)), FedProx (La al (2121a)) q geFedSGD (Li al
(20205)) and FedMA (Wang et all (2220). The results are analyzed as follows.
0.9
1.0
PedMA
FeDEC
quhddcci
FedAug
Accuracy
0.8
Kl-divergence
0.8
Peskeg
Loss
FedFroo
PUDBC
Fodkoy
0.4
FalM
96
192 Epoch
288
18
27
54 63 Time [min)
72 81
(c2r8-110) 10)
Denested11 22
Moboldler (Gatat
Lotie Maait (leter
Figure 4: Convergence of dif- Figure 5: Training accuracy of Figure 6: KL-divergence of dif- ferent algorithms (ResNet18 different algorithms (ResNet18 ferent algorithms. CIFAR-10).
on CIFAR-10).
Convergence: In this experiment we study the convergence of all baselines and our algorithm by showing the total communication epochs versus train loss. Fig. @ shows the result of ResNet18 18 on CIFAR-10. 10. It is shown that the loss of all algorithms tend be stable after a number of epoches. Clearly FeDEC has the lowest loss among all algorithms, which means that FeDEC converges faster that of baselines. The results of more CNN models on different datasets are shown Appendix A.4
Training Efficiency: In this experiment we study the test accuracy versus time during training of CNN model with federated learning. Fig shown the results of training ResNet18 on CIFAR- 10. is shown that FeDEC reaches 0.8 accuracy after 18 minutes, while FedMA, FedProx, and FedAvg take 36 to 63 minutes to reach the same accuracy. FeDEC approaches 0.9 accuracy after 54 minutes, while the accuracy of other algorithms are below 0.85. The results of more CNN models on different datasets are shown in Appendix A. suggests that FeDEC trains much faster than the baseline algorithms and it can reach high accuracy in short time period.

---

Under review as a conference paper at ICLR 2021
Parameter Bias: In this experiment study the parameter bias of federated learning algorithms. Fig. G compares the KL-divergence between the means and variances in BN of global models aggregated by different algorithms and the central model. It is shown that FedAvg, FedProx, and q- FedSGD have exceptional high parameter bias, while FeDEC has significantly lower KI--divergence compared to the baselines for different CNN models different datasets.
Global Model Accuracy: In this experiment, we compare the global model accuracy of different federated parameter aggregation algorithms after training to converge. We repeat the experiment for 20 rounds and show the average results in Table II As shown in the table, the central method yields the highest accuracy. In comparison of different federated learning methods, FeDEC sig- nificantly outperforms the other algorithms in global model accuracy. It performs better than the state-of-the-ar method FedMA with 2.87%, 3.17%, 2.58%, and 3.09% accuracy improvement in R8 8, DenseNet121, MobileNer and 4-1. CNN respectively for CIFAR-10, 1. 1.09% improve- ment in LeNet for F-MNIST, and 0.33% improvement in LeNet for MNIST accordingly. FeDEC achieves the highest accuracy among all baselines, and it performs very close to the centralized method, whose accuracy drop is less than 3% in all cases.
Average test accuracy on non- IID datasets. The "Central" method trains the CNN model in the cen- tral server with global dataset. The "FeDEC(W/o)" method means using the proposed probabilistic-weihhtt. aggregation methodwithout distinguishing NN and mean, var. The "FeDEC" method represents the pro-
Table 1:
posed decoupled probabilistic- weighted aggregation approach.
Model@ Datuset
Central
FedAvg
RNN
FedProx
g-FedSGD
FedMA
FeDEC(w/o)
FcDEC
ResNetBBeCIFARID
92.33%
83.29%
83.47%
81.0%%
87.44%
88.30%
90.31%
DenseNet121 CIFAR-10
93.24%
82.36%
85.03%
82.43%
88.12%
89.25%
91.29%
MobileNetV2@ @ CIIARR-O
92.51%
83.11%
80.68%
79.82%
86.59%
87.93%
89.17%
4-L @
85.53%
79.33%
80.02%
77.77%
81.76%
83.11%
84.85%
LeNet -
90.42%
87.41%
82.32%
86.21%
99.02%
89,52%
90.11 %
LeNet MNIST
98.95%
97.32%
97.06%
97.55%
95.88%
16%
98.37%
98,49% %
Bii 6
72.14%
71.08% 56
68.44%
73.81
77.51 %
77.51 %
Hyperparameter Analysis: We further analyze the influence of two hyperparameters in federated learning: the number of clients involved and the heterogeneity of local datasets.
0.95
Text Accuracy
0.90
Fig. compares the test accuracy of the global model for dif-
0.85
ferent number of involved clients. According to the figure, the performance of FeDEC is stable. When the number of mobile clients increases from 5 to 20, the test accuracy slightly de- creases from 0.909 to 0.893. Other baseline algorithms yield significant performance drop. FeDEC achieves the highest test accuracy among all federated learning algorithms in all cases, and it performs very close to the central model.
0.75
# Clients
Figure 7: ferent number of clients (ResNet18 on CIFAR-10).
Test accuracy with dif-
Central 1000C
TecdMA
Fedhrox
In the experiment, the heterogeneity of local datasets is repre- sented by 7,. the probability that client tends to sample from a particular class. The more 7) approaches to 1, the more hetero- geneous the local datasets are. Fig. & shows the test accuracy under different level of heterogeneity. As 17 increases, the test accuracy of all models decreases. FeDEC yields the highest test accuracy among all algorithms, and its performance drops much slower than the baselines. verifies the effectiveness of the proposed probabilistic- bobiiiistte--eiigheeddd gradient aggregation ap- proach under non-IID conditions.
0.95
Tedt Accuracy
0.90
0.85
0.80
0.75
0.70
9-
0.5
n- D.8
n 1.0
Figure 8:
Test accuracy with
different level of heterogeneity (ResNet18 18 on CIFAR-10).
6 CONCLUSION
Gradient aggregation played an important role in federated learning to form a global model. To address the problem of data and parameter bias federated learning for non-lID dataset, we pro- posed a novel probabilistic parameter aggregation method called FeDEC that decoupled gradient parameters and statistical parameters to aggregate them separatively. The probabilistic weights were optimized with variational inference, and the proposed method was proved to be convergence guar- anteed. Extensive experiments showed that FeDEC significantly outperforms the ofate-of oftte-----e---t- -ht- on variety of performance metrics.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends in Machine Learning. 3(1): 122, 2011.
Tianyi Chen, Georgios Giannakis, Tao Sun, and Wotao Yin. Lag: Lazily aggregated gradient for communication--ffiiiee distributed learning. Advances in Neural Information Processing Sys- tems (NIPS' 18), pp. 5050-5060, 2018.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthicu Devin, Mark Mao, Marc' Ranzato, Andrew Senior, Paul Tucker, Ke Yang, Quoc V. Le, and Andrew Y. Ng. Large scale distributed deep networks. In Advances in Neural Information Processing Systems (NIPS' PP- 1223 1231. 2012.
O. Dekel, R. Gilad-Bachrach O. Shamir, and L. Xiao. Optimal distributed online prediction using mini-batches. Journal of Machine Learning Research, 13:165-202, 2012.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the algorithm. JOURNAL OF THE ROYAL STATISTICAL SOCIETY, SERIES B, 39(1):1-38, 1977.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. IEEE Conference on Computer Vision and Pattern Recognition (CVPR' 16), pp. 770-778, 2016.
Gao Huang, Zhuang Liu, Laurens van der Maaten, and Kilian Q. Weinberger. Densely con- nected convolutional networks. In IEEE Conference Computer Vision and Pattern Recognition (CVPR' pp. 2261-2269, 2017.
Jinlong Ji, Xuhui Chen, Qianlong Wang, Lixing Yu, and Pan Li. Learning to learn gradient aggre- gation by gradient descent. Proceedings of the Eeenty--iggth International Joint Conference 2019.
Artificial Intelligence (ICAAI'9)) 19). pp.
Linshan Jiang, Rui Tan, Xin Lou, and Guosheng Lin. On lightweight privacy-preserving collabo- rative learning for internet-of-things objects. In Proceedings of the International Conference Internet of Things Design and Implementation (loTDI' 19), pp. 70-81, 2019.
James M. Joyce.
Divergence, pp. 720-722. 2011.
Marcel Keller, Valerio Pastro, and Dragos Rotaru. Overdrive: Making SPDZ great again. In Ad- vances in Cryptology (EUROCRYPT volume 10822, pp. 158-189, 2018.
Peter R. Killeen. An alternative null-hypothesis significance tests. Psychological science, 16(5): 345-53, 2005.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR' 14), 2014.
Jakub Konečný, H. Brendan McMahan, and Daniel Ramage. Federated optimization: Distributed optimization beyond the datacenter. NIPS Optimization for Machine Learning Workshop 2015, pp- pp.5, 2015.
Jakub Koneený, H. Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimiza- tion: Distributed machine learning for on-device intelligence. ArXiv, abs/1610.02577, 2016.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Yann Lecun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. In Proceedings of the IEEE, pp. 2278-2324, 1998.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs (Online). 2010.
Available: http://yamm.eecc..cmm/ee/a/

---

Under review as a conference paper at ICLR 2021
Mu Li, David G. Andersen, Jun Woo Park, Alexander J. Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J. Shekita, and Bor- Ying Su. Scaling distributed machine learning with the parameter server. In 11th USENIX Symposium Operating Systems Design and Implementation (OSDI" 14), pp. 583-598, 2014.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems (MLSys 20), pp. 429-450. 2020a.
Tian Li, Maziar Sanjabi, and Virginia Smith. Fair resource allocation in federated learning. In International Conference on Learning Representations (ICLR'20) 2020b.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-efficii learning of deep networks from decentralized data. Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS' 17), 54: 1273 1282, 2017.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Pro- ceedings of the 36th International Conference on Machine Learning (ICML'I9) I9), volume 97, pp. 4625, 2019.
Y. Nesterov. A method of solving convex programming problem with convergence rate o K Dokl. Akad. Nauk SSSR, 269:543-547, 1983.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, and Zachary Viti.. Pytorch: An imperative style. high-performance deep learning library. In Advances Neural Information Processing Systems (NeurIPS' 19), pp. 8024-8035, 2019.
Peter Richtárik and Martin Takáe. Distributed coordinate descent method for learning with big data. Journal of Machine Learning Research, 17(1):1-25, 2016.
M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L.. Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR' 18), pp. 45 4510-4500, 2018.
Ohad Shamir, Nati Srebro, and Tong Zhang. Communication-efficienn distributed optimization using an approximate newton-type method. In Proceedings of the 31st International Conference Machine Learning (ICML'14) volume 32, pp. 1000-1008, 2014.
V. Smith, S. Forte, C. Ma. M. Takac, M. I. Jordan, and M. Jaggi. Cocoa: A general framework for communication-effiicc. distributed optimization. Journal of Machine Learning Research, 18 (230):1-47, 2018.
Shizhao Sun, Wei Chen, Jiang Bian, Xiaoguang Liu, and Tie-Yan Liu. Ensembbe-cooppeession A new method for parallel training of deep neural networks. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases (ECML-KDD) MDDM-LDDDIIT 17), pp. 187-202, 2017.
Hongyi Wang, Mikhail Yurochkin. Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference on Learning Represen- tations (ICLR 20),, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Pashion-mnisse a novel image dataset for benchmark- ing machine learning algorithms. 2017.
Cong Xie, Sanmi Koyejo, and Indranil Gupta. Zeno: Distributed stochastic gradient descent with suspicion based fault-tolerance, In Proceedings of the 36th International Conference on Machine Learning (ICML'19) volume 97. pp. 6893-6901, 2019.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In Pro- ceedings of the 36th International Conference on Machine Learning (ICML' (I),, volume 97, pp. 7252-7261, 2019.
10

---

Under review as a conference paper at ICLR 2021
Sixin Zhang, Anna E Choromanska, and Yann LeCun. Deep learning with elastic averaging sgd. In Advances in Neural Information Processing Systems (NIPS' 15), pp. 685-693, 2015.
Y. Zhang, J. C. Duchi, and M. J. Wainwright. Communication-efficennn algorithms for statistical opti- mization. Journal of Machine Learning Research, 14:33213363, 2013.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and V. Chandra. Federated learning with non-iid data. ArXiv, abs/ 800.01806., 2018.
H. Zhu and Y. Jin. Multi-objective evolutionary federated learning. IEEE Transactions on Neural Networks and Learning Systems, 3104):1300-1322 2020.

---

Under review as a conference paper at ICLR 2021
A APPENDIX
A.1 PROOF OF CONVERGENCE GUARANTEE (THEOREM IN SECTION 4)
We provide the detailed proof of Theorem Section 4. We first restate the necessary equations and the theorem.
In section 3, we propose the usage of the following equations for gradient decent update algorithm:
W NN
W NN
*KSk:
"
(3)
k=1
where 33 is the learning rate for parameter update. And for batch normalization update:
K
W
Wtk
(4)
mean
mean
k=1
W var
(||x|||
1
Wi. ear
(5)
X
n
We restate the theorem in section the following:
Assumption 1 (Unbiased Gradient): We assume that the stochastic gradients gi is an unbiased estimator of the true gradient f(w)). i.e., = where f()) is any convex objective function and w is its variables.
Assumption (Gradient Convex Set): We assume that gradient set is a convex set, where all gradients g1 gK are in G, and any g = - X:g; (V)i >> and Σ ) = 1) is in G.
Lemma 1 (L-Lipschitz Continuity): For function f()) is Lipschitz continuous if there exists positive real constant L. such that, for all real xı and x22
(f(z)) - f(x22) S -22l-
Lemma 2 (Jensen's Inequality): If f(w) is a convex function on W, and B[f(w)] and of f(B[w]) finite, then:
B[f(w)] 2 f(E(w))).
Definition 1 (Projection Operation): Assume w is intermediate result of optimization, define a project operator (w.) to project w. to the domain W, which is computed by:
II (w.))= =ww.)aarg min || l|w-ww -w.ll- W
wEW
Definition 2 (Diameter of Domain): Given function f(w), where w E W, and W is f's domain of definition. The diameter of W is denoted by for every W1, W2 E W: [[] wall S.
Theorem (Guaranteed Convergence Rate): If convex function f(w) is L-Lipschitz continuous function, then [[V](() S L. Let T be the diameter of domain. Applying equations (B))) for gradients aggregation, we have the following convergence rate for the proposed FeDEC algorithm:
2
3 r.?),
f(w?)
min f(w) 0 wEW
(11)
2TT
where T. the average result of w for total training epoch T, B the learning rate equation- (().ionnn())))))))))) and T is the total training epoch. If we let 3 = T
LVT the convergence rate is O +r)
To simplify the analysis, we consider fixed learning rate B3. The proof includes the following steps:
Proof:
(1) According the definition of convex function,
S(w') f(w) < V((w")(w' -w)
12

---

Under review as a conference paper at ICLR 2021
(2) We define G(()) (f"(w"(w''-w and = _== in update time + is denoted by w With assumption and assumption 2, we have:
= Tagi. The intermediate result of f(w)
G(()
(w'- w )w' -w) ww)
(ww)) w'w wf + + ww)
23 (ww')² w'w + w² (w + 2ww ++1 w2 + +(wl)² - 2w'w w+1 + (ww+1)3
1
28 (l/w' w/12 1|w w/l2 + 1/wl 113)
23 (//w w/1 || w/13) - 2 V (V
2B (I/w' w/1 || w/l3) 2 |l..
So we have:
f(w") f(w) I (//w" w/1 1| - + 3 2
1
=11"111.
2
(3) We project w w With definition 1 and using non-expandable property of projection operation of convex set, we have:
1
(w(w) S (llw* w/1 -optt-w/lp + =1l1/11
233
Due to L-Lipschitz Continuity (Lemmaa 1), we have:
G(w) S (1|w' w|l 1| wwl3) 11+)+ 8 L
1
23
2
So have:
23 (l/w' - -// || ++1 w| + B 2
f(w') - f(w) S
2
(4) According to definition 2. summing up all from t 1 to T. have:
T
B 1.2T
Σ
ww/3) + (ww) - T[(() x 23 (/|w] - ww12 - -l|wtt1-w133
2
1=1
23 w' -w11 + 8 L'T
r +
LTT.
23
(5) According to Jensen's Inequality (Lemma 2), we have:
T
f(w") s(w)
f
Σ w²) - f(w)
T
Σ
f( f(w*) - f(w)
T
t=1
22
B I.2
2
2
We can get the result:
1
3
f(wT)
min f(w) < 0
1.2).
wew
23T
2
13

---

Under review as a conference paper at ICLR 2021
(6) Taking 3 B=I/(L)T), the right part of the above equation becomes
r 3 L 12=
I'LNT
r
L' 1²= TL
2
2
2T
2LVT
VT
Therefore we can obtain the simplified expression of the convergence bound O 1
)
/T
14

---

Under review as a conference paper at ICLR 2021
A.2 EXPLANATION UNBIAS PARAMETER AGGREGATION SECTION 3.2.2
We compute the expectation of the aggregated parameters W mean and W bar in Section 3.2.2 as follows.
K
E|
W
mean]
E
W
meen
k=1
A&
W meen.
km
K
W var
[1|x || 1)TH
W var
|
K
(1|xx|1 - 1) T) W
X
[||x||] - 1) [wit
X
var
According the above equations, if the parameters of the local models W unbias, then the aggregated model parameters are unbias as well.
mean and W var
15

---

Under review as a conference paper at ICLR 2021
A.3 STRUCTURE OF THE NEURAL NETWORK MODELS SECTION B
Here we report the detailed model structure used in the experiments. We use LeNet shown in Table 2 and the 4-layer CNN model shown in Table El We adopt slim ResNet18 shown in Table HL where "Conv2d" is convolution layer, "BatchNorm2d" is batch normalization layer, and "Linear" is fully-connected layer. We can observe that every convolution layer followed by batch normal- ization (BN) layer. For all models, we use ReLU layer after every Conv2d layer. The structure of DenseNet1210 and MobileNetV2 can be found in GitHub.
For language model, we consider the sentiment analysis task on tweets from Sentiment140 with 2-layer BiLSTM. The BiLSTM binary classifier containing 256 hidden units with pretrained 100- dimentional GloVe embedding. Each twitter account corresponds to device.
Table 2: Detailed information of the LeNet architecture.
Parameter
Shape
Layer haper- parameter stride= padding=2
Conv2d 16 x 16 x 2
BatchNorm2d
MaxPool2d
stride=2, kernel size=2 stride= padding-2
Conv2d
16 x 16 5 x 5
BatchNorm2d
x
Conv2d
16 x 32 5 x 5
stride= padding=2
BatchNorm2d
32 x
MaxPool2d
stride=2, kernel size=2 fiid==1,padding.
Conv2d
32 x 32 x x 5
BatchNorm2d
32
Linear
1568 x 512
bias=True
Linear
512 128
bias=True
Linear
128 x 10
bias=True
Table 3: Detailed information of the 4-1. CNN architecture.
Parameter
Shape
Layer hyper- parameter
Conv2d
3 32 x x
BatchNorm2d
32
Conv2d
32 x 32 x x 5
BatchNorm2d
32 x
MaxPool2d
stride=1, kernel size=2
Conv2d
32 64 x 3
BatchNorm2d
64 x
Conv2d
64 64 x 5
BatchNorm2d
64 x 2
MaxPool2d
stride=1, kernel size=2
Linear
1024 X 128
bias=True
Linear
128 x
bias=True
com/ http:///ww.....cooooooooocccco/a wttp://w....cttttoomommmmmmmmmem
16

---

Under review as a conference paper at ICLR 2021
Table 4: Detailed information of the ResNer18 architecture.
Parameter
Shape
Layer hyper-parameter
Conv2d
3 64
3
stride=1, padding=1
BatchNorm2d
64 x
Conv2d
64 64
x
stride=1, padding=1
BatchNorm22d
64 x
Conv2d
64 64 x
stride=1, padding=1
BatchNorm2d
64 X
Conv2d
64 64 x
stride=1, padding=1
BatchNorm2d
64 x
Conv2d
64 64 3
stride=1, padding=1
BatchNorm2d
64 x
Conv2d
64 X 128 3 x
stride=2, padding=1
BatchNorm2d
128 2
Conv2d
128 X 128 x 3 x
stride=1, padding=1
BatchNorm2d
128 2
Conv2d
128 x 128 x 3 x
stride=1, padding=1
BatchNorm2d
128 2
Conv2d
128 x 128 x 3 x 3
stride=1, padding=l
BatchNorm2d
128 2
Conv2d
128 x 256 x 3 x 3
stride=2, padding=1
BatchNorm2d
256
Conv2d
256
256 x 3 x 3
stride=1, padding=1
BatchNorm2d
256
Conv2d
256 x 256 x 3 x
stride=1, padding=1
BatchNorm2d
256
Conv2d
256
256 x 3 X
stride= padding=1
BatchNorm2d
256 2
Conv2d
256 512 x 3 3
stride=2, padding=1
BatchNorm2d
512 x 2
Conv2d
512
512 3 x
stride=1, padding=1
BatchNorm2d
512 x 2
Conv2d
512
512 3 3
stride=1, padding=1
BatchNorm2d
512 x 2
Conv2d
512 512 x 3 x
stride=1, padding=1
BatchNorm2d
512 x 2
Linear
512 X 10
bias True
17

---

Under review as a conference paper at ICLR 2021
A.4 CONVERGENCE OF FEDERATED LEARNING ALGORITTMS FOR DIFFERENT MODELS ON DIFFERENT DATASETS
FwDEC FacMM
Feddug q-Fed5GD
Peclug FedProx prC
F-eccC FwiM
Pedlug FeeProx FwBEC
&
FeeMA
FedPrex
96
192
288
106
212
318
424
0
90
180
270
360
Epoch
Epoch
Epoch
CIFAR-IO (ResNet18)
(b) CIFAR-IO (DenseNet121).
(c) CIFAR-I0 (MobileNetV2))
Figure 9: Convergence of different algorithms for different CNN models on CIFAR- 10.
FeDEC
FesAvg
FeDEC
FedAyg o-Fec5GD
FedHA
c-Fes5GD
FedMA Fedivox
Zeno
RMN
AMN
16
48
16
24
Epoch
Epoch
(a) Aishion-MNIST (LeNet).
(b) MNIST (LeNet).
Figure 10: Convergence of different algorithms for LeNet FMNIST and MNIST.
18

---

Under review as a conference paper at ICLR 2021
A.5 TRAINING EFFICIENCY OF FEDERATED LEARNING ALGOIIHTMM FOR DIFFERENT MODELS ON DIFFERENT DATASETS
0.9
Test Accursry
0.9
0.8
0.7
0.7
0.6
18 27 36 45 54 63 72 81 (min)
46
92
138
184
30
60
90
120
Time (min)
(min)
(a) CIFAR-10 (ResNet18..
(b) CIFAR-10 (DenseNet111)
(c) CIFAR-10 (MobileNetV2).
Figure 11: Training accuracy of different algorithms for different CNN models on CIFAR-10 10.
Accuracy
Test Accuracy
0.9
PLDEC
Polkyy elessoo
FcDBC FedMA
hally g
RNN
12
32 36
(min)
(min)
(a) Fashion MNIST (LeNet).
Fashion-MNIST
(b) MNIST (LeNet).
Figure 12: Training accuracy of different algorithms for LeNet on FMNIST and MNIST.
19

---

Under review as a conference paper at ICLR 2021
Algorithm 1: FeDEC Aggregation
Server
for I= to T do
Transmit W to all clients.
Receive g", ft, o from all clients.
Inference o with & based on fa, from Wt
meons W
ear
Update W N. with B to get W NN-
mean with R and W ear with B. Combine WNN, W
Aggregate W
mean and W var to new model W
Client
for 1= to T do
11
Receive server model W
12
Train local model W based on local dataset and get local gradients g'.
13
Transmit g" and f. & to server.
Stop training until received server model W
14
0.4
Client-A
Client-B
Percentage
0.3
0.2
0.1
0.0
True & Inferenced Distribution
Figure 13: Variational inferenced approximate distribution of two clients of CIFAR-10 dataset ex- periment.
20


Reviews:
[{"review_id": "sPEOWJ789pe", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this work, the authors propose new ways of averaging updates received at the server from a subset of clients in a federated scenario. Specifically the authors aim to address issues arising from the non-iid nature of the data that arises in FL and propose to treat BatchNorm parameters differently from other NN parameters. \n\nIntroduction\nReading this paper, I am confused about terminology used by the authors. Specifically, the authors discuss 'data bias' and 'parameter bias'. \nIn in the introduction, the authors claim that 'Conventional approaches average gradients uniformly from the clients, which could cause great bias to the real data distribution'. Assuming the authors understand FedAvg to be a conventional approach, the averaging with $|x_k|/|x|$ (Eq. 1) is not uniform but weighted by the local dataset size. Further, the meaning of 'causing bias to the real data distribution' is not clear to me. Unfortunately, the further explanation and Figure 1 don't help me in understanding what is meant. The 'GroundTruth' distribution of labels for cifar10 is approximately uniform. From the context, I understand that the authors try to describe some consequence of a non-i.i.d sampled distribution of data according to labels, but I cannot understand the point they try to make. \nNext, the authors discuss 'parameter bias'. The authors distinguish between BN parameters and other NN parameters. They term the BN parameters 'statistical parameters such as mean and variance'. Generally speaking, BN contains the 'scale and shift' parameters $\\gamma$ and $\\beta$ (https://arxiv.org/pdf/1502.03167.pdf), which I am assuming the authors reference to. Again, the authors make use of the term 'bias' to say: '[...] bias on the BN layer parameters'. I am not familiar with the notion of 'bias on a parameter' and would like the authors to clarify. Based on Figure 2 I assume they aim to convey that BN parameters in a FL setting converge to different values compared to a centrally trained model. \nIn Section 3.1 the authors further make the distinction explicitly between 'gradient parameters', by which they mean weights and biases as opposed to 'statistical parameters' of BN. Since the scale-and-shift parameters of BN are also updated by gradient descent, I am wondering if the authors mean the mean and variance estimate across data-points for feature-maps, which also plays a central part in (federated) BatchNormalization. The authors in their experimental section make no mention of how they form these global estimates for mean and variance in BN, so they omit that crucial detail there.\n\nThe authors specifically focus on label-skew as source of non-i.i.d-ness in this work, but never make this limitation explicit. Since the non-i.i.d challenges in FL are not limited to label skew, I believe the authors should make this explicit. \n\nRelated work:\nThe issues with BN in the federated setting has been described for example in Section 5 of https://arxiv.org/pdf/1910.00189.pdf. There, the authors propose to replace BN with GroupNorm, an approach that has been adopted in several follow-up and recent works in FL with models that originally contain BN. I would encourage the authors to compare their work against this approach, both in the RW section and also in the experimental section.\n\nMethod Section.\nNotation-wise, I encourage the authors to not use $x$ or $x_k$ to denote a (labeled) dataset, since $x$ is usually reserved for a single data-point with associated label $y$. In Eq. (1) the loss formulation of FL is a bit sloppy since the parameters to optimise for, $W$ do not appear in the RHS of the equation. In FedAvg, we explicitly optimise $min \\sum_k |x_k|/|x| L_k(W,x_k)$, where the local parameter estimates $W_k$ appear as intermediate parameters as a consequence of multiple local optimisation steps. \nThe authors claim that 'data points available locally could be biased from the overall distribution'. Again, I believe to understand the intended meaning to be the non-iid issue, but I encourage the authors to make their understanding of 'bias' more concrete. \n\nIn FedAvg, the individual clients do not transmit gradients $\\nabla L_k ()$ to the server (Section 3.2.1). This is the approach in conventional distributed SGD as employed in a high-speed-connected data-centre for speeding up centralised training. In this centralised setting, the non-iid problem does not exist. Instead, in FL, clients transmit parameters that have been updated through a series of gradient-descent steps. More recent work (https://arxiv.org/abs/2003.00295) makes the role that these transmitted parameters have in an interpretation as a gradient. This distinction is important.\n\nThe derivation in Appendix A.2 for show-casing the unbiased-ness of the variance parameter averaging seems wrong. Going from the third to fourth equation makes a mistake and also if the original expectation was equal to the sum of weighted expectations, then simply averaging would actually be the unbiased estimator. The authors here are falsifying their own argument through a derivation mistake. The last equation should only pull the expectation into the sum in the right-most term and you are done.\n\nSection 3.2.3\nI like the notion of modelling the datasets as GMM and to infer responsibilities at averaging time. The explanation of the approach is confusing to me, however. \nIf I understand it right, then EQ 9 describes a VAE setup per client k. There is no sharing of parameters or latent space between clients. $s_k = [\\mu_k,\\sigma_k]$ are the mean and standard-deviation across the whole local dataset at a client $k$. The authors propose to encode this single vector $s_k$ into a latent space z_k of dimension $C$. The authors do not explicitly specify the prior p(z_k), but given the constraints on $z_k$, I assume it is meant to be a Dirichlet distribution. From context I could imagine that it has something to do with the per-client label distribution. \nSince each ELBO is per-client and each client has just one data-point $s_k$, I do not understand the need for auto-encoding. Simply infer z given a decoder-model for the single data-point.  The authors mention the use of neural networks, but they do not detail their architecture choices anywhere. It is also unclear to me how $\\pi_k$ falls out in equation 10. I imagine it I corresponds to some sort of posterior across all local models' encoding z. Since the latent-spaces across clients k do not share any meaning, I don't see how that can be sensible.  All together, this section is not readable to me. \nI can see the appeal of reweighing updates through a specific formulation pi_k, but since the updates are label-independent (eq. 3,4,5), how does that play into this.\n\nSection 4\nMy understanding of proofs of this form is somewhat limited. From what I can gather, this proof shows that FedAvg converges if the per-client loss-function is pre-multiplied by a constant factor $\\pi_k$. As such, the convergence proof should be analogous to what is presented in e.g. https://arxiv.org/pdf/1907.02189.pdf with the exception of the update in Eq 5. Maybe the other reviewers can comment further on this. \n\nSection 5.\nIn the Federated Setting, 20 clients should be considered not enough for experimental validation generally speaking. I appreciate the breath of experiments in terms of models, algorithms and related algorithms. \nUnfortunately, the authors chose to define their own non-iid-split of the datasets. In general, I would appreciate to see comparisons with existing data-set splits in related work on non-iid data, such as for example https://arxiv.org/abs/2003.00295. This would help avoid the bifurcation of the literature. \nThe reference to q-FedSGD seems to be wrong.\n\nAfter reading the experiment section, some serious questions arise:\nWhat is the fraction of selected clients per communication round? What is the number of local epochs per client? Some of the related works seem to suggest performing FedSGD (also the first paragraph in Section 3.2.1 suggests this). If you have single-gradients per device and equal-size data-sets per client, then I don't see how non-iid data-distribution across clients are an issue as this approaches a global mini batch step, where each mini-batch consists of smaller mini-batches, one from each client. The non-i.i.d issue in FL stems from the fact that each client optimises on its own for a sufficiently long time that the resulting progress is destroyed by averaging in parameter space. The authors need to specify their setup here. Concretely, I would want to at least see Cifar10 split into 100 clients, 10 of which are selected at every round. Each client needs to optimise locally for a full epoch on its own dataset of 45000/100 = 450 data-points. \n\nThe authors propose two things: A new averaging approach by computation of $\\pi_k$, as well as a new approach to estimating the gradient for the scale-parameter of BN using pooled averaging. These two things need to be studied separately by setting $\\pi_k = |x_k|/|x|$ in one scenario. At the moment, my trust into the computation of $\\pi_k$ is very low, since the corresponding section is not understandable to me. The pooled averaging approach seems sensible and I am curious to see if it solves the issue of BN in FL. Additionally, the authors need to clarify how the BN-statistics are computed at test time. I would also like to see a comparison of the proposed models with BN and with the state-of-the-art method which is replacing them with GroupNormalization. \n\nFinally, I would like to thank the authors for the interesting approach to training BN-equiped Neural Networks in a federated setting. This idea seems promising to me. The approach for computing pi_k is not clear to me and I would like the authors to revise it. Please specify the role that knowledge of the label-distribution plays and if the method is applicable when non-iid-ness stems from other sources than label skew. (I propose looking at the FEMNIST dataset for example). \nMany issues with this paper remain and I encourage the authors to overhaul their work. \n\nI see no issues with the Code of Ethics", "extra_scores": {}, "extra_text": {}}, {"review_id": "YsCPK2SAvIW", "reviewer": "AnonReviewer2", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\n\nSummary:\n \nThe paper proposed FedDEC, a novel approach to conduct model updates aggregation in federated learning. The main motivation of this paper is to decouple the aggregation of normal model weights and statistics in BNs separately such that both data and model heterogeneity can be handled. Theoretical analysis indicates that the proposed FedDEC method enjoys a good convergence guarantee. Extensive experimental results are provided to show that FedDEC enjoys high efficiency and better model accuracy under the non-IID environment compared to the considered baseline methods.\n\n##########################################################################\n\nReasons for score: \n \nOverall, I think the current manuscript is marginally above the acceptance threshold of the ICLR conference. Studying the effectiveness of the model aggregation process in federated learning is a promising research direction. The proposed approach to conducting the decoupled model aggregation over the model parameters and statistics aggregated in the BN layers is promising. The proposed FedDEC method is justified both theoretically and empirically. Theoretically, the authors show that FedDEC enjoys the same convergence rate as normal SGD. Extensive experimental results over the simulated federated learning environment indicate that FedDEC enjoys good training effectiveness while reaching to better model accuracy. If my concerns on the current manuscript (please see \"Cons\") are addressed, I will be happy to improve my evaluation score.\n \n##########################################################################\n\nPros: \n \n1. The paper is well written. The research direction on improving the model aggregation in federated learning is promising and practical.\n \n2. The formulation and theoretical analysis of the proposed FedDEC method looks promising.\n \n3. Extensive experimental results are provided for the image classification tasks under the simulated non-iid environment, which indicates that FedDEC enjoys high effectiveness in improving the model test accuracy under the data heterogeneity.\n \n\n##########################################################################\n\n\nCons: \n \n1. A detailed discussion of the algorithmic aspect of FedDEC is missing. Section 3 discusses the formulation of FedDEC, however, it’s not super clear on the exact steps that FedDEC will run e.g. what information will the clients upload to the data center?; what are the exact operations that the data center will conduct? The authors’ are highly encouraged to add an algorithm box to describe the FedDEC algorithm. \n2. From the description in Section 3, it seems FedDEC always assumes the local clients to train only ONE local epoch. If that’s the case, then the communication efficiency of FedDEC can be worse than FedAvg due to higher communication frequency. A more detailed discussion on this can be helpful to understand FedDEC better. \n3. The convergence rate of FedAvg has been explicitly studied in [1]. It would be helpful to make a clearer comparison between the convergence rates of FedDEC and FedAvg. \n4. In Section 3.2.2., the mean and variance in BN are referred to as BN layer parameters. However, in a BN layer, there are two trainable weights (\\gamma + \\beta [2]) and mean + variance statistics collected over the training batches. It would be helpful to make the notations clearer. \n5. Although language models usually do not have BN layers, it’s still easy to imagine that the FedDEC algorithm can be applied over language tasks. Adding a language processing task can make the current manuscript stronger. \n\n[1] https://arxiv.org/pdf/1907.02189.pdf\n\n[2] https://arxiv.org/pdf/1502.03167.pdf\n \n\n#########################################################################\n\n\nMinor Comments: \n\n1. Typos in the draft: (i) in the abstract “suffer form” —> “suffer from“; (ii) “$W_{NN}$ (parameters of parameters of BN layers)” —> “$W_{BN}$ (parameters of parameters of BN layers)” in Section 3.2. \n2. Missing references: [1-4]. \n\n[1] https://arxiv.org/pdf/1907.02189.pdf\n\n[2] https://arxiv.org/pdf/1812.01097.pdf\n\n[3] https://arxiv.org/pdf/1908.07873.pdf\n\n[4] https://arxiv.org/pdf/1912.04977.pdf", "extra_scores": {}, "extra_text": {}}, {"review_id": "TgHqjnG2DX", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "1/ Summary\n\nThis paper introduces an aggregation mechanism designed for neural networks with batch normalisation layers. This mechanism relies on two parts: probabilistic mixing weights of the loss function and the use of a weighted pool estimator for aggregating the BN variance parameters. The mixing weights are derived from a GMM with variational inference. A convergence result in the *convex* case is provided. Experimental results on 3 image datasets show that this approach yields better results than other standard FL algorithms (FedAvg, FedProx, q-FedSGD, FedMA…) as well as a better resilience to heterogeneity (understood as class imbalance).\n\n2/ Acceptance decision\n\nDespite its seemingly good experimental results, I am in favour of rejecting this paper as correcting its weak points would require major changes.\n\n3/ Supporting arguments\n\nA/ The theoretical results (Sec 4, one page) are irrelevant to the problem tackled by the paper.\n\nIndeed, they are proved in the case of a convex function. However, the proposed method is relevant for neural networks with batch normalisation layers, so more than 1 layer, and therefore these networks yield non-convex loss functions, which do not satisfy the hypothesis of the results. Further, the unbiased gradients assumption (assumption 1) seems dubious in the non-iid case tackled by the authors. Last, but not least, it is difficult to understand if this assumption is valid or not (and the same holds for the results) because the notations used to state the theorem in Eq 11 do not correspond to the notations of the loss function Eq 2. \n\nB/ Important technical details are omitted, which makes it difficult to understand the whole algorithm, threatens reproducibility, and shades some doubts on the experimental results.\n\n- In section 3.2.3, the authors introduce a GMM for the mixing weights \\pi_k while dropping the upper script t for convenience, which is understandable. However, it is not stated how the estimation process of \\pi_k^t evolves throughout the different batches. Does one need to do 1 full variational inference at each aggregation step? Can one re-use previous parameters through aggregation rounds? Does one rely on constant \\pi_k?\n- A related question is the behaviour of the proposed method when reducing communication frequency. Indeed, the algorithm is only exposed in the case of communication after each local gradient computation, which is typically avoided in the edge setting considered by the authors. How is it extended in this case? In particular, how does the estimation of the parameters \\pi_k change?\n- In a related note, the number of local computations done before any aggregation is omitted in the experimental section. Are all models compared with the same number of local steps?\n\nC/ Writing is not very clear and should be improved. In particular, some notations do not make sense; notations tend to change a lot across the different sections, making it difficult to see a global picture emerge.\n\n4/ Additional comments\n\n- Delving into appendix A.2 which explains the aggregation rule for variance layers, I do not understand the transition from the 3rd equation to the 4th (from top to bottom). Could the authors elaborate more on it?\n- Although the approach is « decoupled », the mixing weights used in the aggregation rules (Eq 3) (standard layers) and ( Eq 4) (BN means) are the same. Have the authors tried completely decoupling these mixing weights?\n- In Section 3.2.3, if I understood correctly the means \\mu_k and \\sigma_k are related to the empirical means and variances of the batches in the different clients. Sharing means and variances of each batch could yield to large privacy leaks. This remark goes together with the communication frequency remark above.\n- In Section 3.2.3, the authors state that they assume that data are drawn from a Gaussian distribution « without loss of generality ». However, it is well known that e.g. image distributions are not Gaussian.\n- In Section 5.1, I do not understand the notation $pr(x) = \\mathcal{N}(0, 1)$. How can a probability take negative values?\n- How does performance of the different methods vary with respect to the communication frequency? This is an important metric of the training efficiency, and missing from the experiments.\n- Given that the experiments have been run multiple times, it would be nice to provide error bars to understand the significance of the different results.\n- Some typos:\n    - page 4,  $W_{BN}$ instead of $W_{NN}$ at the beginning of Sec 3.2\n    - \"Form\" -> \"from\" in the introduction\n- The presentation of figures 4, 5, 6 make it difficult to read the captions.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "5UndtHcTAac", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "PAPER SUMMARY\n\nThis paper presents a new federated learning scheme for neural network that accounts simultaneously for different layer types and different (heterogeneous) local data distributions. In particular, the proposed method differentiates between parameter updates for conv, fc versus bn (batch-norm) layers. \n\nThe main argument for this differentiation is that aggregating parameter gradient for bn layers is vulnerable to high bias when local data are heterogeneously distributed. This leads to a new proposal on aggregating bn parameters directly (instead of using their gradient).\n\nBoth parameter gradient and parameter updates are weighted by the same set of probabilistic weights that seem to be associated with mixture weights of a GMM model that generate local training data. These weights are learned via a deep generative model that encode local statistics of local data into a latent space, which seems to be associated with random drawn from a (latent) categorical distribution associated with those mixture weights. \n\nA theoretical convergence guarantee for the proposed algorithm is also provided.\n\nNOVELTY & SIGNIFICANCE\n\nThis paper poses an interesting view on the parameter aggregation of federated learning. I appreciate the authors' perspective that perhaps the aggregation of parameters defining different layer types need to be diversified to fit the nature of the data. \n\nHowever, I find the proposed solution and its fundamental idea of modeling such heterogenity somewhat arbitrary and perhaps flawed. To elaborate:\n\nFirst, it is not at all clear to me why the authors associate the weights that combine parameter gradients (for fc and conv) and parameters (for bn) with the mixture weight that generate local data. This seems like an arbitrary choice to me and there seems to be a flaw here: if a local device receives very little data but its data come from a mixture component with large weight, its gradient will likely be biased (due to the lack of data) but will still dominate others (due to its large mixture weight) -- I would like to hear the authors' thoughts on this.\n\nSecond, if I understand correctly, the auto-encoder model devised to learn those mixture weights is supposed to encode the local statistics (i.e., mean & variance estimates of a Gaussian that fit the local training input) to a categorical distribution over label classes. I am not sure if I missed something here but this is a bit strange: by right, the mean and variance of the *input* distribution do not contain any information about its *label* -- how does the embedding manage to generate information that does not exist in the first place? Likewise, how would such *label* information be related to the mixture weights of the *input* distribution via Eq. (10) -- could the authors elaborate further on this?\n\nThird, on a high-level of idea, I am also not fully convinced that the aggregating the bn parameter gradient is vulnerable to high bias. This statement somehow came across as a politically correct statement with no concrete substantiation. It will be better if there is some technical demonstration that explicitly show that for bn layer, the bias of the aggregated gradient is higher than the bias of its aggregated parameters.\n\nLast, I find the theoretical analysis is not at all particular about the inner working of the proposed algorithm. The analysis is on one hand generally applicable to any function that satisfies the stated assumption and, on the other hand, not specific enough to account for the facts that the combination weights are also being updated and that the bn parameters are not updated via gradient aggregation. Given this, I do not think the analysis is applicable to this setting, and it should be evaluated separately in its own setting.   \n\nTECHNICAL SOUNDNESS\n\nI have made high-level check over the main bulk of technical derivation and have not spot any glaring issues. But, as I said above, on the idea level, the three technical contributions here (separating bn aggregation from those of fc and conv, learning the mixture weights from local statistics (barring access to data) and the theoretical analysis) are orthogonal with little connection to one another -- for the first two contributions, there seems to be flaws on the idea level that need further clarification.\n\nEXPERIMENT\n\nThe reported results appear positive but it seems the authors only generate those for one single run. As the improvement margin is relatively small, it is important to average results over multiple runs to make sure the improvement is significantly above the deviation margin.\n\nFurthermore, reporting only the final performance gives very little insight regarding how the invented components help avoid accumulating high bias, and also whether the auto-encoding scheme correctly recover the mixture weight -- perhaps this can be shown by evaluating the proposed mechanism on a synthetic dataset where we know the ground-truth. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors’ feedback has not fully addressed the reviewers’ concerns and the reviewers think that the paper is not ready for the publication. The authors should consider the following issues for the future submission:

1) The concern from Reviewer 1: if a local device receives very little data but its data come from a mixture component with large weight, its gradient will likely be biased (due to the lack of data) but will still dominate others (due to its large mixture weight).

2) Numerical experiments are not consistent with theoretical results. The theory is for convex but experiments are with non-convex loss. The response from authors does not resolve this issue. 

3) Notation is confusing and changing throughout. We strongly suggest the authors revise carefully this and make it clear. 

Although the experimental results are potential, we would like the authors to revise it carefully by addressing the reviewers’ concerns and further improve it by considering theoretical results for non-convex in order to submit to the next venues. 

Decision: Reject