FedBN: Federated Learning on Non-IID Features via Local Batch Normalization

1 INTRODUCTION . Federated learning ( FL ) , has gained popularity for various applications involving learning from distributed data . In FL , a cloud server ( the “ server ” ) can communicate with distributed data sources ( the “ clients ” ) , while the clients hold data separately . A major challenge in FL is the training data statistical heterogeneity among the clients ( Kairouz et al. , 2019 ; Li et al. , 2020b ) . It has been shown that standard federated methods such as FedAvg ( McMahan et al. , 2017 ) which are not designed particularly taking care of non-iid data significantly suffer from performance degradation or even diverge if deployed over non-iid samples ( Karimireddy et al. , 2019 ; Li et al. , 2018 ; 2020a ) . Recent studies have attempted to address the problem of FL on non-iid data . Most variants of FedAvg primarily tackle the issues of stability , client drift and heterogeneous label distribution over clients ( Li et al. , 2020b ; Karimireddy et al. , 2019 ; Zhao et al. , 2018 ) . Instead , we focus on the shift ∗Corresponding author . ro r E in the feature space , which has not yet been explored in the literature . Specifically , we consider that local data deviates in terms of the distribution in feature space , and identify this scenario as feature shift . This type of non-iid data is a critical problem in many real-world scenarios , typically in cases where the local devices are responisble for a heterogeneity in the feature distributions . For example in cancer diagnosis tasks , medical radiology images collected in different hospitals have uniformly distributed labels ( i.e. , the cancer types treated are quite similar across the hospitals ) . However , the image appearance can vary a lot due to different imaging machines and protocols used in hospitals , e.g. , different intensity and contrast . In this example , each hospital is a client and hospitals aim to collaboratively train a cancer detection model without sharing privacy-sensitive data . Tackling non-iid data with feature shift has been explored in classical centralized training in the context of domain adaptation . Here , an effective approach in practice is utilizing Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) : recent work has proposed BN as a tool to mitigate domain shifts in domain adaptation tasks with promising results achieved ( Li et al. , 2016 ; Liu et al. , 2020 ; Chang et al. , 2019 ) . Inspired by this , this paper proposes to apply BN for feature shift FL . To illustrate the idea , we present a toy example that illustrates how BN may help harmonizing local feature distributions . Observation of BN in a FL Toy Example : We consider a simple non-convex learning problem : we generate data x , y ∈ R with y = cos ( wtruex ) + , where x ∈ R is drawn iid from Gaussian distribution and is zero-mean Gaussian noise and consider models of the form fw ( x ) = cos ( wx ) with model parameter w ∈ R. Local data deviates in the variance of x . First , we illustrate that local batch normalization harmonizes local data distributions . We consider a simplified form of BN that normalizes the input by scaling it with γ , i.e. , the local empirical standard deviation , and a setting with 2 clients . As Fig . 1 shows , the local squared loss is very different between the two clients . Thus , averaging the model does not lead to a good model . However when applying local BN , the local training error surfaces become similar and averaging the models can be beneficial . To further illustrate the impact of BN , we plot the error surface for one client with respect to both model parameter w ∈ R and BN parameter γ ∈ R in Fig . 2 . The figure shows that for an optimal weight w∗1 , changing γ deteriorates the model quality . Similarly , for a given optimal BN parameter γ∗1 , changing w deteriorates the quality . In particular , the average model w = ( w ∗ 1 + w ∗ 2 ) /2 and average BN parameters γ = ( γ∗1 + γ ∗ 2 ) /2 has a high generalization error . At the same time , the average model w with local BN parameter γ∗1 performs very well . Motivated by the above insight and observation , this paper proposes a novel federated learning method , called FedBN , for addressing non-iid training data which keeps the client BN layers updated locally , without communicating , and aggregating them at the server . In practice , we can simply update the non-BN layers using FedAvg , without modifying any optimization or aggregation scheme . This approach has zero parameters to tune , requires minimal additional computational resources , and can be easily applied to arbitrary neural network architectures with BN layers in FL . Besides the benefit shown in the toy example , we also show the benefits in accelerating convergence by theoretically analyzing the convergence of FedBN in the over-parameterized regime . In addition , we have conducted extensive experiments on a benchmark and three real-world datasets . Compared to classical FedAvg , as well as the state-of-the-art for non-iid data ( FedProx ) , our novel method , FedBN , demonstrates significant practical improvements on the extensive experiments . 2 RELATED WORK . Techniques for Non-IID Challenges in Federated Learning : The widely known aggregation strategy in FL , FedAvg ( McMahan et al. , 2017 ) , often suffers when data is heterogeneous over local client . Empirical work addressing non-iid issues , mainly focus on label distribution skew , where a non-iid dataset is formed by partitioning a “ flat ” existing dataset based on the labels . FedProx ( Li et al. , 2020b ) , a recent framework tackled the heterogeneity by allowing partial information aggregation and adding a proximal term to FedAvg . Zhao et al . ( 2018 ) assumed a subset of the data is globally shared between all the clients , hence generalizes to the problem at hand . FedMA ( Wang et al. , 2020 ) proposed an aggregation strategy for non-iid data partition that shares global model in a layer-wise manner . However , so far there are only limited attempts considering non-iid induced from feature shift , which is common in medical data collecting from different equipment and natural image collected in various noisy environment . Very recently , FedRobust ( Reisizadeh et al. , 2020 ) assumes data follows an affine distribution shift and tackles this problem by learning the affine transformation . This hampers the generalization when we can not estimate the explicit affine transformation . Concurrently to our work , SiloBN Andreux et al . ( 2020 ) empirically shows that local clients keeping some untrainable BN parameters could improve robustness to data heterogeneity , but provides no theoretical analysis of the approach . FedBN instead keeps all BN parameters strictly local . Recently , an orthogonal approach to the non-iid problem has been proposed that focuses on improving the optimization mechanism ( Reddi et al. , 2020 ; Zhang et al. , 2020 ) . Batch Normalization in Deep Neural Networks : Batch Normalization ( Ioffe & Szegedy , 2015 ) is an indispensable component in many deep neural networks and has shown its success in neural network training . Relevant literature has uncovered a number of benefits given by batch normalization . Santurkar et al . ( 2018 ) showed that BN makes the optimization landscape significantly smoother . Luo et al . ( 2018 ) investigated an explicit regularization form of BN such that improving the robustness of optimization . Morcos et al . ( 2018 ) suggested that BN implicitly discourages single direction reliance , thus improving model generalizability . Li et al . ( 2018 ) took advantage of BN for tackling the domain adaptation problem . However , what a role BN is playing in the scope of federated learning , especially for non-iid training , still remains unexplored to date . 3 PRELIMINARY . Non-IID Data in Federated Learning : We introduce the concept of feature shift in federated learning as a novel category of client ’ s non-iid data distribution . So far , the categories of non-iid data considered according to Kairouz et al . ( 2019 ) ; Hsieh et al . ( 2019 ) can be described by the joint probability between features x and labels y on each client . We can rewrite Pi ( x , y ) as Pi ( y|x ) Pi ( x ) and Pi ( x|y ) Pi ( y ) . We define feature shift as the case that covers : 1 ) covariate shift : the marginal distributions Pi ( x ) varies across clients , even if Pi ( y|x ) is the same for all client ; and 2 ) concept shift : the conditional distribution Pi ( x|y ) varies across clients and P ( y ) is the same . Federated Averaging ( FedAvg ) : We establish our algorithm on FedAvg introduced by McMahan et al . ( 2017 ) which is the most popular existing and easiest to implement federated learning strategy , where clients collaboratively send updates of locally trained models to a global server . Each client runs a local copy of the global model on its local data . The global model ’ s weights are then updated with an average of local clients ’ updates and deployed back to the clients . This builds upon previous distributed learning work by not only supplying local models but also performing training locally on each device . Hence FedAvg potentially empowers clients ( especially clients with small dataset ) to collaboratively learn a shared prediction model while keeping all training data locally . Although FedAvg has shown successes in classical Federated Learning tasks , it suffers from slow convergence and low accuracy in most non-iid contents ( Li et al. , 2020b ; 2019 ) . 4 FEDERATED AVERAGING WITH LOCAL BATCH NORMALIZATION . 4.1 PROPOSED METHOD - FEDBN . We propose an efficient and effective learning strategy denoted FedBN . Similar to FedAvg , FedBN performs local updates and averages local models . However , FedBN assumes local models have BN layers and excludes their parameters from the averaging step . We present the full algorithm in Appendix C. This simple modification results in significant empirical improvements in non-iid settings . We provide an explanation for these improvements in a simplified scenario , in which we show that FedBN improves the convergence rate under feature shift . 4.2 PROBLEM SETUP . We assume N ∈ N clients to jointly train for T ∈ N epochs and to communicate after E ∈ N local iterations . Thus , the system has T/E communication rounds over the T epochs . For simplicity , we assume all clients to have M ∈ N training examples ( a difference in training examples can be account for by weighted averaging ( McMahan et al. , 2017 ) ) for a regression task , i.e. , each client i ∈ [ N ] ( [ N ] = { 1 , . . . , N } ) has training examples { ( xij , yij ) ∈ Rd×R : j ∈ [ M ] } . Furthermore , we assume a two-layer neural network with ReLU activations trained by gradient descent . Let vk ∈ Rd denote the parameters of the first layer , where k ∈ [ m ] and m is the width of the hidden layer . Let ‖ v ‖S , √ v > Sv denote the induced vector norm for a positive definite matrix S. We consider a non-iid setting in FL where local feature distributions differ—not label distribution , as considered , e.g. , in McMahan et al . ( 2017 ) ; Li et al . ( 2019 ) . To be more precise , we make the following assumption . Assumption 4.1 ( Data Distribution ) . For each client i ∈ [ N ] the inputs xij are centered ( Exi = 0 ) with covariance matrix Si = Exixi > , where Si is independent from the label y and may differ for each i ∈ [ N ] e.g. , Si are not all identity matrices , and for each index pair p 6= q , xp 6= κ ·xq for all κ ∈ R \ { 0 } . With Assumption 4.1 , the normalization of the first layer for client i is v > k x i ‖vk‖Si . FedBN with client- specified BN parameters trains a model f∗ : Rd → R parameterized by ( V , γ , c ) ∈ Rm×d × Rm×N × Rm , i.e. , f∗ ( x ; V , γ , c ) = 1√ m m∑ k=1 ck N∑ i=1 σ ( γk , i · v > k x ‖ vk ‖Si ) · 1 { x ∈ client i } , ( 1 ) where γ is the scaling parameter of BN and σ ( s ) = max { s , 0 } is the ReLU activation function , c is the top layer parameters of the network . Here , we omit learning the shift parameter of BN 1 . FedAvg instead trains a function f : Rd → R which is a special case of Eq . 1 with γk , i = γk for ∀i ∈ [ N ] . We take a random initialization of the parameters ( Salimans & Kingma , 2016 ) in our analysis : vk ( 0 ) ∼ N ( 0 , α2I ) , ck ∼ U { −1 , 1 } , and γk = γk , i = ‖vk ( 0 ) ‖2 /α , ( 2 ) where α2 controls the magnitude of vk at initialization . The initialization of the BN parameters γk and γk , i are independent of α . The parameters of the network f∗ ( x ; V , γ , c ) are obtained by minimizing the empirical risk with respect to the squared loss using gradient descent : L ( f∗ ) = 1 NM N∑ i=1 M∑ j=1 ( f∗ ( xij ) − yij ) 2 . ( 3 )

Summary: This paper develops a modified version of FedAvg by local batch normalization that is tailored for federated learning with non-i.i.d. data. Different from most of existing work that consider the unbalanced labels, this paper uses unbalanced features to motivate the non-i.i.d. federated settings. Specifically, the unbalanced features are captured by the difference in local covariances.