Federated Learning with Decoupled Probabilistic-Weighted Gradient Aggregation

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 . 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 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 .

Summary: 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).