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 work proposes an extremely simple approach to a well-know problem within Federated Learning: batch normalisation. Indeed, FL usually imply an averaging of different model parameters trained on different distributed devices. But what happen with the running statistics that some training methods have ? Such as batch norm and some optimisers ? Well, this work proposes to address the first issue by introducing a specific batch norm method named Local Batch Normalization (FedBN). Standard BN statistics cannot simply be aggregated due to strong variations potentially occurring when training data across the clients aren't IID distributed (i.e strong shift in the input representation). To alleviate this issue, the authors propose to exclude the BN parameters from the aggregation (i.e. each client has its own BN excluded from FL).