Year: 2020

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.

heading: 1 INTRODUCTION
text: Federated learning (FL) has emerged as a novel distributed machine learning paradigm that allows a global machine learning model to 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/gradients from local models to form a global model without transferring data to a central server. Federated learning has been drawn much attention in mobile-edge computing (Konecný et al. (2016); Sun et al. (2017)) with its advantages in preserving data privacy (Zhu & Jin (2020); Jiang et al. (2019); Keller et al. (2018)) and enhancing communication efficiency (Shamir et al. (2014); Smith et al. (2018); Zhang et al. (2013); McMahan et al. (2017); Wang et al. (2020)).
Gradient aggregation is the key technology of federated learning, which typically involves the following 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 to 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 et al. (2017)) and its variants such as FedProx (Li et al. (2020a)), Zeno (Xie et al. (2019)) and q-FedSGD (Li et al. (2020b)) applied the synchronous parameter averaging method to the entire model indiscriminately. Agnostic federated learning (AFL) (Mohri et al. (2019)) defined an agnostic and risk-averse objective to optimize a mixture of the client distributions. FedMA (Wang et al. (2020)) constructed the shared global model in a layer-wise manner by matching and averaging hidden elements with similar feature extraction signatures. The recurrent neural network (RNN) based aggregator (Ji et al. (2019)) learned an 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. It was reported in (Zhao et al. (2018)) that the accuracy of a convolutional neural network (CNN) model trained by FedAvg reduces by up to 55% for highly skewed non-IID dataset. The work of
(Wang et al. (2020)) showed that the accuracy of FadAvg (McMahan et al. (2017)) and FedProx (Li et al. (2020a)) 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 we make the following observations.
Data Bias: In the federated learning setting, local datasets are only accessible by the owner and they are typically non-IID. Conventional approaches aggregate gradients uniformly from the clients, which could cause great bias to the real data distribution. Fig. 1 shows the distribution of the real dataset and the distributions of uniformly taking samples from different number of clients in the CIFAR-10 dataset (Krizhevsky (2009)). It is observed that there are great differences between the real data and the sampled distributions. The more clients involved, the more difference occurs.
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 statistical 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. 2 shows the means and variances in BN layer distribution of a centrally-trained CNN model and that of FedAvg-trained models with different number of clients on non-IID local datasets. It is 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 a 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 a decoupled probabilisticweighted method. To the best of our knowledge, we make the first attempt to aggregate gradient parameters and statistical parameters separatively, and adopt a probabilistic mixture model to resolve the problem of aggregation bias for federated learning with heterogeneous clients. (2) We propose a 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-IID conditions. It is shown that FeDEC significantly outperforms the state-of-the-arts in terms of model accuracy and training efficiency.

heading: 2 RELATED WORK
text: We summarize the related work as two categories: parameter/gradient aggregation for distributed learning and federated learning.
Distributed Learning: In distributed learning, the most famous parameter aggregation paradigm is the Parameter Server Framework (Li et al. (2014)). In this framework, multiple servers maintain a 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 al. (2012); Dekel et al. (2012); Richtárik & Takác (2016); Zhang et al. (2015)). Several works focused on
improving the communication-efficiency for distributed learning (Shamir et al. (2014); Smith et al. (2018); Zhang et al. (2013)). To address the issue of model robustness, Zeno (Xie et al. (2019)) was proposed to make distributed machine learning tolerant to an arbitrary number of faulty workers. The RNN based aggregator (Ji et al. (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 (Konečnỳ et al. (2015)) is an emerging edge distributed machine learning paradigm that aims to build machine-learning models based on datasets distributing across multiple clients. One of the standard parameter aggregation methods is FedAvg (McMahan et al. (2017)), which combines local stochastic gradient descent (SGD) on each client with a server that performs parameter averaging. The lazily aggregated gradient (Lag) method (Chen et al. (2018)) allowed clients running multiple epochs before model aggregation to reduce communication cost. For heterogeneous datasets, FedProx (Li et al. (2020a)) modified FedAvg by adding a heterogeneity bound on datasets and devices to tackle heterogeneity. The FedMA (Wang et al. (2020)) method, derived from AFL Mohri et al. (2019) and PFNM (Yurochkin et al. (2019)), demonstrated that permutations of layers can affect the gradient aggregation results, and proposed a layer-wise gradient aggregation method to solve the problem. For fair resources allocation, the q-FedSGD (Li et al. (2020b)) method encouraged a 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 decouple 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.

heading: 3 FEDEC: A DECOUPLED GRADIENT AGGREGATION METHOD

heading: 3.1 OBJECTIVE OF FEDERATED LEARNING WITH NON-IID DATA
text: Consider a federated learning scenario with K clients that train their local CNN models independently based on local datasets x1,x2, . . . ,xK and report their gradients and model parameters to a 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,x2, . . . ,xK}. Conventional federated learning tends to optimize the following loss function:
min W
L(W,x) := K∑
k=1
|xk| |x| Lk(Wk,xk), (1)
where W is the parameters of the global model, Wk (k = 1, 2, · · · ,K) is the parameters of the k-th local model; L(·) and Lk(·) 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-weighted of the local losses.
As discussed in section 1, conventional federated learning has two drawbacks. Firstly, local datasets are collected by mobile devices used by particular users, which are typically non-IID. Training samples on each client may be drawn from a different distribution, therefore the data points available locally could be bias from the overall distribution. Secondly, since a neural network model is typically consists of convolutional (Conv) layers and full-connected (FC) layers that are 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 probabilistic-weighted approach for federated learning that focuses on optimizing the following loss function:
min W∗
L({WtNN ,Wtmean,Wtvar},x) := K∑
k=1
πkLk({Wt−1,kNN ,W t−1,k mean,W t−1,k var },xk), (2)
where ∗ indicates NN,mean and var; WtNN , Wtmean and Wtvar are the parameters of Conv and FC layers of the global model after t-th aggregation epoch; Wt−1,kNN , W t−1,k mean and W t−1,k var are the k-th local model been trained several local epoch based t− 1-th global model; πk(k = 1, . . . ,K) is
the probability that a sample is drawn from the distribution of the k-th client, i.e., πk ∈ [0, 1] (k = 1, . . . ,K) and ∑K k=1 πk = 1.
The above formulation objects to minimize the expected loss over K clients with non-IID datasets. Next, we will introduce a decoupled method called FeDEC to optimize the parameters of different types of layers separatively, and derived the probability weights πk for parameter aggregation.

heading: 3.2 DECOUPLED PROBABILISTIC-WEIGHTED GRADIENT AGGREGATION METHOD
text: In this section, we proposed a decoupled method to derive the global model with respect to WtNN (parameters of Conv and FC layers) and Wtmean,W t var (statistical parameters of BN layers).

heading: 3.2.1 GRADIENT AGGREGATION FOR CONV AND FC LAYERS
text: Since the parameters of Conv and FC layers are neural network weights which are updated by distributed gradient descent method (Nesterov (1983)), they are appropriate to be aggregated with a similar approach that adapts conventional federated average for non-IID datasets. Let gtk = W t−1,k NN∗ − W t−1,k NN (k = 1, . . . ,K), where NN∗ indicates NN parameters after full local training. be the gradient of the k-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.
WtNN = W t−1 NN − β K∑ k=1 πtkg t k, (3)
where β is the learning rate for parameter update, πtk (k = 1, . . . ,K) are the probabilistic weights with ∑K k=1 π t k = 1 that are derived in section 3.2.3.

heading: 3.2.2 PARAMETER AGGREGATION FOR MEANS AND VARIANCES IN BN
text: 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 Wtmean,W t var in BN layer, which are updated by:
Wtmean = K∑ k=1 πtkW t,k mean, (4)
Wtvar = 1
|x| −K K∑ k=1 (|xk| − 1)πtkWt,kvar, (5)
where Wt,kmean and W t,k var indicate the means and variances in BN layers of the k-th client in epoch t; πtk (k = 1, . . . ,K) are probabilistic weights with ∑K k=1 π t k = 1 that are derived in section 3.2.3.
In the above equations, we update the mean with the weighted average of local models, and update the variance with the weighted pooled variance (Killeen (2005)), which can give an unbias estimation of parameters of the whole dataset under non-IID conditions (see Appendix A.2).

heading: 3.2.3 DERIVATION OF PROBABILISTIC WEIGHTS
text: We adopt a mixture probabilistic model to describe non-IID datasets in federated learning. Without loss of generality, in the t-th training epoch, we assume the mini-batch samples of each client follows a Gaussian distribution Nk(µk, σk) (k = 1, . . . ,K), where µk, σk are the mean and standard deviation of the distribution that vary among clients. We omit the upper script t for simplicity thereafter. The whole samples can be described as a Gaussian Mixture Model (GMM) with the following probability function:
p(x|λ) = K∑
k=1
πkp(xk|µk, σk), (6)
where λ = {πk, µk, σk | k = 1, 2, · · · ,K} are the parameters of the GMM model1. 1Noted that the proposed variational inference method can be applied to other non-Gaussian distributions with slight modification.
In federated learning, the local data samples are accessed by particular client and the central server can only observe the statistics of local dataset such as mean and standard variance. Without knowing the overall samples, conventional expectation-maximization (EM) algorithm (Dempster et al. (1977)) cannot be applied to derive λ. Alternatively, we introduce a variational inference method to estimate the parameters of λ.
1
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. 3. The notations are explained as follows.
•stk = {Wt,kmean,Wt,kvar} is the observed statistics from the feature maps of k-th client.
• ztk = {ztk,i|(i = 1, 2, · · · , C)} is a vector of latent variables with length C, where ztk,i ∈ [0, 1], ∑C i=1 z t k,i = 1, and C is the number of classes for a classification task. ztk can be viewed as a data distribution that represents the probability of a sample in client k belonging to the classes.
• θ = {θk} are generative model parameters, and ϕ = {ϕk} are variational parameters.
The solid lines in Fig. 3 denote the generative model pθk(z t k)pθk(s t k|ztk), and the dashed lines denote the variational approximation qϕk(z t k|stk) to the intractable posterior pθk(ztk|stk) We approximate pθk(z t k|stk) with qϕk(ztk|stk) by minimizing their divergence:
ϕ∗k,θ ∗ k = arg min
θk,ϕk divergence(qϕk(z
t k|stk) || pθk(ztk|stk)),
s.t. ∑C
i=1 ztk,i = 1.
(7)
To derive the optimal value of the parameters ϕk and θk, we compute the marginal likelihood of stk:
log p(stk) = DKL(qϕk(z t k|stk) || pθk(ztk|stk)) + Eqϕk (ztk|stk)
[ log pθk(z t k, s t k)
qϕk(z t k|stk)
] . (8)
In Eq. 8, the first term is the KL-divergence (Joyce (2011)) 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(stk)is non-negative, the minimization problem of Eq. 7 can be converted to maximize the ELBO. To solve the problem, we change the form of ELBO as:
Eqϕk (ztk|stk) [ log pθk(z t k, s t k)
qϕk(z t k|stk)
] = Eqϕk (ztk|stk) [ log p(ztk)
qϕk(z t k|stk) ] ︸ ︷︷ ︸
Encoder
+Eqϕk (ztk|stk)[log pθk(s t k|ztk)]︸ ︷︷ ︸
Decoder
.
(9) The above form is a variational encoder-decoder structure: the model qϕk(z t k|stk) can be viewed as a probabilistic encoder that given an observed statistics stk it produces a distribution over the possible values of the latent variables ztk; The model pθk(s t k|ztk) can be refered to as a probabilistic decoder that reconstructs the value of stk based on the code z t k. According to the theory of variational inference (Kingma & Welling (2014)), the problem in Eq. 9 can be solved with stochastic gradient descent (SGD) method using a fully-connected neural network to optimize the mean squared error loss function.
With the derived optimal parameters ϕ∗k,θ ∗ k, we can extract the latent variables z t k that is interpreted as the sample distribution of client-k. Therefore ztk can be used to infer the parameters (πk, µk, σk) of k-th component of the GMM model. Specifically, the probabilistic weights πk can be represented by
πtk = { C∑ i=1 ztk,i∑K j=1 z t j,i }/{ K∑ k=1 C∑ i=1 ztk,i∑K j=1 z t j,i } . (10)

heading: 4 CONVERGENCE ANALYSIS
text: In this section, we will show that the convergence of the proposed FeDEC algorithm is theoretically guaranteed. We use the following assumptions and lemmas, and the convergence guarantee is provided in Theorem 1.
Assumption 1 (Unbiased Gradient): We assume that the stochastic gradients gti is an unbiased estimator of the true gradient ∇f(wti), i.e., E[gti ] = ∇f(wti), where f(·) is any convex objective function and wti is its variables.
Assumption 2 (Gradient Convex Set): We assume that gradient set G is a convex set, where all gradients g1,g2, . . . ,gK are in G, and any g = ∑K i=1 λigi (∀λi > 0 and ∑K i=1 λi = 1) is in G.
Lemma 1 (L-Lipschitz Continuity): For a function f(·) is Lipschitz continuous if there exists a positive real constant L such that, for all real x1 and x2:
|f(x1)− f(x2)| ≤ L|x1 − x2|.
Lemma 2 (Jensen’s Inequality): If f(w) is a convex function on W , and E[f(w)] and f(E[w]) are finite, then:
E[f(w)] ≥ f(E[w])).
Definition 1 (Projection Operation): Assume w∗ is an intermediate result of optimization, we define a project operator ∏ W(w∗) to project w∗ to the domain W , which is computed by:∏
W (w∗) = arg min w∈W ||w −w∗||.
Definition 2 (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 Γ: for every w1,w2 ∈ W: ||w1 −w2|| ≤ Γ. Theorem 1 (Guaranteed Convergence Rate): If a convex function f(w) is L-Lipschitz continuous function, then ||∇f(w)|| ≤ L. Let Γ be the diameter of domain. Applying equations (3)(4)(5) for gradients aggregation, we have the following convergence rate for the proposed FeDEC algorithm:
f(w̄T )− min w∈W
f(w) ≤ O( Γ 2
2βT +
β 2 L2), (11)
where w̄T is the average result of w for total training epoch T , β is the learning rate in equation-(3), and T is the total training epoch. If we let β = Γ
L √ T , the convergence rate is O( 1√ T ).
Proof skeleton: We provide a simple description of the proof skeleton of Theorem 1 with the following steps. (1) Since f(·) is a convex function, we have f(wt)−f(w) ≤ ∇f(wt)(wt−w). (2) With assumption 1 and 2, we have f(wt)− f(w) ≤ 12β (||w t −w||2 − ||wt+1∗ −w||2) + β 2 ||∇f(w
t)||2, where wt+1∗ is the intermediate result of f(w) in update time t+1. (3) With lemma 1 and definition 1, by projecting wt+1∗ to w t+1, we have f(wt)−f(w) ≤ 12β (||w t−w||2−||wt+1−w||2)+ β2L 2.
(4) Summing from t = 1 to T and with definition 1 and 2, we have ∑T
t=1 f(w t) − Tf(w) ≤
1 2βΓ 2 + β2L 2T . (5) According to lemma 2, we have f(w̄T ) − f(w) ≤ Γ
2
2βT + β 2L 2. (6) Taking
β = Γ/(L √ T ), we can obtain the convergence rate in the theorem. The detailed proof of Theorem 1 and explanations are provided in Appendix A.1.
According to Theorem 1, the FeDEC parameter aggregation algorithm is guaranteed to 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 is related to the optimization problem parameters such as lipschitz constant L, and diameter of domain Γ.

heading: 5 PERFORMANCE EVALUATION
text: In this section, we evaluate the performance of the proposed FeDEC method for federated learning.

heading: 5.1 EXPERIMENTAL SETUP
text: Implementation. We implement the proposed FeDEC parameter aggregation approach and the considered baselines in PyTorch (Paszke et al. (2019)). We train the models in a 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 β = 0.01. We conduct experiments on a GPU-equipped personal computer (CPU: Inter Core i78700 3.2GHz, GPU: Nvidia GeForce 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: ResNet18 (He et al. (2016)), LeNet (Lecun et al. (1998)), DenseNet121 (Huang et al. (2017)), MobileNetV2 (Sandler et al. (2018)), and a 4-layer CNN (every CNN layer is followed by a BN layer). The detailed structure of the CNN models are provided in Appendix A.3.
We use three real world datasets: MNIST (LeCun et al. (2010)), Fashion-MNIST (Xiao et al. (2017)), and CIFAR-10 (Krizhevsky (2009)). MNIST is a dataset for hand written digits classification with 60000 samples and each example is a 28 × 28 greyscale image. Fashion-MNIST is a 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 use 80% of the data for training and amalgamate the remaining data into a global test set.
We form non-IID local datasets as follows. Assume there are C classes of samples in a dataset. Each client draw samples form the dataset with probability pr(x) = {
η ∈ [0, 1], if x ∈ classj , N (0.5, 1), otherwise.
It means that the client draw samples from a particular class j with a fixed probability η, 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.

heading: 5.2 PERFORMANCE COMPARISON
text: We compare the performance of FeDEC with 5 state-of-the-art methods: FedAvg (McMahan et al. (2017)), RNN based aggregator (Ji et al. (2019)), FedProx (Li et al. (2020a)), q-FedSGD (Li et al. (2020b)), and FedMA (Wang et al. (2020)). The results are analyzed as follows.
Convergence: In this experiment we study the convergence of all baselines and our algorithm by showing the total communication epochs versus train loss. Fig. 4 shows the result of ResNet18 on CIFAR-10. It is shown that the loss of all algorithms tend to 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 in Appendix A.4.
Training Efficiency: In this experiment we study the test accuracy versus time during training of a CNN model with federated learning. Fig.5 shown the results of training ResNet18 on CIFAR-10. It 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.5. It suggests that FeDEC trains much faster than the baseline algorithms and it can reach high accuracy in a short time period.
Parameter Bias: In this experiment we study the parameter bias of federated learning algorithms. Fig. 6 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 qFedSGD have exceptional high parameter bias, while FeDEC has significantly lower KL-divergence compared to the baselines for different CNN models on 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 1. As shown in the table, the central method yields the highest accuracy. In comparison of different federated learning methods, FeDEC significantly outperforms the other algorithms in global model accuracy. It performs better than the state-of-the-art method FedMA with 2.87%, 3.17%, 2.58%, and 3.09% accuracy improvement in ResNet18, DenseNet121, MobileNetV2, and 4-L CNN respectively for CIFAR-10, 1.09% improvement 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.
Hyperparameter Analysis: We further analyze the influence of two hyperparameters in federated learning: the number of clients involved and the heterogeneity of local datasets.
Fig. 7 compares the test accuracy of the global model for different 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 decreases 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.
In the experiment, the heterogeneity of local datasets is represented by η, the probability that a client tends to sample from a particular class. The more η approaches to 1, the more heterogeneous the local datasets are. Fig. 8 shows the test accuracy under different level of heterogeneity. As η 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. It verifies the effectiveness of the proposed probabilistic-weighted gradient aggregation approach under non-IID conditions.

heading: 6 CONCLUSION
text: Gradient aggregation played an important role in federated learning to form a global model. To address the problem of data and parameter bias in federated learning for non-IID dataset, we proposed 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 guaranteed. Extensive experiments showed that FeDEC significantly outperforms the state-of-the-arts on a variety of performance metrics.

heading: A APPENDIX
text: A.1 PROOF OF CONVERGENCE GUARANTEE (THEOREM 1 IN SECTION 4)
We provide the detailed proof of Theorem 1 in 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:
WtNN = W t−1 NN − β K∑ k=1 πtkg t k, (3)
where β is the learning rate for parameter update. And for batch normalization update:
Wtmean = K∑ k=1 πtkW t,k mean (4)
Wtvar = 1
||X|| − n K∑ k=1 (||xk|| − 1)πtkWt,kvar (5)
We restate the theorem in section 4 in the following:
Assumption 1 (Unbiased Gradient): We assume that the stochastic gradients gti is an unbiased estimator of the true gradient ∇f(wti), i.e., E[gti ] = ∇f(wti), where f(·) is any convex objective function and wti is its variables.
Assumption 2 (Gradient Convex Set): We assume that gradient set G is a convex set, where all gradients g1,g2, . . . ,gK are in G, and any g = ∑K i=1 λigi (∀λi > 0 and ∑K i=1 λi = 1) is in G.
Lemma 1 (L-Lipschitz Continuity): For a function f(·) is Lipschitz continuous if there exists a positive real constant L such that, for all real x1 and x2:
|f(x1)− f(x2)| ≤ L|x1 − x2|.
Lemma 2 (Jensen’s Inequality): If f(w) is a convex function on W , and E[f(w)] and f(E[w]) are finite, then:
E[f(w)] ≥ f(E[w])). Definition 1 (Projection Operation): Assume w∗ is an intermediate result of optimization, we
define a project operator ∏
W(w∗) to project w∗ to the domain W , which is computed by:∏ W (w∗) = arg min w∈W ||w −w∗||.
Definition 2 (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 Γ: for every w1,w2 ∈ W: ||w1 −w2|| ≤ Γ. Theorem 1 (Guaranteed Convergence Rate): If a convex function f(w) is L-Lipschitz continuous function, then ||∇f(w)|| ≤ L. Let Γ be the diameter of domain. Applying equations (3)(4)(5) for gradients aggregation, we have the following convergence rate for the proposed FeDEC algorithm:
f(w̄T )− min w∈W
f(w) ≤ O( Γ 2
2βT +
β 2 L2), (11)
where w̄T is the average result of w for total training epoch T , β is the learning rate in equation-(3), and T is the total training epoch. If we let β = Γ
L √ T , the convergence rate is O( 1√ T ).
Proof: To simplify the analysis, we consider fixed learning rate β. The proof includes the following steps:
(1) According to the definition of convex function,
f(wt)− f(w) ≤ ∇f(wt)(wt −w).
(2) We define G(w) = ∇fT (wt)(wt−w), and gt = ∑K
k=1 πkg t k. The intermediate result of f(w)
in update time t+ 1 is denoted by wt+1∗ . With assumption 1 and assumption 2, we have:
G(w) = 1
β (wt −wt+1∗ )(wt −w)
= 1
β ((wt)2 −wtw −wt+1∗ wt +wt+1∗ w)
= 1
2β ((wt)2 −wtw +w2 − (wt+1∗ )2 + 2wwt+1∗ −w2 + (wt)2 − 2wtwt+1∗ + (wt+1∗ )2)
= 1
2β (||wt −w||2 − ||wt+1∗ −w||2 + ||wt −wt+1∗ ||2)
= 1
2β (||wt −w||2 − ||wt+1∗ −w||2) +
β 2 ||∇f(wt)||2
= 1
2β (||wt −w||2 − ||wt+1∗ −w||2) +
β 2 ||gt||2.
So we have:
f(wt)− f(w) ≤ 1 2β (||wt −w||2 − ||wt+1∗ −w||2) + β 2 ||gt||2.
(3) We project wt+1∗ to w t+1. With definition 1 and using non-expandable property of projection operation of convex set, we have:
G1(w) ≤ 1 2β (||wt −w||2 − ||wt+1 −w||2) + β 2 ||gt||2
Due to L-Lipschitz Continuity (Lemma 1), we have:
G(w) ≤ 1 2β (||wt −w||2 − ||wt+1 −w||2) + β 2 L2
So we have: f(wt)− f(w) ≤ 1
2β (||wt −w||2 − ||wt+1 −w||2) + β 2 L2
(4) According to definition 2, summing up all w from t = 1 to T , we have: T∑
t=1
f(wt)− Tf(w) ≤ 1 2β (||w1 −w||2 − ||wt+1 −w||2) + β 2 L2T
≤ 1 2β |w1 −w||2 + β 2 L2T
≤ 1 2β Γ2 + β 2 L2T.
(5) According to Jensen’s Inequality (Lemma 2), we have:
f(w̄T )− f(w) = f( 1 T T∑ t=1 wt)− f(w)
≤ 1 T T∑ t=1 f(wt)− f(w)
≤ Γ 2
2βT +
β 2 L2.
We can get the result:
f(w̄T )− min w∈W
f(w) ≤ O( Γ 2
2βT +
β 2 L2),
(6) Taking β = Γ/(L √ T ), the right part of the above equation becomes
Γ2
2βT +
β 2 L2 =
Γ2L √ T
2ΓT +
Γ
2L √ T L2 = ΓL√ T .
Therefore we can obtain the simplified expression of the convergence bound O( 1√ T ).
A.2 EXPLANATION OF UNBIAS PARAMETER AGGREGATION IN SECTION 3.2.2
We compute the expectation of the aggregated parameters Wtmean and W t var in Section 3.2.2 as follows.
E[Wtmean] = E
[ K∑
k=1
πkW t,k mean
]
= K∑ k=1 πkE [ Wt,kmean ]
E[Wtvar] = E
[ 1
||X|| −K K∑ k=1 (||xk|| − 1)πkWt,kvar
]
= 1
||X|| −K E
[ K∑
k=1
(||xk|| − 1)πkWt,kvar
]
= 1
||X|| −K K∑ k=1 (||xk|| − 1)πkE [ Wt,kvar ] According to the above equations, if the parameters of the local models Wt,kmean and W t,k var are
unbias, then the aggregated model parameters are unbias as well.
A.3 STRUCTURE OF THE NEURAL NETWORK MODELS IN SECTION 5
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 3. We adopt a slim ResNet18 as shown in Table 4, where “Conv2d” is convolution layer, “BatchNorm2d” is batch normalization layer, and “Linear” is fully-connected layer. We can observe that every convolution layer is followed by a batch normalization (BN) layer. For all models, we use ReLU layer after every Conv2d layer. The structure of DenseNet1212 and MobileNetV23 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 a device.
A.4 CONVERGENCE OF FEDERATED LEARNING ALGORITHMS FOR DIFFERENT MODELS ON DIFFERENT DATASETS
A.5 TRAINING EFFICIENCY OF FEDERATED LEARNING ALGORITHMS FOR DIFFERENT MODELS ON DIFFERENT DATASETS
Algorithm 1: FeDEC Aggregation 1 Server 2 for t= 1 to T do 3 Transmit Wt to all clients. 4 Receive gt, µ̂, σ̂ from all clients. 5 Inference πtk with 9 based on µ̂, σ̂ from W t mean,W t var. 6 Update Wt−1NN with 3 to get W t NN . 7 Aggregate Wtmean with 4 and W t var with 5. 8 Combine WtNN , W t mean and W t var to new model W
t. 9 Client
10 for t= 1 to T do 11 Receive server model Wt. 12 Train local model Wt based on local dataset and get local gradients gt. 13 Transmit gt and µ̂, σ̂ to server. 14 Stop training until received server model Wt.

Review summaries:
The 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.

This 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).